scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

Merge commit '36b8be7d81b' into litellm_mantle_native_anthropic_messages_b4dc

Browse source 
# Conflicts:
#	tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py
This commit is contained in:
mateo-berri 2026-09-21 12:00:37 -07:00
commit 4eed951e6f
1538 changed files with 26360 additions and 11731 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

209
.circleci/config.yml
View file

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

113
.circleci/scripts/run_migration_tests.py Normal file
View file

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

3
.env.example
View file

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

15
.github/e2e-stack/assert_tests_ran.py vendored
View file

@ -18,11 +18,6 @@ def main() -> int:
return 1
cases: Final = tuple(report.iter("testcase"))
expected_count: Final = os.environ.get("E2E_REQUIRED_TEST_COUNT")
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
passed: Final = frozenset(
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
)
@ -43,9 +38,10 @@ 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", "")
@ -53,6 +49,11 @@ def main() -> int:
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

2
.github/e2e-stack/select_tests.py vendored
View file

@ -4,7 +4,7 @@ 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$"

2
.github/template.yaml vendored
View file

@ -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:

60
.github/workflows/_test-unit-base.yml vendored
View file

@ -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'

3
.github/workflows/test-e2e-changed.yml vendored
View file

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

3
.github/workflows/test-mcp-oauth-e2e.yml vendored
View file

@ -12,6 +12,9 @@ on:
- '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'

14
.github/workflows/test-rust.yml vendored
View file

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

3
.github/workflows/test-unit.yml vendored
View file

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

5
CONTRIBUTING.md
View file

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

6
README.md
View file

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

1
cookbook/livekit_agent_sdk/config.example.yaml
View file

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

1
cookbook/misc/config.yaml
View file

@ -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:

3
docker/.env.example
View file

@ -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.

3
helm/litellm/templates/backend/deployment.yaml
View file

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

3
helm/litellm/templates/gateway/deployment.yaml
View file

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

3
helm/litellm/templates/ui/deployment.yaml
View file

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

100
helm/litellm/tests/replica_count_tests.yaml Normal file
View file

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

9
helm/litellm/values.yaml
View file

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

89
litellm-proxy-extras/litellm_proxy_extras/migration_lock.py Normal file
View file

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

158
litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py Normal file
View file

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

5
litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py
View file

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

197
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

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

231
litellm-proxy-extras/tests/test_setup_database_fail_fast.py
View file

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

897
litellm-rust/Cargo.lock generated
View file

File diff suppressed because it is too large Load diff

14
litellm-rust/Cargo.toml
View file

@ -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,10 +29,15 @@ litellm-core-utils = { path = "crates/core-utils" }
litellm-cache = { path = "crates/cache" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-token-counter = { path = "crates/token-counter" }
litellm-token-counter-fast = { path = "crates/token-counter-fast" }
litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" }
litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" }
litellm-host-python = { path = "crates/host-python" }
bytes = "1"
http = "1"
google-cloud-auth = { version = "1.16.0", default-features = false }
jsonwebtoken = { version = "11.1.0", default-features = false }
hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] }
proptest = "1.7.0"
pyo3 = "0.29.2"
@ -45,6 +55,8 @@ serde_with = { version = "=3.16.1", default-features = false, features = ["std",
sha2 = "0.10"
subtle = "2"
thiserror = "2.0"
tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }
tiktoken-rs = "0.12.0"
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] }
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }

2
litellm-rust/crates/auth-aws/Cargo.toml
View file

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

2
litellm-rust/crates/auth-aws/src/constants.rs
View file

@ -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";

6
litellm-rust/crates/auth-aws/src/error.rs
View file

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

2
litellm-rust/crates/auth-azure/Cargo.toml
View file

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

2
litellm-rust/crates/auth-azure/src/credential_provider_cache.rs
View file

@ -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 {

13
litellm-rust/crates/auth-azure/src/native.rs
View file

@ -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
));
}
}
}

16
litellm-rust/crates/auth-azure/src/resolve.rs
View file

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

4
litellm-rust/crates/auth-azure/src/types.rs
View file

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

7
litellm-rust/crates/auth-gcp/Cargo.toml
View file

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

62
litellm-rust/crates/auth-gcp/src/lib.rs
View file

@ -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 =

106
litellm-rust/crates/auth-gcp/src/sdk.rs Normal file
View file

@ -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"));
}
}

15
litellm-rust/crates/auth-types/Cargo.toml Normal file
View file

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

4
litellm-rust/crates/auth/src/credential.rs → litellm-rust/crates/auth-types/src/credential.rs
View file

@ -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 {

0
litellm-rust/crates/auth/src/error.rs → litellm-rust/crates/auth-types/src/error.rs
View file

3
litellm-rust/crates/auth/src/http.rs → litellm-rust/crates/auth-types/src/http.rs
View file

@ -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 {

57
litellm-rust/crates/auth-types/src/lib.rs Normal file
View file

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

6
litellm-rust/crates/auth/src/policy.rs → litellm-rust/crates/auth-types/src/policy.rs
View file

@ -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 {

0
litellm-rust/crates/auth/src/secret.rs → litellm-rust/crates/auth-types/src/secret.rs
View file

4
litellm-rust/crates/auth/src/token.rs → litellm-rust/crates/auth-types/src/token.rs
View file

@ -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 {

17
litellm-rust/crates/auth/Cargo.toml
View file

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

61
litellm-rust/crates/auth/src/lib.rs
View file

@ -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;

33
litellm-rust/crates/auth/tests/facade.rs Normal file
View file

@ -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())]
);
}

