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

.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:
@@ -1650,6 +1673,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 +1694,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 +1769,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 +1866,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 +1956,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 +2017,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 +2095,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 +2178,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 +2201,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 +2279,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 +2354,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 +2437,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 +2529,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 +2711,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 +2855,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 +2894,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 +2939,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 +3034,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 +3142,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 +3209,7 @@ workflows:
                 - main
                 - /litellm_.*/
   build_and_test:
+    unless: << pipeline.parameters.run_migration_tests >>
     jobs:
       - using_litellm_on_windows:
           filters: &main_branches
@@ -3029,6 +3217,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/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/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

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

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.

enterprise/litellm_enterprise/enterprise_callbacks/send_emails/endpoints.py

@@ -60,6 +60,11 @@ async def _get_email_settings(prisma_client) -> Dict[str, bool]:
 
 async def _save_email_settings(prisma_client, settings: Dict[str, bool]):
     """Helper function to save email settings to general_settings in db"""
+    from litellm.proxy.proxy_server import proxy_config
+
+    proxy_config.reject_config_owned_writes(
+        section_name="general_settings", changed_keys={"email_settings": settings}
+    )
     try:
         verbose_proxy_logger.debug(
             f"Saving email settings to general_settings: {settings}"
@@ -168,6 +173,8 @@ async def update_event_settings(
         await _save_email_settings(prisma_client, settings_dict)
 
         return {"message": "Email event settings updated successfully"}
+    except HTTPException:
+        raise
     except Exception as e:
         verbose_proxy_logger.exception(f"Error updating email settings: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))
@@ -197,6 +204,8 @@ async def reset_event_settings(
         await _save_email_settings(prisma_client, default_settings)
 
         return {"message": "Email event settings reset to defaults"}
+    except HTTPException:
+        raise
     except Exception as e:
         verbose_proxy_logger.exception(f"Error resetting email settings: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -20,6 +20,7 @@ from typing import (
 )
 from uuid import NAMESPACE_URL, uuid5
 
+import httpx
 from fastapi import HTTPException
 from pydantic import ValidationError
 
@@ -34,6 +35,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
 )
 from openai.types.file_deleted import FileDeleted
 
+from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
 from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
 from litellm.llms.base_llm.managed_resources.isolation import (
     build_list_page,
@@ -59,6 +61,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
     get_content_type_from_file_object,
     get_model_id_from_unified_batch_id,
     get_original_file_id,
+    is_litellm_executed_batch,
     map_raw_file_ids_to_unified,
     normalize_mime_type_for_provider,
     resolve_managed_output_file_model_name,
@@ -75,6 +78,7 @@ from litellm.types.llms.openai import (  # pyright: ignore[reportAttributeAccess
     CreateFileRequest,
     FileListPage,
     FileObject,
+    HttpxBinaryResponseContent,
     OpenAIFileObject,
     ResponsesAPIResponse,
 )
@@ -86,10 +90,6 @@ from litellm.types.utils import (
     SpecialEnums,
 )
 
-if TYPE_CHECKING:
-    from litellm.types.llms.openai import HttpxBinaryResponseContent
-
-
 if TYPE_CHECKING:
     from opentelemetry.trace import Span as _Span
     from prisma.models import (
@@ -204,6 +204,19 @@ def _managed_object_table(prisma_client: PrismaClient) -> _ManagedObjectTableAct
     return prisma_client.db.litellm_managedobjecttable
 
 
+def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
+    hidden_params: Final = cast(  # cast-ok: _hidden_params is an untyped attribute the upload path sets
+        "Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
+    )
+    return MappingProxyType(
+        {
+            key: value
+            for key in ("storage_backend", "storage_url")
+            if isinstance(value := hidden_params.get(key), str)
+        }
+    )
+
+
 class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
     # Class variables or attributes
     def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient):
@@ -226,6 +239,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         user_api_key_dict: UserAPIKeyAuth,
     ) -> None:
         verbose_logger.info(f"Storing LiteLLM Managed File object with id={file_id} in cache")
+        storage_metadata: Final = _storage_metadata_of(file_object)
         if file_object is not None:
             litellm_managed_file_object = LiteLLM_ManagedFileTable(
                 unified_file_id=file_id,
@@ -235,6 +249,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 created_by=resolve_resource_owner_id(user_api_key_dict),
                 team_id=user_api_key_dict.team_id,
                 updated_by=user_api_key_dict.user_id,
+                storage_backend=storage_metadata.get("storage_backend"),
+                storage_url=storage_metadata.get("storage_url"),
             )
             await self.internal_usage_cache.async_set_cache(
                 key=file_id,
@@ -262,14 +278,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             file_object_json = file_object.model_dump_json()
             db_data["file_object"] = file_object_json
             update_data["file_object"] = file_object_json
-            # Extract storage metadata from hidden params if present
-            hidden_params = getattr(file_object, "_hidden_params", {}) or {}
-            if "storage_backend" in hidden_params:
-                db_data["storage_backend"] = hidden_params["storage_backend"]
-                update_data["storage_backend"] = hidden_params["storage_backend"]
-            if "storage_url" in hidden_params:
-                db_data["storage_url"] = hidden_params["storage_url"]
-                update_data["storage_url"] = hidden_params["storage_url"]
+            db_data.update(storage_metadata)
+            update_data.update(storage_metadata)
 
             verbose_logger.debug(
                 f"Storage metadata: storage_backend={db_data.get('storage_backend')}, "
@@ -314,6 +324,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         request_tags: Sequence[str] | None = None,
         persist_attribution: bool = False,
         create_if_missing: bool = True,
+        batch_processed: bool = False,
     ) -> None:
         """Persist a managed object row, caching it and upserting it in the DB.
 
@@ -328,6 +339,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         row absent from the table is left absent rather than created with the
         observer as its creator, because created_by and team_id are written from
         whoever calls the create branch.
+
+        batch_processed is set by callers that have already billed the batch
+        themselves, so CheckBatchCost skips the row instead of billing it twice.
+        It is written only in the upsert create branch.
         """
         verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
         litellm_managed_object = LiteLLM_ManagedObjectTable(
@@ -379,6 +394,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                     "updated_by": user_api_key_dict.user_id,
                     "status": file_object.status,
                     **attribution_columns,
+                    "batch_processed": batch_processed,
                 },
                 "update": update_columns,
             },
@@ -1343,6 +1359,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
     ) -> LLMResponseTypes:
         if isinstance(response, LiteLLMBatch):
+            decoded_batch_id: Final = _is_base64_encoded_unified_file_id(response.id)
+            if decoded_batch_id and is_litellm_executed_batch(decoded_batch_id):
+                return response
             ## Check if unified_file_id is in the response
             unified_file_id = response._hidden_params.get("unified_file_id")  # managed file id
             unified_batch_id = response._hidden_params.get("unified_batch_id")  # managed batch id
@@ -1794,24 +1813,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         # Check if file deletion should be blocked due to batch references
         await self._check_file_deletion_allowed(file_id)
 
-        # file_id = convert_b64_uid_to_unified_uid(file_id)
-        model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span)
-
-        specific_model_file_id_mapping = model_file_id_mapping.get(file_id)
-        if specific_model_file_id_mapping:
-            # Remove conflicting keys from data to avoid duplicate keyword arguments
-            filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")}
-            for model_id, model_file_id in specific_model_file_id_mapping.items():
-                credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
-                delete_data = {
-                    **{k: v for k, v in filtered_data.items() if k != "_litellm_internal_model_credentials"},
-                    **(
-                        {"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
-                        if credentials is not None
-                        else {}
-                    ),
-                }
-                await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
+        managed_file: Final = await self.get_unified_file_id(file_id, litellm_parent_otel_span)
+        if managed_file is not None and managed_file.storage_backend and managed_file.storage_url:
+            await self._delete_storage_backend_content(managed_file.storage_backend, managed_file.storage_url)
+        else:
+            await self._delete_provider_files(file_id, litellm_parent_otel_span, llm_router, data)
 
         await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
 
@@ -1820,16 +1826,53 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             prom_logger.record_managed_file_deleted(result="success")
         return FileDeleted(id=file_id, object="file", deleted=True)
 
+    async def _delete_storage_backend_content(self, storage_backend_name: str, storage_url: str) -> None:
+        try:
+            storage_backend: Final = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client)
+        except ValueError as e:
+            raise HTTPException(status_code=400, detail=f"Cannot delete the stored file content: {e}") from e
+        await storage_backend.delete_file(storage_url)
+
+    async def _delete_provider_files(
+        self,
+        file_id: str,
+        litellm_parent_otel_span: Span | None,
+        llm_router: Router,
+        data: Mapping[str, object],
+    ) -> None:
+        model_file_id_mapping: Final = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span)
+        specific_model_file_id_mapping: Final = model_file_id_mapping.get(file_id)
+        if not specific_model_file_id_mapping:
+            return
+        filtered_data: Final = {
+            k: v for k, v in data.items() if k not in ("model", "file_id", "_litellm_internal_model_credentials")
+        }
+        for model_id, model_file_id in specific_model_file_id_mapping.items():
+            credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
+            delete_data = {
+                **filtered_data,
+                **(
+                    {"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
+                    if credentials is not None
+                    else {}
+                ),
+            }
+            await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
+
     async def afile_content(
         self,
         file_id: str,
         litellm_parent_otel_span: Optional[Span],
         llm_router: Router,
         **data: Dict,
-    ) -> "HttpxBinaryResponseContent":
+    ) -> HttpxBinaryResponseContent:
         """
         Get the content of a file from first model that has it
         """
