scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

chore: merge main after MCP server ordering fix

Browse source
This commit is contained in:
Joshua Valluru 2026-09-21 16:26:52 -07:00
commit 8136a67410
106 changed files with 10395 additions and 681 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -130,7 +130,7 @@ jobs:
- name: Test secret manager feature combinations
run: |
cargo test -p litellm-auth-gcp --locked --no-default-features
for features in '' aws google azure cyberark aws,google aws,azure google,azure aws,google,azure aws,google,cyberark aws,google,azure,cyberark; do
for features in '' aws google hashicorp azure cyberark aws,google aws,azure google,azure aws,google,azure aws,google,cyberark aws,google,azure,cyberark aws,google,hashicorp,azure,cyberark; do
cargo test -p litellm-secrets --locked --no-default-features --features "$features"
done

11
docker/README.md
View file

@ -2,6 +2,17 @@
This guide provides instructions for building and running the LiteLLM application using Docker and Docker Compose.
> **Just want to run LiteLLM?** This guide builds from source. To run the published
> image instead, use `docker-compose.quickstart.yml` in this directory — the
> two-service stack (gateway + Postgres) that the
> [Docker quickstart](https://docs.litellm.ai/docs/proxy/docker_quick_start) documents:
>
> ```bash
> curl -sSLO https://github.com/BerriAI/litellm/raw/main/docker/docker-compose.quickstart.yml
> printf 'LITELLM_MASTER_KEY=sk-%s\nLITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" "$(openssl rand -hex 32)" > .env
> docker compose -f docker-compose.quickstart.yml up -d
> ```
## Prerequisites
- Docker

41
docker/docker-compose.quickstart.yml Normal file
View file

@ -0,0 +1,41 @@
# LiteLLM quickstart stack: the gateway plus a Postgres database that stores
# models, virtual keys, and spend logs. Used by
# https://docs.litellm.ai/docs/proxy/docker_quick_start
#
# curl -sSLO https://github.com/BerriAI/litellm/raw/main/docker/docker-compose.quickstart.yml
# printf 'LITELLM_MASTER_KEY=sk-%s\nLITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" "$(openssl rand -hex 32)" > .env
# docker compose -f docker-compose.quickstart.yml up -d
#
# Compose reads .env from this directory. Keep it: regenerating LITELLM_SALT_KEY
# makes credentials already stored in the database unreadable. For anything
# beyond local evaluation, pin the image to a specific release tag.
services:
litellm:
image: docker.litellm.ai/berriai/litellm:main-stable
ports:
- "4000:4000"
environment:
LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY:?set it in .env - see the header of this file}
LITELLM_SALT_KEY: ${LITELLM_SALT_KEY:?set it in .env - see the header of this file}
DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm
STORE_MODEL_IN_DB: "True"
depends_on:
db:
condition: service_healthy
db:
image: postgres:16
environment:
POSTGRES_USER: litellm
POSTGRES_PASSWORD: litellm
POSTGRES_DB: litellm
healthcheck:
test: ["CMD-SHELL", "pg_isready -U litellm"]
interval: 5s
timeout: 5s
retries: 10
volumes:
- postgres_data:/var/lib/postgresql/data
volumes:
postgres_data:

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260921000000_add_password_reset_columns/migration.sql Normal file
View file

@ -0,0 +1,3 @@
ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "password_reset_required" BOOLEAN;
ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "last_breach_check_at" TIMESTAMP(3);

2
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -246,6 +246,8 @@ model LiteLLM_UserTable {
organization_id String?
object_permission_id String?
password String?
password_reset_required Boolean?
last_breach_check_at DateTime?
teams String[] @default([])
user_role String?
max_budget Float?

428
litellm-rust/Cargo.lock generated
View file

@ -40,6 +40,12 @@ dependencies = [
"cc",
]
[[package]]
name = "allocator-api2"
version = "0.2.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923"
[[package]]
name = "android_system_properties"
version = "0.1.6"
@ -230,6 +236,7 @@ dependencies = [
"aws-credential-types",
"aws-sigv4",
"aws-smithy-async",
"aws-smithy-eventstream",
"aws-smithy-http",
"aws-smithy-runtime",
"aws-smithy-runtime-api",
@ -238,7 +245,9 @@ dependencies = [
"bytes",
"bytes-utils",
"fastrand",
"http 0.2.12",
"http 1.4.2",
"http-body 0.4.6",
"http-body 1.1.0",
"percent-encoding",
"pin-project-lite",
@ -272,6 +281,43 @@ dependencies = [
"tracing",
]
[[package]]
name = "aws-sdk-s3"
version = "1.146.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2cd651b4400d4011b8927b83a9552bf90ff11e6e5da0b9f0a7583247aceec971"
dependencies = [
"arc-swap",
"aws-credential-types",
"aws-runtime",
"aws-sigv4",
"aws-smithy-async",
"aws-smithy-checksums",
"aws-smithy-eventstream",
"aws-smithy-http",
"aws-smithy-json",
"aws-smithy-observability",
"aws-smithy-runtime",
"aws-smithy-runtime-api",
"aws-smithy-schema",
"aws-smithy-types",
"aws-smithy-xml 0.62.1",
"aws-types",
"bytes",
"fastrand",
"hex",
"hmac",
"http 0.2.12",
"http 1.4.2",
"http-body 1.1.0",
"lru",
"percent-encoding",
"regex-lite",
"sha2 0.11.0",
"tracing",
"url",
]
[[package]]
name = "aws-sdk-secretsmanager"
version = "1.117.0"
@ -316,7 +362,7 @@ dependencies = [
"aws-smithy-runtime-api",
"aws-smithy-schema",
"aws-smithy-types",
"aws-smithy-xml",
"aws-smithy-xml 0.61.1",
"aws-types",
"fastrand",
"http 0.2.12",
@ -332,6 +378,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "31d955e76ff96acd555bf06fa0fa6d5bf9335fa84ae7c64481b20ae61d231f70"
dependencies = [
"aws-credential-types",
"aws-smithy-eventstream",
"aws-smithy-http",
"aws-smithy-runtime-api",
"aws-smithy-types",
@ -359,10 +406,31 @@ dependencies = [
]
[[package]]
name = "aws-smithy-eventstream"
version = "0.61.1"
name = "aws-smithy-checksums"
version = "0.65.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a9381123ab62d20c13082b151f30f962a3b112b727345394536dfa39a482944"
checksum = "b67ecd999972b58e67cab052f5129906c08c25883bd0788ceefc55ef97d61307"
dependencies = [
"aws-smithy-http",
"aws-smithy-types",
"bytes",
"crc-fast",
"hex",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"md-5",
"pin-project-lite",
"sha1 0.11.0",
"sha2 0.11.0",
"tracing",
]
[[package]]
name = "aws-smithy-eventstream"
version = "0.61.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "80c2051c2f1016fb8e6548dd07b8bc2ac9c3fe583721444b92f515e856d31609"
dependencies = [
"aws-smithy-types",
"bytes",
@ -375,6 +443,7 @@ version = "0.64.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "37843d9add67c3aff5856f409c6dc315d3cdff60f9c0cb5b670dab1e9920306d"
dependencies = [
"aws-smithy-eventstream",
"aws-smithy-runtime-api",
"aws-smithy-types",
"bytes",
@ -554,6 +623,18 @@ dependencies = [
"xmlparser",
]
[[package]]
name = "aws-smithy-xml"
version = "0.62.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b932c8d6dc127fc980eecd78f8694ae9b9551b69a93a7def2a199c1c0033daf"
dependencies = [
"aws-smithy-runtime-api",
"aws-smithy-schema",
"aws-smithy-types",
"xmlparser",
]
[[package]]
name = "aws-types"
version = "1.6.0"
@ -980,6 +1061,16 @@ dependencies = [
"libc",
]
[[package]]
name = "crc-fast"
version = "1.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e75b2483e97a5a7da73ac68a05b629f9c53cff58d8ed1c77866079e18b00dba5"
dependencies = [
"digest 0.10.7",
"spin",
]
[[package]]
name = "crc16"
version = "0.4.0"
@ -1348,7 +1439,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
dependencies = [
"libc",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -1377,6 +1468,18 @@ dependencies = [
"pin-project-lite",
]
[[package]]
name = "fallible-iterator"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649"
[[package]]
name = "fallible-streaming-iterator"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a"
[[package]]
name = "fancy-regex"
version = "0.17.0"
@ -1428,6 +1531,12 @@ version = "1.0.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
[[package]]
name = "foldhash"
version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
[[package]]
name = "form_urlencoded"
version = "1.2.2"
@ -1892,11 +2001,34 @@ version = "0.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
[[package]]
name = "hashbrown"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
dependencies = [
"foldhash",
]
[[package]]
name = "hashbrown"
version = "0.17.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
dependencies = [
"allocator-api2",
"equivalent",
"foldhash",
]
[[package]]
name = "hashlink"
version = "0.12.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a596f1b20ed2cc5ecac41a164aaebc7258057060f06c0cf7a2ba3991ee7990fb"
dependencies = [
"hashbrown 0.17.1",
]
[[package]]
name = "heck"
@ -2109,7 +2241,7 @@ dependencies = [
"libc",
"percent-encoding",
"pin-project-lite",
"socket2 0.6.5",
"socket2 0.5.10",
"tokio",
"tower-service",
"tracing",
@ -2277,6 +2409,12 @@ version = "2.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2"
[[package]]
name = "iter-read"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "071ed4cc1afd86650602c7b11aa2e1ce30762a1c27193201cb5cee9c6ebb1294"
[[package]]
name = "itertools"
version = "0.13.0"
@ -2435,6 +2573,17 @@ version = "0.2.186"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "libsqlite3-sys"
version = "0.38.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f1d20bef17f513b9b3004532233187769cd072d790971f4e4da0e346eb6401e8"
dependencies = [
"cc",
"pkg-config",
"vcpkg",
]
[[package]]
name = "linux-raw-sys"
version = "0.12.1"
@ -2540,6 +2689,36 @@ dependencies = [
"url",
]
[[package]]
name = "litellm-cache-disk"
version = "0.1.0"
dependencies = [
"litellm-cache",
"py_literal",
"rand 0.8.7",
"rstest",
"rusqlite",
"serde-pickle",
"serde_json",
"tempfile",
"tokio",
]
[[package]]
name = "litellm-cache-gcs"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-auth-gcp",
"litellm-auth-types",
"litellm-cache",
"percent-encoding",
"reqwest 0.12.28",
"serde_json",
"tokio",
"wiremock",
]
[[package]]
name = "litellm-cache-memory"
version = "0.1.0"
@ -2578,6 +2757,21 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-cache-s3"
version = "0.1.0"
dependencies = [
"aws-credential-types",
"aws-sdk-s3",
"aws-smithy-types",
"aws-types",
"litellm-auth-aws",
"litellm-cache",
"serde_json",
"tokio",
"wiremock",
]
[[package]]
name = "litellm-callbacks-legacy-python"
version = "0.1.0"
@ -2740,12 +2934,16 @@ dependencies = [
"criterion",
"futures-util",
"litellm-auth",
"litellm-auth-aws",
"litellm-auth-gcp",
"litellm-cache",
"litellm-cache-azure-blob",
"litellm-cache-disk",
"litellm-cache-gcs",
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-cache-response",
"litellm-cache-s3",
"litellm-callbacks-legacy-python",
"litellm-core",
"litellm-core-utils",
@ -2778,6 +2976,7 @@ dependencies = [
"litellm-secrets-azure",
"litellm-secrets-cyberark",
"litellm-secrets-google",
"litellm-secrets-hashicorp",
"litellm-secrets-types",
"moka",
"reqwest 0.12.28",
@ -2875,6 +3074,26 @@ dependencies = [
"wiremock",
]
[[package]]
name = "litellm-secrets-hashicorp"
version = "0.1.0"
dependencies = [
"litellm-core-utils",
"litellm-secrets-types",
"moka",
"rstest",
"rustify",
"rustify_derive",
"serde",
"serde_json",
"tempfile",
"thiserror 2.0.19",
"tokio",
"vaultrs",
"veil",
"wiremock",
]
[[package]]
name = "litellm-secrets-types"
version = "0.1.0"
@ -2966,6 +3185,15 @@ version = "0.4.33"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad"
[[package]]
name = "lru"
version = "0.18.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ff9840bcc50b71349309900da0ce7279aa336ae71d73250b07998932c7d97c25"
dependencies = [
"hashbrown 0.17.1",
]
[[package]]
name = "lru-slab"
version = "0.1.2"
@ -2988,6 +3216,16 @@ version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c"
[[package]]
name = "md-5"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98"
dependencies = [
"cfg-if",
"digest 0.11.3",
]
[[package]]
name = "memchr"
version = "2.8.3"
@ -3630,7 +3868,7 @@ dependencies = [
"quinn-udp",
"rustc-hash",
"rustls 0.23.42",
"socket2 0.6.5",
"socket2 0.5.10",
"thiserror 2.0.19",
"tokio",
"tracing",
@ -3669,9 +3907,9 @@ dependencies = [
"cfg_aliases",
"libc",
"once_cell",
"socket2 0.6.5",
"socket2 0.5.10",
"tracing",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -4033,6 +4271,16 @@ dependencies = [
"windows-sys 0.52.0",
]
[[package]]
name = "rsqlite-vfs"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c"
dependencies = [
"hashbrown 0.16.1",
"thiserror 2.0.19",
]
[[package]]
name = "rstest"
version = "0.26.1"
@ -4073,6 +4321,21 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "rusqlite"
version = "0.40.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "23f2a97da3e3873c73cb2a2e71b35c40ff95e0b1eefa8d72d8499a6928c3b5b3"
dependencies = [
"bitflags 2.13.1",
"fallible-iterator",
"fallible-streaming-iterator",
"hashlink",
"libsqlite3-sys",
"smallvec",
"sqlite-wasm-rs",
]
[[package]]
name = "rustc-hash"
version = "2.1.3"
@ -4088,6 +4351,40 @@ dependencies = [
"semver",
]
[[package]]
name = "rustify"
version = "0.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4800ce4c1cc2fec12c559dae2ddbf0e17fcee7569b796e6d75898efef443368b"
dependencies = [
"anyhow",
"async-trait",
"bytes",
"http 1.4.2",
"reqwest 0.13.5",
"rustify_derive",
"serde",
"serde_json",
"serde_urlencoded",
"thiserror 1.0.69",
"tracing",
"url",
]
[[package]]
name = "rustify_derive"
version = "0.5.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "78ea7fda74240f7410d0198b603a8a2f662acc7d76b6667a49f9b162cd8d9b4f"
dependencies = [
"proc-macro2",
"quote",
"regex",
"serde_urlencoded",
"syn 1.0.109",
"synstructure 0.12.6",
]
[[package]]
name = "rustix"
version = "1.1.5"
@ -4098,7 +4395,7 @@ dependencies = [
"errno",
"libc",
"linux-raw-sys",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -4169,7 +4466,7 @@ dependencies = [
"security-framework",
"security-framework-sys",
"webpki-root-certs",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -4330,6 +4627,19 @@ dependencies = [
"serde_derive",
]
[[package]]
name = "serde-pickle"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b641fdc8bcf2781ee78b30c599700d64ad4f412976143e4c5d0b9df906bb4843"
dependencies = [
"byteorder",
"iter-read",
"num-bigint 0.4.8",
"num-traits",
"serde",
]
[[package]]
name = "serde_core"
version = "1.0.229"
@ -4429,6 +4739,17 @@ dependencies = [
"digest 0.10.7",
]
[[package]]
name = "sha1"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214"
dependencies = [
"cfg-if",
"cpufeatures 0.3.0",
"digest 0.11.3",
]
[[package]]
name = "sha1_smol"
version = "1.0.1"
@ -4545,6 +4866,12 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "spin"
version = "0.10.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3"
[[package]]
name = "spm_precompiled"
version = "0.1.4"
@ -4557,6 +4884,18 @@ dependencies = [
"unicode-segmentation",
]
[[package]]
name = "sqlite-wasm-rs"
version = "0.5.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dc3efc0da82635d7e1ced0053bbbfa8c7ab9645d0bf36ceb4f7127bb85315d75"
dependencies = [
"cc",
"js-sys",
"rsqlite-vfs",
"wasm-bindgen",
]
[[package]]
name = "sse-stream"
version = "0.2.6"
@ -4615,6 +4954,17 @@ version = "2.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
[[package]]
name = "syn"
version = "1.0.109"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237"
dependencies = [
"proc-macro2",
"quote",
"unicode-ident",
]
[[package]]
name = "syn"
version = "2.0.119"
@ -4646,6 +4996,18 @@ dependencies = [
"futures-core",
]
[[package]]
name = "synstructure"
version = "0.12.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f36bdaa60a83aca3921b5259d5400cbf5e90fc51931376a9bd4a0eb79aa7210f"
dependencies = [
"proc-macro2",
"quote",
"syn 1.0.109",
"unicode-xid",
]
[[package]]
name = "synstructure"
version = "0.13.2"
@ -4676,10 +5038,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.4.3",
"getrandom 0.3.4",
"once_cell",
"rustix",
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -5060,6 +5422,7 @@ version = "0.1.44"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100"
dependencies = [
"log",
"pin-project-lite",
"tracing-attributes",
"tracing-core",
@ -5142,7 +5505,7 @@ dependencies = [
"rand 0.8.7",
"rustls 0.23.42",
"rustls-pki-types",
"sha1",
"sha1 0.10.7",
"thiserror 1.0.69",
"utf-8",
]
@ -5244,6 +5607,12 @@ version = "1.13.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8"
[[package]]
name = "unicode-xid"
version = "0.2.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
[[package]]
name = "unicode_categories"
version = "0.1.1"
@ -5303,6 +5672,31 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65"
[[package]]
name = "vaultrs"
version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "30ffcc0e81025065dda612ec1e26a3d81bb16ef3062354873d17a35965d68522"
dependencies = [
"async-trait",
"derive_builder",
"http 1.4.2",
"reqwest 0.13.5",
"rustify",
"rustify_derive",
"serde",
"serde_json",
"thiserror 2.0.19",
"tracing",
"url",
]
[[package]]
name = "vcpkg"
version = "0.2.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426"
[[package]]
name = "veil"
version = "0.3.0"
@ -5520,7 +5914,7 @@ version = "0.1.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
dependencies = [
"windows-sys 0.61.2",
"windows-sys 0.52.0",
]
[[package]]
@ -5746,7 +6140,7 @@ dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
"synstructure",
"synstructure 0.13.2",
]
[[package]]
@ -5787,7 +6181,7 @@ dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
"synstructure",
"synstructure 0.13.2",
]
[[package]]

8
litellm-rust/Cargo.toml
View file

@ -22,6 +22,7 @@ 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-secrets-hashicorp = { path = "crates/secrets-hashicorp" }
litellm-secrets-azure = { path = "crates/secrets-azure" }
litellm-secrets-cyberark = { path = "crates/secrets-cyberark" }
litellm-http = { path = "crates/http" }
@ -32,6 +33,9 @@ litellm-cache = { path = "crates/cache" }
litellm-cache-azure-blob = { path = "crates/cache-azure-blob" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-cache-redis = { path = "crates/cache-redis" }
litellm-cache-s3 = { path = "crates/cache-s3" }
litellm-cache-gcs = { path = "crates/cache-gcs" }
litellm-cache-disk = { path = "crates/cache-disk" }
litellm-cache-response = { path = "crates/cache-response" }
litellm-token-counter = { path = "crates/token-counter" }
litellm-token-counter-fast = { path = "crates/token-counter-fast" }
@ -53,6 +57,9 @@ reqwest = { version = "0.12", default-features = false, features = ["json", "mul
rstest = "0.26.1"
rstest_reuse = "0.7.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
rustify = "=0.7.0"
rustify_derive = "=0.5.5"
vaultrs = { version = "=0.8.0", default-features = false, features = ["rustls"] }
rustls-native-certs = "0.8"
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }
@ -69,6 +76,7 @@ base64 = "0.22"
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
percent-encoding = "2.3"
webpki-roots = "1"
time = { version = "0.3.53", features = ["parsing"] }
criterion = "0.8.2"

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

@ -128,6 +128,14 @@ impl VertexAuth {
}
}
pub async fn access_token(
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<String, Error> {
self.load_provider(config, env_lookup).await?.token().await
}
pub async fn validate_environment(
&self,
headers: Vec<(String, String)>,

19
litellm-rust/crates/cache-disk/Cargo.toml Normal file
View file

@ -0,0 +1,19 @@
[package]
name = "litellm-cache-disk"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
py_literal = "0.4.0"
rand.workspace = true
rusqlite = { version = "0.40", features = ["bundled"] }
serde-pickle = "1.2"
serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
rstest.workspace = true
tempfile = "3.27.0"

10
litellm-rust/crates/cache-disk/src/adapter.rs Normal file
View file

@ -0,0 +1,10 @@
use litellm_cache::Error;
use crate::StoredValue;
pub trait ValueAdapter: Send + Sync + 'static {
fn read(&self, value: StoredValue) -> Result<Option<Vec<u8>>, Error>;
fn write(&self, payload: Vec<u8>) -> StoredValue;
fn counter_seed(&self, value: Option<StoredValue>) -> Result<f64, Error>;
fn counter_value(&self, value: f64) -> StoredValue;
}

301
litellm-rust/crates/cache-disk/src/cache.rs Normal file
View file

@ -0,0 +1,301 @@
use std::{
path::Path,
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus,
CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache,
};
use crate::{DiskStore, DiskcacheSqliteStore, PythonDiskCacheAdapter, StoredValue, ValueAdapter};
pub struct DiskCache<S, D = DiskcacheSqliteStore, A = PythonDiskCacheAdapter> {
store: Arc<D>,
adapter: Arc<A>,
codec: S,
}
impl<S: CacheCodec> DiskCache<S> {
pub fn open(directory: impl AsRef<Path>, codec: S) -> Result<Self, Error> {
Ok(Self {
store: Arc::new(DiskcacheSqliteStore::open(directory)?),
adapter: Arc::new(PythonDiskCacheAdapter),
codec,
})
}
}
impl<S: CacheCodec, D: DiskStore> DiskCache<S, D, PythonDiskCacheAdapter> {
pub fn with_store(store: D, codec: S) -> Self {
Self {
store: Arc::new(store),
adapter: Arc::new(PythonDiskCacheAdapter),
codec,
}
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DiskCache<S, D, A> {
pub fn with_adapter(store: D, adapter: A, codec: S) -> Self {
Self {
store: Arc::new(store),
adapter: Arc::new(adapter),
codec,
}
}
pub fn directory(&self) -> &Path {
self.store.directory()
}
fn decode_stored(&self, value: StoredValue) -> Result<Option<S::Value>, Error> {
let Some(bytes) = self.adapter.read(value)? else {
return Ok(None);
};
self.codec.decode(&bytes).map(Some)
}
async fn run_blocking<T, F>(store: Arc<D>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&D) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || operation(&store))
.await
.map_err(|_| Error::Unavailable)?
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BaseCache for DiskCache<S, D, A> {
type Value = S::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &Self::Context,
) -> Result<(), Error> {
let value = self.adapter.write(self.codec.encode(&value)?);
let expire_time = context.ttl.map(|ttl| unix_now() + ttl.as_secs_f64());
self.store.set(key, value, expire_time, unix_now())
}
fn get_cache(&self, key: &str, _: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.store
.get(key, unix_now())?
.map(|value| self.decode_stored(value))
.transpose()
.map(|value| value.flatten())
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: ExactCacheContext,
) -> Result<(), Error> {
let value = self.adapter.write(self.codec.encode(&value)?);
let ttl = context.ttl;
let key = key.to_string();
Self::run_blocking(Arc::clone(&self.store), move |store| {
let expire_time = ttl.map(|ttl| unix_now() + ttl.as_secs_f64());
store.set(&key, value, expire_time, unix_now())
})
.await
}
async fn async_get_cache(
&self,
key: &str,
_: &ExactCacheContext,
) -> Result<Option<Self::Value>, Error> {
let key = key.to_string();
let value = Self::run_blocking(Arc::clone(&self.store), move |store| {
store.get(&key, unix_now())
})
.await?;
value
.map(|value| self.decode_stored(value))
.transpose()
.map(|value| value.flatten())
}
async fn async_set_cache_pipeline(
&self,
entries: Vec<(String, Self::Value)>,
context: ExactCacheContext,
) -> Result<(), Error> {
let entries = entries
.into_iter()
.map(|(key, value)| {
self.codec
.encode(&value)
.map(|value| (key, self.adapter.write(value)))
})
.collect::<Result<Vec<_>, _>>()?;
let expire_after = context.ttl;
Self::run_blocking(Arc::clone(&self.store), move |store| {
for (key, value) in entries {
let expire_time = expire_after.map(|ttl| unix_now() + ttl.as_secs_f64());
store.set(&key, value, expire_time, unix_now())?;
}
Ok(())
})
.await
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
let result = Self::run_blocking(Arc::clone(&self.store), |store| {
store.probe().map(|_| CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Disk cache connection test successful".into(),
error: None,
})
})
.await;
Ok(match result {
Ok(result) => result,
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Disk cache connection failed: {error}"),
error: Some(error.to_string()),
},
})
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D, A> {
fn batch_get_cache(
&self,
keys: &[String],
context: &ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
keys.iter()
.map(|key| match self.get_cache(key, context) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
})
.collect()
}
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
_: ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let values = Self::run_blocking(Arc::clone(&self.store), move |store| {
keys.into_iter()
.map(|key| store.get(&key, unix_now()).map(|value| (key, value)))
.collect::<Result<Vec<_>, _>>()
})
.await?;
values
.into_iter()
.map(|(_, value)| match value {
None => Ok(BatchEntry::Miss),
Some(value) => match self.decode_stored(value) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
},
})
.collect()
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DeleteCache for DiskCache<S, D, A> {
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.store.pop(key, unix_now()).map(|_| ())
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = key.to_string();
Self::run_blocking(Arc::clone(&self.store), move |store| {
store.pop(&key, unix_now()).map(|_| ())
})
.await
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> FlushCache for DiskCache<S, D, A> {
fn flush_cache(&self) -> Result<(), Error> {
self.store.clear()
}
async fn async_flush_cache(&self) -> Result<(), Error> {
Self::run_blocking(Arc::clone(&self.store), |store| store.clear()).await
}
}
impl<S: CacheCodec<Value = f64>, D: DiskStore, A: ValueAdapter> CounterCache
for DiskCache<S, D, A>
{
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
increment(
self.adapter.as_ref(),
self.store.as_ref(),
key,
amount,
context.ttl,
)
}
async fn async_increment(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
let key = key.to_string();
let adapter = Arc::clone(&self.adapter);
Self::run_blocking(Arc::clone(&self.store), move |store| {
increment(adapter.as_ref(), store, &key, amount, context.ttl)
})
.await
}
}
fn increment<A: ValueAdapter, D: DiskStore>(
adapter: &A,
store: &D,
key: &str,
amount: f64,
ttl: Option<Duration>,
) -> Result<f64, Error> {
let mut result = None;
let mut apply = |current: Option<StoredValue>| {
let initial = adapter.counter_seed(current)?;
let value = initial + amount;
let stored = adapter.counter_value(value);
result = Some(value);
Ok((stored, ttl.map(|ttl| unix_now() + ttl.as_secs_f64())))
};
store.update(key, unix_now(), &mut apply)?;
result.ok_or(Error::InvalidEntry)
}
fn unix_now() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs_f64()
}

11
litellm-rust/crates/cache-disk/src/lib.rs Normal file
View file

@ -0,0 +1,11 @@
mod adapter;
mod cache;
mod python;
mod sqlite;
mod store;
pub use adapter::ValueAdapter;
pub use cache::DiskCache;
pub use python::PythonDiskCacheAdapter;
pub use sqlite::DiskcacheSqliteStore;
pub use store::{DiskStore, StoredValue};

77
litellm-rust/crates/cache-disk/src/python/mod.rs Normal file
View file

@ -0,0 +1,77 @@
mod value;
use litellm_cache::Error;
use py_literal::Value;
use crate::{StoredValue, ValueAdapter};
#[derive(Clone, Copy, Debug, Default)]
pub struct PythonDiskCacheAdapter;
impl PythonDiskCacheAdapter {
fn python_get_cache(value: StoredValue) -> Result<Option<Value>, Error> {
let value = match value {
StoredValue::Bytes(value) => Value::Bytes(value),
StoredValue::Text(value) => Value::String(value),
StoredValue::Integer(value) => Value::Integer(value.into()),
StoredValue::Float(value) => Value::Float(value),
StoredValue::Pickle(value) => value::from_pickle(&value)?,
};
if !value::is_truthy(&value) {
return Ok(None);
}
match value {
Value::String(text) => Ok(Some(
value::from_json_text(&text).unwrap_or(Value::String(text)),
)),
Value::Bytes(bytes) => match std::str::from_utf8(&bytes) {
Ok(text) => Ok(Some(
value::from_json_text(text).unwrap_or(Value::Bytes(bytes)),
)),
Err(_) => Ok(Some(Value::Bytes(bytes))),
},
value => Ok(Some(value)),
}
}
}
impl ValueAdapter for PythonDiskCacheAdapter {
fn read(&self, value: StoredValue) -> Result<Option<Vec<u8>>, Error> {
match value {
StoredValue::Text(value) => Ok((!value.is_empty()).then(|| value.into_bytes())),
StoredValue::Bytes(value) => Ok((!value.is_empty()).then_some(value)),
value => {
let Some(value) = Self::python_get_cache(value)? else {
return Ok(None);
};
value::to_json(&value).map(Some)
}
}
}
fn write(&self, payload: Vec<u8>) -> StoredValue {
StoredValue::Bytes(payload)
}
fn counter_seed(&self, value: Option<StoredValue>) -> Result<f64, Error> {
let Some(value) = value else {
return Ok(0.0);
};
let Some(value) = Self::python_get_cache(value)? else {
return Ok(0.0);
};
Ok(if value::is_int(&value) {
value::to_f64(&value).unwrap_or(0.0)
} else {
0.0
})
}
fn counter_value(&self, value: f64) -> StoredValue {
if value.fract() == 0.0 && value >= i64::MIN as f64 && value <= i64::MAX as f64 {
StoredValue::Integer(value as i64)
} else {
StoredValue::Float(value)
}
}
}

173
litellm-rust/crates/cache-disk/src/python/value.rs Normal file
View file

@ -0,0 +1,173 @@
use litellm_cache::Error;
use py_literal::Value;
use serde_json::{Map, Number};
pub(crate) fn from_pickle(bytes: &[u8]) -> Result<Value, Error> {
let value = serde_pickle::value_from_slice(bytes, Default::default())
.map_err(|_| Error::InvalidEntry)?;
from_pickle_value(value)
}
fn from_pickle_value(value: serde_pickle::Value) -> Result<Value, Error> {
match value {
serde_pickle::Value::None => Ok(Value::None),
serde_pickle::Value::Bool(value) => Ok(Value::Boolean(value)),
serde_pickle::Value::I64(value) => integer(value.to_string()),
serde_pickle::Value::Int(value) => integer(value.to_string()),
serde_pickle::Value::F64(value) => Ok(Value::Float(value)),
serde_pickle::Value::String(value) => Ok(Value::String(value)),
serde_pickle::Value::Bytes(value) => Ok(Value::Bytes(value)),
serde_pickle::Value::List(values) => values
.into_iter()
.map(from_pickle_value)
.collect::<Result<Vec<_>, _>>()
.map(Value::List),
serde_pickle::Value::Tuple(values) => values
.into_iter()
.map(from_pickle_value)
.collect::<Result<Vec<_>, _>>()
.map(Value::Tuple),
serde_pickle::Value::Set(values) => values
.into_iter()
.map(from_pickle_hashable)
.collect::<Result<Vec<_>, _>>()
.map(Value::Set),
serde_pickle::Value::FrozenSet(values) => values
.into_iter()
.map(from_pickle_hashable)
.collect::<Result<Vec<_>, _>>()
.map(Value::Set),
serde_pickle::Value::Dict(values) => values
.into_iter()
.map(|(key, value)| Ok((from_pickle_hashable(key)?, from_pickle_value(value)?)))
.collect::<Result<Vec<_>, Error>>()
.map(Value::Dict),
}
}
fn from_pickle_hashable(value: serde_pickle::HashableValue) -> Result<Value, Error> {
Ok(match value {
serde_pickle::HashableValue::None => Value::None,
serde_pickle::HashableValue::Bool(value) => Value::Boolean(value),
serde_pickle::HashableValue::I64(value) => integer(value.to_string())?,
serde_pickle::HashableValue::Int(value) => integer(value.to_string())?,
serde_pickle::HashableValue::F64(value) => Value::Float(value),
serde_pickle::HashableValue::Bytes(value) => Value::Bytes(value),
serde_pickle::HashableValue::String(value) => Value::String(value),
serde_pickle::HashableValue::Tuple(values) => Value::Tuple(
values
.into_iter()
.map(from_pickle_hashable)
.collect::<Result<Vec<_>, _>>()?,
),
serde_pickle::HashableValue::FrozenSet(values) => Value::Set(
values
.into_iter()
.map(from_pickle_hashable)
.collect::<Result<Vec<_>, _>>()?,
),
})
}
fn integer(value: String) -> Result<Value, Error> {
value.parse().map_err(|_| Error::InvalidEntry)
}
pub(crate) fn from_json(value: serde_json::Value) -> Value {
match value {
serde_json::Value::Null => Value::None,
serde_json::Value::Bool(value) => Value::Boolean(value),
serde_json::Value::Number(value) => {
if value.is_i64() || value.is_u64() {
integer(value.to_string())
.unwrap_or(Value::Float(value.as_f64().unwrap_or(f64::NAN)))
} else {
Value::Float(value.as_f64().unwrap_or(f64::NAN))
}
}
serde_json::Value::String(value) => Value::String(value),
serde_json::Value::Array(values) => {
Value::List(values.into_iter().map(from_json).collect())
}
serde_json::Value::Object(values) => Value::Dict(
values
.into_iter()
.map(|(key, value)| (Value::String(key), from_json(value)))
.collect(),
),
}
}
pub(crate) fn from_json_text(value: &str) -> Result<Value, Error> {
serde_json::from_str(value)
.map(from_json)
.map_err(|_| Error::InvalidEntry)
}
pub(crate) fn is_truthy(value: &Value) -> bool {
match value {
Value::None => false,
Value::Boolean(value) => *value,
Value::Integer(value) => value.to_string() != "0",
Value::Float(value) => *value != 0.0,
Value::Complex(value) => value.re != 0.0 || value.im != 0.0,
Value::String(value) => !value.is_empty(),
Value::Bytes(value) => !value.is_empty(),
Value::Tuple(value) | Value::List(value) | Value::Set(value) => !value.is_empty(),
Value::Dict(value) => !value.is_empty(),
}
}
pub(crate) fn is_int(value: &Value) -> bool {
matches!(value, Value::Integer(_) | Value::Boolean(_))
}
pub(crate) fn to_f64(value: &Value) -> Option<f64> {
match value {
Value::Integer(value) => value.to_string().parse().ok(),
Value::Boolean(value) => Some(if *value { 1.0 } else { 0.0 }),
_ => None,
}
}
pub(crate) fn to_json(value: &Value) -> Result<Vec<u8>, Error> {
serde_json::to_vec(&to_json_value(value)?).map_err(|_| Error::InvalidEntry)
}
fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
Ok(match value {
Value::None => serde_json::Value::Null,
Value::Boolean(value) => serde_json::Value::Bool(*value),
Value::Integer(value) => serde_json::Value::Number(
value
.to_string()
.parse::<Number>()
.map_err(|_| Error::InvalidEntry)?,
),
Value::Float(value) => {
serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry)?)
}
Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry),
Value::String(value) => serde_json::Value::String(value.clone()),
Value::Tuple(values) | Value::List(values) | Value::Set(values) => {
serde_json::Value::Array(
values
.iter()
.map(to_json_value)
.collect::<Result<Vec<_>, _>>()?,
)
}
Value::Dict(values) => {
let values = values
.iter()
.map(|(key, value)| {
let Value::String(key) = key else {
return Err(Error::InvalidEntry);
};
Ok((key.clone(), to_json_value(value)?))
})
.collect::<Result<Map<String, serde_json::Value>, _>>()?;
serde_json::Value::Object(values)
}
})
}

817
litellm-rust/crates/cache-disk/src/sqlite.rs Normal file
View file

@ -0,0 +1,817 @@
use std::{
collections::HashMap,
fs::{self, OpenOptions},
io::Write,
path::{Path, PathBuf},
sync::Mutex,
};
use litellm_cache::Error;
use rand::RngCore;
use rusqlite::{Connection, OptionalExtension, params, types::Value};
use crate::{DiskStore, StoredValue};
const MODE_RAW: i64 = 1;
const MODE_BINARY: i64 = 2;
const MODE_TEXT: i64 = 3;
const MODE_PICKLE: i64 = 4;
const DEFAULT_DISK_MIN_FILE_SIZE: i64 = 2_i64.pow(15);
const DEFAULT_SIZE_LIMIT: i64 = 2_i64.pow(30);
const DEFAULT_CULL_LIMIT: i64 = 10;
pub struct DiskcacheSqliteStore {
directory: PathBuf,
connection: Mutex<Connection>,
min_file_size: usize,
eviction_policy: String,
size_limit: i64,
cull_limit: i64,
statistics: bool,
}
struct StoredColumns {
size: i64,
mode: i64,
filename: Option<String>,
value: Option<Value>,
}
struct Row {
rowid: i64,
mode: i64,
filename: Option<String>,
value: Value,
}
impl DiskcacheSqliteStore {
pub fn open(directory: impl AsRef<Path>) -> Result<Self, Error> {
let directory = directory.as_ref().to_path_buf();
fs::create_dir_all(&directory).map_err(|_| Error::Unavailable)?;
let directory = std::path::absolute(&directory).map_err(|_| Error::Unavailable)?;
let database = directory.join("cache.db");
let connection = Connection::open(database).map_err(|_| Error::Unavailable)?;
connection
.busy_timeout(std::time::Duration::from_secs(60))
.map_err(|_| Error::Unavailable)?;
let mut settings = read_settings(&connection)?;
for (key, value) in default_settings() {
settings.entry(key).or_insert(value);
}
for (key, value) in settings
.iter()
.filter(|(key, _)| key.starts_with("sqlite_"))
{
apply_pragma(&connection, key, value)?;
}
connection
.execute_batch(
"CREATE TABLE IF NOT EXISTS Settings (
key TEXT NOT NULL UNIQUE,
value
)",
)
.map_err(|_| Error::Unavailable)?;
for (key, value) in &settings {
if !matches!(key.as_str(), "count" | "size" | "hits" | "misses") {
connection
.execute(
"INSERT OR REPLACE INTO Settings VALUES (?, ?)",
params![key, value],
)
.map_err(|_| Error::Unavailable)?;
}
}
for (key, value) in [
("count", Value::Integer(0)),
("size", Value::Integer(0)),
("hits", Value::Integer(0)),
("misses", Value::Integer(0)),
] {
connection
.execute(
"INSERT OR IGNORE INTO Settings VALUES (?, ?)",
params![key, value],
)
.map_err(|_| Error::Unavailable)?;
}
connection
.execute_batch(
"CREATE TABLE IF NOT EXISTS Cache (
rowid INTEGER PRIMARY KEY,
key BLOB,
raw INTEGER,
store_time REAL,
expire_time REAL,
access_time REAL,
access_count INTEGER DEFAULT 0,
tag BLOB,
size INTEGER DEFAULT 0,
mode INTEGER DEFAULT 0,
filename TEXT,
value BLOB
);
CREATE UNIQUE INDEX IF NOT EXISTS Cache_key_raw ON Cache(key, raw);
CREATE INDEX IF NOT EXISTS Cache_expire_time ON Cache(expire_time);",
)
.map_err(|_| Error::Unavailable)?;
let eviction_policy = setting_string(&settings, "eviction_policy")
.unwrap_or_else(|| "least-recently-stored".to_string());
match eviction_policy.as_str() {
"none" => {}
"least-recently-stored" => {
connection
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_store_time ON Cache(store_time)",
)
.map_err(|_| Error::Unavailable)?;
}
"least-recently-used" => {
connection
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_access_time ON Cache(access_time)",
)
.map_err(|_| Error::Unavailable)?;
}
"least-frequently-used" => {
connection
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_access_count ON Cache(access_count)",
)
.map_err(|_| Error::Unavailable)?;
}
_ => return Err(Error::Unavailable),
}
connection
.execute_batch(
"CREATE TRIGGER IF NOT EXISTS Settings_count_insert
AFTER INSERT ON Cache FOR EACH ROW BEGIN
UPDATE Settings SET value = value + 1
WHERE key = \"count\"; END;
CREATE TRIGGER IF NOT EXISTS Settings_count_delete
AFTER DELETE ON Cache FOR EACH ROW BEGIN
UPDATE Settings SET value = value - 1
WHERE key = \"count\"; END;
CREATE TRIGGER IF NOT EXISTS Settings_size_insert
AFTER INSERT ON Cache FOR EACH ROW BEGIN
UPDATE Settings SET value = value + NEW.size
WHERE key = \"size\"; END;
CREATE TRIGGER IF NOT EXISTS Settings_size_update
AFTER UPDATE ON Cache FOR EACH ROW BEGIN
UPDATE Settings
SET value = value + NEW.size - OLD.size
WHERE key = \"size\"; END;
CREATE TRIGGER IF NOT EXISTS Settings_size_delete
AFTER DELETE ON Cache FOR EACH ROW BEGIN
UPDATE Settings SET value = value - OLD.size
WHERE key = \"size\"; END;",
)
.map_err(|_| Error::Unavailable)?;
let min_file_size = setting_i64(&settings, "disk_min_file_size")
.unwrap_or(DEFAULT_DISK_MIN_FILE_SIZE)
.try_into()
.map_err(|_| Error::Unavailable)?;
let size_limit = setting_i64(&settings, "size_limit").unwrap_or(DEFAULT_SIZE_LIMIT);
let cull_limit = setting_i64(&settings, "cull_limit").unwrap_or(DEFAULT_CULL_LIMIT);
let statistics = setting_i64(&settings, "statistics").unwrap_or_default() != 0;
Ok(Self {
directory,
connection: Mutex::new(connection),
min_file_size,
eviction_policy,
size_limit,
cull_limit,
statistics,
})
}
fn set_locked(
&self,
connection: &Connection,
key: &str,
columns: StoredColumns,
expire_time: Option<f64>,
now: f64,
) -> Result<Vec<String>, Error> {
let mut cleanup = Vec::new();
if let Some(old_filename) = connection
.query_row(
"SELECT filename FROM Cache WHERE key = ? AND raw = 1",
params![key],
|row| row.get::<_, Option<String>>(0),
)
.optional()
.map_err(|_| Error::Unavailable)?
.flatten()
{
cleanup.push(old_filename);
}
let (size, mode, filename, value) =
(columns.size, columns.mode, columns.filename, columns.value);
let rowid = connection
.query_row(
"SELECT rowid FROM Cache WHERE key = ? AND raw = 1",
params![key],
|row| row.get::<_, i64>(0),
)
.optional()
.map_err(|_| Error::Unavailable)?;
if let Some(rowid) = rowid {
connection
.execute(
"UPDATE Cache SET store_time = ?, expire_time = ?, access_time = ?,
access_count = 0, tag = NULL, size = ?, mode = ?, filename = ?, value = ?
WHERE rowid = ?",
params![now, expire_time, now, size, mode, filename, value, rowid],
)
.map_err(|_| Error::Unavailable)?;
} else {
connection
.execute(
"INSERT INTO Cache(
key, raw, store_time, expire_time, access_time, access_count,
tag, size, mode, filename, value
) VALUES (?, 1, ?, ?, ?, 0, NULL, ?, ?, ?, ?)",
params![key, now, expire_time, now, size, mode, filename, value],
)
.map_err(|_| Error::Unavailable)?;
}
cleanup.extend(self.cull(connection, now)?);
Ok(cleanup)
}
fn cull(&self, connection: &Connection, now: f64) -> Result<Vec<String>, Error> {
if self.cull_limit <= 0 {
return Ok(Vec::new());
}
let mut cleanup = Vec::new();
let expired = connection
.prepare(
"SELECT rowid, filename FROM Cache
WHERE expire_time IS NOT NULL AND expire_time < ?
ORDER BY expire_time LIMIT ?",
)
.map_err(|_| Error::Unavailable)?
.query_map(params![now, self.cull_limit], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
for (_, filename) in &expired {
if let Some(filename) = filename {
cleanup.push(filename.clone());
}
}
for (rowid, _) in &expired {
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
.map_err(|_| Error::Unavailable)?;
}
let remaining = self.cull_limit - i64::try_from(expired.len()).unwrap_or(self.cull_limit);
if remaining <= 0 || self.volume(connection)? < self.size_limit {
return Ok(cleanup);
}
let order = match self.eviction_policy.as_str() {
"none" => return Ok(cleanup),
"least-recently-stored" => "store_time",
"least-recently-used" => "access_time",
"least-frequently-used" => "access_count",
_ => return Err(Error::Unavailable),
};
let rows = connection
.prepare(&format!(
"SELECT rowid, filename FROM Cache ORDER BY {order} LIMIT ?"
))
.map_err(|_| Error::Unavailable)?
.query_map(params![remaining], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
for (_, filename) in &rows {
if let Some(filename) = filename {
cleanup.push(filename.clone());
}
}
for (rowid, _) in rows {
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
.map_err(|_| Error::Unavailable)?;
}
Ok(cleanup)
}
fn volume(&self, connection: &Connection) -> Result<i64, Error> {
let page_count: i64 = connection
.query_row("PRAGMA page_count", [], |row| row.get(0))
.map_err(|_| Error::Unavailable)?;
let page_size: i64 = connection
.query_row("PRAGMA page_size", [], |row| row.get(0))
.map_err(|_| Error::Unavailable)?;
let size: i64 = connection
.query_row("SELECT value FROM Settings WHERE key = 'size'", [], |row| {
row.get(0)
})
.map_err(|_| Error::Unavailable)?;
Ok(page_count.saturating_mul(page_size).saturating_add(size))
}
}
impl DiskStore for DiskcacheSqliteStore {
fn directory(&self) -> &Path {
&self.directory
}
fn get(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let select = "SELECT rowid, expire_time, mode, filename, value FROM Cache
WHERE key = ? AND raw = 1 AND (expire_time IS NULL OR expire_time > ?)";
let row = connection
.query_row(select, params![key, now], row_from_query)
.optional()
.map_err(|_| Error::Unavailable)?;
if !self.statistics && !has_get_update(&self.eviction_policy) {
return row
.map(|row| fetch_row(&self.directory, row))
.transpose()
.map(|value| value.flatten());
}
transactional(&connection, |connection| {
let row = connection
.query_row(select, params![key, now], row_from_query)
.optional()
.map_err(|_| Error::Unavailable)?;
let Some(row) = row else {
if self.statistics {
connection
.execute(
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
[],
)
.map_err(|_| Error::Unavailable)?;
}
return Ok(None);
};
let rowid = row.rowid;
let value = fetch_row(&self.directory, row);
let hit = value.as_ref().is_ok_and(Option::is_some);
if hit && self.statistics {
connection
.execute(
"UPDATE Settings SET value = value + 1 WHERE key = 'hits'",
[],
)
.map_err(|_| Error::Unavailable)?;
} else if !hit && self.statistics {
connection
.execute(
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
[],
)
.map_err(|_| Error::Unavailable)?;
}
if has_get_update(&self.eviction_policy) && hit {
let update = match self.eviction_policy.as_str() {
"least-recently-used" => "UPDATE Cache SET access_time = ? WHERE rowid = ?",
"least-frequently-used" => {
"UPDATE Cache SET access_count = access_count + 1 WHERE rowid = ?"
}
_ => return Err(Error::Unavailable),
};
if self.eviction_policy == "least-recently-used" {
connection
.execute(update, params![now, rowid])
.map_err(|_| Error::Unavailable)?;
} else {
connection
.execute(update, params![rowid])
.map_err(|_| Error::Unavailable)?;
}
}
value
})
}
fn set(
&self,
key: &str,
value: StoredValue,
expire_time: Option<f64>,
now: f64,
) -> Result<(), Error> {
let columns = store_value(&self.directory, self.min_file_size, value)?;
let new_filename = columns.filename.clone();
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let result = transactional(&connection, |connection| {
self.set_locked(connection, key, columns, expire_time, now)
});
match result {
Ok(cleanup) => {
cleanup_files(&self.directory, cleanup);
Ok(())
}
Err(error) => {
if let Some(filename) = new_filename {
remove_file(&self.directory, &filename);
}
Err(error)
}
}
}
fn pop(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let selected = transactional(&connection, |connection| {
let row = connection
.query_row(
"SELECT rowid, expire_time, mode, filename, value FROM Cache
WHERE key = ? AND raw = 1
AND (expire_time IS NULL OR expire_time > ?)",
params![key, now],
row_from_query,
)
.optional()
.map_err(|_| Error::Unavailable)?;
let Some(row) = row else {
return Ok(None);
};
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![row.rowid])
.map_err(|_| Error::Unavailable)?;
Ok(Some(row))
})?;
let Some(row) = selected else {
return Ok(None);
};
let filename = row.filename.clone();
let result = fetch_row(&self.directory, row)?;
if let Some(filename) = filename {
remove_file(&self.directory, &filename);
}
Ok(result)
}
fn clear(&self) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let mut last_rowid = 0_i64;
loop {
let batch = transactional(&connection, |connection| {
let rows = connection
.prepare(
"SELECT rowid, filename FROM Cache
WHERE rowid > ? ORDER BY rowid LIMIT 100",
)
.map_err(|_| Error::Unavailable)?
.query_map(params![last_rowid], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
if rows.is_empty() {
return Ok(rows);
}
let ids = rows
.iter()
.map(|(rowid, _)| rowid.to_string())
.collect::<Vec<_>>()
.join(",");
connection
.execute(&format!("DELETE FROM Cache WHERE rowid IN ({ids})"), [])
.map_err(|_| Error::Unavailable)?;
Ok(rows)
})?;
if batch.is_empty() {
return Ok(());
}
last_rowid = batch.last().map(|(rowid, _)| *rowid).unwrap_or(last_rowid);
cleanup_files(
&self.directory,
batch
.into_iter()
.filter_map(|(_, filename)| filename)
.collect(),
);
}
}
fn update(
&self,
key: &str,
now: f64,
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let mut created_filename = None;
let result = transactional(&connection, |connection| {
let current = connection
.query_row(
"SELECT rowid, expire_time, mode, filename, value FROM Cache
WHERE key = ? AND raw = 1
AND (expire_time IS NULL OR expire_time > ?)",
params![key, now],
row_from_query,
)
.optional()
.map_err(|_| Error::Unavailable)?
.map(|row| fetch_row(&self.directory, row))
.transpose()?
.flatten();
let (value, expire_time) = apply(current)?;
let columns = store_value(&self.directory, self.min_file_size, value)?;
created_filename = columns.filename.clone();
let cleanup = self.set_locked(connection, key, columns, expire_time, now)?;
Ok(cleanup)
});
match result {
Ok(cleanup) => {
cleanup_files(&self.directory, cleanup);
Ok(())
}
Err(error) => {
if let Some(filename) = created_filename {
remove_file(&self.directory, &filename);
}
Err(error)
}
}
}
fn probe(&self) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
connection
.query_row(
"SELECT value FROM Settings WHERE key = 'count'",
[],
|row| row.get::<_, i64>(0),
)
.map(|_| ())
.map_err(|_| Error::Unavailable)
}
}
fn default_settings() -> HashMap<String, Value> {
HashMap::from([
("statistics".to_string(), Value::Integer(0)),
("tag_index".to_string(), Value::Integer(0)),
(
"eviction_policy".to_string(),
Value::Text("least-recently-stored".to_string()),
),
("size_limit".to_string(), Value::Integer(DEFAULT_SIZE_LIMIT)),
("cull_limit".to_string(), Value::Integer(DEFAULT_CULL_LIMIT)),
("sqlite_auto_vacuum".to_string(), Value::Integer(1)),
("sqlite_cache_size".to_string(), Value::Integer(8192)),
(
"sqlite_journal_mode".to_string(),
Value::Text("wal".to_string()),
),
(
"sqlite_mmap_size".to_string(),
Value::Integer(2_i64.pow(26)),
),
("sqlite_synchronous".to_string(), Value::Integer(1)),
(
"disk_min_file_size".to_string(),
Value::Integer(DEFAULT_DISK_MIN_FILE_SIZE),
),
("disk_pickle_protocol".to_string(), Value::Integer(5)),
])
}
fn read_settings(connection: &Connection) -> Result<HashMap<String, Value>, Error> {
let mut statement = match connection.prepare("SELECT key, value FROM Settings") {
Ok(statement) => statement,
Err(_) => return Ok(HashMap::new()),
};
statement
.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))
.map_err(|_| Error::Unavailable)?
.collect::<Result<HashMap<_, _>, _>>()
.map_err(|_| Error::Unavailable)
}
fn apply_pragma(connection: &Connection, key: &str, value: &Value) -> Result<(), Error> {
let pragma = key.strip_prefix("sqlite_").ok_or(Error::Unavailable)?;
match value {
Value::Integer(value) => connection
.pragma_update(None, pragma, value)
.map_err(|_| Error::Unavailable),
Value::Text(value) => connection
.pragma_update(None, pragma, value)
.map_err(|_| Error::Unavailable),
_ => Err(Error::Unavailable),
}
}
fn setting_i64(settings: &HashMap<String, Value>, key: &str) -> Option<i64> {
match settings.get(key) {
Some(Value::Integer(value)) => Some(*value),
_ => None,
}
}
fn setting_string(settings: &HashMap<String, Value>, key: &str) -> Option<String> {
match settings.get(key) {
Some(Value::Text(value)) => Some(value.clone()),
_ => None,
}
}
fn has_get_update(policy: &str) -> bool {
matches!(policy, "least-recently-used" | "least-frequently-used")
}
fn row_from_query(row: &rusqlite::Row<'_>) -> rusqlite::Result<Row> {
Ok(Row {
rowid: row.get(0)?,
mode: row.get(2)?,
filename: row.get(3)?,
value: row.get(4)?,
})
}
fn fetch_row(directory: &Path, row: Row) -> Result<Option<StoredValue>, Error> {
match row.mode {
MODE_RAW => match row.value {
Value::Blob(value) => Ok(Some(StoredValue::Bytes(value))),
Value::Text(value) => Ok(Some(StoredValue::Text(value))),
Value::Integer(value) => Ok(Some(StoredValue::Integer(value))),
Value::Real(value) => Ok(Some(StoredValue::Float(value))),
Value::Null => Err(Error::InvalidEntry),
},
MODE_BINARY | MODE_PICKLE => {
let bytes = match row.value {
Value::Blob(value) => value,
Value::Null => {
let Some(value) = read_file(directory, row.filename.as_deref())? else {
return Ok(None);
};
value
}
_ => return Err(Error::InvalidEntry),
};
Ok(Some(if row.mode == MODE_BINARY {
StoredValue::Bytes(bytes)
} else {
StoredValue::Pickle(bytes)
}))
}
MODE_TEXT => {
let bytes = match row.value {
Value::Null => {
let Some(value) = read_file(directory, row.filename.as_deref())? else {
return Ok(None);
};
value
}
Value::Blob(value) => value,
Value::Text(value) => value.into_bytes(),
_ => return Err(Error::InvalidEntry),
};
Ok(Some(StoredValue::Text(
String::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?,
)))
}
_ => Err(Error::InvalidEntry),
}
}
fn read_file(directory: &Path, filename: Option<&str>) -> Result<Option<Vec<u8>>, Error> {
let Some(filename) = filename else {
return Err(Error::InvalidEntry);
};
match fs::read(directory.join(filename)) {
Ok(value) => Ok(Some(value)),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(_) => Err(Error::Unavailable),
}
}
fn store_value(
directory: &Path,
min_file_size: usize,
value: StoredValue,
) -> Result<StoredColumns, Error> {
match value {
StoredValue::Integer(value) => Ok(StoredColumns {
size: 0,
mode: MODE_RAW,
filename: None,
value: Some(Value::Integer(value)),
}),
StoredValue::Float(value) => Ok(StoredColumns {
size: 0,
mode: MODE_RAW,
filename: None,
value: Some(Value::Real(value)),
}),
StoredValue::Text(value) if value.chars().count() < min_file_size => Ok(StoredColumns {
size: 0,
mode: MODE_RAW,
filename: None,
value: Some(Value::Text(value)),
}),
StoredValue::Text(value) => {
let bytes = value.into_bytes();
let filename = write_file(directory, &bytes)?;
Ok(StoredColumns {
size: i64::try_from(bytes.len()).map_err(|_| Error::Unavailable)?,
mode: MODE_TEXT,
filename: Some(filename),
value: None,
})
}
StoredValue::Bytes(value) if value.len() < min_file_size => Ok(StoredColumns {
size: 0,
mode: MODE_RAW,
filename: None,
value: Some(Value::Blob(value)),
}),
StoredValue::Bytes(value) => {
let filename = write_file(directory, &value)?;
Ok(StoredColumns {
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
mode: MODE_BINARY,
filename: Some(filename),
value: None,
})
}
StoredValue::Pickle(value) if value.len() < min_file_size => Ok(StoredColumns {
size: 0,
mode: MODE_PICKLE,
filename: None,
value: Some(Value::Blob(value)),
}),
StoredValue::Pickle(value) => {
let filename = write_file(directory, &value)?;
Ok(StoredColumns {
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
mode: MODE_PICKLE,
filename: Some(filename),
value: None,
})
}
}
}
fn write_file(directory: &Path, bytes: &[u8]) -> Result<String, Error> {
let mut random = [0_u8; 16];
rand::rngs::OsRng.fill_bytes(&mut random);
let hex = random
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>();
let filename = format!("{}/{}/{}.val", &hex[..2], &hex[2..4], &hex[4..]);
let path = directory.join(&filename);
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).map_err(|_| Error::Unavailable)?;
}
let mut file = OpenOptions::new()
.write(true)
.create_new(true)
.open(path)
.map_err(|_| Error::Unavailable)?;
file.write_all(bytes).map_err(|_| Error::Unavailable)?;
Ok(filename)
}
fn cleanup_files(directory: &Path, filenames: Vec<String>) {
for filename in filenames {
remove_file(directory, &filename);
}
}
fn remove_file(directory: &Path, filename: &str) {
let path = directory.join(filename);
let _ = fs::remove_file(&path);
}
fn transactional<T>(
connection: &Connection,
operation: impl FnOnce(&Connection) -> Result<T, Error>,
) -> Result<T, Error> {
connection
.execute_batch("BEGIN IMMEDIATE")
.map_err(|_| Error::Unavailable)?;
match operation(connection) {
Ok(value) => {
connection
.execute_batch("COMMIT")
.map_err(|_| Error::Unavailable)?;
Ok(value)
}
Err(error) => {
let _ = connection.execute_batch("ROLLBACK");
Err(error)
}
}
}

33
litellm-rust/crates/cache-disk/src/store.rs Normal file
View file

@ -0,0 +1,33 @@
use std::path::Path;
use litellm_cache::Error;
#[derive(Clone, Debug, PartialEq)]
pub enum StoredValue {
Bytes(Vec<u8>),
Text(String),
Integer(i64),
Float(f64),
Pickle(Vec<u8>),
}
pub trait DiskStore: Send + Sync + 'static {
fn directory(&self) -> &Path;
fn get(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error>;
fn set(
&self,
key: &str,
value: StoredValue,
expire_time: Option<f64>,
now: f64,
) -> Result<(), Error>;
fn pop(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error>;
fn clear(&self) -> Result<(), Error>;
fn update(
&self,
key: &str,
now: f64,
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
) -> Result<(), Error>;
fn probe(&self) -> Result<(), Error>;
}

431
litellm-rust/crates/cache-disk/tests/cache.rs Normal file
View file

@ -0,0 +1,431 @@
use std::{
fs,
path::{Path, PathBuf},
sync::Arc,
thread,
time::Duration,
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, ExactCacheContext,
FlushCache, JsonCodec,
};
use litellm_cache_disk::{DiskCache, DiskStore, DiskcacheSqliteStore, StoredValue, ValueAdapter};
use rstest::{fixture, rstest};
use rusqlite::Connection;
use serde_json::{Value, json};
use tempfile::TempDir;
struct Sandbox {
directory: TempDir,
}
#[fixture]
fn sandbox() -> Sandbox {
Sandbox {
directory: tempfile::tempdir().unwrap(),
}
}
impl Sandbox {
fn store(&self) -> DiskcacheSqliteStore {
DiskcacheSqliteStore::open(self.directory.path()).unwrap()
}
fn cache<V>(&self) -> DiskCache<JsonCodec<V>>
where
JsonCodec<V>: CacheCodec,
{
DiskCache::open(self.directory.path(), JsonCodec::new()).unwrap()
}
fn db(&self) -> Connection {
Connection::open(self.directory.path().join("cache.db")).unwrap()
}
fn value_files(&self) -> Vec<PathBuf> {
fn visit(directory: &Path, files: &mut Vec<PathBuf>) {
for entry in fs::read_dir(directory).unwrap() {
let path = entry.unwrap().path();
if path.is_dir() {
visit(&path, files);
} else if path.extension().is_some_and(|extension| extension == "val") {
files.push(path);
}
}
}
let mut files = Vec::new();
visit(self.directory.path(), &mut files);
files
}
}
#[rstest]
fn relative_store_directory_is_absolutized(sandbox: Sandbox) {
let relative = PathBuf::from(format!(
".litellm-cache-disk-{}",
sandbox
.directory
.path()
.file_name()
.unwrap()
.to_string_lossy()
));
let store = DiskcacheSqliteStore::open(&relative).unwrap();
assert!(store.directory().is_absolute());
assert!(store.directory().ends_with(&relative));
let directory = store.directory().to_path_buf();
drop(store);
fs::remove_dir_all(directory).unwrap();
}
#[derive(Clone, Copy, Debug, Default)]
struct TextAdapter;
impl ValueAdapter for TextAdapter {
fn read(&self, value: StoredValue) -> Result<Option<Vec<u8>>, litellm_cache::Error> {
match value {
StoredValue::Text(value) => Ok(Some(value.into_bytes())),
_ => Ok(None),
}
}
fn write(&self, payload: Vec<u8>) -> StoredValue {
StoredValue::Text(String::from_utf8(payload).unwrap())
}
fn counter_seed(&self, _: Option<StoredValue>) -> Result<f64, litellm_cache::Error> {
Ok(0.0)
}
fn counter_value(&self, value: f64) -> StoredValue {
if value.fract() == 0.0 {
StoredValue::Integer(value as i64)
} else {
StoredValue::Float(value)
}
}
}
#[rstest]
fn roundtrip_persists_and_reopens(sandbox: Sandbox) {
let context = ExactCacheContext::default();
let opened = sandbox.cache::<Value>();
opened
.set_cache("key", json!({"answer": 42}), &context)
.unwrap();
assert_eq!(
opened.get_cache("key", &context).unwrap(),
Some(json!({"answer": 42}))
);
drop(opened);
let reopened = sandbox.cache::<Value>();
assert_eq!(
reopened.get_cache("key", &context).unwrap(),
Some(json!({"answer": 42}))
);
}
#[rstest]
fn ttl_and_expired_culling_match_cache_contract(sandbox: Sandbox) {
let store = sandbox.store();
store
.set(
"expired",
StoredValue::Bytes(b"old".to_vec()),
Some(10.0),
0.0,
)
.unwrap();
assert_eq!(store.get("expired", 10.0).unwrap(), None);
store
.set("new", StoredValue::Bytes(b"new".to_vec()), None, 11.0)
.unwrap();
assert_eq!(
sandbox
.db()
.query_row("SELECT COUNT(*) FROM Cache", [], |row| row.get::<_, i64>(0))
.unwrap(),
1
);
assert_eq!(
sandbox
.db()
.query_row(
"SELECT value FROM Settings WHERE key = 'count'",
[],
|row| row.get::<_, i64>(0)
)
.unwrap(),
1
);
}
#[rstest]
fn batch_preserves_order_and_classifies_misses_and_invalid_values(sandbox: Sandbox) {
let store = sandbox.store();
store
.set(
"hit",
StoredValue::Bytes(br#"{"ok":true}"#.to_vec()),
None,
0.0,
)
.unwrap();
store
.set(
"invalid",
StoredValue::Pickle(vec![0x80, 0x05, 0x2e]),
None,
0.0,
)
.unwrap();
let entries = sandbox
.cache::<Value>()
.batch_get_cache(
&["hit".into(), "missing".into(), "invalid".into()],
&ExactCacheContext::default(),
)
.unwrap();
assert_eq!(
entries,
vec![
BatchEntry::Hit(json!({"ok": true})),
BatchEntry::Miss,
BatchEntry::Invalid
]
);
}
#[rstest]
#[case(StoredValue::Bytes(Vec::new()))]
#[case(StoredValue::Text(String::new()))]
#[case(StoredValue::Integer(0))]
#[case(StoredValue::Float(0.0))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x89, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x00, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x00, 0x94, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x5d, 0x94, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x7d, 0x94, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x29, 0x2e]))]
fn falsy_values_are_misses(sandbox: Sandbox, #[case] value: StoredValue) {
sandbox.store().set("key", value, None, 0.0).unwrap();
assert_eq!(
sandbox
.cache::<Value>()
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
None
);
}
#[rstest]
#[case(Some(StoredValue::Integer(2)), 1.5, 3.5, "real")]
#[case(Some(StoredValue::Integer(2)), 1.0, 3.0, "integer")]
#[case(Some(StoredValue::Float(3.5)), 1.0, 1.0, "integer")]
#[case(Some(StoredValue::Text("not a number".into())), 2.0, 2.0, "integer")]
#[case(Some(StoredValue::Text("5".into())), 2.0, 7.0, "integer")]
#[case(Some(StoredValue::Text("3.5".into())), 2.0, 2.0, "integer")]
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e])), 1.0, 2.0, "integer")]
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x02, 0x2e])), 1.0, 3.0, "integer")]
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e])), 1.0, 1.0, "integer")]
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e])), 4.0, 4.0, "integer")]
fn counters_follow_python_initialization(
sandbox: Sandbox,
#[case] initial: Option<StoredValue>,
#[case] amount: f64,
#[case] expected: f64,
#[case] sqlite_type: &str,
) {
if let Some(initial) = initial {
sandbox.store().set("counter", initial, None, 0.0).unwrap();
}
let cache = sandbox.cache::<f64>();
assert_eq!(
cache
.increment_cache("counter", amount, ExactCacheContext::default())
.unwrap(),
expected
);
assert_eq!(
sandbox
.db()
.query_row(
"SELECT typeof(value) FROM Cache WHERE key = 'counter'",
[],
|row| row.get::<_, String>(0)
)
.unwrap(),
sqlite_type
);
}
#[rstest]
fn counters_are_atomic_across_concurrent_callers(sandbox: Sandbox) {
let cache = Arc::new(sandbox.cache::<f64>());
let workers = (0..8)
.map(|_| {
let cache = Arc::clone(&cache);
thread::spawn(move || {
for _ in 0..25 {
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap();
}
})
})
.collect::<Vec<_>>();
for worker in workers {
worker.join().unwrap();
}
assert_eq!(
cache
.increment_cache("counter", 0.0, ExactCacheContext::default())
.unwrap(),
200.0
);
}
#[rstest]
fn fractional_then_integer_increment_follows_python_behavior(sandbox: Sandbox) {
let cache = sandbox.cache::<f64>();
assert_eq!(
cache
.increment_cache("counter", 3.5, ExactCacheContext::default())
.unwrap(),
3.5
);
assert_eq!(
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap(),
1.0
);
}
#[rstest]
fn increment_ttl_replacement_clears_expiry_without_ttl(sandbox: Sandbox) {
let cache = sandbox.cache::<f64>();
cache
.increment_cache(
"counter",
1.0,
ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
},
)
.unwrap();
assert!(
sandbox
.db()
.query_row(
"SELECT expire_time IS NOT NULL FROM Cache WHERE key = 'counter'",
[],
|row| row.get::<_, bool>(0)
)
.unwrap()
);
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap();
assert!(
!sandbox
.db()
.query_row(
"SELECT expire_time IS NOT NULL FROM Cache WHERE key = 'counter'",
[],
|row| row.get::<_, bool>(0)
)
.unwrap()
);
}
#[rstest]
fn custom_adapter_controls_storage_and_reads(sandbox: Sandbox) {
let cache = DiskCache::with_adapter(sandbox.store(), TextAdapter, JsonCodec::<Value>::new());
cache
.set_cache("key", json!({"answer": 42}), &ExactCacheContext::default())
.unwrap();
assert!(matches!(
sandbox.store().get("key", 0.0).unwrap(),
Some(StoredValue::Text(_))
));
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(json!({"answer": 42}))
);
}
#[rstest]
fn delete_flush_and_spilled_file_replacement_clean_up_storage(sandbox: Sandbox) {
let large = vec![b'x'; 32 * 1024];
sandbox
.store()
.set("large", StoredValue::Bytes(large.clone()), None, 0.0)
.unwrap();
assert_eq!(sandbox.value_files().len(), 1);
sandbox
.store()
.set(
"large",
StoredValue::Bytes(vec![b'y'; 32 * 1024]),
None,
0.0,
)
.unwrap();
assert_eq!(sandbox.value_files().len(), 1);
sandbox.store().pop("large", 0.0).unwrap();
assert!(sandbox.value_files().is_empty());
sandbox
.store()
.set("a", StoredValue::Bytes(large.clone()), None, 0.0)
.unwrap();
sandbox
.store()
.set("b", StoredValue::Bytes(large), None, 0.0)
.unwrap();
sandbox.store().clear().unwrap();
assert!(sandbox.value_files().is_empty());
}
#[rstest]
#[tokio::test]
async fn async_operations_connection_and_delete_match_sync_operations(sandbox: Sandbox) {
let cache = sandbox.cache::<Value>();
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
};
cache
.async_set_cache("a", json!(1), context.clone())
.await
.unwrap();
cache
.async_set_cache_pipeline(
vec![("b".into(), json!(2)), ("c".into(), json!(3))],
context.clone(),
)
.await
.unwrap();
assert_eq!(
cache.async_get_cache("a", &context).await.unwrap(),
Some(json!(1))
);
assert_eq!(
cache
.async_batch_get_cache(vec!["c".into(), "missing".into()], context.clone())
.await
.unwrap(),
vec![BatchEntry::Hit(json!(3)), BatchEntry::Miss]
);
cache.async_delete_cache("a").await.unwrap();
cache.async_flush_cache().await.unwrap();
assert_eq!(
cache.test_connection().await.unwrap().status,
litellm_cache::CacheConnectionStatus::Success
);
}

113
litellm-rust/crates/cache-disk/tests/python_compat.rs Normal file
View file

@ -0,0 +1,113 @@
use litellm_cache::Error;
use litellm_cache_disk::{PythonDiskCacheAdapter, StoredValue, ValueAdapter};
use rstest::rstest;
enum ReadExpectation {
Bytes(&'static [u8]),
Miss,
Invalid,
}
#[rstest]
#[case::pickled_dictionary_with_string_keys(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e]),
ReadExpectation::Bytes(br#"{"a":1}"#)
)]
#[case::pickled_list_of_integers(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x5d, 0x94, 0x28, 0x4b, 0x01, 0x4b, 0x02, 0x65, 0x2e]),
ReadExpectation::Bytes(br#"[1,2]"#)
)]
#[case::pickled_tuple_of_integers(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x07, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4b, 0x01, 0x4b, 0x02, 0x86, 0x94, 0x2e]),
ReadExpectation::Bytes(br#"[1,2]"#)
)]
#[case::pickled_set_of_integers(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8f, 0x94, 0x28, 0x4b, 0x01, 0x4b, 0x02, 0x90, 0x2e]),
ReadExpectation::Bytes(br#"[1,2]"#)
)]
#[case::pickled_response_envelope(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x28, 0x8c, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x94, 0x47, 0x3f, 0xf8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x08, 0x72, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x94, 0x8c, 0x08, 0x7b, 0x22, 0x61, 0x22, 0x3a, 0x20, 0x31, 0x7d, 0x94, 0x75, 0x2e]),
ReadExpectation::Bytes(br#"{"response":"{\"a\": 1}","timestamp":1.5}"#)
)]
#[case::non_json_text(
StoredValue::Text("not json".into()),
ReadExpectation::Bytes(b"not json")
)]
#[case::json_text(
StoredValue::Text("{\"a\": 1}".into()),
ReadExpectation::Bytes(br#"{"a": 1}"#)
)]
#[case::non_utf8_bytes(
StoredValue::Bytes(vec![0xff, 0xfe]),
ReadExpectation::Bytes(&[0xff, 0xfe])
)]
#[case::integer_seven(StoredValue::Integer(7), ReadExpectation::Bytes(b"7"))]
#[case::float_one_point_five(StoredValue::Float(1.5), ReadExpectation::Bytes(b"1.5"))]
#[case::pickled_true(
StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e]),
ReadExpectation::Bytes(b"true")
)]
#[case::pickled_negative_integer(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4a, 0xfd, 0xff, 0xff, 0xff, 0x2e]),
ReadExpectation::Bytes(b"-3")
)]
#[case::pickled_bytes(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x43, 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, 0x94, 0x2e]),
ReadExpectation::Invalid
)]
#[case::pickled_dictionary_with_integer_key(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x4b, 0x01, 0x8c, 0x01, 0x61, 0x94, 0x73, 0x2e]),
ReadExpectation::Invalid
)]
#[case::pickled_complex(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x2e, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x08, 0x62, 0x75, 0x69, 0x6c, 0x74, 0x69, 0x6e, 0x73, 0x94, 0x8c, 0x07, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x94, 0x93, 0x94, 0x47, 0x3f, 0xf0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x86, 0x94, 0x52, 0x94, 0x2e]),
ReadExpectation::Invalid
)]
#[case::truncated_pickle(
StoredValue::Pickle(vec![0x80, 0x05, 0x2e]),
ReadExpectation::Invalid
)]
#[case::empty_bytes(StoredValue::Bytes(Vec::new()), ReadExpectation::Miss)]
#[case::empty_text(StoredValue::Text(String::new()), ReadExpectation::Miss)]
#[case::zero_integer(StoredValue::Integer(0), ReadExpectation::Miss)]
#[case::zero_float(StoredValue::Float(0.0), ReadExpectation::Miss)]
#[case::pickled_none(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_false(StoredValue::Pickle(vec![0x80, 0x05, 0x89, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_zero(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x00, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_zero_float(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_empty_string(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x00, 0x94, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_empty_list(StoredValue::Pickle(vec![0x80, 0x05, 0x5d, 0x94, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_empty_dictionary(StoredValue::Pickle(vec![0x80, 0x05, 0x7d, 0x94, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_empty_tuple(StoredValue::Pickle(vec![0x80, 0x05, 0x29, 0x2e]), ReadExpectation::Miss)]
fn python_read_cases(#[case] row: StoredValue, #[case] expected: ReadExpectation) {
let result = PythonDiskCacheAdapter.read(row);
match expected {
ReadExpectation::Bytes(expected) => assert_eq!(result.unwrap().unwrap(), expected),
ReadExpectation::Miss => assert_eq!(result.unwrap(), None),
ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry))),
}
}
#[rstest]
#[case::integer_two(Some(StoredValue::Integer(2)), 2.0)]
#[case::float_three_point_five(Some(StoredValue::Float(3.5)), 0.0)]
#[case::text_not_a_number(Some(StoredValue::Text("not a number".into())), 0.0)]
#[case::text_five(Some(StoredValue::Text("5".into())), 5.0)]
#[case::text_three_point_five(Some(StoredValue::Text("3.5".into())), 0.0)]
#[case::pickled_true(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e])), 1.0)]
#[case::pickled_two(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x02, 0x2e])), 2.0)]
#[case::pickled_dictionary(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e])), 0.0)]
#[case::missing(None, 0.0)]
#[case::pickled_none(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e])), 0.0)]
fn python_counter_seed_cases(#[case] row: Option<StoredValue>, #[case] expected: f64) {
assert_eq!(PythonDiskCacheAdapter.counter_seed(row).unwrap(), expected);
}
#[rstest]
#[case::integer_three(3.0, StoredValue::Integer(3))]
#[case::fractional_three_point_five(3.5, StoredValue::Float(3.5))]
#[case::negative_zero(-0.0, StoredValue::Integer(0))]
#[case::large_float(1e300, StoredValue::Float(1e300))]
fn python_counter_value_cases(#[case] value: f64, #[case] expected: StoredValue) {
assert_eq!(PythonDiskCacheAdapter.counter_value(value), expected);
}

20
litellm-rust/crates/cache-gcs/Cargo.toml Normal file
View file

@ -0,0 +1,20 @@
[package]
name = "litellm-cache-gcs"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
futures-util.workspace = true
litellm-auth-gcp.workspace = true
litellm-auth-types.workspace = true
litellm-cache.workspace = true
percent-encoding.workspace = true
reqwest.workspace = true
tokio.workspace = true
[dev-dependencies]
serde_json.workspace = true
tokio.workspace = true
wiremock = "0.6.5"

260
litellm-rust/crates/cache-gcs/src/cache.rs Normal file
View file

@ -0,0 +1,260 @@
use std::{future::Future, sync::Arc, time::Duration};
use futures_util::future::try_join_all;
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, Error, ExactCacheContext,
FlushCache,
};
use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode};
use reqwest::Client;
use crate::{GcpTokenSource, TokenSource};
pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com";
const OBJECT_NAME_ENCODE_SET: &AsciiSet = &NON_ALPHANUMERIC
.remove(b'-')
.remove(b'_')
.remove(b'.')
.remove(b'~');
pub fn key_prefix(gcs_path: Option<&str>) -> String {
match gcs_path {
Some(path) if !path.is_empty() => format!("{}/", path.trim_end_matches('/')),
_ => String::new(),
}
}
#[derive(Clone, Debug)]
pub struct GcsConfig {
pub bucket_name: String,
pub gcs_path: Option<String>,
pub path_service_account: Option<String>,
pub endpoint: String,
}
impl GcsConfig {
pub fn new(bucket_name: impl Into<String>) -> Self {
Self {
bucket_name: bucket_name.into(),
gcs_path: None,
path_service_account: None,
endpoint: DEFAULT_ENDPOINT.to_string(),
}
}
}
pub struct GcsCache<S: CacheCodec> {
config: GcsConfig,
key_prefix: String,
client: Client,
token: Arc<dyn TokenSource>,
codec: S,
}
impl<S: CacheCodec> GcsCache<S> {
pub fn new(config: GcsConfig, codec: S) -> Result<Self, Error> {
let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone()));
Self::with_token_source(config, codec, token)
}
pub fn with_token_source(
config: GcsConfig,
codec: S,
token: Arc<dyn TokenSource>,
) -> Result<Self, Error> {
let client = Client::builder().build().map_err(|_| Error::Unavailable)?;
let key_prefix = key_prefix(config.gcs_path.as_deref());
Ok(Self {
config,
key_prefix,
client,
token,
codec,
})
}
pub fn bucket_name(&self) -> &str {
&self.config.bucket_name
}
pub fn key_prefix(&self) -> &str {
&self.key_prefix
}
pub fn path_service_account(&self) -> Option<&str> {
self.config.path_service_account.as_deref()
}
pub fn object_name(&self, key: &str) -> String {
format!("{}{}", self.key_prefix, key)
}
fn encoded_object_name(&self, key: &str) -> String {
percent_encode(self.object_name(key).as_bytes(), OBJECT_NAME_ENCODE_SET).to_string()
}
fn endpoint(&self, path: &str) -> String {
format!("{}{}", self.config.endpoint.trim_end_matches('/'), path)
}
async fn async_set(&self, key: &str, value: S::Value) -> Result<(), Error> {
let token = self.token.bearer_token().await?;
let payload = self.codec.encode(&value)?;
let url = self.endpoint(&format!(
"/upload/storage/v1/b/{}/o?uploadType=media&name={}",
self.config.bucket_name,
self.encoded_object_name(key)
));
let response = self
.client
.post(url)
.bearer_auth(token)
.header(reqwest::header::CONTENT_TYPE, "application/json")
.body(payload)
.send()
.await
.map_err(|_| Error::Unavailable)?;
if !response.status().is_success() {
return Err(Error::Unavailable);
}
Ok(())
}
async fn async_get(&self, key: &str) -> Result<Option<S::Value>, Error> {
let token = self.token.bearer_token().await?;
let url = self.endpoint(&format!(
"/storage/v1/b/{}/o/{}?alt=media",
self.config.bucket_name,
self.encoded_object_name(key)
));
let response = self
.client
.get(url)
.bearer_auth(token)
.header(reqwest::header::CONTENT_TYPE, "application/json")
.send()
.await
.map_err(|_| Error::Unavailable)?;
if response.status() == reqwest::StatusCode::NOT_FOUND {
return Ok(None);
}
if !response.status().is_success() {
return Err(Error::Unavailable);
}
let body = response.bytes().await.map_err(|_| Error::Unavailable)?;
self.codec
.decode(&body)
.map(Some)
.map_err(|_| Error::InvalidEntry)
}
fn run_sync<T, F>(future: F) -> Result<T, Error>
where
F: Future<Output = Result<T, Error>> + Send,
T: Send,
{
let run = || {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|_| Error::Unavailable)
.and_then(|runtime| runtime.block_on(future))
};
if let Ok(handle) = tokio::runtime::Handle::try_current() {
if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread {
return tokio::task::block_in_place(run);
}
return std::thread::scope(|scope| {
scope
.spawn(run)
.join()
.map_err(|_| Error::Unavailable)
.and_then(|result| result)
});
}
run()
}
}
impl<S: CacheCodec> BaseCache for GcsCache<S> {
type Value = S::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, _: &Self::Context) -> Option<Duration> {
None
}
fn set_cache(&self, key: &str, value: Self::Value, _: &Self::Context) -> Result<(), Error> {
Self::run_sync(self.async_set(key, value))
}
fn get_cache(&self, key: &str, _: &Self::Context) -> Result<Option<Self::Value>, Error> {
Self::run_sync(self.async_get(key))
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
_: Self::Context,
) -> Result<(), Error> {
self.async_set(key, value).await
}
async fn async_get_cache(
&self,
key: &str,
_: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.async_get(key).await
}
async fn async_set_cache_pipeline(
&self,
entries: Vec<(String, Self::Value)>,
context: Self::Context,
) -> Result<(), Error> {
try_join_all(entries.into_iter().map(|(key, value)| {
let context = context.clone();
async move { self.async_set_cache(&key, value, context).await }
}))
.await
.map(|_| ())
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Err(Error::UnsupportedOperation)
}
}
impl<S: CacheCodec> BatchCache for GcsCache<S> {
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
context: Self::Context,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
try_join_all(keys.into_iter().map(|key| {
let context = context.clone();
async move {
match self.async_get_cache(&key, &context).await {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
}
}
}))
.await
}
}
impl<S: CacheCodec> FlushCache for GcsCache<S> {
fn flush_cache(&self) -> Result<(), Error> {
Ok(())
}
}

5
litellm-rust/crates/cache-gcs/src/lib.rs Normal file
View file

@ -0,0 +1,5 @@
mod cache;
mod token;
pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix};
pub use token::{GcpTokenSource, StaticTokenSource, TokenSource};

44
litellm-rust/crates/cache-gcs/src/token.rs Normal file
View file

@ -0,0 +1,44 @@
use std::{future::Future, pin::Pin};
use litellm_auth_gcp::{VertexAuth, VertexConfig};
use litellm_auth_types::{InputSource, SecretValue, Sourced};
use litellm_cache::Error;
pub trait TokenSource: Send + Sync + 'static {
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>>;
}
pub struct GcpTokenSource {
auth: VertexAuth,
config: VertexConfig,
}
impl GcpTokenSource {
pub fn new(path_service_account: Option<String>) -> Self {
let credentials = path_service_account
.map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment));
Self {
auth: VertexAuth::default(),
config: VertexConfig::new(credentials, None, None),
}
}
}
impl TokenSource for GcpTokenSource {
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
Box::pin(async move {
self.auth
.access_token(&self.config, &|name| std::env::var(name).ok())
.await
.map_err(|_| Error::Unavailable)
})
}
}
pub struct StaticTokenSource(pub String);
impl TokenSource for StaticTokenSource {
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
Box::pin(async move { Ok(self.0.clone()) })
}
}

324
litellm-rust/crates/cache-gcs/tests/cache.rs Normal file
View file

@ -0,0 +1,324 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheContext, Error, ExactCacheContext, FlushCache,
JsonCodec,
};
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource, TokenSource, key_prefix};
use serde_json::json;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_bytes, header, method, path, query_param},
};
fn config(server: &MockServer, gcs_path: Option<&str>) -> GcsConfig {
GcsConfig {
bucket_name: "bucket".into(),
gcs_path: gcs_path.map(str::to_string),
path_service_account: None,
endpoint: server.uri(),
}
}
fn cache(server: &MockServer, gcs_path: Option<&str>) -> GcsCache<JsonCodec<serde_json::Value>> {
GcsCache::with_token_source(
config(server, gcs_path),
JsonCodec::new(),
Arc::new(StaticTokenSource("tok".into())),
)
.unwrap()
}
#[tokio::test]
async fn set_writes_encoded_object_and_headers() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.and(query_param("uploadType", "media"))
.and(header("authorization", "Bearer tok"))
.and(header("content-type", "application/json"))
.and(body_bytes(br#"{"value":"entry"}"#))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&server)
.await;
cache(&server, Some("cache/"))
.set_cache(
"team:a b/c",
json!({"value": "entry"}),
&ExactCacheContext::default(),
)
.unwrap();
let requests = server.received_requests().await.unwrap();
assert_eq!(requests.len(), 1);
assert_eq!(
requests[0].url.query(),
Some("uploadType=media&name=cache%2Fteam%3Aa%20b%2Fc")
);
}
#[tokio::test]
async fn get_maps_statuses_and_decode_failures() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/hit"))
.and(query_param("alt", "media"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/missing"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/server-error"))
.respond_with(ResponseTemplate::new(500))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/invalid"))
.respond_with(ResponseTemplate::new(200).set_body_string("not json"))
.mount(&server)
.await;
let cache = cache(&server, None);
assert_eq!(
cache
.get_cache("hit", &ExactCacheContext::default())
.unwrap(),
Some(json!({"value": "entry"}))
);
assert_eq!(
cache
.get_cache("missing", &ExactCacheContext::default())
.unwrap(),
None
);
assert_eq!(
cache
.get_cache("server-error", &ExactCacheContext::default())
.unwrap_err(),
Error::Unavailable
);
assert_eq!(
cache
.get_cache("invalid", &ExactCacheContext::default())
.unwrap_err(),
Error::InvalidEntry
);
}
#[test]
fn key_prefix_normalizes_paths() {
assert_eq!(key_prefix(None), "");
assert_eq!(key_prefix(Some("a/b/")), "a/b/");
assert_eq!(key_prefix(Some("a/b")), "a/b/");
assert_eq!(key_prefix(Some("")), "");
}
#[tokio::test]
async fn object_names_use_python_quote_encoding() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.and(query_param("uploadType", "media"))
.respond_with(ResponseTemplate::new(200))
.expect(2)
.mount(&server)
.await;
let cache = cache(&server, Some("p/"));
cache
.set_cache(
"a~b-c_d.e/f g%h",
json!({"value": "punctuation"}),
&ExactCacheContext::default(),
)
.unwrap();
cache
.set_cache(
"ключ",
json!({"value": "utf8"}),
&ExactCacheContext::default(),
)
.unwrap();
let requests = server.received_requests().await.unwrap();
let queries: Vec<_> = requests
.iter()
.filter_map(|request| request.url.query())
.collect();
assert!(queries.contains(&"uploadType=media&name=p%2Fa~b-c_d.e%2Ff%20g%25h"));
assert!(queries.contains(&"uploadType=media&name=p%2F%D0%BA%D0%BB%D1%8E%D1%87"));
}
#[tokio::test]
async fn ignores_ttl_and_writes_pipeline_concurrently() {
let server = MockServer::start().await;
for key in ["one", "two", "three"] {
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.and(query_param("uploadType", "media"))
.and(query_param("name", key))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&server)
.await;
}
let cache = cache(&server, None);
assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None);
assert_eq!(
cache.get_ttl(&ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5)))),
None
);
cache
.async_set_cache_pipeline(
vec![
("one".into(), json!({"key": "one"})),
("two".into(), json!({"key": "two"})),
("three".into(), json!({"key": "three"})),
],
ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5))),
)
.await
.unwrap();
}
#[tokio::test]
async fn async_batch_get_preserves_hits_misses_and_invalid_entries() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/hit"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/missing"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/invalid"))
.respond_with(ResponseTemplate::new(200).set_body_string("not json"))
.mount(&server)
.await;
assert_eq!(
cache(&server, None)
.async_batch_get_cache(
vec!["hit".into(), "missing".into(), "invalid".into()],
ExactCacheContext::default(),
)
.await
.unwrap(),
vec![
BatchEntry::Hit(json!({"value": "entry"})),
BatchEntry::Miss,
BatchEntry::Invalid,
]
);
}
#[tokio::test]
async fn lifecycle_operations_are_noops_and_connection_test_is_unsupported() {
let server = MockServer::start().await;
let cache = cache(&server, None);
assert_eq!(cache.flush_cache(), Ok(()));
assert_eq!(cache.disconnect().await, Ok(()));
assert_eq!(
cache.test_connection().await,
Err(Error::UnsupportedOperation)
);
}
#[test]
fn sync_operations_work_without_an_active_runtime() {
let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.unwrap();
let server = runtime.block_on(MockServer::start());
runtime.block_on(
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.respond_with(ResponseTemplate::new(200))
.mount(&server),
);
runtime.block_on(
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/key"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.mount(&server),
);
let cache = cache(&server, None);
cache
.set_cache(
"key",
json!({"value": "entry"}),
&ExactCacheContext::default(),
)
.unwrap();
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(json!({"value": "entry"}))
);
}
#[tokio::test(flavor = "multi_thread")]
async fn sync_operations_work_inside_a_multi_thread_runtime() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.respond_with(ResponseTemplate::new(200))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/key"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.mount(&server)
.await;
let cache = cache(&server, None);
cache
.set_cache(
"key",
json!({"value": "entry"}),
&ExactCacheContext::default(),
)
.unwrap();
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(json!({"value": "entry"}))
);
}
struct FailingTokenSource;
impl TokenSource for FailingTokenSource {
fn bearer_token(
&self,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String, Error>> + Send + '_>>
{
Box::pin(async { Err(Error::Unavailable) })
}
}
#[tokio::test]
async fn token_source_failure_skips_http() {
let server = MockServer::start().await;
let cache = GcsCache::with_token_source(
config(&server, None),
JsonCodec::<serde_json::Value>::new(),
Arc::new(FailingTokenSource),
)
.unwrap();
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap_err(),
Error::Unavailable
);
assert_eq!(server.received_requests().await.unwrap().len(), 0);
}

20
litellm-rust/crates/cache-s3/Cargo.toml Normal file
View file

@ -0,0 +1,20 @@
[package]
name = "litellm-cache-s3"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-auth-aws.workspace = true
aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] }
aws-credential-types = "1.3.0"
aws-smithy-types = "1.6.0"
aws-types = "1.6.0"
tokio.workspace = true
[dev-dependencies]
wiremock = "0.6.5"
serde_json.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

101
litellm-rust/crates/cache-s3/src/auth.rs Normal file
View file

@ -0,0 +1,101 @@
use aws_credential_types::{
Credentials as AwsCredentials,
provider::{ProvideCredentials, error::CredentialsError, future},
};
use litellm_auth_aws::{AwsAuthConfig, resolve_credentials};
#[derive(Clone)]
pub(crate) struct Credentials {
config: AwsAuthConfig,
env: fn(&str) -> Option<String>,
}
impl Credentials {
pub(crate) fn new(config: AwsAuthConfig) -> Self {
Self::with_env(config, |name| std::env::var(name).ok())
}
pub(crate) fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
Self { config, env }
}
}
impl ProvideCredentials for Credentials {
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
where
Self: 'a,
{
future::ProvideCredentials::new(async {
if let (Some(access_key_id), Some(secret_access_key)) = (
self.config.access_key_id.clone(),
self.config.secret_access_key.clone(),
) {
return Ok(AwsCredentials::new(
access_key_id,
secret_access_key,
self.config.session_token.clone(),
None,
"litellm-s3-cache",
));
}
resolve_credentials(self.config.clone(), &self.env)
.await
.map_err(|_| CredentialsError::provider_error("S3 cache authentication failed"))
})
}
}
impl std::fmt::Debug for Credentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Credentials").finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn explicit_keys_ignore_an_ambient_session_token() {
let provider = Credentials::with_env(
AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
|name| (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()),
);
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.access_key_id(), "key");
assert_eq!(credentials.secret_access_key(), "secret");
assert_eq!(credentials.session_token(), None);
}
#[tokio::test]
async fn explicit_keys_keep_their_session_token() {
let provider = Credentials::new(AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
session_token: Some("t".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
});
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.session_token(), Some("t"));
}
#[tokio::test]
async fn environment_keys_resolve_with_their_session_token() {
let provider = Credentials::with_env(AwsAuthConfig::default(), |name| match name {
"AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
"AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
"AWS_SESSION_TOKEN" => Some("env-token".to_string()),
_ => None,
});
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.access_key_id(), "env-key");
assert_eq!(credentials.secret_access_key(), "env-secret");
assert_eq!(credentials.session_token(), Some("env-token"));
}
}

220
litellm-rust/crates/cache-s3/src/cache.rs Normal file
View file

@ -0,0 +1,220 @@
use std::{
future::Future,
sync::Arc,
time::{Duration, SystemTime},
};
use aws_sdk_s3::{
config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation},
error::SdkError,
primitives::ByteStream,
};
use aws_smithy_types::{DateTime, date_time::Format};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::{
BaseCache, BatchCache, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
};
use tokio::runtime::Handle;
use crate::auth::Credentials;
pub struct S3Endpoint {
pub url: String,
}
pub struct S3CacheConfig {
pub bucket: String,
pub key_prefix: String,
pub region: String,
pub endpoint: Option<S3Endpoint>,
pub auth: AwsAuthConfig,
}
pub struct S3Cache<C: CacheCodec> {
client: aws_sdk_s3::Client,
codec: C,
runtime: Handle,
bucket: Arc<str>,
key_prefix: Arc<str>,
region: Arc<str>,
endpoint: Option<Arc<str>>,
}
impl<C: CacheCodec> S3Cache<C> {
pub fn new(config: S3CacheConfig, codec: C, runtime: Handle) -> Self {
let endpoint_url: Option<String> = config.endpoint.map(|endpoint| endpoint.url);
let base = aws_sdk_s3::Config::builder()
.behavior_version(BehaviorVersion::latest())
.region(Region::new(config.region.clone()))
.credentials_provider(Credentials::new(config.auth))
.request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
.response_checksum_validation(ResponseChecksumValidation::WhenRequired);
let builder = match &endpoint_url {
Some(url) => base.endpoint_url(url).force_path_style(true),
None => base,
};
Self {
client: aws_sdk_s3::Client::from_conf(builder.build()),
codec,
runtime,
bucket: config.bucket.into(),
key_prefix: config.key_prefix.into(),
region: config.region.into(),
endpoint: endpoint_url.map(Into::into),
}
}
pub fn bucket(&self) -> &str {
&self.bucket
}
pub fn key_prefix(&self) -> &str {
&self.key_prefix
}
pub fn region(&self) -> &str {
&self.region
}
pub fn endpoint(&self) -> Option<&str> {
self.endpoint.as_deref()
}
pub fn to_s3_key(&self, key: &str) -> String {
format!("{}{}", self.key_prefix, key.replace(':', "/"))
}
fn block_on<F: Future>(&self, future: F) -> F::Output {
if Handle::try_current().is_ok() {
tokio::task::block_in_place(|| self.runtime.block_on(future))
} else {
self.runtime.block_on(future)
}
}
async fn put(
&self,
key: &str,
value: C::Value,
context: &ExactCacheContext,
) -> Result<(), Error> {
let s3_key = self.to_s3_key(key);
let body = self.codec.encode(&value)?;
let request = self
.client
.put_object()
.bucket(self.bucket.as_ref())
.key(&s3_key)
.body(ByteStream::from(body))
.content_type("application/json")
.content_language("en")
.content_disposition(format!("inline; filename=\"{s3_key}.json\""));
let request = match context.ttl {
Some(ttl) => {
let seconds = ttl.as_secs_f64();
request
.cache_control(format!("immutable, max-age={seconds}, s-maxage={seconds}"))
.expires(DateTime::from(SystemTime::now() + ttl))
}
None => request.cache_control("immutable, max-age=31536000, s-maxage=31536000"),
};
request.send().await.map_err(|_| Error::Unavailable)?;
Ok(())
}
async fn get(&self, key: &str) -> Result<Option<C::Value>, Error> {
let output = match self
.client
.get_object()
.bucket(self.bucket.as_ref())
.key(self.to_s3_key(key))
.send()
.await
{
Ok(output) => output,
Err(error) => {
if let SdkError::ServiceError(service) = &error {
let status = error
.raw_response()
.map(|response| response.status().as_u16());
let not_found = service.err().is_no_such_key()
|| service.err().meta().code() == Some("AccessDenied")
|| status == Some(404)
|| status == Some(403);
if not_found {
return Ok(None);
}
}
return Err(Error::Unavailable);
}
};
if let Some(expires) = output.expires_string()
&& let Ok(expires) = DateTime::from_str(expires, Format::HttpDate)
&& expires < DateTime::from(SystemTime::now())
{
return Ok(None);
}
let bytes = output
.body
.collect()
.await
.map_err(|_| Error::Unavailable)?
.into_bytes();
self.codec.decode(&bytes).map(Some)
}
}
impl<C: CacheCodec> BaseCache for S3Cache<C> {
type Value = C::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &Self::Context,
) -> Result<(), Error> {
self.block_on(self.put(key, value, context))
}
fn get_cache(&self, key: &str, _context: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.block_on(self.get(key))
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: Self::Context,
) -> Result<(), Error> {
self.put(key, value, &context).await
}
async fn async_get_cache(
&self,
key: &str,
_context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.get(key).await
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Err(Error::UnsupportedOperation)
}
}
impl<C: CacheCodec> BatchCache for S3Cache<C> {}
impl<C: CacheCodec> FlushCache for S3Cache<C> {
fn flush_cache(&self) -> Result<(), Error> {
Ok(())
}
}

4
litellm-rust/crates/cache-s3/src/lib.rs Normal file
View file

@ -0,0 +1,4 @@
mod auth;
mod cache;
pub use cache::{S3Cache, S3CacheConfig, S3Endpoint};

278
litellm-rust/crates/cache-s3/tests/cache.rs Normal file
View file

@ -0,0 +1,278 @@
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, Error, ExactCacheContext, FlushCache, JsonCodec,
};
use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint};
use serde_json::{Value, json};
use tokio::runtime::Handle;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{method, path},
};
fn config(endpoint: String) -> S3CacheConfig {
S3CacheConfig {
bucket: "cache-bucket".to_string(),
key_prefix: "team/".to_string(),
region: "us-east-1".to_string(),
endpoint: Some(S3Endpoint { url: endpoint }),
auth: AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
}
}
fn cache(endpoint: &str) -> S3Cache<JsonCodec<Value>> {
S3Cache::new(
config(endpoint.to_string()),
JsonCodec::<Value>::new(),
Handle::current(),
)
}
async fn mock_server() -> MockServer {
let server = MockServer::start().await;
Mock::given(method("PUT"))
.respond_with(ResponseTemplate::new(200).insert_header("etag", "\"etag\""))
.mount(&server)
.await;
server
}
fn http_date_from(headers: &wiremock::http::HeaderMap, name: &str) -> Option<SystemTime> {
use aws_smithy_types::{DateTime, date_time::Format};
headers
.get(name)
.and_then(|value| DateTime::from_str(value.to_str().ok()?, Format::HttpDate).ok())
.map(|date| UNIX_EPOCH + Duration::new(date.secs() as u64, date.subsec_nanos()))
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn set_writes_python_metadata_with_and_without_ttl() {
let server = mock_server().await;
let cache = cache(&server.uri());
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(90)),
};
cache
.set_cache("alpha:beta", json!({"answer": 1}), &context)
.unwrap();
cache
.set_cache("plain", json!({"answer": 2}), &ExactCacheContext::default())
.unwrap();
let requests = server.received_requests().await.unwrap();
let ttl_request = requests
.iter()
.find(|request| request.url.path() == "/cache-bucket/team/alpha/beta")
.expect("ttl write should hit the converted S3 key");
assert_eq!(
ttl_request.headers["cache-control"].to_str().unwrap(),
"immutable, max-age=90, s-maxage=90"
);
assert_eq!(
ttl_request.headers["content-type"].to_str().unwrap(),
"application/json"
);
assert_eq!(
ttl_request.headers["content-language"].to_str().unwrap(),
"en"
);
assert_eq!(
ttl_request.headers["content-disposition"].to_str().unwrap(),
"inline; filename=\"team/alpha/beta.json\""
);
let expires = http_date_from(&ttl_request.headers, "expires").expect("ttl write sets Expires");
let remaining = expires.duration_since(SystemTime::now()).unwrap();
assert!(remaining > Duration::from_secs(60) && remaining <= Duration::from_secs(91));
assert_eq!(
serde_json::from_slice::<Value>(&ttl_request.body).unwrap(),
json!({"answer": 1})
);
let plain = requests
.iter()
.find(|request| request.url.path() == "/cache-bucket/team/plain")
.expect("no-ttl write should hit the converted S3 key");
assert_eq!(
plain.headers["cache-control"].to_str().unwrap(),
"immutable, max-age=31536000, s-maxage=31536000"
);
assert!(plain.headers.get("expires").is_none());
assert_eq!(
plain.headers["content-disposition"].to_str().unwrap(),
"inline; filename=\"team/plain.json\""
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_hit_miss_expired_and_invalid_entries() {
let server = mock_server().await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/hit"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 3})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/missing"))
.respond_with(
ResponseTemplate::new(404).set_body_string("<Error><Code>NoSuchKey</Code></Error>"),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/denied"))
.respond_with(
ResponseTemplate::new(403).set_body_string("<Error><Code>AccessDenied</Code></Error>"),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/expired"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("expires", "Thu, 01 Jan 1970 00:00:00 GMT")
.set_body_json(json!({"answer": 4})),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/malformed"))
.respond_with(ResponseTemplate::new(200).set_body_string("not a cache entry"))
.mount(&server)
.await;
let cache = cache(&server.uri());
let context = ExactCacheContext::default();
assert_eq!(
cache.get_cache("hit", &context).unwrap(),
Some(json!({"answer": 3}))
);
assert_eq!(cache.get_cache("missing", &context).unwrap(), None);
assert_eq!(cache.get_cache("denied", &context).unwrap(), None);
assert_eq!(cache.get_cache("expired", &context).unwrap(), None);
assert_eq!(
cache.get_cache("malformed", &context),
Err(Error::InvalidEntry)
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn batch_get_preserves_order_with_hits_misses_and_invalid() {
let server = mock_server().await;
for (key, status, body) in [
("first", 200, "{\"answer\": 1}"),
("invalid", 200, "garbage"),
] {
Mock::given(method("GET"))
.and(path(format!("/cache-bucket/team/{key}")))
.respond_with(ResponseTemplate::new(status).set_body_string(body))
.mount(&server)
.await;
}
Mock::given(method("GET"))
.and(path("/cache-bucket/team/miss"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
let cache = cache(&server.uri());
let context = ExactCacheContext::default();
let keys = vec![
"first".to_string(),
"miss".to_string(),
"invalid".to_string(),
];
let entries = cache.batch_get_cache(&keys, &context).unwrap();
assert_eq!(
entries,
vec![
BatchEntry::Hit(json!({"answer": 1})),
BatchEntry::Miss,
BatchEntry::Invalid,
]
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn unsupported_and_noop_capabilities_match_python() {
let server = mock_server().await;
let cache = cache(&server.uri());
assert_eq!(
cache.test_connection().await,
Err(Error::UnsupportedOperation)
);
cache.flush_cache().unwrap();
cache.disconnect().await.unwrap();
assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None);
assert_eq!(
cache.get_ttl(&ExactCacheContext {
ttl: Some(Duration::from_secs(45)),
}),
Some(Duration::from_secs(45))
);
assert!(server.received_requests().await.unwrap().is_empty());
}
#[test]
fn key_conversion_prefixes_and_splits_colons() {
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(1)
.enable_all()
.build()
.unwrap();
let _guard = runtime.enter();
let cache = S3Cache::new(
S3CacheConfig {
key_prefix: "team/".to_string(),
..config("http://localhost".to_string())
},
JsonCodec::<Value>::new(),
runtime.handle().clone(),
);
assert_eq!(cache.bucket(), "cache-bucket");
assert_eq!(cache.key_prefix(), "team/");
assert_eq!(cache.to_s3_key("a:b:c"), "team/a/b/c");
assert_eq!(cache.to_s3_key("plain"), "team/plain");
let unprefixed = S3Cache::new(
S3CacheConfig {
key_prefix: String::new(),
..config("http://localhost".to_string())
},
JsonCodec::<Value>::new(),
runtime.handle().clone(),
);
assert_eq!(unprefixed.to_s3_key("a:b"), "a/b");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn sync_methods_block_inside_and_outside_the_runtime() {
let server = mock_server().await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/key"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 9})))
.mount(&server)
.await;
let uri = server.uri();
let cache = tokio::task::spawn_blocking(move || {
let cache = cache(&uri);
let context = ExactCacheContext::default();
cache
.set_cache("key", json!({"answer": 9}), &context)
.unwrap();
cache.get_cache("key", &context).unwrap()
})
.await
.unwrap();
assert_eq!(cache, Some(json!({"answer": 9})));
}

2
litellm-rust/crates/cache/src/error.rs vendored
View file

@ -6,4 +6,6 @@ pub enum Error {
InvalidEntry,
#[error("flushing Redis requires an explicit namespace")]
UnscopedFlush,
#[error("operation is not supported by this cache")]
UnsupportedOperation,
}

2
litellm-rust/crates/framer/Cargo.toml
View file

@ -11,7 +11,7 @@ aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"]
sse = ["dep:sse-stream"]
[dependencies]
aws-smithy-eventstream = { version = "=0.61.1", optional = true }
aws-smithy-eventstream = { version = "=0.61.4", optional = true }
aws-smithy-types = { version = "1.6.1", optional = true }
bytes = "1"
futures-util.workspace = true

2
litellm-rust/crates/llms/Cargo.toml
View file

@ -34,7 +34,7 @@ tokio = { workspace = true, features = ["sync"] }
url.workspace = true
[dev-dependencies]
aws-smithy-eventstream = "=0.61.1"
aws-smithy-eventstream = "=0.61.4"
aws-smithy-types = "1.6.1"
rstest.workspace = true
tokio.workspace = true

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

@ -24,9 +24,13 @@ litellm-cache.workspace = true
litellm-cache-azure-blob.workspace = true
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-s3.workspace = true
litellm-cache-gcs.workspace = true
litellm-cache-disk.workspace = true
litellm-cache-response.workspace = true
serde.workspace = true
litellm-auth.workspace = true
litellm-auth-aws.workspace = true
litellm-callbacks-legacy-python.workspace = true
litellm-core.workspace = true
litellm-core-utils.workspace = true

589
litellm-rust/crates/python-bridge/src/cache/config.rs vendored
View file

@ -1,11 +1,13 @@
use std::time::Duration;
use std::{path::PathBuf, time::Duration};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::CacheType;
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use pyo3::{
exceptions::{PyTypeError, PyValueError},
exceptions::{PyAttributeError, PyTypeError, PyValueError},
prelude::*,
types::{PyAny, PyDict, PyList, PyString},
types::{PyAny, PyBool, PyDict, PyList, PyString},
};
use super::{native::NativeResponseCache, request::duration};
@ -26,6 +28,10 @@ pub(super) struct MemoryCacheConfig {
pub(super) max_entry_bytes: usize,
}
pub(super) struct DiskCacheConfig {
pub(super) directory: PathBuf,
}
#[derive(Debug, PartialEq)]
pub(super) enum RedisProtocol {
Resp2,
@ -75,6 +81,13 @@ pub(super) struct RedisCacheConfig {
pub(super) connection: RedisConnectionConfig,
}
#[derive(Debug, PartialEq)]
pub(super) struct GcsCacheConfig {
pub(super) bucket_name: String,
pub(super) key_prefix: String,
pub(super) path_service_account: Option<String>,
}
pub(super) struct AzureBlobCacheConfig {
pub(super) account_url: String,
pub(super) container: String,
@ -94,6 +107,9 @@ const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31;
pub(super) enum CacheBackendConfig {
Memory(MemoryCacheConfig),
Redis(Box<RedisCacheConfig>),
S3(Box<S3CacheConfig>),
Gcs(GcsCacheConfig),
Disk(DiskCacheConfig),
AzureBlob(AzureBlobCacheConfig),
}
@ -109,6 +125,11 @@ pub(super) enum UnsupportedCacheConfig {
RedisCredentials,
RedisConnection,
RedisOption,
S3Client,
S3Credentials,
S3Option,
GcsBucket,
DiskStore,
}
impl UnsupportedCacheConfig {
@ -119,6 +140,11 @@ impl UnsupportedCacheConfig {
Self::RedisCredentials => "native Redis credentials require Python",
Self::RedisConnection => "native Redis connection type is not implemented",
Self::RedisOption => "native Redis configuration requires Python",
Self::S3Client => "native S3 client type is not implemented",
Self::S3Credentials => "native S3 credentials require Python",
Self::S3Option => "native S3 configuration requires Python",
Self::GcsBucket => "native GCS cache requires a configured bucket name",
Self::DiskStore => "native disk cache requires the built-in diskcache store",
}
}
}
@ -161,6 +187,27 @@ impl NativeCacheConfig {
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::S3) => match project_s3(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::S3(Box::new(backend)),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::Gcs) => match project_gcs(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::Gcs(backend),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::Disk) => match project_disk(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::Disk(backend),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::AzureBlob) => project_azure_blob(&backend).map(|backend| {
CacheConfigProjection::Native(Box::new(Self {
policy,
@ -168,12 +215,7 @@ impl NativeCacheConfig {
}))
}),
Some(
CacheType::RedisSemantic
| CacheType::ValkeySemantic
| CacheType::S3
| CacheType::Disk
| CacheType::QdrantSemantic
| CacheType::Gcs,
CacheType::RedisSemantic | CacheType::ValkeySemantic | CacheType::QdrantSemantic,
)
| None => Ok(CacheConfigProjection::Unsupported(
UnsupportedCacheConfig::Backend,
@ -185,7 +227,10 @@ impl NativeCacheConfig {
let default_ttl = match &self.backend {
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
CacheBackendConfig::AzureBlob(_) => None,
CacheBackendConfig::S3(_) => None,
CacheBackendConfig::Disk(_)
| CacheBackendConfig::AzureBlob(_)
| CacheBackendConfig::Gcs(_) => None,
};
if service.default_ttl() != default_ttl {
return Some("facade and native backend default TTLs must match");
@ -212,6 +257,66 @@ impl NativeCacheConfig {
CacheBackendConfig::Redis(config) => (service.namespace()
!= config.namespace.as_deref())
.then_some("facade and native backend namespaces must match"),
CacheBackendConfig::S3(_) if service.kind() != "s3" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::S3(config) if service.bucket() != Some(config.bucket.as_str()) => {
Some("facade and native backend buckets must match")
}
CacheBackendConfig::S3(config)
if service.key_prefix() != Some(config.key_prefix.as_str()) =>
{
Some("facade and native backend key prefixes must match")
}
CacheBackendConfig::S3(config) if service.region() != Some(config.region.as_str()) => {
Some("facade and native backend regions must match")
}
CacheBackendConfig::S3(config)
if service.endpoint()
!= config
.endpoint
.as_ref()
.map(|endpoint| endpoint.url.as_str()) =>
{
Some("facade and native backend endpoints must match")
}
CacheBackendConfig::S3(_) => None,
CacheBackendConfig::Gcs(_) if service.kind() != "gcs" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::Gcs(config)
if service
.gcs_backend()
.is_none_or(|backend| backend.bucket_name() != config.bucket_name) =>
{
Some("facade and native backend buckets must match")
}
CacheBackendConfig::Gcs(config)
if service
.gcs_backend()
.is_none_or(|backend| backend.key_prefix() != config.key_prefix) =>
{
Some("facade and native backend key prefixes must match")
}
CacheBackendConfig::Gcs(config)
if service.gcs_backend().is_none_or(|backend| {
backend.path_service_account() != config.path_service_account.as_deref()
}) =>
{
Some("facade and native backend credentials must match")
}
CacheBackendConfig::Gcs(_) => None,
CacheBackendConfig::Disk(_) if service.kind() != "disk" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::Disk(config) => {
let Some(directory) = service.directory() else {
return Some("facade and native backend types must match");
};
let native = std::fs::canonicalize(directory).ok();
let facade = std::fs::canonicalize(&config.directory).ok();
(native != facade).then_some("facade and native backend directories must match")
}
CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() {
None => Some("facade and native backend types must match"),
Some((account_url, container))
@ -252,6 +357,38 @@ fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
})
}
#[inline(never)]
fn project_gcs(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<GcsCacheConfig, UnsupportedCacheConfig>> {
let bucket_name = match backend.getattr("bucket_name")?.extract::<Option<String>>() {
Ok(Some(bucket_name)) if !bucket_name.is_empty() => bucket_name,
_ => return Ok(Err(UnsupportedCacheConfig::GcsBucket)),
};
Ok(Ok(GcsCacheConfig {
bucket_name,
key_prefix: backend.getattr("key_prefix")?.extract::<String>()?,
path_service_account: backend
.getattr("path_service_account")?
.extract::<Option<String>>()?,
}))
}
#[inline(never)]
fn project_disk(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<DiskCacheConfig, UnsupportedCacheConfig>> {
let store = backend.getattr("disk_cache")?;
if !instance_class_is(&store, "diskcache.core", "Cache")?
|| !instance_class_is(&store.getattr("_disk")?, "diskcache.core", "Disk")?
{
return Ok(Err(UnsupportedCacheConfig::DiskStore));
}
Ok(Ok(DiskCacheConfig {
directory: PathBuf::from(store.getattr("directory")?.extract::<String>()?),
}))
}
#[inline(never)]
fn project_redis(
backend: &Bound<'_, PyAny>,
@ -346,6 +483,77 @@ fn project_redis(
}))
}
#[inline(never)]
fn project_s3(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<S3CacheConfig, UnsupportedCacheConfig>> {
let client = backend.getattr("s3_client")?;
if !instance_class_is(&client, "botocore.client", "S3")? {
return Ok(Err(UnsupportedCacheConfig::S3Client));
}
let meta = client.getattr("meta")?;
let Some(region) = optional_string(meta.getattr("region_name")?)? else {
return Ok(Err(UnsupportedCacheConfig::S3Option));
};
let Some(endpoint_url) = optional_string(meta.getattr("endpoint_url")?)? else {
return Ok(Err(UnsupportedCacheConfig::S3Option));
};
let client_config = meta.getattr("config")?;
for name in ["s3", "proxies", "client_cert"] {
if optional_attribute(&client_config, name)?.is_some_and(|value| !value.is_none()) {
return Ok(Err(UnsupportedCacheConfig::S3Option));
}
}
let signature = match optional_attribute(&client_config, "signature_version")? {
Some(value) => value.extract::<Option<String>>()?,
None => None,
};
if signature.as_deref() != Some("s3v4") {
return Ok(Err(UnsupportedCacheConfig::S3Option));
}
let insecure = endpoint_url.starts_with("http://");
let verify = optional_attribute_chain(&client, &["_endpoint", "http_session", "_verify"])?;
let verified = verify
.and_then(|value| value.cast::<PyBool>().ok().map(|value| value.is_true()))
.unwrap_or(false);
if !verified && !insecure {
return Ok(Err(UnsupportedCacheConfig::S3Option));
}
let credentials = optional_attribute_chain(&client, &["_request_signer", "_credentials"])?
.ok_or(UnsupportedCacheConfig::S3Credentials);
let credentials = match credentials {
Ok(credentials) if !credentials.is_none() => credentials,
_ => return Ok(Err(UnsupportedCacheConfig::S3Credentials)),
};
let auth = if credentials.getattr("method")?.extract::<String>()?.as_str() == "explicit" {
AwsAuthConfig {
access_key_id: credentials
.getattr("access_key")?
.extract::<Option<String>>()?,
secret_access_key: credentials
.getattr("secret_key")?
.extract::<Option<String>>()?,
session_token: credentials.getattr("token")?.extract::<Option<String>>()?,
region_name: Some(region.clone()),
..Default::default()
}
} else {
AwsAuthConfig {
region_name: Some(region.clone()),
..Default::default()
}
};
let default_endpoint = endpoint_url == format!("https://s3.{region}.amazonaws.com")
|| (region == "us-east-1" && endpoint_url == "https://s3.amazonaws.com");
Ok(Ok(S3CacheConfig {
bucket: backend.getattr("bucket_name")?.extract::<String>()?,
key_prefix: backend.getattr("key_prefix")?.extract::<String>()?,
region,
endpoint: (!default_endpoint).then_some(S3Endpoint { url: endpoint_url }),
auth,
}))
}
#[inline(never)]
fn project_standalone_client<'py>(
client: &Bound<'py, PyAny>,
@ -549,6 +757,31 @@ fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult<O
}
}
#[inline(never)]
fn optional_attribute<'py>(
value: &Bound<'py, PyAny>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
match value.getattr(name) {
Ok(value) => Ok(Some(value)),
Err(error) if error.is_instance_of::<PyAttributeError>(value.py()) => Ok(None),
Err(error) => Err(error),
}
}
#[inline(never)]
fn optional_attribute_chain<'py>(
value: &Bound<'py, PyAny>,
names: &[&str],
) -> PyResult<Option<Bound<'py, PyAny>>> {
names
.iter()
.try_fold(Some(value.clone()), |current, name| match current {
Some(current) => optional_attribute(&current, name),
None => Ok(None),
})
}
#[inline(never)]
fn optional_string(value: Bound<'_, PyAny>) -> PyResult<Option<String>> {
Ok(value
@ -636,13 +869,16 @@ mod tests {
use pyo3::{prelude::*, types::PyDict};
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use litellm_host_python::run_sync_value;
use super::{
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig,
RedisProtocol,
CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement,
DiskCacheConfig, GcsCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
};
use crate::cache::native::NativeResponseCache;
use litellm_cache_redis::{RedisNode, RedisTopology};
fn cluster_facade<'py>(py: Python<'py>, startup_nodes: &str, hook: &str) -> Bound<'py, PyAny> {
facade(
@ -715,6 +951,71 @@ mod tests {
});
}
#[test]
fn projects_gcs_configuration() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"backend = SimpleNamespace(bucket_name='bucket', key_prefix='cache/', path_service_account='credentials.json')\n\
facade = SimpleNamespace(type='gcs', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("GCS cache should be supported");
};
let CacheBackendConfig::Gcs(gcs) = config.backend else {
panic!("expected GCS configuration");
};
assert_eq!(
gcs,
GcsCacheConfig {
bucket_name: "bucket".into(),
key_prefix: "cache/".into(),
path_service_account: Some("credentials.json".into()),
}
);
let matching = NativeResponseCache::gcs(
litellm_cache_gcs::GcsConfig {
bucket_name: "bucket".into(),
gcs_path: Some("cache/".into()),
path_service_account: Some("credentials.json".into()),
endpoint: litellm_cache_gcs::DEFAULT_ENDPOINT.into(),
},
Some("token".into()),
)
.unwrap();
let matching_config = NativeCacheConfig {
policy: config.policy,
backend: CacheBackendConfig::Gcs(gcs),
};
assert_eq!(matching_config.service_mismatch(&matching), None);
});
}
#[test]
fn rejects_gcs_without_a_bucket_name() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"backend = SimpleNamespace(bucket_name=None, key_prefix='', path_service_account=None)\n\
facade = SimpleNamespace(type='gcs', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("GCS cache without a bucket should be unsupported");
};
assert!(matches!(&reason, UnsupportedCacheConfig::GcsBucket));
assert_eq!(
reason.message(),
"native GCS cache requires a configured bucket name"
);
});
}
#[test]
fn projects_resolved_redis_tls_configuration() {
Python::initialize();
@ -775,6 +1076,268 @@ mod tests {
assert_eq!(reason.message(), "native Redis credentials require Python");
});
}
#[test]
fn projects_builtin_disk_configuration_and_rejects_custom_stores() {
Python::initialize();
Python::attach(|py| {
let root =
std::env::temp_dir().join(format!("litellm-disk-config-{}", std::process::id()));
let directory = root.to_string_lossy();
let disk_facade = facade(
py,
&format!(
"Cache = type('Cache', (), {{'__module__': 'diskcache.core'}})\n\
Disk = type('Disk', (), {{'__module__': 'diskcache.core'}})\n\
store = Cache()\n\
store._disk = Disk()\n\
store.directory = {directory:?}\n\
backend = SimpleNamespace(disk_cache=store)\n\
facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)"
),
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&disk_facade).unwrap()
else {
panic!("disk cache should be supported");
};
let CacheBackendConfig::Disk(disk) = config.backend else {
panic!("expected disk configuration");
};
assert_eq!(disk.directory, root);
let matching = NativeResponseCache::disk(&directory).unwrap();
assert_eq!(
(NativeCacheConfig {
policy: config.policy,
backend: CacheBackendConfig::Disk(disk),
})
.service_mismatch(&matching),
None
);
let other = NativeResponseCache::disk(&root.join("other").to_string_lossy()).unwrap();
let mismatch = NativeCacheConfig {
policy: CachePolicy {
mode: "default-on".into(),
ttl: None,
namespace: None,
supported_call_types: None,
redis_flush_size: None,
semantic_cache_scope: "key".into(),
},
backend: CacheBackendConfig::Disk(DiskCacheConfig {
directory: root.clone(),
}),
};
assert_eq!(
mismatch.service_mismatch(&other),
Some("facade and native backend directories must match")
);
let custom = facade(
py,
&format!(
"CustomCache = type('CustomCache', (), {{'__module__': 'mypkg'}})\n\
CustomDisk = type('CustomDisk', (), {{'__module__': 'mypkg'}})\n\
store = CustomCache()\n\
store._disk = CustomDisk()\n\
store.directory = {directory:?}\n\
backend = SimpleNamespace(disk_cache=store)\n\
facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)"
),
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&custom).unwrap()
else {
panic!("custom disk store must stay on Python");
};
assert_eq!(
reason.message(),
"native disk cache requires the built-in diskcache store"
);
});
}
fn s3_facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> {
let locals = PyDict::new(py);
py.run(
&CString::new(format!(
"from types import SimpleNamespace\n\
S3Client = type('S3', (), {{'__module__': 'botocore.client'}})\n\
client = S3Client()\n\
client.meta = SimpleNamespace(region_name='us-east-1', endpoint_url='https://example.test', config=SimpleNamespace(s3=None, proxies=None, client_cert=None, signature_version='s3v4'))\n\
client._endpoint = SimpleNamespace(http_session=SimpleNamespace(_verify=True))\n\
client._request_signer = SimpleNamespace(_credentials=SimpleNamespace(method='explicit', access_key='key', secret_key='secret', token='token'))\n\
backend = SimpleNamespace(bucket_name='bucket', key_prefix='team/', s3_client=client)\n\
facade = SimpleNamespace(type='s3', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)\n\
{body}"
))
.unwrap(),
None,
Some(&locals),
)
.unwrap();
locals.get_item("facade").unwrap().unwrap()
}
#[test]
fn projects_s3_configuration_with_explicit_credentials_and_custom_endpoint() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(py, "");
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("S3 cache should be supported");
};
let CacheBackendConfig::S3(s3) = config.backend else {
panic!("expected S3 configuration");
};
assert_eq!(s3.bucket, "bucket");
assert_eq!(s3.key_prefix, "team/");
assert_eq!(s3.region, "us-east-1");
assert_eq!(
s3.endpoint.map(|endpoint| endpoint.url).as_deref(),
Some("https://example.test")
);
assert_eq!(s3.auth.access_key_id.as_deref(), Some("key"));
assert_eq!(s3.auth.secret_access_key.as_deref(), Some("secret"));
assert_eq!(s3.auth.session_token.as_deref(), Some("token"));
assert_eq!(s3.auth.region_name.as_deref(), Some("us-east-1"));
});
}
#[test]
fn default_s3_endpoint_projects_no_custom_endpoint() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(
py,
"facade.cache.s3_client.meta.endpoint_url = 'https://s3.us-east-1.amazonaws.com'",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("S3 cache should be supported");
};
let CacheBackendConfig::S3(s3) = config.backend else {
panic!("expected S3 configuration");
};
assert!(s3.endpoint.is_none());
});
}
#[test]
fn non_sigv4_proxies_and_disabled_verification_stay_on_python() {
Python::initialize();
Python::attach(|py| {
for (body, message) in [
(
"facade.cache.s3_client.meta.config.signature_version = 's3'",
"native S3 configuration requires Python",
),
(
"facade.cache.s3_client.meta.config.proxies = {'https': 'proxy'}",
"native S3 configuration requires Python",
),
(
"facade.cache.s3_client._endpoint.http_session._verify = False",
"native S3 configuration requires Python",
),
(
"del facade.cache.s3_client._endpoint.http_session._verify",
"native S3 configuration requires Python",
),
] {
let facade = s3_facade(py, body);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("{body} must stay on Python");
};
assert_eq!(reason.message(), message);
}
let facade = s3_facade(py, "facade.cache.s3_client = SimpleNamespace()");
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("non-botocore client must stay on Python");
};
assert_eq!(reason.message(), "native S3 client type is not implemented");
let facade = s3_facade(
py,
"facade.cache.s3_client._request_signer._credentials = None",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("missing credentials must stay on Python");
};
assert_eq!(reason.message(), "native S3 credentials require Python");
});
}
#[test]
fn non_explicit_s3_credentials_use_the_default_chain() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(
py,
"facade.cache.s3_client._request_signer._credentials = SimpleNamespace(method='sso', access_key=None, secret_key=None, token=None)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("default-chain credentials should be supported");
};
let CacheBackendConfig::S3(s3) = config.backend else {
panic!("expected S3 configuration");
};
assert_eq!(s3.auth.access_key_id, None);
assert_eq!(s3.auth.secret_access_key, None);
assert_eq!(s3.auth.region_name.as_deref(), Some("us-east-1"));
});
}
fn s3_service(py: Python<'_>, region: &str, endpoint: Option<&str>) -> NativeResponseCache {
let config = S3CacheConfig {
bucket: "bucket".to_string(),
key_prefix: "team/".to_string(),
region: region.to_string(),
endpoint: endpoint.map(|url| S3Endpoint {
url: url.to_string(),
}),
auth: AwsAuthConfig::default(),
};
run_sync_value(py, async move { Ok(NativeResponseCache::s3(config).await) }).unwrap()
}
#[test]
fn s3_binding_rejects_region_and_endpoint_mismatches() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(py, "");
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("S3 cache should be supported");
};
assert_eq!(
config.service_mismatch(&s3_service(py, "us-east-1", Some("https://example.test"))),
None
);
assert_eq!(
config.service_mismatch(&s3_service(py, "us-west-2", Some("https://example.test"))),
Some("facade and native backend regions must match")
);
assert_eq!(
config.service_mismatch(&s3_service(py, "us-east-1", Some("https://other.test"))),
Some("facade and native backend endpoints must match")
);
assert_eq!(
config.service_mismatch(&s3_service(py, "us-east-1", None)),
Some("facade and native backend endpoints must match")
);
});
}
#[test]
fn projects_cluster_startup_nodes_as_redis_topology() {

68
litellm-rust/crates/python-bridge/src/cache/facade.rs vendored
View file

@ -34,6 +34,14 @@ struct RedisPoolGuard {
attributes: RedisPoolAttributes,
}
struct S3ClientGuard {
reference: Py<PyAny>,
}
struct DiskStoreGuard {
reference: Py<PyAny>,
directory: String,
}
struct AzureBlobClientGuard {
sync_client: Py<PyAny>,
async_client: Py<PyAny>,
@ -45,8 +53,8 @@ enum ConnectionGuard {
None,
RedisPool(RedisPoolGuard),
AzureBlob(AzureBlobClientGuard),
S3(S3ClientGuard),
}
struct RedisPoolAttributes {
pool: &'static str,
connection_class: &'static str,
@ -64,10 +72,10 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes {
connection_class: "connection_pool_class",
max_connections: None,
};
pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
disk_store: Option<DiskStoreGuard>,
connection: ConnectionGuard,
}
@ -218,6 +226,41 @@ impl RedisPoolGuard {
}
}
impl S3ClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
reference: backend.getattr("s3_client")?.unbind(),
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
Ok(self.reference.bind(py).is(&backend.getattr("s3_client")?))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reference)
}
}
impl DiskStoreGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let store = backend.getattr("disk_cache")?;
Ok(Self {
reference: store.clone().unbind(),
directory: store.getattr("directory")?.extract()?,
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
let store = backend.getattr("disk_cache")?;
Ok(self.reference.bind(py).is(&store)
&& self.directory == store.getattr("directory")?.extract::<String>()?)
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reference)
}
}
impl AzureBlobClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let sync_client = backend.getattr("container_client")?;
@ -252,6 +295,7 @@ impl ConnectionGuard {
("redis", false) => Self::RedisPool(RedisPoolGuard::capture(backend, STANDALONE_POOL)?),
("redis", true) => Self::RedisPool(RedisPoolGuard::capture(backend, CLUSTER_POOL)?),
("azure-blob", _) => Self::AzureBlob(AzureBlobClientGuard::capture(backend)?),
("s3", _) => Self::S3(S3ClientGuard::capture(backend)?),
_ => Self::None,
})
}
@ -261,6 +305,7 @@ impl ConnectionGuard {
Self::None => Ok(true),
Self::RedisPool(guard) => guard.matches(py, backend),
Self::AzureBlob(guard) => guard.matches(py, backend),
Self::S3(guard) => guard.matches(py, backend),
}
}
@ -269,10 +314,10 @@ impl ConnectionGuard {
Self::None => Ok(()),
Self::RedisPool(guard) => guard.traverse(visit),
Self::AzureBlob(guard) => guard.traverse(visit),
Self::S3(guard) => guard.traverse(visit),
}
}
}
impl FacadeGuard {
pub(super) fn capture(
py: Python<'_>,
@ -295,6 +340,9 @@ impl FacadeGuard {
"RedisClusterCache",
"redis",
),
("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"),
("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"),
("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"),
("azure-blob", _) => (
"litellm.caching.azure_blob_cache",
"AzureBlobCache",
@ -343,8 +391,14 @@ impl FacadeGuard {
"max_size_per_item",
"redis_kwargs",
"redis_flush_size",
"bucket_name",
"key_prefix",
"path_service_account",
],
)?,
disk_store: (kind == "disk")
.then(|| DiskStoreGuard::capture(&backend))
.transpose()?,
connection: ConnectionGuard::capture(kind, cluster, &backend)?,
})
}
@ -357,12 +411,20 @@ impl FacadeGuard {
if !self.backend.matches(py, &backend)? {
return Ok(false);
}
if let Some(guard) = &self.disk_store
&& !guard.matches(py, &backend)?
{
return Ok(false);
}
self.connection.matches(py, &backend)
}
pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
self.outer.traverse(&visit)?;
self.backend.traverse(&visit)?;
if let Some(guard) = &self.disk_store {
guard.traverse(&visit)?;
}
self.connection.traverse(&visit)
}
}

76
litellm-rust/crates/python-bridge/src/cache/handle.rs vendored
View file

@ -1,7 +1,11 @@
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use litellm_host_python::{release_gil, run_sync_value};
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration};
#[pyclass(frozen, name = "_CacheTestHandle")]
@ -64,6 +68,77 @@ impl CacheTestHandle {
})
}
#[staticmethod]
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (bucket, *, region, endpoint_url=None, key_prefix="", access_key_id=None, secret_access_key=None, session_token=None))]
fn s3(
py: Python<'_>,
bucket: String,
region: String,
endpoint_url: Option<String>,
key_prefix: &str,
access_key_id: Option<String>,
secret_access_key: Option<String>,
session_token: Option<String>,
) -> PyResult<Self> {
let config = S3CacheConfig {
bucket,
key_prefix: key_prefix.to_string(),
region: region.clone(),
endpoint: endpoint_url.map(|url| S3Endpoint { url }),
auth: AwsAuthConfig {
access_key_id,
secret_access_key,
session_token,
region_name: Some(region),
..Default::default()
},
};
let service = run_sync_value(py, async move { Ok(NativeResponseCache::s3(config).await) })?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (bucket_name, *, gcs_path=None, path_service_account=None, endpoint=None, token=None))]
fn gcs(
py: Python<'_>,
bucket_name: String,
gcs_path: Option<String>,
path_service_account: Option<String>,
endpoint: Option<String>,
token: Option<String>,
) -> PyResult<Self> {
let config = GcsConfig {
bucket_name,
gcs_path,
path_service_account,
endpoint: endpoint.unwrap_or_else(|| DEFAULT_ENDPOINT.to_string()),
};
let service = release_gil(py, move || NativeResponseCache::gcs(config, token))
.map_err(cache_error)?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (directory))]
fn disk(py: Python<'_>, directory: String) -> PyResult<Self> {
let service =
release_gil(py, move || NativeResponseCache::disk(&directory)).map_err(cache_error)?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (account_url, container))]
fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult<Self> {
@ -78,7 +153,6 @@ impl CacheTestHandle {
pid: std::process::id(),
})
}
#[getter]
fn backend(&self) -> &'static str {
self.service.kind()

3
litellm-rust/crates/python-bridge/src/cache/mod.rs vendored
View file

@ -10,7 +10,7 @@ mod resolver;
use litellm_cache::Error;
use pyo3::{
exceptions::{PyRuntimeError, PyValueError},
exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError},
prelude::*,
};
@ -21,6 +21,7 @@ pub(crate) use self::{
fn cache_error(error: Error) -> PyErr {
match error {
Error::InvalidEntry => PyValueError::new_err(error.to_string()),
Error::UnsupportedOperation => PyNotImplementedError::new_err(error.to_string()),
_ => PyRuntimeError::new_err(error.to_string()),
}
}

133
litellm-rust/crates/python-bridge/src/cache/native.rs vendored
View file

@ -1,12 +1,15 @@
use std::{sync::Arc, time::Duration};
use std::{path::Path, sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
use litellm_cache_azure_blob::AzureBlobCache;
use litellm_cache_disk::DiskCache;
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::{RedisCache, RedisTopology};
use litellm_cache_response::{
CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer,
};
use litellm_cache_s3::{S3Cache, S3CacheConfig};
use serde_json::Value;
#[derive(Clone)]
@ -16,6 +19,9 @@ pub(super) enum NativeResponseCache {
cache: Arc<ResponseCache<RedisCache<ResponseCacheCodec>>>,
buffer: Option<Arc<WriteBuffer>>,
},
S3(Arc<ResponseCache<S3Cache<ResponseCacheCodec>>>),
Gcs(Arc<ResponseCache<GcsCache<ResponseCacheCodec>>>),
Disk(Arc<ResponseCache<DiskCache<ResponseCacheCodec>>>),
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
}
@ -47,6 +53,30 @@ impl NativeResponseCache {
buffer: None,
})
}
pub async fn s3(config: S3CacheConfig) -> Self {
let runtime = tokio::runtime::Handle::current();
Self::S3(Arc::new(ResponseCache::new(Arc::new(S3Cache::new(
config,
ResponseCacheCodec,
runtime,
)))))
}
pub fn disk(directory: &str) -> Result<Self, Error> {
let cache = DiskCache::open(directory, ResponseCacheCodec)?;
Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache)))))
}
pub fn gcs(config: GcsConfig, token: Option<String>) -> Result<Self, Error> {
let backend = match token {
Some(token) => GcsCache::with_token_source(
config,
ResponseCacheCodec,
Arc::new(StaticTokenSource(token)),
)?,
None => GcsCache::new(config, ResponseCacheCodec)?,
};
Ok(Self::Gcs(Arc::new(ResponseCache::new(Arc::new(backend)))))
}
pub async fn azure_blob(account_url: &str, container: &str) -> Result<Self, Error> {
let backend = AzureBlobCache::connect(
@ -67,7 +97,9 @@ impl NativeResponseCache {
cache.backend().account_url(),
cache.backend().container_name(),
)),
Self::Memory(_) | Self::Redis { .. } => None,
Self::Memory(_) | Self::Redis { .. } | Self::S3(_) | Self::Disk(_) | Self::Gcs(_) => {
None
}
}
}
}
@ -77,6 +109,9 @@ impl NativeResponseCache {
match self {
Self::Memory(_) => "memory",
Self::Redis { .. } => "redis",
Self::S3(_) => "s3",
Self::Gcs(_) => "gcs",
Self::Disk(_) => "disk",
Self::AzureBlob(_) => "azure-blob",
}
}
@ -85,35 +120,76 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.default_ttl(),
Self::Redis { cache, .. } => cache.default_ttl(),
Self::S3(cache) => cache.default_ttl(),
Self::Gcs(cache) => cache.default_ttl(),
Self::Disk(cache) => cache.default_ttl(),
Self::AzureBlob(cache) => cache.default_ttl(),
}
}
pub fn bucket(&self) -> Option<&str> {
match self {
Self::S3(cache) => Some(cache.backend().bucket()),
_ => None,
}
}
pub fn key_prefix(&self) -> Option<&str> {
match self {
Self::S3(cache) => Some(cache.backend().key_prefix()),
_ => None,
}
}
pub fn region(&self) -> Option<&str> {
match self {
Self::S3(cache) => Some(cache.backend().region()),
_ => None,
}
}
pub fn endpoint(&self) -> Option<&str> {
match self {
Self::S3(cache) => cache.backend().endpoint(),
_ => None,
}
}
pub fn namespace(&self) -> Option<&str> {
match self {
Self::Memory(_) | Self::AzureBlob(_) => None,
Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) => None,
Self::Redis { cache, .. } => cache.backend().namespace(),
Self::S3(_) | Self::Gcs(_) => None,
}
}
pub fn topology(&self) -> Option<&RedisTopology> {
match self {
Self::Memory(_) | Self::AzureBlob(_) => None,
Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None,
Self::Redis { cache, .. } => Some(cache.backend().topology()),
Self::S3(_) => None,
}
}
pub fn capacity(&self) -> Option<usize> {
match self {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
Self::Redis { .. } | Self::AzureBlob(_) => None,
Self::Redis { .. }
| Self::S3(_)
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
}
}
pub fn max_entry_bytes(&self) -> Option<usize> {
match self {
Self::Memory(cache) => cache.backend().max_entry_bytes(),
Self::Redis { .. } | Self::AzureBlob(_) => None,
Self::Redis { .. }
| Self::S3(_)
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
}
}
@ -127,6 +203,17 @@ impl NativeResponseCache {
}
}
pub fn directory(&self) -> Option<&Path> {
match self {
Self::Disk(cache) => Some(cache.backend().directory()),
Self::Memory(_)
| Self::Redis { .. }
| Self::S3(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
}
}
pub fn lookup(
&self,
request: &ResponseCacheRequest,
@ -135,6 +222,9 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.lookup(request, now),
Self::Redis { cache, .. } => cache.lookup(request, now),
Self::S3(cache) => cache.lookup(request, now),
Self::Gcs(cache) => cache.lookup(request, now),
Self::Disk(cache) => cache.lookup(request, now),
Self::AzureBlob(cache) => cache.lookup(request, now),
}
}
@ -148,6 +238,9 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.store(request, response, now),
Self::Redis { cache, .. } => cache.store(request, response, now),
Self::S3(cache) => cache.store(request, response, now),
Self::Gcs(cache) => cache.store(request, response, now),
Self::Disk(cache) => cache.store(request, response, now),
Self::AzureBlob(cache) => cache.store(request, response, now),
}
}
@ -160,6 +253,9 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.lookup_batch(requests, now),
Self::Redis { cache, .. } => cache.lookup_batch(requests, now),
Self::S3(cache) => cache.lookup_batch(requests, now),
Self::Gcs(cache) => cache.lookup_batch(requests, now),
Self::Disk(cache) => cache.lookup_batch(requests, now),
Self::AzureBlob(cache) => cache.lookup_batch(requests, now),
}
}
@ -172,6 +268,9 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.async_lookup(request, now).await,
Self::Redis { cache, .. } => cache.async_lookup(request, now).await,
Self::S3(cache) => cache.async_lookup(request, now).await,
Self::Gcs(cache) => cache.async_lookup(request, now).await,
Self::Disk(cache) => cache.async_lookup(request, now).await,
Self::AzureBlob(cache) => cache.async_lookup(request, now).await,
}
}
@ -192,6 +291,9 @@ impl NativeResponseCache {
cache,
buffer: Some(buffer),
} => buffer.async_store(cache, request, response, now).await,
Self::S3(cache) => cache.async_store(request, response, now).await,
Self::Gcs(cache) => cache.async_store(request, response, now).await,
Self::Disk(cache) => cache.async_store(request, response, now).await,
Self::AzureBlob(cache) => cache.async_store(request, response, now).await,
}
}
@ -204,6 +306,9 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.async_lookup_batch(requests, now).await,
Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await,
Self::S3(cache) => cache.async_lookup_batch(requests, now).await,
Self::Gcs(cache) => cache.async_lookup_batch(requests, now).await,
Self::Disk(cache) => cache.async_lookup_batch(requests, now).await,
Self::AzureBlob(cache) => cache.async_lookup_batch(requests, now).await,
}
}
@ -216,6 +321,9 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.async_store_batch(entries, now).await,
Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await,
Self::S3(cache) => cache.async_store_batch(entries, now).await,
Self::Gcs(cache) => cache.async_store_batch(entries, now).await,
Self::Disk(cache) => cache.async_store_batch(entries, now).await,
Self::AzureBlob(cache) => cache.async_store_batch(entries, now).await,
}
}
@ -229,6 +337,9 @@ impl NativeResponseCache {
}
cache.async_flush().await
}
Self::S3(cache) => cache.async_flush().await,
Self::Gcs(cache) => cache.async_flush().await,
Self::Disk(cache) => cache.async_flush().await,
Self::AzureBlob(cache) => cache.async_flush().await,
}
}
@ -237,7 +348,17 @@ impl NativeResponseCache {
match self {
Self::Memory(cache) => cache.test_connection().await,
Self::Redis { cache, .. } => cache.test_connection().await,
Self::S3(cache) => cache.test_connection().await,
Self::Gcs(cache) => cache.test_connection().await,
Self::Disk(cache) => cache.test_connection().await,
Self::AzureBlob(cache) => cache.test_connection().await,
}
}
pub fn gcs_backend(&self) -> Option<&GcsCache<ResponseCacheCodec>> {
match self {
Self::Gcs(cache) => Some(cache.backend()),
_ => None,
}
}
}

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

@ -0,0 +1,25 @@
[package]
name = "litellm-secrets-hashicorp"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-core-utils.workspace = true
litellm-secrets-types.workspace = true
moka.workspace = true
rustify.workspace = true
rustify_derive.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio.workspace = true
vaultrs.workspace = true
veil.workspace = true
[dev-dependencies]
rstest.workspace = true
tempfile = "3"
tokio.workspace = true
wiremock = "0.6.5"

22
litellm-rust/crates/secrets-hashicorp/src/cert_login.rs Normal file
View file

@ -0,0 +1,22 @@
#[derive(Debug, rustify_derive::Endpoint)]
#[endpoint(path = "/auth/{self.mount}/login", method = "POST")]
pub struct CertLoginRequest {
#[endpoint(skip)]
pub mount: String,
#[endpoint(raw)]
body: Vec<u8>,
}
impl CertLoginRequest {
pub fn new(name: Option<&str>) -> Self {
let body: Vec<u8> = match name {
Some(name) => serde_json::to_vec(&serde_json::json!({ "name": name }))
.expect("json object serialization is infallible"),
None => b"{}".to_vec(),
};
Self {
mount: "cert".to_owned(),
body,
}
}
}

161
litellm-rust/crates/secrets-hashicorp/src/config.rs Normal file
View file

@ -0,0 +1,161 @@
use std::{path::PathBuf, time::Duration};
use litellm_core_utils::settings::Lookup;
use litellm_secrets_types::SecretValue;
use crate::Error;
const DEFAULT_ADDRESS: &str = "http://127.0.0.1:8200";
const DEFAULT_MOUNT: &str = "secret";
const DEFAULT_APPROLE_MOUNT_PATH: &str = "approle";
const DEFAULT_REFRESH_INTERVAL: Duration = Duration::from_secs(86400);
const HCP_VAULT_ADDR: &str = "HCP_VAULT_ADDR";
const HCP_VAULT_TOKEN: &str = "HCP_VAULT_TOKEN";
const HCP_VAULT_NAMESPACE: &str = "HCP_VAULT_NAMESPACE";
const HCP_VAULT_LOGIN_NAMESPACE: &str = "HCP_VAULT_LOGIN_NAMESPACE";
const HCP_VAULT_SECRET_NAMESPACE: &str = "HCP_VAULT_SECRET_NAMESPACE";
const HCP_VAULT_MOUNT_NAME: &str = "HCP_VAULT_MOUNT_NAME";
const HCP_VAULT_PATH_PREFIX: &str = "HCP_VAULT_PATH_PREFIX";
const HCP_VAULT_APPROLE_ROLE_ID: &str = "HCP_VAULT_APPROLE_ROLE_ID";
const HCP_VAULT_APPROLE_SECRET_ID: &str = "HCP_VAULT_APPROLE_SECRET_ID";
const HCP_VAULT_APPROLE_MOUNT_PATH: &str = "HCP_VAULT_APPROLE_MOUNT_PATH";
const HCP_VAULT_CLIENT_CERT: &str = "HCP_VAULT_CLIENT_CERT";
const HCP_VAULT_CLIENT_KEY: &str = "HCP_VAULT_CLIENT_KEY";
const HCP_VAULT_CERT_ROLE: &str = "HCP_VAULT_CERT_ROLE";
const HCP_VAULT_REFRESH_INTERVAL: &str = "HCP_VAULT_REFRESH_INTERVAL";
const SECRET_MANAGER_REFRESH_INTERVAL: &str = "SECRET_MANAGER_REFRESH_INTERVAL";
#[derive(Clone, Debug)]
pub struct AppRoleAuth {
pub role_id: String,
pub secret_id: SecretValue,
pub mount_path: String,
}
#[derive(Clone, Debug)]
pub struct TlsCertAuth {
pub cert_path: PathBuf,
pub key_path: PathBuf,
pub role: Option<String>,
}
#[derive(Clone, Debug)]
pub struct HashicorpVaultConfig {
pub address: String,
pub token: Option<SecretValue>,
pub namespace: Option<String>,
pub login_namespace: Option<String>,
pub secret_namespace: Option<String>,
pub mount: String,
pub path_prefix: Option<String>,
pub approle: Option<AppRoleAuth>,
pub tls_cert: Option<TlsCertAuth>,
pub refresh_interval: Duration,
}
impl HashicorpVaultConfig {
pub fn from_environment(environment: &dyn Lookup) -> Result<Self, Error> {
let address: String = environment
.get(HCP_VAULT_ADDR)
.and_then(|value| nonempty(value.trim()))
.map(|value| value.trim_end_matches('/').to_owned())
.filter(|value| !value.is_empty())
.unwrap_or_else(|| DEFAULT_ADDRESS.to_owned());
let token: Option<SecretValue> = environment
.get(HCP_VAULT_TOKEN)
.and_then(nonempty)
.map(SecretValue::new);
let namespace: Option<String> = path_component(environment.get(HCP_VAULT_NAMESPACE));
let login_namespace: Option<String> =
path_component(environment.get(HCP_VAULT_LOGIN_NAMESPACE));
let secret_namespace: Option<String> =
path_component(environment.get(HCP_VAULT_SECRET_NAMESPACE));
let mount: String = path_component(environment.get(HCP_VAULT_MOUNT_NAME))
.unwrap_or_else(|| DEFAULT_MOUNT.to_owned());
let path_prefix: Option<String> = path_component(environment.get(HCP_VAULT_PATH_PREFIX));
let approle: Option<AppRoleAuth> = match (
environment
.get(HCP_VAULT_APPROLE_ROLE_ID)
.and_then(nonempty),
environment
.get(HCP_VAULT_APPROLE_SECRET_ID)
.and_then(nonempty)
.map(SecretValue::new),
) {
(Some(role_id), Some(secret_id)) => Some(AppRoleAuth {
role_id,
secret_id,
mount_path: path_component(environment.get(HCP_VAULT_APPROLE_MOUNT_PATH))
.unwrap_or_else(|| DEFAULT_APPROLE_MOUNT_PATH.to_owned()),
}),
_ => None,
};
let tls_cert: Option<TlsCertAuth> = match (
environment.get(HCP_VAULT_CLIENT_CERT).and_then(nonempty),
environment.get(HCP_VAULT_CLIENT_KEY).and_then(nonempty),
) {
(Some(cert_path), Some(key_path)) => Some(TlsCertAuth {
cert_path: PathBuf::from(cert_path),
key_path: PathBuf::from(key_path),
role: environment.get(HCP_VAULT_CERT_ROLE).and_then(nonempty),
}),
_ => None,
};
let refresh_interval: Duration = refresh_interval(environment)?;
Ok(Self {
address,
token,
namespace,
login_namespace,
secret_namespace,
mount,
path_prefix,
approle,
tls_cert,
refresh_interval,
})
}
pub fn login_namespace(&self) -> Option<&str> {
self.login_namespace
.as_deref()
.or(self.namespace.as_deref())
}
pub fn secret_namespace(&self) -> Option<&str> {
self.secret_namespace
.as_deref()
.or(self.namespace.as_deref())
}
}
fn nonempty(value: impl AsRef<str>) -> Option<String> {
let value: &str = value.as_ref();
(!value.is_empty()).then(|| value.to_owned())
}
fn path_component(value: Option<String>) -> Option<String> {
value
.and_then(|value| nonempty(value.trim()))
.map(|value| value.trim_matches('/').to_owned())
.filter(|value| !value.is_empty())
}
fn refresh_interval(environment: &dyn Lookup) -> Result<Duration, Error> {
let value: Option<String> = environment
.get(HCP_VAULT_REFRESH_INTERVAL)
.and_then(nonempty)
.or_else(|| {
environment
.get(SECRET_MANAGER_REFRESH_INTERVAL)
.and_then(nonempty)
});
let Some(value) = value else {
return Ok(DEFAULT_REFRESH_INTERVAL);
};
let seconds: i64 = value.parse().map_err(|_| Error::RefreshInterval)?;
if seconds < 0 {
return Err(Error::RefreshInterval);
}
Ok(Duration::from_secs(seconds as u64))
}

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

@ -0,0 +1,34 @@
#[derive(thiserror::Error, veil::Redact)]
pub enum Error {
#[error("HashiCorp Vault requires an enterprise license")]
EnterpriseRequired,
#[error("invalid secret name")]
InvalidSecretName(#[from] litellm_secrets_types::Error),
#[error("HashiCorp Vault client failed")]
Client(
#[from]
#[redact]
vaultrs::error::ClientError,
),
#[error("HashiCorp Vault client settings are invalid: {message}")]
ClientSettings { message: String },
#[error("HashiCorp Vault TLS identity could not be configured for {path}: {message}")]
TlsIdentity {
path: std::path::PathBuf,
message: String,
},
#[error("HashiCorp Vault login returned HTTP {status}")]
LoginStatus { status: u16 },
#[error("HashiCorp Vault login response is malformed")]
MalformedLogin,
#[error("HashiCorp Vault authentication is not configured")]
NoAuthConfigured,
#[error("HashiCorp Vault returned HTTP {status}")]
Status { status: u16 },
#[error("HashiCorp Vault response payload is malformed")]
MalformedPayload,
#[error("HashiCorp Vault secret value is not a string")]
NonStringValue,
#[error("invalid HashiCorp Vault refresh interval")]
RefreshInterval,
}

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

@ -0,0 +1,10 @@
#![forbid(unsafe_code)]
mod cert_login;
mod config;
mod error;
pub mod secret_manager;
pub use config::{AppRoleAuth, HashicorpVaultConfig, TlsCertAuth};
pub use error::Error;
pub use secret_manager::{HashicorpVault, SecretLocation};

359
litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs Normal file
View file

@ -0,0 +1,359 @@
use std::{
collections::HashMap,
fmt,
sync::Arc,
time::{Duration, Instant},
};
use litellm_core_utils::settings::Lookup;
use litellm_secrets_types::{
BaseSecretManager, SecretValue, async_rotate_secret, validate_secret_name,
};
use moka::future::Cache;
use rustify::errors::ClientError as RustifyClientError;
use serde_json::Value;
use tokio::sync::Mutex;
use vaultrs::{
api,
auth::approle,
client::{Identity, VaultClient, VaultClientSettingsBuilder},
error::ClientError,
kv2,
};
use crate::{Error, HashicorpVaultConfig, TlsCertAuth, cert_login::CertLoginRequest};
const CACHE_CAPACITY: u64 = 200;
#[derive(Clone)]
struct CachedClient {
client: Arc<VaultClient>,
expires_at: Option<Instant>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SecretLocation {
pub namespace: Option<String>,
pub mount: String,
pub path: String,
}
#[derive(Clone)]
pub struct HashicorpVault {
config: HashicorpVaultConfig,
cache: Cache<String, SecretValue>,
auth_client: Arc<Mutex<Option<CachedClient>>>,
}
impl fmt::Debug for HashicorpVault {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("HashicorpVault")
.field("config", &self.config)
.finish_non_exhaustive()
}
}
impl HashicorpVault {
pub fn new(
environment: Arc<dyn Lookup + Send + Sync>,
enterprise_enabled: bool,
) -> Result<Self, Error> {
let config: HashicorpVaultConfig =
HashicorpVaultConfig::from_environment(environment.as_ref())?;
Self::from_config(config, enterprise_enabled)
}
pub fn from_config(
config: HashicorpVaultConfig,
enterprise_enabled: bool,
) -> Result<Self, Error> {
if !enterprise_enabled {
return Err(Error::EnterpriseRequired);
}
let cache: Cache<String, SecretValue> = Cache::builder()
.max_capacity(CACHE_CAPACITY)
.time_to_live(config.refresh_interval)
.build();
Ok(Self {
config,
cache,
auth_client: Arc::new(Mutex::new(None)),
})
}
pub fn secret_location(&self, secret_name: &str) -> Result<SecretLocation, Error> {
validate_secret_name(secret_name).map_err(Error::InvalidSecretName)?;
let path: String = [
self.config.path_prefix.clone(),
Some(secret_name.to_owned()),
]
.into_iter()
.flatten()
.collect::<Vec<String>>()
.join("/");
Ok(SecretLocation {
namespace: self.config.secret_namespace().map(str::to_owned),
mount: self.config.mount.clone(),
path,
})
}
pub fn config(&self) -> &HashicorpVaultConfig {
&self.config
}
pub async fn async_read_secret(&self, secret_name: &str) -> Result<Option<SecretValue>, Error> {
let location: SecretLocation = self.secret_location(secret_name)?;
let cache_key: String = cache_key(&location);
if let Some(value) = self.cache.get(&cache_key).await {
return Ok(Some(value));
}
let client: Arc<VaultClient> = self.vault_client().await?;
let data: HashMap<String, Value> =
match kv2::read(client.as_ref(), &location.mount, &location.path).await {
Ok(data) => data,
Err(error) if api_status(&error) == Some(404) => return Ok(None),
Err(error) => return Err(map_api_error(error, ErrorContext::Read)),
};
let Some(value) = data.get("key") else {
return Ok(None);
};
let value: &str = value.as_str().ok_or(Error::NonStringValue)?;
let value: SecretValue = SecretValue::new(value);
self.cache.insert(cache_key, value.clone()).await;
Ok(Some(value))
}
pub async fn async_write_secret(
&self,
secret_name: &str,
value: SecretValue,
description: Option<&str>,
) -> Result<Value, Error> {
let location: SecretLocation = self.secret_location(secret_name)?;
let cache_key: String = cache_key(&location);
let data: HashMap<String, Value> = match description {
Some(description) => [
("key".to_owned(), Value::String(value.expose().to_owned())),
(
"description".to_owned(),
Value::String(description.to_owned()),
),
]
.into_iter()
.collect(),
None => [("key".to_owned(), Value::String(value.expose().to_owned()))]
.into_iter()
.collect(),
};
let client: Arc<VaultClient> = self.vault_client().await?;
let metadata = kv2::set(client.as_ref(), &location.mount, &location.path, &data)
.await
.map_err(|error| map_api_error(error, ErrorContext::Secret))?;
self.cache.invalidate(&cache_key).await;
serde_json::to_value(metadata)
.map_err(|source| Error::Client(ClientError::JsonParseError { source }))
}
pub async fn async_delete_secret(&self, secret_name: &str) -> Result<(), Error> {
let location: SecretLocation = self.secret_location(secret_name)?;
let cache_key: String = cache_key(&location);
let client: Arc<VaultClient> = self.vault_client().await?;
kv2::delete_latest(client.as_ref(), &location.mount, &location.path)
.await
.map_err(|error| map_api_error(error, ErrorContext::Secret))?;
self.cache.invalidate(&cache_key).await;
Ok(())
}
pub async fn async_rotate_secret(
&self,
current_name: &str,
new_name: &str,
value: &SecretValue,
) -> Result<Value, Error> {
async_rotate_secret(self, current_name, new_name, value).await
}
async fn vault_client(&self) -> Result<Arc<VaultClient>, Error> {
let mut cached = self.auth_client.lock().await;
if let Some(entry) = cached.as_ref()
&& entry
.expires_at
.is_none_or(|expires_at| expires_at > Instant::now())
{
return Ok(entry.client.clone());
}
let (client, expires_at): (VaultClient, Option<Instant>) =
match (self.config.approle.as_ref(), self.config.tls_cert.as_ref()) {
(Some(approle), _) => {
let login_client: VaultClient =
self.build_client(self.config.login_namespace(), "")?;
let auth = approle::login(
&login_client,
&approle.mount_path,
&approle.role_id,
approle.secret_id.expose(),
)
.await
.map_err(|error| map_api_error(error, ErrorContext::Login))?;
(
self.build_client(self.config.secret_namespace(), &auth.client_token)?,
token_expiry(auth.lease_duration),
)
}
(None, Some(tls)) => {
let login_client: VaultClient =
self.build_client(self.config.login_namespace(), "")?;
let endpoint: CertLoginRequest = CertLoginRequest::new(tls.role.as_deref());
let auth = api::auth(&login_client, endpoint)
.await
.map_err(|error| map_api_error(error, ErrorContext::Login))?;
(
self.build_client(self.config.secret_namespace(), &auth.client_token)?,
token_expiry(auth.lease_duration),
)
}
(None, None) => {
let token: SecretValue =
self.config.token.clone().ok_or(Error::NoAuthConfigured)?;
(
self.build_client(self.config.secret_namespace(), token.expose())?,
None,
)
}
};
let client: Arc<VaultClient> = Arc::new(client);
*cached = Some(CachedClient {
client: client.clone(),
expires_at,
});
Ok(client)
}
fn build_client(&self, namespace: Option<&str>, token: &str) -> Result<VaultClient, Error> {
let settings = VaultClientSettingsBuilder::default()
.address(&self.config.address)
.token(token.to_owned())
.namespace(namespace.map(str::to_owned))
.identity(identity_for(self.config.tls_cert.as_ref())?)
.ca_certs(Vec::new())
.verify(true)
.build()
.map_err(|message| Error::ClientSettings {
message: message.to_string(),
})?;
VaultClient::new(settings).map_err(Error::Client)
}
}
impl BaseSecretManager for HashicorpVault {
type Error = Error;
type WriteResponse = Value;
type DeleteResponse = ();
async fn async_read_secret(&self, name: &str) -> Result<Option<SecretValue>, Error> {
HashicorpVault::async_read_secret(self, name).await
}
async fn async_write_secret(
&self,
name: &str,
value: &SecretValue,
description: Option<&str>,
) -> Result<Value, Error> {
HashicorpVault::async_write_secret(self, name, value.clone(), description).await
}
async fn async_delete_secret(
&self,
name: &str,
_recovery_window_in_days: i64,
) -> Result<(), Error> {
HashicorpVault::async_delete_secret(self, name).await
}
}
#[derive(Clone, Copy)]
enum ErrorContext {
Login,
Read,
Secret,
}
fn cache_key(location: &SecretLocation) -> String {
format!(
"{:?}/{}/{}",
location.namespace, location.mount, location.path
)
}
fn identity_for(tls: Option<&TlsCertAuth>) -> Result<Option<Identity>, Error> {
tls.map(|tls| {
let cert: Vec<u8> = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity {
path: tls.cert_path.clone(),
message: source.to_string(),
})?;
let key: Vec<u8> = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity {
path: tls.key_path.clone(),
message: source.to_string(),
})?;
Identity::from_pem(&[cert.as_slice(), key.as_slice()].concat()).map_err(|source| {
Error::TlsIdentity {
path: tls.cert_path.clone(),
message: source.to_string(),
}
})
})
.transpose()
}
fn map_api_error(error: ClientError, context: ErrorContext) -> Error {
match error {
ClientError::APIError { code, .. } => match context {
ErrorContext::Login => Error::LoginStatus { status: code },
ErrorContext::Read | ErrorContext::Secret => Error::Status { status: code },
},
ClientError::JsonParseError { source } => match context {
ErrorContext::Login => Error::MalformedLogin,
ErrorContext::Read => Error::MalformedPayload,
ErrorContext::Secret => Error::Client(ClientError::JsonParseError { source }),
},
ClientError::ResponseEmptyError | ClientError::ResponseDataEmptyError => {
malformed_response(context)
}
ClientError::RestClientError { source } => match source {
RustifyClientError::ServerResponseError { code, .. } => match context {
ErrorContext::Login => Error::LoginStatus { status: code },
ErrorContext::Read | ErrorContext::Secret => Error::Status { status: code },
},
RustifyClientError::ResponseParseError { .. } => malformed_response(context),
source => Error::Client(ClientError::RestClientError { source }),
},
error => Error::Client(error),
}
}
fn api_status(error: &ClientError) -> Option<u16> {
match error {
ClientError::APIError { code, .. } => Some(*code),
ClientError::RestClientError {
source: RustifyClientError::ServerResponseError { code, .. },
} => Some(*code),
_ => None,
}
}
fn malformed_response(context: ErrorContext) -> Error {
match context {
ErrorContext::Login => Error::MalformedLogin,
ErrorContext::Read => Error::MalformedPayload,
ErrorContext::Secret => Error::Client(ClientError::ResponseDataEmptyError),
}
}
fn token_expiry(lease_duration: u64) -> Option<Instant> {
(lease_duration > 0).then(|| Instant::now() + Duration::from_secs(lease_duration))
}

602
litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs Normal file
View file

@ -0,0 +1,602 @@
use std::{collections::HashMap, sync::Arc, time::Duration};
use litellm_core_utils::settings::Lookup;
use litellm_secrets_hashicorp::{Error, HashicorpVault, HashicorpVaultConfig};
use litellm_secrets_types::SecretValue;
use serde::Deserialize;
use serde_json::json;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_json, header, method, path},
};
fn config(server: &MockServer, values: &[(&str, &str)]) -> HashicorpVaultConfig {
let mut environment_values: HashMap<String, String> = values
.iter()
.map(|(name, value)| ((*name).to_owned(), (*value).to_owned()))
.collect();
environment_values.insert("HCP_VAULT_ADDR".to_owned(), server.uri());
let environment: Arc<dyn Lookup + Send + Sync> =
Arc::new(move |name: &str| environment_values.get(name).cloned());
HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap()
}
fn manager(server: &MockServer, values: &[(&str, &str)]) -> HashicorpVault {
HashicorpVault::from_config(config(server, values), true).unwrap()
}
fn auth_response(token: &str, lease_duration: u64) -> serde_json::Value {
json!({
"auth": {
"client_token": token,
"accessor": "",
"policies": [],
"token_policies": [],
"metadata": null,
"lease_duration": lease_duration,
"renewable": false,
"entity_id": "",
"token_type": "service",
"orphan": false
},
"lease_id": "",
"lease_duration": lease_duration,
"renewable": false,
"request_id": "",
"warnings": null,
"wrap_info": null
})
}
fn read_response(data: serde_json::Value) -> serde_json::Value {
json!({
"data": {
"data": data,
"metadata": {
"created_time": "",
"deletion_time": "",
"custom_metadata": null,
"destroyed": false,
"version": 1
}
},
"lease_id": "",
"lease_duration": 0,
"renewable": false,
"request_id": "",
"warnings": null,
"wrap_info": null
})
}
#[tokio::test]
async fn token_reads_use_vault_headers_and_cache_values() {
let server: MockServer = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/v1/secret/data/name"))
.and(header("X-Vault-Token", "token"))
.respond_with(
ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))),
)
.expect(1)
.mount(&server)
.await;
let manager: HashicorpVault = manager(&server, &[("HCP_VAULT_TOKEN", "token")]);
assert_eq!(
manager
.async_read_secret("name")
.await
.unwrap()
.unwrap()
.expose(),
"value"
);
let requests = server.received_requests().await.unwrap();
assert!(
requests
.iter()
.all(|request| !request.headers.contains_key("X-Vault-Namespace"))
);
assert_eq!(
manager
.async_read_secret("name")
.await
.unwrap()
.unwrap()
.expose(),
"value"
);
}
#[tokio::test]
async fn namespace_mount_and_prefix_are_sanitized_in_the_url() {
let server: MockServer = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/v1/kv-prod/data/virtual-keys/name"))
.and(header("X-Vault-Namespace", "team-a"))
.respond_with(
ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))),
)
.expect(1)
.mount(&server)
.await;
let manager: HashicorpVault = manager(
&server,
&[
("HCP_VAULT_TOKEN", "token"),
("HCP_VAULT_SECRET_NAMESPACE", " /team-a/ "),
("HCP_VAULT_MOUNT_NAME", " /kv-prod/ "),
("HCP_VAULT_PATH_PREFIX", " /virtual-keys/ "),
],
);
let location = manager.secret_location("name").unwrap();
assert_eq!(location.namespace.as_deref(), Some("team-a"));
assert_eq!(location.mount, "kv-prod");
assert_eq!(location.path, "virtual-keys/name");
assert!(manager.async_read_secret("name").await.unwrap().is_some());
}
#[test]
fn trailing_address_slashes_are_removed() {
let environment: Arc<dyn Lookup + Send + Sync> = Arc::new(|name: &str| match name {
"HCP_VAULT_ADDR" => Some("http://vault.test:8200///".to_owned()),
"HCP_VAULT_TOKEN" => Some("token".to_owned()),
_ => None,
});
let config: HashicorpVaultConfig =
HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap();
let manager: HashicorpVault = HashicorpVault::from_config(config, true).unwrap();
assert_eq!(
manager.secret_location("name").unwrap(),
litellm_secrets_hashicorp::SecretLocation {
namespace: None,
mount: "secret".to_owned(),
path: "name".to_owned(),
}
);
}
#[rstest::rstest]
#[case("-1")]
#[case("not-a-number")]
fn invalid_refresh_intervals_are_rejected(#[case] value: &str) {
let environment: Arc<dyn Lookup + Send + Sync> = Arc::new(move |name: &str| match name {
"HCP_VAULT_REFRESH_INTERVAL" => Some(value.to_owned()),
_ => None,
});
assert!(matches!(
HashicorpVaultConfig::from_environment(environment.as_ref()),
Err(Error::RefreshInterval)
));
}
#[tokio::test]
async fn approle_login_uses_namespace_and_reuses_the_token() {
let server: MockServer = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/auth/custom-approle/login"))
.and(header("X-Vault-Namespace", "login-root"))
.and(body_json(json!({"role_id": "role", "secret_id": "secret"})))
.respond_with(ResponseTemplate::new(200).set_body_json(auth_response("login-token", 3600)))
.expect(1)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/v1/secret/data/name"))
.and(header("X-Vault-Token", "login-token"))
.and(header("X-Vault-Namespace", "secret-root"))
.respond_with(
ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))),
)
.expect(1)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/v1/secret/data/name-2"))
.respond_with(ResponseTemplate::new(404).set_body_json(json!({"errors": ["missing"]})))
.expect(1)
.mount(&server)
.await;
let manager: HashicorpVault = manager(
&server,
&[
("HCP_VAULT_APPROLE_ROLE_ID", "role"),
("HCP_VAULT_APPROLE_SECRET_ID", "secret"),
("HCP_VAULT_APPROLE_MOUNT_PATH", "custom-approle"),
("HCP_VAULT_NAMESPACE", "secret-root"),
("HCP_VAULT_LOGIN_NAMESPACE", "login-root"),
],
);
assert!(manager.async_read_secret("name").await.unwrap().is_some());
assert!(manager.async_read_secret("name-2").await.unwrap().is_none());
}
#[tokio::test]
async fn approle_tokens_expire_after_the_vault_lease() {
let server: MockServer = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/auth/approle/login"))
.respond_with(ResponseTemplate::new(200).set_body_json(auth_response("login-token", 1)))
.expect(2)
.mount(&server)
.await;
Mock::given(method("GET"))
.respond_with(
ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))),
)
.expect(2)
.mount(&server)
.await;
let manager: HashicorpVault = manager(
&server,
&[
("HCP_VAULT_APPROLE_ROLE_ID", "role"),
("HCP_VAULT_APPROLE_SECRET_ID", "secret"),
("HCP_VAULT_REFRESH_INTERVAL", "0"),
],
);
assert!(manager.async_read_secret("first").await.unwrap().is_some());
tokio::time::sleep(Duration::from_secs(1) + Duration::from_millis(50)).await;
assert!(manager.async_read_secret("second").await.unwrap().is_some());
}
#[tokio::test]
async fn tls_login_posts_the_role_and_uses_the_client_identity() {
let server: MockServer = MockServer::start().await;
let directory: tempfile::TempDir = tempfile::tempdir().unwrap();
let cert_path = directory.path().join("client.crt");
let key_path = directory.path().join("client.key");
std::fs::write(&cert_path, TEST_CERTIFICATE).unwrap();
std::fs::write(&key_path, TEST_PRIVATE_KEY).unwrap();
Mock::given(method("POST"))
.and(path("/v1/auth/cert/login"))
.and(header("X-Vault-Namespace", "login-ns"))
.respond_with(ResponseTemplate::new(200).set_body_json(auth_response("cert-token", 0)))
.expect(2)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/v1/secret/data/name"))
.and(header("X-Vault-Token", "cert-token"))
.and(header("X-Vault-Namespace", "secret-ns"))
.respond_with(
ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))),
)
.expect(2)
.mount(&server)
.await;
let role_values: HashMap<String, String> = HashMap::from([
("HCP_VAULT_ADDR".to_owned(), server.uri()),
(
"HCP_VAULT_CLIENT_CERT".to_owned(),
cert_path.to_str().unwrap().to_owned(),
),
(
"HCP_VAULT_CLIENT_KEY".to_owned(),
key_path.to_str().unwrap().to_owned(),
),
("HCP_VAULT_CERT_ROLE".to_owned(), "vault-role".to_owned()),
(
"HCP_VAULT_LOGIN_NAMESPACE".to_owned(),
"login-ns".to_owned(),
),
(
"HCP_VAULT_SECRET_NAMESPACE".to_owned(),
"secret-ns".to_owned(),
),
]);
let role_environment: Arc<dyn Lookup + Send + Sync> =
Arc::new(move |name: &str| role_values.get(name).cloned());
let role_manager: HashicorpVault = HashicorpVault::new(role_environment, true).unwrap();
assert!(
role_manager
.async_read_secret("name")
.await
.unwrap()
.is_some()
);
let no_role_values: HashMap<String, String> = HashMap::from([
("HCP_VAULT_ADDR".to_owned(), server.uri()),
(
"HCP_VAULT_CLIENT_CERT".to_owned(),
cert_path.to_str().unwrap().to_owned(),
),
(
"HCP_VAULT_CLIENT_KEY".to_owned(),
key_path.to_str().unwrap().to_owned(),
),
(
"HCP_VAULT_LOGIN_NAMESPACE".to_owned(),
"login-ns".to_owned(),
),
(
"HCP_VAULT_SECRET_NAMESPACE".to_owned(),
"secret-ns".to_owned(),
),
]);
let no_role_environment: Arc<dyn Lookup + Send + Sync> =
Arc::new(move |name: &str| no_role_values.get(name).cloned());
let no_role_manager: HashicorpVault = HashicorpVault::new(no_role_environment, true).unwrap();
assert!(
no_role_manager
.async_read_secret("name")
.await
.unwrap()
.is_some()
);
let login_bodies: Vec<serde_json::Value> = server
.received_requests()
.await
.unwrap()
.iter()
.filter(|request| request.method.as_str() == "POST")
.map(|request| serde_json::from_slice(&request.body).unwrap())
.collect();
assert!(login_bodies.contains(&json!({"name": "vault-role"})));
assert!(login_bodies.contains(&json!({})));
}
#[rstest::rstest]
#[case::missing(404, json!({"errors": ["missing"]}), 0)]
#[case::malformed(200, json!({"data": "invalid"}), 1)]
#[case::missing_key(200, json!({}), 0)]
#[case::non_string(200, json!({"key": 1}), 2)]
#[tokio::test]
async fn read_responses_distinguish_absence_and_malformed_payloads(
#[case] status: u16,
#[case] body: serde_json::Value,
#[case] expected: u8,
) {
let server: MockServer = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(ResponseTemplate::new(status).set_body_json(
if status == 200 && expected != 1 {
read_response(body)
} else {
body
},
))
.expect(1)
.mount(&server)
.await;
let result: Result<Option<SecretValue>, Error> =
manager(&server, &[("HCP_VAULT_TOKEN", "token")])
.async_read_secret("name")
.await;
match expected {
0 => assert!(result.unwrap().is_none()),
1 => assert!(matches!(result, Err(Error::MalformedPayload))),
2 => assert!(matches!(result, Err(Error::NonStringValue))),
_ => unreachable!(),
}
}
#[tokio::test]
async fn write_and_delete_invalidate_the_read_cache() {
let server: MockServer = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/v1/secret/data/name"))
.respond_with(
ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))),
)
.expect(2)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/v1/secret/data/name"))
.and(body_json(
json!({"data": {"key": "updated", "description": "description"}}),
))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"data": {
"created_time": "",
"deletion_time": "",
"custom_metadata": null,
"destroyed": false,
"version": 2
},
"lease_id": "",
"lease_duration": 0,
"renewable": false,
"request_id": "",
"warnings": null,
"wrap_info": null
})))
.expect(1)
.mount(&server)
.await;
Mock::given(method("DELETE"))
.and(path("/v1/secret/data/name"))
.respond_with(ResponseTemplate::new(204))
.expect(1)
.mount(&server)
.await;
let manager: HashicorpVault = manager(&server, &[("HCP_VAULT_TOKEN", "token")]);
assert!(manager.async_read_secret("name").await.unwrap().is_some());
assert!(
manager
.async_write_secret("name", SecretValue::new("updated"), Some("description"))
.await
.is_ok()
);
assert!(manager.async_read_secret("name").await.unwrap().is_some());
manager.async_delete_secret("name").await.unwrap();
}
#[tokio::test]
async fn no_auth_and_invalid_names_fail_without_requests() {
let server: MockServer = MockServer::start().await;
let manager: HashicorpVault = manager(&server, &[]);
assert!(matches!(
manager.async_read_secret("name").await,
Err(Error::NoAuthConfigured)
));
assert!(matches!(
manager.async_read_secret("../name").await,
Err(Error::InvalidSecretName(_))
));
assert!(server.received_requests().await.unwrap().is_empty());
}
#[tokio::test]
async fn debug_output_redacts_authentication_values() {
let server: MockServer = MockServer::start().await;
let manager: HashicorpVault =
HashicorpVault::from_config(config(&server, &[("HCP_VAULT_TOKEN", "token-value")]), true)
.unwrap();
let debug: String = format!("{manager:?}");
assert!(!debug.contains("token-value"));
assert!(!debug.contains("secret-id"));
}
#[derive(Deserialize)]
struct ParityCase {
env: HashMap<String, String>,
expected_secret_url: String,
expected_login_url: Option<String>,
expected_login_namespace: Option<String>,
expected_secret_namespace: Option<String>,
secret_name: String,
}
#[test]
fn configuration_matches_python_parity_fixture() {
let cases: Vec<ParityCase> = serde_json::from_str(include_str!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../../../tests/test_litellm/secret_managers/hashicorp_vault_parity.json"
)))
.unwrap();
for case in cases {
let values: HashMap<String, String> = case.env.clone();
let environment: Arc<dyn Lookup + Send + Sync> =
Arc::new(move |name: &str| values.get(name).cloned());
let config: HashicorpVaultConfig =
HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap();
let manager: HashicorpVault = HashicorpVault::from_config(config.clone(), true).unwrap();
let location = manager.secret_location(&case.secret_name).unwrap();
let namespace = location
.namespace
.as_deref()
.map(|namespace| format!("{namespace}/"))
.unwrap_or_default();
assert_eq!(
format!(
"{}/v1/{}{}/data/{}",
config.address, namespace, location.mount, location.path
),
case.expected_secret_url
);
let login_url = config.approle.as_ref().map_or_else(
|| {
config
.tls_cert
.as_ref()
.map(|_| format!("{}/v1/auth/cert/login", config.address))
},
|approle| {
Some(format!(
"{}/v1/auth/{}/login",
config.address, approle.mount_path
))
},
);
assert_eq!(login_url, case.expected_login_url);
assert_eq!(
manager.config().login_namespace(),
case.expected_login_namespace.as_deref()
);
assert_eq!(
manager.config().secret_namespace(),
case.expected_secret_namespace.as_deref()
);
}
}
#[tokio::test]
#[ignore]
async fn live_vault_round_trip() {
let environment: Arc<dyn Lookup + Send + Sync> =
Arc::new(litellm_core_utils::settings::ProcessEnvironment);
let manager: HashicorpVault = HashicorpVault::new(environment, true).unwrap();
let name: String = std::env::var("LITELLM_VAULT_LIVE_SECRET_NAME").unwrap();
let value: SecretValue = SecretValue::new("native-live-value");
let location = manager.secret_location(&name).unwrap();
println!(
"native provenance: {} vaultrs {} {:?} {} {}",
module_path!(),
manager.config().address,
location.namespace,
location.mount,
location.path
);
manager
.async_write_secret(&name, value.clone(), None)
.await
.unwrap();
assert_eq!(
manager.async_read_secret(&name).await.unwrap().unwrap(),
value
);
manager.async_delete_secret(&name).await.unwrap();
assert!(manager.async_read_secret(&name).await.unwrap().is_none());
}
const TEST_CERTIFICATE: &str = "-----BEGIN CERTIFICATE-----
MIIDDzCCAfegAwIBAgIUeMzLFLM/mRbPGbNAew5N2UTscocwDQYJKoZIhvcNAQEL
BQAwFzEVMBMGA1UEAwwMbGl0ZWxsbS10ZXN0MB4XDTI2MDkyMTIwMjA1OVoXDTI2
MDkyMjIwMjA1OVowFzEVMBMGA1UEAwwMbGl0ZWxsbS10ZXN0MIIBIjANBgkqhkiG
9w0BAQEFAAOCAQ8AMIIBCgKCAQEAveYoSUJXybmkHmQsBfhBcv2Ob5Oy8ejZu+B3
vTnrPumW4ANi1XXKBSazRGB3fEtAgr+3KhKeHaSKEQeBwJkAEBfdmQv0tpXICwHs
1kFNtU0owy54HVW5/ia+LMszsFcPzVIoMnbUOuiKr9RaV7P+IEFzILPBVuV4DoYH
yocjD3+9QNqokWgNL8LK37JijmNEFVaKFz0X6SyL2VRDlfPWTEBK52Gp/pvDgA6G
eTSfyI+kCm9h5ECTYUAtmatk9WPVS8sWOqV1EXVanFyYBU+mDxoywAS1/6CHeIPh
bNmCOZjPoO9qWBJ7ZyGhOconBigXY8qnlXymev+44IPHrx4urwIDAQABo1MwUTAd
BgNVHQ4EFgQUvaZrZ6HKtbr3ekeZmgy4b5Pq95QwHwYDVR0jBBgwFoAUvaZrZ6HK
tbr3ekeZmgy4b5Pq95QwDwYDVR0TAQH/BAUwAwEB/zANBgkqhkiG9w0BAQsFAAOC
AQEAEejrD8d1qDxW55XxQ4IC31rufoEvDV955jyvh2kALPaN/i5oWsBGI+UAQZna
aaoQXwzlmHrtDUBWl0LztVTUamIleUep2+PLLauqqt43vxppxMX8Jn2mnPO20YE/
hIzGx0jN/LBG8PDyLSvHdlgjP9ofA4Vg4rTQugdXRgOvlCE/epnH/MADcg9KYJtJ
C1RObCIkL3LcdUbjStJRCY/U/FeWcgyncEPz95OFDkbrlNDajb6o6CkYfouqvhTc
8XlgjjAVKIbAbRgbVu3elsquuFM97x2DzWDjkrMNmDt1FJ9ubK36gL6B3o0UMaoQ
00R7x/eqvH+EkWa/2ekW9lpleQ==
-----END CERTIFICATE-----
";
const TEST_PRIVATE_KEY: &str = "-----BEGIN PRIVATE KEY-----
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQC95ihJQlfJuaQe
ZCwF+EFy/Y5vk7Lx6Nm74He9Oes+6ZbgA2LVdcoFJrNEYHd8S0CCv7cqEp4dpIoR
B4HAmQAQF92ZC/S2lcgLAezWQU21TSjDLngdVbn+Jr4syzOwVw/NUigydtQ66Iqv
1FpXs/4gQXMgs8FW5XgOhgfKhyMPf71A2qiRaA0vwsrfsmKOY0QVVooXPRfpLIvZ
VEOV89ZMQErnYan+m8OADoZ5NJ/Ij6QKb2HkQJNhQC2Zq2T1Y9VLyxY6pXURdVqc
XJgFT6YPGjLABLX/oId4g+Fs2YI5mM+g72pYEntnIaE5yicGKBdjyqeVfKZ6/7jg
g8evHi6vAgMBAAECggEAGdJjlP6b8Fa5bdaCM/ebcrbuuNZVJVbb0JPHxGfNSLs7
pE9hj5QaOdQW2Uviw3h6F61ZCzQH4xD+Iy2po5ZKb2XHYKnDB1bboj+LRGER337T
9aJqe9at2VTMVEv3Rdm40NsEk0QcPLxlK16NQFK90gYEUSSQPDAswJDSG2R/zHn+
vADI907mW/goEJHeLn8PWGlNlSiR6x+5JJtq+GXCzUzVvJYQSCLGxCSl2x2H+0g7
NhFI0zPpdzNmO/h+yhzaFb6Rp5U8+ZsnZ3qYjQ/03gw1myTDKJt1YaO9JvArnNYX
hcJQQ8Rt0bHhcrZA16bBOpqZlo5pKCicwI/netgN8QKBgQDcFz7AzdJ26sMSV32V
rwrMgIoggt8qDjO1ARwqW35A1TIge0FoW4M4KpsXQGGfT341uU1esXEcyZ/1L/5X
3ql2gX4DbOYLZLWYzZGR2hq33oi8HkhN98QrEwL9emSH8NqYX3Xxja3PrmCrSYJe
Zbnd9TIm2XkxyMoyXJu6M/QvnwKBgQDc4dzqTbxoGEGa5MuJoGmMwPnqgdG9UM5J
eExVnh7osxc2sOdsiPeRjjQTxs9v2kJwctC359OJoo9yGaaJeSghU4LEWJo1sqnA
fzSCLammYvtVAtniyNv5Mxk/6Uimi4NNDKaAKB+m4K2uSn3U9AmY7KPYMGaSbS9W
XSnobjxm8QKBgC8bPpAvvWs8ZhIn7bY659nLbUT2HeO3dHO6UBf0yzn/J6JyHxbB
93zvCZDZc8uQTRgcmCW7XtVlhjoJUqvl+Wlm39zF0xr/LCsPXKfWAb/2/lcdOCaP
8Emz4QD10EyUTYUtcWYJB/mafhBLRH8F0Nlj4J8WDu2L51MOJTqeYhZLAoGAWffN
icocAbJPlo22sdoa4+/+W5yBF8GAJMDRJtZ+9H1t6SLpQHYRkMIBSETkXUTjZvX9
Ocs9iIQkNW9pO/mTdO+VBfCo71JUfknR02xR+6m5gYjlws/ZeYlssXGN2/hbhNiw
QOcW7Vv6olFJK6Iy/oz0t6wPO3kpnN3Zogi0paECgYEAwo44M1DdYCtV0snhmYM9
5u0mPfYt5P2SVLXyUbr+vFTfrTL/WKnXIJgbsnj3Gvf+GIZv9tKcXhSNmEHQCYX4
X3w9iTPddCHuvZ1fpufi2TyArJh0OkoNtLXJHTKrHjf2N+61AQzFiv5WieJrdE+H
qr32PTUuVGPyO9LyTY4/RL0=
-----END PRIVATE KEY-----
";

2
litellm-rust/crates/secrets/Cargo.toml
View file

@ -9,6 +9,7 @@ repository.workspace = true
default = []
aws = ["dep:litellm-secrets-aws"]
google = ["dep:litellm-secrets-google"]
hashicorp = ["dep:litellm-secrets-hashicorp"]
azure = ["dep:litellm-secrets-azure"]
cyberark = ["dep:litellm-secrets-cyberark"]
@ -16,6 +17,7 @@ cyberark = ["dep:litellm-secrets-cyberark"]
litellm-secrets-types.workspace = true
litellm-secrets-aws = { workspace = true, optional = true }
litellm-secrets-google = { workspace = true, optional = true }
litellm-secrets-hashicorp = { workspace = true, optional = true }
litellm-secrets-azure = { workspace = true, optional = true }
litellm-secrets-cyberark = { workspace = true, optional = true }
litellm-core-utils.workspace = true

2
litellm-rust/crates/secrets/README.md
View file

@ -9,3 +9,5 @@ Backend failures propagate by default. To allow fallback during a backend failur
`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
The HashiCorp Vault backend is enabled with the `hashicorp` feature and reads KV v2 values from `HCP_VAULT_*` environment variables. It supports static tokens, AppRole authentication, and TLS certificate authentication

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

@ -30,6 +30,9 @@ pub enum Error {
#[cfg(feature = "google")]
#[error(transparent)]
Google(#[from] litellm_secrets_google::Error),
#[cfg(feature = "hashicorp")]
#[error(transparent)]
Hashicorp(#[from] litellm_secrets_hashicorp::Error),
#[cfg(feature = "azure")]
#[error(transparent)]
Azure(#[from] litellm_secrets_azure::Error),

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

@ -13,6 +13,8 @@ pub enum SecretManager {
GoogleKms(crate::google::GoogleKms),
#[cfg(feature = "google")]
GoogleSecretManager(crate::google::GoogleSecretManager),
#[cfg(feature = "hashicorp")]
HashicorpVault(crate::hashicorp::HashicorpVault),
#[cfg(feature = "azure")]
AzureKeyVault(crate::azure::AzureKeyVault),
#[cfg(feature = "cyberark")]
@ -31,6 +33,8 @@ impl SecretManager {
Self::GoogleKms(_) => KeyManagementSystem::GoogleKms,
#[cfg(feature = "google")]
Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager,
#[cfg(feature = "hashicorp")]
Self::HashicorpVault(_) => KeyManagementSystem::HashicorpVault,
#[cfg(feature = "azure")]
Self::AzureKeyVault(_) => KeyManagementSystem::AzureKeyVault,
#[cfg(feature = "cyberark")]
@ -86,6 +90,12 @@ pub async fn get_secret_from_manager(
.get_secret_from_google_secret_manager(secret_name)
.await
.map_err(Error::from),
#[cfg(feature = "hashicorp")]
SecretManager::HashicorpVault(client) => client
.async_read_secret(secret_name)
.await
.map(|value| value.map(Secret::String))
.map_err(Error::from),
#[cfg(feature = "azure")]
SecretManager::AzureKeyVault(client) => client
.get_secret_from_azure_key_vault(secret_name)

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

@ -23,3 +23,5 @@ pub use litellm_secrets_azure as azure;
pub use litellm_secrets_cyberark as cyberark;
#[cfg(feature = "google")]
pub use litellm_secrets_google as google;
#[cfg(feature = "hashicorp")]
pub use litellm_secrets_hashicorp as hashicorp;

142
litellm-rust/crates/secrets/tests/handler.rs
View file

@ -105,6 +105,148 @@ async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whites
Err(Error::MissingCiphertext)
));
}
#[cfg(feature = "hashicorp")]
#[tokio::test]
async fn hashicorp_handler_resolves_found_missing_and_failed_values() {
use std::sync::Arc;
use litellm_core_utils::settings::Lookup;
use litellm_secrets::{
Error, FailurePolicy, KeyManagementSettings, SecretManager, SecretManagerState,
SecretResolver, hashicorp::HashicorpVault, hashicorp::HashicorpVaultConfig,
};
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{method, path},
};
let found_server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/v1/secret/data/KEY"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"data": {
"data": {"key": "remote"},
"metadata": {
"created_time": "",
"deletion_time": "",
"custom_metadata": null,
"destroyed": false,
"version": 1
}
},
"lease_id": "",
"lease_duration": 0,
"renewable": false,
"request_id": "",
"warnings": null,
"wrap_info": null
})))
.mount(&found_server)
.await;
let found_environment: Arc<dyn Lookup + Send + Sync> = Arc::new({
let address = found_server.uri();
move |name: &str| match name {
"HCP_VAULT_ADDR" => Some(address.clone()),
"HCP_VAULT_TOKEN" => Some("token".into()),
_ => None,
}
});
let found_config = HashicorpVaultConfig::from_environment(found_environment.as_ref()).unwrap();
let found_manager = HashicorpVault::from_config(found_config, true).unwrap();
let found_resolver = SecretResolver::new(
Arc::new(SecretManagerState::new(
SecretManager::HashicorpVault(found_manager),
KeyManagementSettings {
hosted_keys: Some(vec!["KEY".into()]),
..Default::default()
},
)),
Arc::new(|_: &str| None),
litellm_secrets::OidcResolver::default(),
);
assert_eq!(
found_resolver
.get_secret_str("KEY", None)
.await
.unwrap()
.unwrap()
.expose(),
"remote"
);
let missing_server = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(
ResponseTemplate::new(404).set_body_json(serde_json::json!({"errors": ["missing"]})),
)
.mount(&missing_server)
.await;
let missing_environment: Arc<dyn Lookup + Send + Sync> = Arc::new({
let address = missing_server.uri();
move |name: &str| match name {
"HCP_VAULT_ADDR" => Some(address.clone()),
"HCP_VAULT_TOKEN" => Some("token".into()),
_ => None,
}
});
let missing_config =
HashicorpVaultConfig::from_environment(missing_environment.as_ref()).unwrap();
let missing_manager = HashicorpVault::from_config(missing_config, true).unwrap();
let missing_state = SecretManagerState::new(
SecretManager::HashicorpVault(missing_manager),
KeyManagementSettings {
hosted_keys: Some(vec!["KEY".into()]),
..Default::default()
},
);
let missing = litellm_secrets::get_secret_from_manager(
missing_state.backend().unwrap(),
"KEY",
missing_state.settings().unwrap(),
&|_: &str| None,
)
.await
.unwrap();
assert!(missing.is_none());
let failed_server = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(
ResponseTemplate::new(500).set_body_json(serde_json::json!({"errors": ["failed"]})),
)
.mount(&failed_server)
.await;
let failed_environment: Arc<dyn Lookup + Send + Sync> = Arc::new({
let address = failed_server.uri();
move |name: &str| match name {
"HCP_VAULT_ADDR" => Some(address.clone()),
"HCP_VAULT_TOKEN" => Some("token".into()),
_ => None,
}
});
let failed_config =
HashicorpVaultConfig::from_environment(failed_environment.as_ref()).unwrap();
let failed_manager = HashicorpVault::from_config(failed_config, true).unwrap();
let failed_state = SecretManagerState::new(
SecretManager::HashicorpVault(failed_manager),
KeyManagementSettings {
hosted_keys: Some(vec!["KEY".into()]),
..Default::default()
},
);
let failed_resolver = SecretResolver::new(
Arc::new(failed_state),
Arc::new(|_: &str| None),
litellm_secrets::OidcResolver::default(),
)
.with_failure_policy(FailurePolicy::Propagate);
assert!(matches!(
failed_resolver.get_secret_str("KEY", None).await,
Err(Error::Hashicorp(
litellm_secrets::hashicorp::Error::Status { status: 500 }
))
));
}
#[cfg(feature = "azure")]
#[tokio::test]

9
litellm/constants.py
View file

@ -1860,6 +1860,12 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
"max_ui_session_budget",
"budget_rollover",
"mcp_tool_search",
"turn_off_message_logging",
"datadog_params",
"datadog_llm_observability_params",
"newrelic_params",
"pointfive_params",
"aws_sqs_callback_params",
]
SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"]
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60))
@ -2121,3 +2127,6 @@ BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit"
# Shared read-only empty mapping, for defaulting optional Mapping parameters without
# constructing a fresh mutable dict at each call site.
EMPTY_MAPPING: Final = MappingProxyType({})
# API endpoint for breached password k-anonymity search
HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range"

1
litellm/llms/xai/chat/transformation.py
View file

@ -118,6 +118,7 @@ class XAIChatConfig(OpenAIGPTConfig):
base_openai_params: Final = [
"logit_bias",
"logprobs",
"max_completion_tokens",
"max_tokens",
"n",
"parallel_tool_calls",

55
litellm/model_prices_and_context_window_backup.json
View file

@ -43037,21 +43037,21 @@
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4-pro": {
"input_cost_per_token": 8.92272e-07,
"input_cost_per_token": 8.87226e-07,
"input_cost_per_token_cache_hit": 4.4e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"max_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 1.784544e-06,
"output_cost_per_token": 1.774452e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 7.4356e-08,
"cache_read_input_token_cost": 7.39355e-08,
"supports_audio_input": false,
"supports_pdf_input": false,
"supports_vision": false,
@ -73794,6 +73794,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-1.6": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 2.5e-07,
"input_cost_per_token_above_128k_tokens": 5e-07,
"litellm_provider": "openrouter",
@ -73815,6 +73816,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-1.6-flash": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 7.5e-08,
"input_cost_per_token_above_128k_tokens": 1e-07,
"litellm_provider": "openrouter",
@ -73855,6 +73857,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-2.0-code": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 5e-07,
"input_cost_per_token_above_128k_tokens": 1e-06,
"litellm_provider": "openrouter",
@ -76975,5 +76978,51 @@
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"xiaomi_mimo/mimo-v2.6-pro": {
"cache_read_input_token_cost": 3.6e-09,
"input_cost_per_token": 4.35e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 8.7e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
},
"xiaomi_mimo/mimo-v2.6-flash": {
"cache_read_input_token_cost": 2.8e-09,
"input_cost_per_token": 1.4e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2.8e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
}
}

2
litellm/models/user.py
View file

@ -24,6 +24,8 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase):
organization_id: str | None = None
object_permission_id: str | None = None
password: str | None = Field(default=None, exclude=True)
password_reset_required: bool | None = None
last_breach_check_at: datetime | None = None
teams: list[str] = []
user_role: str | None = None
max_budget: float | None = None

31
litellm/proxy/_types.py
View file

@ -905,6 +905,7 @@ class LiteLLMRoutes(enum.Enum):
"/claude_code_gateway/v1/traces",
"/user/list", # org admins checked in endpoint; non-admins get 403
"/management/v1/users/bulk_delete", # proxy admins delete anyone, org admins only their orgs' users; others 403
"/user/password/change", # endpoint only ever writes the caller's own row
"/model/{model_id}/update",
"/prompt/list",
"/prompt/info",
@ -1865,6 +1866,17 @@ class NewUserRequest(GenerateRequestBase):
send_invite_email: bool | None = None
sso_user_id: str | None = None
organizations: list[str] | None = None
password: str | None = None
@field_validator("password")
@classmethod
def password_not_supported(cls, value: str | None) -> str | None:
if value is not None:
raise ValueError(
"password cannot be set via /user/new. Users set their own password through an "
"invitation link (POST /invitation/new)."
)
return value
class NewUserResponse(GenerateKeyResponse):
@ -1887,7 +1899,8 @@ class NewUserResponse(GenerateKeyResponse):
class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with BulkUpdateUserRequest
password: str | None = None
# repr=False keeps the plaintext out of management-endpoint alerts, which str() the request model
password: str | None = Field(default=None, repr=False)
spend: float | None = None
metadata: dict | None = None
user_alias: str | None = None
@ -1917,6 +1930,16 @@ class UpdateUserRequest(UpdateUserRequestNoUserIDorEmail):
return values
class ChangePasswordRequest(LiteLLMPydanticObjectBase):
current_password: str = Field(repr=False)
new_password: str = Field(repr=False)
class ChangePasswordResponse(LiteLLMPydanticObjectBase):
user_id: str
message: str
class DeleteUserRequest(LiteLLMPydanticObjectBase):
user_ids: list[str] # required
@ -3940,6 +3963,12 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
)
class HTTPExceptionErrorDetail(TypedDict):
"""The `{"error": <message>}` shape most proxy endpoints raise as `HTTPException.detail`."""
error: ReadOnly[str]
class SpendLogsRouterMetadata(TypedDict):
"""
Router provenance stamped on spend logs for deployments flagged with

89
litellm/proxy/auth/login_utils.py
View file

@ -10,14 +10,16 @@ import secrets
from collections.abc import Mapping
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import Final, Literal, cast
from typing import TYPE_CHECKING, Final, Literal, cast
import jwt
from fastapi import HTTPException
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import LITELLM_PROXY_ADMIN_NAME, LITELLM_UI_SESSION_DURATION
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy._types import (
LiteLLM_UserTable,
LitellmUserRoles,
@ -28,6 +30,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.auth_utils import is_sso_provider_fully_configured
from litellm.proxy.auth.login_throttle import LoginAttempt, LoginThrottle
from litellm.proxy.auth.password_policy import is_breach_check_enabled, is_password_breached
from litellm.proxy.management_endpoints.internal_user_endpoints import user_update
from litellm.proxy.management_endpoints.key_management_endpoints import (
generate_key_helper_fn,
@ -50,6 +53,57 @@ INVALID_UI_CREDENTIALS_MESSAGE: Final = (
)
INVALID_USER_PASSWORD_MESSAGE: Final = "Invalid credentials used to access UI. Check the password set for your user"
if TYPE_CHECKING:
from prisma import types as prisma_types
BREACH_RECHECK_INTERVAL: Final = timedelta(hours=24)
PASSWORD_RESET_ALLOWED_ROUTES: Final = ("/user/password/change",)
PASSWORD_SESSION_METADATA: Final = MappingProxyType({"login_method": "username_password"})
def _breach_recheck_due(last_breach_check_at: datetime | None) -> bool:
if last_breach_check_at is None:
return True
last_checked_utc: Final = (
last_breach_check_at
if last_breach_check_at.tzinfo is not None
else last_breach_check_at.replace(tzinfo=timezone.utc)
)
return datetime.now(timezone.utc) - last_checked_utc >= BREACH_RECHECK_INTERVAL
async def screen_login_password_for_breach(
user_id: str,
password: str,
last_breach_check_at: datetime | None,
general_settings: Mapping[str, object],
prisma_client: PrismaClient,
client: AsyncHTTPHandler | None = None,
) -> bool:
"""Screens a successfully verified login password against HIBP, stamps
``password_reset_required`` when breached, and returns whether a breach was
found so the login it runs in can restrict the session it is about to mint.
Fails open (HIBP or DB trouble never fails the login) and rechecks a given
user at most once per ``BREACH_RECHECK_INTERVAL``."""
if not is_breach_check_enabled(general_settings):
return False
if not _breach_recheck_due(last_breach_check_at):
return False
breached: Final = await is_password_breached(password, general_settings, client)
checked_at: Final = datetime.now(timezone.utc)
breached_update: Final[prisma_types.LiteLLM_UserTableUpdateInput] = {
"last_breach_check_at": checked_at,
"password_reset_required": True,
}
recheck_update: Final[prisma_types.LiteLLM_UserTableUpdateInput] = {"last_breach_check_at": checked_at}
update_data: Final = breached_update if breached else recheck_update
find_user: Final[prisma_types.LiteLLM_UserTableWhereInput] = {"user_id": user_id}
try:
await UserRepository(prisma_client).table.update(where=find_user, data=update_data)
except Exception as e: # noqa: BLE001 # a failed stamp must never surface into the login
verbose_proxy_logger.warning("Login-time breach screening could not update user %s: %s", user_id, e)
return breached
async def _rehash_password_if_needed(user_id: str, password: str, stored: str) -> None:
"""Rehash legacy password (SHA256) to scrypt on successful login."""
@ -137,6 +191,7 @@ class LoginResult:
user_email: str | None
user_role: str
login_method: Literal["sso", "username_password"]
password_reset_required: bool
def __init__(
self,
@ -145,12 +200,14 @@ class LoginResult:
user_email: str | None,
user_role: str,
login_method: Literal["sso", "username_password"] = "username_password",
password_reset_required: bool = False,
):
self.user_id = user_id
self.key = key
self.user_email = user_email
self.user_role = user_role
self.login_method = login_method
self.password_reset_required = password_reset_required
async def authenticate_user(
@ -356,20 +413,28 @@ async def _sign_in(
if verify_password(password, _password):
await _rehash_password_if_needed(_user_row.user_id, password, _password)
breached_now: Final = prisma_client is not None and await screen_login_password_for_breach(
user_id=_user_row.user_id,
password=password,
last_breach_check_at=getattr(_user_row, "last_breach_check_at", None),
general_settings=general_settings,
prisma_client=prisma_client,
)
password_reset_required: Final = breached_now or getattr(_user_row, "password_reset_required", None) is True
if os.getenv("DATABASE_URL") is not None:
response = await generate_key_helper_fn(
llm_router=None,
request_type="key",
**{
"user_role": user_role,
"duration": LITELLM_UI_SESSION_DURATION,
"key_max_budget": litellm.max_ui_session_budget,
"models": [],
"aliases": {},
"config": {},
"spend": 0,
"user_id": user_id,
"team_id": "litellm-dashboard",
user_role=user_role,
duration=LITELLM_UI_SESSION_DURATION,
key_max_budget=litellm.max_ui_session_budget,
spend=0,
user_id=user_id,
team_id="litellm-dashboard",
allowed_routes=list(PASSWORD_RESET_ALLOWED_ROUTES) if password_reset_required else None,
metadata={
**PASSWORD_SESSION_METADATA,
**({"password_reset_required": True} if password_reset_required else {}),
},
)
else:
@ -390,6 +455,7 @@ async def _sign_in(
user_email=user_email,
user_role=cast(str, user_role),
login_method="username_password",
password_reset_required=password_reset_required,
)
else:
await attempt.failed()
@ -460,4 +526,5 @@ def create_ui_token_object(
auth_header_name=general_settings.get("litellm_key_header_name", "Authorization"),
disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation,
server_root_path=get_server_root_path(),
password_reset_required=login_result.password_reset_required,
)

128
litellm/proxy/auth/password_policy.py
View file

@ -4,13 +4,28 @@ Applied at every path that persists a new or changed password for a DB-backed
user (``/user/update``, ``/user/bulk_update``, and the invitation onboarding
claim flow), so the strength bar is configured in one place instead of
per-endpoint.
Also screens new passwords against known data breaches via the
haveibeenpwned.com (HIBP) k-anonymity range API: only the first 5 characters
of the password's SHA-1 hash ever leave the proxy, and the check fails open
(allows the password) when HIBP is unreachable.
"""
from collections.abc import Mapping
import asyncio
import hashlib
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from litellm._logging import verbose_proxy_logger
from litellm._version import version
from litellm.constants import HIBP_RANGE_API_BASE
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.types.llms.custom_http import httpxSpecialProvider
HIBP_TIMEOUT_SECONDS: Final = 5.0
DEFAULT_MIN_LENGTH: Final = 12
MIN_ALLOWED_LENGTH: Final = 8
@ -90,3 +105,114 @@ def validate_password_policy(password: str, general_settings: Mapping[str, objec
param="password",
code=400,
)
def _hibp_client() -> AsyncHTTPHandler:
return get_async_httpx_client(
llm_provider=httpxSpecialProvider.PasswordBreachCheck,
params={"timeout": HIBP_TIMEOUT_SECONDS}, # mutable-ok: callee takes a bare dict (PEP 589)
)
def _is_suffix_in_range_response(response_body: str, hash_suffix: str) -> bool:
for line in response_body.upper().splitlines():
entry_suffix, _, count = line.strip().partition(":")
if entry_suffix == hash_suffix:
return int(count.strip() or "0") > 0
return False
async def _is_password_breached(password: str, client: AsyncHTTPHandler) -> bool:
# usedforsecurity=False: SHA-1 is only a lookup key into the HIBP dataset, so no security property rests on it
sha1_hex: Final = hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper()
headers: Final = { # mutable-ok: callee takes a bare dict (PEP 589)
"Add-Padding": "true",
"User-Agent": f"litellm-proxy/{version}",
}
try:
response: Final = await client.get(
f"{HIBP_RANGE_API_BASE}/{sha1_hex[:5]}",
headers=headers,
)
response.raise_for_status()
breached: Final = _is_suffix_in_range_response(response.text, sha1_hex[5:])
except Exception as e: # noqa: BLE001 # fail-open: any HIBP failure skips the check, never breaks the caller
verbose_proxy_logger.warning("Breached-password check skipped, HIBP lookup failed: %s", e)
return False
return breached
def is_breach_check_enabled(general_settings: Mapping[str, object]) -> bool:
return general_settings.get("password_policy_check_breached_passwords", True) is not False
async def is_password_breached(
password: str,
general_settings: Mapping[str, object],
client: AsyncHTTPHandler | None = None,
) -> bool:
"""False when the check is disabled, the password is absent from the HIBP
corpus, or HIBP is unreachable (fail open)."""
if not is_breach_check_enabled(general_settings):
return False
return await _is_password_breached(password, client if client is not None else _hibp_client())
def breached_password_error() -> ProxyException:
return ProxyException(
message=(
"This password appears in known data breaches and cannot be used. Please choose a different password."
),
type=ProxyErrorTypes.validation_error,
param="password",
code=400,
)
async def validate_password_not_breached(
password: str,
general_settings: Mapping[str, object],
client: AsyncHTTPHandler | None = None,
) -> None:
"""Raise ``ProxyException`` (400) if ``password`` appears in a known data breach.
Fails open: an unreachable or misbehaving HIBP allows the password."""
if not await is_password_breached(password, general_settings, client):
return
raise breached_password_error()
def _strength_verdict(password: str, general_settings: Mapping[str, object]) -> ProxyException | None:
try:
validate_password_policy(password, general_settings)
except ProxyException as e:
return e
return None
async def validate_passwords_bulk(
passwords: Sequence[str],
general_settings: Mapping[str, object],
client: AsyncHTTPHandler | None = None,
) -> Mapping[str, ProxyException | None]:
"""Per-unique-password policy verdicts for a batch: the ProxyException to
surface, or None when the password is acceptable.
Deduplicates first, then issues every needed HIBP lookup concurrently, so a
batch caller pays one HIBP timeout window in the worst case instead of one
per password (each lookup still fails open independently)."""
unique_passwords: Final = tuple(dict.fromkeys(passwords))
strength_verdicts: Final[Mapping[str, ProxyException | None]] = MappingProxyType(
{password: _strength_verdict(password, general_settings) for password in unique_passwords}
)
to_screen: Final = tuple(password for password in unique_passwords if strength_verdicts[password] is None)
breached_flags: Final = await asyncio.gather(
*(is_password_breached(password, general_settings, client) for password in to_screen)
)
breached_passwords: Final = frozenset(password for password, breached in zip(to_screen, breached_flags) if breached)
return MappingProxyType(
{
password: breached_password_error() if password in breached_passwords else strength_verdicts[password]
for password in unique_passwords
}
)

27
litellm/proxy/auth/route_checks.py
View file

@ -194,6 +194,16 @@ class RouteChecks:
if denied_auth_enforced_pass_through_route:
raise RouteChecks._auth_pass_through_denied_exception(route=route)
if valid_token.metadata.get("password_reset_required") is True:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=(
"This account's password must be changed before the session can be used: "
"it was either found in a known data breach or set by an admin. "
"Change it via POST /user/password/change (UI: /ui/change-password), then log in again."
),
)
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=f"Virtual key is not allowed to call this route. Only allowed to call routes: {valid_token.allowed_routes}. Tried to call route: {route}",
@ -812,7 +822,8 @@ class RouteChecks:
in the codebase is automatically readable by Admin Viewer
without needing to remember to add it to an allowlist.
3. Unsafe HTTP method (POST/PUT/PATCH/DELETE):
- Allow `/user/update` only when restricted to user_email/password.
- Allow `/user/update` only when restricted to user_email.
- Allow `/user/password/change` (endpoint only writes the caller's own row).
- Block all explicit writes in `_ADMIN_VIEWER_BLOCKED_WRITE_ROUTES`.
- Otherwise allow only if the route is in admin_viewer_routes /
global_spend_tracking_routes (legacy explicit-allow set).
@ -832,10 +843,10 @@ class RouteChecks:
if request_data is not None and isinstance(request_data, dict):
_params_updated: Final = request_data.keys()
for param in _params_updated:
if param not in ["user_email", "password"]:
if param != "user_email":
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email and password can be updated",
detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email can be updated",
)
elif RouteChecks.check_route_access(route=route, allowed_routes=_PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES) or (
route.startswith("/key/") and route.endswith(_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES)
@ -854,21 +865,25 @@ class RouteChecks:
return
# ── Unsafe HTTP method: explicit checks ──────────────────────────
# Allow `/user/update` for self-service email / password change.
# Allow `/user/update` for self-service email change.
if route == "/user/update":
if request_data is not None and isinstance(request_data, dict):
for param in request_data:
if param not in ["user_email", "password"]:
if param != "user_email":
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=(
f"user not allowed to access this route, role= {_user_role}. "
f"Trying to access: {route} and updating invalid param: {param}. "
"only user_email and password can be updated"
"only user_email can be updated"
),
)
return
# Self-service password change; the endpoint only writes the caller's own row.
if route == "/user/password/change":
return
# Hard-block known write routes regardless of HTTP method (defensive
# — these are POSTs in practice, but pinning them here protects
# against future GET-shaped writes).

62
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -28,6 +28,7 @@ from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import (
delete_cache_key_objects,
@ -35,7 +36,11 @@ from litellm.proxy.auth.auth_checks import (
get_team_object,
get_user_object,
)
from litellm.proxy.auth.password_policy import validate_password_policy
from litellm.proxy.auth.password_policy import (
validate_password_not_breached,
validate_password_policy,
validate_passwords_bulk,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
from litellm.proxy.common_utils.user_api_key_cache import (
@ -173,11 +178,23 @@ def _team_membership_table(
return team_membership_table
def _hash_password_in_dict(data: dict, general_settings: Mapping[str, object]) -> None:
"""Validate and hash password field in-place if present."""
async def _hash_password_in_dict(
data: dict, general_settings: Mapping[str, object], password_prevalidated: bool = False
) -> None:
"""Validate and hash password field in-place if present.
``password_prevalidated`` skips the policy checks for callers that already
validated the password (the bulk path screens its whole batch upfront).
An admin-set password is known to whoever set it, so the user is also
flagged for a forced password change at next login."""
if "password" in data and data["password"] is not None:
validate_password_policy(data["password"], general_settings)
if not password_prevalidated:
validate_password_policy(data["password"], general_settings)
await validate_password_not_breached(data["password"], general_settings)
data["password"] = hash_password(data["password"])
data["password_reset_required"] = True
data["last_breach_check_at"] = None
def _strip_password_from_response(response) -> None:
@ -505,6 +522,7 @@ async def new_user(
- prompts: Optional[List[str]] - List of allowed prompts for the user. If specified, the user will only be able to use these specific prompts.
- organizations: List[str] - List of organization id's the user is a member of
- budget_limits: Optional[list] - List of concurrent budget windows for the user. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}].
- password: Optional[str] - Not supported; any value is rejected with a 422. Users set their own password through an invitation link (POST /invitation/new).
Returns:
- key: (str) The generated api key for the user
- expires: (datetime) Datetime object for when key expires.
@ -524,7 +542,7 @@ async def new_user(
```
"""
try:
from litellm.proxy.proxy_server import _license_check, general_settings, prisma_client
from litellm.proxy.proxy_server import _license_check, prisma_client
if prisma_client is None:
raise HTTPException(status_code=400, detail=CommonProxyErrors.db_not_connected_error.value)
@ -572,7 +590,7 @@ async def new_user(
# generate_key_helper_fn only forwards object_permission_id, so without this the entitlement
# the caller sent would be dropped on the floor.
data_json = await _set_object_permission(data_json=data_json, prisma_client=prisma_client)
_hash_password_in_dict(data_json, general_settings)
data_json.pop("password", None)
teams = data.teams
if teams is None:
teams = check_if_default_team_set()
@ -1438,6 +1456,7 @@ async def _update_single_user_helper(
user_request: UpdateUserRequest,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None = None,
password_prevalidated: bool = False,
) -> dict[str, Any]:
"""
Helper function to update a single user.
@ -1460,7 +1479,7 @@ async def _update_single_user_helper(
data_json: Final[dict] = user_request.model_dump(exclude_unset=True)
non_default_values = _update_internal_user_params(data_json=data_json, data=user_request)
_hash_password_in_dict(non_default_values, general_settings)
await _hash_password_in_dict(non_default_values, general_settings, password_prevalidated=password_prevalidated)
existing_user_row: BaseModel | None = None
if user_request.user_id:
@ -1641,7 +1660,7 @@ async def user_update(
Parameters:
- user_id: Optional[str] - Specify a user id. If not set, a unique id will be generated.
- user_email: Optional[str] - Specify a user email.
- password: Optional[str] - Specify a user password.
- password: Optional[str] - Set the user's password (admin only). Must satisfy the configured password policy. The user is required to change it at their next login. Users change their own password with POST /user/password/change.
- user_alias: Optional[str] - A descriptive name for you to know who this user id refers to.
- teams: Optional[list] - specify a list of team id's a user belongs to.
- send_invite_email: Optional[bool] - Specify if an invite email should be sent.
@ -1709,19 +1728,38 @@ async def bulk_update_processed_users(
users_to_update: list[UpdateUserRequest],
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: str | None = None,
hibp_client: AsyncHTTPHandler | None = None,
) -> BulkUpdateUserResponse:
from litellm.proxy.proxy_server import general_settings
results: Final[list[UserUpdateResult]] = []
successful_updates = 0
failed_updates = 0
# Screen the batch's passwords upfront and concurrently: done per-user
# inside the loop below, each HIBP lookup would be awaited serially and a
# degraded-slow HIBP could stretch a full batch to minutes, timing out the
# request after some updates already persisted.
password_verdicts: Final = await validate_passwords_bulk(
tuple(u.password for u in users_to_update if u.password is not None),
general_settings,
client=hibp_client,
)
# Process each user update independently
try:
for user_request in users_to_update:
try:
if (
user_request.password is not None
and (password_error := password_verdicts.get(user_request.password)) is not None
):
raise password_error
response = await _update_single_user_helper(
user_request=user_request,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
password_prevalidated=True,
)
# Record success
results.append(
@ -1859,6 +1897,14 @@ async def bulk_user_update(
status_code=403,
detail="Only proxy admins can update all users at once.",
)
if data.user_updates.password is not None:
bulk_password_error: Final[HTTPExceptionErrorDetail] = {
"error": (
"Setting one password for all users is not supported. "
"Use per-user updates via the 'users' list instead."
)
}
raise HTTPException(status_code=400, detail=bulk_password_error)
# Optimized path for updating all users directly in database
all_users_in_db: Final = await _user_table(prisma_client).find_many(order={"created_at": "desc"})

9
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -1078,6 +1078,9 @@ if MCP_AVAILABLE:
return {"servers": registry_servers}
## FastAPI Routes
def _mcp_server_display_order(server: LiteLLM_MCPServerTable) -> tuple[str, str]:
return ((server.server_name or server.alias or server.server_id).lower(), server.server_id)
def _get_user_mcp_management_mode() -> UserMCPManagementMode:
from litellm.proxy.proxy_server import (
general_settings as proxy_general_settings,
@ -1228,10 +1231,12 @@ if MCP_AVAILABLE:
detail="You do not have permission to view MCP servers for this team.",
)
redacted_mcp_servers = await _get_team_scoped_mcp_server_list(sanitized_team_id)
redacted_mcp_servers = sorted(
await _get_team_scoped_mcp_server_list(sanitized_team_id), key=_mcp_server_display_order
)
else:
servers: Final = await _resolve_accessible_mcp_servers(user_api_key_dict)
redacted_mcp_servers = _redact_mcp_credentials_list(servers)
redacted_mcp_servers = sorted(_redact_mcp_credentials_list(servers), key=_mcp_server_display_order)
if connected_app_view is True and is_ui_session_credential(user_api_key_dict):
reachable_ids: Final = await _connected_app_reachable_server_ids(user_api_key_dict)

154
litellm/proxy/management_endpoints/password_endpoints.py Normal file
View file

@ -0,0 +1,154 @@
"""
Self-service password management.
/user/password/change
Deliberately NOT wrapped in `management_endpoint_wrapper`: the wrapper emits
request kwargs to OTEL spans, which would log plaintext passwords. The audit
signal is emitted by hand below, with field names only, never values.
"""
from typing import TYPE_CHECKING, Annotated, Final
from fastapi import APIRouter, Depends, HTTPException
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
UI_TEAM_ID,
ChangePasswordRequest,
ChangePasswordResponse,
CommonProxyErrors,
HTTPExceptionErrorDetail,
LitellmTableNames,
UserAPIKeyAuth,
)
from litellm.proxy.auth.login_utils import PASSWORD_SESSION_METADATA
from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_helpers.audit_logs import create_object_audit_log
from litellm.proxy.utils import hash_password, verify_password
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.user_repository import UserRepository
if TYPE_CHECKING:
from prisma import models as prisma_models
from prisma import types as prisma_types
from litellm.proxy.utils import PrismaClient
router: Final = APIRouter()
_PASSWORD_CHANGED_AUDIT_VALUES: Final = '{"fields_changed": ["password"]}'
_KEY_METADATA: Final = TypeAdapter(dict[str, object])
def _error_detail(message: str) -> HTTPExceptionErrorDetail:
detail: Final[HTTPExceptionErrorDetail] = {"error": message}
return detail
def _is_password_login_session(user_api_key_dict: UserAPIKeyAuth) -> bool:
if user_api_key_dict.team_id != UI_TEAM_ID:
return False
key_metadata: Final = _KEY_METADATA.validate_python(user_api_key_dict.metadata)
return all(key_metadata.get(k) == v for k, v in PASSWORD_SESSION_METADATA.items())
def _user_table(
prisma_client: "PrismaClient | None",
) -> "TableActions[prisma_models.LiteLLM_UserTable]":
user_table: Final[TableActions[prisma_models.LiteLLM_UserTable]] = UserRepository(prisma_client).table
return user_table
@router.post(
"/user/password/change",
tags=("Internal User management",),
dependencies=(Depends(user_api_key_auth),),
)
async def change_password(
data: ChangePasswordRequest,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> ChangePasswordResponse:
"""
Change the calling user's own password.
Only callable with the dashboard session issued by a username/password
login; SSO sessions and virtual keys are rejected with 403. Requires the
current password. The new password must differ from the
current one and satisfy the configured password policy
(`general_settings.password_policy_*`: minimum length, character classes,
and, when enabled, breached-password screening via haveibeenpwned.com).
A successful change lifts any pending forced password reset
(`password_reset_required`) on the account.
Parameters:
- current_password: str - The user's current password.
- new_password: str - The password to change to.
"""
from litellm.proxy.proxy_server import general_settings, litellm_proxy_admin_name, prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail=_error_detail(CommonProxyErrors.db_not_connected_error.value),
)
if not _is_password_login_session(user_api_key_dict):
raise HTTPException(
status_code=403,
detail=_error_detail(
"Passwords can only be changed from a dashboard session created by logging in with a password."
),
)
user_id: Final = user_api_key_dict.user_id
if user_id is None:
raise HTTPException(
status_code=400,
detail=_error_detail("No user is associated with this session, so there is no password to change."),
)
find_user: Final[prisma_types.LiteLLM_UserTableWhereInput] = {"user_id": user_id}
user_row: Final = await _user_table(prisma_client).find_first(where=find_user)
stored_password: Final = user_row.password if user_row is not None else None
if stored_password is None:
raise HTTPException(
status_code=400,
detail=_error_detail(
"This account has no password set, so there is no password to change. "
"Passwords are set through an invitation link (POST /invitation/new)."
),
)
if not verify_password(data.current_password, stored_password):
raise HTTPException(status_code=400, detail=_error_detail("Current password is incorrect."))
if data.new_password == data.current_password:
raise HTTPException(
status_code=400,
detail=_error_detail("New password must be different from the current password."),
)
validate_password_policy(data.new_password, general_settings)
await validate_password_not_breached(data.new_password, general_settings)
password_update: Final[prisma_types.LiteLLM_UserTableUpdateInput] = {
"password": hash_password(data.new_password),
"password_reset_required": False,
"last_breach_check_at": None,
}
await _user_table(prisma_client).update(where=find_user, data=password_update)
verbose_proxy_logger.info("Password changed via /user/password/change for user_id=%s", user_id)
await create_object_audit_log(
object_id=user_id,
action="updated",
litellm_changed_by=None,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
table_name=LitellmTableNames.USER_TABLE_NAME,
after_value=_PASSWORD_CHANGED_AUDIT_VALUES,
)
return ChangePasswordResponse(user_id=user_id, message="Password updated successfully.")

1
litellm/proxy/management_endpoints/ui_sso.py
View file

@ -3665,6 +3665,7 @@ class SSOAuthenticationHandler:
auth_header_name=general_settings.get("litellm_key_header_name", "Authorization"),
disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation,
server_root_path=get_server_root_path(),
password_reset_required=False,
)
from litellm.proxy.auth.login_utils import encode_ui_session_jwt

16
litellm/proxy/proxy_server.py
View file

@ -359,7 +359,7 @@ from litellm.proxy.auth.model_checks import (
get_mcp_server_ids,
get_team_models,
)
from litellm.proxy.auth.password_policy import validate_password_policy
from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy
from litellm.proxy.auth.user_api_key_auth import (
_fetch_global_spend_with_event_coordination,
user_api_key_auth,
@ -601,6 +601,9 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
from litellm.proxy.management_endpoints.organization_endpoints import (
router as organization_router,
)
from litellm.proxy.management_endpoints.password_endpoints import (
router as password_management_router,
)
from litellm.proxy.management_endpoints.prompt_caching_requests import (
router as prompt_caching_requests_router,
)
@ -16647,6 +16650,7 @@ async def onboarding(invite_link: str, request: Request):
auth_header_name=general_settings.get("litellm_key_header_name", "Authorization"),
disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation,
server_root_path=get_server_root_path(),
password_reset_required=False,
)
jwt_token: Final = jwt.encode(
cast(dict, returned_ui_token_object),
@ -16757,6 +16761,7 @@ async def _generate_onboarding_ui_session_token(user_obj: _UserTableRow) -> str:
auth_header_name=general_settings.get("litellm_key_header_name", "Authorization"),
disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation,
server_root_path=get_server_root_path(),
password_reset_required=False,
)
assert master_key is not None
return jwt.encode(
@ -16827,6 +16832,7 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request):
)
validate_password_policy(data.password, general_settings)
await validate_password_not_breached(data.password, general_settings)
hashed_pw: Final = hash_password(data.password)
current_time = litellm.utils.get_utc_datetime()
async with prisma_client.db.tx() as tx:
@ -16846,7 +16852,12 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request):
### UPDATE USER OBJECT ###
user_obj: Final[_UserTableRow | None] = await tx.litellm_usertable.update(
where={"user_id": invite_obj.user_id}, data={"password": hashed_pw}
where={"user_id": invite_obj.user_id},
data={
"password": hashed_pw,
"password_reset_required": False,
"last_breach_check_at": None,
},
)
if user_obj is None:
@ -19284,6 +19295,7 @@ app.include_router(pass_through_router)
app.include_router(health_router)
app.include_router(key_management_router)
app.include_router(internal_user_router)
app.include_router(password_management_router)
app.include_router(team_router)
app.include_router(ui_sso_router)
app.include_router(organization_router)

2
litellm/proxy/schema.prisma
View file

@ -246,6 +246,8 @@ model LiteLLM_UserTable {
organization_id String?
object_permission_id String?
password String?
password_reset_required Boolean?
last_breach_check_at DateTime?
teams String[] @default([])
user_role String?
max_budget Float?

104
litellm/rust_bridge/_native.pyi
View file

@ -93,6 +93,110 @@ class ResponsesWebSocketConnection:
def recv_text(self) -> Future[str | None]: ...
def close(self) -> Future[None]: ...
@final
class _CacheTestBinding:
@property
def kind(self) -> str: ...
def lookup(
self,
request: object,
*,
callback_kwargs: Mapping[str, object] | Sequence[object] | None = None,
) -> object: ...
def store(
self,
request: object,
response: object,
*,
callback_kwargs: Mapping[str, object] | None = None,
) -> None: ...
def lookup_batch(
self,
requests: Sequence[object],
*,
callback_kwargs: Sequence[object] | None = None,
) -> object: ...
def async_lookup(
self,
request: object,
*,
callback_kwargs: Mapping[str, object] | None = None,
) -> Future[object]: ...
def async_store(
self,
request: object,
response: object,
*,
callback_kwargs: Mapping[str, object] | None = None,
) -> Future[None]: ...
def async_lookup_batch(
self,
requests: Sequence[object],
*,
callback_kwargs: Sequence[object] | None = None,
) -> Future[object]: ...
def async_store_batch(
self,
requests: Sequence[object],
responses: Sequence[object],
*,
callback_result: object = None,
callback_kwargs: Mapping[str, object] | None = None,
) -> Future[object]: ...
def async_flush(self) -> Future[None]: ...
def ping(self) -> Future[object]: ...
@final
class _CacheTestHandle:
def __new__(cls, _uninstantiable: Never, /) -> Never: ...
@staticmethod
def memory(
*,
capacity: int = 200,
ttl_seconds: float = 600.0,
max_entry_bytes: int = 1048576,
) -> _CacheTestHandle: ...
@staticmethod
def redis(
url: str,
*,
ttl_seconds: float = 60.0,
namespace: str | None = None,
startup_nodes: Sequence[tuple[str, int]] | None = None,
) -> _CacheTestHandle: ...
@staticmethod
def disk(directory: str) -> _CacheTestHandle: ...
@staticmethod
def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ...
@staticmethod
def gcs(
bucket_name: str,
*,
gcs_path: str | None = None,
path_service_account: str | None = None,
endpoint: str | None = None,
token: str | None = None,
) -> _CacheTestHandle: ...
@staticmethod
def s3(
bucket: str,
*,
region: str,
endpoint_url: str | None = None,
key_prefix: str = "",
access_key_id: str | None = None,
secret_access_key: str | None = None,
session_token: str | None = None,
) -> _CacheTestHandle: ...
@property
def backend(self) -> str: ...
def _bind_facade(self, facade: object) -> None: ...
@final
class _CacheTestResolver:
def __new__(cls, namespace: object) -> _CacheTestResolver: ...
def resolve(self) -> _CacheTestBinding: ...
@final
class TokenCounter:
def __new__(cls, tokenizer_json: str) -> TokenCounter: ...

1
litellm/types/llms/custom_http.py
View file

@ -31,6 +31,7 @@ class httpxSpecialProvider(str, Enum):
UI = "ui"
Sandbox = "sandbox"
ModelCostMap = "model_cost_map"
PasswordBreachCheck = "password_breach_check"
VerifyTypes = str | bool | ssl.SSLContext

3
litellm/types/proxy/ui_sso.py
View file

@ -1,6 +1,6 @@
from typing import Literal
from typing_extensions import TypedDict
from typing_extensions import ReadOnly, TypedDict
class ReturnedUITokenObject(TypedDict):
@ -17,6 +17,7 @@ class ReturnedUITokenObject(TypedDict):
auth_header_name: str
disabled_non_admin_personal_key_creation: bool
server_root_path: str # e.g. `/litellm`
password_reset_required: ReadOnly[bool]
class ParsedOpenIDResult(TypedDict, total=False):

55
model_prices_and_context_window.json
View file

@ -43037,21 +43037,21 @@
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4-pro": {
"input_cost_per_token": 8.92272e-07,
"input_cost_per_token": 8.87226e-07,
"input_cost_per_token_cache_hit": 4.4e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"max_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 1.784544e-06,
"output_cost_per_token": 1.774452e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 7.4356e-08,
"cache_read_input_token_cost": 7.39355e-08,
"supports_audio_input": false,
"supports_pdf_input": false,
"supports_vision": false,
@ -73794,6 +73794,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-1.6": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 2.5e-07,
"input_cost_per_token_above_128k_tokens": 5e-07,
"litellm_provider": "openrouter",
@ -73815,6 +73816,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-1.6-flash": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 7.5e-08,
"input_cost_per_token_above_128k_tokens": 1e-07,
"litellm_provider": "openrouter",
@ -73855,6 +73857,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-2.0-code": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 5e-07,
"input_cost_per_token_above_128k_tokens": 1e-06,
"litellm_provider": "openrouter",
@ -76975,5 +76978,51 @@
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"xiaomi_mimo/mimo-v2.6-pro": {
"cache_read_input_token_cost": 3.6e-09,
"input_cost_per_token": 4.35e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 8.7e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
},
"xiaomi_mimo/mimo-v2.6-flash": {
"cache_read_input_token_cost": 2.8e-09,
"input_cost_per_token": 1.4e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2.8e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
}
}

2
schema.prisma
View file

@ -246,6 +246,8 @@ model LiteLLM_UserTable {
organization_id String?
object_permission_id String?
password String?
password_reset_required Boolean?
last_breach_check_at DateTime?
teams String[] @default([])
user_role String?
max_budget Float?

1
terraform/provider/tools/endpointaudit/coverage_allowlist.txt
View file

@ -86,6 +86,7 @@ POST /team/key/bulk_update
POST /team/permissions_bulk_update
POST /team/{team_id}/disable_logging
POST /user/bulk_update
POST /user/password/change
# Alternate method or path for functionality the provider already manages elsewhere
GET /credentials/by_model/{model_id}

3
tests/e2e/coverage_registry/llm_conversational.yaml
View file

@ -89,6 +89,9 @@
- {id: llm.chat_completions.together_ai.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together tool result round trip"}
- {id: llm.chat_completions.together_ai.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together cost header and spend row match the registry price"}
- {id: llm.chat_completions.together_ai.thinking.nonstream.effort_none_disables, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [effort_none_disables], source: "llm_translation/test_together_ai_e2e.py", rationale: "reasoning_effort=none maps to Together's reasoning disable toggle on hybrid models"}
- {id: llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "Native MiMo v2.6 rows price the cost header and spend row from the cost map"}
- {id: llm.chat_completions.xiaomi_mimo.thinking.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: thinking, streaming: stream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo reasoning deltas stream as reasoning_content"}
- {id: llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo tool calls are not dropped"}
- {id: llm.chat_completions.together_ai.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: structured_output, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "response_format json_schema reaches Together and constrains the reply"}
- {id: llm.chat_completions.together_ai.prompt_cache_5m.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: prompt_cache_5m, streaming: nonstream, assertions: [cache_hit, cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together prefix-cache reads bill at cache_read_input_token_cost, not full input price"}
- {id: llm.messages.together_ai.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together over /v1/messages streaming"}

1
tests/e2e/coverage_registry/schema.py
View file

@ -58,6 +58,7 @@ LlmRoute = Literal[
"openai",
"together_ai",
"vertex",
"xiaomi_mimo",
]
LlmCapability = Literal[

214
tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py Normal file
View file

@ -0,0 +1,214 @@
"""Live e2e: Xiaomi MiMo v2.6 through the gateway on /chat/completions.
Both native ``xiaomi_mimo/`` v2.6 rows (pro and flash) are registered via
``/model/new`` and driven against Xiaomi's own endpoint. What the gateway owes
us is that the reasoning chain surfaces as ``reasoning_content``, tool calls
survive translation, and the cost header plus spend row follow the proxy's own
cost-map price for the row (read back from ``/model/info``, never pinned here).
Requires XIAOMI_MIMO_API_KEY on the proxy; no skip gate.
"""
from __future__ import annotations
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_http import StreamingResponse, require_successful_call, unwrap
from lifecycle import ResourceManager
from models import (
ChatBody,
ChatMessage,
ChatResponse,
ChatTool,
ChatToolFunction,
CostMapEntry,
LiteLLMParamsBody,
OutMessage,
SpendLogRow,
)
from passthrough_client import PassthroughClient
from pydantic import BaseModel
pytestmark = pytest.mark.e2e
BACKENDS: Final = ("xiaomi_mimo/mimo-v2.6-pro", "xiaomi_mimo/mimo-v2.6-flash")
ARITHMETIC_PROMPT = "What is 17 + 26? Answer with just the number."
WEATHER_PROMPT = "What is the weather in Paris? Use the tool."
COUNTING_PROMPT = "Count from 1 to 50, one number per line."
WEATHER_TOOL = ChatTool(
function=ChatToolFunction(
name="get_weather",
description="Get the current weather for a location.",
parameters={
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
)
)
class _WeatherArgs(BaseModel):
location: str
class _StreamDelta(BaseModel):
content: str | None = None
reasoning_content: str | None = None
class _StreamChoice(BaseModel):
delta: _StreamDelta | None = None
class _StreamChunk(BaseModel):
choices: list[_StreamChoice] = []
def _approx_equal(actual: float, expected: float) -> bool:
return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)
@pytest.fixture(scope="module")
def registry(client: PassthroughClient) -> dict[str, CostMapEntry]:
return client.proxy.model_cost_map()
def _register(client: PassthroughClient, resources: ResourceManager, backend: str) -> tuple[str, str]:
model = f"e2e-xiaomi-{unique_marker()}"
model_id = client.proxy.create_model(
model, LiteLLMParamsBody(model=backend, api_key="os.environ/XIAOMI_MIMO_API_KEY")
)
resources.defer(lambda: client.proxy.delete_model(model_id))
return model, resources.key()
def _message(response: ChatResponse) -> OutMessage:
assert response.choices, f"Xiaomi returned no choices: {response}"
message = response.choices[0].message
assert message is not None, f"Xiaomi choice has no message: {response}"
return message
def _deltas(result: StreamingResponse) -> list[_StreamDelta]:
require_successful_call(result)
assert result.is_streaming, f"response was not streamed: {result.headers}"
assert not result.stream_error, f"stream errored: {result.stream_error}"
assert result.stream_done, f"stream never reached [DONE]: {result.stream_events[-3:]}"
return [
choice.delta
for event in result.stream_events
for choice in _StreamChunk.model_validate_json(event).choices
if choice.delta is not None
]
@pytest.mark.parametrize("backend", BACKENDS)
class TestXiaomiMimoChatCompletions:
@pytest.mark.covers("llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged")
def test_cost_header_and_spend_row_match_the_registry_price(
self,
client: PassthroughClient,
resources: ResourceManager,
registry: dict[str, CostMapEntry],
backend: str,
) -> None:
price = registry.get(backend)
assert price is not None, f"{backend} has no row in the proxy's cost map, so native calls would bill $0"
assert price.litellm_provider == "xiaomi_mimo", f"{backend} is filed under the wrong provider: {price}"
assert price.input_cost_per_token and price.output_cost_per_token, f"{backend} carries no price: {price}"
model, key = _register(client, resources, backend)
result = client.proxy.transport.send(
"/chat/completions",
headers=client.proxy.transport.bearer(key),
json=ChatBody(
model=model,
messages=[ChatMessage(role="user", content=f"{ARITHMETIC_PROMPT} {unique_marker()}")],
max_tokens=1024,
),
)
require_successful_call(result)
response = ChatResponse.model_validate_json(result.body)
message = _message(response)
assert message.content and "43" in message.content, f"answer lost: {message}"
assert message.reasoning_content, f"{backend} reasons, but no reasoning_content came back: {message}"
usage = response.usage
assert usage is not None and usage.prompt_tokens and usage.completion_tokens, (
f"response carries no usage, so the cost cannot be real: {result.body[:300]}"
)
header_cost = result.response_cost
assert header_cost is not None and header_cost > 0, (
f"x-litellm-response-cost header missing or non-positive: {result.headers}"
)
cached = (usage.prompt_tokens_details.cached_tokens or 0) if usage.prompt_tokens_details else 0
expected = (
(usage.prompt_tokens - cached) * price.input_cost_per_token
+ cached * (price.cache_read_input_token_cost or 0.0)
+ usage.completion_tokens * price.output_cost_per_token
)
assert _approx_equal(header_cost, expected), (
f"header cost {header_cost} disagrees with the registry price for {backend} at {usage}: expected {expected}"
)
def _priced(rows: list[SpendLogRow]) -> bool:
return any(row.spend is not None and row.spend > 0 for row in rows)
rows = client.proxy.poll_logs_for_key(key, predicate=_priced)
priced = [row for row in rows if row.spend is not None and row.spend > 0]
assert priced, f"no priced spend row landed for key {key}; got {rows}"
row = priced[0]
assert row.custom_llm_provider == "xiaomi_mimo", f"spend row misattributed: {row}"
assert row.spend is not None and _approx_equal(row.spend, header_cost), (
f"logged spend {row.spend} disagrees with the x-litellm-response-cost header {header_cost}"
)
@pytest.mark.covers("llm.chat_completions.xiaomi_mimo.thinking.stream.works")
def test_reasoning_and_answer_stream_as_deltas(
self, client: PassthroughClient, resources: ResourceManager, backend: str
) -> None:
model, key = _register(client, resources, backend)
deltas = _deltas(
client.proxy.chat_stream(
key,
ChatBody(
model=model,
messages=[ChatMessage(role="user", content=COUNTING_PROMPT)],
max_tokens=2048,
stream=True,
),
)
)
reasoning = "".join(delta.reasoning_content or "" for delta in deltas)
content = "".join(delta.content or "" for delta in deltas)
assert reasoning, f"stream carried no reasoning_content deltas: {deltas[:5]}"
assert "50" in content, f"streamed answer lost: {content[:300]!r}"
@pytest.mark.covers("llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works")
def test_tool_call_is_returned(self, client: PassthroughClient, resources: ResourceManager, backend: str) -> None:
model, key = _register(client, resources, backend)
message = _message(
unwrap(
client.proxy.chat(
key,
ChatBody(
model=model,
messages=[ChatMessage(role="user", content=WEATHER_PROMPT)],
tools=[WEATHER_TOOL],
max_tokens=1024,
),
)
)
)
assert message.tool_calls, f"{backend} dropped the tool call: {message}"
call = message.tool_calls[0]
assert call.id, f"tool call carries no id, so a tool result cannot answer it: {call}"
assert call.function.name == "get_weather", f"wrong tool called: {call}"
assert call.function.arguments, f"tool call carries no arguments: {call}"
args = _WeatherArgs.model_validate_json(call.function.arguments)
assert "paris" in args.location.lower(), f"tool arguments lost the location: {args}"

10
tests/test_litellm/llms/xai/test_xai_chat_transformation.py
View file

@ -90,6 +90,16 @@ class TestXAIReasoningTokenFolding:
assert response.usage.total_tokens == 999
def test_max_completion_tokens_is_accepted_and_mapped_to_max_tokens() -> None:
optional_params = litellm.get_optional_params(
model="grok-4.20",
custom_llm_provider="xai",
max_completion_tokens=64,
)
assert optional_params["max_tokens"] == 64, optional_params
assert "max_completion_tokens" not in optional_params, optional_params
class TestXAIParallelToolCalls:
"""Test suite for XAI parallel tool calls functionality."""

281
tests/test_litellm/proxy/auth/test_login_utils.py
View file

@ -5,12 +5,15 @@ This module tests the refactored login logic that was moved from proxy_server.py
to login_utils.py for better reusability.
"""
import hashlib
import os
from collections.abc import Mapping
from contextlib import ExitStack
from typing import TYPE_CHECKING, Final
from datetime import datetime, timedelta, timezone
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
if TYPE_CHECKING:
@ -34,6 +37,7 @@ def _unlimited_throttle():
from litellm.constants import LITELLM_PROXY_ADMIN_NAME
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy._types import (
LiteLLM_UserTable,
LitellmUserRoles,
@ -46,8 +50,13 @@ from litellm.proxy.auth.login_utils import (
authenticate_user,
get_ui_credentials,
is_env_credential_login_enabled,
screen_login_password_for_breach,
)
# Successful DB-user logins schedule the background HIBP screen; disable it so
# no test ever does live network I/O to haveibeenpwned.com from CI.
_POLICY_NO_BREACH_CHECK = {"password_policy_check_breached_passwords": False}
def test_get_ui_credentials_prefers_explicit_password():
"""The configured UI password should be returned when available."""
@ -326,6 +335,7 @@ async def test_authenticate_user_email_case_insensitive_login():
master_key=master_key,
prisma_client=mock_prisma_client,
throttle=_unlimited_throttle(),
general_settings=_POLICY_NO_BREACH_CHECK,
)
result_lower = await authenticate_user(
username=stored_email,
@ -333,6 +343,7 @@ async def test_authenticate_user_email_case_insensitive_login():
master_key=master_key,
prisma_client=mock_prisma_client,
throttle=_unlimited_throttle(),
general_settings=_POLICY_NO_BREACH_CHECK,
)
assert result_mixed.user_id == result_lower.user_id == "test-user-123"
@ -576,6 +587,7 @@ async def test_authenticate_user_database_login_with_non_ascii_password():
master_key=master_key,
prisma_client=mock_prisma_client,
throttle=_unlimited_throttle(),
general_settings=_POLICY_NO_BREACH_CHECK,
)
assert isinstance(result, LoginResult)
@ -721,7 +733,12 @@ async def _db_login(throttle, username: str, password: str, *, correct: bool):
),
):
return await authenticate_user(
username=username, password=password, master_key="sk-master", prisma_client=MagicMock(), throttle=throttle
username=username,
password=password,
master_key="sk-master",
prisma_client=MagicMock(),
throttle=throttle,
general_settings=_POLICY_NO_BREACH_CHECK,
)
@ -2064,3 +2081,265 @@ class TestIsEnvCredentialLoginEnabled:
with ExitStack() as stack:
_patch_sso_configured(stack, configured=False)
assert is_env_credential_login_enabled({"disable_password_login_when_sso_enabled": True}) is True
def _db_user_row(*, password: str, password_reset_required: bool | None = None, last_breach_check_at=None):
hashed = hash_token(token=password)
row = MagicMock()
row.user_id = "reset-user-1"
row.user_email = "reset@example.com"
row.password = hashed
row.user_role = LitellmUserRoles.INTERNAL_USER
row.password_reset_required = password_reset_required
row.last_breach_check_at = last_breach_check_at
return row
def _prisma_with_user(row) -> MagicMock:
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=row)
mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value=row)
return mock_prisma_client
_DB_LOGIN_ENV = {
"DATABASE_URL": "postgresql://test:test@localhost/test",
"UI_USERNAME": "admin",
"UI_PASSWORD": "admin-password",
}
class TestPasswordResetRequiredSessionMinting:
"""A user flagged `password_reset_required` must receive a UI session key
restricted to the change-password endpoint (server-side enforcement, so a
script driving the management API with the session key is blocked too);
an unflagged user must keep getting an unrestricted key."""
async def _login(self, mock_prisma_client) -> tuple[LoginResult, dict]:
with patch.dict(os.environ, _DB_LOGIN_ENV):
with patch( # test-quality-ok: asserting the minted key's restriction requires seeing its kwargs
"litellm.proxy.auth.login_utils.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"token": "session-token"},
) as mock_generate_key:
result = await authenticate_user(
username="reset@example.com",
password="Str0ng!Passw0rd",
master_key="sk-1234",
prisma_client=mock_prisma_client,
throttle=_unlimited_throttle(),
general_settings=_POLICY_NO_BREACH_CHECK,
)
return result, mock_generate_key.call_args.kwargs
@pytest.mark.asyncio
async def test_flagged_user_gets_key_restricted_to_change_password(self):
row = _db_user_row(password="Str0ng!Passw0rd", password_reset_required=True)
result, key_kwargs = await self._login(_prisma_with_user(row))
assert key_kwargs["allowed_routes"] == ["/user/password/change"]
assert key_kwargs["metadata"] == {"login_method": "username_password", "password_reset_required": True}
assert result.password_reset_required is True
@pytest.mark.asyncio
async def test_unflagged_user_gets_unrestricted_key(self):
row = _db_user_row(password="Str0ng!Passw0rd", password_reset_required=None)
result, key_kwargs = await self._login(_prisma_with_user(row))
assert key_kwargs["allowed_routes"] is None
assert key_kwargs["metadata"] == {"login_method": "username_password"}
assert result.password_reset_required is False
async def _login_with_screen_result(self, mock_prisma_client, breached: bool) -> tuple[LoginResult, dict, dict]:
with patch.dict(os.environ, _DB_LOGIN_ENV):
with patch( # test-quality-ok: asserting the minted key's restriction requires seeing its kwargs
"litellm.proxy.auth.login_utils.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"token": "session-token"},
) as mock_generate_key:
with (
patch( # test-quality-ok: authenticate_user has no HIBP client seam; the screen itself is tested against MockTransport below
"litellm.proxy.auth.login_utils.screen_login_password_for_breach",
new_callable=AsyncMock,
return_value=breached,
) as mock_screen
):
result = await authenticate_user(
username="reset@example.com",
password="Str0ng!Passw0rd",
master_key="sk-1234",
prisma_client=mock_prisma_client,
throttle=_unlimited_throttle(),
general_settings=_POLICY_NO_BREACH_CHECK,
)
return result, mock_generate_key.call_args.kwargs, mock_screen.call_args.kwargs
@pytest.mark.asyncio
async def test_login_screens_with_row_state_before_minting(self):
"""The login must hand the screen the row's recheck timestamp, or the
24h throttle can never work."""
checked_at = datetime.now(timezone.utc) - timedelta(hours=1)
row = _db_user_row(password="Str0ng!Passw0rd", last_breach_check_at=checked_at)
mock_prisma_client = _prisma_with_user(row)
_, _, screen_kwargs = await self._login_with_screen_result(mock_prisma_client, breached=False)
assert screen_kwargs["user_id"] == "reset-user-1"
assert screen_kwargs["password"] == "Str0ng!Passw0rd"
assert screen_kwargs["last_breach_check_at"] == checked_at
assert screen_kwargs["prisma_client"] is mock_prisma_client
@pytest.mark.asyncio
async def test_fresh_breach_hit_restricts_the_current_session(self):
"""A breach found during THIS login must restrict THIS session, not
just the next one."""
row = _db_user_row(password="Str0ng!Passw0rd", password_reset_required=None)
mock_prisma_client = _prisma_with_user(row)
result, key_kwargs, _ = await self._login_with_screen_result(mock_prisma_client, breached=True)
assert key_kwargs["allowed_routes"] == ["/user/password/change"]
assert key_kwargs["metadata"] == {"login_method": "username_password", "password_reset_required": True}
assert result.password_reset_required is True
def _sha1_upper(password: str) -> str:
return hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper()
def _client_with_transport(handler) -> AsyncHTTPHandler:
http_handler = AsyncHTTPHandler()
http_handler.client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
return http_handler
def _client_returning_breach_hit(password: str) -> AsyncHTTPHandler:
body = f"{_sha1_upper(password)[5:]}:42"
return _client_with_transport(lambda request: httpx.Response(200, text=body))
def _client_returning_no_hit() -> AsyncHTTPHandler:
return _client_with_transport(lambda request: httpx.Response(200, text="0000000000000000000000000000000000A:3"))
def _client_never_called() -> AsyncHTTPHandler:
def handler(request: httpx.Request) -> httpx.Response:
raise AssertionError(f"unexpected HTTP call to {request.url}")
return _client_with_transport(handler)
class TestScreenLoginPasswordForBreach:
"""The awaited login-time screen: flags a breached password for a forced
reset, stamps the recheck timestamp, rechecks at most every 24h, returns
the breach verdict so the login can restrict the session it is minting,
and never raises into the login."""
@pytest.mark.asyncio
async def test_breached_password_sets_reset_flag_and_timestamp(self):
password = "Password123!"
mock_prisma_client = _prisma_with_user(None)
breached = await screen_login_password_for_breach(
user_id="reset-user-1",
password=password,
last_breach_check_at=None,
general_settings={},
prisma_client=mock_prisma_client,
client=_client_returning_breach_hit(password),
)
assert breached is True
update_kwargs = mock_prisma_client.db.litellm_usertable.update.call_args.kwargs
assert update_kwargs["where"] == {"user_id": "reset-user-1"}
assert update_kwargs["data"]["password_reset_required"] is True
assert isinstance(update_kwargs["data"]["last_breach_check_at"], datetime)
@pytest.mark.asyncio
async def test_clean_password_stamps_timestamp_without_flag(self):
mock_prisma_client = _prisma_with_user(None)
breached = await screen_login_password_for_breach(
user_id="reset-user-1",
password="Str0ng!Passw0rd",
last_breach_check_at=None,
general_settings={},
prisma_client=mock_prisma_client,
client=_client_returning_no_hit(),
)
assert breached is False
update_kwargs = mock_prisma_client.db.litellm_usertable.update.call_args.kwargs
assert "password_reset_required" not in update_kwargs["data"]
assert isinstance(update_kwargs["data"]["last_breach_check_at"], datetime)
@pytest.mark.asyncio
async def test_skips_hibp_when_checked_within_24_hours(self):
mock_prisma_client = _prisma_with_user(None)
breached = await screen_login_password_for_breach(
user_id="reset-user-1",
password="Password123!",
last_breach_check_at=datetime.now(timezone.utc) - timedelta(hours=23),
general_settings={},
prisma_client=mock_prisma_client,
client=_client_never_called(),
)
assert breached is False
mock_prisma_client.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
async def test_rechecks_when_last_check_is_older_than_24_hours(self):
password = "Password123!"
mock_prisma_client = _prisma_with_user(None)
breached = await screen_login_password_for_breach(
user_id="reset-user-1",
password=password,
last_breach_check_at=datetime.now(timezone.utc) - timedelta(hours=25),
general_settings={},
prisma_client=mock_prisma_client,
client=_client_returning_breach_hit(password),
)
assert breached is True
assert (
mock_prisma_client.db.litellm_usertable.update.call_args.kwargs["data"]["password_reset_required"] is True
)
@pytest.mark.asyncio
async def test_skips_hibp_when_check_disabled(self):
mock_prisma_client = _prisma_with_user(None)
breached = await screen_login_password_for_breach(
user_id="reset-user-1",
password="Password123!",
last_breach_check_at=None,
general_settings=_POLICY_NO_BREACH_CHECK,
prisma_client=mock_prisma_client,
client=_client_never_called(),
)
assert breached is False
mock_prisma_client.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
async def test_db_failure_never_raises_but_still_reports_the_breach(self):
"""A failed flag write must not fail the login, but the breach verdict
still has to restrict the session being minted right now."""
password = "Password123!"
mock_prisma_client = _prisma_with_user(None)
mock_prisma_client.db.litellm_usertable.update = AsyncMock(side_effect=RuntimeError("db down"))
assert (
await screen_login_password_for_breach(
user_id="reset-user-1",
password=password,
last_breach_check_at=None,
general_settings={},
prisma_client=mock_prisma_client,
client=_client_returning_breach_hit(password),
)
is True
)

144
tests/test_litellm/proxy/auth/test_onboarding.py
View file

@ -8,15 +8,20 @@ Covers the security behavior of:
session key only after the password is written
"""
import hashlib
from datetime import timedelta
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import jwt
import pytest
import respx
from fastapi import HTTPException
import litellm
from litellm.proxy._types import InvitationClaim
from litellm.proxy._types import InvitationClaim, ProxyException
_POLICY_NO_BREACH_CHECK = {"password_policy_check_breached_passwords": False}
# ---------------------------------------------------------------------------
# Helpers
@ -386,7 +391,9 @@ async def test_claim_token_rejects_concurrent_reuse_before_password_write():
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma),
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
patch("litellm.proxy.proxy_server.general_settings", {}),
patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
patch(
"litellm.proxy.proxy_server.generate_key_helper_fn",
new_callable=AsyncMock,
@ -426,7 +433,9 @@ async def test_claim_token_sets_accepted_at_after_password_written():
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma),
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
patch("litellm.proxy.proxy_server.general_settings", {}),
patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
patch("litellm.proxy.proxy_server.premium_user", False),
patch(
"litellm.proxy.proxy_server.generate_key_helper_fn",
@ -454,6 +463,10 @@ async def test_claim_token_sets_accepted_at_after_password_written():
call_kwargs = prisma.db.litellm_usertable.update.call_args
assert call_kwargs.kwargs["where"] == {"user_id": "user-123"}
assert "password" in call_kwargs.kwargs["data"]
# A freshly claimed, policy-screened password lifts any pending forced
# reset and re-arms the login-time breach screen.
assert call_kwargs.kwargs["data"]["password_reset_required"] is False
assert call_kwargs.kwargs["data"]["last_breach_check_at"] is None
# is_accepted was flipped to True on the invitation link
prisma.db.litellm_invitationlink.update.assert_called_once()
@ -483,7 +496,9 @@ async def test_claim_token_rolls_back_invite_when_session_key_mint_fails():
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma),
patch("litellm.proxy.proxy_server.master_key", "sk-test"),
patch("litellm.proxy.proxy_server.general_settings", {}),
patch( # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
patch(
"litellm.proxy.proxy_server.generate_key_helper_fn",
new_callable=AsyncMock,
@ -505,3 +520,124 @@ async def test_claim_token_rolls_back_invite_when_session_key_mint_fails():
}
assert rollback_kwargs["data"]["accepted_at"] is None
assert rollback_kwargs["data"]["is_accepted"] is False
# ---------------------------------------------------------------------------
# POST /onboarding/claim_token - password policy
# ---------------------------------------------------------------------------
def _hibp_url_for(password: str) -> str:
sha1 = hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper()
return f"https://api.pwnedpasswords.com/range/{sha1[:5]}"
def _hibp_suffix_for(password: str) -> str:
return hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper()[5:]
@pytest.mark.asyncio
async def test_claim_token_rejects_short_password_before_consuming_invite():
"""Default policy requires 12 characters; the invite must stay claimable."""
from litellm.proxy.proxy_server import claim_onboarding_link
invite = _make_invite(is_accepted=False)
prisma = _make_prisma(invite, _make_user())
request = _make_claim_request(_make_onboarding_token())
data = InvitationClaim(
invitation_link="invite-abc",
user_id="user-123",
password="Sh0rt!pw",
)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
patch("litellm.proxy.proxy_server.master_key", "sk-test"), # test-quality-ok: same as above
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: same as above
):
with pytest.raises(ProxyException) as exc_info:
await claim_onboarding_link(data=data, request=request)
assert exc_info.value.code == "400"
assert "at least 12 characters" in exc_info.value.message
prisma.db.litellm_invitationlink.update_many.assert_not_called()
prisma.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
@respx.mock
async def test_claim_token_rejects_breached_password_before_consuming_invite():
"""A password found in the HIBP corpus must be rejected and never stored."""
from litellm.proxy.proxy_server import claim_onboarding_link
password = "P@ssword123456"
respx.get(_hibp_url_for(password)).mock(
return_value=httpx.Response(200, text=f"{_hibp_suffix_for(password)}:1387")
)
invite = _make_invite(is_accepted=False)
prisma = _make_prisma(invite, _make_user())
request = _make_claim_request(_make_onboarding_token())
data = InvitationClaim(
invitation_link="invite-abc",
user_id="user-123",
password=password,
)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
patch("litellm.proxy.proxy_server.master_key", "sk-test"), # test-quality-ok: same as above
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: same as above
):
with pytest.raises(ProxyException) as exc_info:
await claim_onboarding_link(data=data, request=request)
assert exc_info.value.code == "400"
assert "data breaches" in exc_info.value.message
prisma.db.litellm_invitationlink.update_many.assert_not_called()
prisma.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
@respx.mock
async def test_claim_token_fails_open_when_hibp_unreachable():
"""An HIBP outage must never block onboarding: the claim proceeds."""
from litellm.proxy.proxy_server import claim_onboarding_link
password = "NewP@ssw0rd-2026"
respx.get(_hibp_url_for(password)).mock(side_effect=httpx.ConnectError("no route to host"))
invite = _make_invite(is_accepted=False)
user = _make_user()
prisma = _make_prisma(invite, user)
request = _make_claim_request(_make_onboarding_token())
data = InvitationClaim(
invitation_link="invite-abc",
user_id="user-123",
password=password,
)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: claim_onboarding_link reads proxy_server module globals; no injection seam
patch("litellm.proxy.proxy_server.master_key", "sk-test"), # test-quality-ok: same as above
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: same as above
patch("litellm.proxy.proxy_server.premium_user", False), # test-quality-ok: same as above
patch( # test-quality-ok: same as above
"litellm.proxy.proxy_server.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"token": "sk-generated-key", "user_id": "user-123"},
),
patch( # test-quality-ok: same as above
"litellm.proxy.proxy_server.get_custom_url",
return_value="http://localhost:4000/",
),
patch( # test-quality-ok: same as above
"litellm.proxy.proxy_server.get_disabled_non_admin_personal_key_creation",
return_value=False,
),
patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""), # test-quality-ok: same as above
):
result = await claim_onboarding_link(data=data, request=request)
assert "token" in result
prisma.db.litellm_usertable.update.assert_called_once()

209
tests/test_litellm/proxy/auth/test_password_policy.py
View file

@ -2,22 +2,56 @@
Tests for the configurable password-strength policy in
`litellm.proxy.auth.password_policy`, enforced on every path that persists a
new or changed password for a locally-managed user.
The breach-check (HIBP) tests inject a real AsyncHTTPHandler wrapping an
httpx.MockTransport, so no network is touched and nothing is monkeypatched.
"""
import asyncio
import hashlib
import httpx
import pytest
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy.auth.password_policy import (
DEFAULT_MIN_LENGTH,
MIN_ALLOWED_LENGTH,
PasswordPolicy,
get_password_policy,
validate_password_not_breached,
validate_password_policy,
validate_passwords_bulk,
)
STRONG_PASSWORD = "Str0ng!Passw0rd"
def _sha1_upper(password: str) -> str:
return hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper()
def _client_with_transport(handler) -> AsyncHTTPHandler:
http_handler = AsyncHTTPHandler()
http_handler.client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
return http_handler
def _client_never_called() -> AsyncHTTPHandler:
def handler(request: httpx.Request) -> httpx.Response:
raise AssertionError(f"unexpected HTTP call to {request.url}")
return _client_with_transport(handler)
def _client_returning(body: str, status_code: int = 200) -> AsyncHTTPHandler:
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(status_code, text=body)
return _client_with_transport(handler)
def test_get_password_policy_defaults_to_pif_baseline():
policy = get_password_policy({})
assert policy == PasswordPolicy(
@ -134,3 +168,178 @@ def test_validate_password_policy_rejects_unicode_letter_as_special_character():
def test_validate_password_policy_accepts_real_special_character_with_unicode_letters():
"""Same base password as the rejection test above, plus an actual symbol."""
assert validate_password_policy("Passwörd1234!", {}) is None
@pytest.mark.asyncio
async def test_breach_check_skipped_when_disabled():
result = await validate_password_not_breached(
password="password12345", # breached in reality, but the check is off
general_settings={"password_policy_check_breached_passwords": False},
client=_client_never_called(),
)
assert result is None
@pytest.mark.asyncio
async def test_rejects_breached_password():
password = "correct horse battery staple"
sha1 = _sha1_upper(password)
body = f"AAAA000000000000000000000000000000A:0\r\n{sha1[5:]}:42\r\nBBBB000000000000000000000000000000B:7"
with pytest.raises(ProxyException) as exc_info:
await validate_password_not_breached(password=password, general_settings={}, client=_client_returning(body))
assert exc_info.value.code == "400"
assert exc_info.value.type == ProxyErrorTypes.validation_error
assert exc_info.value.param == "password"
assert "data breaches" in exc_info.value.message
@pytest.mark.asyncio
async def test_only_sha1_prefix_leaves_the_proxy():
password = "a very secret password"
sha1 = _sha1_upper(password)
captured_requests: list[httpx.Request] = []
def handler(request: httpx.Request) -> httpx.Response:
captured_requests.append(request)
return httpx.Response(200, text="0000000000000000000000000000000000A:1")
result = await validate_password_not_breached(
password=password, general_settings={}, client=_client_with_transport(handler)
)
assert result is None
(request,) = captured_requests
assert request.url.path == f"/range/{sha1[:5]}"
assert sha1[5:] not in str(request.url)
assert request.headers["Add-Padding"] == "true"
assert "litellm" in request.headers["User-Agent"]
@pytest.mark.asyncio
async def test_ignores_padding_entries_with_zero_count():
"""HIBP padding entries (requested via Add-Padding) carry count 0 and must
not be treated as breaches when they collide with the password's suffix."""
password = "a padded-away password"
sha1 = _sha1_upper(password)
result = await validate_password_not_breached(
password=password, general_settings={}, client=_client_returning(f"{sha1[5:]}:0")
)
assert result is None
@pytest.mark.asyncio
async def test_accepts_password_absent_from_breach_corpus():
result = await validate_password_not_breached(
password="a genuinely novel password",
general_settings={},
client=_client_returning("0018A45C4D1DEF81644B54AB7F969B88D65:1\r\n00D4F6E8FA6EECAD2A3AA415EEC418D38EC:2"),
)
assert result is None
@pytest.mark.asyncio
async def test_breach_check_fails_open_on_network_error():
def handler(request: httpx.Request) -> httpx.Response:
raise httpx.ConnectError("no route to host")
result = await validate_password_not_breached(
password="password12345", # breached, but HIBP is unreachable
general_settings={},
client=_client_with_transport(handler),
)
assert result is None
@pytest.mark.asyncio
async def test_breach_check_fails_open_on_http_error_status():
result = await validate_password_not_breached(
password="password12345",
general_settings={},
client=_client_returning("service unavailable", status_code=503),
)
assert result is None
@pytest.mark.asyncio
async def test_breach_check_fails_open_on_malformed_response_body():
result = await validate_password_not_breached(
password="password12345",
general_settings={},
client=_client_returning(f"{_sha1_upper('password12345')[5:]}:not-a-number"),
)
assert result is None
@pytest.mark.asyncio
async def test_validate_passwords_bulk_screens_concurrently():
"""All HIBP lookups for a batch must be in flight at once: each handler
call stalls until every expected request has arrived, and a handler that
gives up waiting reports the password as breached. Serial awaiting (the
old per-user behavior) leaves each earlier request waiting forever for the
later ones, so every verdict comes back as a breach and the test fails."""
passwords = ("Uniqu3!Passw0rd-a", "Uniqu3!Passw0rd-b", "Uniqu3!Passw0rd-c")
suffix_by_prefix = {_sha1_upper(p)[:5]: _sha1_upper(p)[5:] for p in passwords}
all_arrived = asyncio.Event()
arrivals: list[str] = []
async def handler(request: httpx.Request) -> httpx.Response:
arrivals.append(request.url.path)
if len(arrivals) == len(passwords):
all_arrived.set()
try:
await asyncio.wait_for(all_arrived.wait(), timeout=5)
except TimeoutError:
return httpx.Response(200, text=f"{suffix_by_prefix[request.url.path.rsplit('/', 1)[-1]]}:1")
return httpx.Response(200, text="0000000000000000000000000000000000A:1")
verdicts = await validate_passwords_bulk(passwords, {}, client=_client_with_transport(handler))
assert set(arrivals) == {f"/range/{prefix}" for prefix in suffix_by_prefix}
assert all(verdicts[p] is None for p in passwords)
@pytest.mark.asyncio
async def test_validate_passwords_bulk_deduplicates_lookups():
"""500 users sharing one password must cost exactly one HIBP lookup."""
password = "Sh@red-Passw0rd!"
request_count = 0
def handler(request: httpx.Request) -> httpx.Response:
nonlocal request_count
request_count += 1
return httpx.Response(200, text="0000000000000000000000000000000000A:1")
verdicts = await validate_passwords_bulk((password,) * 500, {}, client=_client_with_transport(handler))
assert request_count == 1
assert verdicts == {password: None}
@pytest.mark.asyncio
async def test_validate_passwords_bulk_mixed_verdicts():
"""Weak passwords are rejected without an HIBP lookup; breached ones get
the breach error; acceptable ones map to None."""
breached = "Br3ached!Passw0rd"
clean = "Cl3an!!Passw0rd42"
weak = "short1!"
breached_sha1 = _sha1_upper(breached)
looked_up_prefixes: list[str] = []
def handler(request: httpx.Request) -> httpx.Response:
looked_up_prefixes.append(request.url.path.rsplit("/", 1)[-1])
if request.url.path == f"/range/{breached_sha1[:5]}":
return httpx.Response(200, text=f"{breached_sha1[5:]}:99")
return httpx.Response(200, text="0000000000000000000000000000000000A:1")
verdicts = await validate_passwords_bulk((breached, clean, weak), {}, client=_client_with_transport(handler))
assert _sha1_upper(weak)[:5] not in looked_up_prefixes
assert verdicts[clean] is None
assert "data breaches" in verdicts[breached].message
assert verdicts[breached].code == "400"
assert "12 characters" in verdicts[weak].message
@pytest.mark.asyncio
async def test_validate_passwords_bulk_empty_batch_makes_no_lookups():
verdicts = await validate_passwords_bulk((), {}, client=_client_never_called())
assert verdicts == {}

274
tests/test_litellm/proxy/auth/test_route_checks.py
View file

@ -3,7 +3,6 @@ from datetime import datetime
from typing import Final
from unittest.mock import MagicMock, patch
import pytest
from fastapi import HTTPException, Request
@ -39,7 +38,7 @@ def test_non_admin_config_update_route_rejected():
request.query_params = {}
# Test that calling /config/update route raises HTTPException with 403 status
with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info:
with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info:
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER.value,
@ -50,9 +49,8 @@ def test_non_admin_config_update_route_rejected():
)
# Verify the exception is raised with the correct message
assert (
"Only proxy admin can be used to generate, delete, update info for new keys/users/teams"
in str(exc_info.value)
assert "Only proxy admin can be used to generate, delete, update info for new keys/users/teams" in str(
exc_info.value
)
assert "Route=/config/update" in str(exc_info.value)
assert "Your role=internal_user" in str(exc_info.value)
@ -158,7 +156,7 @@ def test_user_banner_update_rejected_for_non_admin():
request = MagicMock(spec=Request)
request.query_params = {}
with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info:
with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info:
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER.value,
@ -733,9 +731,7 @@ def test_virtual_key_llm_api_route_includes_passthrough_prefix(route):
valid_token = UserAPIKeyAuth(user_id="test_user", allowed_routes=["llm_api_routes"])
result = RouteChecks.is_virtual_key_allowed_to_call_route(
route=route, valid_token=valid_token
)
result = RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=valid_token)
assert result is True
@ -760,9 +756,7 @@ def test_virtual_key_llm_api_routes_allows_google_routes(route):
valid_token = UserAPIKeyAuth(user_id="test_user", allowed_routes=["llm_api_routes"])
result = RouteChecks.is_virtual_key_allowed_to_call_route(
route=route, valid_token=valid_token
)
result = RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=valid_token)
assert result is True
@ -832,18 +826,14 @@ def test_google_routes_with_dynamic_model_names_accessible_to_internal_users():
)
# If no exception is raised, the test passes
except Exception as e:
pytest.fail(
f"Internal user should be able to access Google generateContent route. Got error: {str(e)}"
)
pytest.fail(f"Internal user should be able to access Google generateContent route. Got error: {e!s}")
def test_virtual_key_allowed_routes_with_multiple_litellm_routes_member_names():
"""Test that virtual key works with multiple LiteLLMRoutes member names in allowed_routes"""
# Create a UserAPIKeyAuth with multiple LiteLLMRoutes member names
valid_token = UserAPIKeyAuth(
user_id="test_user", allowed_routes=["openai_routes", "info_routes"]
)
valid_token = UserAPIKeyAuth(user_id="test_user", allowed_routes=["openai_routes", "info_routes"])
# Test that routes from both groups are allowed
result1 = RouteChecks.is_virtual_key_allowed_to_call_route(
@ -897,13 +887,9 @@ def test_virtual_key_allowed_routes_with_no_member_names_only_explicit():
)
# Test that explicit routes are allowed
result1 = RouteChecks.is_virtual_key_allowed_to_call_route(
route="/chat/completions", valid_token=valid_token
)
result1 = RouteChecks.is_virtual_key_allowed_to_call_route(route="/chat/completions", valid_token=valid_token)
result2 = RouteChecks.is_virtual_key_allowed_to_call_route(
route="/custom/route", valid_token=valid_token
)
result2 = RouteChecks.is_virtual_key_allowed_to_call_route(route="/custom/route", valid_token=valid_token)
assert result1 is True
assert result2 is True
@ -1301,9 +1287,7 @@ def test_virtual_key_without_llm_api_routes_cannot_access_pass_through():
)
assert exc_info.value.status_code == 403
assert "Virtual key is not allowed to call this route" in str(
exc_info.value.detail
)
assert "Virtual key is not allowed to call this route" in str(exc_info.value.detail)
def test_check_passthrough_route_access_key_metadata_exact_match():
@ -1762,9 +1746,7 @@ def test_videos_route_accessible_to_internal_users():
)
# If no exception is raised, the test passes
except Exception as e:
pytest.fail(
f"Internal user should be able to access /v1/videos route. Got error: {str(e)}"
)
pytest.fail(f"Internal user should be able to access /v1/videos route. Got error: {e!s}")
def test_videos_route_with_virtual_key_llm_api_routes():
@ -1786,12 +1768,8 @@ def test_videos_route_with_virtual_key_llm_api_routes():
]
for route in test_routes:
result = RouteChecks.is_virtual_key_allowed_to_call_route(
route=route, valid_token=valid_token
)
assert (
result is True
), f"Virtual key with llm_api_routes should be able to access {route}"
result = RouteChecks.is_virtual_key_allowed_to_call_route(route=route, valid_token=valid_token)
assert result is True, f"Virtual key with llm_api_routes should be able to access {route}"
def test_non_proxy_admin_wildcard_allowed_routes():
@ -1862,9 +1840,7 @@ def test_proxy_admin_viewer_can_access_global_spend_tags():
)
# If no exception is raised, the test passes
except Exception as e:
pytest.fail(
f"proxy_admin_viewer should be able to access /global/spend/tags route. Got error: {str(e)}"
)
pytest.fail(f"proxy_admin_viewer should be able to access /global/spend/tags route. Got error: {e!s}")
# Routes returning proxy-wide spend across every team / customer / api_key.
@ -1892,7 +1868,7 @@ def test_internal_user_blocked_from_global_spend_routes(route):
request = MagicMock(spec=Request)
request.query_params = {}
with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info:
with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info:
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER.value,
@ -1921,7 +1897,7 @@ def test_internal_user_view_only_blocked_from_global_spend_routes(route):
request = MagicMock(spec=Request)
request.query_params = {}
with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info:
with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info:
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value,
@ -2023,9 +1999,7 @@ def test_proxy_admin_viewer_can_access_audit_logs(route):
request_data={},
)
except Exception as e:
pytest.fail(
f"proxy_admin_viewer should be able to access {route} route. Got error: {str(e)}"
)
pytest.fail(f"proxy_admin_viewer should be able to access {route} route. Got error: {e!s}")
# ── Admin Viewer parity: Logs page endpoints ──────────────────────────────────
@ -2088,9 +2062,7 @@ def test_proxy_admin_viewer_can_access_logs_page_endpoints(route):
request_data={},
)
except Exception as e:
pytest.fail(
f"proxy_admin_viewer should be able to access {route}. Got error: {str(e)}"
)
pytest.fail(f"proxy_admin_viewer should be able to access {route}. Got error: {e!s}")
@pytest.mark.parametrize(
@ -2200,7 +2172,7 @@ def test_internal_user_blocked_from_admin_viewer_logs_routes(route):
if route not in INTERNAL_USER_BLOCKED_SUBSET:
return
with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info:
with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info:
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER.value,
@ -2276,9 +2248,7 @@ def test_proxy_admin_viewer_can_access_settings_read_endpoints(route):
request_data={},
)
except Exception as e:
pytest.fail(
f"proxy_admin_viewer should be able to access {route}. Got error: {str(e)}"
)
pytest.fail(f"proxy_admin_viewer should be able to access {route}. Got error: {e!s}")
# ── Admin Viewer parity: default-allow GET semantics ─────────────────────────
@ -2477,9 +2447,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints:
)
local_file = os.path.abspath(local_file)
spec = importlib.util.spec_from_file_location(
"local_enterprise_route_checks", local_file
)
spec = importlib.util.spec_from_file_location("local_enterprise_route_checks", local_file)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
return mod.EnterpriseRouteChecks
@ -2490,9 +2458,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints:
EnterpriseRouteChecks = self._get_enterprise_route_checks()
with (
patch.object(
EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True
),
patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True),
patch.object(
EnterpriseRouteChecks,
"is_management_routes_disabled",
@ -2508,9 +2474,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints:
EnterpriseRouteChecks = self._get_enterprise_route_checks()
with (
patch.object(
EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True
),
patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True),
patch.object(
EnterpriseRouteChecks,
"is_management_routes_disabled",
@ -2526,9 +2490,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints:
EnterpriseRouteChecks = self._get_enterprise_route_checks()
with (
patch.object(
EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True
),
patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True),
patch.object(
EnterpriseRouteChecks,
"is_management_routes_disabled",
@ -2539,9 +2501,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints:
EnterpriseRouteChecks.should_call_route("/v1/chat/completions")
assert exc_info.value.status_code == 403
assert "LLM API routes are disabled for this instance." in str(
exc_info.value.detail
)
assert "LLM API routes are disabled for this instance." in str(exc_info.value.detail)
@patch("litellm.proxy.proxy_server.premium_user", True)
def test_should_embeddings_still_blocked_when_llm_api_disabled(self):
@ -2549,9 +2509,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints:
EnterpriseRouteChecks = self._get_enterprise_route_checks()
with (
patch.object(
EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True
),
patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=True),
patch.object(
EnterpriseRouteChecks,
"is_management_routes_disabled",
@ -2569,9 +2527,7 @@ class TestModelsRouteExemptFromDisableLLMEndpoints:
EnterpriseRouteChecks = self._get_enterprise_route_checks()
with (
patch.object(
EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=False
),
patch.object(EnterpriseRouteChecks, "is_llm_api_route_disabled", return_value=False),
patch.object(
EnterpriseRouteChecks,
"is_management_routes_disabled",
@ -2590,9 +2546,7 @@ def test_route_in_additional_public_routes_wildcard_match():
from litellm.proxy.auth.auth_utils import route_in_additonal_public_routes
with (
patch(
"litellm.proxy.proxy_server.general_settings", {"public_routes": ["/api/*"]}
),
patch("litellm.proxy.proxy_server.general_settings", {"public_routes": ["/api/*"]}),
patch("litellm.proxy.proxy_server.premium_user", True),
):
# Wildcard should match subpaths
@ -2684,7 +2638,7 @@ def test_non_admin_non_team_admin_cannot_access_config_update_but_can_attempt_re
)
# /config/update is still blocked
with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info:
with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info:
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER.value,
@ -2772,8 +2726,6 @@ def test_available_roles_accessible_to_non_admin_users(user_role):
# ── _user_is_org_admin tests ──────────────────────────────────────────────────
def _make_org_admin_user(org_id: str) -> LiteLLM_UserTable:
membership = LiteLLM_OrganizationMembershipTable(
user_id="org-admin-user",
@ -2896,9 +2848,7 @@ async def test_add_team_org_context_noop_when_org_id_already_present():
raise AssertionError("must not resolve when organization_id is present")
body = {"team_id": "team-1", "organization_id": "org-explicit"}
out = await add_team_org_context_to_request_body(
route="/team/update", request_body=body, fetch_team_org_id=fetch
)
out = await add_team_org_context_to_request_body(route="/team/update", request_body=body, fetch_team_org_id=fetch)
assert out == body
@ -2910,9 +2860,7 @@ async def test_add_team_org_context_noop_for_other_routes():
raise AssertionError("must not resolve for a non-opted-in route")
body = {"team_id": "team-1"}
out = await add_team_org_context_to_request_body(
route="/team/delete", request_body=body, fetch_team_org_id=fetch
)
out = await add_team_org_context_to_request_body(route="/team/delete", request_body=body, fetch_team_org_id=fetch)
assert out == body
@ -2925,9 +2873,7 @@ async def test_add_team_org_context_noop_when_team_has_no_org():
return None
body = {"team_id": "team-1"}
out = await add_team_org_context_to_request_body(
route="/team/update", request_body=body, fetch_team_org_id=fetch
)
out = await add_team_org_context_to_request_body(route="/team/update", request_body=body, fetch_team_org_id=fetch)
assert out == body
@ -3198,9 +3144,7 @@ async def test_initialize_pass_through_registers_wildcard_for_auth_subpath():
# Removing the endpoint should clean up openai_routes
# remove_endpoint_routes takes endpoint_id (UUID portion of
# the route key "{id}:exact:{path}:{methods}")
registered = (
InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes()
)
registered = InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes()
endpoint_ids = {k.split(":")[0] for k in registered}
for eid in endpoint_ids:
InitPassThroughEndpointHelpers.remove_endpoint_routes(eid)
@ -3210,9 +3154,7 @@ async def test_initialize_pass_through_registers_wildcard_for_auth_subpath():
LiteLLMRoutes.openai_routes.value[:] = original_routes
# Clean up any routes registered during this test to avoid
# polluting the module-level _registered_pass_through_routes
registered = (
InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes()
)
registered = InitPassThroughEndpointHelpers.get_all_registered_pass_through_routes()
for k in registered:
InitPassThroughEndpointHelpers.remove_endpoint_routes(k.split(":")[0])
@ -3243,8 +3185,7 @@ def test_provider_name_substring_not_classified_as_llm_route(route):
from litellm.proxy.auth.route_checks import RouteChecks
assert RouteChecks.is_llm_api_route(route=route) is False, (
f"{route!r} should NOT be classified as an LLM API route — "
"provider-name substring match bypass"
f"{route!r} should NOT be classified as an LLM API route — provider-name substring match bypass"
)
@ -3266,9 +3207,7 @@ def test_legitimate_passthrough_routes_still_classified_as_llm_route(route):
"""Legitimate passthrough routes must still pass is_llm_api_route."""
from litellm.proxy.auth.route_checks import RouteChecks
assert (
RouteChecks.is_llm_api_route(route=route) is True
), f"{route!r} should be classified as an LLM API route"
assert RouteChecks.is_llm_api_route(route=route) is True, f"{route!r} should be classified as an LLM API route"
@pytest.mark.parametrize(
@ -3326,7 +3265,7 @@ def test_internal_user_blocked_from_search_tool_writes(route):
request = MagicMock(spec=Request)
request.query_params = {}
with pytest.raises(Exception, match='Only proxy admin can be used to generate, delete, update') as exc_info:
with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update") as exc_info:
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER.value,
@ -3702,12 +3641,7 @@ def test_agent_inference_routes_stay_llm_api(route):
def test_agent_routes_union_still_covers_both_halves(route):
"""Keys configured with allowed_routes=["agent_routes"] must keep both halves."""
assert (
RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.agent_routes.value
)
is True
)
assert RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.agent_routes.value) is True
@pytest.mark.parametrize("route", AGENT_MANAGEMENT_ROUTES)
@ -3761,6 +3695,136 @@ def test_agent_registry_route_gate_open_to_non_admin_roles(user_role, method, ro
valid_token=valid_token,
request_data={},
)
def test_proxy_admin_viewer_user_update_password_param_rejected():
"""The self-service /user/update password carve-out is closed: non-admins
change their own password through /user/password/change, which verifies
the current password. Admin password sets don't pass through this check."""
with pytest.raises(HTTPException) as exc_info:
RouteChecks._check_proxy_admin_viewer_access(
route="/user/update",
_user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
request_data={"password": "hunter2hunter2"},
)
assert exc_info.value.status_code == 403
assert "password" in str(exc_info.value.detail)
def test_proxy_admin_viewer_user_update_user_email_still_allowed():
request = MagicMock(spec=Request)
request.method = "POST"
allowed = RouteChecks._check_proxy_admin_viewer_access(
route="/user/update",
_user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
request_data={"user_email": "viewer@example.com"},
request=request,
)
assert allowed is None
def test_proxy_admin_viewer_can_change_own_password():
request = MagicMock(spec=Request)
request.method = "POST"
allowed = RouteChecks._check_proxy_admin_viewer_access(
route="/user/password/change",
_user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
request_data={"current_password": "a", "new_password": "b"},
request=request,
)
assert allowed is None
@pytest.mark.parametrize(
"user_role",
[
LitellmUserRoles.INTERNAL_USER.value,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value,
],
)
def test_non_admin_roles_can_change_own_password(user_role):
valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role)
request = MagicMock(spec=Request)
request.method = "POST"
request.query_params = {}
allowed = RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=LiteLLM_UserTable(user_id="test_user", user_role=user_role),
_user_role=user_role,
route="/user/password/change",
request=request,
valid_token=valid_token,
request_data={"current_password": "a", "new_password": "b"},
)
assert allowed is None
def _password_reset_session_token() -> UserAPIKeyAuth:
"""The UI session key `authenticate_user` mints for a user flagged
`password_reset_required`."""
return UserAPIKeyAuth(
user_id="flagged_user",
allowed_routes=["/user/password/change"],
metadata={"password_reset_required": True},
)
def test_password_reset_session_can_reach_change_password():
result = RouteChecks.is_virtual_key_allowed_to_call_route(
route="/user/password/change",
valid_token=_password_reset_session_token(),
)
assert result is True
@pytest.mark.parametrize(
"route",
[
"/user/info",
"/key/generate",
"/user/update",
"/chat/completions",
],
)
def test_password_reset_session_is_blocked_everywhere_else_with_reset_message(route):
"""Server-side enforcement of the forced reset: a script that logs in via
/v2/login and drives the management API with the session key must get a 403
naming the remediation endpoint, on every route but the change-password one."""
with pytest.raises(HTTPException) as exc_info:
RouteChecks.is_virtual_key_allowed_to_call_route(
route=route,
valid_token=_password_reset_session_token(),
)
assert exc_info.value.status_code == 403
assert "password must be changed" in str(exc_info.value.detail)
assert "/user/password/change" in str(exc_info.value.detail)
def test_restricted_key_without_reset_marker_keeps_generic_message():
"""The reset-specific 403 must not leak onto ordinary allowed_routes keys."""
valid_token = UserAPIKeyAuth(
user_id="test_user",
allowed_routes=["/chat/completions"],
)
with pytest.raises(HTTPException) as exc_info:
RouteChecks.is_virtual_key_allowed_to_call_route(
route="/user/info",
valid_token=valid_token,
)
assert exc_info.value.status_code == 403
assert "password must be changed" not in str(exc_info.value.detail)
assert "not allowed to call this route" in str(exc_info.value.detail)
TEAM_CALLBACK_ROUTES = (
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/callback",
"/team/06bda574-5ca9-43d3-beb8-3b23c2f17112/callback/langfuse",

623
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

File diff suppressed because it is too large Load diff

56
tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
View file

@ -6,7 +6,7 @@ import logging
from contextlib import ExitStack
from datetime import datetime, timedelta
from types import SimpleNamespace
from typing import List, Optional, cast
from typing import Final, List, Optional, cast
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -1614,6 +1614,60 @@ class TestTeamScopedMCPServerAccess:
result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth, team_id="any-team-id")
assert len(result) == 1
class TestFetchAllMCPServersOrdering:
def test_display_order_is_case_insensitive_name_then_id(self) -> None:
servers: Final = (
LiteLLM_MCPServerTable(server_id="s-2", server_name="GitHub", alias="aaa", transport=MCPTransport.http),
LiteLLM_MCPServerTable(server_id="s-1", alias="github", transport=MCPTransport.http),
LiteLLM_MCPServerTable(server_id="s-0", server_name="Slack", alias="zzz", transport=MCPTransport.http),
LiteLLM_MCPServerTable(server_id="confluence", server_name="", alias="", transport=MCPTransport.http),
)
ordered: Final = sorted(servers, key=mgmt_endpoints._mcp_server_display_order)
assert [s.server_id for s in ordered] == ["confluence", "s-1", "s-2", "s-0"]
@pytest.mark.parametrize("team_id", [None, "team-1"])
@pytest.mark.parametrize("reverse", [False, True])
@pytest.mark.asyncio
async def test_list_is_sorted_by_display_name_regardless_of_resolution_order(
self, team_id: str | None, reverse: bool
) -> None:
mock_user_auth: Final = generate_mock_user_api_key_auth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin_user",
)
servers: Final = (
generate_mock_mcp_server_db_record(server_id="s-zeta", alias="zeta"),
generate_mock_mcp_server_db_record(server_id="s-alpha", alias="Alpha"),
generate_mock_mcp_server_db_record(server_id="s-mid", alias="mid"),
)
resolved: Final = list(reversed(servers) if reverse else servers)
mock_manager: Final = MagicMock()
mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=resolved)
with (
patch( # test-quality-ok: the route reads a module-global manager with no injection seam
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch( # test-quality-ok: admin view is derived from module-global proxy settings
"litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
return_value=True,
),
patch( # test-quality-ok: auth contexts need a live prisma client
"litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
AsyncMock(return_value=[mock_user_auth]),
),
patch( # test-quality-ok: isolate the route's ordering from team database resolution
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_scoped_mcp_server_list",
AsyncMock(return_value=resolved),
),
):
result: Final = await mgmt_endpoints.fetch_all_mcp_servers(
user_api_key_dict=mock_user_auth, team_id=team_id
)
assert [s.server_id for s in result] == ["s-alpha", "s-mid", "s-zeta"]
@pytest.mark.asyncio
async def test_restricted_virtual_key_cannot_use_team_id_filter(self):
"""Restricted virtual keys must not bypass access limits via team_id."""

401
tests/test_litellm/proxy/management_endpoints/test_password_endpoints.py Normal file
View file

@ -0,0 +1,401 @@
"""
Tests for POST /user/password/change (litellm/proxy/management_endpoints/password_endpoints.py).
HIBP traffic is intercepted with respx; no test here touches the network.
"""
import hashlib
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
from fastapi import HTTPException
from litellm.proxy._types import UI_TEAM_ID, LitellmTableNames, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.login_utils import PASSWORD_SESSION_METADATA
from litellm.proxy.management_endpoints.password_endpoints import change_password
from litellm.proxy.utils import hash_password, verify_password
CURRENT_PASSWORD = "OldP@ssw0rd-2026"
NEW_PASSWORD = "NewP@ssw0rd-2026"
_POLICY_NO_BREACH_CHECK = {"password_policy_check_breached_passwords": False}
def _make_user_row(password: str | None) -> MagicMock:
user = MagicMock()
user.user_id = "user-123"
user.password = password
return user
def _make_prisma(user: MagicMock | None) -> MagicMock:
prisma = MagicMock()
prisma.db.litellm_usertable.find_first = AsyncMock(return_value=user)
prisma.db.litellm_usertable.update = AsyncMock(return_value=user)
return prisma
def _caller(user_id: str | None = "user-123") -> UserAPIKeyAuth:
return UserAPIKeyAuth(user_id=user_id, team_id=UI_TEAM_ID, metadata=dict(PASSWORD_SESSION_METADATA))
def _sso_session_caller() -> UserAPIKeyAuth:
return UserAPIKeyAuth(user_id="user-123", team_id=UI_TEAM_ID, metadata={})
def _virtual_key_caller() -> UserAPIKeyAuth:
return UserAPIKeyAuth(user_id="user-123", team_id="team-abc", metadata=dict(PASSWORD_SESSION_METADATA))
def _hibp_url_for(password: str) -> str:
sha1 = hashlib.sha1(password.encode(), usedforsecurity=False).hexdigest().upper()
return f"https://api.pwnedpasswords.com/range/{sha1[:5]}"
def _hibp_suffix_for(password: str) -> str:
return hashlib.sha1(password.encode(), usedforsecurity=False).hexdigest().upper()[5:]
@pytest.mark.asyncio
async def test_change_password_success_writes_new_scrypt_hash():
from litellm.proxy._types import ChangePasswordRequest
prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD)))
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
):
response = await change_password(
data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD),
user_api_key_dict=_caller(),
)
assert response.user_id == "user-123"
update_kwargs = prisma.db.litellm_usertable.update.call_args.kwargs
assert update_kwargs["where"] == {"user_id": "user-123"}
stored = update_kwargs["data"]["password"]
assert stored != NEW_PASSWORD
assert verify_password(NEW_PASSWORD, stored)
# A successful change lifts any pending forced reset and re-arms the
# login-time breach screen for the new password.
assert update_kwargs["data"]["password_reset_required"] is False
assert update_kwargs["data"]["last_breach_check_at"] is None
@pytest.mark.asyncio
async def test_change_password_rejects_wrong_current_password():
from litellm.proxy._types import ChangePasswordRequest
prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD)))
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
):
with pytest.raises(HTTPException) as exc_info:
await change_password(
data=ChangePasswordRequest(current_password="not-the-password", new_password=NEW_PASSWORD),
user_api_key_dict=_caller(),
)
assert exc_info.value.status_code == 400
assert "Current password is incorrect" in exc_info.value.detail["error"]
prisma.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
async def test_change_password_rejects_unchanged_password():
from litellm.proxy._types import ChangePasswordRequest
prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD)))
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
):
with pytest.raises(HTTPException) as exc_info:
await change_password(
data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=CURRENT_PASSWORD),
user_api_key_dict=_caller(),
)
assert exc_info.value.status_code == 400
assert "must be different from the current password" in exc_info.value.detail["error"]
prisma.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"caller",
[
pytest.param(_sso_session_caller(), id="sso_dashboard_session"),
pytest.param(_virtual_key_caller(), id="virtual_key_with_forged_metadata"),
],
)
async def test_change_password_rejects_non_password_login_session(caller: UserAPIKeyAuth):
"""Only the session minted by a password login may change the password, so a
stolen virtual key or an SSO session cannot use the endpoint as a
current_password guessing oracle."""
from litellm.proxy._types import ChangePasswordRequest
prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD)))
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
):
with pytest.raises(HTTPException) as exc_info:
await change_password(
data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD),
user_api_key_dict=caller,
)
assert exc_info.value.status_code == 403
assert "logging in with a password" in exc_info.value.detail["error"]
prisma.db.litellm_usertable.find_first.assert_not_called()
prisma.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
async def test_change_password_rejects_session_without_user():
from litellm.proxy._types import ChangePasswordRequest
prisma = _make_prisma(user=None)
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
):
with pytest.raises(HTTPException) as exc_info:
await change_password(
data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD),
user_api_key_dict=_caller(user_id=None),
)
assert exc_info.value.status_code == 400
prisma.db.litellm_usertable.find_first.assert_not_called()
prisma.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
async def test_change_password_rejects_account_without_password():
"""SSO users and the env-credential admin have no DB password row to change."""
from litellm.proxy._types import ChangePasswordRequest
prisma = _make_prisma(_make_user_row(password=None))
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
):
with pytest.raises(HTTPException) as exc_info:
await change_password(
data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD),
user_api_key_dict=_caller(),
)
assert exc_info.value.status_code == 400
assert "no password set" in exc_info.value.detail["error"]
prisma.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
async def test_change_password_enforces_min_length():
from litellm.proxy._types import ChangePasswordRequest
prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD)))
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
):
with pytest.raises(ProxyException) as exc_info:
await change_password(
data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password="Short1!"),
user_api_key_dict=_caller(),
)
assert exc_info.value.code == "400"
assert exc_info.value.type == ProxyErrorTypes.validation_error
assert exc_info.value.param == "password"
assert "at least 12 characters" in exc_info.value.message
prisma.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
@respx.mock
async def test_change_password_rejects_breached_password():
"""With the default policy, the new password is screened against HIBP."""
from litellm.proxy._types import ChangePasswordRequest
breached_password = "Password123!"
respx.get(_hibp_url_for(breached_password)).mock(
return_value=httpx.Response(200, text=f"{_hibp_suffix_for(breached_password)}:1")
)
prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD)))
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", {}
),
):
with pytest.raises(ProxyException) as exc_info:
await change_password(
data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=breached_password),
user_api_key_dict=_caller(),
)
assert exc_info.value.code == "400"
assert exc_info.value.type == ProxyErrorTypes.validation_error
assert exc_info.value.param == "password"
assert "data breaches" in exc_info.value.message
prisma.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
@respx.mock
async def test_change_password_verifies_current_password_before_hibp_lookup():
"""A caller who fails current-password verification must not trigger any
HIBP traffic. The HIBP check fails open on errors, so an unmocked lookup
could not prove ordering; instead the route is registered and asserted
uncalled."""
from litellm.proxy._types import ChangePasswordRequest
hibp_route = respx.get(_hibp_url_for(NEW_PASSWORD)).mock(return_value=httpx.Response(200, text=""))
prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD)))
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", {}
),
):
with pytest.raises(HTTPException) as exc_info:
await change_password(
data=ChangePasswordRequest(current_password="not-the-password", new_password=NEW_PASSWORD),
user_api_key_dict=_caller(),
)
assert exc_info.value.status_code == 400
assert "Current password is incorrect" in exc_info.value.detail["error"]
assert not hibp_route.called
prisma.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
async def test_change_password_success_emits_redacted_audit_log():
"""A successful change must land in the audit trail as field names only;
the plaintext passwords must never reach the audit call."""
from litellm.proxy._types import ChangePasswordRequest
prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD)))
audit_mock = AsyncMock()
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
patch( # test-quality-ok: audit sink is a module-level import; no injection seam
"litellm.proxy.management_endpoints.password_endpoints.create_object_audit_log", audit_mock
),
):
await change_password(
data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD),
user_api_key_dict=_caller(),
)
audit_mock.assert_awaited_once()
audit_kwargs = audit_mock.await_args.kwargs
assert audit_kwargs["object_id"] == "user-123"
assert audit_kwargs["action"] == "updated"
assert audit_kwargs["table_name"] == LitellmTableNames.USER_TABLE_NAME
assert audit_kwargs["after_value"] == '{"fields_changed": ["password"]}'
assert CURRENT_PASSWORD not in str(audit_kwargs)
assert NEW_PASSWORD not in str(audit_kwargs)
@pytest.mark.asyncio
async def test_change_password_failure_emits_no_audit_log():
from litellm.proxy._types import ChangePasswordRequest
prisma = _make_prisma(_make_user_row(hash_password(CURRENT_PASSWORD)))
audit_mock = AsyncMock()
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", prisma
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
patch( # test-quality-ok: audit sink is a module-level import; no injection seam
"litellm.proxy.management_endpoints.password_endpoints.create_object_audit_log", audit_mock
),
):
with pytest.raises(HTTPException):
await change_password(
data=ChangePasswordRequest(current_password="not-the-password", new_password=NEW_PASSWORD),
user_api_key_dict=_caller(),
)
audit_mock.assert_not_awaited()
@pytest.mark.asyncio
async def test_change_password_requires_db():
from litellm.proxy._types import ChangePasswordRequest
with (
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.prisma_client", None
),
patch( # test-quality-ok: change_password reads proxy_server module globals; no injection seam
"litellm.proxy.proxy_server.general_settings", _POLICY_NO_BREACH_CHECK
),
):
with pytest.raises(HTTPException) as exc_info:
await change_password(
data=ChangePasswordRequest(current_password=CURRENT_PASSWORD, new_password=NEW_PASSWORD),
user_api_key_dict=_caller(),
)
assert exc_info.value.status_code == 500

3
tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py
View file

@ -10,6 +10,7 @@ Routes covered:
from __future__ import annotations
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock
from .conftest import normalize
@ -510,6 +511,8 @@ def _db_user(monkeypatch, email: str):
user.user_email = email
user.user_role = "internal_user"
user.password = "scrypt:stored"
user.password_reset_required = None
user.last_breach_check_at = datetime.now(timezone.utc)
repo = MagicMock()
repo.return_value.table.find_first = AsyncMock(return_value=user)
monkeypatch.setattr(ps, "prisma_client", MagicMock())

42
tests/test_litellm/proxy/test__types.py
View file

@ -5,11 +5,13 @@ from pydantic import ValidationError
from litellm.proxy._types import (
ROLES_WITHIN_ORG,
ChangePasswordRequest,
GenerateKeyRequest,
KeyRequest,
LiteLLM_AuditLogs,
LiteLLM_TeamMembership,
LitellmUserRoles,
NewUserRequest,
OrganizationMemberUpdateRequest,
ResetSpendRequest,
UpdateKeyRequest,
@ -337,3 +339,43 @@ def test_virtual_key_mapping_counts_as_configured_when_any_issuer_sets_the_claim
)
assert jwt_auth.is_virtual_key_mapping_configured() is is_configured
def test_new_user_request_loudly_rejects_a_password():
"""
/user/new has never persisted a password (the field used to be silently
dropped). Sending one must now fail visibly so the dead path cannot be
revived without going through the password policy.
"""
with pytest.raises(ValidationError, match="invitation link"):
NewUserRequest(user_email="alice@example.com", password="hunter2hunter2")
def test_new_user_request_without_password_still_works():
request = NewUserRequest(user_email="alice@example.com")
assert request.password is None
def test_update_user_request_accepts_a_password():
"""Admins set user passwords through /user/update; the value must survive
model validation so the endpoint can policy-check and hash it."""
request = UpdateUserRequest(user_id="user-123", password="hunter2hunter2")
assert request.password == "hunter2hunter2"
def test_update_user_request_password_hidden_from_repr():
"""management_endpoint_wrapper string-formats endpoint kwargs into Slack
alerts, so the model's repr/str must never contain the plaintext password."""
request = UpdateUserRequest(user_id="user-123", password="hunter2hunter2")
assert "hunter2hunter2" not in repr(request)
assert "hunter2hunter2" not in str(request)
def test_change_password_request_passwords_hidden_from_repr():
"""Any accidental str()/repr() of the request model (debug logs, exception
handlers, a future management_endpoint_wrapper) must never contain either
plaintext password."""
request = ChangePasswordRequest(current_password="hunter2hunter2", new_password="NewP@ssw0rd-2026")
for rendered in (repr(request), str(request)):
assert "hunter2hunter2" not in rendered
assert "NewP@ssw0rd-2026" not in rendered

60
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -11766,6 +11766,66 @@ def test_prompt_caching_settings_propagate_on_config_reload(monkeypatch, field_n
assert getattr(litellm, field_name) == db_value
@pytest.mark.asyncio
async def test_db_stored_datadog_redaction_settings_apply_before_logger_init(monkeypatch: pytest.MonkeyPatch):
"""A DB-only litellm_settings row that pairs success_callback: ["datadog"] with
datadog_params.turn_off_message_logging: true must build the DataDogLogger redacted, the
same as the identical block in YAML. Regression for the redaction keys being absent from
the safe-override allowlist while the callback half of the row was honoured."""
import litellm.proxy.proxy_server as ps
from litellm.integrations.datadog.datadog import DataDogLogger
from litellm.litellm_core_utils import litellm_logging
monkeypatch.setenv("DD_API_KEY", "test-key")
monkeypatch.setenv("DD_SITE", "us5.datadoghq.com")
monkeypatch.setattr(litellm, "datadog_params", None)
monkeypatch.setattr(litellm, "turn_off_message_logging", False)
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "_async_success_callback", [])
monkeypatch.setattr(litellm, "failure_callback", [])
monkeypatch.setattr(litellm, "_async_failure_callback", [])
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [])
db_row = {
"success_callback": ["datadog"],
"datadog_params": {"turn_off_message_logging": True},
"turn_off_message_logging": True,
}
pc = ps.ProxyConfig()
pc._apply_litellm_settings_db_values(pc._prepared_db_settings_values("litellm_settings", db_row))
pc._add_callbacks_from_db_config({"litellm_settings": db_row})
datadog_loggers = [cb for cb in litellm.success_callback if isinstance(cb, DataDogLogger)]
assert len(datadog_loggers) == 1
assert datadog_loggers[0].turn_off_message_logging is True
assert litellm.turn_off_message_logging is True
@pytest.mark.parametrize(
"field_name",
[
"datadog_params",
"datadog_llm_observability_params",
"newrelic_params",
"pointfive_params",
"aws_sqs_callback_params",
],
)
def test_db_stored_callback_params_propagate_to_litellm_module(monkeypatch: pytest.MonkeyPatch, field_name: str):
"""Every callback init params block stored in the DB litellm_settings row must land on the
litellm module before the matching logger is built, so the DB row behaves like YAML."""
import litellm.proxy.proxy_server as ps
monkeypatch.setattr(litellm, field_name, None)
db_value = {"turn_off_message_logging": True}
pc = ps.ProxyConfig()
pc._apply_litellm_settings_db_values(pc._prepared_db_settings_values("litellm_settings", {field_name: db_value}))
assert getattr(litellm, field_name) == db_value
def test_get_config_list_marks_untouched_prompt_caching_flag_as_not_set(monkeypatch):
"""The flag defaults to False rather than None, so a plain 'is not None' check would
report the default as 'In Config' and imply an admin had set it."""

85
tests/test_litellm/secret_managers/hashicorp_vault_parity.json Normal file
View file

@ -0,0 +1,85 @@
[
{
"name": "defaults",
"env": {
"HCP_VAULT_TOKEN": "token"
},
"secret_name": "OPENAI_API_KEY",
"expected_secret_url": "http://127.0.0.1:8200/v1/secret/data/OPENAI_API_KEY",
"expected_login_url": null,
"expected_login_namespace": null,
"expected_secret_namespace": null
},
{
"name": "global_namespace",
"env": {
"HCP_VAULT_ADDR": "http://vault.test:8200",
"HCP_VAULT_TOKEN": "token",
"HCP_VAULT_NAMESPACE": "admin"
},
"secret_name": "OPENAI_API_KEY",
"expected_secret_url": "http://vault.test:8200/v1/admin/secret/data/OPENAI_API_KEY",
"expected_login_url": null,
"expected_login_namespace": "admin",
"expected_secret_namespace": "admin"
},
{
"name": "namespace_overrides",
"env": {
"HCP_VAULT_ADDR": "http://vault.test:8200",
"HCP_VAULT_TOKEN": "token",
"HCP_VAULT_NAMESPACE": "admin",
"HCP_VAULT_LOGIN_NAMESPACE": "root",
"HCP_VAULT_SECRET_NAMESPACE": "teams/team-a"
},
"secret_name": "OPENAI_API_KEY",
"expected_secret_url": "http://vault.test:8200/v1/teams/team-a/secret/data/OPENAI_API_KEY",
"expected_login_url": null,
"expected_login_namespace": "root",
"expected_secret_namespace": "teams/team-a"
},
{
"name": "custom_mount_and_prefix",
"env": {
"HCP_VAULT_ADDR": "http://vault.test:8200",
"HCP_VAULT_TOKEN": "token",
"HCP_VAULT_MOUNT_NAME": " /kv-prod/ ",
"HCP_VAULT_PATH_PREFIX": " /virtual-keys/ "
},
"secret_name": "DB_PASSWORD",
"expected_secret_url": "http://vault.test:8200/v1/kv-prod/data/virtual-keys/DB_PASSWORD",
"expected_login_url": null,
"expected_login_namespace": null,
"expected_secret_namespace": null
},
{
"name": "approle_custom_mount",
"env": {
"HCP_VAULT_ADDR": "http://vault.test:8200",
"HCP_VAULT_APPROLE_ROLE_ID": "role-id",
"HCP_VAULT_APPROLE_SECRET_ID": "secret-id",
"HCP_VAULT_APPROLE_MOUNT_PATH": "custom-approle",
"HCP_VAULT_NAMESPACE": "admin"
},
"secret_name": "OPENAI_API_KEY",
"expected_secret_url": "http://vault.test:8200/v1/admin/secret/data/OPENAI_API_KEY",
"expected_login_url": "http://vault.test:8200/v1/auth/custom-approle/login",
"expected_login_namespace": "admin",
"expected_secret_namespace": "admin"
},
{
"name": "tls_cert",
"env": {
"HCP_VAULT_ADDR": "http://vault.test:8200",
"HCP_VAULT_CLIENT_CERT": "/tmp/client.crt",
"HCP_VAULT_CLIENT_KEY": "/tmp/client.key",
"HCP_VAULT_CERT_ROLE": "vault-role",
"HCP_VAULT_NAMESPACE": "admin"
},
"secret_name": "OPENAI_API_KEY",
"expected_secret_url": "http://vault.test:8200/v1/admin/secret/data/OPENAI_API_KEY",
"expected_login_url": "http://vault.test:8200/v1/auth/cert/login",
"expected_login_namespace": "admin",
"expected_secret_namespace": "admin"
}
]

50
tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py
View file

@ -1,4 +1,5 @@
import datetime
import json
from collections.abc import Mapping
from pathlib import Path
from typing import Final
@ -18,6 +19,23 @@ LOGIN_RESPONSE: Final = {"auth": {"client_token": "hvs.login-token", "lease_dura
SECRET_RESPONSE: Final = {"data": {"data": {"key": "sk-from-vault", "password": "pw-from-vault"}}}
NAMESPACE_ENV_VARS: Final = ("HCP_VAULT_NAMESPACE", "HCP_VAULT_LOGIN_NAMESPACE", "HCP_VAULT_SECRET_NAMESPACE")
PARITY_ENV_VARS: Final = (
"HCP_VAULT_ADDR",
"HCP_VAULT_TOKEN",
"HCP_VAULT_NAMESPACE",
"HCP_VAULT_LOGIN_NAMESPACE",
"HCP_VAULT_SECRET_NAMESPACE",
"HCP_VAULT_MOUNT_NAME",
"HCP_VAULT_PATH_PREFIX",
"HCP_VAULT_APPROLE_ROLE_ID",
"HCP_VAULT_APPROLE_SECRET_ID",
"HCP_VAULT_APPROLE_MOUNT_PATH",
"HCP_VAULT_CLIENT_CERT",
"HCP_VAULT_CLIENT_KEY",
"HCP_VAULT_CERT_ROLE",
"HCP_VAULT_REFRESH_INTERVAL",
"SECRET_MANAGER_REFRESH_INTERVAL",
)
def _build_manager(monkeypatch: pytest.MonkeyPatch, env: Mapping[str, str]) -> HashicorpSecretManager:
@ -236,3 +254,35 @@ def test_tls_login_uses_login_namespace(monkeypatch: pytest.MonkeyPatch, tmp_pat
assert manager._auth_via_tls_cert() == "hvs.login-token"
assert login_route.calls.last.request.headers["X-Vault-Namespace"] == "root"
with Path(__file__).with_name("hashicorp_vault_parity.json").open() as parity_file:
PARITY_CASES: Final = json.load(parity_file)
@pytest.mark.parametrize("case", PARITY_CASES, ids=lambda case: case["name"])
def test_configuration_matches_native_parity_fixture(
monkeypatch: pytest.MonkeyPatch, case: Mapping[str, object]
) -> None:
monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True)
for name in PARITY_ENV_VARS:
monkeypatch.delenv(name, raising=False)
for name, value in case["env"].items():
monkeypatch.setenv(name, value)
manager: Final = HashicorpSecretManager()
env: Final = case["env"]
expected_login_url: Final = case["expected_login_url"]
if env.get("HCP_VAULT_APPROLE_ROLE_ID") and env.get("HCP_VAULT_APPROLE_SECRET_ID"):
login_url: str | None = (
f"{manager.vault_addr}/v1/auth/{manager.approle_mount_path}/login"
)
elif env.get("HCP_VAULT_CLIENT_CERT") and env.get("HCP_VAULT_CLIENT_KEY"):
login_url = f"{manager.vault_addr}/v1/auth/cert/login"
else:
login_url = None
assert manager.get_url(case["secret_name"]) == case["expected_secret_url"]
assert manager.vault_login_namespace == case["expected_login_namespace"]
assert manager.vault_secret_namespace == case["expected_secret_namespace"]
assert login_url == expected_login_url

152
tests/test_litellm_rust/support/fake_gcs.py Normal file
View file

@ -0,0 +1,152 @@
from __future__ import annotations
import json
import threading
from collections.abc import Mapping
from dataclasses import dataclass
from functools import partial
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from socket import socket
from types import MappingProxyType
from typing import Final, cast
from urllib.parse import unquote, urlsplit
@dataclass(frozen=True, slots=True)
class RecordedRequest:
method: str
path: str
query: str
headers: Mapping[str, str]
body: bytes
class _FakeGcsHandler(BaseHTTPRequestHandler):
def __init__(
self,
request: socket | tuple[bytes, socket],
client_address: tuple[str, int],
server: ThreadingHTTPServer,
*,
fake: FakeGcs,
) -> None:
self._fake: Final = fake
super().__init__(request, client_address, server)
def _handle(self) -> None:
parsed: Final = urlsplit(self.path)
content_length: Final = int(self.headers.get("Content-Length", "0"))
body: Final = self.rfile.read(content_length) if content_length else b""
headers: Final = MappingProxyType(
{name.title(): value for name, value in self.headers.items()}
)
self._fake.record(
RecordedRequest(
method=self.command,
path=parsed.path,
query=parsed.query,
headers=headers,
body=body,
)
)
if self.headers.get("Authorization") != f"Bearer {self._fake.token}":
self._send_json(401, {"error": "unauthorized"})
return
upload_prefix: Final = "/upload/storage/v1/b/"
download_prefix: Final = "/storage/v1/b/"
if parsed.path.startswith(upload_prefix) and parsed.path.endswith("/o"):
self._upload(parsed.path[len(upload_prefix) : -2], parsed.query, body)
return
if parsed.path.startswith(download_prefix):
self._download(parsed.path[len(download_prefix) :], parsed.query)
return
self._send_json(404, {"error": "not found"})
def _upload(self, path: str, query: str, body: bytes) -> None:
values: Final = {
unquote(pair.partition("=")[0]): unquote(pair.partition("=")[2])
for pair in query.split("&")
if pair
}
if not path or values.get("uploadType") != "media" or "name" not in values:
self._send_json(404, {"error": "not found"})
return
self._fake.put_object(path, values["name"], body)
self._send_json(200, {"name": values["name"], "bucket": path})
def _download(self, path: str, query: str) -> None:
bucket, separator, encoded_name = path.partition("/o/")
if not separator or query != "alt=media":
self._send_json(404, {"error": "not found"})
return
name: Final = unquote(encoded_name)
if name.endswith("/server-error") or name == "server-error":
self._send_json(500, {"error": "server error"})
return
body: Final = self._fake.get_object(bucket, name)
if body is None:
self._send_json(404, {"error": "not found"})
return
self._send(200, body, "application/octet-stream")
def _send_json(self, status: int, value: object) -> None:
payload: Final = json.dumps(value).encode()
self._send(status, payload, "application/json")
def _send(self, status: int, body: bytes, content_type: str) -> None:
self.send_response(status)
self.send_header("Content-Type", content_type)
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def log_message(self, format: str, *args: object) -> None:
pass
do_GET = _handle
do_POST = _handle
class FakeGcs:
def __init__(self) -> None:
self._objects: dict[tuple[str, str], bytes] = {} # mutable-ok: fake object store
self._requests: list[RecordedRequest] = [] # mutable-ok: recorded request history
self._server = ThreadingHTTPServer(
("127.0.0.1", 0),
partial(_FakeGcsHandler, fake=self),
)
self._worker = threading.Thread(target=self._server.serve_forever, daemon=True)
self._worker.start()
self.token: Final = "test-token"
@property
def url(self) -> str:
address: Final = cast(tuple[str, int], self._server.server_address)
host, port = address
return f"http://{host}:{port}"
@property
def objects(self) -> Mapping[tuple[str, str], bytes]:
return MappingProxyType(self._objects)
@property
def requests(self) -> tuple[RecordedRequest, ...]:
return tuple(self._requests)
def put(self, bucket: str, name: str, body: bytes) -> None:
self.put_object(bucket, name, body)
def close(self) -> None:
self._server.shutdown()
self._server.server_close()
self._worker.join(timeout=5)
def record(self, request: RecordedRequest) -> None:
self._requests.append(request)
def put_object(self, bucket: str, name: str, body: bytes) -> None:
self._objects[(bucket, name)] = body
def get_object(self, bucket: str, name: str) -> bytes | None:
return self._objects.get((bucket, name))

112
tests/test_litellm_rust/support/s3_stub.py Normal file
View file

@ -0,0 +1,112 @@
"""In-process path-style S3 stub for native cache parity tests."""
import threading
from dataclasses import dataclass, field
from email.utils import parsedate_to_datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final
from urllib.parse import unquote, urlsplit
_STORED_HEADERS: Final = (
"cache-control",
"content-type",
"content-language",
"content-disposition",
"expires",
)
@dataclass
class S3Object:
body: bytes
headers: dict[str, str] = field(default_factory=dict)
class S3Stub:
"""Minimal path-style S3 endpoint serving PUT and GET object operations."""
def __init__(self) -> None:
self._objects: dict[str, S3Object] = {}
stub: Final = self
class Handler(BaseHTTPRequestHandler):
def _key(self) -> str:
parts: Final = urlsplit(self.path).path.lstrip("/").split("/", 1)
return unquote(parts[1]) if len(parts) == 2 else ""
def _read_body(self) -> bytes:
transfer: Final = self.headers.get("transfer-encoding", "")
if "chunked" not in transfer:
return self.rfile.read(int(self.headers.get("content-length", 0)))
chunks: Final = bytearray()
while True:
size = int(self.rfile.readline().split(b";")[0].strip(), 16)
if size == 0:
while self.rfile.readline().strip():
pass
return bytes(chunks)
chunks.extend(self.rfile.read(size))
self.rfile.readline()
def do_PUT(self) -> None:
body: Final = self._read_body()
headers: Final = {name: self.headers[name] for name in _STORED_HEADERS if name in self.headers}
stub._objects = {**stub._objects, self._key(): S3Object(body=body, headers=headers)}
self.send_response(200)
self.send_header("ETag", '"stub"')
self.send_header("Content-Length", "0")
self.end_headers()
def do_HEAD(self) -> None:
self._object(send_body=False)
def do_GET(self) -> None:
self._object(send_body=True)
def _object(self, send_body: bool) -> None:
entry: Final = stub._objects.get(self._key())
if entry is None:
self.send_response(404)
self.send_header("Content-Type", "application/xml")
body: Final = b'<?xml version="1.0" encoding="UTF-8"?><Error><Code>NoSuchKey</Code></Error>'
self.send_header("Content-Length", str(len(body)))
self.end_headers()
if send_body:
self.wfile.write(body)
return
self.send_response(200)
for name, value in entry.headers.items():
self.send_header(name, value)
self.send_header("ETag", '"stub"')
self.send_header("Content-Length", str(len(entry.body)))
self.end_headers()
if send_body:
self.wfile.write(entry.body)
def log_message(self, format: str, *args: object) -> None:
pass
self._server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
self._worker: Final = threading.Thread(target=self._server.serve_forever, daemon=True)
self._worker.start()
@property
def url(self) -> str:
host, port = self._server.server_address[:2]
return f"http://{host}:{port}"
@property
def objects(self) -> dict[str, S3Object]:
return self._objects
def put_object(self, key: str, body: bytes, headers: dict[str, str] | None = None) -> None:
self._objects = {**self._objects, key: S3Object(body=body, headers=headers or {})}
def expires(self, key: str) -> object:
header: Final = self._objects[key].headers.get("expires")
return parsedate_to_datetime(header) if header else None
def close(self) -> None:
self._server.shutdown()
self._server.server_close()
self._worker.join(timeout=5)

544
tests/test_litellm_rust/test_cache.py
View file

@ -8,10 +8,16 @@ import time
import uuid
import weakref
from collections.abc import Generator
from datetime import datetime
from pathlib import Path
from types import SimpleNamespace
from typing import Final, Protocol, cast
from unittest.mock import Mock
from urllib.parse import urlparse
import boto3
import botocore.config
import diskcache
import fakeredis
import pytest
import redis
@ -20,17 +26,23 @@ from azure.storage.blob import ContainerClient
import litellm
from litellm.caching.azure_blob_cache import AzureBlobCache
from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache
from litellm.caching.disk_cache import DiskCache
from litellm.caching.gcs_cache import GCSCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.caching.s3_cache import S3Cache
from litellm.rust_bridge import _native
from litellm.types.caching import LiteLLMCacheType
from tests.test_litellm_rust.support.fake_gcs import FakeGcs
from tests.test_litellm_rust.support.isolation import rebound
from tests.test_litellm_rust.support.s3_stub import S3Stub
pytestmark: Final = pytest.mark.requires_rust_extension
class CacheLookup(Protocol):
def get_cache(self, **kwargs: object) -> object: ...
def flush_cache(self) -> object: ...
def request(key: str = "key") -> dict[str, object]:
@ -50,6 +62,15 @@ def redis_url() -> Generator[str]:
worker.join(timeout=5)
@pytest.fixture
def fake_gcs() -> Generator[FakeGcs]:
server: Final = FakeGcs()
try:
yield server
finally:
server.close()
@pytest.fixture
def azure_blob_facade() -> Generator[Cache]:
account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL")
@ -411,15 +432,20 @@ def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure
assert handle.backend == "azure-blob"
account_url: Final = backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}")
with pytest.raises(TypeError, match="containers must match"):
_native._CacheTestHandle.azure_blob(account_url, f"{backend.container_client.container_name}-other")._bind_facade(
azure_blob_facade
)
_native._CacheTestHandle.azure_blob(
account_url, f"{backend.container_client.container_name}-other"
)._bind_facade(azure_blob_facade)
handle._bind_facade(azure_blob_facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=azure_blob_facade))
native: Final = resolver.resolve()
assert native.kind == "native"
response: Final = {"choices": [{"text": "caf\u00e9 \u2603"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None}
response: Final = {
"choices": [{"text": "caf\u00e9 \u2603"}],
"usage": {"total_tokens": 3},
"flag": True,
"empty": None,
}
native.store({**request("sync"), "ttl_seconds": 0.001}, response)
native.store(request("sync"), {"choices": [{"text": "second"}]})
time.sleep(0.01)
@ -474,7 +500,9 @@ async def test_azure_blob_native_async_writes_overwrite_batch_and_flush_like_pyt
await binding.async_store({**request("async"), "ttl_seconds": 0.001}, {"value": 2})
time.sleep(0.01)
assert await binding.async_lookup(request("async")) == {"value": 2}
assert await backend.async_get_cache("async") == json.loads(backend.container_client.download_blob("async").readall())
assert await backend.async_get_cache("async") == json.loads(
backend.container_client.download_blob("async").readall()
)
assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="async") == {"value": 2}
await binding.async_store_batch([request("first"), request("second")], [{"value": 3}, {"value": 4}])
@ -521,6 +549,508 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None:
client.close()
async def test_disk_reads_python_entries_and_python_reads_native_entries(tmp_path: Path) -> None:
disk_cache: Final = DiskCache(disk_cache_dir=str(tmp_path))
response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}}
disk_cache.disk_cache.set(
"sync",
{"timestamp": time.time(), "response": json.dumps(response)},
)
disk_cache.disk_cache.set("async", json.dumps({"timestamp": time.time(), "response": response}))
disk_cache.disk_cache.set("raw", json.dumps(response))
disk_cache.disk_cache.set("invalid", "not a cache entry")
disk_cache.disk_cache.set(
"large",
{"timestamp": time.time(), "response": {"text": "x" * 70_000}},
)
binding: Final = _native._CacheTestResolver(
SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path)))
).resolve()
assert binding.lookup(request("sync")) == response
assert await binding.async_lookup(request("async")) == response
assert binding.lookup(request("raw")) == response
assert await binding.async_lookup(request("invalid")) is None
assert binding.lookup(request("large")) == {"text": "x" * 70_000}
await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response)
stored_response: Final = disk_cache.get_cache("native")
assert isinstance(stored_response, dict)
assert stored_response["response"] == response
stored, expire_time = disk_cache.disk_cache.get("native", expire_time=True)
assert stored is not None
assert time.time() < expire_time <= time.time() + 12.0
await binding.async_store(request("no-ttl"), response)
_, no_expiry = disk_cache.disk_cache.get("no-ttl", expire_time=True)
assert no_expiry is None
async def test_disk_entries_survive_a_fresh_handle_and_expire_on_time(tmp_path: Path) -> None:
first: Final = _native._CacheTestResolver(
SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path)))
).resolve()
await first.async_store(request("persistent"), {"value": "persistent"})
await first.async_store({**request("expiring"), "ttl_seconds": 0.3}, {"value": "expiring"})
fresh: Final = _native._CacheTestResolver(
SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path)))
).resolve()
assert fresh.lookup(request("persistent")) == {"value": "persistent"}
assert fresh.lookup(request("expiring")) == {"value": "expiring"}
await asyncio.sleep(0.4)
assert fresh.lookup(request("expiring")) is None
assert fresh.lookup(request("persistent")) == {"value": "persistent"}
def test_disk_facade_registers_and_store_changes_fall_back(tmp_path: Path) -> None:
facade: Final = Cache(type=LiteLLMCacheType.DISK, disk_cache_dir=str(tmp_path))
with pytest.raises(TypeError, match="directories must match"):
_native._CacheTestHandle.disk(str(tmp_path / "other"))._bind_facade(facade)
handle: Final = _native._CacheTestHandle.disk(str(tmp_path))
handle._bind_facade(facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
binding: Final = resolver.resolve()
assert binding.kind == "native"
binding.store(request("native"), {"value": "native"})
assert facade.get_cache(cache_key="native") == {"value": "native"}
with rebound(facade.cache, "disk_cache", diskcache.Cache(str(tmp_path))):
assert resolver.resolve().kind == "python_callback"
assert resolver.resolve().kind == "native"
class CustomDiskCache(DiskCache):
pass
with rebound(facade, "cache", CustomDiskCache(disk_cache_dir=str(tmp_path))):
assert resolver.resolve().kind == "python_callback"
class CustomStore(diskcache.Cache):
pass
custom_facade: Final = Cache(type=LiteLLMCacheType.DISK, disk_cache_dir=str(tmp_path))
custom_facade.cache.disk_cache = CustomStore(str(tmp_path))
with pytest.raises(TypeError, match="built-in diskcache store"):
_native._CacheTestHandle.disk(str(tmp_path))._bind_facade(custom_facade)
async def test_disk_native_batch_lookup_and_store_report_partial_hits(tmp_path: Path) -> None:
binding: Final = _native._CacheTestResolver(
SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path)))
).resolve()
requests: Final = [request("hit"), request("miss"), request("disabled")]
requests[2]["controls"] = {
"supported_call_type": True,
"configured": True,
"native_backend": True,
"default_on": True,
"caching": False,
"no_cache": False,
"no_store": False,
"use_cache": False,
}
await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}])
partial: Final = await binding.async_lookup_batch(requests)
assert partial == {
"values": [{"value": 1}, {"value": 2}, None],
"missing_indices": [2],
}
@pytest.fixture
def s3_stub() -> Generator[S3Stub]:
stub: Final = S3Stub()
try:
yield stub
finally:
stub.close()
def python_s3(url: str) -> S3Cache:
return S3Cache(
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
)
async def test_s3_reads_python_entries_and_writes_with_python_metadata(s3_stub: S3Stub) -> None:
python_cache: Final = python_s3(s3_stub.url)
response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}}
python_cache.set_cache("sync:key", {"timestamp": time.time(), "response": response}, ttl=90)
python_cache.set_cache("plain", {"timestamp": time.time(), "response": response})
s3_stub.put_object("team/malformed", b"not a cache entry")
s3_stub.put_object(
"team/expired",
json.dumps({"timestamp": time.time(), "response": response}).encode(),
{"expires": "Thu, 01 Jan 1970 00:00:00 GMT"},
)
binding: Final = _native._CacheTestResolver(
SimpleNamespace(
cache=_native._CacheTestHandle.s3(
"cache-bucket",
region="us-east-1",
endpoint_url=s3_stub.url,
key_prefix="team/",
access_key_id="key",
secret_access_key="secret",
)
)
).resolve()
assert binding.lookup(request("sync:key")) == response
assert await binding.async_lookup(request("plain")) == response
assert binding.lookup(request("malformed")) is None
assert binding.lookup(request("expired")) is None
assert binding.lookup(request("absent")) is None
binding.store({**request("native:key"), "ttl_seconds": 90.0}, response)
await binding.async_store(request("no_ttl"), response)
stored: Final = s3_stub.objects["team/native/key"]
assert stored.headers["content-type"] == "application/json"
assert stored.headers["content-language"] == "en"
assert stored.headers["content-disposition"] == 'inline; filename="team/native/key.json"'
assert stored.headers["cache-control"] == "immutable, max-age=90, s-maxage=90"
expires: Final = cast(datetime, s3_stub.expires("team/native/key"))
remaining: Final = (expires - datetime.now(expires.tzinfo)).total_seconds()
assert 60 < remaining <= 91
no_ttl: Final = s3_stub.objects["team/no_ttl"]
assert no_ttl.headers["cache-control"] == "immutable, max-age=31536000, s-maxage=31536000"
assert "expires" not in no_ttl.headers
assert python_cache.get_cache("native:key")["response"] == response
partial: Final = await binding.async_lookup_batch([request("native:key"), request("absent"), request("malformed")])
assert partial == {"values": [response, None, None], "missing_indices": [1, 2]}
def test_s3_facade_binds_only_exact_configuration_and_falls_back_on_mutation(s3_stub: S3Stub) -> None:
facade: Final = Cache(
type=LiteLLMCacheType.S3,
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=s3_stub.url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
)
handle: Final = _native._CacheTestHandle.s3(
"cache-bucket",
region="us-east-1",
endpoint_url=s3_stub.url,
key_prefix="team/",
access_key_id="key",
secret_access_key="secret",
)
with pytest.raises(TypeError, match="buckets must match"):
_native._CacheTestHandle.s3("other", region="us-east-1", endpoint_url=s3_stub.url)._bind_facade(facade)
with pytest.raises(TypeError, match="key prefixes must match"):
_native._CacheTestHandle.s3(
"cache-bucket", region="us-east-1", endpoint_url=s3_stub.url, key_prefix="other/"
)._bind_facade(facade)
handle._bind_facade(facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
binding: Final = resolver.resolve()
assert binding.kind == "native"
handler: Final = Mock()
facade.cache.s3_client.meta.events.register("before-call.s3.*", handler)
binding.store(request("native"), {"answer": 1})
assert binding.lookup(request("native")) == {"answer": 1}
assert handler.call_count == 0
assert "team/native" in s3_stub.objects
with rebound(facade.cache, "bucket_name", "other"):
assert resolver.resolve().kind == "python_callback"
other_client: Final = boto3.client(
"s3",
region_name="us-east-1",
endpoint_url=s3_stub.url,
aws_access_key_id="key",
aws_secret_access_key="secret",
)
with rebound(facade.cache, "s3_client", other_client):
assert resolver.resolve().kind == "python_callback"
class CustomS3Cache(S3Cache):
pass
subclassed: Final = Cache(
type=LiteLLMCacheType.S3,
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=s3_stub.url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
)
subclassed.cache = CustomS3Cache(
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=s3_stub.url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
)
with pytest.raises(TypeError):
handle._bind_facade(subclassed)
assert _native._CacheTestResolver(SimpleNamespace(cache=subclassed)).resolve().kind == "python_callback"
def test_s3_facade_rejects_configurations_that_require_python(s3_stub: S3Stub) -> None:
handle: Final = _native._CacheTestHandle.s3(
"cache-bucket",
region="us-east-1",
endpoint_url=s3_stub.url,
key_prefix="team/",
access_key_id="key",
secret_access_key="secret",
)
unverified: Final = Cache(
type=LiteLLMCacheType.S3,
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url="https://s3.example.test",
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
s3_verify=False,
)
with pytest.raises(TypeError, match="requires Python"):
handle._bind_facade(unverified)
proxied: Final = Cache(
type=LiteLLMCacheType.S3,
s3_bucket_name="cache-bucket",
s3_region_name="us-east-1",
s3_endpoint_url=s3_stub.url,
s3_aws_access_key_id="key",
s3_aws_secret_access_key="secret",
s3_path="team",
s3_config=botocore.config.Config(proxies={"https": "http://proxy.test"}),
)
with pytest.raises(TypeError, match="requires Python"):
handle._bind_facade(proxied)
async def test_gcs_reads_python_entries_and_writes_python_compatible_objects(
fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False)
monkeypatch.delenv("GCS_BUCKET_NAME", raising=False)
response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None}
fake_gcs.put(
"bucket",
"cache/sync",
json.dumps({"timestamp": time.time(), "response": json.dumps(response)}).encode(),
)
fake_gcs.put("bucket", "cache/async", json.dumps({"timestamp": time.time(), "response": response}).encode())
fake_gcs.put("bucket", "cache/raw", json.dumps(response).encode())
fake_gcs.put("bucket", "cache/invalid", b"not a cache entry")
binding: Final = _native._CacheTestResolver(
SimpleNamespace(
cache=_native._CacheTestHandle.gcs(
"bucket",
gcs_path="cache",
endpoint=fake_gcs.url,
token=fake_gcs.token,
)
)
).resolve()
assert binding.lookup(request("sync")) == response
assert await binding.async_lookup(request("async")) == response
assert binding.lookup(request("raw")) == response
assert await binding.async_lookup(request("invalid")) is None
assert binding.lookup(request("missing")) is None
await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response)
stored: Final = fake_gcs.objects[("bucket", "cache/native")]
stored_value: Final = cast(dict[str, object], json.loads(stored))
assert stored_value["response"] == response
assert isinstance(stored_value["timestamp"], float)
upload: Final = next(item for item in fake_gcs.requests if item.method == "POST")
assert upload.path == "/upload/storage/v1/b/bucket/o"
assert upload.query == "uploadType=media&name=cache%2Fnative"
assert upload.headers["Authorization"] == f"Bearer {fake_gcs.token}"
assert upload.headers["Content-Type"] == "application/json"
upload_text: Final = f"{upload.path}?{upload.query}{upload.headers}"
assert "ttl" not in upload_text.lower()
assert "expiry" not in upload_text.lower()
download: Final = next(item for item in fake_gcs.requests if item.path.endswith("/cache%2Fsync"))
assert download.path == "/storage/v1/b/bucket/o/cache%2Fsync"
assert download.query == "alt=media"
binding.store(request("sync2"), response)
assert binding.lookup(request("sync2")) == response
assert GCSCache(bucket_name="bucket", gcs_path="cache").key_prefix == "cache/"
assert GCSCache(bucket_name="bucket", gcs_path="cache/").key_prefix == "cache/"
assert GCSCache(bucket_name="bucket").key_prefix == ""
async def test_gcs_batch_lookup_preserves_order_and_treats_malformed_entries_as_misses(fake_gcs: FakeGcs) -> None:
fake_gcs.put("bucket", "cache/hit", json.dumps({"timestamp": time.time(), "response": {"value": 1}}).encode())
fake_gcs.put("bucket", "cache/invalid", b"not a cache entry")
binding: Final = _native._CacheTestResolver(
SimpleNamespace(
cache=_native._CacheTestHandle.gcs(
"bucket",
gcs_path="cache",
endpoint=fake_gcs.url,
token=fake_gcs.token,
)
)
).resolve()
requests: Final = [request("hit"), request("missing"), request("invalid")]
expected: Final = {"values": [{"value": 1}, None, None], "missing_indices": [1, 2]}
assert await binding.async_lookup_batch(requests) == expected
assert binding.lookup_batch(requests) == expected
await binding.async_store_batch([request("first"), request("second")], [{"value": 1}, {"value": 2}])
assert ("bucket", "cache/first") in fake_gcs.objects
assert ("bucket", "cache/second") in fake_gcs.objects
async def test_gcs_facade_binds_only_exact_matching_configuration(
fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False)
monkeypatch.delenv("GCS_BUCKET_NAME", raising=False)
monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", "/nonexistent")
facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/")
assert type(facade.cache) is GCSCache
mismatched_bucket: Final = _native._CacheTestHandle.gcs(
"other",
gcs_path="cache",
endpoint=fake_gcs.url,
token=fake_gcs.token,
)
with pytest.raises(TypeError, match="buckets must match"):
mismatched_bucket._bind_facade(facade)
mismatched_prefix: Final = _native._CacheTestHandle.gcs(
"bucket",
gcs_path="x",
endpoint=fake_gcs.url,
token=fake_gcs.token,
)
with pytest.raises(TypeError, match="key prefixes must match"):
mismatched_prefix._bind_facade(facade)
mismatched_credentials: Final = _native._CacheTestHandle.gcs(
"bucket",
gcs_path="cache",
path_service_account="sa.json",
endpoint=fake_gcs.url,
token=fake_gcs.token,
)
with pytest.raises(TypeError, match="credentials must match"):
mismatched_credentials._bind_facade(facade)
with pytest.raises(TypeError, match="types must match"):
_native._CacheTestHandle.memory()._bind_facade(facade)
matching: Final = _native._CacheTestHandle.gcs(
"bucket",
gcs_path="cache",
endpoint=fake_gcs.url,
token=fake_gcs.token,
)
matching._bind_facade(facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
binding: Final = resolver.resolve()
assert binding.kind == "native"
await binding.async_store(request("native"), {"value": "native"})
assert await binding.async_lookup(request("native")) == {"value": "native"}
assert cast(CacheLookup, facade).get_cache(cache_key="native") is None
with rebound(facade.cache, "bucket_name", "other"):
assert resolver.resolve().kind == "python_callback"
with rebound(facade.cache, "key_prefix", "x/"):
assert resolver.resolve().kind == "python_callback"
with rebound(facade.cache, "path_service_account", "sa.json"):
assert resolver.resolve().kind == "python_callback"
def no_get_cache(*args: object, **kwargs: object) -> None:
return None
with rebound(facade.cache, "get_cache", no_get_cache):
assert resolver.resolve().kind == "python_callback"
with rebound(facade, "ttl", 12):
assert resolver.resolve().kind == "python_callback"
class CustomGcs(GCSCache):
pass
with rebound(facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")):
assert resolver.resolve().kind == "python_callback"
custom_facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/")
with rebound(custom_facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")):
with pytest.raises(TypeError, match="types must match"):
matching._bind_facade(custom_facade)
missing_bucket: Final = Cache(type=LiteLLMCacheType.GCS)
with pytest.raises(TypeError, match="requires a configured bucket name"):
matching._bind_facade(missing_bucket)
async def test_gcs_flush_is_a_no_op_and_ping_is_not_implemented(
fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False)
monkeypatch.delenv("GCS_BUCKET_NAME", raising=False)
binding: Final = _native._CacheTestResolver(
SimpleNamespace(
cache=_native._CacheTestHandle.gcs(
"bucket",
gcs_path="cache",
endpoint=fake_gcs.url,
token=fake_gcs.token,
)
)
).resolve()
await binding.async_store(request("key"), {"value": "stored"})
await binding.async_flush()
assert ("bucket", "cache/key") in fake_gcs.objects
assert await binding.async_lookup(request("key")) == {"value": "stored"}
with pytest.raises(NotImplementedError):
await binding.ping()
facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/")
with pytest.raises(AttributeError):
await facade.ping()
assert cast(CacheLookup, facade.cache).flush_cache() is None
async def test_gcs_unauthorized_and_server_errors_surface_as_runtime_errors(fake_gcs: FakeGcs) -> None:
wrong_token: Final = _native._CacheTestResolver(
SimpleNamespace(
cache=_native._CacheTestHandle.gcs(
"bucket",
gcs_path="cache",
endpoint=fake_gcs.url,
token="wrong-token",
)
)
).resolve()
with pytest.raises(RuntimeError):
wrong_token.lookup(request("missing"))
assert not fake_gcs.objects
binding: Final = _native._CacheTestResolver(
SimpleNamespace(
cache=_native._CacheTestHandle.gcs(
"bucket",
gcs_path="cache",
endpoint=fake_gcs.url,
token=fake_gcs.token,
)
)
).resolve()
with pytest.raises(RuntimeError):
binding.lookup(request("server-error"))
assert binding.lookup(request("missing")) is None
async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_natively(
cluster_nodes: tuple[tuple[str, int], ...],
) -> None:
@ -574,7 +1104,9 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n
await binding.async_flush()
remaining: Final = tuple(sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node)))
remaining: Final = tuple(
sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node))
)
assert remaining == (), remaining
assert client.get("unscoped") == b"stays"
client.delete("unscoped")

6
tests/unit/models/test_models.py
View file

@ -324,6 +324,8 @@ class TestUser:
assert user_no_models.has_model_access("any-model")
def test_password_hash_excluded_from_serialization(self):
import json
from litellm.proxy._types import LiteLLM_UserTableWithKeyCount
secret = "$2b$12$abcdefghijklmnopqrstuv"
@ -331,12 +333,12 @@ class TestUser:
assert user.password == secret
assert "password" not in user.model_dump()
assert "password" not in user.model_dump_json()
assert "password" not in json.loads(user.model_dump_json())
with_keys = LiteLLM_UserTableWithKeyCount(user_id="u1", user_email="a@b.c", password=secret, key_count=2)
assert with_keys.password == secret
assert "password" not in with_keys.model_dump()
assert "password" not in with_keys.model_dump_json()
assert "password" not in json.loads(with_keys.model_dump_json())
class TestVerificationToken:

110
ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.integration.test.tsx Normal file
View file

@ -0,0 +1,110 @@
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
import { beforeEach, describe, expect, it, vi } from "vitest";
import ChangePasswordForm from "./ChangePasswordForm";
const mockChangePasswordCall = vi.fn();
const mockToastSuccess = vi.fn();
const mockClearTokenCookies = vi.fn();
let mockPasswordResetRequired = false;
vi.mock("@/components/networking", () => ({
changePasswordCall: (...args: unknown[]) => mockChangePasswordCall(...args),
getProxyBaseUrl: () => "",
}));
vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
default: () => ({ accessToken: "sk-session-token", passwordResetRequired: mockPasswordResetRequired }),
}));
vi.mock("@/lib/toast", () => ({
toast: {
success: (...args: unknown[]) => mockToastSuccess(...args),
fromError: vi.fn(),
},
}));
vi.mock("@/utils/cookieUtils", () => ({
clearTokenCookies: (...args: unknown[]) => mockClearTokenCookies(...args),
}));
const fillForm = (values: { current: string; next: string; confirm: string }) => {
fireEvent.change(screen.getByLabelText("Current Password"), { target: { value: values.current } });
fireEvent.change(screen.getByLabelText("New Password"), { target: { value: values.next } });
fireEvent.change(screen.getByLabelText("Confirm New Password"), { target: { value: values.confirm } });
};
const submit = () => fireEvent.click(screen.getByRole("button", { name: "Change Password" }));
describe("ChangePasswordForm", () => {
beforeEach(() => {
vi.clearAllMocks();
mockPasswordResetRequired = false;
});
it("sends the current and new password to the change endpoint and resets on success", async () => {
mockChangePasswordCall.mockResolvedValue({ user_id: "user-123", message: "Password updated successfully." });
render(<ChangePasswordForm />);
fillForm({ current: "OldP@ssw0rd-2026", next: "NewP@ssw0rd-2026", confirm: "NewP@ssw0rd-2026" });
submit();
expect(await screen.findByLabelText("Current Password")).toHaveValue("");
expect(mockChangePasswordCall).toHaveBeenCalledWith("sk-session-token", "OldP@ssw0rd-2026", "NewP@ssw0rd-2026");
expect(mockToastSuccess).toHaveBeenCalled();
});
it("blocks submission when the confirmation does not match", async () => {
render(<ChangePasswordForm />);
fillForm({ current: "OldP@ssw0rd-2026", next: "NewP@ssw0rd-2026", confirm: "Different-2026" });
submit();
expect(await screen.findByText("New passwords do not match")).toBeInTheDocument();
expect(mockChangePasswordCall).not.toHaveBeenCalled();
});
it("shows the proxy's rejection message unwrapped", async () => {
mockChangePasswordCall.mockRejectedValue(new Error("{'error': 'Current password is incorrect.'}"));
render(<ChangePasswordForm />);
fillForm({ current: "wrong-password", next: "NewP@ssw0rd-2026", confirm: "NewP@ssw0rd-2026" });
submit();
expect(await screen.findByText("Current password is incorrect.")).toBeInTheDocument();
expect(mockToastSuccess).not.toHaveBeenCalled();
});
describe("forced password reset", () => {
it("shows the forced-reset warning only when the session is flagged", () => {
mockPasswordResetRequired = true;
render(<ChangePasswordForm />);
expect(screen.getByText(/must be changed before you can use the dashboard/)).toBeInTheDocument();
});
it("hides the forced-reset warning for a normal session", () => {
render(<ChangePasswordForm />);
expect(screen.queryByText(/must be changed before you can use the dashboard/)).not.toBeInTheDocument();
});
it("signs the user out to re-login after a successful forced change", async () => {
mockPasswordResetRequired = true;
mockChangePasswordCall.mockResolvedValue({ user_id: "user-123", message: "Password updated successfully." });
const replaceMock = vi.fn();
const realLocation = window.location;
Object.defineProperty(window, "location", { configurable: true, value: { replace: replaceMock } });
try {
render(<ChangePasswordForm />);
fillForm({ current: "OldP@ssw0rd-2026", next: "NewP@ssw0rd-2026", confirm: "NewP@ssw0rd-2026" });
submit();
await waitFor(() => expect(replaceMock).toHaveBeenCalledWith("/ui/login/"));
expect(mockClearTokenCookies).toHaveBeenCalled();
} finally {
Object.defineProperty(window, "location", { configurable: true, value: realLocation });
}
});
});
});

120
ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx Normal file
View file

@ -0,0 +1,120 @@
"use client";
import React, { useState } from "react";
import { CircleAlert } from "lucide-react";
import { z } from "zod/v4";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import { Alert, AlertTitle } from "@/components/shared/Alert";
import { PasswordInput } from "@/components/shared/PasswordInput";
import { FormField } from "@/components/shared/form/FormField";
import { Button } from "@/components/ui/button";
import { Card, CardContent } from "@/components/ui/card";
import { FieldGroup } from "@/components/ui/field";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import { changePasswordCall, getProxyBaseUrl } from "@/components/networking";
import { extractProxyErrorMessage } from "@/lib/http/client";
import { useZodForm } from "@/lib/forms/useZodForm";
import { toast } from "@/lib/toast";
import { clearTokenCookies } from "@/utils/cookieUtils";
import { getLoginUrl } from "@/utils/returnUrlUtils";
const changePasswordSchema = z
.object({
currentPassword: z.string().min(1, "Current password is required"),
newPassword: z.string().min(1, "New password is required"),
confirmNewPassword: z.string().min(1, "Confirm your new password"),
})
.refine((values) => values.newPassword === values.confirmNewPassword, {
message: "New passwords do not match",
path: ["confirmNewPassword"],
});
type ChangePasswordValues = z.infer<typeof changePasswordSchema>;
export function ChangePasswordForm() {
const { accessToken, passwordResetRequired } = useAuthorized();
const form = useZodForm(changePasswordSchema, {
defaultValues: { currentPassword: "", newPassword: "", confirmNewPassword: "" },
});
const [isPending, setIsPending] = useState(false);
const [submitError, setSubmitError] = useState<string | null>(null);
const handleSubmit = async (values: ChangePasswordValues) => {
if (!accessToken) return;
setSubmitError(null);
setIsPending(true);
try {
await changePasswordCall(accessToken, values.currentPassword, values.newPassword);
if (passwordResetRequired) {
// The session key was minted restricted; only a fresh login lifts it.
toast.success("Password updated. Please log in with your new password.");
clearTokenCookies();
window.location.replace(getLoginUrl(getProxyBaseUrl()));
return;
}
toast.success("Password updated");
form.reset();
} catch (error) {
setSubmitError(extractProxyErrorMessage(error));
} finally {
setIsPending(false);
}
};
return (
<div className="mx-auto mt-10 w-full max-w-md">
<Card>
<CardContent>
<h3 className="text-2xl font-semibold text-foreground">Change Password</h3>
<p className="text-sm text-muted-foreground">
Enter your current password and choose a new one. The new password must meet this proxy&apos;s password
policy.
</p>
{passwordResetRequired && (
<Alert variant="warning" className="mt-4">
<CircleAlert />
<AlertTitle>
Your password must be changed before you can use the dashboard: it was either found in a known data
breach or set by an administrator as a temporary password. After updating it, you will be signed out to
log in again.
</AlertTitle>
</Alert>
)}
<form className="mb-2 mt-8" onSubmit={form.handleSubmit(handleSubmit)}>
<FieldGroup>
<FormField control={form.control} name="currentPassword" label="Current Password">
{({ ref, ...field }) => <PasswordInput {...field} ref={ref} autoComplete="current-password" />}
</FormField>
<FormField control={form.control} name="newPassword" label="New Password">
{({ ref, ...field }) => <PasswordInput {...field} ref={ref} autoComplete="new-password" />}
</FormField>
<FormField control={form.control} name="confirmNewPassword" label="Confirm New Password">
{({ ref, ...field }) => <PasswordInput {...field} ref={ref} autoComplete="new-password" />}
</FormField>
</FieldGroup>
{submitError && (
<Alert variant="error" className="mt-6">
<CircleAlert />
<AlertTitle>{submitError}</AlertTitle>
</Alert>
)}
<div className="mt-8">
<Button type="submit" disabled={isPending}>
{isPending && <UiLoadingSpinner className="size-4" role="img" aria-label="loading" />}
Change Password
</Button>
</div>
</form>
</CardContent>
</Card>
</div>
);
}
export default ChangePasswordForm;

7
ui/litellm-dashboard/src/app/(dashboard)/change-password/page.tsx Normal file
View file

@ -0,0 +1,7 @@
"use client";
import ChangePasswordForm from "./ChangePasswordForm";
export default function ChangePasswordPage() {
return <ChangePasswordForm />;
}

2
ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts
View file

@ -50,7 +50,9 @@ const useAuthorized = () => {
isViewOnly: isViewOnlySessionRole(decoded?.user_role),
premiumUser: decoded?.premium_user ?? null,
disabledPersonalKeyCreation: decoded?.disabled_non_admin_personal_key_creation ?? null,
loginMethod: decoded?.login_method ?? null,
showSSOBanner: decoded?.login_method === "username_password",
passwordResetRequired: decoded?.password_reset_required === true,
};
};

58
ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx
View file

@ -1,4 +1,4 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { render, screen, waitFor } from "@testing-library/react";
import { AuthProvider } from "@/contexts/AuthContext";
import Layout from "./layout";
@ -121,4 +121,60 @@ describe("(dashboard) Layout", () => {
expect(screen.queryByTestId("dashboard-header")).not.toBeInTheDocument();
expect(screen.queryByTestId("sidebar")).not.toBeInTheDocument();
});
describe("forced password reset routing", () => {
const sessionCookie = (claims: Record<string, unknown>) => {
const encode = (part: Record<string, unknown>) =>
btoa(JSON.stringify(part)).replace(/\+/g, "-").replace(/\//g, "_").replace(/=+$/, "");
const exp = Math.floor(Date.now() / 1000) + 3600;
return `${encode({ alg: "HS256", typ: "JWT" })}.${encode({ ...claims, exp })}.sig`;
};
afterEach(() => {
document.cookie = "token=; Max-Age=0; Path=/";
});
it("routes a session flagged password_reset_required to the change-password page", async () => {
const flaggedClaims = {
user_id: "flagged-user",
key: "sk-session",
login_method: "username_password",
password_reset_required: true,
};
document.cookie = `token=${sessionCookie(flaggedClaims)}; Path=/`;
render(
<AuthProvider>
<Layout>
<div data-testid="page-content" />
</Layout>
</AuthProvider>,
);
pendingUiConfig.resolve();
await waitFor(() => expect(replaceMock).toHaveBeenCalledWith(expect.stringContaining("/change-password")));
});
it("does not reroute an unflagged session", async () => {
document.cookie = `token=${sessionCookie({
user_id: "normal-user",
key: "sk-session",
login_method: "username_password",
})}; Path=/`;
render(
<AuthProvider>
<Layout>
<div data-testid="page-content" />
</Layout>
</AuthProvider>,
);
pendingUiConfig.resolve();
expect(await screen.findByTestId("page-content")).toBeInTheDocument();
expect(replaceMock).not.toHaveBeenCalled();
});
});
});

13
ui/litellm-dashboard/src/app/(dashboard)/layout.tsx
View file

@ -7,7 +7,7 @@ import LoadingScreen from "@/components/common_components/LoadingScreen";
import { ThemeProvider } from "@/contexts/ThemeContext";
import { useAuth } from "@/contexts/AuthContext";
import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider";
import { useRouter, useSearchParams } from "next/navigation";
import { usePathname, useRouter, useSearchParams } from "next/navigation";
import { DebugWarningBanner } from "@/components/DebugWarningBanner";
import { NoRedisWarningBanner } from "@/components/NoRedisWarningBanner";
import { EnvCredentialLoginWarningBanner } from "@/components/EnvCredentialLoginWarningBanner";
@ -149,7 +149,8 @@ function DashboardShell({ children }: { children: React.ReactNode }) {
function LayoutContent({ children }: { children: React.ReactNode }) {
const router = useRouter();
const searchParams = useSearchParams();
const { accessToken, authLoading } = useAuth();
const pathname = usePathname();
const { accessToken, authLoading, passwordResetRequired } = useAuth();
const isInvitationFlow = Boolean(searchParams.get("invitation_id"));
// Legacy invitation links point at /ui/?invitation_id=; the onboarding form now lives at its own
@ -160,6 +161,14 @@ function LayoutContent({ children }: { children: React.ReactNode }) {
}
}, [authLoading, isInvitationFlow, router, searchParams]);
// A session flagged for a forced password reset can only reach the change-password
// endpoint server-side; keep the UI on the matching page.
useEffect(() => {
if (!authLoading && passwordResetRequired && !pathname?.endsWith("/change-password")) {
router.replace(uiHref("change-password"));
}
}, [authLoading, passwordResetRequired, pathname, router]);
if (authLoading || isInvitationFlow) {
return <LoadingScreen />;
}

227
ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx
View file

@ -3,7 +3,8 @@ import { render, waitFor, screen, act, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { describe, it, expect, vi, beforeEach } from "vitest";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import MCPServers from "./mcp_servers";
import MCPServers, { compareServers, type SortKey } from "./mcp_servers";
import type { MCPServer } from "@/components/mcp_tools/types";
import * as networking from "@/components/networking";
// Mock the networking module
@ -31,6 +32,98 @@ const createQueryClient = () =>
},
});
describe("compareServers", () => {
const server = (server_id: string, name: string, created_at = ""): MCPServer => ({
server_id,
server_name: name,
created_at,
updated_at: created_at,
created_by: "user",
updated_by: "user",
});
const shuffled = [server("c", "github"), server("a", "slack"), server("b", "Jira")];
it("orders servers without timestamps by name so config.yaml servers render in a stable order", () => {
const byCreated = [...shuffled].sort((a, b) => compareServers(a, b, "created_desc")).map((s) => s.server_id);
const byUpdated = [...shuffled].sort((a, b) => compareServers(a, b, "updated_desc")).map((s) => s.server_id);
const byHealth = [...shuffled].sort((a, b) => compareServers(a, b, "health")).map((s) => s.server_id);
expect(byCreated).toEqual(["c", "b", "a"]);
expect(byUpdated).toEqual(["c", "b", "a"]);
expect(byHealth).toEqual(["c", "b", "a"]);
});
it("keeps newest-first when timestamps differ", () => {
const newest = server("new", "zzz", "2026-02-01T00:00:00Z");
const oldest = server("old", "aaa", "2026-01-01T00:00:00Z");
expect([oldest, newest].sort((a, b) => compareServers(a, b, "created_desc")).map((s) => s.server_id)).toEqual([
"new",
"old",
]);
});
it.each<SortKey>(["created_desc", "updated_desc", "name_asc", "health"])(
"breaks equal timestamps and names by ID for %s regardless of input order",
(sort) => {
const servers = [
server("b", "GitHub", "2026-01-01T00:00:00Z"),
server("c", "Slack", "2026-01-01T00:00:00Z"),
server("a", "github", "2026-01-01T00:00:00Z"),
];
for (const input of [servers, [...servers].reverse()]) {
expect([...input].sort((a, b) => compareServers(a, b, sort)).map((s) => s.server_id)).toEqual(["a", "b", "c"]);
}
},
);
it("uses the display name before alias, then falls back to alias and ID", () => {
const servers: MCPServer[] = [
{ ...server("s-slack", "Slack"), alias: "aaa" },
{ ...server("s-github", ""), server_name: null, alias: "GitHub" },
{ ...server("confluence", ""), alias: "" },
];
for (const input of [servers, [...servers].reverse()]) {
expect([...input].sort((a, b) => compareServers(a, b, "name_asc")).map((s) => s.server_id)).toEqual([
"confluence",
"s-github",
"s-slack",
]);
}
});
it.each<SortKey>(["created_desc", "updated_desc", "health"])(
"keeps timestamped servers before missing timestamps for %s",
(sort) => {
const servers = [
server("config", "aaa"),
server("older", "bbb", "2026-01-01T00:00:00Z"),
server("newer", "zzz", "2026-02-01T00:00:00Z"),
];
for (const input of [servers, [...servers].reverse()]) {
expect([...input].sort((a, b) => compareServers(a, b, sort)).map((s) => s.server_id)).toEqual([
"newer",
"older",
"config",
]);
}
},
);
it("sorts health before recency and display name", () => {
const servers: MCPServer[] = [
{ ...server("healthy", "aaa", "2026-03-01T00:00:00Z"), status: "healthy" },
{ ...server("unknown", "bbb", "2026-02-01T00:00:00Z"), status: "unknown" },
{ ...server("unhealthy", "zzz", "2026-01-01T00:00:00Z"), status: "unhealthy" },
];
expect(servers.sort((a, b) => compareServers(a, b, "health")).map((s) => s.server_id)).toEqual([
"unhealthy",
"unknown",
"healthy",
]);
});
});
describe("MCPServers", () => {
const defaultProps = {
accessToken: "123",
@ -74,6 +167,134 @@ describe("MCPServers", () => {
const myConnections = await screen.findByRole("link", { name: "My Connections" });
expect(myConnections).toBeVisible();
expect(myConnections).toHaveAttribute("href", "/ui/connect");
for (const name of ["Semantic Filter", "Tool Search", "Network Settings", "Submitted MCPs"]) {
const tab = screen.queryByRole("tab", { name });
if (userRole === "Admin") {
expect(tab).toBeVisible();
} else {
expect(tab).not.toBeInTheDocument();
}
}
expect(
screen.getByRole("button", {
name: userRole === "Admin" ? "+ Add New MCP Server" : "+ Submit MCP Server",
}),
).toBeVisible();
});
it.each(["cancel", "success", "failure", "unnamed"])("preserves delete confirmation on %s", async (outcome) => {
const server: MCPServer = {
created_at: "",
updated_at: "",
server_id: "delete-server",
server_name: outcome === "unnamed" ? null : "Delete fixture",
alias: "delete-alias",
url: outcome === "unnamed" ? null : "https://example.com/mcp",
created_by: "user",
updated_by: "user",
};
vi.mocked(networking.fetchMCPServers).mockResolvedValue([server]);
vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]);
let finishDelete: () => void = () => {};
vi.mocked(networking.deleteMCPServer).mockImplementation(
() =>
new Promise((resolve, reject) => {
finishDelete = () => (outcome === "failure" ? reject(new Error("Delete failed")) : resolve(undefined));
}),
);
render(
<QueryClientProvider client={createQueryClient()}>
<MCPServers {...defaultProps} />
</QueryClientProvider>,
);
await userEvent.click(await screen.findByRole("button", { name: "Server actions" }));
await userEvent.click(await screen.findByRole("menuitem", { name: "Delete" }));
const dialog = await screen.findByRole("alertdialog", { name: "Delete MCP Server?" });
expect(within(dialog).getByText("delete-server")).toBeVisible();
if (outcome === "unnamed") {
expect(within(dialog).queryByText("Name")).not.toBeInTheDocument();
expect(within(dialog).queryByText("URL")).not.toBeInTheDocument();
} else {
expect(within(dialog).getByText("Delete fixture")).toBeVisible();
expect(within(dialog).getByText("https://example.com/mcp")).toBeVisible();
}
if (outcome === "cancel") {
await userEvent.click(within(dialog).getByRole("button", { name: "Cancel" }));
expect(networking.deleteMCPServer).not.toHaveBeenCalled();
} else {
await userEvent.click(within(dialog).getByRole("button", { name: "Delete" }));
expect(within(dialog).getByRole("button", { name: "Deleting..." })).toBeDisabled();
expect(within(dialog).getByRole("button", { name: "Cancel" })).toBeDisabled();
expect(networking.deleteMCPServer).toHaveBeenCalledWith("123", "delete-server");
await act(async () => finishDelete());
}
await waitFor(() => expect(screen.queryByRole("alertdialog")).not.toBeInTheDocument());
});
it("filters servers by access group", async () => {
const server = { created_by: "user", updated_by: "user" };
vi.mocked(networking.fetchMCPServers).mockResolvedValue([
{
...server,
server_id: "string-group",
server_name: "String group",
alias: "string-alias",
mcp_access_groups: ["shared"],
},
{
...server,
server_id: "legacy-group",
server_name: "Legacy group",
alias: "legacy-alias",
mcp_access_groups: ["shared"],
},
{
...server,
server_id: "other-group",
server_name: "Other group",
alias: "other-alias",
mcp_access_groups: ["different"],
},
]);
vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]);
render(
<QueryClientProvider client={createQueryClient()}>
<MCPServers {...defaultProps} />
</QueryClientProvider>,
);
await screen.findByText("String group");
await userEvent.click(screen.getByRole("combobox", { name: "Access Group" }));
await userEvent.click(await screen.findByRole("option", { name: "shared" }));
expect(screen.getByText("String group")).toBeVisible();
expect(screen.getByText("Legacy group")).toBeVisible();
expect(screen.queryByText("Other group")).not.toBeInTheDocument();
});
it.each(["server_name", "alias", "url", "server_id"] as const)("searches by %s case-insensitively", async (field) => {
const server: MCPServer = {
created_at: "",
updated_at: "",
server_id: "search-server",
server_name: "Search fixture",
created_by: "user",
updated_by: "user",
[field]: "Needle",
};
vi.mocked(networking.fetchMCPServers).mockResolvedValue([server]);
vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]);
render(
<QueryClientProvider client={createQueryClient()}>
<MCPServers {...defaultProps} />
</QueryClientProvider>,
);
await screen.findByTestId("mcp-servers-grid");
const search = screen.getByPlaceholderText("Search by name, alias, URL, or ID");
await userEvent.type(search, " NEEDLE ");
expect(screen.getByTestId("mcp-servers-grid")).toBeVisible();
await userEvent.clear(search);
await userEvent.type(search, "no-match");
expect(screen.queryByTestId("mcp-servers-grid")).not.toBeInTheDocument();
expect(screen.getByText("No servers match the current filters or search.")).toBeVisible();
});
it("should render mocked MCP servers data in the table", async () => {
@ -316,9 +537,7 @@ describe("MCPServers", () => {
expect(screen.getByText("Team B Server")).toBeInTheDocument();
expect(screen.getByText("Team A Server 2")).toBeInTheDocument();
// Find the team select by its "Team" label, then the combobox it labels
const teamLabel = screen.getByText("Team");
const teamSelect = within(teamLabel.parentElement!).getByRole("combobox");
const teamSelect = screen.getByRole("combobox", { name: "Team" });
await userEvent.click(teamSelect);

244
ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx
View file

@ -49,7 +49,7 @@ import { cn } from "@/lib/cva.config";
import UserEnvVarsModal from "./UserEnvVarsModal";
import { listMCPUserEnvVarStatus } from "@/components/networking";
type SortKey = "created_desc" | "updated_desc" | "name_asc" | "health";
export type SortKey = "created_desc" | "updated_desc" | "name_asc" | "health";
const SORT_OPTIONS: { value: SortKey; label: string }[] = [
{ value: "created_desc", label: "Recently created" },
@ -64,32 +64,33 @@ const HEALTH_RANK: Record<string, number> = {
healthy: 2,
};
const compareServers = (a: MCPServer, b: MCPServer, sort: SortKey): number => {
const compareByName = (a: MCPServer, b: MCPServer): number => {
const nameA = (a.server_name || a.alias || a.server_id).toLowerCase();
const nameB = (b.server_name || b.alias || b.server_id).toLowerCase();
return nameA.localeCompare(nameB) || a.server_id.localeCompare(b.server_id);
};
const compareByTimestampDesc = (a: string | null | undefined, b: string | null | undefined): number => {
const ta = a ? new Date(a).getTime() : 0;
const tb = b ? new Date(b).getTime() : 0;
return tb - ta;
};
export const compareServers = (a: MCPServer, b: MCPServer, sort: SortKey): number => {
switch (sort) {
case "name_asc": {
const nameA = (a.server_name || a.alias || a.server_id).toLowerCase();
const nameB = (b.server_name || b.alias || b.server_id).toLowerCase();
return nameA.localeCompare(nameB);
}
case "updated_desc": {
const ta = a.updated_at ? new Date(a.updated_at).getTime() : 0;
const tb = b.updated_at ? new Date(b.updated_at).getTime() : 0;
return tb - ta;
}
case "name_asc":
return compareByName(a, b);
case "updated_desc":
return compareByTimestampDesc(a.updated_at, b.updated_at) || compareByName(a, b);
case "health": {
const ra = HEALTH_RANK[a.status ?? "unknown"] ?? 1;
const rb = HEALTH_RANK[b.status ?? "unknown"] ?? 1;
if (ra !== rb) return ra - rb;
const ta = a.created_at ? new Date(a.created_at).getTime() : 0;
const tb = b.created_at ? new Date(b.created_at).getTime() : 0;
return tb - ta;
return compareByTimestampDesc(a.created_at, b.created_at) || compareByName(a, b);
}
case "created_desc":
default: {
const ta = a.created_at ? new Date(a.created_at).getTime() : 0;
const tb = b.created_at ? new Date(b.created_at).getTime() : 0;
return tb - ta;
}
default:
return compareByTimestampDesc(a.created_at, b.created_at) || compareByName(a, b);
}
};
@ -112,6 +113,62 @@ const readToolsOAuthServerId = (): string | null => {
}
};
function DeleteServerDialog({
open,
onOpenChange,
server,
isDeleting,
onConfirm,
}: {
open: boolean;
onOpenChange: (open: boolean) => void;
server: MCPServer | undefined;
isDeleting: boolean;
onConfirm: () => Promise<void>;
}) {
return (
<AlertDialog open={open} onOpenChange={onOpenChange}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>Delete MCP Server?</AlertDialogTitle>
</AlertDialogHeader>
<div className="space-y-4">
<p className="text-sm text-muted-foreground">
This action is permanent and cannot be undone. All associated configurations will be removed.
</p>
{server && (
<dl className="mt-3 space-y-1 rounded-lg border border-border bg-muted p-4">
{server.server_name && (
<div className="flex gap-2">
<dt className="text-sm text-muted-foreground">Name</dt>
<dd className="text-sm font-semibold">{server.server_name}</dd>
</div>
)}
<div className="flex gap-2">
<dt className="text-sm text-muted-foreground">ID</dt>
<dd className="font-mono text-xs">{server.server_id}</dd>
</div>
{server.url && (
<div className="flex gap-2">
<dt className="text-sm text-muted-foreground">URL</dt>
<dd className="font-mono text-xs break-all">{server.url}</dd>
</div>
)}
</dl>
)}
</div>
<AlertDialogFooter>
<AlertDialogCancel disabled={isDeleting}>Cancel</AlertDialogCancel>
<Button variant="destructive" disabled={isDeleting} onClick={onConfirm}>
{isDeleting ? "Deleting..." : "Delete"}
</Button>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
);
}
const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, isViewOnly = false }) => {
const { data: mcpServers, isLoading: isLoadingServers, refetch } = useMCPServers();
@ -298,16 +355,12 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
}
if (group !== "all") {
filtered = filtered.filter((server) =>
server.mcp_access_groups?.some((g: any) => (typeof g === "string" ? g === group : g && g.name === group)),
server.mcp_access_groups?.some((g: string | { name?: string } | null) =>
typeof g === "string" ? g === group : g?.name === group,
),
);
}
const sorted = [...filtered].sort((a, b) => {
if (!a.created_at && !b.created_at) return 0;
if (!a.created_at) return 1;
if (!b.created_at) return -1;
return new Date(b.created_at).getTime() - new Date(a.created_at).getTime();
});
setFilteredServers(sorted);
setFilteredServers(filtered);
},
[serversWithHealth],
);
@ -338,7 +391,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
const alias = (s.alias || "").toLowerCase();
const url = (s.url || "").toLowerCase();
const id = s.server_id.toLowerCase();
return name.includes(q) || alias.includes(q) || url.includes(q) || id.includes(q);
return [name, alias, url, id].some((value) => value.includes(q));
})
: filteredServers;
return [...matches].sort((a, b) => compareServers(a, b, sortKey));
@ -381,9 +434,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
};
// Find the server to delete from the servers list
const serverToDelete = serverIdToDelete
? (mcpServers || []).find((server) => server.server_id === serverIdToDelete)
: null;
const serverToDelete = mcpServers?.find((server) => server.server_id === serverIdToDelete);
const handleCreateSuccess = (newMcpServer: MCPServer) => {
setFilteredServers((prev) => [...prev, newMcpServer]);
@ -425,45 +476,13 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
return (
<TooltipProvider>
<div className="h-full w-full p-6">
<AlertDialog open={isDeleteModalOpen} onOpenChange={(open) => !open && cancelDelete()}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>Delete MCP Server?</AlertDialogTitle>
</AlertDialogHeader>
<div className="space-y-4">
<p className="text-sm text-muted-foreground">
This action is permanent and cannot be undone. All associated configurations will be removed.
</p>
{serverToDelete && (
<dl className="mt-3 space-y-1 rounded-lg border border-border bg-muted p-4">
{serverToDelete.server_name && (
<div className="flex gap-2">
<dt className="text-sm text-muted-foreground">Name</dt>
<dd className="text-sm font-semibold">{serverToDelete.server_name}</dd>
</div>
)}
<div className="flex gap-2">
<dt className="text-sm text-muted-foreground">ID</dt>
<dd className="font-mono text-xs">{serverToDelete.server_id}</dd>
</div>
{serverToDelete.url && (
<div className="flex gap-2">
<dt className="text-sm text-muted-foreground">URL</dt>
<dd className="font-mono text-xs break-all">{serverToDelete.url}</dd>
</div>
)}
</dl>
)}
</div>
<AlertDialogFooter>
<AlertDialogCancel disabled={isDeletingServer}>Cancel</AlertDialogCancel>
<Button variant="destructive" disabled={isDeletingServer} onClick={confirmDelete}>
{isDeletingServer ? "Deleting..." : "Delete"}
</Button>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
<DeleteServerDialog
open={isDeleteModalOpen}
onOpenChange={(open) => !open && cancelDelete()}
server={serverToDelete}
isDeleting={isDeletingServer}
onConfirm={confirmDelete}
/>
<CreateMCPServer
userRole={userRole}
userID={userID}
@ -492,7 +511,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
<Plug />
My Connections
</Link>
{isAdminRole(userRole) && (
{isAdminRole(userRole) ? (
<>
<Button className="shrink-0" variant="secondary" onClick={() => setImportVisible(true)}>
Import from JSON
@ -501,8 +520,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
+ Add New MCP Server
</Button>
</>
)}
{!isAdminRole(userRole) && (
) : (
<Button
className="shrink-0"
onClick={() => {
@ -549,24 +567,23 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
Connect
</TabsTrigger>
{isAdminRole(userRole) && (
<TabsTrigger value="semantic-filter" className="flex-none rounded-none px-4 py-2">
Semantic Filter
</TabsTrigger>
)}
{isAdminRole(userRole) && (
<TabsTrigger value="tool-search" className="flex-none rounded-none px-4 py-2">
Tool Search
</TabsTrigger>
)}
{isAdminRole(userRole) && (
<TabsTrigger value="network-settings" className="flex-none rounded-none px-4 py-2">
Network Settings
</TabsTrigger>
)}
{isAdminRole(userRole) && (
<TabsTrigger value="submitted" className="flex-none rounded-none px-4 py-2">
Submitted MCPs
</TabsTrigger>
<>
<TabsTrigger value="semantic-filter" className="flex-none rounded-none px-4 py-2">
Semantic Filter
</TabsTrigger>
<TabsTrigger value="tool-search" className="flex-none rounded-none px-4 py-2">
Tool Search
</TabsTrigger>
<TabsTrigger value="network-settings" className="flex-none rounded-none px-4 py-2">
Network Settings
</TabsTrigger>
<TabsTrigger value="submitted" className="flex-none rounded-none px-4 py-2">
Submitted MCPs
</TabsTrigger>
</>
)}
{isProxyAdminTierRole(userRole) && (
<TabsTrigger value="connections" className="flex-none rounded-none px-4 py-2">
@ -601,13 +618,11 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
value={selectedTeam}
onValueChange={(v: string | null) => handleTeamChange(v ?? "all")}
>
<SelectTrigger className="w-55">
<SelectTrigger className="w-55" aria-label="Team">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="all">
{isInternalUser ? "All Available Servers" : "All Servers"}
</SelectItem>
<SelectItem value="all">{teamSelectItems.all}</SelectItem>
<SelectItem value="personal">Personal</SelectItem>
{uniqueTeams.map((team) => (
<SelectItem key={team.team_id} value={team.team_id}>
@ -641,7 +656,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
value={selectedMcpAccessGroup}
onValueChange={(v: string | null) => handleMcpAccessGroupChange(v ?? "all")}
>
<SelectTrigger className="w-55">
<SelectTrigger className="w-55" aria-label="Access Group">
<SelectValue />
</SelectTrigger>
<SelectContent>
@ -742,24 +757,23 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, i
<MCPConnect />
</TabsContent>
{isAdminRole(userRole) && (
<TabsContent value="semantic-filter" keepMounted>
<MCPSemanticFilterSettings accessToken={accessToken} />
</TabsContent>
)}
{isAdminRole(userRole) && (
<TabsContent value="tool-search" keepMounted>
<MCPToolSearchSettings accessToken={accessToken} />
</TabsContent>
)}
{isAdminRole(userRole) && (
<TabsContent value="network-settings" keepMounted>
<MCPNetworkSettings accessToken={accessToken} />
</TabsContent>
)}
{isAdminRole(userRole) && (
<TabsContent value="submitted" keepMounted>
<MCPSubmissionsTab accessToken={accessToken} />
</TabsContent>
<>
<TabsContent value="semantic-filter" keepMounted>
<MCPSemanticFilterSettings accessToken={accessToken} />
</TabsContent>
<TabsContent value="tool-search" keepMounted>
<MCPToolSearchSettings accessToken={accessToken} />
</TabsContent>
<TabsContent value="network-settings" keepMounted>
<MCPNetworkSettings accessToken={accessToken} />
</TabsContent>
<TabsContent value="submitted" keepMounted>
<MCPSubmissionsTab accessToken={accessToken} />
</TabsContent>
</>
)}
{isProxyAdminTierRole(userRole) && (
<TabsContent value="connections">

52
ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.test.tsx
View file

@ -3,13 +3,25 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
import { renderWithProviders, screen, waitFor } from "../../../../tests/test-utils";
import UserDropdown from "./UserDropdown";
let mockUseAuthorizedImpl = () => ({
let mockUseAuthorizedImpl: () => {
userId: string | null;
userEmail: string | null;
userRoleLabel: string;
premiumUser: boolean;
loginMethod?: string | null;
} = () => ({
userId: "test-user-id",
userEmail: "test@example.com",
userRoleLabel: "Admin",
premiumUser: false,
});
const mockRouterPush = vi.fn();
vi.mock("next/navigation", () => ({
useRouter: () => ({ push: mockRouterPush }),
}));
let mockUseDisableShowPromptsImpl = () => false;
let mockGetLocalStorageItemImpl = (key: string): string | null => {
@ -143,6 +155,44 @@ describe("UserDropdown", () => {
expect(mockOnLogout).toHaveBeenCalledTimes(1);
});
it("should navigate to the change-password page for username/password sessions", async () => {
mockUseAuthorizedImpl = () => ({
userId: "test-user-id",
userEmail: "test@example.com",
userRoleLabel: "Admin",
premiumUser: false,
loginMethod: "username_password",
});
const user = userEvent.setup();
renderWithProviders(<UserDropdown onLogout={mockOnLogout} />);
await user.click(getAccountTrigger());
await user.click(await screen.findByText("Change Password"));
expect(mockRouterPush).toHaveBeenCalledWith(expect.stringContaining("change-password"));
});
it("should hide the change-password entry for SSO sessions", async () => {
mockUseAuthorizedImpl = () => ({
userId: "test-user-id",
userEmail: "test@example.com",
userRoleLabel: "Admin",
premiumUser: false,
loginMethod: "sso",
});
const user = userEvent.setup();
renderWithProviders(<UserDropdown onLogout={mockOnLogout} />);
await user.click(getAccountTrigger());
await waitFor(() => {
expect(screen.getAllByText("test@example.com").length).toBeGreaterThan(0);
});
expect(screen.queryByText("Change Password")).not.toBeInTheDocument();
});
it("should toggle hide new feature indicators switch", async () => {
const user = userEvent.setup();
renderWithProviders(<UserDropdown onLogout={mockOnLogout} />);

23
ui/litellm-dashboard/src/components/Navbar/UserDropdown/UserDropdown.tsx
View file

@ -9,7 +9,9 @@ import {
setLocalStorageItem,
} from "@/utils/localStorageUtils";
import { navAccountDisplayName } from "@/components/Navbar/navDisplayName";
import { ChevronDown, ChevronsUpDown, Crown, LogOut, Mail, ShieldCheck, User } from "lucide-react";
import { uiHref } from "@/utils/uiHref";
import { ChevronDown, ChevronsUpDown, Crown, KeyRound, LogOut, Mail, ShieldCheck, User } from "lucide-react";
import { useRouter } from "next/navigation";
import { Avatar, AvatarFallback } from "@/components/ui/avatar";
import { Badge } from "@/components/ui/badge";
import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
@ -63,7 +65,9 @@ interface UserDropdownProps {
}
const UserDropdown: React.FC<UserDropdownProps> = ({ onLogout, variant = "navbar", collapsed = false }) => {
const { userId, userEmail, userRoleLabel: userRole, premiumUser } = useAuthorized();
const { userId, userEmail, userRoleLabel: userRole, premiumUser, loginMethod } = useAuthorized();
const router = useRouter();
const [open, setOpen] = useState(false);
const disableShowPrompts = useDisableShowPrompts();
const disableBlogPosts = useDisableBlogPosts();
const disableBouncingIcon = useDisableBouncingIcon();
@ -197,7 +201,7 @@ const UserDropdown: React.FC<UserDropdownProps> = ({ onLogout, variant = "navbar
const displayName = navAccountDisplayName(userEmail, userId);
return (
<Popover>
<Popover open={open} onOpenChange={setOpen}>
{variant === "sidebar" ? (
<PopoverTrigger
render={
@ -258,6 +262,19 @@ const UserDropdown: React.FC<UserDropdownProps> = ({ onLogout, variant = "navbar
>
{renderUserInfoSection()}
<Separator />
{loginMethod === "username_password" && (
<button
type="button"
onClick={() => {
setOpen(false);
router.push(uiHref("change-password"));
}}
className="flex w-full items-center gap-2 rounded-sm px-2 py-1.5 text-sm hover:bg-accent"
>
<KeyRound className="size-4" />
Change Password
</button>
)}
<button
type="button"
onClick={onLogout}

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