0
litellm-rust/crates/callbacks-legacy/AGENTS.md → litellm-rust/crates/callbacks-legacy-python/AGENTS.md
View file

2
litellm-rust/crates/callbacks-legacy/Cargo.toml → litellm-rust/crates/callbacks-legacy-python/Cargo.toml
View file

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

0
litellm-rust/crates/callbacks-legacy/python_contract.json → litellm-rust/crates/callbacks-legacy-python/python_contract.json
View file

6
litellm-rust/crates/callbacks-legacy/src/adapter.rs → litellm-rust/crates/callbacks-legacy-python/src/adapter.rs
View file

@ -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.

0
litellm-rust/crates/callbacks-legacy/src/call.rs → litellm-rust/crates/callbacks-legacy-python/src/call.rs
View file

2
litellm-rust/crates/callbacks-legacy/src/callbacks.rs → litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs
View file

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

0
litellm-rust/crates/callbacks-legacy/src/deferred.rs → litellm-rust/crates/callbacks-legacy-python/src/deferred.rs
View file

2
litellm-rust/crates/callbacks-legacy/src/lib.rs → litellm-rust/crates/callbacks-legacy-python/src/lib.rs
View file

@ -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;

8
litellm-rust/crates/callbacks-legacy/src/logger.rs → litellm-rust/crates/callbacks-legacy-python/src/logger.rs
View file

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

2
litellm-rust/crates/callbacks-legacy/src/preparation.rs → litellm-rust/crates/callbacks-legacy-python/src/preparation.rs
View file

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

4
litellm-rust/crates/callbacks-legacy/src/legacy_python.rs → litellm-rust/crates/callbacks-legacy-python/src/python.rs
View file

@ -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 {

0
litellm-rust/crates/callbacks-legacy/tests/deferred.rs → litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs
View file

0
litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs → litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs
View file

0
litellm-rust/crates/callbacks-legacy/tests/payload.rs → litellm-rust/crates/callbacks-legacy-python/tests/payload.rs
View file

10
litellm-rust/crates/callbacks-legacy/tests/support.rs → litellm-rust/crates/callbacks-legacy-python/tests/support.rs
View file

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

0
litellm-rust/crates/callbacks-legacy/tests/terminal.rs → litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs
View file

4
litellm-rust/crates/python-bridge/AGENTS.md
View file

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

9
litellm-rust/crates/python-bridge/Cargo.toml
View file

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

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

@ -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::*,

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

@ -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::{

71
litellm-rust/crates/python-bridge/src/token_counter.rs
View file

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

24
litellm-rust/crates/secrets-aws/Cargo.toml Normal file
View file

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

79
litellm-rust/crates/secrets-aws/src/auth.rs Normal file
View file

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

31
litellm-rust/crates/secrets-aws/src/error.rs Normal file
View file

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

63
litellm-rust/crates/secrets-aws/src/kms.rs Normal file
View file

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

10
litellm-rust/crates/secrets-aws/src/lib.rs Normal file
View file

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

287
litellm-rust/crates/secrets-aws/src/secret_manager.rs Normal file
View file

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

59
litellm-rust/crates/secrets-aws/tests/kms.rs Normal file
View file

@ -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()
);
}
}

312
litellm-rust/crates/secrets-aws/tests/secret_manager.rs Normal file
View file

@ -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(_))
));
}

28
litellm-rust/crates/secrets-google/Cargo.toml Normal file
View file

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

21
litellm-rust/crates/secrets-google/src/auth.rs Normal file
View file

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

43
litellm-rust/crates/secrets-google/src/error.rs Normal file
View file

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

67
litellm-rust/crates/secrets-google/src/kms.rs Normal file
View file

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

10
litellm-rust/crates/secrets-google/src/lib.rs Normal file
View file

@ -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;

157
litellm-rust/crates/secrets-google/src/secret_manager.rs Normal file
View file

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

49
litellm-rust/crates/secrets-google/tests/kms.rs Normal file
View file

@ -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()
);
}
}

188
litellm-rust/crates/secrets-google/tests/secret_manager.rs Normal file
View file

@ -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)
);
}
}

17
litellm-rust/crates/secrets-types/Cargo.toml Normal file
View file

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

58
litellm-rust/crates/secrets-types/src/base_secret_manager.rs Normal file
View file

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

92
litellm-rust/crates/secrets-types/src/config.rs Normal file
View file

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

9
litellm-rust/crates/secrets-types/src/error.rs Normal file
View file

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

12
litellm-rust/crates/secrets-types/src/lib.rs Normal file
View file

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

31
litellm-rust/crates/secrets-types/src/value.rs Normal file
View file

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

60
litellm-rust/crates/secrets-types/tests/config.rs Normal file
View file

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

105
litellm-rust/crates/secrets-types/tests/rotation.rs Normal file
View file

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

34
litellm-rust/crates/secrets/Cargo.toml Normal file
View file

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

11
litellm-rust/crates/secrets/README.md Normal file
View file

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

33
litellm-rust/crates/secrets/src/error.rs Normal file
View file

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

117
litellm-rust/crates/secrets/src/handler.rs Normal file
View file

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

21
litellm-rust/crates/secrets/src/lib.rs Normal file
View file

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

Some files were not shown because too many files have changed in this diff Show more