+        managed_file: Final = await self.get_unified_file_id(file_id, litellm_parent_otel_span)
+        if managed_file is not None and managed_file.storage_backend and managed_file.storage_url:
+            return await self._storage_backend_content(managed_file.storage_backend, managed_file.storage_url)
+
         model_file_id_mapping = data.pop("model_file_id_mapping", None)
         model_file_id_mapping = model_file_id_mapping or await self.get_model_file_id_mapping(
             [file_id], litellm_parent_otel_span
@@ -1859,6 +1902,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         else:
             raise Exception(f"LiteLLM Managed File object with id={file_id} not found")
 
+    async def _storage_backend_content(self, storage_backend_name: str, storage_url: str) -> HttpxBinaryResponseContent:
+        storage_backend: Final = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client)
+        content: Final = await storage_backend.download_file(storage_url)
+        return HttpxBinaryResponseContent(response=httpx.Response(status_code=httpx.codes.OK, content=content))
+
     async def _convert_storage_files_to_base64(
         self,
         messages: List[AllMessageValues],
@@ -1889,16 +1937,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
 
             # File is stored in a storage backend, download and convert to base64
             try:
-                from litellm.llms.base_llm.files.storage_backend_factory import (
-                    get_storage_backend,
-                )
-
                 storage_backend_name = db_file.storage_backend
                 storage_url = db_file.storage_url
 
                 # Get storage backend (uses same env vars as callback)
                 try:
-                    storage_backend = get_storage_backend(storage_backend_name)
+                    storage_backend = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client)
                 except ValueError as e:
                     verbose_logger.warning(
                         f"Storage backend '{storage_backend_name}' error for file {file_id}: {str(e)}"

enterprise/pyproject.toml

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

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/migrations/20260915000000_add_autorouter_savings_estimate_coverage/migration.sql

@@ -0,0 +1,5 @@
+ALTER TABLE "LiteLLM_AutoRouterSession"
+ADD COLUMN IF NOT EXISTS "savings_estimated_turns" INTEGER NOT NULL DEFAULT 0,
+ADD COLUMN IF NOT EXISTS "savings_estimated_actual_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
+ADD COLUMN IF NOT EXISTS "savings_estimated_saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
+ADD COLUMN IF NOT EXISTS "savings_estimated_baseline_models" JSONB NOT NULL DEFAULT '{}';

litellm-proxy-extras/litellm_proxy_extras/migrations/20260915010000_add_autorouter_baseline_state/migration.sql

@@ -0,0 +1,36 @@
+CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterBaselineComparison" (
+    "scope" TEXT PRIMARY KEY,
+    "api_key" TEXT NOT NULL,
+    "session_id" TEXT NOT NULL,
+    "router_name" TEXT NOT NULL,
+    "initial_equivalent" BOOLEAN NOT NULL,
+    "revision" BIGINT NOT NULL DEFAULT 0,
+    "published_revision" BIGINT NOT NULL DEFAULT 0,
+    "history" TEXT,
+    "attempted_at" TIMESTAMP(3),
+    "retired" BOOLEAN NOT NULL DEFAULT FALSE,
+    "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP
+);
+
+CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_scope"
+    ON "LiteLLM_AutoRouterBaselineComparison" ("api_key", "session_id", "router_name");
+CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_updated"
+    ON "LiteLLM_AutoRouterBaselineComparison" ("updated_at");
+CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_dirty"
+    ON "LiteLLM_AutoRouterBaselineComparison" ("attempted_at", "updated_at", "scope")
+    WHERE NOT "retired" AND "revision" <> "published_revision";
+
+CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterBaselineObservation" (
+    "request_id" TEXT PRIMARY KEY,
+    "scope" TEXT NOT NULL,
+    "started_at" DOUBLE PRECISION NOT NULL,
+    "revision" BIGINT NOT NULL,
+    "data" TEXT NOT NULL,
+    "publication" TEXT,
+    "conflicted" BOOLEAN NOT NULL DEFAULT FALSE
+);
+
+CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_event_order"
+    ON "LiteLLM_AutoRouterBaselineObservation" ("scope", "started_at", "request_id");
+CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_event_revision"
+    ON "LiteLLM_AutoRouterBaselineObservation" ("scope", "revision", "started_at");

