diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 6f8599daf75..4bf41dc249c 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -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 diff --git a/docker/README.md b/docker/README.md index 26d8c9a37b0..376dc7b2d97 100644 --- a/docker/README.md +++ b/docker/README.md @@ -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 diff --git a/docker/docker-compose.quickstart.yml b/docker/docker-compose.quickstart.yml new file mode 100644 index 00000000000..11631603a72 --- /dev/null +++ b/docker/docker-compose.quickstart.yml @@ -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: diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260921000000_add_password_reset_columns/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260921000000_add_password_reset_columns/migration.sql new file mode 100644 index 00000000000..960b0d4d7eb --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260921000000_add_password_reset_columns/migration.sql @@ -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); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 2d7e557a9d1..368864f9cd1 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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? diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b37202dce1b..de234fa3afa 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -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]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 7d94d31db3e..be86240ac46 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index 8aeddae9efc..534d85acdb0 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -128,6 +128,14 @@ impl VertexAuth { } } + pub async fn access_token( + &self, + config: &VertexConfig, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result { + self.load_provider(config, env_lookup).await?.token().await + } + pub async fn validate_environment( &self, headers: Vec<(String, String)>, diff --git a/litellm-rust/crates/cache-disk/Cargo.toml b/litellm-rust/crates/cache-disk/Cargo.toml new file mode 100644 index 00000000000..b96994b3b55 --- /dev/null +++ b/litellm-rust/crates/cache-disk/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/cache-disk/src/adapter.rs b/litellm-rust/crates/cache-disk/src/adapter.rs new file mode 100644 index 00000000000..b6d318d5509 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/adapter.rs @@ -0,0 +1,10 @@ +use litellm_cache::Error; + +use crate::StoredValue; + +pub trait ValueAdapter: Send + Sync + 'static { + fn read(&self, value: StoredValue) -> Result>, Error>; + fn write(&self, payload: Vec) -> StoredValue; + fn counter_seed(&self, value: Option) -> Result; + fn counter_value(&self, value: f64) -> StoredValue; +} diff --git a/litellm-rust/crates/cache-disk/src/cache.rs b/litellm-rust/crates/cache-disk/src/cache.rs new file mode 100644 index 00000000000..8e1223309b4 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/cache.rs @@ -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 { + store: Arc, + adapter: Arc, + codec: S, +} + +impl DiskCache { + pub fn open(directory: impl AsRef, codec: S) -> Result { + Ok(Self { + store: Arc::new(DiskcacheSqliteStore::open(directory)?), + adapter: Arc::new(PythonDiskCacheAdapter), + codec, + }) + } +} + +impl DiskCache { + pub fn with_store(store: D, codec: S) -> Self { + Self { + store: Arc::new(store), + adapter: Arc::new(PythonDiskCacheAdapter), + codec, + } + } +} + +impl DiskCache { + 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, Error> { + let Some(bytes) = self.adapter.read(value)? else { + return Ok(None); + }; + self.codec.decode(&bytes).map(Some) + } + + async fn run_blocking(store: Arc, operation: F) -> Result + where + T: Send + 'static, + F: FnOnce(&D) -> Result + Send + 'static, + { + tokio::task::spawn_blocking(move || operation(&store)) + .await + .map_err(|_| Error::Unavailable)? + } +} + +impl BaseCache for DiskCache { + type Value = S::Value; + type Context = ExactCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + 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, 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, 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::, _>>()?; + 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 { + 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 BatchCache for DiskCache { + fn batch_get_cache( + &self, + keys: &[String], + context: &ExactCacheContext, + ) -> Result>, 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, + _: ExactCacheContext, + ) -> Result>, 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::, _>>() + }) + .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 DeleteCache for DiskCache { + 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 FlushCache for DiskCache { + 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, D: DiskStore, A: ValueAdapter> CounterCache + for DiskCache +{ + fn increment_cache( + &self, + key: &str, + amount: f64, + context: ExactCacheContext, + ) -> Result { + 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 { + 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( + adapter: &A, + store: &D, + key: &str, + amount: f64, + ttl: Option, +) -> Result { + let mut result = None; + let mut apply = |current: Option| { + 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() +} diff --git a/litellm-rust/crates/cache-disk/src/lib.rs b/litellm-rust/crates/cache-disk/src/lib.rs new file mode 100644 index 00000000000..9b2ffc24915 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/lib.rs @@ -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}; diff --git a/litellm-rust/crates/cache-disk/src/python/mod.rs b/litellm-rust/crates/cache-disk/src/python/mod.rs new file mode 100644 index 00000000000..7a370db357c --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/python/mod.rs @@ -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, 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>, 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) -> StoredValue { + StoredValue::Bytes(payload) + } + + fn counter_seed(&self, value: Option) -> Result { + 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) + } + } +} diff --git a/litellm-rust/crates/cache-disk/src/python/value.rs b/litellm-rust/crates/cache-disk/src/python/value.rs new file mode 100644 index 00000000000..645eba7b757 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/python/value.rs @@ -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 { + 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 { + 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::, _>>() + .map(Value::List), + serde_pickle::Value::Tuple(values) => values + .into_iter() + .map(from_pickle_value) + .collect::, _>>() + .map(Value::Tuple), + serde_pickle::Value::Set(values) => values + .into_iter() + .map(from_pickle_hashable) + .collect::, _>>() + .map(Value::Set), + serde_pickle::Value::FrozenSet(values) => values + .into_iter() + .map(from_pickle_hashable) + .collect::, _>>() + .map(Value::Set), + serde_pickle::Value::Dict(values) => values + .into_iter() + .map(|(key, value)| Ok((from_pickle_hashable(key)?, from_pickle_value(value)?))) + .collect::, Error>>() + .map(Value::Dict), + } +} + +fn from_pickle_hashable(value: serde_pickle::HashableValue) -> Result { + 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::, _>>()?, + ), + serde_pickle::HashableValue::FrozenSet(values) => Value::Set( + values + .into_iter() + .map(from_pickle_hashable) + .collect::, _>>()?, + ), + }) +} + +fn integer(value: String) -> Result { + 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 { + 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 { + 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, Error> { + serde_json::to_vec(&to_json_value(value)?).map_err(|_| Error::InvalidEntry) +} + +fn to_json_value(value: &Value) -> Result { + 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::() + .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::, _>>()?, + ) + } + 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::, _>>()?; + serde_json::Value::Object(values) + } + }) +} diff --git a/litellm-rust/crates/cache-disk/src/sqlite.rs b/litellm-rust/crates/cache-disk/src/sqlite.rs new file mode 100644 index 00000000000..9a36f8af6ad --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/sqlite.rs @@ -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, + min_file_size: usize, + eviction_policy: String, + size_limit: i64, + cull_limit: i64, + statistics: bool, +} + +struct StoredColumns { + size: i64, + mode: i64, + filename: Option, + value: Option, +} + +struct Row { + rowid: i64, + mode: i64, + filename: Option, + value: Value, +} + +impl DiskcacheSqliteStore { + pub fn open(directory: impl AsRef) -> Result { + 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, + now: f64, + ) -> Result, 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>(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, 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>(1)?)) + }) + .map_err(|_| Error::Unavailable)? + .collect::, _>>() + .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>(1)?)) + }) + .map_err(|_| Error::Unavailable)? + .collect::, _>>() + .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 { + 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, 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, + 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, 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>(1)?)) + }) + .map_err(|_| Error::Unavailable)? + .collect::, _>>() + .map_err(|_| Error::Unavailable)?; + if rows.is_empty() { + return Ok(rows); + } + let ids = rows + .iter() + .map(|(rowid, _)| rowid.to_string()) + .collect::>() + .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) -> Result<(StoredValue, Option), 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 { + 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, 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::, _>>() + .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, key: &str) -> Option { + match settings.get(key) { + Some(Value::Integer(value)) => Some(*value), + _ => None, + } +} + +fn setting_string(settings: &HashMap, key: &str) -> Option { + 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 { + 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, 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>, 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 { + 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 { + let mut random = [0_u8; 16]; + rand::rngs::OsRng.fill_bytes(&mut random); + let hex = random + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + 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) { + 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( + connection: &Connection, + operation: impl FnOnce(&Connection) -> Result, +) -> Result { + 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) + } + } +} diff --git a/litellm-rust/crates/cache-disk/src/store.rs b/litellm-rust/crates/cache-disk/src/store.rs new file mode 100644 index 00000000000..ed167317cf0 --- /dev/null +++ b/litellm-rust/crates/cache-disk/src/store.rs @@ -0,0 +1,33 @@ +use std::path::Path; + +use litellm_cache::Error; + +#[derive(Clone, Debug, PartialEq)] +pub enum StoredValue { + Bytes(Vec), + Text(String), + Integer(i64), + Float(f64), + Pickle(Vec), +} + +pub trait DiskStore: Send + Sync + 'static { + fn directory(&self) -> &Path; + fn get(&self, key: &str, now: f64) -> Result, Error>; + fn set( + &self, + key: &str, + value: StoredValue, + expire_time: Option, + now: f64, + ) -> Result<(), Error>; + fn pop(&self, key: &str, now: f64) -> Result, Error>; + fn clear(&self) -> Result<(), Error>; + fn update( + &self, + key: &str, + now: f64, + apply: &mut dyn FnMut(Option) -> Result<(StoredValue, Option), Error>, + ) -> Result<(), Error>; + fn probe(&self) -> Result<(), Error>; +} diff --git a/litellm-rust/crates/cache-disk/tests/cache.rs b/litellm-rust/crates/cache-disk/tests/cache.rs new file mode 100644 index 00000000000..dd1f2b1f04e --- /dev/null +++ b/litellm-rust/crates/cache-disk/tests/cache.rs @@ -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(&self) -> DiskCache> + where + JsonCodec: 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 { + fn visit(directory: &Path, files: &mut Vec) { + 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>, litellm_cache::Error> { + match value { + StoredValue::Text(value) => Ok(Some(value.into_bytes())), + _ => Ok(None), + } + } + + fn write(&self, payload: Vec) -> StoredValue { + StoredValue::Text(String::from_utf8(payload).unwrap()) + } + + fn counter_seed(&self, _: Option) -> Result { + 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::(); + 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::(); + 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::() + .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::() + .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, + #[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::(); + 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::()); + 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::>(); + 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::(); + 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::(); + 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::::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::(); + 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 + ); +} diff --git a/litellm-rust/crates/cache-disk/tests/python_compat.rs b/litellm-rust/crates/cache-disk/tests/python_compat.rs new file mode 100644 index 00000000000..9cbef8573bd --- /dev/null +++ b/litellm-rust/crates/cache-disk/tests/python_compat.rs @@ -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, #[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); +} diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml new file mode 100644 index 00000000000..4ec60bcfa3b --- /dev/null +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs new file mode 100644 index 00000000000..65282ac99d5 --- /dev/null +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -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, + pub path_service_account: Option, + pub endpoint: String, +} + +impl GcsConfig { + pub fn new(bucket_name: impl Into) -> Self { + Self { + bucket_name: bucket_name.into(), + gcs_path: None, + path_service_account: None, + endpoint: DEFAULT_ENDPOINT.to_string(), + } + } +} + +pub struct GcsCache { + config: GcsConfig, + key_prefix: String, + client: Client, + token: Arc, + codec: S, +} + +impl GcsCache { + pub fn new(config: GcsConfig, codec: S) -> Result { + 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, + ) -> Result { + 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, 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(future: F) -> Result + where + F: Future> + 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 BaseCache for GcsCache { + type Value = S::Value; + type Context = ExactCacheContext; + + fn get_ttl(&self, _: &Self::Context) -> Option { + 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, 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, 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 { + Err(Error::UnsupportedOperation) + } +} + +impl BatchCache for GcsCache { + async fn async_batch_get_cache( + &self, + keys: Vec, + context: Self::Context, + ) -> Result>, 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 FlushCache for GcsCache { + fn flush_cache(&self) -> Result<(), Error> { + Ok(()) + } +} diff --git a/litellm-rust/crates/cache-gcs/src/lib.rs b/litellm-rust/crates/cache-gcs/src/lib.rs new file mode 100644 index 00000000000..cbb61cf0685 --- /dev/null +++ b/litellm-rust/crates/cache-gcs/src/lib.rs @@ -0,0 +1,5 @@ +mod cache; +mod token; + +pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix}; +pub use token::{GcpTokenSource, StaticTokenSource, TokenSource}; diff --git a/litellm-rust/crates/cache-gcs/src/token.rs b/litellm-rust/crates/cache-gcs/src/token.rs new file mode 100644 index 00000000000..adb601c276c --- /dev/null +++ b/litellm-rust/crates/cache-gcs/src/token.rs @@ -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> + Send + '_>>; +} + +pub struct GcpTokenSource { + auth: VertexAuth, + config: VertexConfig, +} + +impl GcpTokenSource { + pub fn new(path_service_account: Option) -> 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> + 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> + Send + '_>> { + Box::pin(async move { Ok(self.0.clone()) }) + } +} diff --git a/litellm-rust/crates/cache-gcs/tests/cache.rs b/litellm-rust/crates/cache-gcs/tests/cache.rs new file mode 100644 index 00000000000..45eecf01cec --- /dev/null +++ b/litellm-rust/crates/cache-gcs/tests/cache.rs @@ -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> { + 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> + 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::::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); +} diff --git a/litellm-rust/crates/cache-s3/Cargo.toml b/litellm-rust/crates/cache-s3/Cargo.toml new file mode 100644 index 00000000000..cdc17e732cb --- /dev/null +++ b/litellm-rust/crates/cache-s3/Cargo.toml @@ -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"] } diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/cache-s3/src/auth.rs new file mode 100644 index 00000000000..b7ca722cea3 --- /dev/null +++ b/litellm-rust/crates/cache-s3/src/auth.rs @@ -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, +} + +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) -> 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")); + } +} diff --git a/litellm-rust/crates/cache-s3/src/cache.rs b/litellm-rust/crates/cache-s3/src/cache.rs new file mode 100644 index 00000000000..9c791f42c3b --- /dev/null +++ b/litellm-rust/crates/cache-s3/src/cache.rs @@ -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, + pub auth: AwsAuthConfig, +} + +pub struct S3Cache { + client: aws_sdk_s3::Client, + codec: C, + runtime: Handle, + bucket: Arc, + key_prefix: Arc, + region: Arc, + endpoint: Option>, +} + +impl S3Cache { + pub fn new(config: S3CacheConfig, codec: C, runtime: Handle) -> Self { + let endpoint_url: Option = 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(&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, 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 BaseCache for S3Cache { + type Value = C::Value; + type Context = ExactCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + 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, 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, Error> { + self.get(key).await + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Err(Error::UnsupportedOperation) + } +} + +impl BatchCache for S3Cache {} + +impl FlushCache for S3Cache { + fn flush_cache(&self) -> Result<(), Error> { + Ok(()) + } +} diff --git a/litellm-rust/crates/cache-s3/src/lib.rs b/litellm-rust/crates/cache-s3/src/lib.rs new file mode 100644 index 00000000000..f6126dfa908 --- /dev/null +++ b/litellm-rust/crates/cache-s3/src/lib.rs @@ -0,0 +1,4 @@ +mod auth; +mod cache; + +pub use cache::{S3Cache, S3CacheConfig, S3Endpoint}; diff --git a/litellm-rust/crates/cache-s3/tests/cache.rs b/litellm-rust/crates/cache-s3/tests/cache.rs new file mode 100644 index 00000000000..9a71656286b --- /dev/null +++ b/litellm-rust/crates/cache-s3/tests/cache.rs @@ -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> { + S3Cache::new( + config(endpoint.to_string()), + JsonCodec::::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 { + 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::(&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("NoSuchKey"), + ) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/cache-bucket/team/denied")) + .respond_with( + ResponseTemplate::new(403).set_body_string("AccessDenied"), + ) + .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::::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::::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}))); +} diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index ff3ff6572d4..1a381d0afd8 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -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, } diff --git a/litellm-rust/crates/framer/Cargo.toml b/litellm-rust/crates/framer/Cargo.toml index d22502f871a..62bfcc7da3d 100644 --- a/litellm-rust/crates/framer/Cargo.toml +++ b/litellm-rust/crates/framer/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 0cc7af1836f..f04b78feee1 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 5528e12ee2d..9191e61a25c 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 5c962f2bc7a..7316b3e259c 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -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, +} + 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), + S3(Box), + 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 { }) } +#[inline(never)] +fn project_gcs( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let bucket_name = match backend.getattr("bucket_name")?.extract::>() { + 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::()?, + path_service_account: backend + .getattr("path_service_account")? + .extract::>()?, + })) +} + +#[inline(never)] +fn project_disk( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + 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::()?), + })) +} + #[inline(never)] fn project_redis( backend: &Bound<'_, PyAny>, @@ -346,6 +483,77 @@ fn project_redis( })) } +#[inline(never)] +fn project_s3( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + 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::>()?, + 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::().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::()?.as_str() == "explicit" { + AwsAuthConfig { + access_key_id: credentials + .getattr("access_key")? + .extract::>()?, + secret_access_key: credentials + .getattr("secret_key")? + .extract::>()?, + session_token: credentials.getattr("token")?.extract::>()?, + 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::()?, + key_prefix: backend.getattr("key_prefix")?.extract::()?, + 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( + value: &Bound<'py, PyAny>, + name: &str, +) -> PyResult>> { + match value.getattr(name) { + Ok(value) => Ok(Some(value)), + Err(error) if error.is_instance_of::(value.py()) => Ok(None), + Err(error) => Err(error), + } +} + +#[inline(never)] +fn optional_attribute_chain<'py>( + value: &Bound<'py, PyAny>, + names: &[&str], +) -> PyResult>> { + names + .iter() + .try_fold(Some(value.clone()), |current, name| match current { + Some(current) => optional_attribute(¤t, name), + None => Ok(None), + }) +} + #[inline(never)] fn optional_string(value: Bound<'_, PyAny>) -> PyResult> { 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() { diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 0e20676938f..e89d0fc9715 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -34,6 +34,14 @@ struct RedisPoolGuard { attributes: RedisPoolAttributes, } +struct S3ClientGuard { + reference: Py, +} + +struct DiskStoreGuard { + reference: Py, + directory: String, +} struct AzureBlobClientGuard { sync_client: Py, async_client: Py, @@ -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, connection: ConnectionGuard, } @@ -218,6 +226,41 @@ impl RedisPoolGuard { } } +impl S3ClientGuard { + fn capture(backend: &Bound<'_, PyAny>) -> PyResult { + Ok(Self { + reference: backend.getattr("s3_client")?.unbind(), + }) + } + + fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { + 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 { + 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 { + let store = backend.getattr("disk_cache")?; + Ok(self.reference.bind(py).is(&store) + && self.directory == store.getattr("directory")?.extract::()?) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.reference) + } +} impl AzureBlobClientGuard { fn capture(backend: &Bound<'_, PyAny>) -> PyResult { 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) } } diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 002cd6b7a33..d6e622a5c49 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -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, + key_prefix: &str, + access_key_id: Option, + secret_access_key: Option, + session_token: Option, + ) -> PyResult { + 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, + path_service_account: Option, + endpoint: Option, + token: Option, + ) -> PyResult { + 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 { + 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 { @@ -78,7 +153,6 @@ impl CacheTestHandle { pid: std::process::id(), }) } - #[getter] fn backend(&self) -> &'static str { self.service.kind() diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index aec08610f6e..ac4494150d9 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -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()), } } diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 2cd72a7ce14..39496890009 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -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>>, buffer: Option>, }, + S3(Arc>>), + Gcs(Arc>>), + Disk(Arc>>), AzureBlob(Arc>>), } @@ -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 { + let cache = DiskCache::open(directory, ResponseCacheCodec)?; + Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache))))) + } + + pub fn gcs(config: GcsConfig, token: Option) -> Result { + 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 { 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 { 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 { 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> { + match self { + Self::Gcs(cache) => Some(cache.backend()), + _ => None, + } + } } diff --git a/litellm-rust/crates/secrets-hashicorp/Cargo.toml b/litellm-rust/crates/secrets-hashicorp/Cargo.toml new file mode 100644 index 00000000000..c646e02ef09 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/secrets-hashicorp/src/cert_login.rs b/litellm-rust/crates/secrets-hashicorp/src/cert_login.rs new file mode 100644 index 00000000000..f99df62db12 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/cert_login.rs @@ -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, +} + +impl CertLoginRequest { + pub fn new(name: Option<&str>) -> Self { + let body: Vec = 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, + } + } +} diff --git a/litellm-rust/crates/secrets-hashicorp/src/config.rs b/litellm-rust/crates/secrets-hashicorp/src/config.rs new file mode 100644 index 00000000000..d32491199c4 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/config.rs @@ -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, +} + +#[derive(Clone, Debug)] +pub struct HashicorpVaultConfig { + pub address: String, + pub token: Option, + pub namespace: Option, + pub login_namespace: Option, + pub secret_namespace: Option, + pub mount: String, + pub path_prefix: Option, + pub approle: Option, + pub tls_cert: Option, + pub refresh_interval: Duration, +} + +impl HashicorpVaultConfig { + pub fn from_environment(environment: &dyn Lookup) -> Result { + 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 = environment + .get(HCP_VAULT_TOKEN) + .and_then(nonempty) + .map(SecretValue::new); + let namespace: Option = path_component(environment.get(HCP_VAULT_NAMESPACE)); + let login_namespace: Option = + path_component(environment.get(HCP_VAULT_LOGIN_NAMESPACE)); + let secret_namespace: Option = + 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 = path_component(environment.get(HCP_VAULT_PATH_PREFIX)); + let approle: Option = 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 = 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) -> Option { + let value: &str = value.as_ref(); + (!value.is_empty()).then(|| value.to_owned()) +} + +fn path_component(value: Option) -> Option { + 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 { + let value: Option = 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)) +} diff --git a/litellm-rust/crates/secrets-hashicorp/src/error.rs b/litellm-rust/crates/secrets-hashicorp/src/error.rs new file mode 100644 index 00000000000..e26033af085 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/error.rs @@ -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, +} diff --git a/litellm-rust/crates/secrets-hashicorp/src/lib.rs b/litellm-rust/crates/secrets-hashicorp/src/lib.rs new file mode 100644 index 00000000000..0c2b05647f8 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/lib.rs @@ -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}; diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs new file mode 100644 index 00000000000..3ad9a549438 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -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, + expires_at: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SecretLocation { + pub namespace: Option, + pub mount: String, + pub path: String, +} + +#[derive(Clone)] +pub struct HashicorpVault { + config: HashicorpVaultConfig, + cache: Cache, + auth_client: Arc>>, +} + +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, + enterprise_enabled: bool, + ) -> Result { + 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 { + if !enterprise_enabled { + return Err(Error::EnterpriseRequired); + } + let cache: Cache = 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 { + 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::>() + .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, 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 = self.vault_client().await?; + let data: HashMap = + 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 { + let location: SecretLocation = self.secret_location(secret_name)?; + let cache_key: String = cache_key(&location); + let data: HashMap = 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 = 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 = 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 { + async_rotate_secret(self, current_name, new_name, value).await + } + + async fn vault_client(&self) -> Result, 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) = + 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 = Arc::new(client); + *cached = Some(CachedClient { + client: client.clone(), + expires_at, + }); + Ok(client) + } + + fn build_client(&self, namespace: Option<&str>, token: &str) -> Result { + 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, Error> { + HashicorpVault::async_read_secret(self, name).await + } + + async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result { + 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, Error> { + tls.map(|tls| { + let cert: Vec = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity { + path: tls.cert_path.clone(), + message: source.to_string(), + })?; + let key: Vec = 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 { + 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 { + (lease_duration > 0).then(|| Instant::now() + Duration::from_secs(lease_duration)) +} diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs new file mode 100644 index 00000000000..c52db46e41e --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs @@ -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 = 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 = + 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 = 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 = 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 = 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 = + 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 = 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 = + 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 = 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, 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, + expected_secret_url: String, + expected_login_url: Option, + expected_login_namespace: Option, + expected_secret_namespace: Option, + secret_name: String, +} + +#[test] +fn configuration_matches_python_parity_fixture() { + let cases: Vec = 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 = case.env.clone(); + let environment: Arc = + 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 = + 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----- +"; diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index 962f66c92d1..e5f30025976 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/secrets/README.md b/litellm-rust/crates/secrets/README.md index 183a39e15bb..0d333c7d116 100644 --- a/litellm-rust/crates/secrets/README.md +++ b/litellm-rust/crates/secrets/README.md @@ -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 diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs index 7e03f1f8cbf..1be0adc2bf5 100644 --- a/litellm-rust/crates/secrets/src/error.rs +++ b/litellm-rust/crates/secrets/src/error.rs @@ -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), diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index 71214ffc9ce..5ab2caa75d4 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -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) diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs index dec924abdd0..1acb5269e66 100644 --- a/litellm-rust/crates/secrets/src/lib.rs +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -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; diff --git a/litellm-rust/crates/secrets/tests/handler.rs b/litellm-rust/crates/secrets/tests/handler.rs index fbe13721491..2a8b7070522 100644 --- a/litellm-rust/crates/secrets/tests/handler.rs +++ b/litellm-rust/crates/secrets/tests/handler.rs @@ -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 = 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 = 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 = 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] diff --git a/litellm/constants.py b/litellm/constants.py index 88cc5b04743..5e86927ef56 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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" diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 91bf697487d..33ee727dfab 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -118,6 +118,7 @@ class XAIChatConfig(OpenAIGPTConfig): base_openai_params: Final = [ "logit_bias", "logprobs", + "max_completion_tokens", "max_tokens", "n", "parallel_tool_calls", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index beb90c2fb8c..35c61871e09 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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 } } diff --git a/litellm/models/user.py b/litellm/models/user.py index 82f78c28078..92aca87d303 100644 --- a/litellm/models/user.py +++ b/litellm/models/user.py @@ -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 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6bf846cd2de..4b3c12b3cb4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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": }` shape most proxy endpoints raise as `HTTPException.detail`.""" + + error: ReadOnly[str] + + class SpendLogsRouterMetadata(TypedDict): """ Router provenance stamped on spend logs for deployments flagged with diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index e0d599b0017..4c2b5d3d0fe 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -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, ) diff --git a/litellm/proxy/auth/password_policy.py b/litellm/proxy/auth/password_policy.py index ab7a565894a..7f06a0993d3 100644 --- a/litellm/proxy/auth/password_policy.py +++ b/litellm/proxy/auth/password_policy.py @@ -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 + } + ) diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 1b9fd7c42bf..afddb4866a9 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -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). diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index c6b89096ca9..60ac7e55eaf 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -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"}) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index c1388e8bb81..ee0e4db291b 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -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) diff --git a/litellm/proxy/management_endpoints/password_endpoints.py b/litellm/proxy/management_endpoints/password_endpoints.py new file mode 100644 index 00000000000..03a8b4c4010 --- /dev/null +++ b/litellm/proxy/management_endpoints/password_endpoints.py @@ -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.") diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 00cf357d89d..7859c678c07 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dac1a8dd001..8111e21e114 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 2d7e557a9d1..368864f9cd1 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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? diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 05a6df6d5af..90d876b71e6 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -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: ... diff --git a/litellm/types/llms/custom_http.py b/litellm/types/llms/custom_http.py index d80d7410aae..793893451df 100644 --- a/litellm/types/llms/custom_http.py +++ b/litellm/types/llms/custom_http.py @@ -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 diff --git a/litellm/types/proxy/ui_sso.py b/litellm/types/proxy/ui_sso.py index 0d7e0b99cf0..03b0b92a4d1 100644 --- a/litellm/types/proxy/ui_sso.py +++ b/litellm/types/proxy/ui_sso.py @@ -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): diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index beb90c2fb8c..35c61871e09 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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 } } diff --git a/schema.prisma b/schema.prisma index 2d7e557a9d1..368864f9cd1 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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? diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index 4ea64b152f1..f8277c83a64 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -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} diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 49d4d92ff0b..64335aa560c 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -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"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index 8b0d38a083c..d9c20d5c588 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -58,6 +58,7 @@ LlmRoute = Literal[ "openai", "together_ai", "vertex", + "xiaomi_mimo", ] LlmCapability = Literal[ diff --git a/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py b/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py new file mode 100644 index 00000000000..efca216634b --- /dev/null +++ b/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py @@ -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}" diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py index 290cd3dcb3a..3fd666e4f50 100644 --- a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py @@ -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.""" diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index 55ece36252d..3786169c320 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -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 + ) diff --git a/tests/test_litellm/proxy/auth/test_onboarding.py b/tests/test_litellm/proxy/auth/test_onboarding.py index 524b655b465..0454aea1239 100644 --- a/tests/test_litellm/proxy/auth/test_onboarding.py +++ b/tests/test_litellm/proxy/auth/test_onboarding.py @@ -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() diff --git a/tests/test_litellm/proxy/auth/test_password_policy.py b/tests/test_litellm/proxy/auth/test_password_policy.py index f6e7d443907..f9f5025b57f 100644 --- a/tests/test_litellm/proxy/auth/test_password_policy.py +++ b/tests/test_litellm/proxy/auth/test_password_policy.py @@ -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 == {} diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 87bf4595af5..c9b8d6764c9 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -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", diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index ede0ecb2790..75545a574e3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -1,13 +1,18 @@ +import hashlib import json from datetime import datetime, timezone from types import SimpleNamespace from typing import Final +import httpx import pytest +import respx from fastapi import HTTPException from fastapi.testclient import TestClient from pytest_mock import MockerFixture +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.proxy._types import ( LiteLLM_UserTableFiltered, LitellmUserRoles, @@ -67,9 +72,7 @@ async def test_ui_view_users_with_null_email(mocker, caplog): # Proxy admin: no org filter, no get_user_object call response = await ui_view_users( - user_api_key_dict=UserAPIKeyAuth( - user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN), user_id="test_user", user_email=None, team_id=None, @@ -77,9 +80,7 @@ async def test_ui_view_users_with_null_email(mocker, caplog): page_size=50, ) - assert response == [ - LiteLLM_UserTableFiltered(user_id="test-user-null-email", user_email=None) - ] + assert response == [LiteLLM_UserTableFiltered(user_id="test-user-null-email", user_email=None)] @pytest.mark.asyncio @@ -103,9 +104,7 @@ async def test_ui_view_users_proxy_admin_no_org_filter(mocker): mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) await ui_view_users( - user_api_key_dict=UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), user_id=None, user_email="foo", team_id=None, @@ -128,9 +127,7 @@ async def test_ui_view_users_org_admin_filtered_by_org(mocker): async def mock_find_many(*args, **kwargs): where = kwargs.get("where") or {} assert "organization_memberships" in where - assert where["organization_memberships"] == { - "some": {"organization_id": {"in": [org_id]}} - } + assert where["organization_memberships"] == {"some": {"organization_id": {"in": [org_id]}}} return [] mock_prisma_client.db.litellm_usertable.find_many = mock_find_many @@ -268,9 +265,7 @@ async def test_ui_view_users_flag_on_team_admin_org_team(mocker): async def mock_find_many(*args, **kwargs): where = kwargs.get("where") or {} assert "organization_memberships" in where - assert where["organization_memberships"] == { - "some": {"organization_id": {"in": [org_id]}} - } + assert where["organization_memberships"] == {"some": {"organization_id": {"in": [org_id]}}} return [] mock_prisma_client.db.litellm_usertable.find_many = mock_find_many @@ -401,9 +396,7 @@ async def test_ui_view_users_flag_on_team_admin_org_member_no_team_id(mocker): async def mock_find_many(*args, **kwargs): where = kwargs.get("where") or {} assert "organization_memberships" in where - assert where["organization_memberships"] == { - "some": {"organization_id": {"in": [org_id]}} - } + assert where["organization_memberships"] == {"some": {"organization_id": {"in": [org_id]}}} return [] mock_prisma_client.db.litellm_usertable.find_many = mock_find_many @@ -462,9 +455,7 @@ async def test_ui_view_users_flag_on_team_admin_not_in_org_resolves_via_key_team async def mock_find_many(*args, **kwargs): where = kwargs.get("where") or {} assert "organization_memberships" in where - assert where["organization_memberships"] == { - "some": {"organization_id": {"in": [org_id]}} - } + assert where["organization_memberships"] == {"some": {"organization_id": {"in": [org_id]}}} return [] mock_prisma_client.db.litellm_usertable.find_many = mock_find_many @@ -507,9 +498,7 @@ async def test_ui_view_users_flag_on_team_admin_not_in_org_resolves_via_key_team # No team_id query param, but team_id on the API key response = await ui_view_users( - user_api_key_dict=UserAPIKeyAuth( - user_id="team-admin-no-org", user_role=None, team_id=tid - ), + user_api_key_dict=UserAPIKeyAuth(user_id="team-admin-no-org", user_role=None, team_id=tid), user_id=None, user_email="u", team_id=None, @@ -538,13 +527,9 @@ def test_user_daily_activity_types(): # Assert all fields in SpendMetrics are reported in DailySpendMetadata as "total_" for field in spend_metrics.__dict__: if field.startswith("total_"): - assert hasattr( - daily_spend_metadata, field - ), f"Field {field} is not reported in DailySpendMetadata" + assert hasattr(daily_spend_metadata, field), f"Field {field} is not reported in DailySpendMetadata" else: - assert not hasattr( - daily_spend_metadata, field - ), f"Field {field} is reported in DailySpendMetadata" + assert not hasattr(daily_spend_metadata, field), f"Field {field} is reported in DailySpendMetadata" @pytest.mark.asyncio @@ -591,9 +576,7 @@ async def test_get_users_includes_timestamps(mocker): # Call get_users function directly with proxy admin auth admin_key = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) - response = await get_users( - page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None - ) + response = await get_users(page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None) print("user /list response: ", response) @@ -654,14 +637,10 @@ async def test_get_users_redacts_scim_enterprise_metadata(mocker): ) admin_key = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) - response = await get_users( - page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None - ) + response = await get_users(page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None) listed = response["users"][0] - assert listed.metadata == { - "scim_metadata": {"givenName": "Jane", "familyName": "Doe"} - } + assert listed.metadata == {"scim_metadata": {"givenName": "Jane", "familyName": "Doe"}} assert "scim_enterprise" not in (listed.metadata or {}) @@ -853,9 +832,7 @@ async def test_new_user_license_over_limit(mocker): mocker.patch("litellm.proxy.proxy_server._license_check", mock_license_check) # Create test request data - user_request = NewUserRequest( - user_email="test@example.com", user_role="internal_user" - ) + user_request = NewUserRequest(user_email="test@example.com", user_role="internal_user") # Mock user_api_key_dict mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_admin") @@ -916,9 +893,7 @@ async def test_new_user_license_gate_counts_only_billable_users(mocker): request = NewUserRequest(user_role="internal_user") # 2 active + 3 deactivated -> billable 2, not over max_users 2: gate passes - mocker.patch( - "litellm.proxy.proxy_server.prisma_client", _prisma(total=5, deactivated=3) - ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", _prisma(total=5, deactivated=3)) with pytest.raises(ProxyException) as passed: await new_user(data=request, user_api_key_dict=admin) assert key_gen.call_count == 1 @@ -926,9 +901,7 @@ async def test_new_user_license_gate_counts_only_billable_users(mocker): # 3 active, 0 deactivated -> billable 3, over max_users 2: gate blocks key_gen.reset_mock() - mocker.patch( - "litellm.proxy.proxy_server.prisma_client", _prisma(total=3, deactivated=0) - ) + mocker.patch("litellm.proxy.proxy_server.prisma_client", _prisma(total=3, deactivated=0)) with pytest.raises(ProxyException) as blocked: await new_user(data=request, user_api_key_dict=admin) assert blocked.value.code == 403 or blocked.value.code == "403" @@ -978,14 +951,10 @@ async def test_new_user_non_admin_cannot_create_admin(mocker): mocker.patch("litellm.proxy.proxy_server._license_check", mock_license_check) # Test Case 1: INTERNAL_USER trying to create PROXY_ADMIN - user_request = NewUserRequest( - user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN - ) + user_request = NewUserRequest(user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN) # Mock user_api_key_dict with non-admin role - mock_user_api_key_dict = UserAPIKeyAuth( - user_id="test_internal_user", user_role=LitellmUserRoles.INTERNAL_USER - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_internal_user", user_role=LitellmUserRoles.INTERNAL_USER) # Call new_user function and expect ProxyException with pytest.raises(ProxyException) as exc_info: @@ -993,9 +962,7 @@ async def test_new_user_non_admin_cannot_create_admin(mocker): # Verify the exception details assert exc_info.value.code == 403 or exc_info.value.code == "403" - assert "Only proxy admins can create administrative users" in str( - exc_info.value.message - ) + assert "Only proxy admins can create administrative users" in str(exc_info.value.message) assert "proxy_admin" in str(exc_info.value.message) assert "proxy_admin_viewer" in str(exc_info.value.message) assert str(LitellmUserRoles.PROXY_ADMIN) in str(exc_info.value.message) @@ -1008,15 +975,11 @@ async def test_new_user_non_admin_cannot_create_admin(mocker): ) with pytest.raises(ProxyException) as exc_info2: - await new_user( - data=user_request_viewer, user_api_key_dict=mock_user_api_key_dict - ) + await new_user(data=user_request_viewer, user_api_key_dict=mock_user_api_key_dict) # Verify the exception details assert exc_info2.value.code == 403 or exc_info2.value.code == "403" - assert "Only proxy admins can create administrative users" in str( - exc_info2.value.message - ) + assert "Only proxy admins can create administrative users" in str(exc_info2.value.message) assert str(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) in str(exc_info2.value.message) @@ -1055,9 +1018,7 @@ async def test_new_user_non_admin_permissions_non_empty_rejected(mocker): user_role=LitellmUserRoles.INTERNAL_USER, permissions={"get_spend_routes": True}, ) - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(ProxyException) as exc_info: await new_user(data=data, user_api_key_dict=caller) @@ -1101,9 +1062,7 @@ async def test_new_user_non_admin_permissions_explicit_empty_rejected(mocker): permissions={}, ) assert "permissions" in data.model_fields_set - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(ProxyException) as exc_info: await new_user(data=data, user_api_key_dict=caller) @@ -1156,9 +1115,7 @@ async def test_new_user_non_admin_omits_permissions_succeeds(mocker): user_role=LitellmUserRoles.INTERNAL_USER, ) assert "permissions" not in data.model_fields_set - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) result = await new_user(data=data, user_api_key_dict=caller) assert result is not None @@ -1232,14 +1189,10 @@ async def test_update_single_user_non_admin_permissions_rejected(mocker): user_id="alice", permissions={"get_spend_routes": True}, ) - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(HTTPException) as exc_info: - await _update_single_user_helper( - user_request=data, user_api_key_dict=caller - ) + await _update_single_user_helper(user_request=data, user_api_key_dict=caller) assert exc_info.value.status_code == 403 assert "permissions" in str(exc_info.value.detail) @@ -1261,14 +1214,10 @@ async def test_update_single_user_non_admin_permissions_explicit_empty_rejected( data = UpdateUserRequest(user_id="alice", permissions={}) assert "permissions" in data.model_fields_set - caller = UserAPIKeyAuth( - user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN - ) + caller = UserAPIKeyAuth(user_id="org-admin", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(HTTPException) as exc_info: - await _update_single_user_helper( - user_request=data, user_api_key_dict=caller - ) + await _update_single_user_helper(user_request=data, user_api_key_dict=caller) assert exc_info.value.status_code == 403 assert "permissions" in str(exc_info.value.detail) @@ -1324,15 +1273,11 @@ async def test_user_info_url_encoding_plus_character(mocker): mock_request.url.query = "user_id=machine-user+alp-air-admin-b58-b@tempus.com" # Mock user_api_key_dict - mock_user_api_key_dict = UserAPIKeyAuth( - user_id="test_admin", user_role="proxy_admin" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_admin", user_role="proxy_admin") # Call user_info function with the URL-decoded user_id (as FastAPI would pass it) # FastAPI would normally convert + to space, but our fix should handle this - decoded_user_id = ( - "machine-user alp-air-admin-b58-b@tempus.com" # What FastAPI gives us - ) + decoded_user_id = "machine-user alp-air-admin-b58-b@tempus.com" # What FastAPI gives us expected_user_id = "machine-user+alp-air-admin-b58-b@tempus.com" response = await user_info( @@ -1383,9 +1328,7 @@ async def test_user_info_nonexistent_user(mocker): mock_request = mocker.MagicMock(spec=Request) # Mock user_api_key_dict - mock_user_api_key_dict = UserAPIKeyAuth( - user_id="test_admin", user_role="proxy_admin" - ) + mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_admin", user_role="proxy_admin") # Call user_info function with a non-existent user_id nonexistent_user_id = "nonexistent-user@example.com" @@ -1423,14 +1366,10 @@ async def test_user_info_no_user_id_view_only_admin_gets_proxy_admin_payload(moc mock_get_user_info_for_proxy_admin, ) - viewer = UserAPIKeyAuth( - user_id="viewer", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value - ) + viewer = UserAPIKeyAuth(user_id="viewer", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value) mock_request = mocker.MagicMock(spec=Request) - response = await user_info( - user_id=None, user_api_key_dict=viewer, request=mock_request - ) + response = await user_info(user_id=None, user_api_key_dict=viewer, request=mock_request) mock_get_user_info_for_proxy_admin.assert_awaited_once_with(user_api_key_dict=viewer) assert response is admin_payload @@ -1457,9 +1396,7 @@ async def test_new_user_default_teams_flow(mocker): mock_prisma_client.db.litellm_usertable.count = mock_count persisted_user_row = mocker.MagicMock() persisted_user_row.teams = ["96fed65b-0182-4ff4-8429-2721cd7d42af"] - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - return_value=persisted_user_row - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(return_value=persisted_user_row) # Mock duplicate checks to pass async def mock_check_duplicate_user_email(*args, **kwargs): @@ -1527,26 +1464,20 @@ async def test_new_user_default_teams_flow(mocker): ) # Create test request data WITHOUT teams (teams should come from defaults) - user_request = NewUserRequest( - user_email="test@example.com", user_role="internal_user" - ) + user_request = NewUserRequest(user_email="test@example.com", user_role="internal_user") # Mock user_api_key_dict mock_user_api_key_dict = UserAPIKeyAuth(user_id="test_admin") # Call new_user function - response = await new_user( - data=user_request, user_api_key_dict=mock_user_api_key_dict - ) + response = await new_user(data=user_request, user_api_key_dict=mock_user_api_key_dict) # Verify generate_key_helper_fn was called WITHOUT teams mock_generate_key_helper_fn.assert_called_once() call_kwargs = mock_generate_key_helper_fn.call_args.kwargs # Teams should be removed from the data passed to generate_key_helper_fn - assert ( - "teams" not in call_kwargs - ), "Teams should not be passed to generate_key_helper_fn" + assert "teams" not in call_kwargs, "Teams should not be passed to generate_key_helper_fn" assert call_kwargs["request_type"] == "user" assert call_kwargs["user_email"] == "test@example.com" assert call_kwargs["user_role"] == "internal_user" @@ -1591,24 +1522,16 @@ def test_update_internal_new_user_params_proxy_admin_role(): try: # Create test data with PROXY_ADMIN role - data = NewUserRequest( - user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN.value - ) + data = NewUserRequest(user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN.value) data_json = data.model_dump(exclude_unset=True) # Call the function result = _update_internal_new_user_params(data_json=data_json, data=data) # Assertions - default params should NOT be applied for PROXY_ADMIN - assert ( - "max_budget" not in result - ), "Default max_budget should NOT be applied to PROXY_ADMIN" - assert ( - "models" not in result - ), "Default models should NOT be applied to PROXY_ADMIN" - assert ( - "tpm_limit" not in result - ), "Default tpm_limit should NOT be applied to PROXY_ADMIN" + assert "max_budget" not in result, "Default max_budget should NOT be applied to PROXY_ADMIN" + assert "models" not in result, "Default models should NOT be applied to PROXY_ADMIN" + assert "tpm_limit" not in result, "Default tpm_limit should NOT be applied to PROXY_ADMIN" # These should still work assert result["user_email"] == "admin@example.com" @@ -1722,15 +1645,9 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): user_email_clause = where_clause.get("user_email", {}) # Check that the query structure is correct for case insensitive search - assert ( - "equals" in user_email_clause - ), "Query should use 'equals' for case insensitive search" - assert ( - user_email_clause.get("mode") == "insensitive" - ), "Query should use 'insensitive' mode" - assert ( - user_email_clause.get("equals") == "user@example.com" - ), "Query should search for the provided email" + assert "equals" in user_email_clause, "Query should use 'equals' for case insensitive search" + assert user_email_clause.get("mode") == "insensitive", "Query should use 'insensitive' mode" + assert user_email_clause.get("equals") == "user@example.com", "Query should search for the provided email" return mock_existing_user # Return existing user to simulate duplicate @@ -1741,9 +1658,7 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): await _check_duplicate_user_email("user@example.com", mock_prisma_client) assert exc_info.value.status_code == 409 - assert "User with email User@Example.com already exists" in str( - exc_info.value.detail - ) + assert "User with email User@Example.com already exists" in str(exc_info.value.detail) # Test Case 2: No duplicate found async def mock_find_first_no_duplicate(*args, **kwargs): @@ -1768,9 +1683,7 @@ async def test_check_duplicate_user_email_case_insensitive(mocker): pytest.fail(f"Should not raise exception when no duplicate found, but got: {e}") # Test Case 3: None email should not cause issues - await _check_duplicate_user_email( - None, mock_prisma_client - ) # Should not raise exception + await _check_duplicate_user_email(None, mock_prisma_client) # Should not raise exception @pytest.mark.asyncio @@ -1880,9 +1793,7 @@ def test_process_keys_for_user_info_filters_dashboard_keys(monkeypatch): # Verify dashboard key is not in results result_team_ids = [key.get("team_id") for key in result] - assert ( - UI_SESSION_TOKEN_TEAM_ID not in result_team_ids - ), "Dashboard key should be filtered out" + assert UI_SESSION_TOKEN_TEAM_ID not in result_team_ids, "Dashboard key should be filtered out" # Verify regular keys are included assert "regular-team" in result_team_ids, "Regular team key should be included" @@ -1892,9 +1803,7 @@ def test_process_keys_for_user_info_filters_dashboard_keys(monkeypatch): result_tokens = [key.get("token") for key in result] assert "sk-regular-token" in result_tokens, "Regular key should be included" assert "sk-no-team-token" in result_tokens, "No-team key should be included" - assert ( - "sk-dashboard-token" not in result_tokens - ), "Dashboard key should not be included" + assert "sk-dashboard-token" not in result_tokens, "Dashboard key should not be included" def test_process_keys_for_user_info_handles_none_keys(monkeypatch): @@ -2419,9 +2328,7 @@ async def test_get_user_daily_activity_non_admin_cannot_view_other_users(monkeyp ) assert exc_info.value.status_code == 403 - assert "Non-admin users can only view their own spend data" in str( - exc_info.value.detail - ) + assert "Non-admin users can only view their own spend data" in str(exc_info.value.detail) # Case 2: Non-admin omits user_id — should default to their own user_id mock_response = MagicMock() @@ -2708,14 +2615,10 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker): async def mock_find_unique(*args, **kwargs): return mock_user_row - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) # Mock find_many for teams (no teams) - mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[]) # Mock all delete_many calls mock_prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock( @@ -2741,9 +2644,7 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker): # Call delete_user data = DeleteUserRequest(user_ids=["admin-creator"]) - user_api_key_dict = UserAPIKeyAuth( - user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN - ) + user_api_key_dict = UserAPIKeyAuth(user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN) await delete_user(data=data, user_api_key_dict=user_api_key_dict) @@ -2752,9 +2653,7 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker): call_kwargs = mock_prisma_client.db.litellm_invitationlink.delete_many.call_args where_clause = call_kwargs.kwargs.get("where") or call_kwargs[1].get("where") - assert ( - "OR" in where_clause - ), "Should use OR to match user_id, created_by, and updated_by" + assert "OR" in where_clause, "Should use OR to match user_id, created_by, and updated_by" or_conditions = where_clause["OR"] assert len(or_conditions) == 3, "Should have 3 OR conditions" @@ -2875,9 +2774,7 @@ async def test_delete_user_rejects_org_admin_deleting_outside_scope(mocker): async def mock_find_unique(*args, **kwargs): return mock_target_user - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) # Caller (org_admin_user) administers org-A. caller_membership = mocker.MagicMock() @@ -2903,16 +2800,12 @@ async def test_delete_user_rejects_org_admin_deleting_outside_scope(mocker): return [caller_membership] return [] - mock_prisma_client.db.litellm_organizationmembership.find_many = mocker.AsyncMock( - side_effect=mock_find_memberships - ) + mock_prisma_client.db.litellm_organizationmembership.find_many = mocker.AsyncMock(side_effect=mock_find_memberships) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) data = DeleteUserRequest(user_ids=["victim"]) - user_api_key_dict = UserAPIKeyAuth( - user_id="org_admin_user", user_role=LitellmUserRoles.ORG_ADMIN - ) + user_api_key_dict = UserAPIKeyAuth(user_id="org_admin_user", user_role=LitellmUserRoles.ORG_ADMIN) with pytest.raises(HTTPException) as exc: await delete_user(data=data, user_api_key_dict=user_api_key_dict) @@ -2920,11 +2813,8 @@ async def test_delete_user_rejects_org_admin_deleting_outside_scope(mocker): # Critical: no delete_many calls should have executed. assert ( - not hasattr( - mock_prisma_client.db.litellm_verificationtoken.delete_many, "mock_calls" - ) - or len(mock_prisma_client.db.litellm_verificationtoken.delete_many.mock_calls) - == 0 + not hasattr(mock_prisma_client.db.litellm_verificationtoken.delete_many, "mock_calls") + or len(mock_prisma_client.db.litellm_verificationtoken.delete_many.mock_calls) == 0 ) @@ -2943,9 +2833,7 @@ async def test_user_update_rejects_silent_create_for_non_proxy_admin(mocker): mock_prisma_client = mocker.MagicMock() # user_email lookup yields None → would silently create pre-fix. - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=None) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) user_request = UpdateUserRequest( @@ -2959,9 +2847,7 @@ async def test_user_update_rejects_silent_create_for_non_proxy_admin(mocker): ) with pytest.raises(HTTPException) as exc: - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=org_admin - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=org_admin) assert exc.value.status_code == 404 @@ -3005,17 +2891,13 @@ async def test_user_info_v2_proxy_admin_can_query_any_user(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) response = await user_info_v2( request=mock_request, @@ -3069,17 +2951,13 @@ async def test_user_info_v2_redacts_scim_enterprise_metadata(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) response = await user_info_v2( request=mock_request, @@ -3088,9 +2966,7 @@ async def test_user_info_v2_redacts_scim_enterprise_metadata(mocker): ) assert isinstance(response, UserInfoV2Response) - assert response.metadata == { - "scim_metadata": {"givenName": "Jane", "familyName": "Doe"} - } + assert response.metadata == {"scim_metadata": {"givenName": "Jane", "familyName": "Doe"}} assert "scim_enterprise" not in (response.metadata or {}) @@ -3159,17 +3035,13 @@ async def test_user_info_v2_internal_user_can_query_self(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - user_key = UserAPIKeyAuth( - user_id="self-user", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_key = UserAPIKeyAuth(user_id="self-user", user_role=LitellmUserRoles.INTERNAL_USER) response = await user_info_v2( request=mock_request, @@ -3204,17 +3076,13 @@ async def test_user_info_v2_internal_user_cannot_query_other(mocker): return mock_caller_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - user_key = UserAPIKeyAuth( - user_id="caller-user", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_key = UserAPIKeyAuth(user_id="caller-user", user_role=LitellmUserRoles.INTERNAL_USER) with pytest.raises(ProxyException) as exc_info: await user_info_v2( @@ -3261,17 +3129,13 @@ async def test_user_info_v2_no_user_id_defaults_to_self(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - user_key = UserAPIKeyAuth( - user_id="my-user-id", user_role=LitellmUserRoles.INTERNAL_USER - ) + user_key = UserAPIKeyAuth(user_id="my-user-id", user_role=LitellmUserRoles.INTERNAL_USER) # Call without user_id response = await user_info_v2( @@ -3299,17 +3163,13 @@ async def test_user_info_v2_nonexistent_user_returns_404(mocker): async def mock_find_unique(*args, **kwargs): return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) with pytest.raises(ProxyException) as exc_info: await user_info_v2( @@ -3357,17 +3217,13 @@ async def test_user_info_v2_response_shape(mocker): async def mock_find_unique(*args, **kwargs): return mock_user_row - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) response = await user_info_v2( request=mock_request, @@ -3402,9 +3258,7 @@ async def test_user_info_v2_response_shape(mocker): # The dashboard's user edit form hydrates its per-model budget rows from # these two, so dropping them makes a save replace the user's budgets. - assert response_dict["model_max_budget"] == { - "gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"} - } + assert response_dict["model_max_budget"] == {"gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"}} assert response_dict["model_max_budget_usage"] == { "gpt-3.5-turbo": {"current_spend": 0.0, "budget_limit": 5.0, "time_period": "30d"} } @@ -3463,9 +3317,7 @@ async def test_user_info_v2_team_admin_can_query_team_member(mocker): return mock_target return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) # Mock team with caller as admin mock_team = mocker.MagicMock() @@ -3482,17 +3334,13 @@ async def test_user_info_v2_team_admin_can_query_team_member(mocker): async def mock_find_many_teams(*args, **kwargs): return [mock_team] - mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock( - side_effect=mock_find_many_teams - ) + mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(side_effect=mock_find_many_teams) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - team_admin_key = UserAPIKeyAuth( - user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER - ) + team_admin_key = UserAPIKeyAuth(user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER) response = await user_info_v2( request=mock_request, @@ -3532,9 +3380,7 @@ async def test_user_info_v2_team_admin_cannot_query_non_team_member(mocker): return mock_target return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) # Mock team where caller is admin mock_team = mocker.MagicMock() @@ -3550,17 +3396,13 @@ async def test_user_info_v2_team_admin_cannot_query_non_team_member(mocker): async def mock_find_many_teams(*args, **kwargs): return [mock_team] - mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock( - side_effect=mock_find_many_teams - ) + mock_prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(side_effect=mock_find_many_teams) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) - team_admin_key = UserAPIKeyAuth( - user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER - ) + team_admin_key = UserAPIKeyAuth(user_id="team-admin-user", user_role=LitellmUserRoles.INTERNAL_USER) with pytest.raises(ProxyException) as exc_info: await user_info_v2( @@ -3610,18 +3452,14 @@ async def test_user_info_v2_url_encoding_plus_character(mocker): return mock_user_row return None - mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( - side_effect=mock_find_unique - ) + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=mock_find_unique) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_request = mocker.MagicMock(spec=Request) mock_request.url.query = f"user_id={expected_user_id}" - admin_key = UserAPIKeyAuth( - user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_key = UserAPIKeyAuth(user_id="admin-user", user_role=LitellmUserRoles.PROXY_ADMIN) # Simulate FastAPI converting + to space decoded_user_id = "machine-user admin@example.com" @@ -3718,9 +3556,7 @@ def test_enforce_user_info_access_admin_bypass(): _enforce_user_info_access, ) - admin = UserAPIKeyAuth( - user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN.value - ) + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN.value) # Should not raise even when querying a different user _enforce_user_info_access(user_id="someone_else", user_api_key_dict=admin) @@ -3759,9 +3595,7 @@ def test_enforce_user_info_access_owner_allowed(): _enforce_user_info_access, ) - user = UserAPIKeyAuth( - user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value - ) + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value) _enforce_user_info_access(user_id="alice", user_api_key_dict=user) @@ -3773,9 +3607,7 @@ def test_enforce_user_info_access_no_user_id_allowed(): _enforce_user_info_access, ) - user = UserAPIKeyAuth( - user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value - ) + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value) _enforce_user_info_access(user_id=None, user_api_key_dict=user) @@ -3832,9 +3664,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker, budge "max_budget": 100, } existing_user.user_id = "user-1" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) user_request = UpdateUserRequest.model_validate({"user_id": "user-1", budget_field: budget_value}) @@ -3844,9 +3674,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_max_budget(mocker, budge ) with pytest.raises(HTTPException) as exc: - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=caller - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=caller) assert exc.value.status_code == 403 assert budget_field in str(exc.value.detail) mock_prisma_client.update_data.assert_not_called() @@ -3868,9 +3696,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_spend(mocker): "spend": 50.0, } existing_user.user_id = "user-1" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) user_request = UpdateUserRequest( @@ -3883,9 +3709,7 @@ async def test_ghsa_wvg4_non_admin_cannot_self_escalate_spend(mocker): ) with pytest.raises(HTTPException) as exc: - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=caller - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=caller) assert exc.value.status_code == 403 assert "spend" in str(exc.value.detail) @@ -3904,12 +3728,8 @@ async def test_ghsa_wvg4_proxy_admin_can_update_user_budget(mocker): "max_budget": 100, } existing_user.user_id = "target-user" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) - mock_prisma_client.update_data = mocker.AsyncMock( - return_value={"user_id": "target-user", "max_budget": 500} - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "target-user", "max_budget": 500}) mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") @@ -3923,9 +3743,7 @@ async def test_ghsa_wvg4_proxy_admin_can_update_user_budget(mocker): user_role=LitellmUserRoles.PROXY_ADMIN, ) - result = await _update_single_user_helper( - user_request=user_request, user_api_key_dict=admin_caller - ) + result = await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) assert result is not None @@ -3941,12 +3759,8 @@ async def test_admin_user_update_spend_invalidates_counter(mocker): existing_user = mocker.MagicMock() existing_user.model_dump.return_value = {"user_id": "target-user", "spend": 50.0} existing_user.user_id = "target-user" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) - mock_prisma_client.update_data = mocker.AsyncMock( - return_value={"user_id": "target-user", "spend": -25.0} - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "target-user", "spend": -25.0}) mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") @@ -3961,13 +3775,9 @@ async def test_admin_user_update_spend_invalidates_counter(mocker): # without raising the recurring budget ceiling. Future changes should # continue allowing negative spend counters. user_request = UpdateUserRequest(user_id="target-user", spend=-25) - admin_caller = UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=admin_caller - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) mock_invalidate.assert_awaited_once_with(counter_key="spend:user:target-user") @@ -3984,9 +3794,7 @@ async def test_user_update_rejects_non_finite_spend(mocker): existing_user = mocker.MagicMock() existing_user.model_dump.return_value = {"user_id": "target-user", "spend": 50.0} existing_user.user_id = "target-user" - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) mock_prisma_client.update_data = mocker.AsyncMock() mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") @@ -3996,14 +3804,10 @@ async def test_user_update_rejects_non_finite_spend(mocker): ) user_request = UpdateUserRequest(user_id="target-user", spend=float("nan")) - admin_caller = UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ) + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) with pytest.raises(HTTPException) as exc: - await _update_single_user_helper( - user_request=user_request, user_api_key_dict=admin_caller - ) + await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) assert exc.value.status_code == 400 mock_prisma_client.update_data.assert_not_called() mock_invalidate.assert_not_awaited() @@ -4023,9 +3827,7 @@ async def test_resolve_user_email_metadata_maps_page_user_ids_to_email(mocker): mock_prisma_client = mocker.MagicMock() find_many = mocker.AsyncMock( return_value=[ - SimpleNamespace( - user_id="u1", user_email="alice@example.com", user_alias="Alice" - ), + SimpleNamespace(user_id="u1", user_email="alice@example.com", user_alias="Alice"), SimpleNamespace(user_id="u2", user_email=None, user_alias="bob-alias"), ] ) @@ -4224,19 +4026,13 @@ def _object_permission_mocks(mocker, existing_object_permission_id=None): } existing_user.user_id = "target-user" existing_user.object_permission_id = existing_object_permission_id - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=existing_user - ) - mock_prisma_client.db.litellm_objectpermissiontable.find_unique = mocker.AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = mocker.AsyncMock(return_value=None) mock_prisma_client.db.litellm_objectpermissiontable.upsert = mocker.AsyncMock( return_value=SimpleNamespace(object_permission_id="perm-new") ) mock_prisma_client.db.litellm_mcpservertable.find_many = mocker.AsyncMock(return_value=[]) - mock_prisma_client.update_data = mocker.AsyncMock( - return_value={"user_id": "target-user"} - ) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "target-user"}) mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") @@ -4272,9 +4068,7 @@ async def test_user_update_persists_mcp_entitlement_and_links_it(mocker): "mcp_tool_permissions": {"github": ["list_issues"]}, }, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) upsert_kwargs = mock_prisma_client.db.litellm_objectpermissiontable.upsert.call_args.kwargs @@ -4308,9 +4102,7 @@ async def test_user_update_invalidates_the_cached_entitlement(mocker): user_id="target-user", object_permission={"mcp_tool_permissions": {"github": []}}, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} @@ -4343,9 +4135,7 @@ async def test_admin_can_clear_a_users_mcp_entitlement(mocker): await _update_single_user_helper( user_request=UpdateUserRequest(user_id="target-user", object_permission={}), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) written = mock_prisma_client.update_data.call_args.kwargs["data"] @@ -4382,9 +4172,7 @@ async def test_user_update_invalidates_both_the_old_and_new_permission_rows(mock user_id="target-user", object_permission={"mcp_tool_permissions": {"github": ["list_issues"]}}, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list} @@ -4415,9 +4203,7 @@ async def test_non_admin_cannot_clear_their_own_mcp_entitlement(mocker): with pytest.raises(HTTPException) as exc: await _update_single_user_helper( user_request=UpdateUserRequest(user_id="target-user", object_permission={}), - user_api_key_dict=UserAPIKeyAuth( - user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER - ), + user_api_key_dict=UserAPIKeyAuth(user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER), ) assert exc.value.status_code == 403 @@ -4445,9 +4231,7 @@ async def test_non_admin_cannot_rewrite_their_own_mcp_entitlement(mocker): user_id="target-user", object_permission={"mcp_servers": [], "mcp_tool_permissions": {}}, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER - ), + user_api_key_dict=UserAPIKeyAuth(user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER), ) assert exc.value.status_code == 403 @@ -4464,9 +4248,7 @@ async def test_new_user_persists_the_requested_mcp_entitlement(mocker): return_value=SimpleNamespace(object_permission_id="perm-created") ) mock_prisma_client.db.litellm_mcpservertable.find_many = mocker.AsyncMock(return_value=[]) - mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=None) mock_prisma_client.db.litellm_usertable.count = mocker.AsyncMock(return_value=0) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( @@ -4475,9 +4257,7 @@ async def test_new_user_persists_the_requested_mcp_entitlement(mocker): ) mock_generate = mocker.patch( "litellm.proxy.management_endpoints.internal_user_endpoints.generate_key_helper_fn", - new=mocker.AsyncMock( - return_value={"user_id": "new-human", "token": "sk-x", "expires": None} - ), + new=mocker.AsyncMock(return_value={"user_id": "new-human", "token": "sk-x", "expires": None}), ) mocker.patch( "litellm.proxy.hooks.user_management_event_hooks.UserManagementEventHooks.async_user_created_hook", @@ -4489,9 +4269,7 @@ async def test_new_user_persists_the_requested_mcp_entitlement(mocker): user_id="new-human", object_permission={"mcp_tool_permissions": {"github": ["list_issues"]}}, ), - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) created = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] @@ -4533,16 +4311,12 @@ async def test_user_info_v2_returns_the_mcp_entitlement(mocker): response = await user_info_v2( request=SimpleNamespace(query_params={}), user_id="human-1", - user_api_key_dict=UserAPIKeyAuth( - user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN - ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), ) assert response.object_permission is not None assert response.object_permission.mcp_servers == ["github"] - assert response.object_permission.mcp_tool_permissions == { - "github": ["list_issues"] - } + assert response.object_permission.mcp_tool_permissions == {"github": ["list_issues"]} @pytest.mark.asyncio @@ -4558,9 +4332,7 @@ async def test_user_info_v2_returns_the_mcp_entitlement(mocker): ], ids=["supplied", "omitted", "empty"], ) -async def test_user_new_persists_model_max_budget( - monkeypatch, model_max_budget, expected_written -): +async def test_user_new_persists_model_max_budget(monkeypatch, model_max_budget, expected_written): """ /user/new used to echo model_max_budget back while writing {} to the user row, so a per-model budget looked configured and was read by nothing. @@ -4674,6 +4446,11 @@ async def test_user_update_hashes_and_persists_strong_password(_admin_prisma, mo _update_single_user_helper, ) + mocker.patch( # test-quality-ok: same module-global mocking every test in this file already uses + "litellm.proxy.proxy_server.general_settings", + {"password_policy_check_breached_passwords": False}, + ) + mock_prisma_client = _admin_prisma existing_user = mocker.MagicMock() existing_user.model_dump.return_value = {"user_id": "target-user"} @@ -4691,6 +4468,152 @@ async def test_user_update_hashes_and_persists_strong_password(_admin_prisma, mo written_data = mock_prisma_client.update_data.call_args.kwargs["data"] assert written_data.get("password") is not None assert written_data["password"] != strong_password + # An admin-set password is known to the admin, so the user must be forced + # to change it at next login and the breach screen re-armed. + assert written_data["password_reset_required"] is True + assert written_data["last_breach_check_at"] is None + + +@pytest.mark.asyncio +@respx.mock +async def test_user_update_rejects_breached_password(_admin_prisma): + """A strength-passing password found in the HIBP corpus must be rejected + before it ever reaches the DB write.""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_single_user_helper, + ) + + password = "Str0ng!Passw0rd" + sha1 = hashlib.sha1(password.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + respx.get(f"https://api.pwnedpasswords.com/range/{sha1[:5]}").mock( + return_value=httpx.Response(200, text=f"{sha1[5:]}:1387") + ) + + user_request = UpdateUserRequest(user_id="target-user", password=password) + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + + with pytest.raises(ProxyException) as exc_info: + await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) + + assert exc_info.value.code == "400" + assert "data breaches" in exc_info.value.message + _admin_prisma.db.litellm_usertable.find_first.assert_not_called() + + +@pytest.mark.asyncio +async def test_bulk_update_all_users_rejects_a_password(_admin_prisma): + """The all_users fast path writes user_updates straight to update_many, + bypassing _update_single_user_helper. A password riding along would be + stored as unvalidated plaintext on every row, so it must be rejected + before any DB access.""" + from fastapi import HTTPException + + from litellm.proxy._types import UpdateUserRequestNoUserIDorEmail + from litellm.proxy.management_endpoints.internal_user_endpoints import bulk_user_update + from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( + BulkUpdateUserRequest, + ) + + data = BulkUpdateUserRequest( + all_users=True, + user_updates=UpdateUserRequestNoUserIDorEmail(password="Str0ng!Passw0rd"), + ) + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + + with pytest.raises(HTTPException) as exc_info: + await bulk_user_update(data=data, user_api_key_dict=admin_caller) + + assert exc_info.value.status_code == 400 + assert "not supported" in str(exc_info.value.detail) + _admin_prisma.db.litellm_usertable.find_many.assert_not_called() + _admin_prisma.db.litellm_usertable.update_many.assert_not_called() + + +def _hibp_client_with_handler(handler) -> AsyncHTTPHandler: + """A real AsyncHTTPHandler over httpx.MockTransport (the DI seam used + throughout test_password_policy.py), so no network is touched.""" + http_handler = AsyncHTTPHandler() + http_handler.client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + return http_handler + + +@pytest.mark.asyncio +async def test_bulk_update_breached_password_fails_only_that_user(_admin_prisma, mocker): + """In a bulk batch, a breached password fails only its own entry, before + any DB write for it; sibling entries with acceptable passwords persist.""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + bulk_update_processed_users, + ) + + breached = "Br3ached!Passw0rd" + clean = "NewP@ssw0rd123" + breached_sha1 = hashlib.sha1(breached.encode("utf-8"), usedforsecurity=False).hexdigest().upper() + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == f"/range/{breached_sha1[:5]}": + return httpx.Response(200, text=f"{breached_sha1[5:]}:1387") + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + mock_prisma_client = _admin_prisma + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "user-clean"} + existing_user.user_id = "user-clean" + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "user-clean"}) + mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) + + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + response = await bulk_update_processed_users( + users_to_update=[ + UpdateUserRequest(user_id="user-breached", password=breached), + UpdateUserRequest(user_id="user-clean", password=clean), + ], + user_api_key_dict=admin_caller, + hibp_client=_hibp_client_with_handler(handler), + ) + + assert response.successful_updates == 1 + assert response.failed_updates == 1 + by_user = {r.user_id: r for r in response.results} + assert by_user["user-breached"].success is False + assert "data breaches" in by_user["user-breached"].error + assert by_user["user-clean"].success is True + (write_call,) = mock_prisma_client.update_data.call_args_list + assert write_call.kwargs["user_id"] == "user-clean" + + +@pytest.mark.asyncio +async def test_bulk_update_screens_shared_password_with_single_lookup(_admin_prisma, mocker): + """A batch where every user gets the same password costs one HIBP lookup, + not one per user (the serial per-user checks this regresses against).""" + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + bulk_update_processed_users, + ) + + lookup_count = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal lookup_count + lookup_count += 1 + return httpx.Response(200, text="0000000000000000000000000000000000A:1") + + mock_prisma_client = _admin_prisma + existing_user = mocker.MagicMock() + existing_user.model_dump.return_value = {"user_id": "user-0"} + existing_user.user_id = "user-0" + mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=existing_user) + mock_prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": "user-0"}) + mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x) + + admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) + response = await bulk_update_processed_users( + users_to_update=[UpdateUserRequest(user_id=f"user-{i}", password="NewP@ssw0rd123") for i in range(5)], + user_api_key_dict=admin_caller, + hibp_client=_hibp_client_with_handler(handler), + ) + + assert response.successful_updates == 5 + assert lookup_count == 1 @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 80773f314d8..c66d490f122 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -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.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_password_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_password_endpoints.py new file mode 100644 index 00000000000..c04353fec99 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_password_endpoints.py @@ -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 diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py index 88f8be4e49a..8daea8d0ad2 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_login_sso.py @@ -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()) diff --git a/tests/test_litellm/proxy/test__types.py b/tests/test_litellm/proxy/test__types.py index 8f7859129bb..6bc40658b3c 100644 --- a/tests/test_litellm/proxy/test__types.py +++ b/tests/test_litellm/proxy/test__types.py @@ -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 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 950a6cc3c40..2792e176e0e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -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.""" diff --git a/tests/test_litellm/secret_managers/hashicorp_vault_parity.json b/tests/test_litellm/secret_managers/hashicorp_vault_parity.json new file mode 100644 index 00000000000..f1faefd48e8 --- /dev/null +++ b/tests/test_litellm/secret_managers/hashicorp_vault_parity.json @@ -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" + } +] diff --git a/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py b/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py index 1676540e4ec..fc18cb8b8f7 100644 --- a/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py +++ b/tests/test_litellm/secret_managers/test_hashicorp_secret_manager.py @@ -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 diff --git a/tests/test_litellm_rust/support/fake_gcs.py b/tests/test_litellm_rust/support/fake_gcs.py new file mode 100644 index 00000000000..67eb61798b9 --- /dev/null +++ b/tests/test_litellm_rust/support/fake_gcs.py @@ -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)) diff --git a/tests/test_litellm_rust/support/s3_stub.py b/tests/test_litellm_rust/support/s3_stub.py new file mode 100644 index 00000000000..5a683fb78f3 --- /dev/null +++ b/tests/test_litellm_rust/support/s3_stub.py @@ -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'NoSuchKey' + 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) diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index a0f2396a25d..e0d92a2e957 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -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") diff --git a/tests/unit/models/test_models.py b/tests/unit/models/test_models.py index b8bf55f1b4a..ab456bb1624 100644 --- a/tests/unit/models/test_models.py +++ b/tests/unit/models/test_models.py @@ -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: diff --git a/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.integration.test.tsx new file mode 100644 index 00000000000..c4cf8ebcf9d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.integration.test.tsx @@ -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(); + + 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(); + + 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(); + + 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(); + + expect(screen.getByText(/must be changed before you can use the dashboard/)).toBeInTheDocument(); + }); + + it("hides the forced-reset warning for a normal session", () => { + render(); + + 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(); + 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 }); + } + }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx new file mode 100644 index 00000000000..05a6bf3ae94 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/change-password/ChangePasswordForm.tsx @@ -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; + +export function ChangePasswordForm() { + const { accessToken, passwordResetRequired } = useAuthorized(); + const form = useZodForm(changePasswordSchema, { + defaultValues: { currentPassword: "", newPassword: "", confirmNewPassword: "" }, + }); + const [isPending, setIsPending] = useState(false); + const [submitError, setSubmitError] = useState(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 ( +
+ + +

Change Password

+

+ Enter your current password and choose a new one. The new password must meet this proxy's password + policy. +

+ + {passwordResetRequired && ( + + + + 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. + + + )} + +
+ + + {({ ref, ...field }) => } + + + + {({ ref, ...field }) => } + + + + {({ ref, ...field }) => } + + + + {submitError && ( + + + {submitError} + + )} + +
+ +
+
+ + +
+ ); +} + +export default ChangePasswordForm; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/change-password/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/change-password/page.tsx new file mode 100644 index 00000000000..0a6ae926ceb --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/change-password/page.tsx @@ -0,0 +1,7 @@ +"use client"; + +import ChangePasswordForm from "./ChangePasswordForm"; + +export default function ChangePasswordPage() { + return ; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts index 40d1ec09d1f..581ee8b2580 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts @@ -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, }; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx index d7e1bb82564..d854befa197 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx @@ -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) => { + const encode = (part: Record) => + 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( + + +
+ + , + ); + + 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( + + +
+ + , + ); + + pendingUiConfig.resolve(); + + expect(await screen.findByTestId("page-content")).toBeInTheDocument(); + expect(replaceMock).not.toHaveBeenCalled(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 6d326f5280e..2e903c7b150 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -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 ; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx index 091b2f1403f..1217d878489 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx @@ -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(["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(["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( + + + , + ); + 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( + + + , + ); + 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( + + + , + ); + 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); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx index 2df304eff96..818b5150650 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx @@ -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 = { 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; +}) { + return ( + + + + Delete MCP Server? + +
+

+ This action is permanent and cannot be undone. All associated configurations will be removed. +

+ + {server && ( +
+ {server.server_name && ( +
+
Name
+
{server.server_name}
+
+ )} +
+
ID
+
{server.server_id}
+
+ {server.url && ( +
+
URL
+
{server.url}
+
+ )} +
+ )} +
+ + Cancel + + + + + ); +} + const MCPServers: React.FC = ({ accessToken, userRole, userID, isViewOnly = false }) => { const { data: mcpServers, isLoading: isLoadingServers, refetch } = useMCPServers(); @@ -298,16 +355,12 @@ const MCPServers: React.FC = ({ 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 = ({ 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 = ({ 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 = ({ accessToken, userRole, userID, i return (
- !open && cancelDelete()}> - - - Delete MCP Server? - -
-

- This action is permanent and cannot be undone. All associated configurations will be removed. -

- - {serverToDelete && ( -
- {serverToDelete.server_name && ( -
-
Name
-
{serverToDelete.server_name}
-
- )} -
-
ID
-
{serverToDelete.server_id}
-
- {serverToDelete.url && ( -
-
URL
-
{serverToDelete.url}
-
- )} -
- )} -
- - Cancel - - - - + !open && cancelDelete()} + server={serverToDelete} + isDeleting={isDeletingServer} + onConfirm={confirmDelete} + /> = ({ accessToken, userRole, userID, i My Connections - {isAdminRole(userRole) && ( + {isAdminRole(userRole) ? ( <> - )} - {!isAdminRole(userRole) && ( + ) : ( + )} + )} +