litellm-proxy-extras/litellm_proxy_extras/migrations/20260918000000_add_managed_file_content_table/migration.sql

@@ -0,0 +1,8 @@
+-- CreateTable
+CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedFileContentTable" (
+    "id" TEXT NOT NULL,
+    "content" BYTEA NOT NULL,
+    "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+
+    CONSTRAINT "LiteLLM_ManagedFileContentTable_pkey" PRIMARY KEY ("id")
+);

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/schema.prisma

@@ -1107,6 +1107,12 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
   @@index([team_id, created_at(sort: Desc)])
 }
 
+model LiteLLM_ManagedFileContentTable {
+  id         String   @id @default(uuid())
+  content    Bytes
+  created_at DateTime @default(now())
+}
+
 model LiteLLM_ManagedVectorStoreTable {
   id String @id @default(uuid())
   unified_resource_id String @unique // The base64 encoded unified vector store ID
@@ -1545,6 +1551,36 @@ model LiteLLM_AdaptiveRouterSession {
   @@index([last_activity_at], map: "idx_adaptive_router_session_activity")
 }
 
+model LiteLLM_AutoRouterBaselineComparison {
+  scope               String   @id
+  api_key             String
+  session_id          String
+  router_name         String
+  initial_equivalent  Boolean
+  revision            BigInt   @default(0)
+  published_revision  BigInt   @default(0)
+  history             String?
+  attempted_at        DateTime?
+  retired             Boolean  @default(false)
+  updated_at          DateTime @default(now())
+
+  @@index([api_key, session_id, router_name], map: "idx_autorouter_baseline_scope")
+  @@index([updated_at], map: "idx_autorouter_baseline_updated")
+}
+
+model LiteLLM_AutoRouterBaselineObservation {
+  request_id  String  @id
+  scope       String
+  started_at  Float
+  revision    BigInt
+  data        String
+  publication String?
+  conflicted  Boolean @default(false)
+
+  @@index([scope, started_at, request_id], map: "idx_autorouter_baseline_event_order")
+  @@index([scope, revision, started_at], map: "idx_autorouter_baseline_event_revision")
+}
+
 model LiteLLM_AutoRouterSession {
   api_key               String
   session_id            String
@@ -1571,6 +1607,10 @@ model LiteLLM_AutoRouterSession {
   total_tokens          BigInt   @default(0)
   spend                 Float    @default(0)
   saved_spend           Float    @default(0)
+  savings_estimated_turns Int    @default(0)
+  savings_estimated_actual_spend Float @default(0)
+  savings_estimated_saved_spend Float @default(0)
+  savings_estimated_baseline_models Json @default("{}")
   classifier_cost       Float    @default(0)
   classifier_cost_recorded_turns Int @default(0)
   tier_turns            Json     @default("{}")

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/pyproject.toml

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

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

@@ -11,12 +11,17 @@ repository = "https://github.com/BerriAI/litellm"
 [workspace.dependencies]
 litellm-core = { path = "crates/core" }
 litellm-host = { path = "crates/host" }
-litellm-callbacks-legacy = { path = "crates/callbacks-legacy" }
+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,17 +29,22 @@ 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"
 pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
 pythonize = "0.29.0"
 rand = "0.8"
-reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
+reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
 rstest = "0.26.1"
 rstest_reuse = "0.7.0"
 rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
@@ -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/clippy.toml

@@ -0,0 +1,10 @@
+# The Tokio runtime is reached only through `host-python/src/execution.rs`, whose fork gate
+# must see every entry. Going around it makes a fork-after-use hang instead of raising.
+disallowed-methods = [
+    { path = "pyo3_async_runtimes::tokio::get_runtime", reason = "use litellm_host_python::run_sync / run_sync_value" },
+    { path = "pyo3_async_runtimes::tokio::future_into_py", reason = "use litellm_host_python::run_async / run_async_value" },
+    { path = "pyo3_async_runtimes::tokio::future_into_py_with_locals", reason = "use litellm_host_python::run_async / run_async_value" },
+    { path = "pyo3_async_runtimes::tokio::local_future_into_py", reason = "use litellm_host_python::run_async / run_async_value" },
+    { path = "pyo3_async_runtimes::tokio::run", reason = "use litellm_host_python::run_sync / run_sync_value" },
+    { path = "pyo3_async_runtimes::tokio::run_until_complete", reason = "use litellm_host_python::run_sync / run_sync_value" },
+]

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/callbacks-legacy-python/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "litellm-callbacks-legacy"
+name = "litellm-callbacks-legacy-python"
 version = "0.1.0"
 edition.workspace = true
 license.workspace = true

litellm-rust/crates/callbacks-legacy-python/src/adapter.rs

@@ -19,9 +19,9 @@ use serde_json::Value;
 use crate::{
     DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger,
     deferred::{PendingLogging, PendingSuccess},
-    finalize, is_internal_call,
-    legacy_python::Streaming,
-    prepare, setup,
+    finalize, is_internal_call, prepare,
+    python::Streaming,
+    setup,
 };
 
 /// What the legacy contract needs to know about the route it is logging.

litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs

@@ -6,8 +6,8 @@ use litellm_host::event::{RequestContext, WireRequest};
 use litellm_host_python::to_py;
 use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict};
 
-use crate::legacy_python::{Logging, Wrapper};
 use crate::logger::PythonLogger;
+use crate::python::{Logging, Wrapper};
 
 pub trait LegacyCallbacks {
     /// `Logging.update_from_kwargs`: what the logger is told about the request it is

litellm-rust/crates/callbacks-legacy-python/src/lib.rs

@@ -13,9 +13,9 @@ mod adapter;
 mod call;
 mod callbacks;
 mod deferred;
-mod legacy_python;
 mod logger;
 mod preparation;
+mod python;
 #[cfg(test)]
 #[path = "../tests/support.rs"]
 mod test_support;

litellm-rust/crates/callbacks-legacy-python/src/logger.rs

@@ -5,7 +5,7 @@ use pyo3::{
     types::{PyDict, PyTuple},
 };
 
-use crate::legacy_python::{self, Wrapper};
+use crate::python::{self, Wrapper};
 
 /// The `Logging` instance one call fans out through.
 pub struct PythonLogger {
@@ -90,7 +90,7 @@ impl DeploymentHooks {
         kwargs: &Py<PyDict>,
         call_type: &str,
     ) -> PyResult<Py<PyAny>> {
-        legacy_python::DeploymentHooks::BeforeDeploymentCall
+        python::DeploymentHooks::BeforeDeploymentCall
             .call(py, (kwargs, call_type))
             .map(Bound::unbind)
     }
@@ -101,7 +101,7 @@ impl DeploymentHooks {
         response: &Option<Py<PyAny>>,
         call_type: &str,
     ) -> PyResult<Py<PyAny>> {
-        legacy_python::DeploymentHooks::AfterDeploymentSuccess
+        python::DeploymentHooks::AfterDeploymentSuccess
             .call(py, (kwargs, response, call_type))
             .map(Bound::unbind)
     }
@@ -112,7 +112,7 @@ impl DeploymentHooks {
         error: &Py<PyBaseException>,
         call_type: &str,
     ) -> PyResult<Py<PyAny>> {
-        legacy_python::DeploymentHooks::AfterDeploymentFailure
+        python::DeploymentHooks::AfterDeploymentFailure
             .call(py, (kwargs, error, call_type))
             .map(Bound::unbind)
     }

litellm-rust/crates/callbacks-legacy-python/src/preparation.rs

@@ -3,7 +3,7 @@ use pyo3::{
     types::{PyDict, PyList},
 };
 
-use crate::legacy_python::Wrapper;
+use crate::python::Wrapper;
 
 struct CredentialEntry<'py>(Bound<'py, PyAny>);
 

litellm-rust/crates/callbacks-legacy-python/src/python.rs

@@ -1,7 +1,7 @@
 use pyo3::prelude::*;
 use strum::{IntoStaticStr, VariantArray};
 
-const MODULE: &str = "litellm.rust_bridge.legacy_callbacks";
+const MODULE: &str = "litellm.rust_bridge.callbacks_legacy_python";
 
 /// Every litellm Python internal the native call still borrows, grouped by the subsystem it
 /// belongs to. Rust drives the call; these exist only so behaviour that Python owns today
@@ -9,7 +9,7 @@ const MODULE: &str = "litellm.rust_bridge.legacy_callbacks";
 /// A group is deleted once Rust owns that subsystem, so this enum only shrinks. Calling a
 /// user's own callback is not borrowing and does not belong here.
 ///
-/// `litellm/rust_bridge/legacy_callbacks.py` is the only Python module behind it, and
+/// `litellm/rust_bridge/callbacks_legacy_python.py` is the only Python module behind it, and
 /// `python_contract.json` pins each function's parameters on both sides.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub(crate) enum LegacyPython {

litellm-rust/crates/callbacks-legacy-python/tests/support.rs

@@ -5,12 +5,12 @@ use pyo3::types::{PyDict, PyTuple};
 
 use crate::{LegacyLogging, LegacySurface, PublicCall};
 
-/// The parameters of every `legacy_callbacks` function, as the real module declares them.
-/// `tests/test_litellm/rust_bridge/test_legacy_callbacks.py` pins this file to the Python
+/// The parameters of every `callbacks_legacy_python` function, as the real module declares them.
+/// `tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python
 /// signatures, and [`namespace`] binds every fake call against it.
 pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json");
 
-/// Stand-ins for `legacy_callbacks`, the only Python module the crate calls. Tests
+/// Stand-ins for `callbacks_legacy_python`, the only Python module the crate calls. Tests
 /// share one interpreter and run concurrently, so each fake is installed idempotently and
 /// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`).
 /// Every fake is bound against the contract first, so a call the real module would reject
@@ -23,10 +23,10 @@ import sys
 import traceback
 import types
 
-for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'):
+for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'):
     sys.modules.setdefault(name, types.ModuleType(name))
 
-legacy = sys.modules['litellm.rust_bridge.legacy_callbacks']
+legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python']
 CONTRACT = json.loads(python_contract)
 
 

litellm-rust/crates/host-python/src/execution.rs

@@ -4,6 +4,7 @@ use std::pin::Pin;
 use std::task::{Context, Poll, Waker};
 use std::time::Duration;
 
+use crate::fork_gate::{ForkGate, Refused, RuntimeAlreadyStarted};
 use crate::{Pythonized, panic_to_pyerr, release_gil};
 use futures_util::FutureExt;
 use pyo3::exceptions::PyRuntimeError;
@@ -12,6 +13,67 @@ use serde::Serialize;
 use tokio::runtime::{Handle, Runtime};
 use tokio::time::{self, MissedTickBehavior};
 
+pyo3::create_exception!(
+    _native,
+    ForkedAfterNativeRuntimeStarted,
+    PyRuntimeError,
+    "This process was forked after the native runtime started. Runtime threads do not survive fork(), so native routes cannot run here."
+);
+
+pyo3::create_exception!(
+    _native,
+    ProcessReservedForForking,
+    PyRuntimeError,
+    "This process was reserved for forking workers, so native routes cannot run here."
+);
+
+static FORK_GATE: ForkGate = ForkGate::new();
+
+/// Whether this process has started the Tokio runtime.
+pub fn runtime_started() -> bool {
+    FORK_GATE.started(std::process::id())
+}
+
+/// Declares that this process exists to fork workers, so it must never start the runtime.
+/// Fails if it already has. Workers are unaffected: the reservation is keyed by pid.
+pub fn reserve_process_for_forking() -> Result<(), RuntimeAlreadyStarted> {
+    FORK_GATE.reserve(std::process::id())
+}
+
+/// The only door to the Tokio runtime: every route reaches it through this module, which is
+/// what lets the gate speak for the whole extension. `clippy.toml` disallows going around it.
+fn enter_runtime() -> PyResult<()> {
+    FORK_GATE
+        .enter(std::process::id())
+        .map_err(|refused| match refused {
+            Refused::ReservedForForking => ProcessReservedForForking::new_err(
+                "this process is reserved for forking workers and cannot run native routes; \
+                 move the call into a worker, after the fork",
+            ),
+            Refused::ForkedAfterStart => ForkedAfterNativeRuntimeStarted::new_err(
+                "this process was forked after the native runtime started, and runtime threads \
+                 do not survive fork(); start workers with spawn or forkserver, or fork before \
+                 the first native call",
+            ),
+        })
+}
+
+#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
+fn runtime() -> PyResult<&'static Runtime> {
+    enter_runtime()?;
+    Ok(pyo3_async_runtimes::tokio::get_runtime())
+}
+
+#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
+fn future_into_py<F, T>(py: Python<'_>, future: F) -> PyResult<Bound<'_, PyAny>>
+where
+    F: Future<Output = PyResult<T>> + Send + 'static,
+    T: for<'py> IntoPyObject<'py> + Send + 'static,
+{
+    enter_runtime()?;
+    pyo3_async_runtimes::tokio::future_into_py(py, future)
+}
+
 pub fn run_sync<T, E, F>(
     py: Python<'_>,
     future: F,
@@ -22,12 +84,7 @@ where
     E: Send + 'static,
     F: Future<Output = Result<T, E>> + Send + 'static,
 {
-    run_sync_on(
-        py,
-        pyo3_async_runtimes::tokio::get_runtime(),
-        future,
-        map_error,
-    )
+    run_sync_on(py, runtime()?, future, map_error)
 }
 
 pub fn run_sync_value<T, F>(py: Python<'_>, future: F) -> PyResult<T>
@@ -35,7 +92,7 @@ where
     T: Send + 'static,
     F: Future<Output = PyResult<T>> + Send + 'static,
 {
-    run_sync_value_on(py, pyo3_async_runtimes::tokio::get_runtime(), future)
+    run_sync_value_on(py, runtime()?, future)
 }
 
 fn run_sync_value_on<T, F>(py: Python<'_>, runtime: &Runtime, future: F) -> PyResult<T>
@@ -83,7 +140,7 @@ where
     E: Send + 'static,
     F: Future<Output = Result<T, E>> + Send + 'static,
 {
-    pyo3_async_runtimes::tokio::future_into_py(py, async move {
+    future_into_py(py, async move {
         let result = catch_future_panic(future).await?;
         let result = map_core_result(result, map_error)?;
         Ok(Pythonized(result))
@@ -95,7 +152,7 @@ where
     T: for<'py> IntoPyObject<'py> + Send + 'static,
     F: Future<Output = PyResult<T>> + Send + 'static,
 {
-    pyo3_async_runtimes::tokio::future_into_py(py, async move { catch_future_panic(future).await? })
+    future_into_py(py, async move { catch_future_panic(future).await? })
 }
 
 pub fn poll_async_value<T, F>(py: Python<'_>, future: Pin<&mut F>) -> PyResult<Poll<T>>
@@ -103,8 +160,9 @@ where
     T: Send,
     F: Future<Output = PyResult<T>> + Send,
 {
+    let runtime = runtime()?;
     let result = release_gil(py, || {
-        let _runtime = pyo3_async_runtimes::tokio::get_runtime().enter();
+        let _runtime = runtime.enter();
         std::panic::catch_unwind(AssertUnwindSafe(|| {
             future.poll(&mut Context::from_waker(Waker::noop()))
         }))
@@ -286,27 +344,25 @@ mod tests {
     }
 
     #[pyfunction]
-    fn runtime_worker_count() -> usize {
-        pyo3_async_runtimes::tokio::get_runtime()
-            .metrics()
-            .num_workers()
+    fn runtime_worker_count() -> PyResult<usize> {
+        Ok(runtime()?.metrics().num_workers())
     }
 
     #[pyfunction]
-    fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
+    fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> PyResult<bool> {
         let completion_deadline = Instant::now() + Duration::from_secs(2);
         while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
             if Instant::now() >= completion_deadline {
-                return false;
+                return Ok(false);
             }
             thread::sleep(Duration::from_millis(1));
         }
 
         let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
-        pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
+        runtime()?.spawn(async move {
             let _ = heartbeat_tx.send(());
         });
-        heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
+        Ok(heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok())
     }
 
     fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
@@ -317,6 +373,16 @@ mod tests {
             .expect("result should convert")
     }
 
+    #[rstest]
+    fn reaching_the_runtime_marks_the_process_as_started(
+        #[from(initialized_python)] python: &InitializedPython,
+    ) {
+        python.attach(|py| {
+            run_sync_value(py, async { Ok(()) }).unwrap();
+            assert!(runtime_started());
+        });
+    }
+
     #[rstest]
     fn inline_poll_releases_gil_and_enters_runtime(
         #[from(initialized_python)] python: &InitializedPython,

litellm-rust/crates/host-python/src/fork_gate.rs

@@ -0,0 +1,139 @@
+use std::sync::atomic::{AtomicU32, Ordering};
+
+const UNSET: u32 = 0;
+
+/// Decides which process may use the Tokio runtime. Its worker threads do not survive
+/// `fork()`: a child forked after they started hangs on its first native call. The gate turns
+/// both halves of that hazard into errors, keyed by pid so a fork needs no hook to be seen:
+/// a process reserved for forking can never start the runtime, and a child of a process that
+/// did start it is refused instead of hanging.
+pub(crate) struct ForkGate {
+    runtime_pid: AtomicU32,
+    fork_only_pid: AtomicU32,
+}
+
+#[derive(Debug, PartialEq, Eq)]
+pub(crate) enum Refused {
+    ReservedForForking,
+    ForkedAfterStart,
+}
+
+#[derive(Debug, PartialEq, Eq)]
+pub struct RuntimeAlreadyStarted;
+
+impl ForkGate {
+    pub(crate) const fn new() -> Self {
+        Self {
+            runtime_pid: AtomicU32::new(UNSET),
+            fork_only_pid: AtomicU32::new(UNSET),
+        }
+    }
+
+    /// Claims the runtime for `pid`. Claim first, then look for a reservation: `reserve` does
+    /// the mirror image, so when the two race at least one of them sees the other.
+    pub(crate) fn enter(&self, pid: u32) -> Result<(), Refused> {
+        match self
+            .runtime_pid
+            .compare_exchange(UNSET, pid, Ordering::SeqCst, Ordering::SeqCst)
+        {
+            Err(owner) if owner != pid => return Err(Refused::ForkedAfterStart),
+            _ => {}
+        }
+
+        if self.fork_only_pid.load(Ordering::SeqCst) == pid {
+            // Nothing was started, so the workers forked from here must still find it unclaimed.
+            let _ =
+                self.runtime_pid
+                    .compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst);
+            return Err(Refused::ReservedForForking);
+        }
+
+        Ok(())
+    }
+
+    /// Reserves `pid` for forking. Reserve first, then look for a started runtime: `enter` does
+    /// the mirror image, so when the two race at least one of them sees the other. A refused
+    /// reservation leaves the gate exactly as it was, so a process already running the runtime
+    /// keeps refusing the children it forks.
+    pub(crate) fn reserve(&self, pid: u32) -> Result<(), RuntimeAlreadyStarted> {
+        self.fork_only_pid.store(pid, Ordering::SeqCst);
+        if self.runtime_pid.load(Ordering::SeqCst) == pid {
+            let _ =
+                self.fork_only_pid
+                    .compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst);
+            return Err(RuntimeAlreadyStarted);
+        }
+        Ok(())
+    }
+
+    pub(crate) fn started(&self, pid: u32) -> bool {
+        self.runtime_pid.load(Ordering::SeqCst) == pid
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const MASTER: u32 = 100;
+    const WORKER: u32 = 101;
+
+    #[test]
+    fn unreserved_process_starts_the_runtime_and_stays_started() {
+        let gate = ForkGate::new();
+
+        assert!(!gate.started(MASTER));
+        assert_eq!(gate.enter(MASTER), Ok(()));
+        assert_eq!(gate.enter(MASTER), Ok(()));
+        assert!(gate.started(MASTER));
+    }
+
+    #[test]
+    fn reserved_process_can_never_start_the_runtime() {
+        let gate = ForkGate::new();
+
+        assert_eq!(gate.reserve(MASTER), Ok(()));
+        assert_eq!(gate.enter(MASTER), Err(Refused::ReservedForForking));
+        assert_eq!(gate.enter(MASTER), Err(Refused::ReservedForForking));
+        assert!(!gate.started(MASTER));
+    }
+
+    #[test]
+    fn workers_forked_from_a_reserved_process_start_their_own_runtime() {
+        let gate = ForkGate::new();
+        gate.reserve(MASTER).unwrap();
+        gate.enter(MASTER).unwrap_err();
+
+        assert_eq!(gate.enter(WORKER), Ok(()));
+        assert!(gate.started(WORKER));
+    }
+
+    #[test]
+    fn reserving_after_the_runtime_started_is_refused() {
+        let gate = ForkGate::new();
+        gate.enter(MASTER).unwrap();
+
+        assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted));
+    }
+
+    #[test]
+    fn a_refused_reservation_leaves_the_runtime_claimed_and_its_children_refused() {
+        let gate = ForkGate::new();
+        gate.enter(MASTER).unwrap();
+
+        assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted));
+        assert_eq!(gate.enter(MASTER), Ok(()));
+        assert!(gate.started(MASTER));
+        assert_eq!(gate.enter(WORKER), Err(Refused::ForkedAfterStart));
+    }
+
+    #[test]
+    fn child_forked_after_the_runtime_started_is_refused_instead_of_hanging() {
+        let gate = ForkGate::new();
+        gate.enter(MASTER).unwrap();
+
+        assert_eq!(gate.enter(WORKER), Err(Refused::ForkedAfterStart));
+        assert!(!gate.started(WORKER));
+        assert_eq!(gate.enter(MASTER), Ok(()));
+    }
+}

litellm-rust/crates/host-python/src/lib.rs

@@ -8,6 +8,7 @@ mod argument;
 mod callable;
 mod driver;
 mod execution;
+mod fork_gate;
 mod gil;
 mod handle;
 mod marshal;
@@ -18,7 +19,12 @@ pub use adapter::{
 pub use argument::lookup;
 pub use callable::wrap_failure;
 pub use driver::run_call;
-pub use execution::{poll_async_value, run_async, run_async_value, run_sync, run_sync_value};
+pub use execution::{
+    ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, poll_async_value,
+    reserve_process_for_forking, run_async, run_async_value, run_sync, run_sync_value,
+    runtime_started,
+};
+pub use fork_gate::RuntimeAlreadyStarted;
 pub use gil::{release_count, release_gil};
 pub use handle::{Execution, ExecutionBody, ExecutionStep};
 pub use marshal::{Pythonized, from_py, from_py_argument, panic_to_pyerr, to_py};

litellm-rust/crates/python-bridge/AGENTS.md

@@ -1,9 +1,9 @@
 - Target invariants, not completion claims; these supersede the crate guidance below where they conflict
 - Keep this crate the product-specific PyO3 consumer of `litellm-host-python`
   - Own registration, input projection, the route host and the caller callables it answers operations with (file readers, token providers), public response/error construction and the per-call composition of machine, route host and callback contract
-  - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, re-aliasing unchanged body keys) lives in `litellm-callbacks-legacy` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy
+  - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, re-aliasing unchanged body keys) lives in `litellm-callbacks-legacy-python` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy
   - Value-oriented execution, sync waiting, nested-runtime checks, signal polling and panic containment live in `litellm-host-python`; native async work uses `pyo3-async-runtimes`, Serde output uses `Pythonized<T>`
-  - Core owns typed native state, the route machine, provider preparation/I/O and normalization; the host driver owns terminal events; the legacy adapter in `litellm-callbacks-legacy` owns `Logging` dispatch policy
+  - Core owns typed native state, the route machine, provider preparation/I/O and normalization; the host driver owns terminal events; the legacy adapter in `litellm-callbacks-legacy-python` owns `Logging` dispatch policy
   - Python, Rust SDK and gateway use one lifecycle-bearing core route entrypoint; provider helpers stay private, never bridge-accessible transport drivers
   - Built-in provider/config/secret/auth/document preparation stays in Rust; caller-authored callbacks and focused Python-file reads run only at core-selected points
 - Target GIL-enabled CPython explicitly with `#[pymodule(gil_used = true)]`; detach Rust-only work

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

@@ -10,15 +10,18 @@ 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
 litellm-auth.workspace = true
-litellm-callbacks-legacy.workspace = true
+litellm-callbacks-legacy-python.workspace = true
 litellm-core.workspace = true
 litellm-core-utils.workspace = true
 litellm-auth-gcp.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/diagnostics.rs

@@ -1,5 +1,5 @@
-use litellm_host_python::release_count;
-use pyo3::{prelude::*, types::PyDict};
+use litellm_host_python::{release_count, runtime_started};
+use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict};
 
 #[pyfunction]
 pub(crate) fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
@@ -8,6 +8,20 @@ pub(crate) fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
     Ok(stats.into_any().unbind())
 }
 
+/// True once this process has started the native runtime, which does not survive `fork()`.
+#[pyfunction]
+pub(crate) fn process_state_started() -> bool {
+    runtime_started()
+}
+
+/// Declares that this process only forks workers: from now on every native route raises here,
+/// so the runtime can never start. Raises if it already has. Forked workers are unaffected.
+#[pyfunction]
+pub(crate) fn reserve_process_for_forking() -> PyResult<()> {
+    litellm_host_python::reserve_process_for_forking()
+        .map_err(|_| PyRuntimeError::new_err("the native runtime already started in this process"))
+}
+
 #[cfg(feature = "panic-test")]
 #[pyfunction]
 pub(crate) fn _panic_for_test() {

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

@@ -13,7 +13,7 @@ mod _native {
     #[pymodule_export]
     use crate::diagnostics::_panic_for_test;
     #[pymodule_export]
-    use crate::diagnostics::gil_stats;
+    use crate::diagnostics::{gil_stats, process_state_started, reserve_process_for_forking};
     #[pymodule_export]
     use crate::errors::{RustBridgeDeclined, RustUpstreamError};
     #[pymodule_export]
@@ -30,6 +30,8 @@ mod _native {
     use crate::routes::responses::ResponsesWebSocketConnection;
     #[pymodule_export]
     use crate::token_counter::TokenCounter;
+    #[pymodule_export]
+    use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
 }
 
 use pyo3::prelude::*;
@@ -50,6 +52,8 @@ mod tests {
             let mut expected = vec![
                 "RustBridgeDeclined",
                 "RustUpstreamError",
+                "ForkedAfterNativeRuntimeStarted",
+                "ProcessReservedForForking",
                 "ocr",
                 "aocr",
                 "transcription",
@@ -62,6 +66,8 @@ mod tests {
                 "ResponsesWebSocketConnection",
                 "TokenCounter",
                 "gil_stats",
+                "process_state_started",
+                "reserve_process_for_forking",
             ];
             expected.sort_unstable();
 

litellm-rust/crates/python-bridge/src/routes/messages/mod.rs

@@ -1,7 +1,9 @@
 mod host;
 
 use host::MessagesRouteHost;
-use litellm_callbacks_legacy::{LegacySurface, PassThroughStream, PublicCall, run_legacy_call};
+use litellm_callbacks_legacy_python::{
+    LegacySurface, PassThroughStream, PublicCall, run_legacy_call,
+};
 use litellm_core::messages::route::{messages_machine, supports};
 use pyo3::{
     prelude::*,

litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs

@@ -7,7 +7,7 @@ use std::sync::{Arc, LazyLock};
 
 use host::OcrRouteHost;
 use litellm_auth_gcp::VertexAuth;
-use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call};
+use litellm_callbacks_legacy_python::{LegacySurface, PublicCall, run_legacy_call};
 use litellm_core::ocr::route::ocr_machine;
 use litellm_core_utils::settings::ProcessEnvironment;
 use litellm_llms::base_llm::ocr::{

litellm-rust/crates/python-bridge/src/routes/responses.rs

@@ -25,7 +25,7 @@ impl ResponsesWebSocketConnection {
     ) -> PyResult<Bound<'py, PyAny>> {
         let headers = marshal_headers(headers)?;
         let timeout = optional_timeout(timeout_seconds);
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        litellm_host_python::run_async_value(py, async move {
             let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout)
                 .await
                 .map_err(responses_error_to_pyerr)?;
@@ -35,7 +35,7 @@ impl ResponsesWebSocketConnection {
 
     fn send_text<'py>(&self, py: Python<'py>, text: String) -> PyResult<Bound<'py, PyAny>> {
         let inner = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        litellm_host_python::run_async_value(py, async move {
             inner
                 .send_text(text)
                 .await
@@ -45,14 +45,14 @@ impl ResponsesWebSocketConnection {
 
     fn recv_text<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let inner = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        litellm_host_python::run_async_value(py, async move {
             inner.recv_text().await.map_err(responses_error_to_pyerr)
         })
     }
 
     fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let inner = self.inner.clone();
-        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        litellm_host_python::run_async_value(py, async move {
             inner.close().await.map_err(responses_error_to_pyerr)
         })
     }
@@ -68,6 +68,10 @@ mod tests {
     use tokio_tungstenite::{accept_async, tungstenite::Message};
 
     #[test]
+    #[expect(
+        clippy::disallowed_methods,
+        reason = "the test server shares the routes' runtime"
+    )]
     fn responses_websocket_connection_round_trips_through_python() {
         Python::initialize();
         let runtime = pyo3_async_runtimes::tokio::get_runtime();

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