From a1b87bcca8d7d7558276036fa9ceca98caa1d58e Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 6 Nov 2026 12:30:07 +0900 Subject: [PATCH] feat(redis): extend the native Redis foundation for the store families (#50169) Reworks the redis-client, connect, resolve and testing crates, adds python-compat json loads and TTL seconds, replaces the redis-store memory tier with a DualCache plus a shared failure log, and gives the Python side the Redis timeout log throttle and a request OTEL parent span for service events Co-authored-by: Claude Opus 5.5 --- .github/workflows/test-unit.yml | 1 + .../skills/rust-python-boundary/SKILL.md | 53 + litellm-rust/AGENTS.md | 4 + litellm-rust/Cargo.lock | 207 ++- litellm-rust/Cargo.toml | 4 + litellm-rust/crates/cache-memory/src/cache.rs | 53 +- litellm-rust/crates/config/AGENTS.md | 3 + litellm-rust/crates/config/Cargo.toml | 7 +- litellm-rust/crates/config/src/error.rs | 7 - litellm-rust/crates/config/src/lib.rs | 5 +- litellm-rust/crates/config/src/redis.rs | 658 +------- litellm-rust/crates/config/tests/config.rs | 26 + litellm-rust/crates/config/tests/redis.rs | 405 ----- litellm-rust/crates/gateway/AGENTS.md | 7 +- litellm-rust/crates/gateway/Cargo.toml | 7 + litellm-rust/crates/gateway/src/error.rs | 11 +- litellm-rust/crates/gateway/src/lib.rs | 22 +- litellm-rust/crates/gateway/src/main.rs | 4 +- .../crates/gateway/src/redis/configured.rs | 144 ++ .../crates/gateway/src/redis/layers.rs | 226 +++ .../gateway/src/{redis.rs => redis/mod.rs} | 409 +++-- .../crates/gateway/src/redis/secrets.rs | 61 + .../crates/gateway/src/redis/surface_tests.rs | 964 +++++++++++ litellm-rust/crates/host-python/src/lib.rs | 2 +- .../crates/host-python/src/runtime.rs | 5 + litellm-rust/crates/python-bridge/Cargo.toml | 8 + litellm-rust/crates/python-bridge/src/lib.rs | 15 + .../src/redis/client/credentials.rs | 47 + .../python-bridge/src/redis/client/mod.rs | 385 +++++ .../crates/python-bridge/src/redis/events.rs | 391 +++++ .../crates/python-bridge/src/redis/mod.rs | 11 + .../python-bridge/src/redis/store/mod.rs | 248 +++ .../python-bridge/src/redis/store/request.rs | 60 + .../python-bridge/src/redis/store/value.rs | 125 ++ .../crates/python-compat/benches/formats.rs | 11 +- .../crates/python-compat/src/error.rs | 2 + litellm-rust/crates/python-compat/src/json.rs | 352 +++- litellm-rust/crates/python-compat/src/lib.rs | 5 +- .../crates/python-compat/src/number.rs | 121 ++ .../crates/python-compat/tests/fixtures.rs | 21 + litellm-rust/crates/redis-client/AGENTS.md | 56 +- litellm-rust/crates/redis-client/Cargo.toml | 19 +- litellm-rust/crates/redis-client/src/batch.rs | 147 ++ .../crates/redis-client/src/client.rs | 788 +++------ .../crates/redis-client/src/commands.rs | 457 +++++- .../crates/redis-client/src/config.rs | 498 +++++- .../crates/redis-client/src/connection.rs | 617 ++++++- litellm-rust/crates/redis-client/src/error.rs | 185 ++- .../crates/redis-client/src/events.rs | 78 + litellm-rust/crates/redis-client/src/lib.rs | 26 +- .../crates/redis-client/src/pipeline.rs | 339 ++-- .../crates/redis-client/src/pubsub.rs | 225 +++ .../crates/redis-client/src/script.rs | 76 +- .../crates/redis-client/src/test_support.rs | 126 -- .../crates/redis-client/tests/batch.rs | 330 ++++ .../crates/redis-client/tests/client.rs | 1039 ++++++++++++ .../crates/redis-client/tests/config.rs | 290 ++++ .../crates/redis-client/tests/connection.rs | 329 ++++ .../crates/redis-client/tests/pipeline.rs | 362 +++++ .../crates/redis-client/tests/pubsub.rs | 189 +++ .../crates/redis-client/tests/script.rs | 133 ++ .../crates/redis-client/tests/support/mod.rs | 206 +++ litellm-rust/crates/redis-connect/AGENTS.md | 10 + litellm-rust/crates/redis-connect/Cargo.toml | 39 + .../crates/redis-connect/src/connect.rs | 154 ++ .../crates/redis-connect/src/error.rs | 102 ++ .../crates/redis-connect/src/events.rs | 129 ++ litellm-rust/crates/redis-connect/src/lib.rs | 11 + .../crates/redis-connect/src/refresh.rs | 783 +++++++++ .../crates/redis-connect/src/sources/aws.rs | 188 +++ .../crates/redis-connect/src/sources/azure.rs | 64 + .../crates/redis-connect/src/sources/gcp.rs | 126 ++ .../crates/redis-connect/src/sources/mod.rs | 7 + .../crates/redis-connect/tests/connect.rs | 230 +++ .../crates/redis-connect/tests/live/aws.rs | 17 + .../crates/redis-connect/tests/live/azure.rs | 17 + .../crates/redis-connect/tests/live/gcp.rs | 17 + .../crates/redis-connect/tests/live/main.rs | 111 ++ litellm-rust/crates/redis-resolve/AGENTS.md | 31 + litellm-rust/crates/redis-resolve/Cargo.toml | 20 + .../crates/redis-resolve/redis-py-kwargs.json | 118 ++ .../crates/redis-resolve/redis_py_kwargs.py | 110 ++ litellm-rust/crates/redis-resolve/src/auth.rs | 95 ++ .../crates/redis-resolve/src/endpoint.rs | 180 +++ .../crates/redis-resolve/src/error.rs | 80 + litellm-rust/crates/redis-resolve/src/keys.rs | 304 ++++ litellm-rust/crates/redis-resolve/src/lib.rs | 95 ++ .../crates/redis-resolve/src/options.rs | 281 ++++ litellm-rust/crates/redis-resolve/src/tls.rs | 189 +++ .../crates/redis-resolve/tests/catalog.rs | 129 ++ .../crates/redis-resolve/tests/interpret.rs | 790 +++++++++ litellm-rust/crates/redis-store/AGENTS.md | 13 + litellm-rust/crates/redis-store/Cargo.toml | 25 + litellm-rust/crates/redis-store/src/codec.rs | 74 + .../redis-store/src/dual_cache/batch_read.rs | 242 +++ .../redis-store/src/dual_cache/increment.rs | 129 ++ .../crates/redis-store/src/dual_cache/mod.rs | 279 ++++ .../crates/redis-store/src/dual_cache/read.rs | 55 + .../redis-store/src/dual_cache/write.rs | 178 +++ litellm-rust/crates/redis-store/src/error.rs | 63 + .../crates/redis-store/src/failure_log.rs | 18 + litellm-rust/crates/redis-store/src/lib.rs | 30 + litellm-rust/crates/redis-store/src/memory.rs | 122 ++ litellm-rust/crates/redis-store/src/number.rs | 83 + litellm-rust/crates/redis-store/src/remote.rs | 40 + .../crates/redis-store/src/request.rs | 219 +++ litellm-rust/crates/redis-store/src/value.rs | 79 + .../crates/redis-store/tests/request.rs | 241 +++ .../crates/redis-store/tests/support/mod.rs | 149 ++ litellm-rust/crates/redis-testing/AGENTS.md | 7 + litellm-rust/crates/redis-testing/Cargo.toml | 5 +- .../crates/redis-testing/src/cluster.rs | 69 +- .../crates/redis-testing/src/deployment.rs | 113 ++ .../crates/redis-testing/src/error.rs | 12 + litellm-rust/crates/redis-testing/src/exec.rs | 31 + litellm-rust/crates/redis-testing/src/lib.rs | 9 + .../crates/redis-testing/src/ports.rs | 11 + .../crates/redis-testing/src/proxy.rs | 65 + .../crates/redis-testing/src/sentinel.rs | 170 ++ .../crates/redis-testing/src/standalone.rs | 20 +- litellm-rust/crates/redis-testing/src/tls.rs | 148 ++ litellm-rust/crates/redis-types/AGENTS.md | 2 + litellm-rust/crates/redis-types/Cargo.toml | 9 + litellm-rust/crates/redis-types/src/lib.rs | 175 ++ litellm/_internal_context.py | 19 +- litellm/_redis.py | 139 +- litellm/caching/redis_batch.py | 154 +- litellm/caching/redis_cache.py | 24 +- litellm/redis/__init__.py | 0 litellm/redis/store_settings.py | 11 + litellm/redis/stores/__init__.py | 0 litellm/redis/stores/_route.py | 50 + litellm/redis/stores/connection.py | 57 + litellm/rust_bridge/_native.pyi | 55 + litellm/rust_bridge/host/AGENTS.md | 2 +- litellm/rust_bridge/host/redis_credentials.py | 19 + litellm/rust_bridge/host/redis_events.py | 124 ++ litellm/rust_bridge/logger.py | 32 + tests/unit/caching/test_caching.py | 6 +- tests/unit/caching/test_dual_cache.py | 4 +- tests/unit/caching/test_redis_batch.py | 64 +- tests/unit/caching/test_redis_cache.py | 10 +- tests/unit/redis/__init__.py | 0 tests/unit/redis/stores/__init__.py | 0 tests/unit/redis/stores/test_connection.py | 74 + .../host/test_redis_credentials.py | 38 + .../rust_bridge/host/test_redis_events.py | 139 ++ tests/unit/rust_bridge/test_logger.py | 26 + tests/unit/test_redis.py | 1420 ++++++++++++++++- 149 files changed, 18839 insertions(+), 2552 deletions(-) create mode 100644 litellm-rust/.agents/skills/rust-python-boundary/SKILL.md create mode 100644 litellm-rust/crates/config/AGENTS.md delete mode 100644 litellm-rust/crates/config/tests/redis.rs create mode 100644 litellm-rust/crates/gateway/src/redis/configured.rs create mode 100644 litellm-rust/crates/gateway/src/redis/layers.rs rename litellm-rust/crates/gateway/src/{redis.rs => redis/mod.rs} (60%) create mode 100644 litellm-rust/crates/gateway/src/redis/secrets.rs create mode 100644 litellm-rust/crates/gateway/src/redis/surface_tests.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/client/credentials.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/client/mod.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/events.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/mod.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/mod.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/request.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/value.rs create mode 100644 litellm-rust/crates/python-compat/src/number.rs create mode 100644 litellm-rust/crates/redis-client/src/batch.rs create mode 100644 litellm-rust/crates/redis-client/src/events.rs create mode 100644 litellm-rust/crates/redis-client/src/pubsub.rs delete mode 100644 litellm-rust/crates/redis-client/src/test_support.rs create mode 100644 litellm-rust/crates/redis-client/tests/batch.rs create mode 100644 litellm-rust/crates/redis-client/tests/client.rs create mode 100644 litellm-rust/crates/redis-client/tests/config.rs create mode 100644 litellm-rust/crates/redis-client/tests/connection.rs create mode 100644 litellm-rust/crates/redis-client/tests/pipeline.rs create mode 100644 litellm-rust/crates/redis-client/tests/pubsub.rs create mode 100644 litellm-rust/crates/redis-client/tests/script.rs create mode 100644 litellm-rust/crates/redis-client/tests/support/mod.rs create mode 100644 litellm-rust/crates/redis-connect/AGENTS.md create mode 100644 litellm-rust/crates/redis-connect/Cargo.toml create mode 100644 litellm-rust/crates/redis-connect/src/connect.rs create mode 100644 litellm-rust/crates/redis-connect/src/error.rs create mode 100644 litellm-rust/crates/redis-connect/src/events.rs create mode 100644 litellm-rust/crates/redis-connect/src/lib.rs create mode 100644 litellm-rust/crates/redis-connect/src/refresh.rs create mode 100644 litellm-rust/crates/redis-connect/src/sources/aws.rs create mode 100644 litellm-rust/crates/redis-connect/src/sources/azure.rs create mode 100644 litellm-rust/crates/redis-connect/src/sources/gcp.rs create mode 100644 litellm-rust/crates/redis-connect/src/sources/mod.rs create mode 100644 litellm-rust/crates/redis-connect/tests/connect.rs create mode 100644 litellm-rust/crates/redis-connect/tests/live/aws.rs create mode 100644 litellm-rust/crates/redis-connect/tests/live/azure.rs create mode 100644 litellm-rust/crates/redis-connect/tests/live/gcp.rs create mode 100644 litellm-rust/crates/redis-connect/tests/live/main.rs create mode 100644 litellm-rust/crates/redis-resolve/AGENTS.md create mode 100644 litellm-rust/crates/redis-resolve/Cargo.toml create mode 100644 litellm-rust/crates/redis-resolve/redis-py-kwargs.json create mode 100644 litellm-rust/crates/redis-resolve/redis_py_kwargs.py create mode 100644 litellm-rust/crates/redis-resolve/src/auth.rs create mode 100644 litellm-rust/crates/redis-resolve/src/endpoint.rs create mode 100644 litellm-rust/crates/redis-resolve/src/error.rs create mode 100644 litellm-rust/crates/redis-resolve/src/keys.rs create mode 100644 litellm-rust/crates/redis-resolve/src/lib.rs create mode 100644 litellm-rust/crates/redis-resolve/src/options.rs create mode 100644 litellm-rust/crates/redis-resolve/src/tls.rs create mode 100644 litellm-rust/crates/redis-resolve/tests/catalog.rs create mode 100644 litellm-rust/crates/redis-resolve/tests/interpret.rs create mode 100644 litellm-rust/crates/redis-store/AGENTS.md create mode 100644 litellm-rust/crates/redis-store/Cargo.toml create mode 100644 litellm-rust/crates/redis-store/src/codec.rs create mode 100644 litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs create mode 100644 litellm-rust/crates/redis-store/src/dual_cache/increment.rs create mode 100644 litellm-rust/crates/redis-store/src/dual_cache/mod.rs create mode 100644 litellm-rust/crates/redis-store/src/dual_cache/read.rs create mode 100644 litellm-rust/crates/redis-store/src/dual_cache/write.rs create mode 100644 litellm-rust/crates/redis-store/src/error.rs create mode 100644 litellm-rust/crates/redis-store/src/failure_log.rs create mode 100644 litellm-rust/crates/redis-store/src/lib.rs create mode 100644 litellm-rust/crates/redis-store/src/memory.rs create mode 100644 litellm-rust/crates/redis-store/src/number.rs create mode 100644 litellm-rust/crates/redis-store/src/remote.rs create mode 100644 litellm-rust/crates/redis-store/src/request.rs create mode 100644 litellm-rust/crates/redis-store/src/value.rs create mode 100644 litellm-rust/crates/redis-store/tests/request.rs create mode 100644 litellm-rust/crates/redis-store/tests/support/mod.rs create mode 100644 litellm-rust/crates/redis-testing/src/deployment.rs create mode 100644 litellm-rust/crates/redis-testing/src/exec.rs create mode 100644 litellm-rust/crates/redis-testing/src/proxy.rs create mode 100644 litellm-rust/crates/redis-testing/src/sentinel.rs create mode 100644 litellm-rust/crates/redis-testing/src/tls.rs create mode 100644 litellm-rust/crates/redis-types/AGENTS.md create mode 100644 litellm-rust/crates/redis-types/Cargo.toml create mode 100644 litellm-rust/crates/redis-types/src/lib.rs create mode 100644 litellm/redis/__init__.py create mode 100644 litellm/redis/store_settings.py create mode 100644 litellm/redis/stores/__init__.py create mode 100644 litellm/redis/stores/_route.py create mode 100644 litellm/redis/stores/connection.py create mode 100644 litellm/rust_bridge/host/redis_credentials.py create mode 100644 litellm/rust_bridge/host/redis_events.py create mode 100644 tests/unit/redis/__init__.py create mode 100644 tests/unit/redis/stores/__init__.py create mode 100644 tests/unit/redis/stores/test_connection.py create mode 100644 tests/unit/rust_bridge/host/test_redis_credentials.py create mode 100644 tests/unit/rust_bridge/host/test_redis_events.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 8ac157c94f5..41198887cab 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -182,6 +182,7 @@ jobs: tests/unit/google_genai tests/unit/router_strategy tests/unit/router_utils + tests/unit/redis tests/unit/proxy/common_utils/test_cache_aware_routing.py tests/unit/enterprise/enterprise_callbacks/send_emails tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py diff --git a/litellm-rust/.agents/skills/rust-python-boundary/SKILL.md b/litellm-rust/.agents/skills/rust-python-boundary/SKILL.md new file mode 100644 index 00000000000..5ead8f287ab --- /dev/null +++ b/litellm-rust/.agents/skills/rust-python-boundary/SKILL.md @@ -0,0 +1,53 @@ +--- +name: rust-python-boundary +description: Decide where Python-shaped input becomes typed Rust in litellm-rust. Use when reading config.yaml values, os.environ/ references, REDIS_* or other environment layering, get_secret fallbacks, or kwargs and objects handed over from Python, when coercing text to a bool, int or float the way Python does, or when adding a dependency on litellm-python-compat +metadata: + internal: true +--- + +# Python boundary + +Python-shaped input becomes typed Rust in exactly two crates. Every crate past them handles typed values, so no crate re-guesses how Python would have read a value + +## The two boundaries + +- `crates/config` reads config.yaml, which is Python-shaped: `os.environ/` references, scalars that Helm or the environment turned into text, raw kwargs maps. It also owns the environment layering and `get_secret` fallbacks Python applies while it loads that config. It has no Python runtime, so it reproduces Python value semantics with `litellm-python-compat` +- `crates/python-bridge` receives live Python objects, such as `RedisCache.redis_kwargs`. It reads them through PyO3 with the real protocol (`is_truthy()`, `str()`, `extract`), so it never imitates Python +- Both produce the same typed input, defined in the data crate the consumer already depends on (`redis-types` for Redis), so a downstream crate cannot tell which host it runs under + +## Python value semantics + +Reading a value the way Python itself reads it. This lives only at the two boundaries + +- Truthiness and `x or fallback`, `str(x)`, `int(text)` and `float(text)` +- Each surface's true and false text: `_str_to_bool`, pydantic's lax bool, `get_secret`'s `true` and `false` +- `os.environ/` references and `REDIS_` style environment layering + +A boundary knows which surface a value came from. `general_settings.coordination_redis` is a pydantic model and `cache_params` goes through `_str_to_bool`, so the boundary applies that surface's rule exactly instead of refusing the text two surfaces disagree on + +Each typed field keeps where its value came from (the key, `REDIS_` or the fallback name), so an error or warning downstream still names it + +## What stays downstream + +A crate that stands in for a Python library keeps that library's semantics, because they are its compatibility contract and not value coercion. `redis-resolve` stands in for redis-py, so it keeps + +- Which target wins, and the defaults redis-py applies +- How `parse_url` reads a url, including `int()` on `?db=` and redis-py's `to_bool` on the TLS query parameters +- Which kwargs redis-py accepts, and whether the Rust client reads, warns about or refuses each one + +Its input is typed. A rule it applies reads typed fields, never a `serde_json::Value` it has to coerce first. A value it only carries back to Python untouched, like `passthrough`, may stay untyped + +## litellm-python-compat + +The rule is about where config and kwargs get coerced, not about who may depend on the crate + +- Its data formats (`pickle`, `json.dumps`, `repr`) are fine anywhere Rust reads or writes state Python also reads or writes, such as cache values in a shared Redis +- Its coercions (`truthy`, `number`) belong at the two boundaries, and in a crate that stands in for a Python library when that library calls them itself + +## Existing code + +Some code still coerces Python-shaped values past the boundary, for example the `REDIS_*` merge in `gateway/src/redis/layers.rs` and the flag and number readers in `redis-resolve/src/options.rs`. Do not add more. When you change one, move the coercion to its boundary instead of extending it in place + +## Tests + +Pin a boundary's coercion with named cases next to the Python function it mirrors, in the Rust test and in its Python twin, like the Redis contract cases in `redis-resolve/tests/interpret.rs` and `tests/unit/test_redis.py`. A downstream crate's tests build its typed input directly and never feed it raw text to coerce diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index 5eaf876e3b2..5fcf05f61f7 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -8,6 +8,10 @@ For string-valued enums and their Serde conversions, follow [.agents/skills/rust For error types, error variants and conversions between crate errors, follow [.agents/skills/rust-error/SKILL.md](.agents/skills/rust-error/SKILL.md) +For where config.yaml values, environment layering and kwargs from Python become typed Rust, follow [.agents/skills/rust-python-boundary/SKILL.md](.agents/skills/rust-python-boundary/SKILL.md) + +The process-wide rustls `CryptoProvider` (ring) is installed once by the host. rust-proxy does it at the top of `main` through `litellm_gateway::install_crypto_provider`, and python-bridge must do the same in its module init before it builds any rustls config. Library crates never install it. The AWS SDK links aws-lc-rs next to ring, so without an installed default the first rustls config a library builds, such as redis-rs TLS, panics + For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array. Derive Strum string conversions instead of hand-written variant-to-string matches Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all` diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 107b55e89cb..da8f2690649 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -8,6 +8,17 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures 0.2.17", +] + [[package]] name = "ahash" version = "0.8.12" @@ -982,6 +993,15 @@ dependencies = [ "hybrid-array", ] +[[package]] +name = "block-padding" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8894febbff9f758034a5b8e12d87918f56dfc64a8e1fe757d65e29041538d93" +dependencies = [ + "generic-array", +] + [[package]] name = "bollard" version = "0.20.2" @@ -1134,6 +1154,15 @@ dependencies = [ "rustversion", ] +[[package]] +name = "cbc" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26b52a9543ae338f279b96b0b9fed9c8093744685043739079ce85cd58f289a6" +dependencies = [ + "cipher", +] + [[package]] name = "cc" version = "1.3.0" @@ -1210,6 +1239,16 @@ dependencies = [ "half", ] +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common 0.1.7", + "inout", +] + [[package]] name = "clap" version = "4.6.7" @@ -1774,6 +1813,15 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "des" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffdd80ce8ce993de27e9f063a444a4d53ce8e8db4c1f00cc03af5ad5a9867a1e" +dependencies = [ + "cipher", +] + [[package]] name = "digest" version = "0.10.7" @@ -3105,6 +3153,16 @@ dependencies = [ "serde_core", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "block-padding", + "generic-array", +] + [[package]] name = "ipnet" version = "2.12.0" @@ -3663,8 +3721,7 @@ dependencies = [ "strum", "tempfile", "thiserror 2.0.19", - "tracing", - "tracing-subscriber", + "veil", ] [[package]] @@ -3767,6 +3824,7 @@ dependencies = [ "http-body-util", "litellm-circuit-breaker", "litellm-config", + "litellm-core-utils", "litellm-error", "litellm-gateway-auth", "litellm-gateway-forward", @@ -3777,11 +3835,15 @@ dependencies = [ "litellm-inference", "litellm-llms", "litellm-redis-client", + "litellm-redis-connect", + "litellm-redis-resolve", "litellm-redis-testing", + "litellm-redis-types", "litellm-secret-value", "litellm-secrets", "litellm-tracing", "rstest", + "rustls 0.23.42", "serde", "serde_json", "tempfile", @@ -3792,6 +3854,7 @@ dependencies = [ "tower-sessions-moka-store", "tracing", "uuid", + "veil", ] [[package]] @@ -4323,6 +4386,7 @@ dependencies = [ "litellm-cache", "litellm-cache-response", "litellm-callbacks-legacy-python", + "litellm-circuit-breaker", "litellm-core-utils", "litellm-error", "litellm-host", @@ -4337,6 +4401,12 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-model-catalog", + "litellm-python-compat", + "litellm-redis-client", + "litellm-redis-connect", + "litellm-redis-resolve", + "litellm-redis-store", + "litellm-redis-types", "litellm-router-core", "litellm-router-types", "litellm-secret-value", @@ -4387,23 +4457,108 @@ version = "0.1.0" dependencies = [ "futures-util", "litellm-circuit-breaker", + "litellm-python-compat", "litellm-redis-testing", + "litellm-redis-types", "litellm-secret-value", + "pem", + "pkcs8", "redis", "rstest", + "rustls 0.23.42", + "rustls-native-certs", + "socket2 0.6.5", "thiserror 2.0.19", "tokio", + "tracing", +] + +[[package]] +name = "litellm-redis-connect" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-auth-aws", + "litellm-auth-azure", + "litellm-auth-gcp", + "litellm-auth-types", + "litellm-circuit-breaker", + "litellm-core-utils", + "litellm-error", + "litellm-http", + "litellm-redis-client", + "litellm-redis-resolve", + "litellm-redis-testing", + "litellm-redis-types", + "litellm-secret-value", + "rstest", + "rustls 0.23.42", + "serde", + "serde_json", + "strum", + "thiserror 2.0.19", + "time", + "tokio", + "tracing", + "url", + "veil", +] + +[[package]] +name = "litellm-redis-resolve" +version = "0.1.0" +dependencies = [ + "litellm-python-compat", + "litellm-redis-types", + "litellm-secret-value", + "rstest", + "serde", + "serde_json", + "strum", + "thiserror 2.0.19", + "tracing", + "url", +] + +[[package]] +name = "litellm-redis-store" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-cache", + "litellm-cache-memory", + "litellm-circuit-breaker", + "litellm-python-compat", + "litellm-redis-client", + "litellm-redis-testing", + "litellm-secret-value", + "litellm-tracing", + "rstest", + "strum", + "thiserror 2.0.19", + "tokio", + "tracing", ] [[package]] name = "litellm-redis-testing" version = "0.1.0" dependencies = [ + "pkcs8", + "rcgen", + "tempfile", "testcontainers-modules", "thiserror 2.0.19", "tokio", ] +[[package]] +name = "litellm-redis-types" +version = "0.1.0" +dependencies = [ + "litellm-secret-value", +] + [[package]] name = "litellm-router" version = "0.1.0" @@ -5286,6 +5441,16 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" +[[package]] +name = "pbkdf2" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2" +dependencies = [ + "digest 0.10.7", + "hmac 0.12.1", +] + [[package]] name = "pem" version = "4.0.0" @@ -5396,6 +5561,22 @@ dependencies = [ "spki", ] +[[package]] +name = "pkcs5" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6" +dependencies = [ + "aes", + "cbc", + "der", + "des", + "pbkdf2", + "scrypt", + "sha2 0.10.9", + "spki", +] + [[package]] name = "pkcs8" version = "0.10.2" @@ -5403,6 +5584,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" dependencies = [ "der", + "pkcs5", + "rand_core 0.6.4", "spki", ] @@ -6411,6 +6594,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "salsa20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213" +dependencies = [ + "cipher", +] + [[package]] name = "same-file" version = "1.0.6" @@ -6474,6 +6666,17 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "scrypt" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0516a385866c09368f0b5bcd1caff3366aace790fcd46e2bb032697bb172fd1f" +dependencies = [ + "pbkdf2", + "salsa20", + "sha2 0.10.9", +] + [[package]] name = "sct" version = "0.7.1" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 7121df06bc0..72d6694ba79 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -73,7 +73,11 @@ litellm-host-python = { path = "crates/host-python" } litellm-python-compat = { path = "crates/python-compat" } litellm-circuit-breaker = { path = "crates/circuit-breaker" } litellm-redis-client = { path = "crates/redis-client" } +litellm-redis-connect = { path = "crates/redis-connect" } +litellm-redis-store = { path = "crates/redis-store" } litellm-redis-testing = { path = "crates/redis-testing" } +litellm-redis-resolve = { path = "crates/redis-resolve" } +litellm-redis-types = { path = "crates/redis-types" } litellm-cost = { path = "crates/cost" } litellm-model-catalog = { path = "crates/model-catalog" } litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" } diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 73fbf14b07a..809ffe3044e 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -2,7 +2,10 @@ use std::{ cmp::Reverse, collections::{BinaryHeap, HashMap, HashSet}, hash::Hash, - sync::{Arc, Mutex, MutexGuard}, + sync::{ + Arc, Mutex, MutexGuard, + atomic::{AtomicUsize, Ordering}, + }, time::{Duration, SystemTime, UNIX_EPOCH}, }; @@ -31,7 +34,7 @@ struct CacheState { pub struct InMemoryCache { state: Mutex>, - max_size_in_memory: usize, + max_size_in_memory: AtomicUsize, default_ttl: Duration, max_entry_bytes: Option, measure_value: Option>, @@ -74,7 +77,9 @@ impl InMemoryCache { expirations: HashMap::new(), expiration_heap: BinaryHeap::new(), }), - max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), + max_size_in_memory: AtomicUsize::new( + max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), + ), default_ttl: default_ttl .filter(|ttl| !ttl.is_zero()) .unwrap_or(DEFAULT_TTL), @@ -94,7 +99,7 @@ impl InMemoryCache { value: V, ttl: Option, ) -> Result { - if self.max_size_in_memory == 0 { + if self.max_size_in_memory() == 0 { return Ok(CacheWrite::Disabled); } let now = (self.now)(); @@ -129,7 +134,7 @@ impl InMemoryCache { pub fn evict_cache(&self) -> Result<(), Error> { let now = (self.now)(); let mut state = self.lock()?; - Self::evict(&mut state, self.max_size_in_memory, now, None); + Self::evict(&mut state, self.max_size_in_memory(), now, None); Ok(()) } @@ -170,7 +175,16 @@ impl InMemoryCache { } pub fn max_size_in_memory(&self) -> usize { - self.max_size_in_memory + self.max_size_in_memory.load(Ordering::Relaxed) + } + + /// `DualCache.update_in_memory_max_size`: `None` restores the default. Entries already held stay + /// until the next write evicts down to the new size. + pub fn update_max_size_in_memory(&self, max_size_in_memory: Option) { + self.max_size_in_memory.store( + max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), + Ordering::Relaxed, + ); } pub fn max_entry_bytes(&self) -> Option { @@ -200,6 +214,23 @@ impl InMemoryCache { .collect()) } + /// Every live entry, for callers that inspect the whole tier. + pub fn entries(&self) -> Result, Error> { + let now = (self.now)(); + let state = self.lock()?; + Ok(state + .values + .iter() + .filter(|(key, _)| { + state + .expirations + .get(*key) + .is_none_or(|expiration| *expiration >= now) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect()) + } + pub fn delete_cache(&self, key: &str) -> Result<(), Error> { let mut state = self.lock()?; Self::remove(&mut state, key); @@ -276,7 +307,7 @@ impl InMemoryCache { ttl: Option, now: Duration, ) -> Result { - Self::evict(state, self.max_size_in_memory, now, None); + Self::evict(state, self.max_size_in_memory(), now, None); if !self.check_value_size(&value)? { return Ok(CacheWrite::TooLarge); } @@ -300,12 +331,12 @@ where eligible: &[V], context: ExactCacheContext, ) -> Result { - if self.max_size_in_memory == 0 { + if self.max_size_in_memory() == 0 { return Ok(candidate); } let now = (self.now)(); let mut state = self.lock()?; - Self::evict(&mut state, self.max_size_in_memory, now, Some(key)); + Self::evict(&mut state, self.max_size_in_memory(), now, Some(key)); let existing = state .values .get(key) @@ -335,7 +366,7 @@ impl CounterCache for InMemoryCache { amount: f64, context: ExactCacheContext, ) -> Result { - if self.max_size_in_memory == 0 { + if self.max_size_in_memory() == 0 { return Ok(amount); } let now = (self.now)(); @@ -408,7 +439,7 @@ where values: Vec, ttl: Option, ) -> Result { - if self.max_size_in_memory == 0 { + if self.max_size_in_memory() == 0 { return Ok(values); } let now = (self.now)(); diff --git a/litellm-rust/crates/config/AGENTS.md b/litellm-rust/crates/config/AGENTS.md new file mode 100644 index 00000000000..9f6f1db9a67 --- /dev/null +++ b/litellm-rust/crates/config/AGENTS.md @@ -0,0 +1,3 @@ +## Redis settings + +- This crate only deserializes YAML. `RedisSettings` is the raw `coordination_redis` or `cache_params` map and `CacheSettings` tells a `type: redis` cache apart from the rest. Reading those maps the way `litellm/_redis.py` does belongs to `gateway/src/redis`, and turning merged kwargs into a `Resolution` belongs to `redis-resolve` diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml index 278aea6200e..a433aacc1c9 100644 --- a/litellm-rust/crates/config/Cargo.toml +++ b/litellm-rust/crates/config/Cargo.toml @@ -6,17 +6,16 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-secret-value.workspace = true litellm-router-types.workspace = true +litellm-secret-value.workspace = true serde.workspace = true -serde_json.workspace = true serde_yaml_ng = "0.10.0" strum.workspace = true thiserror.workspace = true -tracing.workspace = true +veil.workspace = true [dev-dependencies] litellm-error.workspace = true rstest.workspace = true +serde_json.workspace = true tempfile.workspace = true -tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } diff --git a/litellm-rust/crates/config/src/error.rs b/litellm-rust/crates/config/src/error.rs index e5219ac992c..300f15380bd 100644 --- a/litellm-rust/crates/config/src/error.rs +++ b/litellm-rust/crates/config/src/error.rs @@ -11,11 +11,4 @@ pub enum Error { path: PathBuf, source: serde_yaml_ng::Error, }, - #[error("Redis host does not yet support configured setting {0}")] - RedisSetting(String), - #[error("invalid Redis setting {setting}: {reason}")] - RedisInvalid { - setting: &'static str, - reason: &'static str, - }, } diff --git a/litellm-rust/crates/config/src/lib.rs b/litellm-rust/crates/config/src/lib.rs index b8af6f11d2e..46657aa91e7 100644 --- a/litellm-rust/crates/config/src/lib.rs +++ b/litellm-rust/crates/config/src/lib.rs @@ -14,10 +14,7 @@ pub use error::Error; pub use litellm_router_types::{LitellmParams, Spelled}; pub use mcp::{McpAuth, McpServer, McpTransport}; pub use model::Model; -pub use redis::{ - CacheSettings, Lookup, RedisCluster, RedisEndpoint, RedisHost, RedisNode, RedisSettings, - ResolvedRedis, environment_has_target, -}; +pub use redis::{CacheSettings, RedisSettings}; pub use settings::{ ClickHouseStoreSettings, GeneralSettings, LiteLlmSettings, RouterSettings, TracingSettings, TracingStoreSettings, diff --git a/litellm-rust/crates/config/src/redis.rs b/litellm-rust/crates/config/src/redis.rs index 7a91691c54b..d0ba107242c 100644 --- a/litellm-rust/crates/config/src/redis.rs +++ b/litellm-rust/crates/config/src/redis.rs @@ -1,195 +1,21 @@ -use std::{collections::BTreeMap, fmt, time::Duration}; +use std::collections::BTreeMap; -use litellm_router_types::Spelled; -use litellm_secret_value::SecretValue; use serde::{Deserialize, Deserializer, de::Error as _}; +use veil::Redact; -use crate::{AdditionalFields, Error, Object, Value}; +use crate::{Object, Value}; -pub type Lookup<'a> = &'a dyn Fn(&str) -> Option; - -const SUPPORTED: &[&str] = &[ - "url", - "host", - "port", - "username", - "password", - "ssl", - "db", - "socket_timeout", - "startup_nodes", - "namespace", -]; - -const IGNORED: &[&str] = &[ - "max_connections", - "single_connection_client", - "health_check_interval", - "socket_keepalive", - "socket_keepalive_options", - "decode_responses", - "encoding", - "encoding_errors", - "charset", - "errors", - "client_name", - "lib_name", - "lib_version", - "sentinel_password", - "gcp_ssl_ca_certs", - "azure_client_id", - "azure_tenant_id", - "azure_client_secret", - "aws_iam_user_name", - "aws_iam_cache_name", - "aws_iam_region", - "aws_iam_serverless", -]; - -const WARN_IGNORED: &[&str] = &[ - "max_connections", - "health_check_interval", - "socket_keepalive", - "socket_keepalive_options", -]; - -const REDIS_PY_KWARGS: &[&str] = &[ - "aws_iam_auth", - "aws_iam_cache_name", - "aws_iam_region", - "aws_iam_serverless", - "aws_iam_user_name", - "azure_client_id", - "azure_client_secret", - "azure_redis_ad_token", - "azure_tenant_id", - "cache", - "cache_config", - "charset", - "client_name", - "db", - "decode_responses", - "encoding", - "encoding_errors", - "errors", - "event_dispatcher", - "gcp_service_account", - "gcp_ssl_ca_certs", - "health_check_interval", - "host", - "lib_name", - "lib_version", - "max_connections", - "password", - "port", - "protocol", - "redis_connect_func", - "retry_on_error", - "retry_on_timeout", - "single_connection_client", - "socket_connect_timeout", - "socket_keepalive", - "socket_keepalive_options", - "socket_timeout", - "ssl", - "ssl_ca_certs", - "ssl_ca_data", - "ssl_ca_path", - "ssl_cert_reqs", - "ssl_certfile", - "ssl_check_hostname", - "ssl_ciphers", - "ssl_keyfile", - "ssl_min_version", - "ssl_ocsp_context", - "ssl_ocsp_expected_cert", - "ssl_password", - "ssl_validate_ocsp", - "ssl_validate_ocsp_stapled", - "unix_socket_path", - "url", - "username", -]; - -const UNSUPPORTED_EXTRA: &[&str] = &[ - "credential_provider", - "sentinel_nodes", - "retry", - "connection_pool", -]; - -const CACHE_ONLY: &[&str] = &[ - "type", - "mode", - "ttl", - "default_in_memory_ttl", - "default_in_redis_ttl", - "supported_call_types", - "similarity_threshold", - "max_messages", - "_backend", -]; - -#[derive(Clone, Default)] +/// A Redis section as `config.yaml` spells it, before any surface reads it. +#[derive(Clone, Default, Deserialize, Redact)] +#[serde(transparent)] pub struct RedisSettings { - pub url: Option, - pub host: Option>, - pub port: Option>, - pub username: Option, - pub password: Option, - pub ssl: Option>, - pub db: Option>, - pub socket_timeout: Option>, - pub startup_nodes: Option, - pub namespace: Option>, - pub additional_fields: AdditionalFields, - unparsed: DeferredInputs, + #[redact(with = "[REDACTED]")] + fields: BTreeMap, } -impl<'de> Deserialize<'de> for RedisSettings { - fn deserialize>(deserializer: D) -> Result { - let fields = BTreeMap::::deserialize(deserializer)?; - let scalar = |key: &str| fields.get(key).and_then(as_text); - Ok(Self { - url: scalar("url").map(SecretValue::new), - host: scalar("host").map(Spelled::Value), - port: spelled(scalar("port")), - username: scalar("username").map(SecretValue::new), - password: scalar("password").map(SecretValue::new), - ssl: spelled(scalar("ssl")), - db: spelled(scalar("db")), - socket_timeout: spelled(scalar("socket_timeout")), - startup_nodes: fields - .get("startup_nodes") - .filter(|value| !value.is_null()) - .cloned(), - namespace: scalar("namespace").map(Spelled::Value), - additional_fields: cloned(&fields, |key, _| !SUPPORTED.contains(&key)), - unparsed: deferred_inputs(&fields), - }) - } -} - -impl fmt::Debug for RedisSettings { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("RedisSettings") - .field("url", &self.url) - .field("host", &self.host) - .field("port", &self.port) - .field("username", &self.username) - .field("password", &self.password) - .field("ssl", &self.ssl) - .field("db", &self.db) - .field("socket_timeout", &self.socket_timeout) - .field( - "startup_nodes", - &self.startup_nodes.as_ref().map(|_| ""), - ) - .field("namespace", &self.namespace) - .field("additional_fields", &self.additional_fields.keys()) - .field("unparsed", &self.unparsed.keys()) - .finish() +impl RedisSettings { + pub fn fields(&self) -> &BTreeMap { + &self.fields } } @@ -216,452 +42,22 @@ impl<'de> Deserialize<'de> for CacheSettings { } } -#[derive(Clone, PartialEq)] -pub struct ResolvedRedis { - pub endpoint: RedisEndpoint, - pub namespace: Option, - pub timeout: Option, -} +#[cfg(test)] +mod tests { + use super::RedisSettings; + use rstest::rstest; -impl fmt::Debug for ResolvedRedis { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("ResolvedRedis") - .field("endpoint", &self.endpoint) - .field("namespace", &self.namespace) - .field("timeout", &self.timeout) - .finish() + #[rstest] + #[case::scalar(serde_json::json!("credential-value"))] + #[case::object(serde_json::json!({"secret": "credential-value"}))] + #[case::array(serde_json::json!(["credential-value"]))] + fn credentials_are_redacted(#[case] password: serde_json::Value) { + let settings: RedisSettings = serde_json::from_value(serde_json::json!({ + "host": "redis", "password": password, + })) + .unwrap(); + let debug = format!("{settings:?}"); + assert!(!debug.contains("credential-value")); + assert!(debug.contains("[REDACTED]")); } } - -#[derive(Clone, PartialEq)] -pub enum RedisEndpoint { - Url(SecretValue), - Host(RedisHost), - Cluster(RedisCluster), -} - -impl fmt::Debug for RedisEndpoint { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Url(_) => formatter.debug_tuple("Url").field(&"").finish(), - Self::Host(host) => host.fmt(formatter), - Self::Cluster(cluster) => cluster.fmt(formatter), - } - } -} - -#[derive(Clone, Debug, PartialEq)] -pub struct RedisHost { - pub host: String, - pub port: u16, - pub username: Option, - pub password: Option, - pub tls: bool, - pub db: i64, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct RedisCluster { - pub nodes: Vec, - pub username: Option, - pub password: Option, - pub tls: bool, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct RedisNode { - pub host: String, - pub port: u16, -} - -#[derive(Clone)] -struct Entry { - value: Value, - source: String, -} - -type Kwargs = BTreeMap; - -impl RedisSettings { - pub fn resolve(&self, lookup: Lookup<'_>) -> Result { - self.resolve_with(lookup, false) - } - - pub fn resolve_cache(&self, lookup: Lookup<'_>) -> Result { - self.resolve_with(lookup, true) - } - - fn resolve_with(&self, lookup: Lookup<'_>, cache: bool) -> Result { - let block = self.kwargs(cache); - let block: Kwargs = block - .into_iter() - .filter_map(|(key, entry)| { - let value = environment_reference(&entry.value, lookup)?; - Some((key, Entry { value, ..entry })) - }) - .collect(); - let environment = environment_kwargs(lookup); - let names_a_target = ["host", "startup_nodes", "sentinel_nodes"] - .iter() - .any(|key| block.contains_key(*key)); - let drop_environment_url = names_a_target && !block.contains_key("url"); - let merged: Kwargs = environment - .into_iter() - .filter(|(key, _)| !(key == "url" && drop_environment_url)) - .chain(block) - .collect(); - reject_unsupported(&merged, lookup)?; - warn_ignored(&merged); - - let startup_nodes = match merged.get("startup_nodes") { - Some(entry) => Some(cluster_nodes(&entry.value)?), - None => lookup("REDIS_CLUSTER_NODES") - .map(|nodes| cluster_nodes(&Value::String(nodes))) - .transpose()?, - }; - let username = secret_value(&merged, "username"); - let password = secret_value(&merged, "password"); - let tls = bool_value(&merged, "ssl"); - let endpoint = match (startup_nodes, text_value(&merged, "url")) { - (Some(nodes), _) if !nodes.is_empty() => RedisEndpoint::Cluster(RedisCluster { - nodes, - username, - password, - tls, - }), - (_, Some(url)) => RedisEndpoint::Url(SecretValue::new(url)), - _ => RedisEndpoint::Host(RedisHost { - host: text_value(&merged, "host").ok_or(Error::RedisInvalid { - setting: "host", - reason: "either host or url must be specified for redis", - })?, - port: integer(&merged, "port")?.unwrap_or(6379), - username, - password, - tls, - db: integer(&merged, "db")?.unwrap_or(0), - }), - }; - let timeout = merged - .get("socket_timeout") - .and_then(|entry| as_text(&entry.value)) - .map(|seconds| { - seconds - .parse::() - .ok() - .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) - .ok_or(Error::RedisInvalid { - setting: "socket_timeout", - reason: "must be a non-negative number of seconds", - }) - }) - .transpose()?; - Ok(ResolvedRedis { - endpoint, - namespace: text_value(&merged, "namespace"), - timeout, - }) - } - - fn kwargs(&self, cache: bool) -> Kwargs { - let known = [ - ( - "url", - self.url.as_ref().map(|value| value.expose().to_owned()), - ), - ("host", spelled_text(self.host.as_ref())), - ("port", spelled_text(self.port.as_ref())), - ( - "username", - self.username - .as_ref() - .map(|value| value.expose().to_owned()), - ), - ( - "password", - self.password - .as_ref() - .map(|value| value.expose().to_owned()), - ), - ("ssl", spelled_text(self.ssl.as_ref())), - ("db", spelled_text(self.db.as_ref())), - ("socket_timeout", spelled_text(self.socket_timeout.as_ref())), - ("namespace", spelled_text(self.namespace.as_ref())), - ]; - let known = known.into_iter().filter_map(|(key, value)| { - value.map(|value| { - ( - key.to_owned(), - Entry { - value: Value::String(value), - source: key.to_owned(), - }, - ) - }) - }); - let startup_nodes = self.startup_nodes.clone().map(|value| { - ( - "startup_nodes".to_owned(), - Entry { - value, - source: "startup_nodes".to_owned(), - }, - ) - }); - let extras = self.additional_fields.iter().filter_map(|(key, value)| { - if cache && CACHE_ONLY.contains(&key.as_str()) { - return None; - } - let key = match (cache, key.as_str()) { - (true, "redis_startup_nodes") => "startup_nodes".to_owned(), - _ => key.clone(), - }; - Some(( - key.clone(), - Entry { - value: value.clone(), - source: key, - }, - )) - }); - let unparsed = self.unparsed.iter().map(|(key, value)| { - ( - key.clone(), - Entry { - value: value.0.clone(), - source: key.clone(), - }, - ) - }); - extras - .chain(unparsed) - .chain(known) - .chain(startup_nodes) - .collect() - } -} - -pub fn environment_has_target(lookup: Lookup<'_>) -> bool { - [ - "REDIS_HOST", - "REDIS_URL", - "REDIS_CLUSTER_NODES", - "REDIS_SENTINEL_NODES", - ] - .iter() - .any(|name| lookup(name).is_some()) -} - -fn reject_unsupported(merged: &Kwargs, lookup: Lookup<'_>) -> Result<(), Error> { - let active_auth = [ - merged - .get("gcp_service_account") - .map(|entry| entry.source.as_str()) - .or_else(|| lookup("REDIS_GCP_SERVICE_ACCOUNT").map(|_| "REDIS_GCP_SERVICE_ACCOUNT")), - text_value(merged, "azure_redis_ad_token") - .or_else(|| lookup("REDIS_AZURE_AD_TOKEN")) - .filter(|value| value.eq_ignore_ascii_case("true")) - .map(|_| { - merged - .get("azure_redis_ad_token") - .map_or("REDIS_AZURE_AD_TOKEN", |entry| entry.source.as_str()) - }), - bool_value(merged, "aws_iam_auth").then(|| merged["aws_iam_auth"].source.as_str()), - merged - .get("sentinel_nodes") - .map(|entry| entry.source.as_str()) - .or_else(|| lookup("REDIS_SENTINEL_NODES").map(|_| "REDIS_SENTINEL_NODES")), - ]; - if let Some(setting) = active_auth.into_iter().flatten().next() { - return Err(Error::RedisSetting(setting.to_owned())); - } - match merged - .iter() - .find(|(key, _)| is_unsupported(key)) - .map(|(_, entry)| entry.source.clone()) - { - Some(setting) => Err(Error::RedisSetting(setting)), - None => Ok(()), - } -} - -fn warn_ignored(merged: &Kwargs) { - for setting in WARN_IGNORED { - if merged.contains_key(*setting) { - tracing::warn!(setting, "Redis setting is accepted but ignored"); - } - } -} - -fn is_unsupported(key: &str) -> bool { - let known = REDIS_PY_KWARGS.contains(&key) || UNSUPPORTED_EXTRA.contains(&key); - let handled = SUPPORTED.contains(&key) - || IGNORED.contains(&key) - || matches!( - key, - "gcp_service_account" | "azure_redis_ad_token" | "aws_iam_auth" - ); - known && !handled -} - -fn environment_kwargs(lookup: Lookup<'_>) -> Kwargs { - REDIS_PY_KWARGS - .iter() - .filter_map(|key| { - let source = format!("REDIS_{}", key.to_ascii_uppercase()); - let value = lookup(&source)?; - Some(( - (*key).to_owned(), - Entry { - value: Value::String(value), - source, - }, - )) - }) - .collect() -} - -fn environment_reference(value: &Value, lookup: Lookup<'_>) -> Option { - match value { - Value::String(text) => match text.strip_prefix("os.environ/") { - Some(name) => lookup(name).map(Value::String), - None => Some(value.clone()), - }, - Value::Null => None, - _ => Some(value.clone()), - } -} - -fn as_text(value: &Value) -> Option { - match value { - Value::String(text) => Some(text.clone()), - Value::Number(number) => Some(number.to_string()), - Value::Bool(flag) => Some(flag.to_string()), - _ => None, - } -} - -fn text_value(kwargs: &Kwargs, key: &str) -> Option { - kwargs.get(key).and_then(|entry| as_text(&entry.value)) -} - -fn secret_value(kwargs: &Kwargs, key: &str) -> Option { - text_value(kwargs, key).map(SecretValue::new) -} - -fn bool_value(kwargs: &Kwargs, key: &str) -> bool { - match kwargs.get(key).map(|entry| &entry.value) { - Some(Value::Bool(flag)) => *flag, - Some(value) => as_text(value) - .is_some_and(|text| matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes")), - None => false, - } -} - -fn integer(kwargs: &Kwargs, key: &'static str) -> Result, Error> { - kwargs - .get(key) - .and_then(|entry| as_text(&entry.value)) - .map(|text| { - text.parse().map_err(|_| Error::RedisInvalid { - setting: key, - reason: "must be an integer", - }) - }) - .transpose() -} - -const INVALID_NODES: Error = Error::RedisInvalid { - setting: "startup_nodes", - reason: "must be a list of {host, port} objects", -}; - -fn cluster_nodes(value: &Value) -> Result, Error> { - match value { - Value::String(json) => { - let nodes: Vec> = - serde_json::from_str(json).map_err(|_| INVALID_NODES)?; - nodes - .iter() - .map(|node| { - cluster_node( - node.get("host").and_then(json_text), - node.get("port").and_then(json_text), - ) - }) - .collect() - } - Value::Sequence(nodes) => nodes - .iter() - .map(|node| { - cluster_node( - node.get("host").and_then(as_text), - node.get("port").and_then(as_text), - ) - }) - .collect(), - _ => Err(INVALID_NODES), - } -} - -fn cluster_node(host: Option, port: Option) -> Result { - match (host, port.and_then(|port| port.parse().ok())) { - (Some(host), Some(port)) => Ok(RedisNode { host, port }), - _ => Err(INVALID_NODES), - } -} - -fn json_text(value: &serde_json::Value) -> Option { - match value { - serde_json::Value::String(text) => Some(text.clone()), - serde_json::Value::Number(number) => Some(number.to_string()), - _ => None, - } -} - -fn spelled_text(value: Option<&Spelled>) -> Option { - value.map(|value| match value { - Spelled::Value(value) => value.to_string(), - Spelled::Text(value) => value.clone(), - }) -} - -fn spelled(text: Option) -> Option> { - text.map(|text| match text.parse() { - Ok(value) => Spelled::Value(value), - Err(_) => Spelled::Text(text), - }) -} - -fn cloned( - fields: &BTreeMap, - keep: impl Fn(&str, &Value) -> bool, -) -> AdditionalFields { - fields - .iter() - .filter(|(key, value)| keep(key, value)) - .map(|(key, value)| (key.clone(), value.clone())) - .collect() -} - -#[derive(Clone)] -struct DeferredInput(Value); - -type DeferredInputs = BTreeMap; - -fn deferred_inputs(fields: &BTreeMap) -> DeferredInputs { - fields - .iter() - .filter(|(key, value)| is_unparsed(key, value)) - .map(|(key, value)| (key.clone(), DeferredInput(value.clone()))) - .collect() -} - -fn is_unparsed(key: &str, value: &Value) -> bool { - SUPPORTED.contains(&key) - && match key { - "startup_nodes" => value.is_null(), - _ => as_text(value).is_none(), - } -} diff --git a/litellm-rust/crates/config/tests/config.rs b/litellm-rust/crates/config/tests/config.rs index 3c7d02203e5..fec042e41be 100644 --- a/litellm-rust/crates/config/tests/config.rs +++ b/litellm-rust/crates/config/tests/config.rs @@ -440,3 +440,29 @@ fn mcp_config_redacts_nested_credentials_and_preserves_policy_for_validation() { fn rejects_invalid_typed_mcp_settings(#[case] setting: &str) { assert!(Config::from_yaml(&format!("mcp_servers:\n docs:\n {setting}\n")).is_err()); } + +#[rstest] +fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() { + let config = Config::from_yaml( + "litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid", + ) + .unwrap(); + + assert!(matches!( + config.litellm_settings.cache_params, + Some(litellm_config::CacheSettings::Redis(_)) + )); +} + +#[rstest] +fn non_redis_cache_settings_remain_opaque() { + let config = Config::from_yaml( + "litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true", + ) + .unwrap(); + + assert!(matches!( + config.litellm_settings.cache_params, + Some(litellm_config::CacheSettings::Other(_)) + )); +} diff --git a/litellm-rust/crates/config/tests/redis.rs b/litellm-rust/crates/config/tests/redis.rs deleted file mode 100644 index ab409541c8e..00000000000 --- a/litellm-rust/crates/config/tests/redis.rs +++ /dev/null @@ -1,405 +0,0 @@ -use std::{ - collections::HashMap, - fmt::Write, - sync::{Arc, Mutex}, -}; - -use litellm_config::{CacheSettings, Config, Error, RedisEndpoint, RedisSettings, ResolvedRedis}; -use rstest::rstest; -use tracing::{Event, Subscriber, field::Visit}; -use tracing_subscriber::{Layer, layer::Context, prelude::*}; - -fn resolve( - yaml: &str, - environment: &[(&str, &str)], -) -> Result { - let environment: HashMap = environment - .iter() - .map(|(key, value)| ((*key).to_owned(), (*value).to_owned())) - .collect(); - let config = Config::from_yaml(yaml).unwrap(); - config - .general_settings - .coordination_redis - .as_ref() - .unwrap_or(&RedisSettings::default()) - .resolve(&|name| environment.get(name).cloned()) -} - -fn resolve_cache(yaml: &str) -> Result { - let config = Config::from_yaml(yaml).unwrap(); - match config.litellm_settings.cache_params.unwrap() { - CacheSettings::Redis(settings) => settings.resolve_cache(&|_| None), - CacheSettings::Other(_) => panic!("expected a Redis cache"), - } -} - -fn cluster_ports(resolved: &ResolvedRedis) -> Vec { - match &resolved.endpoint { - RedisEndpoint::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(), - endpoint => panic!("expected a cluster, got {endpoint:?}"), - } -} - -#[rstest] -#[case::host( - "general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5", - &[], - "host:redis:6380:2:true:app:1500" -)] -#[case::environment_url( - "model_list: []", - &[("REDIS_URL", "redis://environment:6379")], - "url:redis://environment:6379" -)] -#[case::explicit_host_drops_environment_url( - "general_settings:\n coordination_redis:\n host: configured", - &[("REDIS_URL", "redis://environment:6379")], - "host:configured:6379:0:false::none" -)] -#[case::yaml_cluster( - "general_settings:\n coordination_redis:\n startup_nodes:\n - host: node-a\n port: 7000\n - host: node-b\n port: \"7001\"", - &[], - "cluster:node-a:7000,node-b:7001" -)] -#[case::json_environment_cluster( - "model_list: []", - &[("REDIS_CLUSTER_NODES", r#"[{"host":"node-a","port":"7000"},{"host":"node-b","port":7001}]"#)], - "cluster:node-a:7000,node-b:7001" -)] -#[case::missing_reference_keeps_environment_fallback( - "general_settings:\n coordination_redis:\n host: os.environ/MISSING_HOST", - &[("REDIS_HOST", "environment")], - "host:environment:6379:0:false::none" -)] -#[case::configured_values_beat_the_same_environment_values( - "general_settings:\n coordination_redis:\n host: configured\n port: 6380\n db: 3", - &[("REDIS_HOST", "environment"), ("REDIS_PORT", "6390"), ("REDIS_DB", "9")], - "host:configured:6380:3:false::none" -)] -#[case::environment_url_beats_the_environment_host( - "model_list: []", - &[("REDIS_URL", "redis://environment:6379"), ("REDIS_HOST", "ignored")], - "url:redis://environment:6379" -)] -#[case::startup_nodes_beat_a_url( - "general_settings:\n coordination_redis:\n url: redis://single:6379\n startup_nodes:\n - host: node-a\n port: 7000", - &[], - "cluster:node-a:7000" -)] -#[case::environment_references_resolve_the_host_and_port( - "general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT", - &[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")], - "host:referenced:6400:0:false::none" -)] -#[case::a_null_host_falls_back_to_the_environment( - "general_settings:\n coordination_redis:\n host: null", - &[("REDIS_HOST", "environment")], - "host:environment:6379:0:false::none" -)] -fn resolves_redis_through_the_parsed_config( - #[case] yaml: &str, - #[case] environment: &[(&str, &str)], - #[case] expected: &str, -) { - let resolved = resolve(yaml, environment).unwrap(); - let description = match resolved.endpoint { - RedisEndpoint::Url(url) => format!("url:{}", url.expose()), - RedisEndpoint::Host(host) => format!( - "host:{}:{}:{}:{}:{}:{}", - host.host, - host.port, - host.db, - host.tls, - resolved.namespace.unwrap_or_default(), - resolved - .timeout - .map(|timeout| timeout.as_millis().to_string()) - .unwrap_or_else(|| "none".to_owned()) - ), - RedisEndpoint::Cluster(cluster) => format!( - "cluster:{}", - cluster - .nodes - .iter() - .map(|node| format!("{}:{}", node.host, node.port)) - .collect::>() - .join(",") - ), - }; - - assert_eq!(description, expected); -} - -#[rstest] -#[case::custom_tls_file( - "general_settings:\n coordination_redis:\n host: redis\n ssl_ca_certs: /ca.pem", - &[], - "ssl_ca_certs" -)] -#[case::gcp_auth_from_config( - "general_settings:\n coordination_redis:\n host: redis\n gcp_service_account: account", - &[], - "gcp_service_account" -)] -#[case::gcp_auth_from_environment( - "model_list: []", - &[("REDIS_HOST", "redis"), ("REDIS_GCP_SERVICE_ACCOUNT", "account")], - "REDIS_GCP_SERVICE_ACCOUNT" -)] -#[case::azure_auth_from_config( - "general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: true", - &[], - "azure_redis_ad_token" -)] -#[case::azure_auth_from_environment( - "model_list: []", - &[("REDIS_HOST", "redis"), ("REDIS_AZURE_AD_TOKEN", "true")], - "REDIS_AZURE_AD_TOKEN" -)] -#[case::aws_auth_from_config( - "general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true", - &[], - "aws_iam_auth" -)] -#[case::aws_auth_from_environment( - "model_list: []", - &[("REDIS_HOST", "redis"), ("REDIS_AWS_IAM_AUTH", "true")], - "REDIS_AWS_IAM_AUTH" -)] -#[case::sentinel_from_config( - "general_settings:\n coordination_redis:\n sentinel_nodes: [[redis, 26379]]", - &[], - "sentinel_nodes" -)] -#[case::sentinel_from_environment( - "model_list: []", - &[("REDIS_SENTINEL_NODES", "[[\"redis\",26379]]")], - "REDIS_SENTINEL_NODES" -)] -fn rejected_settings_name_their_source( - #[case] yaml: &str, - #[case] environment: &[(&str, &str)], - #[case] expected: &str, -) { - let error = resolve(yaml, environment).unwrap_err(); - - assert!(matches!(error, Error::RedisSetting(setting) if setting == expected)); -} - -#[rstest] -fn an_explicit_host_keeps_environment_credentials() { - let resolved = resolve( - "general_settings:\n coordination_redis:\n host: configured", - &[ - ("REDIS_PASSWORD", "environment-secret"), - ("REDIS_SSL", "True"), - ], - ) - .unwrap(); - - let RedisEndpoint::Host(host) = resolved.endpoint else { - panic!("expected a host endpoint"); - }; - assert_eq!( - host.password.as_ref().map(|password| password.expose()), - Some("environment-secret") - ); - assert!(host.tls); -} - -#[rstest] -#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])] -#[case::a_structured_host_with_an_environment_target( - "general_settings:\n coordination_redis:\n host: []", - &[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")] -)] -fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) { - let error = resolve(yaml, environment).unwrap_err(); - - assert!(matches!( - error, - Error::RedisInvalid { - setting: "host", - .. - } - )); -} - -#[rstest] -fn the_cache_alias_names_cluster_nodes() { - let resolved = resolve_cache( - "litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000", - ) - .unwrap(); - - assert_eq!(cluster_ports(&resolved), vec![7000]); -} - -#[rstest] -fn canonical_startup_nodes_beat_the_cache_alias() { - let resolved = resolve_cache( - "litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101", - ) - .unwrap(); - - assert_eq!(cluster_ports(&resolved), vec![7000]); -} - -#[rstest] -fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() { - let resolved = resolve_cache( - "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100", - ) - .unwrap(); - - assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_))); -} - -#[rstest] -fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() { - let resolved = resolve_cache( - "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100", - ) - .unwrap(); - - assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_))); -} - -#[rstest] -fn rejected_environment_settings_never_print_the_value() { - let error = resolve( - "model_list: []", - &[ - ("REDIS_HOST", "redis"), - ("REDIS_SSL_CERT_REQS", "private-option-value"), - ], - ) - .unwrap_err(); - let message = error.to_string(); - - assert!(message.contains("REDIS_SSL_CERT_REQS")); - assert!(!message.contains("private-option-value")); -} - -#[rstest] -#[case::negative("-1")] -#[case::not_a_number("slow")] -fn rejects_bad_socket_timeout(#[case] timeout: &str) { - let yaml = format!( - "general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}" - ); - let error = resolve(&yaml, &[]).unwrap_err(); - - assert!(matches!( - error, - Error::RedisInvalid { - setting: "socket_timeout", - .. - } - )); -} - -#[rstest] -fn defaults_socket_timeout_for_the_gateway() { - let resolved = resolve( - "general_settings:\n coordination_redis:\n host: redis", - &[], - ) - .unwrap(); - - assert_eq!(resolved.timeout, None); -} - -#[rstest] -fn debug_redacts_redis_secrets_and_url() { - let config = Config::from_yaml( - "general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password", - ) - .unwrap(); - let settings = config.general_settings.coordination_redis.unwrap(); - let resolved = settings.resolve(&|_| None).unwrap(); - let debug = format!("{settings:?} {resolved:?}"); - - assert!(!debug.contains("redis://user:password")); - assert!(!debug.contains("credential-user")); - assert!(!debug.contains("credential-password")); - assert!(debug.contains("redacted") || debug.contains("REDACTED")); -} - -#[rstest] -fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() { - let config = Config::from_yaml( - "litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid", - ) - .unwrap(); - - assert!(matches!( - config.litellm_settings.cache_params, - Some(litellm_config::CacheSettings::Redis(_)) - )); -} - -#[rstest] -fn non_redis_cache_settings_remain_opaque() { - let config = Config::from_yaml( - "litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true", - ) - .unwrap(); - - assert!(matches!( - config.litellm_settings.cache_params, - Some(litellm_config::CacheSettings::Other(_)) - )); -} - -#[rstest] -fn legacy_unknown_settings_do_not_change_resolution() { - let resolved = resolve( - "general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary", - &[], - ) - .unwrap(); - - assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_))); -} - -struct WarningCapture(Arc>>); - -impl Layer for WarningCapture { - fn on_event(&self, event: &Event<'_>, _context: Context<'_, S>) { - let mut fields = Fields::default(); - event.record(&mut fields); - self.0.lock().unwrap().push(fields.0); - } -} - -#[derive(Default)] -struct Fields(String); - -impl Visit for Fields { - fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) { - write!(self.0, "{}={value:?} ", field.name()).unwrap(); - } - - fn record_str(&mut self, field: &tracing::field::Field, value: &str) { - write!(self.0, "{}={value} ", field.name()).unwrap(); - } -} - -#[rstest] -fn ignored_setting_warning_names_the_setting_without_its_value() { - let events = Arc::new(Mutex::new(Vec::new())); - let subscriber = tracing_subscriber::registry().with(WarningCapture(events.clone())); - tracing::subscriber::with_default(subscriber, || { - resolve( - "general_settings:\n coordination_redis:\n host: redis\n max_connections: 47", - &[], - ) - .unwrap(); - }); - let warnings = events.lock().unwrap().join(" "); - - assert!(warnings.contains("max_connections")); - assert!(!warnings.contains("47")); -} diff --git a/litellm-rust/crates/gateway/AGENTS.md b/litellm-rust/crates/gateway/AGENTS.md index a33bd19cc80..16714cd9a1f 100644 --- a/litellm-rust/crates/gateway/AGENTS.md +++ b/litellm-rust/crates/gateway/AGENTS.md @@ -5,8 +5,11 @@ - Test mounting and server lifecycle through public HTTP behavior. Test endpoint semantics in the owning crate ## Redis -- `build_redis` builds the Redis clients once at boot (`src/redis.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored +- `build_redis` builds the Redis clients once at boot (`src/redis/mod.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored - Which Redis backs which role mirrors the proxy: coordination is `general_settings.coordination_redis`, else the `cache_params` Redis when `cache: true` and `type: redis`, else `REDIS_*` env. The transaction buffer reuses coordination, else env when `use_redis_transaction_buffer` is on. The response cache gets its own client so large values never queue in front of coordination commands -- `litellm-config` follows `litellm/_redis.py`'s merge of config blocks and `REDIS_*` env. A redis-py setting we can't honor (sentinel, IAM, custom TLS files, ...) fails boot with its name, never silently dropped +- `src/redis/layers.rs` is the Rust copy of `_merge_redis_kwargs` in `litellm/_redis.py`, which only this host needs because the Python host merges in Python. It puts the configured kwargs over their `REDIS_*` values and fills the `get_secret` fallbacks from the `SETTINGS` rows in `redis-resolve`, under the same conditions Python reads them, then hands the result to `litellm_redis_resolve::interpret`. Values come from a `SecretSnapshot` filled once at boot, read like `get_secret` (secret manager, then environment, with `true` and `false` read as bools), and the only straight environment read is Azure's `REDIS_USERNAME` +- A surface is a `ConfiguredRedis` constructor: `coordination` for `general_settings.coordination_redis`, which rejects a malformed `port` like the proxy's `CoordinationRedisParams` pydantic model, `cache_params` for `litellm_settings.cache_params`, which drops the cache-only keys and reads the `redis_startup_nodes` alias, and `redis_cache_kwargs` for a `RedisCache.redis_kwargs` dict in the contract tests. A new surface is a new constructor, never a flag +- Nulls depend on the surface. `coordination_redis` drops every null like `model_dump(exclude_none=True)`, `cache_params` keeps a null as an explicit `None` that hides the `REDIS_*` value unless `Cache` or `RedisCache` names that parameter, and `RedisCache.redis_kwargs` keeps every null. A missing `os.environ/` reference is a null, as `load_config` resolves it +- A setting the Rust client cannot honor warns and is dropped, and only a security check it cannot enforce fails boot, per `redis-resolve/AGENTS.md` - Boot never waits for a configured Redis (clients connect lazily). Only the env-only coordination fallback is pinged, for 2s, and treated as absent when unreachable, as in Python - `/health/readiness` never touches Redis, matching the proxy diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index d0cedbb3b36..a14df9b9c53 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -20,9 +20,14 @@ futures-util.workspace = true litellm-gateway-auth.workspace = true litellm-gateway-ui.workspace = true litellm-config.workspace = true +litellm-core-utils.workspace = true litellm-error.workspace = true litellm-circuit-breaker.workspace = true litellm-redis-client.workspace = true +litellm-redis-connect.workspace = true +litellm-redis-resolve.workspace = true +litellm-redis-types.workspace = true +rustls.workspace = true litellm-http.workspace = true litellm-llms.workspace = true litellm-secrets.workspace = true @@ -32,10 +37,12 @@ serde.workspace = true tracing.workspace = true tokio = { workspace = true, features = ["signal"] } uuid.workspace = true +veil.workspace = true tower-sessions-moka-store = "0.15.0" [dev-dependencies] litellm-redis-testing.workspace = true +litellm-http = { workspace = true, features = ["test-support"] } base64.workspace = true rstest.workspace = true tempfile.workspace = true diff --git a/litellm-rust/crates/gateway/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index e2f646c5448..c267e79353d 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -12,13 +12,12 @@ pub enum Error { Mcp(#[from] litellm_gateway_mcp::ConnectError), #[error(transparent)] Config(#[from] litellm_config::Error), - #[error("invalid Redis setting {setting}: {reason}")] - RedisInvalid { - setting: &'static str, - reason: &'static str, - }, #[error(transparent)] - RedisConnect(#[from] litellm_redis_client::ConnectError), + Redis(#[from] litellm_redis_resolve::Error), + #[error("Redis setting {setting} is not plain data")] + RedisNotPlainData { setting: String }, + #[error(transparent)] + RedisConnect(#[from] litellm_redis_connect::Error), #[error("MCP requires working gateway authentication")] Auth(#[from] litellm_gateway_auth::Error), } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 12a9bf9110a..0a0191b0499 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -68,15 +68,29 @@ pub fn build_inference(config: &Config) -> Result, Error> { Ok(Arc::new(Gateway::new(resources, http, secrets, router)?)) } -pub async fn build_redis(config: &Config) -> Result { - let environment = secrets::environment_values(&config.environment_variables)?; - let lookup = |name: &str| { +pub fn install_crypto_provider() { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } +} + +pub async fn build_redis( + config: &Config, + pool: &HttpClientPool, + http: &litellm_http::HttpClientConfig, +) -> Result { + let environment = Arc::new(secrets::environment_values(&config.environment_variables)?); + let lookup: litellm_redis_connect::SharedEnvironment = Arc::new(move |name: &str| { environment .get(name) .map(|value| value.expose().to_owned()) .or_else(|| std::env::var(name).ok()) + }); + let resources = redis::Resources { + http: pool.client(http, ClientVariant::Provider)?, + lookup: lookup.clone(), }; - redis::connect(config, &lookup).await + redis::connect(config, &*lookup, &resources).await } pub async fn build_mcp( diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index ff528723095..dc77cfe28eb 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -41,6 +41,7 @@ impl Sink for StderrSink { #[tokio::main] async fn main() -> Result<(), Box> { + litellm_gateway::install_crypto_provider(); if let Some(forward) = litellm_gateway_forward::Settings::from_env()? { let log: LogSettings = settings::from_iter(std::env::vars_os())?; Logger::new(StderrSink { @@ -55,7 +56,8 @@ async fn main() -> Result<(), Box> { Logger::new(StderrSink { level }).install_global()?; let config = Config::load(settings.litellm_config)?; let inference = litellm_gateway::build_inference(&config)?; - let _redis = litellm_gateway::build_redis(&config).await?; + let _redis = + litellm_gateway::build_redis(&config, &inference.resources.pool, &inference.http).await?; let ui = match std::env::var_os("LITELLM_UI_PATH") { Some(directory) => { let ui_settings: UiSettings = settings::from_iter(std::env::vars_os())?; diff --git a/litellm-rust/crates/gateway/src/redis/configured.rs b/litellm-rust/crates/gateway/src/redis/configured.rs new file mode 100644 index 00000000000..93297d0e57f --- /dev/null +++ b/litellm-rust/crates/gateway/src/redis/configured.rs @@ -0,0 +1,144 @@ +use std::collections::BTreeMap; + +use litellm_config::{RedisSettings, Value}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_resolve::{Expected, Kwargs}; +use litellm_redis_types::{ProviderHandle, Resolution}; +use veil::Redact; + +use super::{layers, secrets::SecretSnapshot}; +use crate::Error; + +const CACHE_ONLY: &[&str] = &[ + "type", + "mode", + "ttl", + "default_in_memory_ttl", + "default_in_redis_ttl", + "supported_call_types", + "similarity_threshold", + "max_messages", + "_backend", +]; + +/// The redis-py kwargs one surface hands `_get_redis_client_logic`, before `REDIS_*` and secrets. +#[derive(Clone, Redact)] +pub(crate) struct ConfiguredRedis { + #[redact(with = "[REDACTED]")] + pub(super) fields: BTreeMap, + pub(super) nulls: Nulls, + pub(super) credential_provider: Option, +} + +/// Whether a null setting reaches redis-py as an explicit `None`, which outranks its `REDIS_*` value. +#[derive(Clone, Copy, Debug)] +pub(super) enum Nulls { + /// `coordination_redis` reaches redis-py through `model_dump(exclude_none=True)`. + Dropped, + /// `Cache` drops a null only for the parameters it names and passes any other key on. + KeptUnlessNamedByCache, + /// `RedisCache.redis_kwargs` already holds only the nulls redis-py receives. + #[cfg(test)] + Kept, +} + +impl ConfiguredRedis { + /// `general_settings.coordination_redis`, which the proxy validates as a pydantic model first. + /// An `os.environ/` port is only known after merging, where it is read like any other port. + pub(crate) fn coordination( + settings: &RedisSettings, + ) -> Result { + if let Some(port) = settings.fields().get("port").and_then(text) + && !port.starts_with("os.environ/") + && port.parse::().is_err() + { + return Err(litellm_redis_resolve::Error::Invalid { + setting: "port".to_owned(), + expected: Expected::Integer, + }); + } + Ok(Self { + fields: settings.fields().clone(), + nulls: Nulls::Dropped, + credential_provider: None, + }) + } + + /// `litellm_settings.cache_params` with `type: redis`, as `litellm.Cache` hands it to `RedisCache`. + pub(crate) fn cache_params(settings: &RedisSettings) -> Self { + let fields = settings.fields(); + let alias = fields + .get("redis_startup_nodes") + .filter(|_| !fields.contains_key("startup_nodes")) + .map(|nodes| ("startup_nodes".to_owned(), nodes.clone())); + let fields = fields + .iter() + .filter(|(key, _)| { + key.as_str() != "redis_startup_nodes" && !CACHE_ONLY.contains(&key.as_str()) + }) + .map(|(key, value)| (key.clone(), value.clone())) + .chain(alias) + .collect(); + Self { + fields, + nulls: Nulls::KeptUnlessNamedByCache, + credential_provider: None, + } + } + + /// `RedisCache.redis_kwargs`, the dict every Python surface ends in. + #[cfg(test)] + pub(crate) fn redis_cache_kwargs(fields: BTreeMap) -> Self { + Self { + fields, + nulls: Nulls::Kept, + credential_provider: None, + } + } + + /// The handle the host holds for the `credential_provider` object in these settings. + #[cfg(test)] + pub(crate) fn with_credential_provider(self, handle: ProviderHandle) -> Self { + Self { + credential_provider: Some(handle), + ..self + } + } + + /// No configured settings, so only `REDIS_*` and the other environment fallbacks apply. + pub(crate) fn environment_only() -> Self { + Self { + fields: BTreeMap::new(), + nulls: Nulls::Dropped, + credential_provider: None, + } + } + + /// `_merge_redis_kwargs` in `litellm/_redis.py`: the kwargs over their `REDIS_*` values with every + /// `get_secret` fallback filled. + pub(crate) fn merged(&self, secrets: &SecretSnapshot) -> Result { + layers::merged(self, secrets) + } + + /// The environment is only for what Python reads straight from `os.environ`, `REDIS_USERNAME` for Azure. + pub(crate) fn resolve( + &self, + secrets: &SecretSnapshot, + environment: &dyn Lookup, + ) -> Result { + Ok(litellm_redis_resolve::interpret( + &self.merged(secrets)?, + self.credential_provider, + environment.get("REDIS_USERNAME").as_deref(), + )?) + } +} + +pub(super) fn text(value: &Value) -> Option { + match value { + Value::String(text) => Some(text.clone()), + Value::Number(number) => Some(number.to_string()), + Value::Bool(flag) => Some(flag.to_string()), + _ => None, + } +} diff --git a/litellm-rust/crates/gateway/src/redis/layers.rs b/litellm-rust/crates/gateway/src/redis/layers.rs new file mode 100644 index 00000000000..1cf68b2cd4f --- /dev/null +++ b/litellm-rust/crates/gateway/src/redis/layers.rs @@ -0,0 +1,226 @@ +use litellm_config::Value; +use litellm_core_utils::settings::Lookup; +use litellm_redis_resolve::{Entry, Kwargs, SETTINGS, is_present, setting}; + +use super::configured::{ConfiguredRedis, Nulls}; +use crate::Error; + +const NAMED_BY_CACHE: &[&str] = &[ + "host", + "port", + "password", + "socket_timeout", + "startup_nodes", + "namespace", + "gcp_service_account", + "gcp_ssl_ca_certs", +]; + +/// How `_merge_redis_kwargs` reads a fallback: `get_secret` turns `true` and `false` into a bool, and +/// `get_secret_str` treats that bool as unset. +#[derive(Clone, Copy)] +enum Read { + Secret, + SecretStr, +} + +const TARGETS: &[(&str, Read)] = &[ + ("startup_nodes", Read::Secret), + ("sentinel_nodes", Read::Secret), + ("sentinel_password", Read::SecretStr), + ("service_name", Read::Secret), +]; + +const IDENTITIES: &[(&str, Read)] = &[ + ("gcp_service_account", Read::SecretStr), + ("gcp_ssl_ca_certs", Read::SecretStr), + ("azure_redis_ad_token", Read::Secret), +]; + +const AZURE: &[(&str, Read)] = &[ + ("azure_client_id", Read::SecretStr), + ("azure_tenant_id", Read::SecretStr), + ("azure_client_secret", Read::SecretStr), +]; + +const AWS: &[(&str, Read)] = &[("aws_iam_region", Read::SecretStr)]; + +/// The configured settings over their `REDIS_*` values, `{**environment_kwargs, **env_overrides}` +/// in `_merge_redis_kwargs`, with the `get_secret` fallbacks it reads filled in. A null left in the result is +/// an explicit `None` that hid its environment value. +pub(super) fn merged(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result { + let block: Kwargs = configured + .fields + .iter() + .map(|(key, value)| (key, environment_reference(value, lookup))) + .filter(|(key, value)| !value.is_null() || keeps_null(configured.nulls, key)) + .map(|(key, value)| { + let value = serde_json::to_value(&value).map_err(|_| Error::RedisNotPlainData { + setting: key.clone(), + })?; + let entry = Entry { + value, + source: key.clone(), + }; + Ok((key.clone(), entry)) + }) + .collect::>()?; + let names_a_target = ["host", "startup_nodes", "sentinel_nodes"] + .iter() + .any(|key| is_set(&block, key)); + let drop_environment_url = names_a_target && !is_set(&block, "url"); + let layered: Kwargs = environment(lookup) + .filter(|(key, _)| !(key == "url" && drop_environment_url)) + .chain(block) + .collect(); + let fallbacks = fallbacks(&layered, configured.credential_provider.is_some(), lookup); + Ok(layered.into_iter().chain(fallbacks).collect()) +} + +// The fallbacks are read under the same conditions as _merge_redis_kwargs, so a secret Python never asks +// for is never filled in +fn fallbacks(layered: &Kwargs, has_provider: bool, lookup: &dyn Lookup) -> Kwargs { + let fill = |kwargs: &[(&str, Read)]| -> Kwargs { + kwargs + .iter() + .filter_map(|(kwarg, read)| fallback(layered, kwarg, *read, lookup)) + .collect() + }; + let targets = fill(TARGETS); + if has_provider || is_set(layered, "credential_provider") { + return targets; + } + let identities = fill(IDENTITIES); + let value = |kwarg: &str| { + layered + .get(kwarg) + .filter(|entry| is_present(&entry.value)) + .or_else(|| identities.get(kwarg)) + .map(|entry| &entry.value) + }; + let gcp = value("gcp_service_account").is_some(); + let azure = value("azure_redis_ad_token") + .is_some_and(|token| text(token).is_some_and(|token| token.eq_ignore_ascii_case("true"))); + let aws = layered + .get("aws_iam_auth") + .is_some_and(|entry| coerces_to_true(&entry.value)); + let azure_settings = match azure && !gcp { + true => fill(AZURE), + false => Kwargs::new(), + }; + let aws_settings = match aws && !gcp && !azure { + true => fill(AWS), + false => Kwargs::new(), + }; + targets + .into_iter() + .chain(identities) + .chain(azure_settings) + .chain(aws_settings) + .collect() +} + +fn fallback( + layered: &Kwargs, + kwarg: &str, + read: Read, + lookup: &dyn Lookup, +) -> Option<(String, Entry)> { + if layered + .get(kwarg) + .is_some_and(|entry| is_present(&entry.value)) + { + return None; + } + let names = setting(kwarg).map_or(&[][..], |setting| setting.fallback); + names.iter().find_map(|name| { + let value = lookup + .get(name) + .filter(|value| !value.is_empty()) + .map(secret)?; + match (read, &value) { + (Read::SecretStr, serde_json::Value::Bool(_)) => None, + _ => { + let entry = Entry { + value, + source: (*name).to_owned(), + }; + Some((kwarg.to_owned(), entry)) + } + } + }) +} + +// get_secret reads an environment value of true or false, in any case and padding, as a bool +fn secret_flag(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "true" => Some(true), + "false" => Some(false), + _ => None, + } +} + +fn secret(value: String) -> serde_json::Value { + secret_flag(&value).map_or(serde_json::Value::String(value), serde_json::Value::Bool) +} + +// _coerces_to_true: text is read like _str_to_bool, anything else by its truthiness +fn coerces_to_true(value: &serde_json::Value) -> bool { + match value { + serde_json::Value::String(text) => { + matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes") + } + serde_json::Value::Null => false, + serde_json::Value::Bool(flag) => *flag, + serde_json::Value::Number(number) => number.as_f64().is_some_and(|number| number != 0.0), + serde_json::Value::Array(items) => !items.is_empty(), + serde_json::Value::Object(entries) => !entries.is_empty(), + } +} + +fn text(value: &serde_json::Value) -> Option { + match value { + serde_json::Value::String(text) => Some(text.clone()), + serde_json::Value::Number(number) => Some(number.to_string()), + serde_json::Value::Bool(flag) => Some(flag.to_string()), + _ => None, + } +} + +fn keeps_null(nulls: Nulls, key: &str) -> bool { + match nulls { + Nulls::Dropped => false, + Nulls::KeptUnlessNamedByCache => !NAMED_BY_CACHE.contains(&key), + #[cfg(test)] + Nulls::Kept => true, + } +} + +fn environment(lookup: &dyn Lookup) -> impl Iterator + '_ { + SETTINGS.iter().filter_map(|setting| { + let source = setting.environment_name()?; + let value = lookup.get(&source)?; + let entry = Entry { + value: secret(value), + source, + }; + Some((setting.kwarg.to_owned(), entry)) + }) +} + +// load_config resolves os.environ/ to None when unset, so a missing reference is a null like any other +fn environment_reference(value: &Value, lookup: &dyn Lookup) -> Value { + match value { + Value::String(text) => match text.strip_prefix("os.environ/") { + Some(name) => lookup.get(name).map_or(Value::Null, |value| { + secret_flag(&value).map_or(Value::String(value), Value::Bool) + }), + None => value.clone(), + }, + _ => value.clone(), + } +} + +fn is_set(kwargs: &Kwargs, key: &str) -> bool { + kwargs.get(key).is_some_and(|entry| !entry.value.is_null()) +} diff --git a/litellm-rust/crates/gateway/src/redis.rs b/litellm-rust/crates/gateway/src/redis/mod.rs similarity index 60% rename from litellm-rust/crates/gateway/src/redis.rs rename to litellm-rust/crates/gateway/src/redis/mod.rs index c40b289029b..17b111231e7 100644 --- a/litellm-rust/crates/gateway/src/redis.rs +++ b/litellm-rust/crates/gateway/src/redis/mod.rs @@ -1,19 +1,31 @@ -use std::time::Duration; +use std::{sync::Arc, time::Duration}; -use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, State}; -use litellm_config::{ - CacheSettings, Config, Lookup, RedisEndpoint, RedisSettings, ResolvedRedis, Value, - environment_has_target, -}; -use litellm_redis_client::{ - ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, Namespace, RedisClient, - RedisConfig, -}; +use litellm_circuit_breaker::Config as BreakerConfig; +use litellm_config::{CacheSettings, Config, RedisSettings, Value}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_client::{RedisClient, RedisEvents}; +use litellm_redis_connect::{HostServices, SharedEnvironment, TracingEvents}; +use litellm_redis_resolve::Expected; +use litellm_redis_types::Resolution; use crate::Error; +mod configured; +mod layers; +mod secrets; +#[cfg(test)] +mod surface_tests; + +use configured::ConfiguredRedis; +use secrets::SecretSnapshot; + +pub(crate) struct Resources { + pub(crate) http: litellm_http::Client, + pub(crate) lookup: SharedEnvironment, +} + const ENVIRONMENT_PING_TIMEOUT: Duration = Duration::from_secs(2); -type Params = ResolvedRedis; +type Params = Resolution; pub struct RedisClients { pub coordination: Option, @@ -34,12 +46,15 @@ enum Target { Environment(Params), } -pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result { +pub(crate) fn plan(config: &Config, lookup: &dyn Lookup) -> Result { let response_cache = redis_cache_params(config) - .map(|settings| settings.resolve_cache(lookup)) + .map(|settings| resolve(&ConfiguredRedis::cache_params(settings), lookup)) .transpose()?; let coordination = match &config.general_settings.coordination_redis { - Some(settings) => Some(Target::Configured(settings.resolve(lookup)?)), + Some(settings) => Some(Target::Configured(resolve( + &ConfiguredRedis::coordination(settings)?, + lookup, + )?)), None => match (&response_cache, config.litellm_settings.cache == Some(true)) { (Some(cache), _) => Some(Target::Configured(cache.clone())), (None, true) => environment(lookup)?.map(Target::Configured), @@ -57,12 +72,29 @@ pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result { }) } -pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result { +pub(crate) async fn connect( + config: &Config, + lookup: &dyn Lookup, + resources: &Resources, +) -> Result { let plan = plan(config, lookup)?; let breaker = breaker_config(lookup)?; + let events: Arc = Arc::new(TracingEvents::new(lookup)); + let client = |params: Params| -> Result { + Ok(litellm_redis_connect::connect( + params, + HostServices { + events: Some(events.clone()), + http: resources.http.clone(), + environment: resources.lookup.clone(), + credential_provider: None, + }, + breaker, + )?) + }; let coordination = match plan.coordination { - Some(Target::Configured(params)) => Some(client(params, breaker)?), - Some(Target::Environment(params)) => match client(params, breaker) { + Some(Target::Configured(params)) => Some(client(params)?), + Some(Target::Environment(params)) => match client(params) { Ok(client) => reachable(client).await, Err(_) => { tracing::warn!( @@ -75,13 +107,10 @@ pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result Some(coordination.clone()), - (None, Some(params)) => Some(client(params, breaker)?), + (None, Some(params)) => Some(client(params)?), (None, None) => None, }; - let response_cache = plan - .response_cache - .map(|params| client(params, breaker)) - .transpose()?; + let response_cache = plan.response_cache.map(client).transpose()?; Ok(RedisClients { coordination, transaction_buffer, @@ -98,17 +127,38 @@ fn redis_cache_params(config: &Config) -> Option<&RedisSettings> { } } -fn environment(lookup: Lookup<'_>) -> Result, Error> { +fn resolve(configured: &ConfiguredRedis, lookup: &dyn Lookup) -> Result { + let secrets = SecretSnapshot::read(configured.secret_names(), lookup); + let resolution = configured.resolve(&secrets, lookup)?; + for ignored in &resolution.ignored { + tracing::warn!( + "Redis setting {} is not supported by the Rust client and is ignored: {}", + ignored.setting, + ignored.consequence + ); + } + Ok(resolution) +} + +fn environment(lookup: &dyn Lookup) -> Result, Error> { match environment_has_target(lookup) { - true => RedisSettings::default() - .resolve(lookup) - .map(Some) - .map_err(Into::into), + true => Ok(Some(resolve(&ConfiguredRedis::environment_only(), lookup)?)), false => Ok(None), } } -fn optional_environment(lookup: Lookup<'_>) -> Option { +fn environment_has_target(lookup: &dyn Lookup) -> bool { + [ + "REDIS_HOST", + "REDIS_URL", + "REDIS_CLUSTER_NODES", + "REDIS_SENTINEL_NODES", + ] + .iter() + .any(|name| lookup.get(name).is_some()) +} + +fn optional_environment(lookup: &dyn Lookup) -> Option { match environment(lookup) { Ok(params) => params, Err(error) => { @@ -149,122 +199,101 @@ async fn reachable(client: RedisClient) -> Option { } } -fn client(params: Params, breaker: Option) -> Result { - tracing::info!(endpoint = ?params.endpoint, namespace = ?params.namespace, "redis client configured"); - let endpoint = match params.endpoint { - RedisEndpoint::Url(url) => Endpoint::Url(url), - RedisEndpoint::Host(host) => Endpoint::Host(HostEndpoint { - host: host.host, - port: host.port, - username: host.username, - password: host.password, - tls: host.tls, - db: host.db, - }), - RedisEndpoint::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { - nodes: cluster - .nodes - .into_iter() - .map(|node| ClusterNode { - host: node.host, - port: node.port, - }) - .collect(), - username: cluster.username, - password: cluster.password, - tls: cluster.tls, - }), - }; - Ok(RedisClient::new(RedisConfig { - endpoint, - namespace: params.namespace.map(Namespace::new).unwrap_or_default(), - timeout: params.timeout.unwrap_or(DEFAULT_TIMEOUT), - breaker: breaker.map(|config| Breaker::new(config, Box::new(log_transition))), - })?) +fn parse_seconds(text: &str) -> Option { + text.parse::() + .ok() + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) } -fn log_transition(transition: &litellm_circuit_breaker::Transition) { - match transition.to { - State::Open { .. } => tracing::warn!( - generation = transition.generation, - "redis circuit breaker opened, fast-failing redis calls" - ), - State::HalfOpen => tracing::info!( - generation = transition.generation, - "redis circuit breaker half open, probing redis" - ), - State::Closed { .. } => tracing::info!( - generation = transition.generation, - "redis circuit breaker closed, redis recovered" - ), - } -} - -fn breaker_config(lookup: Lookup<'_>) -> Result, Error> { - let enabled = lookup("REDIS_CIRCUIT_BREAKER_ENABLED") +fn breaker_config(lookup: &dyn Lookup) -> Result, Error> { + let enabled = lookup + .get("REDIS_CIRCUIT_BREAKER_ENABLED") .is_none_or(|value| value.eq_ignore_ascii_case("true")); if !enabled { return Ok(None); } let defaults = BreakerConfig::default(); - let seconds = |name: &'static str, default: Duration| -> Result { - lookup(name) - .map(|value| { - value - .parse::() - .ok() - .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) - .ok_or(Error::RedisInvalid { - setting: name, - reason: "must be a non-negative number of seconds", - }) - }) - .unwrap_or(Ok(default)) - }; - let failure_threshold = lookup("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD") + Ok(Some(BreakerConfig { + failure_threshold: environment_value( + lookup, + "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", + Expected::Integer, + |text| text.parse().ok(), + )? + .unwrap_or(defaults.failure_threshold), + recovery_timeout: environment_value( + lookup, + "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", + Expected::Seconds, + parse_seconds, + )? + .unwrap_or(defaults.recovery_timeout), + timeout_min_duration: environment_value( + lookup, + "REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", + Expected::Seconds, + parse_seconds, + )? + .unwrap_or(defaults.timeout_min_duration), + })) +} + +fn environment_value( + lookup: &dyn Lookup, + name: &str, + expected: Expected, + parse: impl Fn(&str) -> Option, +) -> Result, litellm_redis_resolve::Error> { + lookup + .get(name) .map(|value| { - value.parse().map_err(|_| Error::RedisInvalid { - setting: "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", - reason: "must be an integer", + parse(&value).ok_or_else(|| litellm_redis_resolve::Error::Invalid { + setting: name.to_owned(), + expected, }) }) - .unwrap_or(Ok(defaults.failure_threshold))?; - Ok(Some(BreakerConfig { - failure_threshold, - recovery_timeout: seconds( - "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", - defaults.recovery_timeout, - )?, - timeout_min_duration: seconds( - "REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", - defaults.timeout_min_duration, - )?, - })) + .transpose() } #[cfg(test)] mod tests { + use litellm_core_utils::settings::Lookup; use std::{collections::HashMap, time::Duration}; - use litellm_config::{Config, RedisEndpoint, RedisHost}; + use litellm_config::Config; + use litellm_redis_client::{Standalone, StaticAuth, commands}; use litellm_redis_testing::RedisServer; + use litellm_redis_types::{Auth, HealthCheck}; use rstest::rstest; - use super::{Params, Plan, Target, connect, plan}; + use std::sync::Arc; + + use super::{Params, Plan, Resources, Target, plan}; use crate::Error; + async fn connect(config: &Config, lookup: &dyn Lookup) -> Result { + let resources = Resources { + http: litellm_http::Client::for_test(), + lookup: Arc::new(|_: &str| None), + }; + super::connect(config, lookup, &resources).await + } + fn host(name: &str, port: u16) -> Params { Params { - endpoint: RedisEndpoint::Host(RedisHost { + target: litellm_redis_client::Target::Standalone(Standalone { host: name.to_owned(), port, - username: None, - password: None, - tls: false, + tls: None, db: 0, }), + auth: Auth::Static(StaticAuth::default()), namespace: None, timeout: None, + connect_timeout: None, + health_check: HealthCheck::Unset, + keepalive: true, + ignored: Vec::new(), } } @@ -273,7 +302,7 @@ mod tests { .iter() .map(|(key, value)| ((*key).to_owned(), (*value).to_owned())) .collect(); - plan(&Config::from_yaml(yaml).unwrap(), &|name| { + plan(&Config::from_yaml(yaml).unwrap(), &|name: &str| { environment.get(name).cloned() }) } @@ -350,36 +379,44 @@ mod tests { port_of(&server) )) .unwrap(); - let clients = connect(&config, &|_| None).await.unwrap(); + let clients = connect(&config, &|_: &str| None).await.unwrap(); let coordination = clients.coordination.unwrap(); let buffer = clients.transaction_buffer.unwrap(); let cache = clients.response_cache.unwrap(); - let unnamespaced = connect(&coordination_config(&server), &|_| None) + let unnamespaced = connect(&coordination_config(&server), &|_: &str| None) .await .unwrap() .coordination .unwrap(); coordination - .set("rate-limit", "1", Duration::from_secs(60)) + .send(commands::set("rate-limit", "1", Duration::from_secs(60))) .await .unwrap(); buffer - .set("spend", "2", Duration::from_secs(60)) + .send(commands::set("spend", "2", Duration::from_secs(60))) .await .unwrap(); assert_eq!( - cache.get::("rate-limit").await.unwrap().as_deref(), + cache + .send(commands::get::("rate-limit")) + .await + .unwrap() + .as_deref(), Some("1") ); assert_eq!( - cache.get::("spend").await.unwrap().as_deref(), + cache + .send(commands::get::("spend")) + .await + .unwrap() + .as_deref(), Some("2") ); assert_eq!( unnamespaced - .get::("responses:rate-limit") + .send(commands::get::("responses:rate-limit")) .await .unwrap() .as_deref(), @@ -387,7 +424,7 @@ mod tests { ); assert_eq!( unnamespaced - .get::("responses:spend") + .send(commands::get::("responses:spend")) .await .unwrap() .as_deref(), @@ -395,16 +432,21 @@ mod tests { ); assert!( unnamespaced - .get::("rate-limit") + .send(commands::get::("rate-limit")) + .await + .unwrap() + .is_none() + ); + assert!( + unnamespaced + .send(commands::get::("spend")) .await .unwrap() .is_none() ); - assert!(unnamespaced.get::("spend").await.unwrap().is_none()); } #[rstest] - #[case::malformed_port(&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])] #[case::malformed_cluster_nodes(&[("REDIS_CLUSTER_NODES", "[{\"host\": ")])] fn redis_migration_regression_malformed_optional_environment_is_absent( #[case] environment: &[(&str, &str)], @@ -413,54 +455,45 @@ mod tests { } #[rstest] - fn redis_migration_regression_explicit_malformed_setting_still_fails() { + #[case::optional_environment("model_list: []", &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])] + #[case::transaction_buffer_environment( + "general_settings:\n use_redis_transaction_buffer: true", + &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")] + )] + #[case::response_cache( + "litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: env\n port: not-a-port", + &[] + )] + fn a_malformed_port_outside_the_coordination_block_falls_back_to_6379_like_redis_py( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + ) { + let plan = resolved(yaml, environment).unwrap(); + + let target = match (plan.coordination, plan.transaction_buffer_fallback) { + (Some(Target::Configured(params) | Target::Environment(params)), _) + | (None, Some(params)) => params, + (None, None) => panic!("expected a Redis target"), + }; + assert_eq!(target.target, host("env", 6379).target); + } + + #[rstest] + #[case::port("port: not-a-port", "port")] + #[case::ssl("ssl: maybe", "ssl")] + fn a_malformed_coordination_block_fails_like_the_pydantic_model( + #[case] setting: &str, + #[case] expected: &str, + ) { let error = resolved( - "general_settings:\n coordination_redis:\n host: coord\n port: not-a-port", + &format!("general_settings:\n coordination_redis:\n host: coord\n {setting}"), &[], ) .unwrap_err(); - assert!(matches!( - error, - Error::Config(litellm_config::Error::RedisInvalid { - setting: "port", - .. - }) - )); - } - - #[rstest] - fn redis_migration_regression_enabled_transaction_buffer_rejects_invalid_environment() { - let error = resolved( - "general_settings:\n use_redis_transaction_buffer: true", - &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")], - ) - .unwrap_err(); - - assert!(matches!( - error, - Error::Config(litellm_config::Error::RedisInvalid { - setting: "port", - .. - }) - )); - } - - #[rstest] - fn redis_migration_regression_invalid_response_cache_still_fails() { - let error = resolved( - "litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: cache\n port: not-a-port", - &[], - ) - .unwrap_err(); - - assert!(matches!( - error, - Error::Config(litellm_config::Error::RedisInvalid { - setting: "port", - .. - }) - )); + assert!( + matches!(error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == expected) + ); } #[rstest] @@ -468,7 +501,7 @@ mod tests { async fn redis_migration_regression_invalid_optional_environment_url_is_absent() { let config = Config::from_yaml("model_list: []").unwrap(); - let clients = connect(&config, &|name| { + let clients = connect(&config, &|name: &str| { (name == "REDIS_URL").then(|| "not a redis URL".to_owned()) }) .await @@ -506,7 +539,7 @@ mod tests { let server = RedisServer::start().await.unwrap(); server.stop().await.unwrap(); - let clients = connect(&coordination_config(&server), &|_| None) + let clients = connect(&coordination_config(&server), &|_: &str| None) .await .unwrap(); let coordination = clients.coordination.unwrap(); @@ -521,12 +554,12 @@ mod tests { #[tokio::test] async fn the_coordination_client_survives_a_redis_restart_and_backs_the_transaction_buffer() { let server = RedisServer::start().await.unwrap(); - let clients = connect(&coordination_config(&server), &|_| None) + let clients = connect(&coordination_config(&server), &|_: &str| None) .await .unwrap(); let coordination = clients.coordination.unwrap(); coordination - .set("written", "1", Duration::from_secs(60)) + .send(commands::set("written", "1", Duration::from_secs(60))) .await .unwrap(); @@ -537,6 +570,30 @@ mod tests { assert!(buffer.ping().await.is_ok()); } + #[rstest] + #[tokio::test] + async fn an_environment_only_tls_redis_is_reached_with_its_ca() { + crate::install_crypto_provider(); + let server = litellm_redis_testing::RedisTlsServer::start(false) + .await + .unwrap(); + let ca = server.files().ca.display().to_string(); + let port = server.port().to_string(); + let config = Config::from_yaml("model_list: []").unwrap(); + + let clients = connect(&config, &|name: &str| match name { + "REDIS_HOST" => Some(server.host().to_owned()), + "REDIS_PORT" => Some(port.clone()), + "REDIS_SSL" => Some("true".to_owned()), + "REDIS_SSL_CA_CERTS" => Some(ca.clone()), + _ => None, + }) + .await + .unwrap(); + + assert!(clients.coordination.is_some()); + } + #[rstest] #[tokio::test] async fn an_unreachable_environment_only_redis_leaves_coordination_in_memory() { @@ -545,7 +602,7 @@ mod tests { server.stop().await.unwrap(); let config = Config::from_yaml("model_list: []").unwrap(); - let clients = connect(&config, &|name| match name { + let clients = connect(&config, &|name: &str| match name { "REDIS_HOST" => Some("127.0.0.1".to_owned()), "REDIS_PORT" => Some(port.clone()), _ => None, @@ -564,7 +621,7 @@ mod tests { let port = port_of(&server); let config = Config::from_yaml("model_list: []").unwrap(); - let clients = connect(&config, &|name| match name { + let clients = connect(&config, &|name: &str| match name { "REDIS_HOST" => Some("127.0.0.1".to_owned()), "REDIS_PORT" => Some(port.clone()), _ => None, diff --git a/litellm-rust/crates/gateway/src/redis/secrets.rs b/litellm-rust/crates/gateway/src/redis/secrets.rs new file mode 100644 index 00000000000..a969497f4c3 --- /dev/null +++ b/litellm-rust/crates/gateway/src/redis/secrets.rs @@ -0,0 +1,61 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; + +use litellm_config::Value; +use litellm_redis_resolve::SETTINGS; + +use super::configured::ConfiguredRedis; + +#[derive(Clone, Debug, Default)] +pub(crate) struct SecretSnapshot(BTreeMap>); + +impl SecretSnapshot { + pub(crate) fn read(names: BTreeSet, lookup: &dyn Lookup) -> Self { + Self( + names + .into_iter() + .map(|name| { + let value = lookup.get(&name).map(SecretValue::new); + (name, value) + }) + .collect(), + ) + } +} + +impl Lookup for SecretSnapshot { + fn get(&self, name: &str) -> Option { + match self.0.get(name) { + Some(value) => value.as_ref().map(|value| value.expose().to_owned()), + None => { + debug_assert!( + false, + "{name} is missing from ConfiguredRedis::secret_names" + ); + tracing::error!( + name, + "a Redis secret was read without being in the snapshot" + ); + None + } + } + } +} + +impl ConfiguredRedis { + pub(crate) fn secret_names(&self) -> BTreeSet { + let references = self.fields.values().filter_map(|value| match value { + Value::String(text) => text.strip_prefix("os.environ/").map(str::to_owned), + _ => None, + }); + let declared = SETTINGS.iter().flat_map(|setting| { + setting + .environment_name() + .into_iter() + .chain(setting.fallback.iter().map(|name| (*name).to_owned())) + }); + declared.chain(references).collect() + } +} diff --git a/litellm-rust/crates/gateway/src/redis/surface_tests.rs b/litellm-rust/crates/gateway/src/redis/surface_tests.rs new file mode 100644 index 00000000000..9ee881e2e04 --- /dev/null +++ b/litellm-rust/crates/gateway/src/redis/surface_tests.rs @@ -0,0 +1,964 @@ +use std::collections::{BTreeMap, HashMap}; + +use litellm_config::{CacheSettings, Config, Value}; +use litellm_redis_resolve::Expected; +use litellm_redis_types::{ + Auth, AwsIamAuth, AzureAdAuth, Node, ProviderHandle, Resolution, StaticAuth, Target, TlsOptions, +}; +use litellm_secret_value::SecretValue; +use rstest::rstest; +use serde_json::json; + +use super::{configured::ConfiguredRedis, secrets::SecretSnapshot}; +use crate::Error; + +fn resolution( + configured: &ConfiguredRedis, + environment: &HashMap, +) -> Result { + let lookup = |name: &str| environment.get(name).cloned(); + configured.resolve( + &SecretSnapshot::read(configured.secret_names(), &lookup), + &lookup, + ) +} + +const PROVIDER: ProviderHandle = ProviderHandle::new(7); + +fn redis_cache_kwargs(kwargs: serde_json::Value) -> ConfiguredRedis { + let has_provider = kwargs + .get("credential_provider") + .is_some_and(|provider| !provider.is_null()); + let configured = ConfiguredRedis::redis_cache_kwargs( + serde_json::from_value::>(kwargs).unwrap(), + ); + match has_provider { + true => configured.with_credential_provider(PROVIDER), + false => configured, + } +} + +fn coordination_resolution(yaml: &str, environment: &[(&str, &str)]) -> Result { + let environment: HashMap = environment + .iter() + .map(|(key, value)| ((*key).to_owned(), (*value).to_owned())) + .collect(); + let config = Config::from_yaml(yaml).unwrap(); + let configured = config + .general_settings + .coordination_redis + .unwrap_or_default(); + let configured = ConfiguredRedis::coordination(&configured)?; + resolution(&configured, &environment) +} + +fn resolve(yaml: &str, environment: &[(&str, &str)]) -> Result { + coordination_resolution(yaml, environment) +} + +fn resolve_cache(yaml: &str) -> Result { + resolve_cache_in(yaml, &[]) +} + +fn resolve_cache_in(yaml: &str, environment: &[(&str, &str)]) -> Result { + let lookup = |name: &str| { + environment + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| (*value).to_owned()) + }; + let config = Config::from_yaml(yaml).unwrap(); + let configured = match config.litellm_settings.cache_params.unwrap() { + CacheSettings::Redis(settings) => ConfiguredRedis::cache_params(&settings), + CacheSettings::Other(_) => panic!("expected a Redis cache"), + }; + configured.resolve( + &SecretSnapshot::read(configured.secret_names(), &lookup), + &lookup, + ) +} + +fn cluster_ports(resolved: &Resolution) -> Vec { + match &resolved.target { + Target::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(), + target => panic!("expected a cluster, got {target:?}"), + } +} + +#[rstest] +#[case::host( + "general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5", + &[], + "host:redis:6380:2:true:app:1500" +)] +#[case::environment_references_resolve_the_host_and_port( + "general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT", + &[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")], + "host:referenced:6400:0:false::none" +)] +#[case::a_null_host_falls_back_to_the_environment( + "general_settings:\n coordination_redis:\n host: null", + &[("REDIS_HOST", "environment")], + "host:environment:6379:0:false::none" +)] +fn resolves_redis_through_the_parsed_config( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: &str, +) { + let resolved = resolve(yaml, environment).unwrap(); + let description = match resolved.target { + Target::Unix(socket) => format!("unix:{}", socket.path.display()), + Target::Url(url) => format!("url:{}", url.url.expose()), + Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name), + Target::Standalone(host) => format!( + "host:{}:{}:{}:{}:{}:{}", + host.host, + host.port, + host.db, + host.tls.is_some(), + resolved.namespace.unwrap_or_default(), + resolved + .timeout + .map(|timeout| timeout.as_millis().to_string()) + .unwrap_or_else(|| "none".to_owned()) + ), + Target::Cluster(cluster) => format!( + "cluster:{}", + cluster + .nodes + .iter() + .map(|node| format!("{}:{}", node.host, node.port)) + .collect::>() + .join(",") + ), + }; + + assert_eq!(description, expected); +} + +#[rstest] +#[case::tls_ciphers( + "general_settings:\n coordination_redis:\n host: redis\n ssl_ciphers: HIGH", + &[], + "ssl_ciphers" +)] +#[case::ssl_ciphers_from_environment( + "model_list: []", + &[("REDIS_HOST", "redis"), ("REDIS_SSL_CIPHERS", "HIGH")], + "REDIS_SSL_CIPHERS" +)] +#[case::connection_pool( + "general_settings:\n coordination_redis:\n host: redis\n connection_pool: shared", + &[], + "connection_pool" +)] +fn unsupported_settings_resolve_and_are_named_by_their_source( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: &str, +) { + let resolution = coordination_resolution(yaml, environment).unwrap(); + + assert_eq!( + resolution.target, + Target::Standalone(litellm_redis_types::Standalone { + host: "redis".into(), + port: 6379, + db: 0, + tls: None, + }) + ); + let [ignored] = resolution.ignored.as_slice() else { + panic!("expected one ignored setting, got {:?}", resolution.ignored); + }; + assert_eq!(ignored.setting, expected); + assert!(!ignored.consequence.is_empty()); +} + +#[rstest] +#[case::kwarg( + "general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_validate_ocsp: true", + &[], + "ssl_validate_ocsp" +)] +#[case::stapled_from_environment( + "model_list: []", + &[("REDIS_HOST", "redis"), ("REDIS_SSL", "true"), ("REDIS_SSL_VALIDATE_OCSP_STAPLED", "yes")], + "REDIS_SSL_VALIDATE_OCSP_STAPLED" +)] +#[case::rediss_url_query( + "general_settings:\n coordination_redis:\n url: rediss://cache:6380/0?ssl_validate_ocsp=true", + &[], + "url ?ssl_validate_ocsp" +)] +fn an_ocsp_check_fails_instead_of_connecting_without_it( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] setting: &str, +) { + let error = coordination_resolution(yaml, environment).unwrap_err(); + + assert!( + matches!(&error, Error::Redis(litellm_redis_resolve::Error::Unenforceable { setting: named, .. }) if named == setting), + "{error:?}" + ); +} + +#[rstest] +#[case::kwarg(" host: redis\n ssl: true\n ssl_validate_ocsp: false")] +#[case::rediss_url_query(" url: rediss://cache:6380/0?ssl_validate_ocsp=false")] +fn an_ocsp_flag_set_to_false_asks_for_nothing(#[case] block: &str) { + let yaml = format!("general_settings:\n coordination_redis:\n{block}"); + + let resolution = coordination_resolution(&yaml, &[]).unwrap(); + + assert!(resolution.ignored.is_empty(), "{:?}", resolution.ignored); +} + +fn static_auth(username: Option<&str>, password: Option<&str>) -> Auth { + Auth::Static(StaticAuth { + username: username.map(litellm_secret_value::SecretValue::new), + password: password.map(litellm_secret_value::SecretValue::new), + }) +} + +#[rstest] +fn an_explicit_host_keeps_environment_credentials() { + let resolved = resolve( + "general_settings:\n coordination_redis:\n host: configured", + &[ + ("REDIS_PASSWORD", "environment-secret"), + ("REDIS_SSL", "True"), + ], + ) + .unwrap(); + + assert_eq!(resolved.auth, static_auth(None, Some("environment-secret"))); + assert!(matches!(resolved.target, Target::Standalone(host) if host.tls.is_some())); +} + +fn azure(username: Option<&str>) -> Auth { + Auth::AzureAd(AzureAdAuth { + client_id: Some("client".to_owned()), + tenant_id: None, + client_secret: None, + username: username.map(str::to_owned), + }) +} + +fn aws(serverless: bool) -> Auth { + Auth::AwsIam(AwsIamAuth { + user_name: "iam-user".to_owned(), + cache_name: "cache-1".to_owned(), + region: "us-east-1".to_owned(), + serverless, + }) +} + +#[rstest] +#[case::static_credentials( + "general_settings:\n coordination_redis:\n host: redis\n username: user\n password: pass", + &[], + static_auth(Some("user"), Some("pass")) +)] +#[case::gcp_beats_azure_and_aws( + "general_settings:\n coordination_redis:\n host: redis\n password: pass\n azure_redis_ad_token: \"True\"\n aws_iam_auth: true", + &[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com")], + Auth::GcpIam { service_account: "projects/-/serviceAccounts/sa@p.iam.gserviceaccount.com".to_owned() } +)] +#[case::azure_beats_aws_and_takes_the_environment_username( + "general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client\n aws_iam_auth: true", + &[("REDIS_USERNAME", "environment-user")], + azure(Some("environment-user")) +)] +#[case::azure_without_an_environment_username_sends_only_the_token( + "general_settings:\n coordination_redis:\n host: redis\n username: configured-user\n azure_redis_ad_token: true\n azure_client_id: client", + &[], + azure(None) +)] +#[case::azure_needs_the_literal_true( + "general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: \"1\"", + &[], + static_auth(None, None) +)] +#[case::aws_with_a_region_from_the_environment( + "general_settings:\n coordination_redis:\n host: redis\n ssl: true\n password: dropped\n aws_iam_auth: \"yes\"\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1", + &[("AWS_DEFAULT_REGION", "us-east-1")], + aws(false) +)] +#[case::aws_serverless_over_a_rediss_url( + "general_settings:\n coordination_redis:\n url: rediss://cache-1:6379\n aws_iam_auth: true\n aws_iam_user_name: iam-user\n aws_iam_cache_name: cache-1\n aws_iam_region: us-east-1\n aws_iam_serverless: true", + &[], + aws(true) +)] +fn auth_follows_python_precedence( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: Auth, +) { + assert_eq!(resolve(yaml, environment).unwrap().auth, expected); +} + +#[rstest] +#[case::coordination( + "general_settings:\n coordination_redis:\n host: redis\n credential_provider: host-object", + &[] +)] +#[case::without_a_target( + "general_settings:\n coordination_redis:\n credential_provider: host-object", + &[("REDIS_HOST", "redis")] +)] +fn a_credential_provider_without_a_host_handle_fails( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], +) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!( + error, + Error::Redis(litellm_redis_resolve::Error::ProviderWithoutHandle) + ), + "{error}" + ); +} + +#[rstest] +#[case::without_tls( + "general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r", + "requires TLS" +)] +#[case::a_plain_url_ignores_ssl( + "general_settings:\n coordination_redis:\n url: redis://cache:6379\n ssl: true\n aws_iam_auth: true\n aws_iam_user_name: u\n aws_iam_cache_name: c\n aws_iam_region: r", + "requires TLS" +)] +#[case::missing_settings( + "general_settings:\n coordination_redis:\n host: redis\n ssl: true\n aws_iam_auth: true\n aws_iam_cache_name: c", + "requires: aws_iam_user_name, aws_iam_region" +)] +fn incomplete_aws_iam_auth_fails(#[case] yaml: &str, #[case] message: &str) { + let error = resolve(yaml, &[]).unwrap_err(); + + assert!(error.to_string().contains(message), "{error}"); +} + +#[rstest] +#[case::defaults_verify_the_certificate_but_not_the_hostname( + " ssl: true", + &[], + Some(TlsOptions { ca_certs: None, ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false }) +)] +#[case::every_option( + " ssl: true\n ssl_ca_certs: /ca.pem\n ssl_certfile: /client.pem\n ssl_keyfile: /client.key\n ssl_password: key-secret\n ssl_cert_reqs: none\n ssl_check_hostname: true", + &[], + Some(TlsOptions { ca_certs: Some("/ca.pem".into()), ca_data: None, ca_path: None, certfile: Some("/client.pem".into()), keyfile: Some("/client.key".into()), key_password: Some(SecretValue::new("key-secret")), verify_certificate: false, check_hostname: true }) +)] +#[case::options_from_the_environment( + " ssl: true", + &[("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional"), ("REDIS_SSL_PASSWORD", "env-key-secret")], + Some(TlsOptions { ca_certs: Some("/env-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: Some(SecretValue::new("env-key-secret")), verify_certificate: true, check_hostname: false }) +)] +#[case::gcp_ca_certs_replace_the_configured_ones( + " ssl: true\n ssl_ca_certs: /ca.pem\n gcp_service_account: account\n gcp_ssl_ca_certs: /gcp-ca.pem", + &[], + Some(TlsOptions { ca_certs: Some("/gcp-ca.pem".into()), ca_data: None, ca_path: None, certfile: None, keyfile: None, key_password: None, verify_certificate: true, check_hostname: false }) +)] +#[case::options_without_ssl_stay_plaintext(" ssl_ca_certs: /ca.pem", &[], None)] +fn tls_options_resolve_like_redis_py( + #[case] block: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: Option, +) { + let yaml = format!("general_settings:\n coordination_redis:\n host: redis\n{block}"); + + let resolved = resolve(&yaml, environment).unwrap(); + + assert!(matches!(resolved.target, Target::Standalone(host) if host.tls == expected)); +} + +#[rstest] +fn an_unknown_cert_requirement_fails() { + let error = resolve( + "general_settings:\n coordination_redis:\n host: redis\n ssl: true\n ssl_cert_reqs: strict", + &[], + ) + .unwrap_err(); + + assert!(matches!( + error, + Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::CertRequirement }) if setting == "ssl_cert_reqs" + )); +} + +#[rstest] +#[case::from_config( + "general_settings:\n coordination_redis:\n sentinel_nodes: [[sentinel-a, 26379], [sentinel-b, \"26380\"]]\n service_name: primary\n sentinel_password: sentinel-secret\n db: 2", + &[] +)] +#[case::from_the_environment( + "model_list: []", + &[ + ("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379], ["sentinel-b", "26380"]]"#), + ("REDIS_SERVICE_NAME", "primary"), + ("REDIS_SENTINEL_PASSWORD", "sentinel-secret"), + ("REDIS_DB", "2"), + ] +)] +fn sentinel_nodes_resolve_to_a_sentinel_endpoint( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], +) { + let Target::Sentinel(sentinel) = resolve(yaml, environment).unwrap().target else { + panic!("expected a sentinel endpoint"); + }; + + assert_eq!( + sentinel.monitors, + vec![ + Node { + host: "sentinel-a".to_owned(), + port: 26379 + }, + Node { + host: "sentinel-b".to_owned(), + port: 26380 + }, + ] + ); + assert_eq!(sentinel.service_name, "primary"); + assert_eq!( + sentinel.monitor_auth, + Some(StaticAuth { + username: None, + password: Some(litellm_secret_value::SecretValue::new("sentinel-secret")), + }) + ); + assert_eq!(sentinel.db, 2); +} + +#[rstest] +#[case::a_url_beats_sentinels( + "general_settings:\n coordination_redis:\n url: redis://single:6379\n sentinel_nodes: [[s, 26379]]\n service_name: primary", + "url:redis://single:6379" +)] +fn sentinel_precedence(#[case] yaml: &str, #[case] expected: &str) { + let resolved = resolve(yaml, &[]).unwrap(); + let description = match resolved.target { + Target::Url(url) => format!("url:{}", url.url.expose()), + Target::Sentinel(sentinel) => format!("sentinel:{}", sentinel.service_name), + other => panic!("unexpected {other:?}"), + }; + + assert_eq!(description, expected); +} + +#[rstest] +fn sentinels_without_a_service_name_fail() { + let error = resolve( + "general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]", + &[], + ) + .unwrap_err(); + + assert!( + matches!( + error, + Error::Redis(litellm_redis_resolve::Error::NoSentinelService) + ), + "{error:?}" + ); +} + +#[rstest] +#[case::configured_host( + "general_settings:\n coordination_redis:\n host: ignored\n sentinel_nodes: [[s, 26379]]\n service_name: primary", + &[], + "host" +)] +#[case::environment_port( + "general_settings:\n coordination_redis:\n sentinel_nodes: [[s, 26379]]\n service_name: primary", + &[("REDIS_PORT", "6380")], + "REDIS_PORT" +)] +fn a_host_or_port_cannot_reach_sentinel( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] source: &str, +) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!(&error, Error::Redis(litellm_redis_resolve::Error::SentinelWithHost { setting }) if setting == source), + "{error:?}" + ); +} + +#[rstest] +#[case::not_pairs( + "general_settings:\n coordination_redis:\n sentinel_nodes: [{host: s, port: 26379}]\n service_name: primary", + &[], + "sentinel_nodes" +)] +#[case::from_the_environment( + "model_list: []", + &[("REDIS_SENTINEL_NODES", "not-json"), ("REDIS_SERVICE_NAME", "primary")], + "REDIS_SENTINEL_NODES" +)] +fn invalid_sentinel_nodes_name_their_source( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] source: &str, +) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::SentinelNodes }) if setting == source), + "{error:?}" + ); +} + +#[rstest] +#[case::from_config("general_settings:\n coordination_redis:\n startup_nodes: [{host: a}]", &[], "startup_nodes")] +#[case::from_the_environment("model_list: []", &[("REDIS_CLUSTER_NODES", "not-json")], "REDIS_CLUSTER_NODES")] +#[case::out_of_range_port_from_the_environment("model_list: []", &[("REDIS_HOST", "redis"), ("REDIS_PORT", "70000")], "REDIS_PORT")] +fn invalid_values_name_their_source( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] source: &str, +) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!(&error, Error::Redis(litellm_redis_resolve::Error::Invalid { setting, .. }) if setting == source), + "{error:?}" + ); +} + +#[rstest] +#[case::falls_back_to_the_socket_timeout(" socket_timeout: 1.5", Some(1500), None)] +#[case::explicit_connect_timeout( + " socket_timeout: 1.5\n socket_connect_timeout: 0.25", + Some(1500), + Some(250) +)] +fn connect_timeout_is_resolved_separately( + #[case] block: &str, + #[case] timeout_ms: Option, + #[case] connect_ms: Option, +) { + let resolved = resolve( + &format!("general_settings:\n coordination_redis:\n host: redis\n{block}"), + &[], + ) + .unwrap(); + + assert_eq!( + resolved.timeout.map(|timeout| timeout.as_millis()), + timeout_ms + ); + assert_eq!( + resolved.connect_timeout.map(|timeout| timeout.as_millis()), + connect_ms + ); +} + +#[rstest] +#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])] +#[case::a_structured_host_with_an_environment_target( + "general_settings:\n coordination_redis:\n host: []", + &[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")] +)] +fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!( + matches!(error, Error::Redis(litellm_redis_resolve::Error::NoTarget)), + "{error:?}" + ); +} + +#[rstest] +fn the_cache_alias_names_cluster_nodes() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000", + ) + .unwrap(); + + assert_eq!(cluster_ports(&resolved), vec![7000]); +} + +#[rstest] +fn canonical_startup_nodes_beat_the_cache_alias() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101", + ) + .unwrap(); + + assert_eq!(cluster_ports(&resolved), vec![7000]); +} + +#[rstest] +fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100", + ) + .unwrap(); + + assert!(matches!(resolved.target, Target::Standalone(_))); +} + +#[rstest] +fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100", + ) + .unwrap(); + + assert!(matches!(resolved.target, Target::Standalone(_))); +} + +#[rstest] +fn unsupported_environment_settings_never_carry_the_value() { + let resolution = coordination_resolution( + "model_list: []", + &[ + ("REDIS_HOST", "redis"), + ("REDIS_SSL_CIPHERS", "private-option-value"), + ], + ) + .unwrap(); + let ignored = format!("{:?}", resolution.ignored); + + assert!(ignored.contains("REDIS_SSL_CIPHERS"), "{ignored}"); + assert!(!ignored.contains("private-option-value"), "{ignored}"); +} + +#[rstest] +#[case::negative("-1")] +#[case::infinite("inf")] +fn rejects_bad_socket_timeout(#[case] timeout: &str) { + let yaml = format!( + "general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}" + ); + let error = resolve(&yaml, &[]).unwrap_err(); + + assert!(matches!( + error, + Error::Redis(litellm_redis_resolve::Error::Invalid { setting, expected: Expected::Seconds }) if setting == "socket_timeout" + )); +} + +#[rstest] +fn debug_redacts_redis_secrets_and_url() { + let config = Config::from_yaml( + "general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password\n sentinel_kwargs: {password: credential-passthrough}", + ) + .unwrap(); + let settings = config.general_settings.coordination_redis.unwrap(); + let configured = ConfiguredRedis::coordination(&settings).unwrap(); + let resolved = resolution(&configured, &HashMap::new()).unwrap(); + let debug = format!("{settings:?} {configured:?} {resolved:?}"); + + assert!(!debug.contains("redis://user:password")); + assert!(!debug.contains("credential-user")); + assert!(!debug.contains("credential-password")); + assert!(!debug.contains("credential-passthrough")); + assert!(debug.contains("redacted") || debug.contains("REDACTED")); +} + +#[rstest] +fn legacy_unknown_settings_do_not_change_resolution() { + let resolved = resolve( + "general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary", + &[], + ) + .unwrap(); + + assert!(matches!(resolved.target, Target::Standalone(_))); +} + +#[rstest] +#[case::configured( + "general_settings:\n coordination_redis:\n host: redis\n max_connections: 47", + &[], + "max_connections" +)] +#[case::from_the_environment( + "general_settings:\n coordination_redis:\n host: redis", + &[("REDIS_MAX_CONNECTIONS", "47")], + "REDIS_MAX_CONNECTIONS" +)] +fn ignored_settings_are_named_by_their_source( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] source: &str, +) { + let resolution = coordination_resolution(yaml, environment).unwrap(); + + let settings: Vec<&str> = resolution + .ignored + .iter() + .map(|ignored| ignored.setting.as_str()) + .collect(); + assert_eq!(settings, vec![source]); +} + +fn snapshot(values: &[(&str, &str)]) -> SecretSnapshot { + let values: HashMap<&str, &str> = values.iter().copied().collect(); + SecretSnapshot::read( + ConfiguredRedis::environment_only().secret_names(), + &|name: &str| values.get(name).map(|value| (*value).to_owned()), + ) +} + +#[rstest] +fn secret_manager_values_reach_the_resolver_through_the_snapshot() { + let settings = redis_cache_kwargs(json!({"host": "os.environ/COORD_HOST"})); + let secrets = SecretSnapshot::read(settings.secret_names(), &|name: &str| match name { + "COORD_HOST" => Some("from-the-manager".to_owned()), + "REDIS_PASSWORD" => Some("manager-secret".to_owned()), + _ => None, + }); + + let resolved = settings.resolve(&secrets, &|_: &str| None).unwrap(); + + assert!( + matches!(&resolved.target, Target::Standalone(host) if host.host == "from-the-manager") + ); + assert_eq!(resolved.auth, static_auth(None, Some("manager-secret"))); +} + +#[rstest] +#[case::from_the_environment(&[], &[("REDIS_USERNAME", "environment-user")], Some("environment-user"))] +#[case::never_from_the_secret_manager(&[("REDIS_USERNAME", "manager-user")], &[], None)] +fn the_azure_username_comes_from_the_process_environment_like_python( + #[case] secrets: &[(&str, &str)], + #[case] environment: &[(&str, &str)], + #[case] expected: Option<&str>, +) { + let settings = redis_cache_kwargs(json!({"host": "redis", "azure_redis_ad_token": "true"})); + let environment: HashMap<&str, &str> = environment.iter().copied().collect(); + + let resolved = settings + .resolve(&snapshot(secrets), &|name: &str| { + environment.get(name).map(|value| (*value).to_owned()) + }) + .unwrap(); + + assert!( + matches!(&resolved.auth, Auth::AzureAd(azure) if azure.username.as_deref() == expected), + "{:?}", + resolved.auth + ); +} + +#[rstest] +#[case::coordination_drops_the_null( + "general_settings:\n coordination_redis:\n host: redis\n ssl: null", + false, + true +)] +#[case::cache_passes_an_unnamed_null_on( + "litellm_settings:\n cache_params:\n type: redis\n host: redis\n ssl: null", + true, + false +)] +#[case::cache_drops_a_null_it_names( + "litellm_settings:\n cache_params:\n type: redis\n host: null\n ssl: true", + true, + true +)] +fn a_null_overrides_the_environment_only_where_python_passes_it_on( + #[case] yaml: &str, + #[case] cache: bool, + #[case] tls: bool, +) { + let environment = [("REDIS_HOST", "environment"), ("REDIS_SSL", "true")]; + let resolved = match cache { + true => resolve_cache_in(yaml, &environment), + false => resolve(yaml, &environment), + } + .unwrap(); + + let Target::Standalone(host) = resolved.target else { + panic!("expected a standalone target, got {:?}", resolved.target); + }; + assert_eq!(host.tls.is_some(), tls); +} + +#[rstest] +#[case::empty_cluster_nodes_use_the_environment( + json!({"startup_nodes": []}), + &[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": 7000}]"#)], + json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": 7000}]"}), + &[("startup_nodes", "REDIS_CLUSTER_NODES")], +)] +#[case::empty_sentinel_nodes_use_the_environment( + json!({"sentinel_nodes": [], "service_name": "primary"}), + &[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#)], + json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "service_name": "primary"}), + &[("sentinel_nodes", "REDIS_SENTINEL_NODES")], +)] +#[case::empty_service_name_uses_the_environment( + json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": ""}), + &[("REDIS_SERVICE_NAME", "primary")], + json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}), + &[("service_name", "REDIS_SERVICE_NAME")], +)] +#[case::environment_host_port_and_password( + json!({}), + &[("REDIS_HOST", "environment"), ("REDIS_PASSWORD", "environment-secret"), ("REDIS_PORT", "6390"), ("REDIS_SSL", "True")], + json!({"port": "6390", "ssl": true, "host": "environment", "password": "environment-secret"}), + &[("host", "REDIS_HOST"), ("password", "REDIS_PASSWORD"), ("port", "REDIS_PORT"), ("ssl", "REDIS_SSL")], +)] +#[case::padded_environment_numbers_parse_like_python_int( + json!({}), + &[("REDIS_HOST", "redis"), ("REDIS_PORT", "6_380"), ("REDIS_DB", "3\n")], + json!({"port": "6_380", "host": "redis", "db": "3\n"}), + &[("db", "REDIS_DB"), ("host", "REDIS_HOST"), ("port", "REDIS_PORT")], +)] +#[case::socket_timeout_from_the_environment( + json!({"host": "redis"}), + &[("REDIS_SOCKET_TIMEOUT", "0.75")], + json!({"socket_timeout": "0.75", "host": "redis"}), + &[("socket_timeout", "REDIS_SOCKET_TIMEOUT")], +)] +#[case::environment_url( + json!({}), + &[("REDIS_URL", "redis://environment:6379/1")], + json!({"url": "redis://environment:6379/1"}), + &[("url", "REDIS_URL")], +)] +#[case::environment_url_beats_the_environment_host( + json!({}), + &[("REDIS_HOST", "ignored"), ("REDIS_URL", "redis://from-url:6379")], + json!({"host": "ignored", "url": "redis://from-url:6379"}), + &[("host", "REDIS_HOST"), ("url", "REDIS_URL")], +)] +#[case::an_explicit_host_drops_the_environment_url( + json!({"host": "configured"}), + &[("REDIS_URL", "redis://environment:6379")], + json!({"host": "configured"}), + &[], +)] +#[case::environment_cluster_nodes( + json!({}), + &[("REDIS_CLUSTER_NODES", r#"[{"host": "node-a", "port": "7000"}, {"host": "node-b", "port": 7001}]"#)], + json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": \"7000\"}, {\"host\": \"node-b\", \"port\": 7001}]"}), + &[("startup_nodes", "REDIS_CLUSTER_NODES")], +)] +#[case::sentinel_from_the_environment( + json!({}), + &[("REDIS_SENTINEL_NODES", r#"[["sentinel-a", 26379]]"#), ("REDIS_SENTINEL_PASSWORD", "sentinel-secret"), ("REDIS_SERVICE_NAME", "primary")], + json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "sentinel_password": "sentinel-secret", "service_name": "primary"}), + &[("sentinel_nodes", "REDIS_SENTINEL_NODES"), ("sentinel_password", "REDIS_SENTINEL_PASSWORD"), ("service_name", "REDIS_SERVICE_NAME")], +)] +#[case::an_explicit_credential_provider_beats_cloud_auth( + json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}), + &[("REDIS_AZURE_AD_TOKEN", "true")], + json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}), + &[], +)] +#[case::gcp_beats_azure_and_aws( + json!({"aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}), + &[("REDIS_GCP_SERVICE_ACCOUNT", "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com")], + json!({"gcp_service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", "aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}), + &[("gcp_service_account", "REDIS_GCP_SERVICE_ACCOUNT")], +)] +#[case::azure_beats_aws_and_takes_the_environment_username( + json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}), + &[("REDIS_USERNAME", "environment-user")], + json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}), + &[], +)] +#[case::an_empty_environment_username_sends_only_the_token( + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis"}), + &[("REDIS_USERNAME", "")], + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": ""}), + &[("username", "REDIS_USERNAME")], +)] +#[case::aws_iam_over_tls_with_the_default_region( + json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true}), + &[("AWS_DEFAULT_REGION", "us-east-1")], + json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true, "aws_iam_region": "us-east-1"}), + &[("aws_iam_region", "AWS_DEFAULT_REGION")], +)] +#[case::aws_region_beats_the_default_region( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true}), + &[("AWS_DEFAULT_REGION", "us-east-1"), ("AWS_REGION", "eu-central-1")], + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true, "aws_iam_region": "eu-central-1"}), + &[("aws_iam_region", "AWS_REGION")], +)] +#[case::a_null_cert_requirement_skips_verification( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}), + &[("REDIS_SSL_CERT_REQS", "required")], + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}), + &[], +)] +#[case::a_cert_requirement_spelled_like_python_none_fails( + json!({"host": "redis"}), + &[("REDIS_SSL", "true"), ("REDIS_SSL_CERT_REQS", "None")], + json!({"ssl": true, "ssl_cert_reqs": "None", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")], +)] +#[case::tls_options_from_the_environment( + json!({"host": "redis"}), + &[("REDIS_SSL", "true"), ("REDIS_SSL_CA_CERTS", "/env-ca.pem"), ("REDIS_SSL_CERT_REQS", "optional")], + json!({"ssl_ca_certs": "/env-ca.pem", "ssl": true, "ssl_cert_reqs": "optional", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_ca_certs", "REDIS_SSL_CA_CERTS"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")], +)] +#[case::environment_references_resolve( + json!({"host": "os.environ/COORD_HOST", "password": "os.environ/COORD_PASSWORD"}), + &[("COORD_HOST", "referenced"), ("COORD_PASSWORD", "referenced-secret")], + json!({"host": "referenced", "password": "referenced-secret"}), + &[], +)] +#[case::malformed_port_from_the_environment( + json!({}), + &[("REDIS_HOST", "redis"), ("REDIS_PORT", "abc")], + json!({"port": "abc", "host": "redis"}), + &[("host", "REDIS_HOST"), ("port", "REDIS_PORT")], +)] +#[case::retry_on_error_from_the_environment( + json!({"host": "redis"}), + &[("REDIS_RETRY_ON_ERROR", "true")], + json!({"retry_on_error": true, "host": "redis"}), + &[("retry_on_error", "REDIS_RETRY_ON_ERROR")], +)] +#[case::ssl_ca_path_from_the_environment( + json!({"host": "redis"}), + &[("REDIS_SSL", "true"), ("REDIS_SSL_CA_PATH", "/etc/ssl/private-cas")], + json!({"ssl": true, "ssl_ca_path": "/etc/ssl/private-cas", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_ca_path", "REDIS_SSL_CA_PATH")], +)] +fn merges_like_python_redis( + #[case] kwargs: serde_json::Value, + #[case] environment: &[(&str, &str)], + #[case] merged: serde_json::Value, + #[case] sources: &[(&str, &str)], +) { + let configured = redis_cache_kwargs(kwargs); + let lookup = |name: &str| { + environment + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| (*value).to_owned()) + }; + + let kwargs = configured + .merged(&SecretSnapshot::read(configured.secret_names(), &lookup)) + .unwrap(); + + let values: serde_json::Map = kwargs + .iter() + .map(|(kwarg, entry)| (kwarg.clone(), entry.value.clone())) + .collect(); + let renamed: Vec<(&str, &str)> = kwargs + .iter() + .filter(|(kwarg, entry)| *kwarg != &entry.source) + .map(|(kwarg, entry)| (kwarg.as_str(), entry.source.as_str())) + .collect(); + assert_eq!(serde_json::Value::Object(values), merged); + assert_eq!(renamed, sources); +} diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 8e37ab02669..3b22f74d95e 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -37,7 +37,7 @@ pub use owned::PythonOwned; pub use runtime::{ ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, enter_native, poll_async_value, ready_future, reserve_process_for_forking, run_async, run_async_value, run_sync, - run_sync_value, runtime_started, + run_sync_value, runtime_handle, runtime_started, }; pub use services::PythonHostCalls; diff --git a/litellm-rust/crates/host-python/src/runtime.rs b/litellm-rust/crates/host-python/src/runtime.rs index e8872c9340a..4f91e2e085b 100644 --- a/litellm-rust/crates/host-python/src/runtime.rs +++ b/litellm-rust/crates/host-python/src/runtime.rs @@ -68,6 +68,11 @@ fn runtime() -> PyResult<&'static Runtime> { Ok(pyo3_async_runtimes::tokio::get_runtime()) } +/// The native runtime, for work that must outlive the call that schedules it. +pub fn runtime_handle() -> PyResult { + Ok(runtime()?.handle().clone()) +} + #[expect(clippy::disallowed_methods, reason = "this is the gated door")] fn future_into_py(py: Python<'_>, future: F) -> PyResult> where diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index ba892199dcf..ff8f540c806 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -30,6 +30,14 @@ bytes.workspace = true futures-util.workspace = true litellm-cache.workspace = true litellm-error.workspace = true +litellm-circuit-breaker.workspace = true +litellm-redis-client.workspace = true +litellm-redis-connect.workspace = true +litellm-redis-store.workspace = true +litellm-redis-resolve.workspace = true +litellm-python-compat.workspace = true +litellm-redis-types.workspace = true +tracing.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index f4a4f7c2b26..0a9efbdcac3 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -12,6 +12,7 @@ mod marshal; mod model_registry; mod preflight; mod python_settings; +mod redis; mod router; mod routes; mod secrets; @@ -37,6 +38,14 @@ mod _native { #[pymodule_export] use crate::model_registry::bundled_model_cost_map; #[pymodule_export] + use crate::redis::client::NativeRedisClient; + #[pymodule_export] + use crate::redis::client::native_redis_breaker_states; + #[pymodule_export] + use crate::redis::store::request::{NativeRequestRedis, NativeSend}; + #[pymodule_export] + use crate::redis::store::{NativePending, NativeRedisStoreError}; + #[pymodule_export] use crate::router::{RouterCall, RouterCoreHandle}; #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; @@ -111,6 +120,12 @@ mod tests { "aresponses", "ResponsesWebSocketConnection", "NativeDiagnosticProcessor", + "NativeRedisClient", + "NativeRequestRedis", + "NativeSend", + "NativePending", + "NativeRedisStoreError", + "native_redis_breaker_states", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "NativeDatabase", diff --git a/litellm-rust/crates/python-bridge/src/redis/client/credentials.rs b/litellm-rust/crates/python-bridge/src/redis/client/credentials.rs new file mode 100644 index 00000000000..21a37947c30 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/client/credentials.rs @@ -0,0 +1,47 @@ +use std::sync::Arc; + +use litellm_host_python::{PythonContext, attach_blocking}; +use litellm_redis_connect::{CredentialSource, CredentialStage, Error, Expiring, SourceFuture}; +use litellm_secret_value::SecretValue; +use pyo3::prelude::*; + +const MODULE: &str = "litellm.rust_bridge.host.redis_credentials"; + +/// A redis-py `CredentialProvider` the SDK caller passed in `cache_kwargs`, asked through the host +/// module the way redis-py's async connection asks it. +pub(crate) struct PythonCredentialSource { + provider: Arc>, + context: PythonContext, +} + +impl PythonCredentialSource { + pub(crate) fn new(provider: Py, context: PythonContext) -> Self { + Self { + provider: Arc::new(provider), + context, + } + } +} + +impl CredentialSource for PythonCredentialSource { + fn fetch(&self) -> SourceFuture<'_> { + let provider = Arc::clone(&self.provider); + let context = self.context.clone(); + Box::pin(async move { + let (username, password) = attach_blocking(context, move |py| { + py.import(MODULE)? + .getattr("credentials")? + .call1((provider.bind(py),))? + .extract::<(Option, String)>() + }) + .await + .and_then(|credentials| credentials) + .map_err(Error::credential(CredentialStage::Provider))?; + Ok(Expiring { + username, + password: SecretValue::new(password), + expires_at: None, + }) + }) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/client/mod.rs b/litellm-rust/crates/python-bridge/src/redis/client/mod.rs new file mode 100644 index 00000000000..c47e9aec671 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/client/mod.rs @@ -0,0 +1,385 @@ +mod credentials; + +use std::{ + collections::BTreeMap, + sync::{Arc, LazyLock, Mutex, PoisonError}, + time::Duration, +}; + +use litellm_circuit_breaker::Config as BreakerConfig; +use litellm_core_utils::settings::ProcessEnvironment; +use litellm_host_python::PythonContext; +use litellm_python_compat::number::{self, SecondsError}; +use litellm_redis_client::{RedisClient, RedisEvents, WeakRedisClient}; +use litellm_redis_connect::{CredentialSource, HostServices, SharedEnvironment}; +use litellm_redis_resolve::{Kwargs, PythonException}; +use litellm_redis_types::{ProviderHandle, Resolution}; +use pyo3::{ + exceptions::{PyOverflowError, PyTypeError, PyValueError}, + prelude::*, + sync::MutexExt, + types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple}, +}; + +use self::credentials::PythonCredentialSource; +use crate::redis::events::{ServiceLoggingEvents, state_label}; + +pyo3::import_exception!(redis.exceptions, RedisClusterException); + +/// The kwargs `_merge_redis_kwargs` returns for a `RedisCache`, which the resolver reads the way +/// redis-py's constructors do. A key whose value is a Python object, such as `retry`, reaches the +/// resolver by name only, so it can name the setting it cannot honor. A `credential_provider` object +/// is also kept, so the connector can ask it for credentials, and reaches the resolver as a handle. +pub(crate) struct RedisSpec { + kwargs: Kwargs, + credential_provider: Option>, +} + +impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec { + type Error = PyErr; + + fn extract(kwargs: Borrowed<'a, 'py, PyAny>) -> PyResult { + let py = kwargs.py(); + let plain = PyDict::new(py); + let mut host_objects = Vec::new(); + for (name, value) in kwargs.cast::()?.iter() { + match value.is_none() || is_plain(&value) { + true => plain.set_item(&name, value)?, + false => host_objects.push(name.extract::()?), + } + } + let credential_provider = kwargs + .cast::()? + .get_item("credential_provider")? + .filter(|provider| !provider.is_none() && !is_plain(provider)) + .map(Bound::unbind); + let fields: BTreeMap = + litellm_host_python::from_py_argument(plain.as_any())?; + let objects = host_objects + .into_iter() + .map(|name| (name, serde_json::Value::Bool(true))); + Ok(Self { + kwargs: litellm_redis_resolve::kwargs(fields.into_iter().chain(objects)), + credential_provider, + }) + } +} + +/// `litellm.constants`' `REDIS_CIRCUIT_BREAKER_*`, which Python reads from the environment once. +pub(crate) struct RedisBreaker(BreakerConfig); + +impl<'a, 'py> FromPyObject<'a, 'py> for RedisBreaker { + type Error = PyErr; + + fn extract(breaker: Borrowed<'a, 'py, PyAny>) -> PyResult { + let (failure_threshold, recovery_timeout, timeout_min_duration): (u32, f64, f64) = + breaker.extract()?; + let seconds = |seconds: f64| { + Duration::try_from_secs_f64(seconds) + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string())) + }; + Ok(Self(BreakerConfig { + failure_threshold, + recovery_timeout: seconds(recovery_timeout)?, + timeout_min_duration: seconds(timeout_min_duration)?, + })) + } +} + +fn is_plain(value: &Bound<'_, PyAny>) -> bool { + value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() +} + +/// A client the process built, held weakly so a discarded `Router` drops its client and provider. +struct Shared { + resolved: Resolution, + client: WeakRedisClient, +} + +static CLIENTS: LazyLock>> = LazyLock::new(|| Mutex::new(Vec::new())); + +fn live_clients(py: Python<'_>) -> std::sync::MutexGuard<'static, Vec> { + let mut clients = CLIENTS + .lock_py_attached(py) + .unwrap_or_else(PoisonError::into_inner); + clients.retain(|shared| shared.client.upgrade().is_some()); + clients +} + +/// The one client for a resolved Redis config and credential provider, so every family in this process shares its pool +/// and its breaker. A config redis-py refuses building its client raises the same class here. One it only fails on +/// once a command connects, where `RedisCache` logs every failed call, leaves the family on memory with an error. +fn shared_client( + py: Python<'_>, + spec: RedisSpec, + namespace: Option, + breaker: Option, +) -> PyResult> { + let environment: SharedEnvironment = Arc::new(ProcessEnvironment); + let RedisSpec { + kwargs, + credential_provider, + } = spec; + // The registry entry's client keeps the provider alive, so no live entry shares its address + let handle = credential_provider + .as_ref() + .map(|provider| ProviderHandle::new(provider.as_ptr() as usize as u64)); + let logger = crate::logger::capture(py); + let environ = py.import("os")?.getattr("environ")?; + let resolved = match logger.scope(|| resolve(&kwargs, handle, &environ))? { + Ok(resolution) => { + logger.scope(|| { + for ignored in &resolution.ignored { + litellm_tracing::warn!( + "Redis setting {} is not supported by the Rust client and is ignored: {}", + ignored.setting, + ignored.consequence + ); + } + }); + Resolution { + namespace, + ..resolution + } + } + Err(error) => return unresolved(py, &error), + }; + let runtime = litellm_host_python::runtime_handle()?; + let events: Arc = Arc::new(ServiceLoggingEvents::new(py, runtime.clone())?); + let host = HostServices { + events: Some(events), + http: crate::http::host_client(py, litellm_http::ClientVariant::Provider)?, + environment, + credential_provider: credential_provider + .as_ref() + .map(|provider| -> PyResult> { + Ok(Arc::new(PythonCredentialSource::new( + provider.clone_ref(py), + PythonContext::capture(py)?, + ))) + }) + .transpose()?, + }; + let mut clients = live_clients(py); + if let Some(client) = clients + .iter() + .filter(|shared| shared.resolved == resolved) + .find_map(|shared| shared.client.upgrade()) + { + return Ok(Some(client)); + } + let connected = { + let _runtime = runtime.enter(); + litellm_redis_connect::connect(resolved.clone(), host, breaker.map(|breaker| breaker.0)) + }; + let client = match connected { + Ok(client) => client, + Err(error) => { + drop(clients); + return Ok(unconnected(py, &error)); + } + }; + clients.push(Shared { + resolved, + client: client.downgrade(), + }); + Ok(Some(client)) +} + +/// The Rust client for one `RedisCache`, built once and handed to every native store, so a store +/// never sees the redis-py kwargs or the namespace. +#[pyclass(frozen, name = "NativeRedisClient")] +pub(crate) struct NativeRedisClient { + pub(crate) client: RedisClient, + /// `RedisCache.default_ttl`, the TTL a store's write takes when its caller gives none. + pub(crate) default_ttl: Duration, +} + +#[pymethods] +impl NativeRedisClient { + /// `None` when redis-py would fail only once a command connects, which leaves every store given it on memory. + #[staticmethod] + #[pyo3(signature = (redis_kwargs, *, default_ttl_seconds, namespace = None, breaker = None))] + fn connect( + py: Python<'_>, + redis_kwargs: RedisSpec, + default_ttl_seconds: f64, + namespace: Option, + breaker: Option, + ) -> PyResult> { + let default_ttl = number::seconds(default_ttl_seconds).map_err(seconds_error)?; + Ok( + shared_client(py, redis_kwargs, namespace, breaker)?.map(|client| Self { + client, + default_ttl, + }), + ) + } + + #[getter] + fn default_ttl_seconds(&self) -> f64 { + self.default_ttl.as_secs_f64() + } +} + +/// The breaker state of every native Redis client this process built, so the Prometheus state +/// gauge counts them next to the Python `RedisCache` breakers. +#[pyfunction] +pub(crate) fn native_redis_breaker_states(py: Python<'_>) -> Vec<&'static str> { + live_clients(py) + .iter() + .filter_map(|shared| shared.client.upgrade()) + .filter_map(|client| client.breaker_state()) + .map(|state| state_label(&state)) + .collect() +} + +/// What redis-py makes of the kwargs, with `REDIS_USERNAME` read from `environ` the way Python's Azure AD auth +/// reads `os.environ`. +fn resolve( + kwargs: &Kwargs, + provider: Option, + environ: &Bound<'_, PyAny>, +) -> PyResult> { + let redis_username: Option = environ + .call_method1("get", ("REDIS_USERNAME",))? + .extract()?; + Ok(litellm_redis_resolve::interpret( + kwargs, + provider, + redis_username.as_deref(), + )) +} + +fn unresolved( + py: Python<'_>, + error: &litellm_redis_resolve::Error, +) -> PyResult> { + let message = || litellm_error::Report::new(error).to_string(); + match error.raised_building_the_python_client() { + Some(PythonException::ValueError) => Err(PyValueError::new_err(message())), + Some(PythonException::TypeError) => Err(PyTypeError::new_err(message())), + Some(PythonException::RedisClusterException) => { + Err(RedisClusterException::new_err(message())) + } + None => Ok(memory_only(py, error)), + } +} + +// redis-py reads TLS files and fetches credentials only once a command connects, so none of these fails building +// its client +fn unconnected(py: Python<'_>, error: &litellm_redis_connect::Error) -> Option { + match error { + litellm_redis_connect::Error::Connect(_) + | litellm_redis_connect::Error::NoHostCredentialProvider + | litellm_redis_connect::Error::ProcessEnvironmentRequired { .. } + | litellm_redis_connect::Error::Credential { .. } + | litellm_redis_connect::Error::CredentialTimeout { .. } + | litellm_redis_connect::Error::ExpiredCredential + | litellm_redis_connect::Error::MissingCredential { .. } => memory_only(py, error), + } +} + +fn memory_only(py: Python<'_>, error: &(dyn std::error::Error + 'static)) -> Option { + let report = litellm_error::Report::new(error).to_string(); + crate::logger::capture(py).scope(|| { + litellm_tracing::error!( + "Redis is not usable from the Rust stores, their state stays in memory: {report}" + ) + }); + None +} + +fn seconds_error(error: SecondsError) -> PyErr { + match error { + SecondsError::Overflow => PyOverflowError::new_err(error.to_string()), + SecondsError::NotANumber => PyValueError::new_err(error.to_string()), + } +} + +#[cfg(test)] +mod tests { + use std::ffi::CStr; + + use litellm_redis_types::Auth; + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn kwargs(fields: serde_json::Value) -> Kwargs { + let fields: BTreeMap = serde_json::from_value(fields).unwrap(); + litellm_redis_resolve::kwargs(fields) + } + + #[rstest] + #[case::set(c"environ = {'REDIS_USERNAME': 'azure-user'}", Some("azure-user"))] + #[case::empty_is_unset(c"environ = {'REDIS_USERNAME': ''}", None)] + fn azure_auth_takes_its_username_from_the_environ_like_python( + #[case] source: &CStr, + #[case] expected: Option<&str>, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run(source, Some(&locals), Some(&locals)).unwrap(); + let environ = locals.get_item("environ").unwrap().unwrap(); + let kwargs = kwargs(json!({"host": "redis", "azure_redis_ad_token": "true"})); + + let resolved = resolve(&kwargs, None, &environ).unwrap().unwrap(); + + assert!( + matches!(&resolved.auth, Auth::AzureAd(azure) if azure.username.as_deref() == expected), + "{:?}", + resolved.auth + ); + }); + } + + #[rstest] + #[case::no_target(json!({"port": 6379}), Some("ValueError"))] + #[case::unenforceable(json!({"host": "redis", "ssl": true, "ssl_validate_ocsp": true}), Some("TypeError"))] + #[case::fails_only_once_a_command_connects( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": "None"}), + None + )] + fn a_config_raises_where_redis_py_raises_building_its_client( + #[case] fields: serde_json::Value, + #[case] raised: Option<&str>, + ) { + Python::initialize(); + Python::attach(|py| { + let error = litellm_redis_resolve::interpret(&kwargs(fields), None, None).unwrap_err(); + + let outcome = unresolved(py, &error) + .map(|client| client.is_some()) + .map_err(|error| error.get_type(py).name().unwrap().to_string()); + + assert_eq!( + outcome, + raised.map_or(Ok(false), |class| Err(class.to_owned())) + ); + }); + } + + #[rstest] + #[case::overflow(f64::INFINITY, "OverflowError")] + #[case::not_a_number(f64::NAN, "ValueError")] + fn a_default_ttl_python_cannot_make_an_int_raises_like_int( + #[case] seconds: f64, + #[case] raised: &str, + ) { + Python::initialize(); + Python::attach(|py| { + let error = number::seconds(seconds).map_err(seconds_error).unwrap_err(); + + assert_eq!(error.get_type(py).name().unwrap().to_string(), raised); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/events.rs b/litellm-rust/crates/python-bridge/src/redis/events.rs new file mode 100644 index 00000000000..04b0c9dbf32 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/events.rs @@ -0,0 +1,391 @@ +use std::{ + sync::Arc, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_host_python::{PythonContext, attach_blocking}; +use litellm_redis_client::{ + CallEvent, CallTag, FailureKind, Operation, Outcome, RedisEvents, current_tag, tagged, +}; +use pyo3::{prelude::*, types::PyModule}; +use tokio::runtime::Handle; + +/// The asyncio loop a Python caller ran on. `RedisCache` scheduled its service events there with +/// `asyncio.create_task`, so the native events go there too instead of a new loop per event. +#[derive(Clone)] +pub(crate) struct EventLoop(Arc>); + +impl EventLoop { + pub(crate) fn running(py: Python<'_>) -> Option { + pyo3_async_runtimes::get_running_loop(py) + .ok() + .map(|event_loop| Self(Arc::new(event_loop.unbind()))) + } +} + +/// The Python caller of one store operation. Its context carries `with_service_target` and the +/// ambient OTel span, and its `call_type` is the label `RedisCache` reports, so dashboards keep +/// their series. +pub(crate) struct RedisCall { + pub(crate) context: PythonContext, + pub(crate) event_loop: Option, + pub(crate) call_type: &'static str, +} + +impl RedisCall { + pub(crate) fn tag(self) -> CallTag { + Arc::new(self) + } + + pub(crate) async fn scope(self, operation: F) -> F::Output { + tagged(self.tag(), operation).await + } +} + +/// Reports every native Redis call and breaker transition to `ServiceLogging`. +pub(crate) struct ServiceLoggingEvents { + record_call: Arc>, + record_breaker: Arc>, + context: PythonContext, + /// The loop the client was built on, for calls whose caller had none running. + event_loop: Option, + runtime: Handle, +} + +impl ServiceLoggingEvents { + pub(crate) fn new(py: Python<'_>, runtime: Handle) -> PyResult { + const MODULE: &str = "litellm.rust_bridge.host.redis_events"; + Self::from_module(py, &py.import(MODULE)?, runtime) + } + + fn from_module( + py: Python<'_>, + module: &Bound<'_, PyModule>, + runtime: Handle, + ) -> PyResult { + Ok(Self { + record_call: Arc::new(module.getattr("record_redis_call")?.unbind()), + record_breaker: Arc::new(module.getattr("record_redis_breaker")?.unbind()), + context: PythonContext::capture(py)?, + event_loop: EventLoop::running(py), + runtime, + }) + } + + fn report(&self, context: PythonContext, record: Arc>, args: ReportArgs) { + self.runtime.spawn(async move { + let reported = attach_blocking(context, move |py| match args { + ReportArgs::Call(call) => record + .bind(py) + .call1(( + call.event_loop + .as_ref() + .map(|event_loop| event_loop.0.bind(py)), + call.call_type, + call.start_time, + call.end_time, + call.failure, + )) + .map(drop), + ReportArgs::Breaker(state) => record.bind(py).call1((state,)).map(drop), + }) + .await + .and_then(|reported| reported); + if let Err(error) = reported { + tracing::warn!(error = %error, "reporting a Redis event to ServiceLogging failed"); + } + }); + } +} + +/// One call as `record_redis_call` takes it, its times in `time.time()` seconds. +struct CallReport { + event_loop: Option, + call_type: &'static str, + start_time: f64, + end_time: f64, + failure: Option<&'static str>, +} + +enum ReportArgs { + Call(CallReport), + Breaker(&'static str), +} + +/// The wall-clock start and end of a call that just finished after `duration`. +fn wall_clock_span(duration: Duration) -> (f64, f64) { + let end = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default(); + ( + end.saturating_sub(duration).as_secs_f64(), + end.as_secs_f64(), + ) +} + +fn operation_call_type(operation: Operation) -> &'static str { + match operation { + Operation::Command(name) => name, + Operation::Script => "eval", + Operation::Pipeline { .. } => "pipeline", + } +} + +fn failure_label(outcome: Outcome) -> Option<&'static str> { + match outcome { + Outcome::Succeeded => None, + Outcome::Failed(FailureKind::Timeout) => Some("timeout"), + Outcome::Failed(FailureKind::Transport) => Some("transport"), + Outcome::Failed(FailureKind::Server) => Some("server"), + Outcome::Failed(FailureKind::Decode) => Some("decode"), + Outcome::Failed(FailureKind::BreakerOpen) => Some("breaker_open"), + } +} + +pub(crate) fn state_label(state: &State) -> &'static str { + match state { + State::Closed { .. } => "closed", + State::Open { .. } => "open", + State::HalfOpen => "half_open", + } +} + +impl RedisEvents for ServiceLoggingEvents { + fn call(&self, event: &CallEvent) { + let scoped = current_tag().and_then(|tag| tag.downcast::().ok()); + let (context, event_loop, call_type) = match scoped { + Some(call) => ( + call.context.clone(), + call.event_loop.clone().or_else(|| self.event_loop.clone()), + call.call_type, + ), + None => ( + self.context.clone(), + self.event_loop.clone(), + operation_call_type(event.operation), + ), + }; + let (start_time, end_time) = wall_clock_span(event.duration); + self.report( + context, + self.record_call.clone(), + ReportArgs::Call(CallReport { + event_loop, + call_type, + start_time, + end_time, + failure: failure_label(event.outcome), + }), + ); + } + + fn breaker(&self, transition: &Transition) { + self.report( + self.context.clone(), + self.record_breaker.clone(), + ReportArgs::Breaker(state_label(&transition.to)), + ); + } +} + +#[cfg(test)] +mod tests { + use std::{sync::Arc, time::Duration}; + + use litellm_host_python::PythonContext; + use litellm_redis_client::{CallEvent, FailureKind, Operation, Outcome, RedisEvents}; + use pyo3::{ + prelude::*, + types::{PyCFunction, PyModule, PyString}, + }; + use rstest::rstest; + use tokio::runtime::Handle; + + use super::{EventLoop, RedisCall, ServiceLoggingEvents}; + + const RECORDER: &str = " +import contextvars, time +target = contextvars.ContextVar('target', default=None) +seen = [] +def record_redis_call(event_loop, call_type, start_time, end_time, failure): + just_now = abs(time.time() - end_time) < 60 + seen.append((event_loop, call_type, round(end_time - start_time, 3), just_now, failure, target.get())) +def record_redis_breaker(state): + seen.append(('breaker', state)) +"; + + fn stand_in_loop(py: Python<'_>, name: &str) -> EventLoop { + EventLoop(Arc::new(PyString::new(py, name).into_any().unbind())) + } + + fn recorder(name: &str, client_loop: Option<&str>) -> (ServiceLoggingEvents, Py) { + Python::initialize(); + Python::attach(|py| { + let name = std::ffi::CString::new(format!("redis_events_recorder_{name}")).unwrap(); + let module = + PyModule::from_code(py, &std::ffi::CString::new(RECORDER).unwrap(), &name, &name) + .unwrap(); + let events = ServiceLoggingEvents::from_module(py, &module, Handle::current()).unwrap(); + ( + ServiceLoggingEvents { + event_loop: client_loop.map(|client_loop| stand_in_loop(py, client_loop)), + ..events + }, + module.unbind(), + ) + }) + } + + async fn seen(module: &Py, count: usize) -> String { + let deadline = tokio::time::Instant::now() + Duration::from_secs(10); + loop { + let seen: (usize, String) = Python::attach(|py| { + let seen = module.bind(py).getattr("seen").unwrap(); + (seen.len().unwrap(), seen.repr().unwrap().to_string()) + }); + if seen.0 >= count || tokio::time::Instant::now() > deadline { + return seen.1; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + } + + fn scoped_call(namespace: &Py, caller_loop: Option<&str>) -> RedisCall { + Python::attach(|py| { + namespace + .bind(py) + .getattr("target") + .unwrap() + .call_method1("set", ("spend_counter",)) + .unwrap(); + RedisCall { + context: PythonContext::capture(py).unwrap(), + event_loop: caller_loop.map(|caller_loop| stand_in_loop(py, caller_loop)), + call_type: "async_increment", + } + }) + } + + #[rstest] + #[case::on_the_callers_loop(Some("caller loop"), "caller loop")] + #[case::on_the_clients_loop_when_the_caller_had_none(None, "client loop")] + #[tokio::test(flavor = "multi_thread")] + async fn a_scoped_call_reports_the_store_labels_and_its_times_in_the_callers_context( + #[case] caller_loop: Option<&str>, + #[case] expected_loop: &str, + ) { + let (events, namespace) = recorder( + expected_loop.replace(' ', "_").as_str(), + Some("client loop"), + ); + let call = scoped_call(&namespace, caller_loop); + let event = CallEvent { + operation: Operation::Command("INCRBYFLOAT"), + duration: Duration::from_millis(250), + outcome: Outcome::Failed(FailureKind::Timeout), + }; + + call.scope(async { events.call(&event) }).await; + + assert_eq!( + seen(&namespace, 1).await, + format!( + "[('{expected_loop}', 'async_increment', 0.25, True, 'timeout', 'spend_counter')]" + ) + ); + } + + #[rstest] + #[case::opened(litellm_circuit_breaker::State::Open { since: std::time::Instant::now() }, "open")] + #[case::probing(litellm_circuit_breaker::State::HalfOpen, "half_open")] + #[tokio::test(flavor = "multi_thread")] + async fn breaker_transitions_report_python_state_names( + #[case] to: litellm_circuit_breaker::State, + #[case] expected: &str, + ) { + let (events, namespace) = recorder(expected, None); + + events.breaker(&litellm_circuit_breaker::Transition { + from: litellm_circuit_breaker::State::HalfOpen, + to, + generation: 1, + }); + + assert_eq!( + seen(&namespace, 1).await, + format!("[('breaker', '{expected}')]") + ); + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn a_transition_on_a_thread_without_a_runtime_is_still_reported() { + let (events, namespace) = recorder("plain_thread", None); + + std::thread::spawn(move || { + events.breaker(&litellm_circuit_breaker::Transition { + from: litellm_circuit_breaker::State::HalfOpen, + to: litellm_circuit_breaker::State::HalfOpen, + generation: 1, + }); + }) + .join() + .unwrap(); + + assert_eq!(seen(&namespace, 1).await, "[('breaker', 'half_open')]"); + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn an_unscoped_call_falls_back_to_the_command_name_and_the_clients_loop() { + let (events, namespace) = recorder("unscoped", None); + + events.call(&CallEvent { + operation: Operation::Pipeline { ops: 3 }, + duration: Duration::from_millis(5), + outcome: Outcome::Succeeded, + }); + + assert_eq!( + seen(&namespace, 1).await, + "[(None, 'pipeline', 0.005, True, None, None)]" + ); + } + + #[rstest] + fn running_is_the_loop_the_caller_runs_on_and_none_outside_one() { + Python::initialize(); + Python::attach(|py| { + let module = PyModule::from_code( + py, + c"import asyncio\nasync def probe_inside(probe):\n return probe() is asyncio.get_running_loop()\n", + c"event_loop_probe", + c"event_loop_probe", + ) + .unwrap(); + let probe = PyCFunction::new_closure(py, None, None, |_, _| { + Python::attach(|py| -> PyResult>> { + Ok(EventLoop::running(py).map(|event_loop| event_loop.0.clone_ref(py))) + }) + }) + .unwrap(); + let coroutine = module + .getattr("probe_inside") + .unwrap() + .call1((probe,)) + .unwrap(); + + let inside: bool = py + .import("asyncio") + .unwrap() + .call_method1("run", (coroutine,)) + .unwrap() + .extract() + .unwrap(); + + assert!(inside); + assert!(EventLoop::running(py).is_none()); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/mod.rs b/litellm-rust/crates/python-bridge/src/redis/mod.rs new file mode 100644 index 00000000000..ad91f8349ac --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/mod.rs @@ -0,0 +1,11 @@ +#[expect( + dead_code, + reason = "the native stores that use these helpers land in the next PR" +)] +pub(crate) mod client; +mod events; +#[expect( + dead_code, + reason = "the native stores that use these helpers land in the next PR" +)] +pub(crate) mod store; diff --git a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs new file mode 100644 index 00000000000..2c48bf2c689 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs @@ -0,0 +1,248 @@ +pub(crate) mod request; +mod value; + +use value::PyValue; + +use std::{ + pin::Pin, + sync::{Mutex, PoisonError}, + time::Duration, +}; + +use litellm_host_python::PythonContext; +use litellm_python_compat::number::{SecondsError, seconds}; +use litellm_redis_store::{Error, FromValue, IntoValue, Number, Pending, Remote}; +use pyo3::{ + IntoPyObjectExt, + exceptions::{PyOverflowError, PyRuntimeError, PyTypeError, PyValueError}, + prelude::*, + types::{PyBool, PyFloat, PyInt}, +}; + +use crate::execution::{run_async_value, run_sync_value}; +use crate::redis::{ + client::NativeRedisClient, + events::{EventLoop, RedisCall}, +}; + +/// The labels `RedisCache` reports for this operation, in the caller's context so its +/// `service_target` reaches the service event, dispatched on the loop the caller runs on. +fn call(py: Python<'_>, call_type: &'static str) -> PyResult { + Ok(RedisCall { + context: PythonContext::capture(py)?, + event_loop: EventLoop::running(py), + call_type, + }) +} + +/// Runs one Redis operation as a Python awaitable, reported to `ServiceLogging` as `call_type`. +fn redis_async<'py, T, E: 'static>( + py: Python<'py>, + call_type: &'static str, + error: fn(E) -> PyErr, + operation: impl Future> + Send + 'static, +) -> PyResult> +where + T: for<'a> IntoPyObject<'a> + Send + 'static, +{ + let scope = call(py, call_type)?; + run_async_value( + py, + async move { scope.scope(operation).await.map_err(error) }, + ) +} + +/// Runs one Redis operation with the GIL released, for the router's sync paths. +fn redis_sync( + py: Python<'_>, + call_type: &'static str, + error: fn(E) -> PyErr, + operation: impl Future> + Send + 'static, +) -> PyResult +where + T: Send + 'static, +{ + let scope = call(py, call_type)?; + run_sync_value( + py, + async move { scope.scope(operation).await.map_err(error) }, + ) +} + +/// A TTL in seconds, raising what `int(seconds)` raises for infinity and NaN. +fn duration(value: f64) -> PyResult { + seconds(value).map_err(|error| match error { + SecondsError::Overflow => PyOverflowError::new_err(error.to_string()), + SecondsError::NotANumber => PyValueError::new_err(error.to_string()), + }) +} + +fn ttl(value: Option) -> PyResult> { + value.map(duration).transpose() +} + +fn remote(client: &NativeRedisClient) -> Remote { + Remote { + client: client.client.clone(), + default_ttl: client.default_ttl, + } +} + +type Resolve = Pin>> + Send>>; + +/// A Redis call declared on a request's batch. `resolve` flushes that batch if it has not gone +/// out and answers the call, once. +#[pyclass(frozen, name = "NativePending")] +pub(crate) struct NativePending { + resolve: Mutex>, + call_type: &'static str, +} + +impl NativePending { + pub(crate) fn new( + call_type: &'static str, + error: fn(E) -> PyErr, + reply: impl Future> + Send + 'static, + ) -> Self + where + T: for<'a> IntoPyObject<'a> + Send + 'static, + { + let resolve = async move { + let reply = reply.await.map_err(error)?; + Python::attach(|py| reply.into_py_any(py)) + }; + Self { + resolve: Mutex::new(Some(Box::pin(resolve))), + call_type, + } + } +} + +#[pymethods] +impl NativePending { + fn resolve<'py>(&self, py: Python<'py>) -> PyResult> { + let resolve = self + .resolve + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take() + .ok_or_else(|| PyRuntimeError::new_err("this Redis call was already resolved"))?; + redis_async(py, self.call_type, std::convert::identity, resolve) + } +} + +pyo3::create_exception!( + _native, + NativeRedisStoreError, + pyo3::exceptions::PyException, + "A native Redis store operation failed. Args are (kind, message), kind spelled as litellm-redis-store's ErrorKind." +); + +/// What every native store raises, its kind telling a caller which failures it may drop. +fn kind_error(error: Error) -> PyErr { + let kind: &'static str = error.kind().into(); + NativeRedisStoreError::new_err((kind, litellm_error::Report::new(&error).to_string())) +} + +/// A store op declared on the batch the caller's `NativeSend` picked, answered as `reply` shapes it. +fn pending( + call_type: &'static str, + error: fn(Error) -> PyErr, + declared: Pending, + reply: fn(T) -> R, +) -> NativePending +where + T: Send + 'static, + R: for<'a> IntoPyObject<'a> + Send + 'static, +{ + NativePending::new(call_type, error, async move { declared.await.map(reply) }) +} + +/// A Python `int` or `float` across the bridge, keeping which one it is. +pub(crate) struct PyNumber(Number); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyNumber { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + if value.is_instance_of::() && !value.is_instance_of::() { + return Ok(Self(Number::Int(value.extract()?))); + } + if value.is_instance_of::() { + return Ok(Self(Number::Float(value.extract()?))); + } + Err(PyTypeError::new_err("expected an int or a float")) + } +} + +impl<'py> IntoPyObject<'py> for PyNumber { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + Ok(match self.0 { + Number::Int(value) => value.into_pyobject(py)?.into_any(), + Number::Float(value) => value.into_pyobject(py)?.into_any(), + }) + } +} + +/// A store type whose Python side is plain data, converted through its `FromValue` and `IntoValue` +/// so the bridge never names a field. A shape the store refuses is `ValueError`. +pub(crate) struct PyData(pub(crate) T); + +impl<'a, 'py, T: FromValue> FromPyObject<'a, 'py> for PyData { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let PyValue(value) = value.extract()?; + T::from_value(&value) + .map(Self) + .map_err(|error| PyValueError::new_err(error.to_string())) + } +} + +impl<'py, T: IntoValue> IntoPyObject<'py> for PyData { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + PyValue(self.0.into_value()).into_pyobject(py) + } +} + +fn unwrap_all(values: Vec>) -> Vec { + values.into_iter().map(|value| value.0).collect() +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use pyo3::prelude::*; + use rstest::rstest; + + use super::ttl; + + #[rstest] + #[case::infinity(f64::INFINITY, "OverflowError")] + #[case::nan(f64::NAN, "ValueError")] + fn a_ttl_int_refuses_raises_what_int_raises(#[case] seconds: f64, #[case] raised: &str) { + Python::initialize(); + Python::attach(|py| { + let error = ttl(Some(seconds)).unwrap_err(); + + assert_eq!(error.get_type(py).name().unwrap().to_string(), raised); + }); + } + + #[rstest] + #[case::fraction(Some(1.5), Some(Duration::from_millis(1500)))] + #[case::negative_clamps_to_zero(Some(-1.0), Some(Duration::ZERO))] + #[case::absent(None, None)] + fn a_finite_ttl_converts(#[case] seconds: Option, #[case] expected: Option) { + assert_eq!(ttl(seconds).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/request.rs b/litellm-rust/crates/python-bridge/src/redis/store/request.rs new file mode 100644 index 00000000000..3436b26a16c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/request.rs @@ -0,0 +1,60 @@ +use litellm_redis_store::{RequestRedis, Send}; +use pyo3::prelude::*; + +use super::{call, duration}; +use crate::execution::run_async_value; + +/// One request's native Redis batches. Python builds one per request scope and passes it to each +/// request-aware store call. +#[pyclass(frozen, name = "NativeRequestRedis")] +pub(crate) struct NativeRequestRedis(pub(crate) RequestRedis); + +#[pymethods] +impl NativeRequestRedis { + #[new] + fn new(py: Python<'_>, post_call_deadline_seconds: f64) -> PyResult { + let runtime = litellm_host_python::runtime_handle()?; + let deadline = duration(post_call_deadline_seconds)?; + let deadline_tag = call(py, "post_call_redis_batch")?.tag(); + Ok(Self(crate::logger::capture(py).scope(|| { + RequestRedis::with_deadline_tag(runtime, deadline, Some(deadline_tag)) + }))) + } + + fn flush_pre<'py>(&self, py: Python<'py>) -> PyResult> { + let request = self.0.clone(); + let scope = call(py, "request_redis_batch")?; + run_async_value(py, async move { + scope.scope(request.flush_pre()).await; + Ok(()) + }) + } + + fn flush_post<'py>(&self, py: Python<'py>) -> PyResult> { + let request = self.0.clone(); + let scope = call(py, "post_call_redis_batch")?; + run_async_value(py, async move { + scope.scope(request.flush_post()).await; + Ok(()) + }) + } +} + +/// The Python face of `Send`: which batch a store op rides. Each variant carries exactly what it +/// needs, so there is no phase string and no phase without a request. +#[pyclass(frozen, name = "NativeSend")] +pub(crate) enum NativeSend { + Now(), + PreCall(Py), + PostCall(Py), +} + +impl NativeSend { + pub(crate) fn send(&self) -> Send<'_> { + match self { + Self::Now() => Send::Now, + Self::PreCall(request) => Send::PreCall(&request.get().0), + Self::PostCall(request) => Send::PostCall(&request.get().0), + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/value.rs b/litellm-rust/crates/python-bridge/src/redis/store/value.rs new file mode 100644 index 00000000000..c86c9e9bb20 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/value.rs @@ -0,0 +1,125 @@ +use litellm_python_compat::{BigInt, MAX_DEPTH, Value}; +use pyo3::{ + exceptions::{PyTypeError, PyValueError}, + prelude::*, + types::{ + PyBool, PyBytes, PyComplex, PyDict, PyFloat, PyFrozenSet, PyInt, PyList, PySet, PyString, + PyTuple, + }, +}; + +/// A plain Python value (what `json.loads` or `ast.literal_eval` returns) across the bridge. +/// Anything else, a live object included, is refused with `TypeError`. +pub(crate) struct PyValue(pub(crate) Value); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyValue { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + to_value(&value, 0).map(Self) + } +} + +impl<'py> IntoPyObject<'py> for PyValue { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + to_python(py, self.0) + } +} + +fn to_value(value: &Bound<'_, PyAny>, depth: usize) -> PyResult { + if depth > MAX_DEPTH { + return Err(PyValueError::new_err("value is nested too deeply")); + } + let nested = |items: Bound<'_, PyAny>| -> PyResult> { + items + .try_iter()? + .map(|item| to_value(&item?, depth + 1)) + .collect() + }; + if value.is_none() { + return Ok(Value::None); + } + if let Ok(flag) = value.cast::() { + return Ok(Value::Bool(flag.is_true())); + } + if value.is_instance_of::() { + return Ok(Value::Int(match value.extract::() { + Ok(small) => small.into(), + Err(_) => value + .str()? + .to_str()? + .parse::() + .map_err(|error| PyValueError::new_err(error.to_string()))?, + })); + } + if let Ok(number) = value.cast::() { + return Ok(Value::Float(number.value())); + } + if let Ok(text) = value.cast::() { + return Ok(Value::Str(text.to_str()?.to_owned())); + } + if let Ok(bytes) = value.cast::() { + return Ok(Value::Bytes(bytes.as_bytes().to_vec())); + } + if let Ok(complex) = value.cast::() { + return Ok(Value::Complex { + re: complex.real(), + im: complex.imag(), + }); + } + if let Ok(entries) = value.cast::() { + return entries + .iter() + .map(|(key, item)| Ok((to_value(&key, depth + 1)?, to_value(&item, depth + 1)?))) + .collect::>() + .map(Value::Dict); + } + if value.is_instance_of::() { + return nested(value.clone()).map(Value::List); + } + if value.is_instance_of::() { + return nested(value.clone()).map(Value::Tuple); + } + if value.is_instance_of::() || value.is_instance_of::() { + return nested(value.clone()).map(Value::Set); + } + Err(PyTypeError::new_err(format!( + "{} is not a plain data value", + value.get_type().name()? + ))) +} + +fn to_python(py: Python<'_>, value: Value) -> PyResult> { + let all = |items: Vec| -> PyResult>> { + items.into_iter().map(|item| to_python(py, item)).collect() + }; + Ok(match value { + Value::None => py.None().into_bound(py), + Value::Bool(flag) => PyBool::new(py, flag).to_owned().into_any(), + Value::Int(number) => match i64::try_from(&number) { + Ok(small) => small.into_pyobject(py)?.into_any(), + Err(_) => py + .get_type::() + .call1((number.to_string(),))? + .into_any(), + }, + Value::Float(number) => number.into_pyobject(py)?.into_any(), + Value::Complex { re, im } => PyComplex::from_doubles(py, re, im).into_any(), + Value::Str(text) => text.into_pyobject(py)?.into_any(), + Value::Bytes(bytes) => PyBytes::new(py, &bytes).into_any(), + Value::Tuple(items) => PyTuple::new(py, all(items)?)?.into_any(), + Value::List(items) => PyList::new(py, all(items)?)?.into_any(), + Value::Set(items) => PySet::new(py, all(items)?)?.into_any(), + Value::Dict(entries) => { + let dict = PyDict::new(py); + for (key, item) in entries { + dict.set_item(to_python(py, key)?, to_python(py, item)?)?; + } + dict.into_any() + } + }) +} diff --git a/litellm-rust/crates/python-compat/benches/formats.rs b/litellm-rust/crates/python-compat/benches/formats.rs index f0b83cb1658..2f23b882e19 100644 --- a/litellm-rust/crates/python-compat/benches/formats.rs +++ b/litellm-rust/crates/python-compat/benches/formats.rs @@ -10,7 +10,7 @@ use std::{hint::black_box, time::Duration}; use criterion::{ - BatchSize, BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main, + BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main, measurement::WallTime, }; use litellm_python_compat::{Value, json, literal::literal_eval, pickle, repr::repr}; @@ -72,14 +72,9 @@ fn formats(c: &mut Criterion) { let mut group = c.benchmark_group("cached_completion"); text_formats(&mut group, &text, &value); binary_formats(&mut group, &value, &pickled); - // `from_json` consumes its input, so each iteration gets a freshly parsed one. group.throughput(Throughput::Bytes(dumped.len() as u64)); - group.bench_function("from_json", |bencher| { - bencher.iter_batched( - || serde_json::from_str::(&dumped).expect("dumps output parses"), - json::from_json, - BatchSize::SmallInput, - ) + group.bench_function("json_loads", |bencher| { + bencher.iter(|| json::loads(black_box(&dumped))) }); group.finish(); } diff --git a/litellm-rust/crates/python-compat/src/error.rs b/litellm-rust/crates/python-compat/src/error.rs index 6547907056e..db09925dde9 100644 --- a/litellm-rust/crates/python-compat/src/error.rs +++ b/litellm-rust/crates/python-compat/src/error.rs @@ -6,6 +6,8 @@ use crate::MAX_DEPTH; pub enum Error { #[error("malformed Python literal at byte {0}")] InvalidLiteral(usize), + #[error("malformed JSON at byte {0}")] + InvalidJson(usize), #[error("unhashable type: '{0}'")] Unhashable(&'static str), #[error("value nests deeper than {MAX_DEPTH} levels")] diff --git a/litellm-rust/crates/python-compat/src/json.rs b/litellm-rust/crates/python-compat/src/json.rs index d3436519ac4..37cc88fec78 100644 --- a/litellm-rust/crates/python-compat/src/json.rs +++ b/litellm-rust/crates/python-compat/src/json.rs @@ -1,14 +1,19 @@ -//! `json.dumps` with CPython's default options, and the JSON value `json.loads` returns. +//! `json.dumps` with CPython's default options, and `json.loads` with its default decoder. //! //! Defaults are `ensure_ascii=True`, `allow_nan=True`, separators `(", ", ": ")`, and no key -//! sorting. Python's `json.loads` is mapped by [`from_json`]; a `serde_json` number that does -//! not fit `i64` or `u64` arrives as a float, where Python would keep an `int`. +//! sorting on the way out, and `strict=True` with `NaN`, `Infinity` and `-Infinity` read as floats +//! on the way in. -use std::fmt::Write; +use std::{ + collections::{HashMap, hash_map::Entry}, + fmt::Write, + str::FromStr, +}; +use num_bigint::BigInt; use serde_json::{Map, Number}; -use crate::{Error, Value, repr::float_repr}; +use crate::{Error, MAX_DEPTH, Value, repr::float_repr}; /// `json.dumps(value)`. pub fn dumps(value: &Value) -> Result { @@ -44,6 +49,260 @@ pub fn from_json(value: serde_json::Value) -> Value { } } +/// `json.loads(text)`. Integers keep full precision, and a repeated key keeps its first position +/// with its last value, as `dict(pairs)` does. An escape that leaves a lone surrogate fails, since +/// a Rust `String` cannot hold one. +pub fn loads(text: &str) -> Result { + let mut decoder = Decoder { + text, + bytes: text.as_bytes(), + pos: 0, + depth: 0, + }; + decoder.skip_whitespace(); + let value = decoder.value()?; + decoder.skip_whitespace(); + if decoder.pos != decoder.bytes.len() { + return Err(decoder.error()); + } + Ok(value) +} + +/// CPython's `json.scanner` and `scanstring`, one pass over the bytes. +struct Decoder<'a> { + text: &'a str, + bytes: &'a [u8], + pos: usize, + depth: usize, +} + +impl Decoder<'_> { + fn error(&self) -> Error { + Error::InvalidJson(self.pos) + } + + fn peek(&self) -> Option { + self.bytes.get(self.pos).copied() + } + + fn eat(&mut self, token: &str) -> bool { + let found = self.bytes[self.pos..].starts_with(token.as_bytes()); + if found { + self.pos += token.len(); + } + found + } + + fn expect(&mut self, token: &str) -> Result<(), Error> { + if self.eat(token) { + Ok(()) + } else { + Err(self.error()) + } + } + + fn skip_whitespace(&mut self) { + while matches!(self.peek(), Some(b' ' | b'\t' | b'\n' | b'\r')) { + self.pos += 1; + } + } + + fn value(&mut self) -> Result { + match self.peek() { + Some(b'"') => self.string().map(Value::Str), + Some(b'[') => self.nested(Self::array), + Some(b'{') => self.nested(Self::object), + _ if self.eat("null") => Ok(Value::None), + _ if self.eat("true") => Ok(Value::Bool(true)), + _ if self.eat("false") => Ok(Value::Bool(false)), + _ if self.eat("NaN") => Ok(Value::Float(f64::NAN)), + _ if self.eat("Infinity") => Ok(Value::Float(f64::INFINITY)), + _ if self.eat("-Infinity") => Ok(Value::Float(f64::NEG_INFINITY)), + _ => self.number(), + } + } + + fn nested(&mut self, parse: fn(&mut Self) -> Result) -> Result { + self.depth += 1; + if self.depth > MAX_DEPTH { + return Err(Error::TooDeep); + } + let value = parse(self)?; + self.depth -= 1; + Ok(value) + } + + fn array(&mut self) -> Result { + self.expect("[")?; + self.skip_whitespace(); + let mut values = Vec::new(); + if self.eat("]") { + return Ok(Value::List(values)); + } + loop { + values.push(self.value()?); + self.skip_whitespace(); + if self.eat("]") { + return Ok(Value::List(values)); + } + self.expect(",")?; + self.skip_whitespace(); + } + } + + fn object(&mut self) -> Result { + self.expect("{")?; + self.skip_whitespace(); + let mut entries = Vec::new(); + let mut positions = HashMap::<_, usize>::new(); + if self.eat("}") { + return Ok(Value::Dict(entries)); + } + loop { + if self.peek() != Some(b'"') { + return Err(self.error()); + } + let key = self.string()?; + self.skip_whitespace(); + self.expect(":")?; + self.skip_whitespace(); + let value = self.value()?; + match positions.entry(key) { + Entry::Occupied(position) => entries[*position.get()].1 = value, + Entry::Vacant(position) => { + entries.push((Value::Str(position.key().clone()), value)); + position.insert(entries.len() - 1); + } + } + self.skip_whitespace(); + if self.eat("}") { + return Ok(Value::Dict(entries)); + } + self.expect(",")?; + self.skip_whitespace(); + } + } + + /// `scanstring` with `strict=True`: control characters must be escaped. + fn string(&mut self) -> Result { + self.expect("\"")?; + let mut out = String::new(); + loop { + let start = self.pos; + while self + .peek() + .is_some_and(|byte| byte != b'"' && byte != b'\\' && byte >= 0x20) + { + self.pos += 1; + } + out.push_str(&self.text[start..self.pos]); + match self.peek() { + Some(b'"') => { + self.pos += 1; + return Ok(out); + } + Some(b'\\') => { + self.pos += 1; + out.push(self.escape()?); + } + _ => return Err(self.error()), + } + } + } + + fn escape(&mut self) -> Result { + let escaped = match self.peek() { + Some(b'"') => '"', + Some(b'\\') => '\\', + Some(b'/') => '/', + Some(b'b') => '\u{8}', + Some(b'f') => '\u{c}', + Some(b'n') => '\n', + Some(b'r') => '\r', + Some(b't') => '\t', + Some(b'u') => { + self.pos += 1; + return self.unicode_escape(); + } + _ => return Err(self.error()), + }; + self.pos += 1; + Ok(escaped) + } + + /// `\uXXXX`, joining a high surrogate with the low one escaped right after it. + fn unicode_escape(&mut self) -> Result { + let unit = self.hex4()?; + let code = if (0xd800..0xdc00).contains(&unit) { + self.expect("\\u")?; + let low = self.hex4()?; + if !(0xdc00..0xe000).contains(&low) { + return Err(self.error()); + } + 0x10000 + ((unit - 0xd800) << 10) + (low - 0xdc00) + } else { + unit + }; + char::from_u32(code).ok_or_else(|| self.error()) + } + + fn hex4(&mut self) -> Result { + let digits = self + .text + .get(self.pos..self.pos + 4) + .filter(|digits| digits.bytes().all(|byte| byte.is_ascii_hexdigit())) + .ok_or_else(|| self.error())?; + self.pos += 4; + u32::from_str_radix(digits, 16).map_err(|_| self.error()) + } + + /// `NUMBER_RE`: `(-?(?:0|[1-9][0-9]*))(\.[0-9]+)?([eE][-+]?[0-9]+)?`, an `int` unless a + /// fraction or exponent matched. + fn number(&mut self) -> Result { + let start = self.pos; + self.eat("-"); + match self.peek() { + Some(b'0') => self.pos += 1, + Some(b'1'..=b'9') => self.digits(), + _ => return Err(Error::InvalidJson(start)), + } + let integer_end = self.pos; + if self.peek() == Some(b'.') && self.bytes.get(self.pos + 1).is_some_and(u8::is_ascii_digit) + { + self.pos += 1; + self.digits(); + } + if matches!(self.peek(), Some(b'e' | b'E')) { + let exponent = self.pos; + self.pos += 1; + if matches!(self.peek(), Some(b'+' | b'-')) { + self.pos += 1; + } + if self.peek().is_some_and(|byte| byte.is_ascii_digit()) { + self.digits(); + } else { + self.pos = exponent; + } + } + let text = &self.text[start..self.pos]; + if self.pos == integer_end { + BigInt::from_str(text) + .map(Value::Int) + .map_err(|_| Error::InvalidJson(start)) + } else { + f64::from_str(text) + .map(Value::Float) + .map_err(|_| Error::InvalidJson(start)) + } + } + + fn digits(&mut self) { + while self.peek().is_some_and(|byte| byte.is_ascii_digit()) { + self.pos += 1; + } + } +} + /// `json.loads(json.dumps(value))` as a `serde_json` value: tuples become arrays and dict /// keys are coerced to strings as `json.dumps` does. Non-finite floats, which Python writes /// as `NaN` and `Infinity`, have no `serde_json` form and fail with [`Error::NonFiniteFloat`]. @@ -166,3 +425,86 @@ fn write_string(out: &mut String, text: &str) { } out.push('"'); } + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::{dumps, loads}; + use crate::{Error, Value, repr::repr}; + + // Expected values are the repr of CPython 3.14 json.loads(text) + #[rstest] + #[case::nan("NaN", "nan")] + #[case::infinity("Infinity", "inf")] + #[case::negative_infinity("-Infinity", "-inf")] + #[case::non_finite_in_a_list("[NaN, Infinity, -Infinity]", "[nan, inf, -inf]")] + #[case::beyond_u64("18446744073709551616", "18446744073709551616")] + #[case::exponent_overflow("1e400", "inf")] + #[case::upper_exponent("1E+2", "100.0")] + #[case::negative_zero_int("-0", "0")] + #[case::negative_zero_float("-0.0", "-0.0")] + #[case::repeated_key(r#"{"a": 1, "b": 2, "a": 3}"#, "{'a': 3, 'b': 2}")] + #[case::surrogate_pair(r#""😀""#, "'😀'")] + #[case::upper_hex_surrogate_pair(r#""😀""#, "'😀'")] + #[case::escapes( + r#"{"k": [true, false, null, 1.5, "é\/"]}"#, + "{'k': [True, False, None, 1.5, 'é/']}" + )] + #[case::delete_is_not_a_control_character("\"\x7f\"", "'\\x7f'")] + #[case::surrounding_whitespace(" \n[1]\t", "[1]")] + #[case::empty_containers("[{}, []]", "[{}, []]")] + fn loads_reads_like_python(#[case] text: &str, #[case] expected: &str) { + assert_eq!(repr(&loads(text).unwrap()), expected); + } + + // CPython 3.14 json.loads raises JSONDecodeError on each + #[rstest] + #[case::lowercase_nan("nan")] + #[case::signed_nan("-NaN")] + #[case::plus_sign("+1")] + #[case::bare_minus("-")] + #[case::leading_zero("01")] + #[case::bare_point("1.")] + #[case::bare_exponent("1e")] + #[case::trailing_data("Infinity1")] + #[case::trailing_comma("[1,]")] + #[case::missing_comma("[1 2]")] + #[case::missing_colon(r#"{"a" 1}"#)] + #[case::unquoted_key("{1: 2}")] + #[case::raw_newline_in_string("\"a\nb\"")] + #[case::unknown_escape(r#""\x41""#)] + #[case::empty("")] + #[case::byte_order_mark("\u{feff}1")] + fn loads_rejects_what_python_rejects(#[case] text: &str) { + assert!(matches!(loads(text), Err(Error::InvalidJson(_)))); + } + + #[rstest] + fn a_lone_surrogate_is_a_known_limit() { + assert!(matches!(loads(r#""\ud800""#), Err(Error::InvalidJson(_)))); + } + + #[rstest] + fn loads_reads_back_the_non_finite_floats_dumps_writes() { + let written = dumps(&Value::List(vec![ + Value::Float(f64::NAN), + Value::Float(f64::INFINITY), + Value::Float(f64::NEG_INFINITY), + ])) + .unwrap(); + + assert_eq!(repr(&loads(&written).unwrap()), "[nan, inf, -inf]"); + } + + #[rstest] + fn loads_rejects_nesting_past_max_depth() { + let text = format!( + "{}{}", + "[".repeat(crate::MAX_DEPTH + 1), + "]".repeat(crate::MAX_DEPTH + 1) + ); + + assert!(matches!(loads(&text), Err(Error::TooDeep))); + } +} diff --git a/litellm-rust/crates/python-compat/src/lib.rs b/litellm-rust/crates/python-compat/src/lib.rs index b23aff409f2..2cc0eeac397 100644 --- a/litellm-rust/crates/python-compat/src/lib.rs +++ b/litellm-rust/crates/python-compat/src/lib.rs @@ -7,9 +7,10 @@ //! |---|---| //! | [`literal`] | `ast.literal_eval(text)` | //! | [`repr`] | `repr(value)` and `str(value)` | -//! | [`json`] | `json.dumps(value)`, and `json.loads(json.dumps(value))` as a JSON value | +//! | [`json`] | `json.dumps(value)` and `json.loads(text)` | //! | [`pickle`] | `pickle.loads(data)` and `pickle.dumps(value)` for plain data | //! | [`truthy`] | `bool(value)` | +//! | [`number`] | `int(text)`, `float(text)`, and `int(seconds)` as a `Duration` | //! //! [`Value`] is the closed data model these formats share. Live Python objects //! (descriptors, `__bool__`, `__str__`, callbacks) are out of scope: those belong to the @@ -20,10 +21,12 @@ //! (string hashes are randomized per process). Sets keep their literal order. //! - `str` values are Rust `String`s, so lone surrogates cannot be represented. //! - [`pickle::loads`] decodes `tuple`, `set`, and `frozenset` as lists. +//! - [`number`] reads ASCII digits only, while Python's `int` and `float` accept any Unicode decimal digit. mod error; pub mod json; pub mod literal; +pub mod number; pub mod pickle; pub mod repr; pub mod truthy; diff --git a/litellm-rust/crates/python-compat/src/number.rs b/litellm-rust/crates/python-compat/src/number.rs new file mode 100644 index 00000000000..6e2f067ae25 --- /dev/null +++ b/litellm-rust/crates/python-compat/src/number.rs @@ -0,0 +1,121 @@ +use std::{str::FromStr, time::Duration}; + +use num_bigint::BigInt; + +/// `int(text)` in base 10, or `None` where Python raises `ValueError`. +pub fn int(text: &str) -> Option { + parse(text) +} + +/// `float(text)`, or `None` where Python raises `ValueError`. +pub fn float(text: &str) -> Option { + parse(text) +} + +/// Why [`seconds`] refused a value, as `int(value)` would. Messages are CPython's. +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum SecondsError { + /// `OverflowError`. Finite values past `Duration::MAX` land here too, where Python builds a big int. + #[error("cannot convert float infinity to integer")] + Overflow, + /// `ValueError`. + #[error("cannot convert float NaN to integer")] + NotANumber, +} + +/// A float of seconds as a `Duration`, refusing what `int(value)` refuses. Finite negative values clamp to zero. +pub fn seconds(value: f64) -> Result { + match value { + value if value.is_nan() => Err(SecondsError::NotANumber), + value if value.is_infinite() => Err(SecondsError::Overflow), + value => Duration::try_from_secs_f64(value.max(0.0)).map_err(|_| SecondsError::Overflow), + } +} + +fn parse(text: &str) -> Option { + let text = text.trim(); + let digit_at = |index: Option| { + index + .and_then(|index| text.as_bytes().get(index)) + .is_some_and(u8::is_ascii_digit) + }; + let underscores_between_digits = text + .match_indices('_') + .all(|(index, _)| digit_at(index.checked_sub(1)) && digit_at(Some(index + 1))); + if !underscores_between_digits { + return None; + } + text.replace('_', "").parse().ok() +} + +#[cfg(test)] +mod tests { + use num_bigint::BigInt; + use rstest::rstest; + + use std::time::Duration; + + use super::{SecondsError, float, int, seconds}; + + // Expected values are CPython 3.14 int(text) + #[rstest] + #[case::trailing_newline("2\n", Some(2))] + #[case::surrounding_spaces_and_underscore(" 6_380 ", Some(6380))] + #[case::plus_sign("+5", Some(5))] + #[case::negative_zero("-0", Some(0))] + #[case::leading_zeros("007", Some(7))] + #[case::unicode_spaces("\u{a0}7\u{2003}", Some(7))] + #[case::beyond_i64("99999999999999999999", Some(99_999_999_999_999_999_999_i128))] + #[case::double_underscore("1__0", None)] + #[case::leading_underscore("_1", None)] + #[case::trailing_underscore("1_", None)] + #[case::decimal_point("5.0", None)] + #[case::information_separator("\x1c7", None)] + #[case::empty("", None)] + #[case::sign_only("+", None)] + #[case::non_ascii_digit_is_a_known_limit("\u{667}", None)] + fn int_parses_like_python(#[case] text: &str, #[case] expected: Option) { + assert_eq!(int(text), expected.map(BigInt::from)); + } + + // Expected values are CPython 3.14 float(text) + #[rstest] + #[case::underscore_and_newline(" 1_0.5\n", Some(10.5))] + #[case::exponent("1e3", Some(1000.0))] + #[case::leading_point(".5", Some(0.5))] + #[case::trailing_point("5.", Some(5.0))] + #[case::plus_sign("+1.5", Some(1.5))] + #[case::infinity_word("Infinity", Some(f64::INFINITY))] + #[case::infinity_with_space("infinity ", Some(f64::INFINITY))] + #[case::underscore_before_exponent("1_e3", None)] + #[case::underscore_after_exponent("1e_3", None)] + #[case::underscore_after_point("1._5", None)] + #[case::underscore_after_word("inf_", None)] + #[case::hexadecimal("0x10", None)] + #[case::bare_exponent("1.5e", None)] + fn float_parses_like_python(#[case] text: &str, #[case] expected: Option) { + assert_eq!(float(text), expected); + } + + #[rstest] + fn float_reads_a_signed_nan_like_python() { + assert!(float("-NaN").is_some_and(f64::is_nan)); + } + + // CPython 3.14: int(inf) and int(-inf) raise OverflowError, int(nan) raises ValueError + #[rstest] + #[case::fraction(1.5, Ok(Duration::from_millis(1500)))] + #[case::zero(0.0, Ok(Duration::ZERO))] + #[case::negative_clamps_to_zero(-3.0, Ok(Duration::ZERO))] + #[case::infinity(f64::INFINITY, Err(SecondsError::Overflow))] + #[case::negative_infinity(f64::NEG_INFINITY, Err(SecondsError::Overflow))] + #[case::past_duration_max(1e20, Err(SecondsError::Overflow))] + #[case::nan(f64::NAN, Err(SecondsError::NotANumber))] + #[case::negative_nan(-f64::NAN, Err(SecondsError::NotANumber))] + fn seconds_refuses_what_int_refuses( + #[case] value: f64, + #[case] expected: Result, + ) { + assert_eq!(seconds(value), expected); + } +} diff --git a/litellm-rust/crates/python-compat/tests/fixtures.rs b/litellm-rust/crates/python-compat/tests/fixtures.rs index 9c202fdc5f1..3f9c6fd7010 100644 --- a/litellm-rust/crates/python-compat/tests/fixtures.rs +++ b/litellm-rust/crates/python-compat/tests/fixtures.rs @@ -76,6 +76,11 @@ const KNOWN: &[(&str, &str, &str)] = &[ "literal_eval", "deeper than MAX_DEPTH: rejected for stack safety, where CPython still parses it", ), + ( + "nested_150", + "json.loads", + "deeper than MAX_DEPTH: rejected for stack safety, where CPython still parses it", + ), ( "nested_150", "pickle.loads", @@ -341,3 +346,19 @@ fn to_json_matches_python_json_round_trip(fixtures: &Fixtures) { } mismatches.finish(); } + +/// CPython gives back the same text from `json.dumps(json.loads(text))` for every text its +/// `json.dumps` wrote here, non-finite floats included, so `loads` must too. +#[rstest] +fn loads_reads_back_what_python_json_dumps_wrote(fixtures: &Fixtures) { + let mut mismatches = Mismatches::default(); + for row in &fixtures.rows { + let Some(text) = &row.json else { continue }; + let actual = match json::loads(text).and_then(|value| json::dumps(&value)) { + Ok(text) => text, + Err(error) => format!("error: {error}"), + }; + mismatches.check(row, "json.loads", text, &actual); + } + mismatches.finish(); +} diff --git a/litellm-rust/crates/redis-client/AGENTS.md b/litellm-rust/crates/redis-client/AGENTS.md index c73a4ed8c4d..6135919c539 100644 --- a/litellm-rust/crates/redis-client/AGENTS.md +++ b/litellm-rust/crates/redis-client/AGENTS.md @@ -1,12 +1,12 @@ - The one async Redis client for the gateway and the Rust cache crates. It owns connecting, timeouts, error classes and the key namespace - It knows nothing about domains, so key builders, TTL policy, value formats and fallback decisions stay in the crate that owns the state. Every write takes its TTL from the caller, and there is no default - Async only. Build one `RedisClient` per Redis instance at the process boundary and clone it. Never add a pool, a sync API, `spawn_blocking`, or a client built inside a route -- Set every redis-rs timeout and retry explicitly. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's (5s socket timeout from `REDIS_SOCKET_TIMEOUT`) +- Every connection sets TCP_NODELAY, like redis-py's `socket.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1)`. Set every redis-rs timeout and retry explicitly, citing the redis-py 5.3.1 or `litellm/` source next to each constant. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's. `Timeouts` takes `socket_timeout` and `socket_connect_timeout`, and the connect timeout falls back to the socket timeout like redis-py. Standalone reconnects on the next command with no retry of the failed one, matching redis-py's `Retry(NoBackoff(), 0)`. redis-rs keeps a socket whose command timed out, so `Renewing` in `connection.rs` swaps in a fresh `ConnectionManager`, or a fresh cluster connection like redis-py's cluster `aclose`, after a timeout. Without it a socket that went silent times out every command until TCP gives up on it. Cluster keeps redis-py's 3 `cluster_error_retry_attempts` and waits a flat 250ms between them, which is redis-py's CLUSTERDOWN sleep. redis-py's jittered per-node reconnect backoff has no redis-rs equivalent. The retry set also differs, redis-rs retries LOADING and READONLY internally, while redis-py 5.3.1 raises them. Do not claim exact retry parity - Resolving config from `config.yaml` / `REDIS_*` env the way `litellm/_redis.py` does belongs to the consumer that reads that config, not here. `RedisConfig` takes already-resolved values ## Errors - `Transport`, `Server`, `Decode`, each wrapping the redis-rs error as `#[source]`. `ConnectError` is separate because only `connect` returns it -- Classify in `From for Error` by `ErrorKind`. LOADING, CLUSTERDOWN, MASTERDOWN, TRYAGAIN and READONLY are typed server replies, but they mean the node can't serve right now, so they are `Transport`. Codes redis-rs doesn't type (WRONGTYPE) come back as `ErrorKind::Extension` and are `Server`. Each code gets its own rstest case built from a parsed RESP error +- Classify with `Error::classify(error, deployment)`, never a bare `From`, because redis-py (5.3.1) decides by topology. `Transport` is what Python's `_redis_health_error_types` counts: LOADING everywhere, CLUSTERDOWN and MASTERDOWN only on cluster (only `ClusterParser` maps them), READONLY only under Sentinel (where `SentinelManagedConnection` turns it into `ConnectionError` and rediscovers), plus NOAUTH, WRONGPASS and the `ERR` messages redis-py maps to `ConnectionError`, matched exactly like `parse_error`. TRYAGAIN is `Server`, since the cluster client already retried it. Each code gets its own rstest case built from a parsed RESP error - A `Server` error never drops the connection - `Transport` includes timeouts, where the command may still have applied. Callers of counters and locks treat it as an unknown outcome - Never swallow an error or turn it into a miss. The caller picks the fallback @@ -17,27 +17,61 @@ - Every public call (each command, `eval`, one `run` of a pipeline) takes one admission through `guarded` / `admit` in `client.rs`. Never call the connection around it - `Transport` counts as a failure (`Timeout` when `is_timeout()`, else `Hard`). `Server` and `Decode` count as nothing, since Redis answered and the request was bad. A probe that gets one still reopens, through the dropped-probe rule in `circuit-breaker`. A pipeline counts one failure if any op came back `Transport`, else one success - An open breaker returns `Error::BreakerOpen` without touching the socket -- Python doesn't count MASTERDOWN, TRYAGAIN or READONLY as health failures. Here they are `Transport` and do count ## Namespace - Applied once, here, to every key. A key that already starts with `{namespace}:` is left alone, matching `RedisCache.check_and_fix_namespace` - The rule never changes, because Python and Rust pods share keys in the same Redis ## Commands -- One typed method per command, built with `redis::cmd`. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET +- A command is a value, one constructor in `commands.rs` returning `Command`: its alone and pipelined forms and its reply decode. `RedisClient::send`, `Pipeline::with` and `RedisBatch::add` are the only entry points, and each applies its namespace to the keys once. A new command is one constructor, never a method on an executor. `ping` and `subscribe` stay methods. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET +- A store attaches its value codec with `Command::map`, never a keyspace codec inside this crate - On cluster, redis-rs splits MGET and DEL by slot and returns MGET values in input order. Other multi-key commands are not split, so a new one needs a cross-slot test -- TTLs are `Duration` and are sent in milliseconds (`PX`, `PEXPIRE`, `PTTL`) +- Every command is the one `RedisCache` sends, argument for argument, so `MONITOR` on a mixed fleet shows one shape. A difference from main's Python is a defect even when the result is the same +- TTLs are `Duration` and are sent in whole seconds (`SET ... EX`, `EXPIRE`, `TTL`), like redis-py's `int(ttl)`. `set_nx` sends `EX` before `NX`, the order redis-py builds +- Numbers go out as `Amount`, written the way redis-py encodes them: an `int` with `str()` and a `float` with `repr()` from `python-compat` +- Increments are plain commands, never a script. `incr_by` is `increment_cache`: `INCRBY`, `TTL`, then `EXPIRE` only on a key with no TTL, each its own round trip. `incr_by_float` alone is `async_increment`: `INCRBYFLOAT` with `EXPIRE` (`Expiry::Refresh`) or `TTL` (`Expiry::KeepExisting`) in one round trip, and for `KeepExisting` a third, `EXPIRE`, only on a key with no TTL. Pipelined it is `RedisBatch`'s `_Increment`, `INCRBYFLOAT` then `EXPIRE` when an expiry is given, `KeepExisting` lowered to `Refresh`. `Plan` keeps only the cases whose alone and pipelined forms differ ## Scripts and pipelines - A `Script` declares every key it touches in KEYS. On cluster an undeclared key works or fails depending on which node owns its slot, so it can pass tests and break in production -- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline nothing reloads, so a NOSCRIPT op reruns alone with `invoke_async` +- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline a script is sent as `EVALSHA`, like `RedisBatch`. A `NOSCRIPT` reply means that op ran nothing, so it alone runs again after the pipeline through `invoke_async` (`EVALSHA`, `SCRIPT LOAD`, `EVALSHA`), like `RedisBatch`'s `_settle_alone`. It then lands after the ops around it, as it does in Python +- `Plan::Script` keeps a script call apart from plain commands, and the one resend takes only a `ScriptCall` whose reply was `NOSCRIPT`. Nothing can resend a pipeline or a command that already ran, so a retry can never apply an increment twice +- `RedisBatch::pipelined` is one `Pipeline` whose ops answer as typed `Declared`s. A store that sends its own pipeline uses it rather than decoding `run`'s raw `Value`s - A pipeline is never atomic. Every op gets its own `Result`, from `.ignore_errors()`, so one rejected op leaves the others alone. A transport failure fails every op of that round trip -- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key, runs the groups concurrently and returns results in input order -- Pipeline ops take one key each. A multi-key op would span slots inside one group -- Command builders live in `commands.rs` and are shared by `RedisClient` and `Pipeline`, so a command is written once +- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key and runs the slots concurrently, returning results in input order. Within a slot, runs of idempotent ops (GET, SET, DEL, EXPIRE) share a request and every script op gets its own, sent in order. redis-rs retries a cluster request whole, so a script op batched with others would be applied twice when only a neighbour failed +- Pipeline ops take one key each. A multi-key op would span slots inside one group. An op can send more than one command for that key (`INCRBYFLOAT` then `EXPIRE`), and settles from the first reply, failing if any of its replies failed +- Command constructors live in `commands.rs` and every executor sends them, so a command is written once -## Tests +## Pub/sub +- Classic `PUBLISH` / `SUBSCRIBE`, never sharded, since redis-py publishes and subscribes through one plain client and cluster broadcasts classic messages to every node. Payloads are raw bytes, and the codec belongs to the store that owns the channel +- Channel names never get the namespace. Python builds `{namespace}:{channel}` itself, so the store does too +- `publish` is a command like any other, guarded by the breaker and reported to events +- `subscribe` returns a `Stream` that opens its own RESP2 `PubSub` connection when first polled, like redis-py's `PubSub`, to the endpoint, the first reachable cluster node, or the master Sentinel names at that moment. It never touches the breaker or events, like Python's subscriber, which reads off its own connection +- After a lost connection or a failed subscribe it waits and resubscribes, 5s doubling to 60s and reset by a successful subscribe, like `AuthCacheInvalidationSubscriber._run`. Messages published while it is down are lost. The connection has no idle ping, so a socket that dies silently is only noticed by TCP +- redis-rs's `PubSub` takes no credentials provider, so a provider is asked for its current credentials once when the subscription connects, which is when redis-py's pubsub connection asks it. The connect is bounded by the connect timeout + +## Batches +- `RedisBatch` is litellm's `RedisBatch` (`litellm/caching/redis_batch.py`). It keeps each declared op as data and, on a standalone or Sentinel Redis, sends them later as one `Pipeline`, so it inherits every pipeline rule above. On a cluster it never pipelines and runs each op on its own, concurrently, the way `_settle_alone` calls the `RedisCache` method behind it. An increment there keeps a TTL the key already has, as `async_increment` does, where the pipelined one refreshes it like `_Increment` +- `mget` is one `MGET` of the keys in order with repeats read once, like `_MGet`, and answers one value per key it was given, in that order. Use it instead of one `get` per key whenever Python reads several keys in one call +- It knows nothing about requests. The host that owns a request (Axum handler, pyo3 adapter) builds one and passes it down as a value, never through a task-local +- `RedisBatch::alone` runs every op on its own like the cluster path, so a fresh alone batch is "send it now" and `Declared` is the one reply handle on every route +- Each op returns a `Declared` backed by a oneshot. `flush` sends everything declared so far, one flush at a time, and an empty flush makes no round trip. `resolve` flushes only when its op has not gone out yet, so an op declared after another caller resolved costs a second round trip. Stages that know their reads up front declare all of them and call `flush` once +- An op whose flush was cancelled resolves as `Transport`, because it may already have reached Redis +- `settled` waits for whichever flush carries the op and never sends one itself, for follow-ups that must not force a flush. `RedisClient::same_client` tells clones of one client apart from a second connection, so a host can keep one batch per client - Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes +- Tests that only need the public API live in `tests/.rs`, with client builders in `tests/support`. Inline `#[cfg(test)]` modules are only for private items (`error.rs`, `namespace.rs`) +- Observe connections from the server side (`RedisDeployment::clients`, `cli_on_every_node`), never by reaching into `RedisClient`'s connection - Every command test runs on each topology through `#[values(Topology::Standalone, Topology::Cluster)]`. A test is single-topology only when the behavior exists on one (a `db` number has no cluster meaning) -- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables +- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables. A socket that never answers again comes from `TcpProxy::blackhole_open_connections`, since `thaw` revives the same socket - Use a short client timeout in outage tests so they stay fast. Never sleep to wait for a state; poll with a deadline + +## TLS, credentials and Sentinel +- `TlsOptions` follows redis-py 5.3.1. It verifies the certificate by default, skips the hostname check by default (`ssl_check_hostname=False`), and trusts the system roots plus `ca_certs` +- rustls only reads unencrypted keys, so a PKCS#8 `ENCRYPTED PRIVATE KEY` is decrypted with `key_password` before it reaches redis-rs. A legacy `Proc-Type: 4,ENCRYPTED` PEM key is refused with `KeyError::LegacyEncryption` rather than decrypted by hand-rolled OpenSSL key derivation +- redis-rs loads the system trust store for every TLS connection and fails if that fails, even with `ca_certs` set. Where the macOS keychain is blocked (agent sandboxes), run the TLS tests with `SSL_CERT_FILE=/etc/ssl/cert.pem` +- `Credentials::Static` are sent at connect, and URL userinfo wins over them, like redis-py's `from_url` +- A url's `?protocol=` is ignored and every connection speaks RESP2, which every reply decode assumes +- `Credentials::Provider` wraps a redis-rs `StreamingCredentialsProvider`. redis-rs takes the first item at connect and re-`AUTH`s live connections on every later item, on each cluster node. Providers push, redis-rs never polls. With a provider, redis-rs overwrites URL userinfo at connect, which matches Python stripping `user:pass@` from the URL +- Sentinel discovers the master on each new connection and drops it on a transport failure or a READONLY reply, to one command or inside a pipeline, like redis-py's `SentinelConnectionPool` +- redis-rs sends `ROLE` to the master a Sentinel names before using it, which redis-py never does. That probe gets the data credentials (a provider's current pair) and TLS, so the Sentinel client is built on each discovery. redis-rs gives the probe no way to skip the hostname check, and drops the client certificate when the certificate is not verified, so a TLS master whose certificate does not name the announced address, or that requires a client certificate while verification is off, fails discovery where Python connects +- Sentinel monitors get `sentinel_password`, the data username when that password is set, and the TLS mode and certificates. Python's `_sentinel_auth_kwargs` forwards every other connection kwarg too, and it sends the username even without a sentinel password +- A `Debug` impl never shows a secret. Follow `.agents/skills/rust-redaction` diff --git a/litellm-rust/crates/redis-client/Cargo.toml b/litellm-rust/crates/redis-client/Cargo.toml index 3e9447a6bfa..d3d3226abc6 100644 --- a/litellm-rust/crates/redis-client/Cargo.toml +++ b/litellm-rust/crates/redis-client/Cargo.toml @@ -6,13 +6,28 @@ license.workspace = true repository.workspace = true [dependencies] +tracing.workspace = true litellm-circuit-breaker.workspace = true litellm-secret-value.workspace = true +litellm-python-compat.workspace = true +litellm-redis-types.workspace = true futures-util.workspace = true -redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] } +redis = { workspace = true, features = [ + "cluster-async", + "connection-manager", + "sentinel", + "tls-rustls-insecure", + "token-based-authentication", + "tokio-rustls-comp", +] } +pem = "4.0.0" +pkcs8 = { version = "0.10.2", features = ["3des", "encryption", "std"] } +rustls-native-certs.workspace = true +socket2 = "0.6.5" thiserror.workspace = true +tokio = { workspace = true, features = ["rt", "sync", "time"] } [dev-dependencies] litellm-redis-testing.workspace = true +rustls.workspace = true rstest.workspace = true -tokio.workspace = true diff --git a/litellm-rust/crates/redis-client/src/batch.rs b/litellm-rust/crates/redis-client/src/batch.rs new file mode 100644 index 00000000000..84fc1643424 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/batch.rs @@ -0,0 +1,147 @@ +use std::{ + io, + sync::{Arc, Mutex, MutexGuard}, +}; + +use futures_util::future::join_all; +use redis::{RedisError, Value}; +use tokio::sync::oneshot::{self, error::TryRecvError}; + +use crate::{Command, Error, Pipeline, RedisClient, commands::Op}; + +/// Ops declared by many callers and sent together, like litellm's `RedisBatch`: one pipeline on a +/// standalone or Sentinel Redis, and each op on its own on a cluster or in an alone batch. +#[derive(Clone)] +pub struct RedisBatch(Arc); + +#[derive(Clone, Copy)] +enum Mode { + /// litellm's `RedisBatch`: one pipeline, or each op on its own on a cluster. + Batch, + /// Each op on its own, the way the `RedisCache` method behind it sends it. + Alone, + /// One redis-py pipeline on every topology, split by slot on a cluster. + Pipeline, +} + +struct Shared { + client: RedisClient, + mode: Mode, + pending: Mutex, + flushing: tokio::sync::Mutex<()>, +} + +#[derive(Default)] +struct Pending { + ops: Vec, + replies: Vec>>, +} + +type Decode = Box Result + Send>; + +#[must_use] +pub struct Declared { + batch: RedisBatch, + reply: oneshot::Receiver>, + decode: Decode, +} + +impl RedisBatch { + pub(crate) fn new(client: RedisClient) -> Self { + Self::build(client, Mode::Batch) + } + + /// A batch whose every op goes out on its own, the way the `RedisCache` method behind it sends + /// it, so a fresh one is "send it now": the first `resolve` sends it. + pub fn alone(client: RedisClient) -> Self { + Self::build(client, Mode::Alone) + } + + /// A batch sent as one redis-py pipeline, like `async_increment_pipeline`, on every topology. + pub fn pipelined(client: RedisClient) -> Self { + Self::build(client, Mode::Pipeline) + } + + fn build(client: RedisClient, mode: Mode) -> Self { + Self(Arc::new(Shared { + client, + mode, + pending: Mutex::new(Pending::default()), + flushing: tokio::sync::Mutex::new(()), + })) + } + + pub fn add(&self, command: Command) -> Declared { + let (op, decode) = command.prepare(self.0.client.namespace()); + let (sender, reply) = oneshot::channel(); + let mut pending = self.pending(); + pending.ops.push(op); + pending.replies.push(sender); + Declared { + batch: self.clone(), + reply, + decode, + } + } + + pub async fn flush(&self) { + let _flushing = self.0.flushing.lock().await; + let Pending { ops, replies } = std::mem::take(&mut *self.pending()); + if replies.is_empty() { + return; + } + let client = &self.0.client; + let alone = match self.0.mode { + Mode::Alone => true, + Mode::Batch => client.is_cluster(), + Mode::Pipeline => false, + }; + let results = match alone { + true => join_all(ops.into_iter().map(|op| client.run_alone(op))).await, + false => { + client + .run(&Pipeline::of(client.namespace().clone(), ops)) + .await + } + }; + for (reply, result) in replies.into_iter().zip(results) { + let _ = reply.send(result); + } + } + + fn pending(&self) -> MutexGuard<'_, Pending> { + self.0 + .pending + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } +} + +impl Declared { + pub async fn resolve(mut self) -> Result { + match self.reply.try_recv() { + Ok(reply) => return reply.and_then(self.decode), + Err(TryRecvError::Closed) => return Err(abandoned()), + Err(TryRecvError::Empty) => {} + } + self.batch.flush().await; + self.settled().await + } + + pub async fn settled(self) -> Result { + let reply = self.reply.await.unwrap_or_else(|_| Err(abandoned())); + reply.and_then(self.decode) + } + + /// Every reply in order, flushing their batch once if nobody has yet. + pub async fn all(declared: impl IntoIterator) -> Vec> { + join_all(declared.into_iter().map(Self::resolve)).await + } +} + +fn abandoned() -> Error { + Error::Transport(RedisError::from(io::Error::new( + io::ErrorKind::Interrupted, + "the batch flush carrying this op was cancelled", + ))) +} diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index 48dbc2b2dfb..3cf468e0d3f 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -1,21 +1,43 @@ use std::{ future::Future, + sync::{Arc, Mutex, PoisonError, Weak}, time::{Duration, Instant}, }; -use litellm_circuit_breaker::{Admission, Breaker, Failure}; -use redis::{Cmd, FromRedisValue, IntoConnectionInfo, ToRedisArgs, ToSingleRedisArg, Value}; +use futures_util::stream::BoxStream; +use litellm_circuit_breaker::{Admission, Breaker, Failure, State}; +use redis::{Cmd, RedisError, Value}; use crate::{ - ConnectError, Endpoint, Error, Namespace, Pipeline, RedisConfig, Script, commands, + CallEvent, Command, ConnectError, Error, FailureKind, Namespace, Operation, Outcome, Pipeline, + RedisBatch, RedisConfig, RedisEvents, + commands::{self, Op, Plan, Trips}, connection::Connection, + pubsub::Subscriber, }; +/// One Redis client. Clones share the connection, the breaker and the subscriber. #[derive(Clone)] -pub struct RedisClient { - pub(crate) connection: Connection, +pub struct RedisClient(Arc); + +/// A handle that does not keep the client alive, for a registry that must not outlive its users. +#[derive(Clone)] +pub struct WeakRedisClient(Weak); + +impl WeakRedisClient { + pub fn upgrade(&self) -> Option { + self.0.upgrade().map(RedisClient) + } +} + +struct Shared { + connection: Connection, + subscriber: Subscriber, namespace: Namespace, breaker: Option, + events: Option>, + health_check: Option, + last_used: Mutex, } #[derive(Debug, PartialEq, Eq)] @@ -27,122 +49,157 @@ pub enum Ttl { impl RedisClient { pub fn new(config: RedisConfig) -> Result { - let connection = match &config.endpoint { - Endpoint::Url(url) => { - Connection::standalone(url.expose().into_connection_info()?, config.timeout)? - } - Endpoint::Host(host) => { - Connection::standalone(host.connection_info()?, config.timeout)? - } - Endpoint::Cluster(cluster) => Connection::cluster(cluster, config.timeout)?, - }; - Ok(Self { + let connection = Connection::new( + &config.target, + &config.credentials, + config.timeouts, + config.keepalive, + )?; + Ok(Self(Arc::new(Shared { + subscriber: Subscriber::new( + &config.target, + &config.credentials, + config.timeouts, + config.keepalive, + &connection, + )?, connection, namespace: config.namespace, breaker: config.breaker, - }) + events: config.events, + health_check: config.health_check, + last_used: Mutex::new(Instant::now()), + }))) + } + + pub fn downgrade(&self) -> WeakRedisClient { + WeakRedisClient(Arc::downgrade(&self.0)) + } + + /// The breaker's current state, or none when the breaker is disabled. + pub fn breaker_state(&self) -> Option { + self.0.breaker.as_ref().map(Breaker::state) } pub async fn ping(&self) -> Result<(), Error> { - self.query(&commands::ping()).await - } - - pub async fn get(&self, key: &str) -> Result, Error> { - self.query(&commands::get(self.key(key))).await - } - - pub async fn mget(&self, keys: &[&str]) -> Result>, Error> { - if keys.is_empty() { - return Ok(Vec::new()); - } - self.query(&commands::mget(self.keys(keys))).await - } - - pub async fn set( - &self, - key: &str, - value: impl ToSingleRedisArg, - ttl: Duration, - ) -> Result<(), Error> { - self.query(&commands::set(self.key(key), value, ttl)).await - } - - pub async fn set_nx( - &self, - key: &str, - value: impl ToSingleRedisArg, - ttl: Duration, - ) -> Result { - let reply: Option = self - .query(&commands::set_nx(self.key(key), value, ttl)) - .await?; - Ok(reply.is_some()) - } - - pub async fn del(&self, keys: &[&str]) -> Result { - if keys.is_empty() { - return Ok(0); - } - self.query(&commands::del(self.keys(keys))).await - } - - pub async fn expire(&self, key: &str, ttl: Duration) -> Result { - self.query(&commands::expire(self.key(key), ttl)).await - } - - pub async fn ttl(&self, key: &str) -> Result { - let reply: i64 = self.query(&commands::ttl(self.key(key))).await?; - Ok(match reply { - -2 => Ttl::Missing, - -1 => Ttl::Persistent, - millis => Ttl::Expires(Duration::from_millis(millis.unsigned_abs())), - }) - } - - pub async fn incr_by(&self, key: &str, delta: i64) -> Result { - self.query(&commands::incr_by(self.key(key), delta)).await - } - - pub async fn incr_by_float(&self, key: &str, delta: f64) -> Result { - self.query(&commands::incr_by_float(self.key(key), delta)) - .await - } - - pub async fn rpush(&self, key: &str, values: impl ToRedisArgs) -> Result { - self.query(&commands::rpush(self.key(key), values)).await - } - - pub async fn eval( - &self, - script: &Script, - keys: &[&str], - args: impl ToRedisArgs, - ) -> Result { - let keys = self.keys(keys); - let args = args.to_redis_args(); - self.guarded(async { - Ok(script - .invocation(&keys, &args) - .invoke_async(&mut self.connection.clone()) - .await?) + self.guarded(Operation::Command("PING"), async { + commands::ping() + .query_async(&mut self.0.connection.clone()) + .await }) .await } + /// Sends one command on its own, the way the `RedisCache` method behind it sends it. + pub async fn send(&self, command: Command) -> Result { + let (op, decode) = command.prepare(&self.0.namespace); + self.run_alone(op).await.and_then(decode) + } + + pub(crate) async fn run_alone(&self, op: Op) -> Result { + let Op { + operation, plan, .. + } = op; + self.guarded(operation, async { + let mut connection = self.0.connection.clone(); + match plan { + Plan::Reply(value) => Ok(value), + Plan::Commands(commands) => match <[Cmd; 1]>::try_from(commands) { + Ok([command]) => command.query_async(&mut connection).await, + Err(commands) => { + let mut pipeline = redis::pipe(); + for command in commands { + pipeline.add_command(command); + } + let replies: Vec = pipeline.query_async(&mut connection).await?; + Ok(replies.into_iter().next().unwrap_or(Value::Nil)) + } + }, + Plan::Script(call) => { + call.script + .invocation(&call.keys, &call.args) + .invoke_async(&mut connection) + .await + } + Plan::KeepingTtl { + increment, + key, + ttl, + trips, + } => { + let (value, current): (Value, i64) = match trips { + Trips::Shared => { + redis::pipe() + .add_command(increment) + .add_command(commands::ttl_cmd(&key)) + .query_async(&mut connection) + .await? + } + Trips::Separate => ( + increment.query_async(&mut connection).await?, + commands::ttl_cmd(&key).query_async(&mut connection).await?, + ), + }; + if current == -1 { + redis::cmd("EXPIRE") + .arg(&key) + .arg(ttl.as_secs()) + .exec_async(&mut connection) + .await?; + } + Ok(value) + } + } + }) + .await + } + + /// Messages published on `channel`, on a connection of its own. After the connection is lost + /// it reconnects and resubscribes, and messages published meanwhile are not delivered. + pub fn subscribe(&self, channel: &str) -> BoxStream<'static, Vec> { + self.0.subscriber.clone().messages(channel.to_owned()) + } + pub fn pipeline(&self) -> Pipeline { - Pipeline::new(self.namespace.clone()) + Pipeline::new(self.0.namespace.clone()) + } + + pub(crate) fn namespace(&self) -> &Namespace { + &self.0.namespace + } + + pub fn is_cluster(&self) -> bool { + matches!(self.0.connection, Connection::Cluster(_)) + } + + pub fn same_client(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } + + pub fn batch(&self) -> RedisBatch { + RedisBatch::new(self.clone()) } pub async fn run(&self, pipeline: &Pipeline) -> Vec> { + let operation = Operation::Pipeline { + ops: pipeline.len(), + }; + let started = Instant::now(); let admission = match self.admit() { Ok(admission) => admission, Err(_) => { + self.observe( + operation, + started, + Outcome::Failed(FailureKind::BreakerOpen), + ); return (0..pipeline.len()) .map(|_| Err(Error::BreakerOpen)) .collect(); } }; - let results = pipeline.run(&self.connection).await; + self.check_health().await; + let results = pipeline.run(&self.0.connection).await; let transport = results.iter().find_map(|result| { result .as_ref() @@ -150,31 +207,76 @@ impl RedisClient { .filter(|error| matches!(error, Error::Transport(_))) }); settle(admission, transport); + let failure = transport.or_else(|| results.iter().find_map(|result| result.as_ref().err())); + let outcome = failure.map_or(Outcome::Succeeded, |error| { + Outcome::Failed(FailureKind::of(error)) + }); + self.observe(operation, started, outcome); results } - fn key(&self, key: &str) -> String { - self.namespace.apply(key) - } - - fn keys(&self, keys: &[&str]) -> Vec { - keys.iter().map(|key| self.key(key)).collect() - } - - async fn query(&self, command: &Cmd) -> Result { - self.guarded(async { Ok(command.query_async(&mut self.connection.clone()).await?) }) - .await - } - - async fn guarded(&self, call: impl Future>) -> Result { - let admission = self.admit()?; - let result = call.await; - settle(admission, result.as_ref().err()); + async fn guarded( + &self, + operation: Operation, + call: impl Future>, + ) -> Result { + let started = Instant::now(); + let result = match self.admit() { + Ok(admission) => { + self.check_health().await; + let result = call + .await + .map_err(|error| Error::classify(error, self.0.connection.deployment())); + settle(admission, result.as_ref().err()); + result + } + Err(error) => Err(error), + }; + self.observe(operation, started, Outcome::of(&result)); result } + fn observe(&self, operation: Operation, started: Instant, outcome: Outcome) { + *self + .0 + .last_used + .lock() + .unwrap_or_else(PoisonError::into_inner) = Instant::now(); + if let Some(events) = &self.0.events { + events.call(&CallEvent { + operation, + duration: started.elapsed(), + outcome, + }); + } + } + + /// redis-py pings a connection that sat idle longer than `health_check_interval` before + /// sending the next command on it. A ping that times out or finds the socket dropped leaves + /// that connection behind, as a command would, so the command then goes out on a fresh one, + /// where redis-py 5.3.1 disconnects and fails the command with the ping's error. The ping is + /// not a call of the caller's and never reaches the breaker or the events. + async fn check_health(&self) { + let Some(interval) = self.0.health_check else { + return; + }; + let idle = self + .0 + .last_used + .lock() + .unwrap_or_else(PoisonError::into_inner) + .elapsed(); + if idle < interval { + return; + } + if let Err(error) = self.0.connection.ping_every_node().await { + tracing::debug!(%error, "idle Redis connection failed its health check, replacing it"); + } + } + fn admit(&self) -> Result, Error> { - self.breaker + self.0 + .breaker .as_ref() .map(|breaker| { breaker @@ -198,455 +300,3 @@ fn settle(admission: Option, error: Option<&Error>) { Some(Error::Server(_) | Error::Decode(_) | Error::BreakerOpen) => drop(admission), } } - -#[cfg(test)] -mod tests { - use std::time::{Duration, Instant}; - - use litellm_redis_testing::RedisServer; - use litellm_secret_value::SecretValue; - use rstest::rstest; - - use litellm_circuit_breaker::{Breaker, Config, State}; - - use super::{RedisClient, Ttl}; - use crate::{ - DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace, RedisConfig, - test_support::{Server, Topology, client_id}, - }; - - #[rstest] - #[tokio::test] - async fn namespaces_isolate_keys_on_the_server( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let a = server.client("a", DEFAULT_TIMEOUT).await; - let b = server.client("b", DEFAULT_TIMEOUT).await; - let bare = server.client("", DEFAULT_TIMEOUT).await; - - a.set("x", "1", Duration::from_secs(60)).await.unwrap(); - - assert_eq!(a.get::("x").await.unwrap().as_deref(), Some("1")); - assert_eq!(a.get::("a:x").await.unwrap().as_deref(), Some("1")); - assert_eq!( - bare.get::("a:x").await.unwrap().as_deref(), - Some("1") - ); - assert_eq!(b.get::("x").await.unwrap(), None); - } - - #[rstest] - #[tokio::test] - async fn a_rejected_command_keeps_the_connection( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client.rpush("list", "item").await.unwrap(); - let before = client_id(&client, "list").await; - - let rejected = client.incr_by("list", 1).await; - - assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}"); - assert_eq!(client.rpush("list", "next").await.unwrap(), 2); - assert_eq!(client_id(&client, "list").await, before); - } - - #[rstest] - #[tokio::test] - async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client - .set("float", "15.0", Duration::from_secs(60)) - .await - .unwrap(); - - let decoded = client.get::("float").await; - - assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}"); - assert_eq!(client.get::("missing").await.unwrap(), None); - } - - #[rstest] - #[tokio::test] - async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_millis(500)).await; - client - .set("first", "1", Duration::from_secs(60)) - .await - .unwrap(); - client - .set("second", "2", Duration::from_secs(60)) - .await - .unwrap(); - - server.freeze().await; - let timed_out = client.get::("first").await; - server.thaw().await; - - assert!( - matches!(timed_out, Err(Error::Transport(_))), - "{timed_out:?}" - ); - assert_eq!( - client.get::("second").await.unwrap().as_deref(), - Some("2") - ); - } - - #[rstest] - #[tokio::test] - async fn the_same_client_recovers_after_redis_restarts( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_millis(500)).await; - - server.restart().await; - - let deadline = Instant::now() + Duration::from_secs(30); - let recovered = loop { - match client.incr_by("after_restart", 1).await { - Ok(value) => break Some(value), - Err(_) if Instant::now() < deadline => { - tokio::time::sleep(Duration::from_millis(200)).await - } - Err(_) => break None, - } - }; - assert_eq!(recovered, Some(1)); - } - - #[rstest] - #[tokio::test] - async fn a_client_created_while_redis_is_down_works_once_redis_is_up( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - server.stop().await; - - let client = server.client("", Duration::from_millis(500)).await; - let while_down = client.ping().await; - server.resume().await; - - assert!( - matches!(while_down, Err(Error::Transport(_))), - "{while_down:?}" - ); - let deadline = Instant::now() + Duration::from_secs(30); - let recovered = loop { - match client.ping().await { - Ok(()) => break true, - Err(_) if Instant::now() < deadline => { - tokio::time::sleep(Duration::from_millis(200)).await - } - Err(_) => break false, - } - }; - assert!(recovered); - } - - #[rstest] - #[tokio::test] - async fn empty_batches_succeed_without_asking_redis( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - assert_eq!( - client.mget::(&[]).await.unwrap(), - Vec::>::new() - ); - assert_eq!(client.del(&[]).await.unwrap(), 0); - } - - #[rstest] - #[tokio::test] - async fn mget_and_del_span_slots_in_input_order( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client.set("a", "1", Duration::from_secs(60)).await.unwrap(); - client.set("c", "3", Duration::from_secs(60)).await.unwrap(); - - let values = client.mget::(&["c", "b", "a"]).await.unwrap(); - let deleted = client.del(&["a", "b", "c"]).await.unwrap(); - - assert_eq!( - values, - vec![Some("3".to_owned()), None, Some("1".to_owned())] - ); - assert_eq!(deleted, 2); - assert_eq!( - client.mget::(&["a", "c"]).await.unwrap(), - vec![None, None] - ); - } - - #[rstest] - #[tokio::test] - async fn set_nx_only_writes_a_missing_key( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - let first = client - .set_nx("lock", "owner-1", Duration::from_secs(60)) - .await - .unwrap(); - let second = client - .set_nx("lock", "owner-2", Duration::from_secs(60)) - .await - .unwrap(); - - assert!(first); - assert!(!second); - assert_eq!( - client.get::("lock").await.unwrap().as_deref(), - Some("owner-1") - ); - } - - #[rstest] - #[tokio::test] - async fn ttl_reports_missing_persistent_and_remaining_time( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client.incr_by("persistent", 1).await.unwrap(); - client - .set("expiring", "1", Duration::from_secs(60)) - .await - .unwrap(); - client.incr_by_float("expired_later", 0.5).await.unwrap(); - assert!( - client - .expire("expired_later", Duration::from_secs(30)) - .await - .unwrap() - ); - - assert_eq!(client.ttl("missing").await.unwrap(), Ttl::Missing); - assert_eq!(client.ttl("persistent").await.unwrap(), Ttl::Persistent); - assert!(matches!( - client.ttl("expiring").await.unwrap(), - Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) - )); - assert!(matches!( - client.ttl("expired_later").await.unwrap(), - Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30) - )); - assert!( - !client - .expire("missing", Duration::from_secs(30)) - .await - .unwrap() - ); - } - - #[rstest] - #[tokio::test] - async fn incr_by_float_accumulates_fractions( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - client.incr_by_float("spend", 0.25).await.unwrap(); - - assert_eq!(client.incr_by_float("spend", 0.5).await.unwrap(), 0.75); - } - - #[rstest] - #[tokio::test] - async fn concurrent_commands_on_one_cloned_client_all_apply( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - - let tasks: Vec<_> = (0..100) - .map(|_| { - let client = client.clone(); - tokio::spawn(async move { - for _ in 0..10 { - client.incr_by("shared", 1).await.unwrap(); - } - }) - }) - .collect(); - for task in tasks { - task.await.unwrap(); - } - - assert_eq!(client.get::("shared").await.unwrap(), Some(1000)); - } - - fn breaker(failure_threshold: u32) -> Breaker { - Breaker::new( - Config { - failure_threshold, - recovery_timeout: Duration::from_millis(200), - timeout_min_duration: Duration::ZERO, - }, - Box::new(|_| {}), - ) - } - - #[rstest] - #[tokio::test] - async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let breaker = breaker(2); - let client = server - .guarded_client("", Duration::from_millis(300), Some(breaker.clone())) - .await; - - server.freeze().await; - let first = client.get::("k").await; - let second = client.get::("k").await; - let started = Instant::now(); - let third = client.get::("k").await; - let rejected_in = started.elapsed(); - server.thaw().await; - - assert!(matches!(first, Err(Error::Transport(_))), "{first:?}"); - assert!(matches!(second, Err(Error::Transport(_))), "{second:?}"); - assert!(matches!(third, Err(Error::BreakerOpen)), "{third:?}"); - assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); - - let deadline = Instant::now() + Duration::from_secs(30); - let recovered = loop { - match client.get::("k").await { - Ok(value) => break Some(value), - Err(_) if Instant::now() < deadline => { - tokio::time::sleep(Duration::from_millis(50)).await - } - Err(_) => break None, - } - }; - assert_eq!(recovered, Some(None)); - assert!( - matches!(breaker.state(), State::Closed { .. }), - "{:?}", - breaker.state() - ); - } - - #[rstest] - #[tokio::test] - async fn rejected_commands_never_open_the_breaker( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let breaker = breaker(2); - let client = server - .guarded_client("", DEFAULT_TIMEOUT, Some(breaker.clone())) - .await; - client.rpush("list", "item").await.unwrap(); - - let rejected = [ - client.incr_by("list", 1).await, - client.incr_by("list", 1).await, - client.incr_by("list", 1).await, - ]; - - assert!( - matches!(rejected[2], Err(Error::Server(_))), - "{:?}", - rejected[2] - ); - assert!( - matches!(breaker.state(), State::Closed { hard: 0, .. }), - "{:?}", - breaker.state() - ); - } - - #[rstest] - #[tokio::test] - async fn an_open_breaker_fails_every_pipeline_op_without_sending( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server - .guarded_client("", Duration::from_millis(300), Some(breaker(1))) - .await; - let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1); - - server.freeze().await; - let failed = client.run(&pipeline).await; - let rejected = client.run(&pipeline).await; - server.thaw().await; - - assert!( - matches!(failed[0], Err(Error::Transport(_))), - "{:?}", - failed[0] - ); - assert!( - matches!(rejected[0], Err(Error::BreakerOpen)), - "{:?}", - rejected[0] - ); - assert!( - matches!(rejected[1], Err(Error::BreakerOpen)), - "{:?}", - rejected[1] - ); - } - - #[rstest] - #[tokio::test] - async fn a_host_endpoint_connects_to_its_database() { - let server = RedisServer::start().await.unwrap(); - let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); - let connect = |endpoint| { - RedisClient::new(RedisConfig { - endpoint, - namespace: Namespace::default(), - timeout: DEFAULT_TIMEOUT, - breaker: None, - }) - }; - let in_db_3 = connect(Endpoint::Host(HostEndpoint { - host: "127.0.0.1".to_owned(), - port, - username: None, - password: None, - tls: false, - db: 3, - })) - .unwrap(); - let in_db_0 = connect(Endpoint::Url(SecretValue::new(server.url().to_owned()))).unwrap(); - let via_url = connect(Endpoint::Url(SecretValue::new(format!( - "{}/3", - server.url() - )))) - .unwrap(); - - in_db_3 - .set("k", "3", Duration::from_secs(60)) - .await - .unwrap(); - - assert_eq!(in_db_0.get::("k").await.unwrap(), None); - assert_eq!( - via_url.get::("k").await.unwrap().as_deref(), - Some("3") - ); - } -} diff --git a/litellm-rust/crates/redis-client/src/commands.rs b/litellm-rust/crates/redis-client/src/commands.rs index fb8a4a28c94..5f8c1eff3f7 100644 --- a/litellm-rust/crates/redis-client/src/commands.rs +++ b/litellm-rust/crates/redis-client/src/commands.rs @@ -1,62 +1,407 @@ use std::time::Duration; -use redis::{Cmd, ToRedisArgs, ToSingleRedisArg}; +use litellm_python_compat::repr::float_repr; +use redis::{ + Cmd, ErrorKind, FromRedisValue, RedisError, RedisWrite, ToRedisArgs, ToSingleRedisArg, Value, +}; + +use crate::{Error, Namespace, Operation, Script, Ttl, script::ScriptCall}; + +/// How an increment treats the key's TTL, the two ways `RedisCache` increments with one. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Expiry { + /// `async_increment(refresh_ttl=True)` and every pipelined increment: `EXPIRE` after each write. + Refresh(Duration), + /// `async_increment(refresh_ttl=False)`: read the `TTL` and `EXPIRE` only a key that has none. + /// Pipelined it is lowered to [`Expiry::Refresh`], like `RedisBatch.increment`. + KeepExisting(Duration), +} + +/// A number written the way redis-py encodes it: an `int` with `str()`, a `float` with `repr()`. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Amount { + Int(i64), + Float(f64), +} + +impl ToRedisArgs for Amount { + fn write_redis_args(&self, out: &mut W) { + match self { + Self::Int(value) => out.write_arg(value.to_string().as_bytes()), + Self::Float(value) => out.write_arg(float_repr(*value).as_bytes()), + } + } +} + +impl ToSingleRedisArg for Amount {} + +type Build = Box) -> Plan + Send>; +type Decode = Box Result + Send>; + +/// One Redis command as a value: what it sends alone and pipelined, and how its reply decodes. Every +/// executor (`RedisClient::send`, `Pipeline::with`, `RedisBatch::add`) applies its namespace to the +/// keys once. +#[must_use] +pub struct Command { + operation: Operation, + keys: Vec, + idempotent: bool, + build: Build, + decode: Decode, +} + +/// What a command sends once its keys carry the namespace. Only the cases whose alone and pipelined +/// forms differ get a variant of their own. +pub(crate) enum Plan { + /// Answered without a round trip, like an empty `MGET` or `DEL`. + Reply(Value), + /// Sent together; the reply is the first command's, failing if any failed. + Commands(Vec), + /// Alone through `invoke_async`, pipelined as `EVALSHA` resent once on `NOSCRIPT`. + Script(ScriptCall), + /// Alone the increment reads the `TTL` and sends `EXPIRE` only to a key without one, pipelined it + /// is the increment then `EXPIRE`. + KeepingTtl { + increment: Cmd, + key: String, + ttl: Duration, + trips: Trips, + }, +} + +#[derive(Clone, Copy)] +pub(crate) enum Trips { + /// `async_increment`: the increment and its `TTL` in one round trip. + Shared, + /// `increment_cache`: the increment and its `TTL` each on its own round trip. + Separate, +} + +/// A command with the namespace applied, ready for an executor to send. +pub(crate) struct Op { + pub(crate) operation: Operation, + pub(crate) key: String, + pub(crate) idempotent: bool, + pub(crate) plan: Plan, +} + +impl Plan { + /// What a pipeline sends for this op. + pub(crate) fn pipelined(&self) -> Vec { + match self { + Self::Reply(_) => Vec::new(), + Self::Commands(commands) => commands.clone(), + Self::Script(call) => vec![call.script.evalsha(&call.keys, &call.args)], + Self::KeepingTtl { + increment, + key, + ttl, + .. + } => vec![increment.clone(), expire_cmd(key, *ttl)], + } + } +} + +impl Command { + fn new( + operation: Operation, + keys: Vec, + idempotent: bool, + build: impl FnOnce(Vec) -> Plan + Send + 'static, + decode: impl FnOnce(Value) -> Result + Send + 'static, + ) -> Self { + Self { + operation, + keys, + idempotent, + build: Box::new(build), + decode: Box::new(decode), + } + } + + /// Attaches a codec over the decoded reply, such as a store's value format. + pub fn map(self, decode: fn(T) -> Result) -> Command + where + T: 'static, + { + let inner = self.decode; + Command { + operation: self.operation, + keys: self.keys, + idempotent: self.idempotent, + build: self.build, + decode: Box::new(move |value| inner(value).and_then(decode)), + } + } + + pub(crate) fn prepare(self, namespace: &Namespace) -> (Op, Decode) { + let keys: Vec = self.keys.iter().map(|key| namespace.apply(key)).collect(); + let key = keys.first().cloned().unwrap_or_default(); + let op = Op { + operation: self.operation, + key, + idempotent: self.idempotent, + plan: (self.build)(keys), + }; + (op, self.decode) + } +} + +fn decode(value: Value) -> Result { + T::from_redis_value(value).map_err(|error| Error::Decode(error.into())) +} + +fn owned(keys: &[&str]) -> Vec { + keys.iter().map(|key| (*key).to_owned()).collect() +} + +fn single(keys: Vec) -> String { + keys.into_iter().next().unwrap_or_default() +} + +fn expire_cmd(key: &str, ttl: Duration) -> Cmd { + redis::cmd("EXPIRE").arg(key).arg(ttl.as_secs()).to_owned() +} + +pub(crate) fn ttl_cmd(key: &str) -> Cmd { + redis::cmd("TTL").arg(key).to_owned() +} + +/// A reply whose shape a codec over it refuses, as `what` describes. +pub fn unexpected_reply(what: &'static str) -> Error { + Error::Decode(RedisError::from((ErrorKind::UnexpectedReturnType, what))) +} + +pub fn get(key: &str) -> Command> { + Command::new( + Operation::Command("GET"), + vec![key.to_owned()], + true, + |keys| Plan::Commands(vec![redis::cmd("GET").arg(single(keys)).to_owned()]), + decode, + ) +} + +/// One `MGET` of `keys` in order, repeated keys read once, like `RedisBatch.mget`, answering one +/// value per key of `keys` in their order. No keys is no round trip, since Redis rejects an empty +/// `MGET`. Pipelined, every key must share a hash slot on a cluster, and a batch never pipelines +/// there. +pub fn mget(keys: &[&str]) -> Command>> { + let unique = keys.iter().fold(Vec::::new(), |mut unique, key| { + if !unique.iter().any(|seen| seen == key) { + unique.push((*key).to_owned()); + } + unique + }); + let positions: Vec = keys + .iter() + .filter_map(|key| unique.iter().position(|seen| seen == key)) + .collect(); + Command::new( + Operation::Command("MGET"), + unique, + true, + |keys| match keys.is_empty() { + true => Plan::Reply(Value::Array(Vec::new())), + false => Plan::Commands(vec![redis::cmd("MGET").arg(keys).to_owned()]), + }, + move |reply| { + let values: Vec = decode(reply)?; + let in_key_order = positions + .iter() + .map(|&position| values.get(position).cloned()) + .collect::>>() + .ok_or_else(|| unexpected_reply("MGET answered fewer values than keys"))?; + decode(Value::Array(in_key_order)) + }, + ) +} + +// redis-py's `set(ex=...)` takes whole seconds, as `int(ttl)` or a `timedelta`'s `int(total_seconds())` +pub fn set(key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Command<()> { + let value = value.to_redis_args().concat(); + Command::new( + Operation::Command("SET"), + vec![key.to_owned()], + true, + move |keys| { + Plan::Commands(vec![ + redis::cmd("SET") + .arg(single(keys)) + .arg(value) + .arg("EX") + .arg(ttl.as_secs()) + .to_owned(), + ]) + }, + decode, + ) +} + +/// `SET ... EX ... NX`, `EX` before `NX` as redis-py builds it. `true` when this call wrote the key. +pub fn set_nx(key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Command { + let value = value.to_redis_args().concat(); + Command::new( + Operation::Command("SET"), + vec![key.to_owned()], + false, + move |keys| { + Plan::Commands(vec![ + redis::cmd("SET") + .arg(single(keys)) + .arg(value) + .arg("EX") + .arg(ttl.as_secs()) + .arg("NX") + .to_owned(), + ]) + }, + |value| decode::>(value).map(|reply| reply.is_some()), + ) +} + +/// One multi-key `DEL`, no round trip for no keys. Pipelined, every key must share a hash slot. +pub fn del(keys: &[&str]) -> Command { + Command::new( + Operation::Command("DEL"), + owned(keys), + true, + |keys| match keys.is_empty() { + true => Plan::Reply(Value::Int(0)), + false => Plan::Commands(vec![redis::cmd("DEL").arg(keys).to_owned()]), + }, + decode, + ) +} + +pub fn expire(key: &str, ttl: Duration) -> Command { + Command::new( + Operation::Command("EXPIRE"), + vec![key.to_owned()], + true, + move |keys| Plan::Commands(vec![expire_cmd(&single(keys), ttl)]), + decode, + ) +} + +pub fn ttl(key: &str) -> Command { + Command::new( + Operation::Command("TTL"), + vec![key.to_owned()], + true, + |keys| Plan::Commands(vec![ttl_cmd(&single(keys))]), + |value| { + Ok(match decode::(value)? { + -2 => Ttl::Missing, + -1 => Ttl::Persistent, + seconds => Ttl::Expires(Duration::from_secs(seconds.unsigned_abs())), + }) + }, + ) +} + +/// `RedisCache.increment_cache`: `INCRBY`, then `TTL`, then `EXPIRE` only when the key has no TTL, +/// each its own round trip. +pub fn incr_by(key: &str, delta: i64, ttl: Duration) -> Command { + Command::new( + Operation::Command("INCRBY"), + vec![key.to_owned()], + false, + move |keys| { + let key = single(keys); + Plan::KeepingTtl { + increment: redis::cmd("INCRBY").arg(&key).arg(delta).to_owned(), + key, + ttl, + trips: Trips::Separate, + } + }, + decode, + ) +} + +/// `INCRBYFLOAT`, with the TTL handled as `expiry` says, or left alone without one. Alone it is +/// `RedisCache.async_increment`, pipelined `RedisBatch.increment`. +pub fn incr_by_float(key: &str, amount: Amount, expiry: Option) -> Command { + Command::new( + Operation::Command("INCRBYFLOAT"), + vec![key.to_owned()], + false, + move |keys| { + let key = single(keys); + let increment = redis::cmd("INCRBYFLOAT").arg(&key).arg(amount).to_owned(); + match expiry { + None => Plan::Commands(vec![increment]), + Some(Expiry::Refresh(ttl)) => { + Plan::Commands(vec![increment, expire_cmd(&key, ttl)]) + } + Some(Expiry::KeepExisting(ttl)) => Plan::KeepingTtl { + increment, + key, + ttl, + trips: Trips::Shared, + }, + } + }, + decode, + ) +} + +/// A registered script, like redis-py's `Script`. Every key it touches must be in `keys`. +pub fn eval( + script: &Script, + keys: &[&str], + args: impl ToRedisArgs, +) -> Command { + let script = script.clone(); + let args = args.to_redis_args(); + Command::new( + Operation::Script, + owned(keys), + false, + move |keys| Plan::Script(ScriptCall { script, keys, args }), + decode, + ) +} + +/// A script sent with its source as `EVAL` every time, like redis-py's `Redis.eval`, for the scripts +/// `RedisCache` runs that way rather than through `register_script`. +pub fn eval_source( + source: &'static str, + keys: &[&str], + args: impl ToRedisArgs, +) -> Command { + let args = args.to_redis_args(); + Command::new( + Operation::Script, + owned(keys), + false, + move |keys| { + Plan::Commands(vec![ + redis::cmd("EVAL") + .arg(source) + .arg(keys.len()) + .arg(keys) + .arg(args) + .to_owned(), + ]) + }, + decode, + ) +} + +/// The channel is sent as given, without the namespace, like redis-py's `publish`. +pub fn publish(channel: &str, payload: &[u8]) -> Command<()> { + let command = redis::cmd("PUBLISH").arg(channel).arg(payload).to_owned(); + Command::new( + Operation::Command("PUBLISH"), + Vec::new(), + false, + move |_| Plan::Commands(vec![command]), + |_| Ok(()), + ) +} pub(crate) fn ping() -> Cmd { redis::cmd("PING") } - -pub(crate) fn get(key: String) -> Cmd { - redis::cmd("GET").arg(key).to_owned() -} - -pub(crate) fn mget(keys: Vec) -> Cmd { - redis::cmd("MGET").arg(keys).to_owned() -} - -pub(crate) fn set(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd { - redis::cmd("SET") - .arg(key) - .arg(value) - .arg("PX") - .arg(millis(ttl)) - .to_owned() -} - -pub(crate) fn set_nx(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd { - redis::cmd("SET") - .arg(key) - .arg(value) - .arg("NX") - .arg("PX") - .arg(millis(ttl)) - .to_owned() -} - -pub(crate) fn del(keys: Vec) -> Cmd { - redis::cmd("DEL").arg(keys).to_owned() -} - -pub(crate) fn expire(key: String, ttl: Duration) -> Cmd { - redis::cmd("PEXPIRE").arg(key).arg(millis(ttl)).to_owned() -} - -pub(crate) fn ttl(key: String) -> Cmd { - redis::cmd("PTTL").arg(key).to_owned() -} - -pub(crate) fn incr_by(key: String, delta: i64) -> Cmd { - redis::cmd("INCRBY").arg(key).arg(delta).to_owned() -} - -pub(crate) fn incr_by_float(key: String, delta: f64) -> Cmd { - redis::cmd("INCRBYFLOAT").arg(key).arg(delta).to_owned() -} - -pub(crate) fn rpush(key: String, values: impl ToRedisArgs) -> Cmd { - redis::cmd("RPUSH").arg(key).arg(values).to_owned() -} - -fn millis(duration: Duration) -> u64 { - u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) -} diff --git a/litellm-rust/crates/redis-client/src/config.rs b/litellm-rust/crates/redis-client/src/config.rs index c40b9eb6b36..f43ab4feb85 100644 --- a/litellm-rust/crates/redis-client/src/config.rs +++ b/litellm-rust/crates/redis-client/src/config.rs @@ -1,118 +1,458 @@ -use std::time::Duration; +use std::{fmt, path::PathBuf, pin::Pin, sync::Arc, time::Duration}; -use redis::{ConnectionAddr, ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, RedisResult}; +use futures_util::{Stream, StreamExt}; + +use redis::{ + BasicAuth, ClientTlsConfig, ConnectionAddr, ConnectionInfo, ErrorKind, IntoConnectionInfo, + ProtocolVersion, RedisConnectionInfo, RedisError, RedisResult, StreamingCredentialsProvider, + TlsCertificates, TlsMode, io::tcp::TcpSettings, +}; use litellm_circuit_breaker::Breaker; use litellm_secret_value::SecretValue; -use crate::Namespace; +use pkcs8::EncryptedPrivateKeyInfo; +use crate::{ConnectError, KeyError, Namespace, RedisEvents}; +use litellm_redis_types::{Cluster, RedisUrl, Target, TlsOptions}; + +// litellm/constants.py REDIS_SOCKET_TIMEOUT, which _coerce_redis_kwargs_types sets when socket_timeout is unset pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5); +// below the idle timeouts of common cloud load balancers (AWS NLB 350s, Azure Load Balancer 240s), which the +// OS default of two hours never probes within +const KEEPALIVE_IDLE: Duration = Duration::from_secs(60); pub struct RedisConfig { - pub endpoint: Endpoint, + pub target: Target, + pub credentials: Credentials, pub namespace: Namespace, - pub timeout: Duration, + pub timeouts: Timeouts, pub breaker: Option, + pub events: Option>, + /// Ping every node before the next command once the client sat idle this long, like + /// redis-py's `health_check_interval`, so a connection the server dropped silently is + /// replaced before a real command pays for it. + pub health_check: Option, + /// Probe a connection with TCP keepalive once it sat idle, like redis-py's `socket_keepalive`. + pub keepalive: bool, } -#[derive(Clone, Debug, PartialEq)] -pub enum Endpoint { - Url(SecretValue), - Host(HostEndpoint), - Cluster(ClusterEndpoint), +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Timeouts { + pub response: Duration, + pub connect: Duration, } -#[derive(Clone, Debug, PartialEq)] -pub struct HostEndpoint { - pub host: String, - pub port: u16, - pub username: Option, - pub password: Option, - pub tls: bool, - pub db: i64, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct ClusterEndpoint { - pub nodes: Vec, - pub username: Option, - pub password: Option, - pub tls: bool, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct ClusterNode { - pub host: String, - pub port: u16, -} - -impl ClusterEndpoint { - pub(crate) fn addrs(&self) -> impl Iterator + '_ { - self.nodes - .iter() - .map(|node| addr(&node.host, node.port, self.tls)) +impl Timeouts { + pub fn new(socket_timeout: Duration, connect_timeout: Option) -> Self { + Self { + response: socket_timeout, + connect: connect_timeout.unwrap_or(socket_timeout), + } } } -impl HostEndpoint { - pub(crate) fn connection_info(&self) -> RedisResult { - let settings = RedisConnectionInfo::default().set_db(self.db); - let settings = match &self.username { - Some(username) => settings.set_username(username.expose()), - None => settings, - }; - let settings = match &self.password { - Some(password) => settings.set_password(password.expose()), - None => settings, - }; - Ok(addr(&self.host, self.port, self.tls) - .into_connection_info()? - .set_redis_settings(settings)) +impl Default for Timeouts { + fn default() -> Self { + Self::new(DEFAULT_TIMEOUT, None) } } -fn addr(host: &str, port: u16, tls: bool) -> ConnectionAddr { +#[derive(Clone, Debug)] +pub enum Credentials { + Static { + username: Option, + password: Option, + }, + Provider(CredentialsProvider), +} + +#[derive(Clone)] +pub struct CredentialsProvider(pub Arc); + +impl fmt::Debug for CredentialsProvider { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("CredentialsProvider") + .finish_non_exhaustive() + } +} + +impl StreamingCredentialsProvider for CredentialsProvider { + fn subscribe(&self) -> Pin> + Send + 'static>> { + self.0.subscribe() + } +} + +impl Credentials { + pub const NONE: Self = Self::Static { + username: None, + password: None, + }; +} + +pub(crate) fn tls_mode(tls: &TlsOptions) -> TlsMode { + match tls.verify_certificate { + true => TlsMode::Secure, + false => TlsMode::Insecure, + } +} + +pub(crate) fn tls_certificates(tls: &TlsOptions) -> Result { + let client_tls = match (&tls.certfile, &tls.keyfile) { + (Some(certfile), keyfile) => Some(ClientTlsConfig { + client_cert: read(certfile)?, + client_key: client_key( + keyfile.as_ref().unwrap_or(certfile), + tls.key_password.as_ref(), + )?, + }), + (None, _) => None, + }; + let from_file: Option> = tls.ca_certs.as_ref().map(read).transpose()?; + let from_directory: Vec> = tls + .ca_path + .as_ref() + .map(read_directory) + .transpose()? + .unwrap_or_default(); + let configured: Vec> = from_file + .into_iter() + .chain(tls.ca_data.iter().map(|pem| pem.clone().into_bytes())) + .chain(from_directory) + .collect(); + Ok(TlsCertificates { + client_tls, + root_cert: (!configured.is_empty()).then(|| with_system_roots(configured.join(&b'\n'))), + }) +} + +/// Every regular file in `path`, in name order, as redis-py hands `ssl_ca_path` to +/// `load_verify_locations(capath=...)` +fn read_directory(path: &PathBuf) -> Result>, ConnectError> { + let tls_file = |source| ConnectError::TlsFile { + path: path.display().to_string(), + source, + }; + let mut files: Vec = std::fs::read_dir(path) + .map_err(tls_file)? + .map(|entry| entry.map(|entry| entry.path())) + .collect::>() + .map_err(tls_file)?; + files.sort(); + files + .iter() + .filter(|file| file.is_file()) + .map(read) + .collect() +} + +pub(crate) fn apply_tls( + tls: &TlsOptions, + info: ConnectionInfo, +) -> Result { + let tls_info = redis::Client::build_with_tls(info, tls_certificates(tls)?)? + .get_connection_info() + .clone(); + let mut addr = match tls_info.addr().clone() { + ConnectionAddr::TcpTls { + host, + port, + tls_params, + .. + } => ConnectionAddr::TcpTls { + host, + port, + insecure: !tls.verify_certificate, + tls_params, + }, + addr => addr, + }; + addr.set_danger_accept_invalid_hostnames(!tls.check_hostname); + Ok(tls_info.set_addr(addr)) +} + +fn with_system_roots(configured: Vec) -> Vec { + let system = rustls_native_certs::load_native_certs(); + let system_pem = system + .certs + .iter() + .map(|cert| pem::encode(&pem::Pem::new("CERTIFICATE", cert.to_vec()))); + configured + .into_iter() + .chain(b"\n".iter().copied()) + .chain(system_pem.flat_map(String::into_bytes)) + .collect() +} + +/// The key as unencrypted PEM, decrypting a PKCS#8 `ENCRYPTED PRIVATE KEY` with `ssl_password` the way +/// `load_cert_chain(password=...)` does, since rustls only reads unencrypted keys. Like OpenSSL, a password next to +/// an unencrypted key is unused +fn client_key(path: &PathBuf, password: Option<&SecretValue>) -> Result, ConnectError> { + let contents = read(path)?; + let key_error = |reason| ConnectError::TlsKey { + path: path.display().to_string(), + reason, + }; + let Ok(sections) = pem::parse_many(&contents) else { + return Ok(contents); + }; + let encrypted = sections + .iter() + .find(|section| section.tag() == "ENCRYPTED PRIVATE KEY"); + match (encrypted, password) { + (Some(_), None) => Err(key_error(KeyError::MissingPassword)), + (Some(section), Some(password)) => { + let decrypted = EncryptedPrivateKeyInfo::try_from(section.contents()) + .and_then(|info| info.decrypt(password.expose())) + .map_err(|source| key_error(KeyError::Decrypt(source)))?; + Ok(pem::encode(&pem::Pem::new("PRIVATE KEY", decrypted.as_bytes())).into_bytes()) + } + (None, _) if sections.iter().any(is_legacy_encrypted) => { + Err(key_error(KeyError::LegacyEncryption)) + } + (None, _) => Ok(contents), + } +} + +fn is_legacy_encrypted(section: &pem::Pem) -> bool { + section + .headers() + .get("Proc-Type") + .is_some_and(|proc_type| proc_type.contains("ENCRYPTED")) +} + +fn read(path: &PathBuf) -> Result, ConnectError> { + std::fs::read(path).map_err(|source| ConnectError::TlsFile { + path: path.display().to_string(), + source, + }) +} + +pub(crate) fn addr(host: &str, port: u16, tls: Option<&TlsOptions>) -> ConnectionAddr { match tls { - true => ConnectionAddr::TcpTls { + Some(tls) => ConnectionAddr::TcpTls { host: host.to_owned(), port, - insecure: false, + insecure: !tls.verify_certificate, tls_params: None, }, - false => ConnectionAddr::Tcp(host.to_owned(), port), + None => ConnectionAddr::Tcp(host.to_owned(), port), } } +pub(crate) fn cluster_addrs(cluster: &Cluster) -> impl Iterator + '_ { + cluster + .nodes + .iter() + .map(|node| addr(&node.host, node.port, cluster.tls.as_ref())) +} + +pub(crate) fn tcp_settings(keepalive: bool) -> TcpSettings { + let settings = TcpSettings::default().set_nodelay(true); + match keepalive { + true => settings.set_keepalive(socket2::TcpKeepalive::new().with_time(KEEPALIVE_IDLE)), + false => settings, + } +} + +pub(crate) fn connection_info( + target: &Target, + credentials: &Credentials, + keepalive: bool, +) -> Result { + let (info, tls) = match target { + Target::Url(url) => (url_info(url, credentials)?, url.tls.as_ref()), + Target::Standalone(host) => { + let info = addr(&host.host, host.port, host.tls.as_ref()).into_connection_info()?; + let settings = info.redis_settings().clone().set_db(host.db); + (info.set_redis_settings(settings), host.tls.as_ref()) + } + Target::Unix(socket) => { + let info = ConnectionAddr::Unix(socket.path.clone()).into_connection_info()?; + let settings = info.redis_settings().clone().set_db(socket.db); + (info.set_redis_settings(settings), None) + } + Target::Cluster(_) | Target::Sentinel(_) => { + unreachable!("cluster and sentinel endpoints build their own connection info") + } + }; + let info = with_static_credentials(info, credentials).set_tcp_settings(tcp_settings(keepalive)); + match (tls, info.addr()) { + (Some(tls), ConnectionAddr::TcpTls { .. }) => apply_tls(tls, info), + _ => Ok(info), + } +} + +// parse_url reads ?db= over the path, and the url's userinfo over ?username= and ?password=. A credential +// provider supersedes both, as it does in redis-py. ?protocol=3 would switch redis-rs to RESP3, which no +// reply decode here expects, so the protocol is ignored as the resolver warns +fn url_info(url: &RedisUrl, credentials: &Credentials) -> Result { + let info = url.url.expose().into_connection_info()?; + let settings = info + .redis_settings() + .clone() + .set_protocol(ProtocolVersion::RESP2); + let settings = match url.db { + Some(db) => settings.set_db(db), + None => settings, + }; + let info = info.set_redis_settings(settings); + Ok(match credentials { + Credentials::Static { .. } => with_static_credentials( + info, + &Credentials::Static { + username: url.username.clone(), + password: url.password.clone(), + }, + ), + Credentials::Provider(_) => info, + }) +} + +fn with_static_credentials(info: ConnectionInfo, credentials: &Credentials) -> ConnectionInfo { + let settings = static_credentials(info.redis_settings().clone(), credentials); + info.set_redis_settings(settings) +} + +fn static_credentials( + settings: RedisConnectionInfo, + credentials: &Credentials, +) -> RedisConnectionInfo { + let Credentials::Static { username, password } = credentials else { + return settings; + }; + let settings = match (username, settings.username()) { + (Some(username), None) => settings.set_username(username.expose()), + _ => settings, + }; + match (password, settings.password()) { + (Some(password), None) => settings.set_password(password.expose()), + _ => settings, + } +} + +/// `settings` with the credentials a connection opened now sends. A provider is asked once for its +/// current pair, which is when redis-py asks its `credential_provider`. +pub(crate) async fn authenticated( + settings: RedisConnectionInfo, + credentials: &Credentials, +) -> RedisResult { + let Credentials::Provider(provider) = credentials else { + return Ok(static_credentials(settings, credentials)); + }; + let current = provider.subscribe().next().await.ok_or_else(|| { + RedisError::from(( + ErrorKind::AuthenticationFailed, + "the Redis credential provider ended before giving credentials", + )) + })??; + Ok(settings + .set_username(current.username()) + .set_password(current.password())) +} + #[cfg(test)] mod tests { - use litellm_secret_value::SecretValue; + use litellm_redis_types::{RedisUrl, UnixSocket}; use rstest::rstest; - use super::{ClusterEndpoint, Endpoint, HostEndpoint}; + use super::*; #[rstest] - #[case::url(Endpoint::Url(SecretValue::new("redis://credential-value@cache:6379/0")))] - #[case::host( - Endpoint::Host(HostEndpoint { - host: "cache".to_owned(), - port: 6379, - username: Some(SecretValue::new("credential-value")), - password: Some(SecretValue::new("credential-value")), - tls: false, - db: 0, - }) - )] - #[case::cluster(Endpoint::Cluster(ClusterEndpoint { - nodes: Vec::new(), - username: Some(SecretValue::new("credential-value")), - password: Some(SecretValue::new("credential-value")), - tls: false, + #[case::host(Target::Standalone(crate::Standalone { host: "redis".into(), port: 6379, tls: None, db: 0 }))] + #[case::url(Target::Url(RedisUrl { + url: SecretValue::new("redis://cache:6379".to_owned()), + db: None, + username: None, + password: None, + tls: None, }))] - fn debug_output_never_shows_the_secret(#[case] endpoint: Endpoint) { - let debug = format!("{endpoint:?}"); + fn keepalive_probes_an_idle_socket_only_when_enabled(#[case] target: Target) { + let probed = |keepalive| { + connection_info(&target, &Credentials::NONE, keepalive) + .unwrap() + .tcp_settings() + .keepalive() + .is_some() + }; - assert!(!debug.contains("credential-value"), "{debug}"); - assert!(debug.contains("[REDACTED]"), "{debug}"); + assert!(probed(true)); + assert!(!probed(false)); + } + + #[test] + fn a_unix_socket_replaces_host_port_and_tls_like_redis_py() { + let target = Target::Unix(UnixSocket { + path: "/tmp/redis.sock".into(), + db: 3, + }); + + let info = connection_info(&target, &Credentials::NONE, true).unwrap(); + + assert_eq!(*info.addr(), ConnectionAddr::Unix("/tmp/redis.sock".into())); + assert_eq!(info.redis_settings().db(), 3); + } + + #[rstest] + #[case::the_query_db_wins_over_the_path("redis://cache:6379/1", Some(3), None, 3, None)] + #[case::the_path_db_without_a_query_db("redis://cache:6379/1", None, None, 1, None)] + #[case::the_query_password_fills_a_missing_userinfo( + "redis://cache:6379", + None, + Some("query"), + 0, + Some("query") + )] + #[case::the_userinfo_wins_over_the_query_password( + "redis://:userinfo@cache:6379", + None, + Some("query"), + 0, + Some("userinfo") + )] + fn a_url_reads_its_query_like_parse_url( + #[case] url: &str, + #[case] db: Option, + #[case] password: Option<&str>, + #[case] expected_db: i64, + #[case] expected_password: Option<&str>, + ) { + let target = Target::Url(RedisUrl { + url: SecretValue::new(url), + tls: None, + db, + username: None, + password: password.map(SecretValue::new), + }); + let static_password = Credentials::Static { + username: None, + password: Some(SecretValue::new("static")), + }; + + let info = connection_info(&target, &static_password, true).unwrap(); + + assert_eq!(info.redis_settings().db(), expected_db); + assert_eq!( + info.redis_settings().password(), + expected_password.or(Some("static")) + ); + } + + #[rstest] + #[case::by_number("redis://cache:6379/?protocol=3")] + #[case::by_name("redis://cache:6379/?protocol=resp3")] + fn a_url_never_switches_the_protocol_from_resp2(#[case] url: &str) { + let target = Target::Url(RedisUrl { + url: SecretValue::new(url), + tls: None, + db: None, + username: None, + password: None, + }); + + let info = connection_info(&target, &Credentials::NONE, true).unwrap(); + + assert_eq!(info.redis_settings().protocol(), ProtocolVersion::RESP2); } } diff --git a/litellm-rust/crates/redis-client/src/connection.rs b/litellm-rust/crates/redis-client/src/connection.rs index 41729bfd107..c1f82490b3b 100644 --- a/litellm-rust/crates/redis-client/src/connection.rs +++ b/litellm-rust/crates/redis-client/src/connection.rs @@ -1,53 +1,161 @@ -use std::time::Duration; - -use redis::{ - Cmd, ConnectionInfo, Pipeline, RedisFuture, RedisResult, Value, - aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig}, - cluster::{ClusterClientBuilder, ClusterConfig}, - cluster_async::ClusterConnection, +use std::{ + sync::{ + Arc, Mutex, PoisonError, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, }; -use crate::ClusterEndpoint; +use redis::{ + AsyncConnectionConfig, Cmd, ConnectionAddr, ConnectionInfo, ErrorKind, Pipeline, + RedisConnectionInfo, RedisError, RedisFuture, RedisResult, ServerErrorKind, TlsCertificates, + TlsMode, Value, + aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig, MultiplexedConnection}, + cluster::{ClusterClientBuilder, ClusterConfig}, + cluster_async::ClusterConnection, + cluster_routing::{MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo}, + sentinel::{SentinelClient, SentinelClientBuilder, SentinelServerType}, +}; -const RETRIES: u32 = 3; +use crate::{ + Cluster, ConnectError, Credentials, Sentinel, Target, Timeouts, + config::{self, addr}, + error::Deployment, +}; + +// redis-py 5.3.1 redis/asyncio/connection.py: without retry options a connection gets Retry(NoBackoff(), 0) +const RECONNECT_RETRIES: usize = 0; +// redis-py 5.3.1 redis/asyncio/cluster.py: cluster_error_retry_attempts=3 +const CLUSTER_RETRIES: u32 = 3; +// redis-py 5.3.1 redis/asyncio/cluster.py sleeps a flat 0.25s before retrying after CLUSTERDOWN. +// redis-rs draws the wait from [min, max), so a 1ms window pins it to 250ms +const CLUSTER_RETRY_WAIT_MS: u64 = 250; #[derive(Clone)] pub(crate) enum Connection { - Standalone(ConnectionManager), - Cluster(ClusterConnection), + Standalone(Renewing), + Cluster(Renewing), + Sentinel(SentinelConnection), } impl Connection { - pub(crate) fn standalone(info: ConnectionInfo, timeout: Duration) -> RedisResult { - let config = ConnectionManagerConfig::new() - .set_connection_timeout(Some(timeout)) - .set_response_timeout(Some(timeout)) - .set_number_of_retries(RETRIES as usize); - let manager = ConnectionManager::new_lazy_with_config(redis::Client::open(info)?, config)?; - Ok(Self::Standalone(manager)) + pub(crate) fn new( + target: &Target, + credentials: &Credentials, + timeouts: Timeouts, + keepalive: bool, + ) -> Result { + match target { + Target::Url(_) | Target::Standalone(_) | Target::Unix(_) => Self::standalone( + config::connection_info(target, credentials, keepalive)?, + credentials, + timeouts, + ), + Target::Cluster(cluster) => Self::cluster(cluster, credentials, timeouts, keepalive), + Target::Sentinel(sentinel) => Ok(Self::Sentinel(SentinelConnection::new( + sentinel, + credentials, + timeouts, + keepalive, + )?)), + } } - pub(crate) fn cluster(endpoint: &ClusterEndpoint, timeout: Duration) -> RedisResult { - let builder = ClusterClientBuilder::new(endpoint.addrs()) - .connection_timeout(timeout) - .response_timeout(timeout) - .retries(RETRIES); - let builder = match &endpoint.username { - Some(username) => builder.username(username.expose()), + pub(crate) fn deployment(&self) -> Deployment { + match self { + Self::Standalone(_) => Deployment::Standalone, + Self::Cluster(_) => Deployment::Cluster, + Self::Sentinel(_) => Deployment::Sentinel, + } + } + + fn standalone( + info: ConnectionInfo, + credentials: &Credentials, + timeouts: Timeouts, + ) -> Result { + let config = ConnectionManagerConfig::new() + .set_connection_timeout(Some(timeouts.connect)) + .set_response_timeout(Some(timeouts.response)) + .set_number_of_retries(RECONNECT_RETRIES); + let config = match credentials { + Credentials::Provider(provider) => config.set_credentials_provider(provider.clone()), + Credentials::Static { .. } => config, + }; + let client = redis::Client::open(info)?; + Ok(Self::Standalone(Renewing::new(move || { + ConnectionManager::new_lazy_with_config(client.clone(), config.clone()) + })?)) + } + + fn cluster( + endpoint: &Cluster, + credentials: &Credentials, + timeouts: Timeouts, + keepalive: bool, + ) -> Result { + let builder = ClusterClientBuilder::new(config::cluster_addrs(endpoint)) + .tcp_settings(config::tcp_settings(keepalive)) + .connection_timeout(timeouts.connect) + .response_timeout(timeouts.response) + .retries(CLUSTER_RETRIES) + .retry_wait_formula(0, 2) + .min_retry_wait(CLUSTER_RETRY_WAIT_MS) + .max_retry_wait(CLUSTER_RETRY_WAIT_MS + 1); + let builder = match &endpoint.tls { + Some(tls) => builder + .tls(config::tls_mode(tls)) + .certs(config::tls_certificates(tls)?) + .danger_accept_invalid_hostnames(!tls.check_hostname), None => builder, }; - let builder = match &endpoint.password { - Some(password) => builder.password(password.expose()), - None => builder, + let builder = match credentials { + Credentials::Static { username, password } => { + let builder = match username { + Some(username) => builder.username(username.expose()), + None => builder, + }; + match password { + Some(password) => builder.password(password.expose()), + None => builder, + } + } + Credentials::Provider(provider) => builder.set_credentials_provider(provider.clone()), }; - let config = ClusterConfig::new() - .set_connection_timeout(timeout) - .set_response_timeout(timeout); - Ok(Self::Cluster( - builder - .build()? - .get_pending_async_connection_with_config(config), - )) + // The builder's response timeout bounds the whole request, but with_config overwrites each node's + // response timeout with this one even when unset, so it has to be repeated here + let config = ClusterConfig::new().set_response_timeout(timeouts.response); + let client = builder.build()?; + Ok(Self::Cluster(Renewing::new(move || { + Ok(client.get_pending_async_connection_with_config(config.clone())) + })?)) + } +} + +impl Connection { + /// `PING` on the one multiplexed connection, or on every cluster node, so a dead socket is + /// replaced before a real command uses it. + pub(crate) async fn ping_every_node(&self) -> RedisResult<()> { + let ping = redis::cmd("PING"); + match self { + Self::Cluster(renewing) => { + let (generation, mut connection) = renewing.current(); + let reply = connection + .route_command( + ping, + RoutingInfo::MultiNode(( + MultipleNodeRoutingInfo::AllNodes, + Some(ResponsePolicy::AllSucceeded), + )), + ) + .await; + renewing.settle(generation, &reply); + reply.map(drop) + } + Self::Standalone(_) | Self::Sentinel(_) => { + ping.query_async::(&mut self.clone()).await.map(drop) + } + } } } @@ -56,6 +164,7 @@ impl ConnectionLike for Connection { match self { Self::Standalone(connection) => connection.req_packed_command(cmd), Self::Cluster(connection) => connection.req_packed_command(cmd), + Self::Sentinel(connection) => connection.req_packed_command(cmd), } } @@ -68,6 +177,7 @@ impl ConnectionLike for Connection { match self { Self::Standalone(connection) => connection.req_packed_commands(pipeline, offset, count), Self::Cluster(connection) => connection.req_packed_commands(pipeline, offset, count), + Self::Sentinel(connection) => connection.req_packed_commands(pipeline, offset, count), } } @@ -75,6 +185,441 @@ impl ConnectionLike for Connection { match self { Self::Standalone(connection) => connection.get_db(), Self::Cluster(connection) => connection.get_db(), + Self::Sentinel(connection) => connection.get_db(), } } } + +/// A connection swapped for a fresh one once a command on it times out, like redis-py disconnecting +/// a connection whose read timed out. redis-rs keeps the socket after a timeout and the cluster +/// client retries on it, so a socket that went silent would time out every command until TCP gives +/// up on it. +#[derive(Clone)] +pub(crate) struct Renewing(Arc>); + +type Open = Box RedisResult + Send + Sync>; + +struct RenewingState { + open: Open, + current: Mutex<(u64, C)>, +} + +impl Renewing { + fn new(open: impl Fn() -> RedisResult + Send + Sync + 'static) -> RedisResult { + let first = open()?; + Ok(Self(Arc::new(RenewingState { + open: Box::new(open), + current: Mutex::new((0, first)), + }))) + } + + fn current(&self) -> (u64, C) { + self.0 + .current + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } + + /// Replaces the connection of `generation` when its reply timed out, unless another timeout + /// already replaced it. + fn settle(&self, generation: u64, reply: &RedisResult) { + if !reply.as_ref().is_err_and(RedisError::is_timeout) { + return; + } + let mut current = self + .0 + .current + .lock() + .unwrap_or_else(PoisonError::into_inner); + if current.0 != generation { + return; + } + match (self.0.open)() { + Ok(fresh) => *current = (generation + 1, fresh), + Err(error) => { + tracing::warn!(%error, "replacing a timed-out Redis connection failed") + } + } + } +} + +impl ConnectionLike for Renewing { + fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> { + Box::pin(async move { + let (generation, mut connection) = self.current(); + let reply = connection.req_packed_command(cmd).await; + self.settle(generation, &reply); + reply + }) + } + + fn req_packed_commands<'a>( + &'a mut self, + pipeline: &'a Pipeline, + offset: usize, + count: usize, + ) -> RedisFuture<'a, Vec> { + Box::pin(async move { + let (generation, mut connection) = self.current(); + let replies = connection + .req_packed_commands(pipeline, offset, count) + .await; + self.settle(generation, &replies); + replies + }) + } + + fn get_db(&self) -> i64 { + self.current().1.get_db() + } +} + +#[derive(Clone)] +pub(crate) struct SentinelConnection(Arc); + +struct SentinelState { + discovery: tokio::sync::Mutex<()>, + current: Mutex>, + generations: AtomicU64, + endpoint: Sentinel, + credentials: Credentials, + tls: Option<(TlsMode, TlsCertificates)>, + config: AsyncConnectionConfig, + connect_timeout: Duration, + keepalive: bool, +} + +impl SentinelConnection { + fn new( + endpoint: &Sentinel, + credentials: &Credentials, + timeouts: Timeouts, + keepalive: bool, + ) -> Result { + let tls = match &endpoint.tls { + Some(tls) => Some((config::tls_mode(tls), config::tls_certificates(tls)?)), + None => None, + }; + let state = SentinelState { + discovery: tokio::sync::Mutex::new(()), + current: Mutex::new(None), + generations: AtomicU64::new(0), + endpoint: endpoint.clone(), + credentials: credentials.clone(), + tls, + config: async_config(credentials, timeouts), + connect_timeout: timeouts.connect, + keepalive, + }; + state.sentinel_client(&RedisConnectionInfo::default())?; + Ok(Self(Arc::new(state))) + } + + fn cached(&self) -> Option<(u64, MultiplexedConnection)> { + self.0 + .current + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clone() + } + + async fn connection(&self) -> RedisResult<(u64, MultiplexedConnection)> { + if let Some(current) = self.cached() { + return Ok(current); + } + let _discovering = self.0.discovery.lock().await; + if let Some(current) = self.cached() { + return Ok(current); + } + let connection = self + .discover() + .await? + .get_multiplexed_async_connection_with_config(&self.0.config) + .await?; + let mut current = self + .0 + .current + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let generation = self.0.generations.fetch_add(1, Ordering::Relaxed); + *current = Some((generation, connection.clone())); + Ok((generation, connection)) + } + + pub(crate) async fn master(&self) -> RedisResult { + self.discover().await + } + + async fn discover(&self) -> RedisResult { + let master = tokio::time::timeout(self.0.connect_timeout, self.ask_sentinels()) + .await + .map_err(|_| timed_out("asking the Sentinels for the master timed out"))??; + self.master_client(master.get_connection_info().addr()) + } + + /// redis-rs sends `ROLE` to the master a Sentinel names before trusting it, so that probe gets + /// the data credentials and TLS a connection opened now would use. The Sentinel client is built + /// per discovery since a provider's credentials change after `new`. + async fn ask_sentinels(&self) -> RedisResult { + let data_auth = + config::authenticated(RedisConnectionInfo::default(), &self.0.credentials).await?; + self.0.sentinel_client(&data_auth)?.async_get_client().await + } + + fn master_client(&self, master: &ConnectionAddr) -> Result { + let (host, port) = match master { + ConnectionAddr::Tcp(host, port) | ConnectionAddr::TcpTls { host, port, .. } => { + (host.as_str(), *port) + } + _ => { + return Err(RedisError::from(( + ErrorKind::InvalidClientConfig, + "Sentinel named a master that is not a TCP address", + ))); + } + }; + let endpoint = Target::Standalone(crate::Standalone { + host: host.to_owned(), + port, + tls: self.0.endpoint.tls.clone(), + db: self.0.endpoint.db, + }); + let info = config::connection_info(&endpoint, &self.0.credentials, self.0.keepalive) + .map_err(|error| match error { + ConnectError::Redis(error) => error, + ConnectError::TlsFile { path, .. } => RedisError::from(( + ErrorKind::InvalidClientConfig, + "reading a Redis TLS file", + path, + )), + ConnectError::TlsKey { path, reason } => RedisError::from(( + ErrorKind::InvalidClientConfig, + "reading the Redis TLS key", + format!("{path}: {reason}"), + )), + })?; + redis::Client::open(info) + } + + fn forget(&self, generation: u64) { + let mut current = self + .0 + .current + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if current + .as_ref() + .is_some_and(|(cached, _)| *cached == generation) + { + *current = None; + } + } +} + +impl SentinelState { + fn sentinel_client(&self, data_auth: &RedisConnectionInfo) -> RedisResult { + let endpoint = &self.endpoint; + let sentinels = endpoint + .monitors + .iter() + .map(|node| addr(&node.host, node.port, endpoint.tls.as_ref())); + let builder = SentinelClientBuilder::new( + sentinels, + &endpoint.service_name, + SentinelServerType::Master, + )?; + let (username, password) = match &endpoint.monitor_auth { + Some(auth) => (auth.username.as_ref(), auth.password.as_ref()), + None => (None, None), + }; + let builder = match username { + Some(username) => builder.set_client_to_sentinel_username(username.expose()), + None => builder, + }; + let builder = match password { + Some(password) => builder.set_client_to_sentinel_password(password.expose()), + None => builder, + }; + let builder = match data_auth.username() { + Some(username) => builder.set_client_to_redis_username(username), + None => builder, + }; + let builder = match data_auth.password() { + Some(password) => builder.set_client_to_redis_password(password), + None => builder, + }; + let builder = match &self.tls { + Some((mode, certificates)) => builder + .set_client_to_sentinel_tls_mode(*mode) + .set_client_to_sentinel_certificates(certificates.clone()) + .set_client_to_redis_tls_mode(*mode) + .set_client_to_redis_certificates(certificates.clone()), + None => builder, + }; + builder.build() + } +} + +pub(crate) fn async_config(credentials: &Credentials, timeouts: Timeouts) -> AsyncConnectionConfig { + let config = AsyncConnectionConfig::new() + .set_connection_timeout(Some(timeouts.connect)) + .set_response_timeout(Some(timeouts.response)); + match credentials { + Credentials::Provider(provider) => config.set_credentials_provider(provider.clone()), + Credentials::Static { .. } => config, + } +} + +fn timed_out(message: &'static str) -> RedisError { + RedisError::from(std::io::Error::new(std::io::ErrorKind::TimedOut, message)) +} + +fn needs_rediscovery(error: &RedisError) -> bool { + matches!( + crate::Error::classify(error.clone(), Deployment::Sentinel), + crate::Error::Transport(_) + ) +} + +fn replied_read_only(values: &[Value]) -> bool { + values.iter().any(|value| { + matches!(value, Value::ServerError(error) if error.kind() == Some(ServerErrorKind::ReadOnly)) + }) +} + +impl ConnectionLike for SentinelConnection { + fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> { + Box::pin(async move { + let (generation, mut connection) = self.connection().await?; + let reply = connection.req_packed_command(cmd).await; + let rediscover = match &reply { + Ok(value) => replied_read_only(std::slice::from_ref(value)), + Err(error) => needs_rediscovery(error), + }; + if rediscover { + self.forget(generation); + } + reply + }) + } + + fn req_packed_commands<'a>( + &'a mut self, + pipeline: &'a Pipeline, + offset: usize, + count: usize, + ) -> RedisFuture<'a, Vec> { + Box::pin(async move { + let (generation, mut connection) = self.connection().await?; + let replies = connection + .req_packed_commands(pipeline, offset, count) + .await; + let rediscover = match &replies { + Ok(values) => replied_read_only(values), + Err(error) => needs_rediscovery(error), + }; + if rediscover { + self.forget(generation); + } + replies + }) + } + + fn get_db(&self) -> i64 { + self.0.endpoint.db + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Node; + use litellm_redis_testing::{RedisSentinel, SERVICE_NAME}; + + fn endpoint(server: &RedisSentinel) -> Sentinel { + let (host, port) = server.sentinel(); + Sentinel { + monitors: vec![Node { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + monitor_auth: None, + tls: None, + db: 0, + } + } + + #[rstest::rstest] + #[tokio::test] + async fn a_late_failure_does_not_discard_the_reconnected_sentinel_master() { + let server = RedisSentinel::start().await.unwrap(); + let mut client = SentinelConnection::new( + &endpoint(&server), + &Credentials::NONE, + Timeouts::default(), + true, + ) + .unwrap(); + let (old_generation, _) = client.connection().await.unwrap(); + let original: i64 = redis::cmd("CLIENT") + .arg("ID") + .query_async(&mut client) + .await + .unwrap(); + client.forget(old_generation); + let replacement: i64 = redis::cmd("CLIENT") + .arg("ID") + .query_async(&mut client) + .await + .unwrap(); + assert_ne!(original, replacement); + + client.forget(old_generation); + let after_late_failure: i64 = redis::cmd("CLIENT") + .arg("ID") + .query_async(&mut client) + .await + .unwrap(); + + assert_eq!(replacement, after_late_failure); + } + + #[rstest::rstest] + #[tokio::test] + async fn a_read_only_reply_to_one_command_rediscovers_the_master() { + let server = RedisSentinel::start().await.unwrap(); + let mut client = SentinelConnection::new( + &endpoint(&server), + &Credentials::NONE, + Timeouts::default(), + true, + ) + .unwrap(); + let replica = redis::Client::open(format!("redis://127.0.0.1:{}", server.replica_port())) + .unwrap() + .get_multiplexed_async_connection() + .await + .unwrap(); + let stale = client.0.generations.fetch_add(1, Ordering::Relaxed); + *client.0.current.lock().unwrap() = Some((stale, replica)); + + let on_replica = redis::cmd("SET") + .arg("k") + .arg("1") + .exec_async(&mut client) + .await; + let on_master = redis::cmd("SET") + .arg("k") + .arg("1") + .exec_async(&mut client) + .await; + + assert_eq!( + on_replica.unwrap_err().kind(), + ErrorKind::Server(ServerErrorKind::ReadOnly) + ); + on_master.unwrap(); + } +} diff --git a/litellm-rust/crates/redis-client/src/error.rs b/litellm-rust/crates/redis-client/src/error.rs index ed3b2c74004..54e55615e60 100644 --- a/litellm-rust/crates/redis-client/src/error.rs +++ b/litellm-rust/crates/redis-client/src/error.rs @@ -13,32 +13,93 @@ pub enum Error { } #[derive(Debug, thiserror::Error)] -#[error("connecting to Redis")] -pub struct ConnectError(#[from] RedisError); +pub enum ConnectError { + #[error("connecting to Redis")] + Redis(#[from] RedisError), + #[error("reading the Redis TLS file {path}")] + TlsFile { + path: String, + #[source] + source: std::io::Error, + }, + #[error("reading the Redis TLS key {path}")] + TlsKey { + path: String, + #[source] + reason: KeyError, + }, +} -impl From for Error { - fn from(error: RedisError) -> Self { - match error.kind() { - ErrorKind::Server( - ServerErrorKind::BusyLoading - | ServerErrorKind::ClusterDown - | ServerErrorKind::MasterDown - | ServerErrorKind::TryAgain - | ServerErrorKind::ReadOnly, - ) => Self::Transport(error), - ErrorKind::Server(_) | ErrorKind::Extension => Self::Server(error), - ErrorKind::Parse | ErrorKind::UnexpectedReturnType => Self::Decode(error), +#[derive(Debug, thiserror::Error)] +pub enum KeyError { + #[error("the key is encrypted, so ssl_password is required")] + MissingPassword, + #[error("ssl_password does not decrypt the key, or the key uses an unsupported cipher")] + Decrypt(#[source] pkcs8::Error), + #[error( + "the key uses legacy OpenSSL PEM encryption, which Rust cannot read. Convert it to PKCS#8 with `openssl pkcs8 -topk8`" + )] + LegacyEncryption, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum Deployment { + Standalone, + Cluster, + Sentinel, +} + +// redis-py 5.3.1 redis/_parsers/base.py EXCEPTION_CLASSES["ERR"] entries that subclass ConnectionError +const CONNECTION_ERR_MESSAGES: &[&str] = &[ + "max number of clients reached", + "invalid password", + "Client sent AUTH, but no password is set", + "AUTH called without any password configured for the default user. Are you sure your configuration is correct?", +]; + +impl Error { + pub(crate) fn classify(error: RedisError, deployment: Deployment) -> Self { + match (error.kind(), deployment) { + (ErrorKind::Parse | ErrorKind::UnexpectedReturnType, _) => Self::Decode(error), + (ErrorKind::Server(kind), _) if is_unavailable(kind, deployment) => { + Self::Transport(error) + } + (ErrorKind::Server(ServerErrorKind::ResponseError), _) + if error + .detail() + .is_some_and(|detail| CONNECTION_ERR_MESSAGES.contains(&detail)) => + { + Self::Transport(error) + } + (ErrorKind::Extension, _) if matches!(error.code(), Some("NOAUTH" | "WRONGPASS")) => { + Self::Transport(error) + } + (ErrorKind::Server(_) | ErrorKind::Extension | ErrorKind::InvalidClientConfig, _) => { + Self::Server(error) + } _ => Self::Transport(error), } } } +fn is_unavailable(kind: ServerErrorKind, deployment: Deployment) -> bool { + matches!( + (kind, deployment), + (ServerErrorKind::BusyLoading, _) + | ( + ServerErrorKind::ClusterDown | ServerErrorKind::MasterDown, + Deployment::Cluster + ) + | (ServerErrorKind::ReadOnly, Deployment::Sentinel) + ) +} + #[cfg(test)] mod tests { use redis::{RedisError, Value}; use rstest::rstest; - use super::Error; + use super::{Deployment, Error}; fn reply(resp: &str) -> RedisError { match redis::parse_redis_value(resp.as_bytes()).unwrap() { @@ -47,23 +108,101 @@ mod tests { } } + fn is_transport(resp: &str, deployment: Deployment) -> bool { + match Error::classify(reply(resp), deployment) { + Error::Transport(_) => true, + Error::Server(_) => false, + other => panic!("{resp:?} classified as {other:?}"), + } + } + #[rstest] #[case::wrong_type("-WRONGTYPE Operation against a key holding the wrong kind of value\r\n")] #[case::no_script("-NOSCRIPT No matching script\r\n")] #[case::cross_slot("-CROSSSLOT Keys in request don't hash to the same slot\r\n")] #[case::generic("-ERR value is not an integer or out of range\r\n")] #[case::no_perm("-NOPERM this user has no permissions to access one of the keys\r\n")] - fn rejections_are_server_errors(#[case] resp: &str) { - assert!(matches!(Error::from(reply(resp)), Error::Server(_))); + #[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")] + #[case::out_of_memory("-OOM command not allowed when used memory > 'maxmemory'.\r\n")] + #[case::auth_arity("-ERR wrong number of arguments for 'auth' command\r\n")] + #[case::max_clients_with_suffix("-ERR max number of clients reached, try later\r\n")] + fn rejections_are_server_errors_on_every_deployment( + #[case] resp: &str, + #[values(Deployment::Standalone, Deployment::Cluster, Deployment::Sentinel)] + deployment: Deployment, + ) { + assert!(!is_transport(resp, deployment)); } #[rstest] #[case::loading("-LOADING Redis is loading the dataset in memory\r\n")] - #[case::cluster_down("-CLUSTERDOWN The cluster is down\r\n")] - #[case::master_down("-MASTERDOWN Link with MASTER is down\r\n")] - #[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")] - #[case::read_only("-READONLY You can't write against a read only replica.\r\n")] - fn unavailable_nodes_are_transport_errors(#[case] resp: &str) { - assert!(matches!(Error::from(reply(resp)), Error::Transport(_))); + #[case::no_auth("-NOAUTH Authentication required.\r\n")] + #[case::wrong_pass("-WRONGPASS invalid username-password pair or user is disabled.\r\n")] + #[case::max_clients("-ERR max number of clients reached\r\n")] + #[case::invalid_password("-ERR invalid password\r\n")] + #[case::no_password_set("-ERR Client sent AUTH, but no password is set\r\n")] + #[case::no_default_password( + "-ERR AUTH called without any password configured for the default user. Are you sure your configuration is correct?\r\n" + )] + fn unavailable_nodes_are_transport_errors_on_every_deployment( + #[case] resp: &str, + #[values(Deployment::Standalone, Deployment::Cluster, Deployment::Sentinel)] + deployment: Deployment, + ) { + assert!(is_transport(resp, deployment)); + } + + #[rstest] + #[case::master_down_on_standalone( + "-MASTERDOWN Link with MASTER is down\r\n", + Deployment::Standalone, + false + )] + #[case::master_down_on_cluster( + "-MASTERDOWN Link with MASTER is down\r\n", + Deployment::Cluster, + true + )] + #[case::cluster_down_on_standalone( + "-CLUSTERDOWN The cluster is down\r\n", + Deployment::Standalone, + false + )] + #[case::cluster_down_on_cluster( + "-CLUSTERDOWN The cluster is down\r\n", + Deployment::Cluster, + true + )] + #[case::read_only_on_standalone( + "-READONLY You can't write against a read only replica.\r\n", + Deployment::Standalone, + false + )] + #[case::read_only_on_cluster( + "-READONLY You can't write against a read only replica.\r\n", + Deployment::Cluster, + false + )] + #[case::master_down_on_sentinel( + "-MASTERDOWN Link with MASTER is down\r\n", + Deployment::Sentinel, + false + )] + #[case::cluster_down_on_sentinel( + "-CLUSTERDOWN The cluster is down\r\n", + Deployment::Sentinel, + false + )] + #[case::read_only_on_sentinel( + "-READONLY You can't write against a read only replica.\r\n", + Deployment::Sentinel, + true + )] + fn topology_decides_like_redis_py( + #[case] resp: &str, + #[case] deployment: Deployment, + #[case] transport: bool, + ) { + assert_eq!(is_transport(resp, deployment), transport); } } diff --git a/litellm-rust/crates/redis-client/src/events.rs b/litellm-rust/crates/redis-client/src/events.rs new file mode 100644 index 00000000000..9ea9ff4defa --- /dev/null +++ b/litellm-rust/crates/redis-client/src/events.rs @@ -0,0 +1,78 @@ +use std::{any::Any, future::Future, sync::Arc, time::Duration}; + +use litellm_circuit_breaker::{State, Transition}; + +use crate::Error; + +/// What the host attached to the calls made inside `tagged`, opaque to the client. +pub type CallTag = Arc; + +tokio::task_local! { + static TAG: CallTag; +} + +pub async fn tagged(tag: CallTag, operation: F) -> F::Output { + TAG.scope(tag, operation).await +} + +/// The tag of the `tagged` scope this task runs in. `RedisEvents::call` runs on the calling task, +/// so a sink reads its call's tag here. +pub fn current_tag() -> Option { + TAG.try_with(Arc::clone).ok() +} + +pub trait RedisEvents: Send + Sync { + fn call(&self, event: &CallEvent); + + fn breaker(&self, transition: &Transition); +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Operation { + Command(&'static str), + Script, + Pipeline { ops: usize }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Outcome { + Succeeded, + Failed(FailureKind), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FailureKind { + Timeout, + Transport, + Server, + Decode, + BreakerOpen, +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct CallEvent { + pub operation: Operation, + pub duration: Duration, + pub outcome: Outcome, +} + +impl FailureKind { + pub fn of(error: &Error) -> Self { + match error { + Error::Transport(error) if error.is_timeout() => Self::Timeout, + Error::Transport(_) => Self::Transport, + Error::Server(_) => Self::Server, + Error::Decode(_) => Self::Decode, + Error::BreakerOpen => Self::BreakerOpen, + } + } +} + +impl Outcome { + pub(crate) fn of(result: &Result) -> Self { + match result { + Ok(_) => Self::Succeeded, + Err(error) => Self::Failed(FailureKind::of(error)), + } + } +} diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs index 47ff8677e06..2754666c329 100644 --- a/litellm-rust/crates/redis-client/src/lib.rs +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -1,20 +1,30 @@ +mod batch; mod client; -mod commands; +pub mod commands; mod config; mod connection; mod error; +mod events; mod namespace; mod pipeline; +mod pubsub; mod script; -#[cfg(test)] -mod test_support; -pub use client::{RedisClient, Ttl}; -pub use config::{ - ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, RedisConfig, +pub use batch::{Declared, RedisBatch}; +pub use client::{RedisClient, Ttl, WeakRedisClient}; +pub use commands::{Amount, Command, Expiry}; +pub use config::{Credentials, CredentialsProvider, DEFAULT_TIMEOUT, RedisConfig, Timeouts}; +pub use error::{ConnectError, Error, KeyError}; +pub use events::{ + CallEvent, CallTag, FailureKind, Operation, Outcome, RedisEvents, current_tag, tagged, +}; +pub use litellm_redis_types::{ + Cluster, Node, RedisUrl, Sentinel, Standalone, StaticAuth, Target, TlsOptions, UnixSocket, }; -pub use error::{ConnectError, Error}; pub use namespace::Namespace; pub use pipeline::Pipeline; -pub use redis::Value; +pub use redis::{ + BasicAuth, ErrorKind as RedisErrorKind, RedisError, StreamingCredentialsProvider, Value, + cluster_routing::Slot, +}; pub use script::Script; diff --git a/litellm-rust/crates/redis-client/src/pipeline.rs b/litellm-rust/crates/redis-client/src/pipeline.rs index e54c65e9605..0b6bb6a6673 100644 --- a/litellm-rust/crates/redis-client/src/pipeline.rs +++ b/litellm-rust/crates/redis-client/src/pipeline.rs @@ -1,30 +1,20 @@ -use std::{collections::HashMap, time::Duration}; +use std::collections::HashMap; use futures_util::future::join_all; -use redis::{ - Cmd, ErrorKind, RedisResult, ServerErrorKind, ToRedisArgs, ToSingleRedisArg, Value, - cluster_routing::Slot, -}; +use redis::{ErrorKind, RedisError, RedisResult, ServerErrorKind, Value, cluster_routing::Slot}; -use crate::{Error, Namespace, Script, commands, connection::Connection}; +use crate::{ + Command, Error, Namespace, + commands::{Op, Plan}, + connection::Connection, + script::ScriptCall, +}; pub struct Pipeline { namespace: Namespace, ops: Vec, } -struct Op { - command: Cmd, - key: String, - script: Option, -} - -struct ScriptCall { - script: Script, - keys: Vec, - args: Vec>, -} - type Settled = (usize, Result); impl Pipeline { @@ -35,60 +25,13 @@ impl Pipeline { } } - pub fn get(self, key: &str) -> Self { - let key = self.namespace.apply(key); - self.command(commands::get(key.clone()), key) + pub(crate) fn of(namespace: Namespace, ops: Vec) -> Self { + Self { namespace, ops } } - pub fn set(self, key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Self { - let key = self.namespace.apply(key); - self.command(commands::set(key.clone(), value, ttl), key) - } - - pub fn del(self, key: &str) -> Self { - let key = self.namespace.apply(key); - self.command(commands::del(vec![key.clone()]), key) - } - - pub fn expire(self, key: &str, ttl: Duration) -> Self { - let key = self.namespace.apply(key); - self.command(commands::expire(key.clone(), ttl), key) - } - - pub fn incr_by(self, key: &str, delta: i64) -> Self { - let key = self.namespace.apply(key); - self.command(commands::incr_by(key.clone(), delta), key) - } - - pub fn incr_by_float(self, key: &str, delta: f64) -> Self { - let key = self.namespace.apply(key); - self.command(commands::incr_by_float(key.clone(), delta), key) - } - - pub fn eval(self, script: &Script, keys: &[&str], args: impl ToRedisArgs) -> Self { - let keys: Vec = keys.iter().map(|key| self.namespace.apply(key)).collect(); - let args = args.to_redis_args(); - let op = Op { - command: script.evalsha(&keys, &args), - key: keys.first().cloned().unwrap_or_default(), - script: Some(ScriptCall { - script: script.clone(), - keys, - args, - }), - }; - self.push(op) - } - - fn command(self, command: Cmd, key: String) -> Self { - self.push(Op { - command, - key, - script: None, - }) - } - - fn push(mut self, op: Op) -> Self { + /// Appends one command in its pipelined form; `RedisClient::run` answers it in order, undecoded. + pub fn with(mut self, command: Command) -> Self { + let (op, _) = command.prepare(&self.namespace); self.ops.push(op); self } @@ -102,7 +45,7 @@ impl Pipeline { let settled = join_all( groups .iter() - .map(|group| self.run_group(connection.clone(), group)), + .map(|batches| self.run_in_order(connection, batches)), ) .await; let mut settled: Vec = settled.into_iter().flatten().collect(); @@ -110,183 +53,131 @@ impl Pipeline { settled.into_iter().map(|(_, result)| result).collect() } - fn groups(&self, connection: &Connection) -> Vec> { + async fn run_in_order(&self, connection: &Connection, batches: &[Vec]) -> Vec { + let mut settled = Vec::with_capacity(batches.iter().map(Vec::len).sum()); + for batch in batches { + settled.extend(self.run_group(connection.clone(), batch).await); + } + settled + } + + fn groups(&self, connection: &Connection) -> Vec>> { match connection { _ if self.ops.is_empty() => Vec::new(), - Connection::Standalone(_) => vec![(0..self.ops.len()).collect()], + Connection::Standalone(_) | Connection::Sentinel(_) => { + vec![vec![(0..self.ops.len()).collect()]] + } Connection::Cluster(_) => { let mut by_slot: HashMap> = HashMap::new(); - for (index, op) in self.ops.iter().enumerate() { + for index in 0..self.ops.len() { by_slot - .entry(Slot::for_key(&op.key)) + .entry(Slot::for_key(&self.ops[index].key)) .or_default() .push(index); } - by_slot.into_values().collect() + by_slot + .into_values() + .map(|indices| self.retry_safe_batches(indices)) + .collect() } } } + fn retry_safe_batches(&self, indices: Vec) -> Vec> { + indices + .chunk_by(|&left, &right| self.ops[left].idempotent && self.ops[right].idempotent) + .map(<[usize]>::to_vec) + .collect() + } + + /// A script the server had not cached ran nothing, so it alone runs again, the way litellm's + /// `RedisBatch` reruns it through redis-py's `Script`: `EVALSHA`, then `SCRIPT LOAD` and + /// `EVALSHA` on `NOSCRIPT`. + async fn settle_script( + &self, + connection: Connection, + index: usize, + reply: RedisResult, + ) -> (usize, RedisResult) { + match (&self.ops[index].plan, reply) { + (Plan::Script(call), Err(error)) if is_missing_script(&error) => { + (index, rerun_missing_script(connection, call).await) + } + (_, reply) => (index, reply), + } + } + async fn run_group(&self, mut connection: Connection, indices: &[usize]) -> Vec { + let deployment = connection.deployment(); let mut pipeline = redis::pipe(); pipeline.ignore_errors(); - for &index in indices { - pipeline.add_command(self.ops[index].command.clone()); - } - let replies: RedisResult>> = - pipeline.query_async(&mut connection).await; + let widths: Vec = indices + .iter() + .map(|&index| { + let commands = self.ops[index].plan.pipelined(); + let width = commands.len(); + for command in commands { + pipeline.add_command(command); + } + width + }) + .collect(); + let replies: RedisResult>> = match pipeline.is_empty() { + true => Ok(Vec::new()), + false => pipeline.query_async(&mut connection).await, + }; let replies = match replies { Ok(replies) => replies, Err(error) => { return indices .iter() - .map(|&index| (index, Err(Error::from(error.clone())))) + .map(|&index| (index, Err(Error::classify(error.clone(), deployment)))) .collect(); } }; - let mut settled = Vec::with_capacity(indices.len()); - for (&index, reply) in indices.iter().zip(replies) { - settled.push((index, self.settle(&mut connection, index, reply).await)); - } - settled - } - - async fn settle( - &self, - connection: &mut Connection, - index: usize, - reply: RedisResult, - ) -> Result { - match (reply, &self.ops[index].script) { - (Err(error), Some(call)) - if error.kind() == ErrorKind::Server(ServerErrorKind::NoScript) => - { - Ok(call - .script - .invocation(&call.keys, &call.args) - .invoke_async(connection) - .await?) - } - (reply, _) => Ok(reply?), - } + let mut replies = replies.into_iter(); + let per_op: Vec<(usize, RedisResult)> = indices + .iter() + .zip(widths) + .map(|(&index, width)| match &self.ops[index].plan { + Plan::Reply(value) => (index, Ok(value.clone())), + _ => (index, first_reply(replies.by_ref().take(width))), + }) + .collect(); + join_all( + per_op + .into_iter() + .map(|(index, reply)| self.settle_script(connection.clone(), index, reply)), + ) + .await + .into_iter() + .map(|(index, reply)| { + ( + index, + reply.map_err(|error| Error::classify(error, deployment)), + ) + }) + .collect() } } -#[cfg(test)] -mod tests { - use std::time::{Duration, Instant}; - - use redis::Value; - use rstest::rstest; - - use crate::{ - DEFAULT_TIMEOUT, Error, Script, - test_support::{Server, Topology, flush_scripts}, - }; - - #[rstest] - #[tokio::test] - async fn each_op_settles_on_its_own_in_input_order( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - client - .set("text", "not a number", Duration::from_secs(60)) - .await - .unwrap(); - let pipeline = client - .pipeline() - .incr_by("a", 1) - .incr_by("text", 1) - .incr_by("c", 5) - .set("b", "2", Duration::from_secs(60)) - .get("b") - .incr_by("a", 1) - .get("missing"); - - let results = client.run(&pipeline).await; - - assert_eq!(results.len(), 7); - assert_eq!(results[0].as_ref().unwrap(), &Value::Int(1)); - assert!( - matches!(results[1], Err(Error::Server(_))), - "{:?}", - results[1] - ); - assert_eq!(results[2].as_ref().unwrap(), &Value::Int(5)); - assert_eq!(results[3].as_ref().unwrap(), &Value::Okay); - assert_eq!( - results[4].as_ref().unwrap(), - &Value::BulkString(b"2".to_vec()) - ); - assert_eq!(results[5].as_ref().unwrap(), &Value::Int(2)); - assert_eq!(results[6].as_ref().unwrap(), &Value::Nil); - } - - #[rstest] - #[tokio::test] - async fn a_flushed_script_in_a_pipeline_reruns_alone( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("ns", DEFAULT_TIMEOUT).await; - let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); - client.eval::(&incr, &["counter"], 1).await.unwrap(); - flush_scripts(&client).await; - let pipeline = client - .pipeline() - .eval(&incr, &["counter"], 2) - .incr_by("other", 7); - - let results = client.run(&pipeline).await; - - assert_eq!(results[0].as_ref().unwrap(), &Value::Int(3)); - assert_eq!(results[1].as_ref().unwrap(), &Value::Int(7)); - assert_eq!(client.get::("counter").await.unwrap(), Some(3)); - } - - #[rstest] - #[tokio::test] - async fn an_unreachable_server_fails_every_op_as_transport( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_millis(500)).await; - let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1); - - server.freeze().await; - let results = client.run(&pipeline).await; - server.thaw().await; - - assert_eq!(results.len(), 2); - assert!( - matches!(results[0], Err(Error::Transport(_))), - "{:?}", - results[0] - ); - assert!( - matches!(results[1], Err(Error::Transport(_))), - "{:?}", - results[1] - ); - } - - #[rstest] - #[tokio::test] - async fn an_empty_pipeline_returns_without_asking_redis( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", Duration::from_secs(2)).await; - - server.freeze().await; - let started = Instant::now(); - let results = client.run(&client.pipeline()).await; - let elapsed = started.elapsed(); - server.thaw().await; - - assert!(results.is_empty()); - assert!(elapsed < Duration::from_secs(1), "{elapsed:?}"); - } +async fn rerun_missing_script(mut connection: Connection, call: &ScriptCall) -> RedisResult { + call.script + .invocation(&call.keys, &call.args) + .invoke_async(&mut connection) + .await +} + +fn is_missing_script(error: &RedisError) -> bool { + error.kind() == ErrorKind::Server(ServerErrorKind::NoScript) +} + +// RedisBatch settles an op from its first reply and fails it on any reply that failed +fn first_reply(replies: impl Iterator>) -> RedisResult { + replies + .reduce(|first, next| match (first, next) { + (Err(error), _) | (Ok(_), Err(error)) => Err(error), + (first, Ok(_)) => first, + }) + .unwrap_or(Ok(Value::Nil)) } diff --git a/litellm-rust/crates/redis-client/src/pubsub.rs b/litellm-rust/crates/redis-client/src/pubsub.rs new file mode 100644 index 00000000000..b6b18d3e481 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/pubsub.rs @@ -0,0 +1,225 @@ +use std::{sync::Arc, time::Duration}; + +use futures_util::{StreamExt, stream::BoxStream}; +use redis::{ConnectionInfo, RedisResult, aio::PubSubStream}; + +use crate::{ + ConnectError, Credentials, Standalone, Target, Timeouts, config, + connection::{Connection, SentinelConnection}, +}; + +// litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py: _BACKOFF_INITIAL_SECONDS, _BACKOFF_MAX_SECONDS +const BACKOFF_INITIAL: Duration = Duration::from_secs(5); +const BACKOFF_MAX: Duration = Duration::from_secs(60); + +/// A subscription the way redis-py's `PubSub` holds one: a dedicated RESP2 connection that sends +/// `SUBSCRIBE`, authenticated once when it connects. +#[derive(Clone)] +pub(crate) struct Subscriber { + source: Source, + credentials: Credentials, + connect_timeout: Duration, +} + +#[derive(Clone)] +enum Source { + Nodes(Arc<[ConnectionInfo]>), + Sentinel(SentinelConnection), +} + +impl Subscriber { + pub(crate) fn new( + endpoint: &Target, + credentials: &Credentials, + timeouts: Timeouts, + keepalive: bool, + connection: &Connection, + ) -> Result { + let source = match (endpoint, connection) { + (_, Connection::Sentinel(sentinel)) => Source::Sentinel(sentinel.clone()), + (Target::Cluster(cluster), _) => Source::Nodes( + cluster + .nodes + .iter() + .map(|node| { + let node = Target::Standalone(Standalone { + host: node.host.clone(), + port: node.port, + tls: cluster.tls.clone(), + db: 0, + }); + config::connection_info(&node, credentials, keepalive) + }) + .collect::>()?, + ), + _ => Source::Nodes(Arc::new([config::connection_info( + endpoint, + credentials, + keepalive, + )?])), + }; + Ok(Self { + source, + credentials: credentials.clone(), + connect_timeout: timeouts.connect, + }) + } + + pub(crate) fn messages(self, channel: String) -> BoxStream<'static, Vec> { + futures_util::stream::unfold( + (self, channel, None::, BACKOFF_INITIAL), + |(subscriber, channel, live, backoff)| async move { + let mut live = live; + let mut backoff = backoff; + loop { + let Some(current) = live.as_mut() else { + match subscriber.subscribe(&channel).await { + Ok(subscribed) => { + live = Some(subscribed); + backoff = BACKOFF_INITIAL; + } + Err(error) => { + tracing::warn!(%error, ?backoff, "Redis subscribe failed, retrying"); + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(BACKOFF_MAX); + } + } + continue; + }; + match current.next().await { + Some(message) => { + let payload = message.get_payload_bytes().to_vec(); + return Some((payload, (subscriber, channel, live, backoff))); + } + None => live = None, + } + tracing::warn!(?backoff, "Redis subscription connection lost, resubscribing"); + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(BACKOFF_MAX); + } + }, + ) + .boxed() + } + + async fn subscribe(&self, channel: &str) -> RedisResult { + match &self.source { + Source::Sentinel(sentinel) => { + let master = sentinel.master().await?; + self.subscribe_on(master.get_connection_info().clone(), channel) + .await + } + Source::Nodes(nodes) => { + let mut last = None; + for node in nodes.iter() { + match self.subscribe_on(node.clone(), channel).await { + Ok(live) => return Ok(live), + Err(error) => last = Some(error), + } + } + Err(last.expect("an endpoint always has at least one node")) + } + } + } + + async fn subscribe_on(&self, info: ConnectionInfo, channel: &str) -> RedisResult { + let settings = + config::authenticated(info.redis_settings().clone(), &self.credentials).await?; + let client = redis::Client::open(info.set_redis_settings(settings))?; + let mut pubsub = tokio::time::timeout(self.connect_timeout, client.get_async_pubsub()) + .await + .map_err(|_| { + redis::RedisError::from(std::io::Error::new( + std::io::ErrorKind::TimedOut, + "connecting the Redis subscription timed out", + )) + })??; + subscribe_connected(&mut pubsub, channel, self.connect_timeout).await?; + Ok(pubsub.into_on_message()) + } +} + +async fn subscribe_connected( + pubsub: &mut redis::aio::PubSub, + channel: &str, + connect_timeout: Duration, +) -> RedisResult<()> { + tokio::time::timeout(connect_timeout, pubsub.subscribe(channel)) + .await + .map_err(|_| { + redis::RedisError::from(std::io::Error::new( + std::io::ErrorKind::TimedOut, + "subscribing to the Redis channel timed out", + )) + })? +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use futures_util::StreamExt; + use litellm_redis_testing::{RedisServer, TcpProxy}; + use rstest::rstest; + + use super::subscribe_connected; + + #[rstest] + #[tokio::test] + async fn a_stalled_pubsub_subscribe_times_out_and_a_fresh_pubsub_recovers() { + let server = RedisServer::start().await.unwrap(); + let proxy = TcpProxy::start(server.port()).await.unwrap(); + let client = redis::Client::open(format!("redis://127.0.0.1:{}", proxy.port())).unwrap(); + let mut pubsub = tokio::time::timeout(Duration::from_secs(2), client.get_async_pubsub()) + .await + .expect("connecting the initial pubsub should finish") + .unwrap(); + tokio::time::timeout( + Duration::from_secs(2), + subscribe_connected(&mut pubsub, "existing", Duration::from_secs(1)), + ) + .await + .expect("the initial subscription acknowledgement should arrive") + .unwrap(); + + proxy.blackhole_open_connections(); + let result = tokio::time::timeout( + Duration::from_secs(1), + subscribe_connected(&mut pubsub, "stalled", Duration::from_millis(100)), + ) + .await + .expect("the stalled subscription should return before the safety deadline"); + assert!( + matches!(&result, Err(error) if error.is_timeout()), + "{result:?}" + ); + drop(pubsub); + + let mut recovered = tokio::time::timeout(Duration::from_secs(2), client.get_async_pubsub()) + .await + .expect("a fresh pubsub connection should finish") + .unwrap(); + tokio::time::timeout( + Duration::from_secs(2), + subscribe_connected(&mut recovered, "recovered", Duration::from_secs(1)), + ) + .await + .expect("the fresh subscription acknowledgement should arrive") + .unwrap(); + + let publisher = redis::Client::open(server.url()).unwrap(); + let mut connection = publisher.get_multiplexed_async_connection().await.unwrap(); + let subscriber_count: i64 = redis::cmd("PUBLISH") + .arg("recovered") + .arg("payload") + .query_async(&mut connection) + .await + .unwrap(); + assert_eq!(subscriber_count, 1); + let message = tokio::time::timeout(Duration::from_secs(2), recovered.on_message().next()) + .await + .expect("the fresh subscription should receive the published payload") + .unwrap(); + assert_eq!(message.get_payload::().unwrap(), "payload"); + } +} diff --git a/litellm-rust/crates/redis-client/src/script.rs b/litellm-rust/crates/redis-client/src/script.rs index c5d73133ecc..6dbdbafe9c9 100644 --- a/litellm-rust/crates/redis-client/src/script.rs +++ b/litellm-rust/crates/redis-client/src/script.rs @@ -1,13 +1,18 @@ +use std::sync::Arc; + use redis::{Cmd, ScriptInvocation}; +/// A Lua script, cheap to clone into every call that runs it. #[derive(Clone)] -pub struct Script(redis::Script); +pub struct Script(Arc); impl Script { pub fn new(source: &str) -> Self { - Self(redis::Script::new(source)) + Self(Arc::new(redis::Script::new(source))) } + /// The `EVALSHA` a pipeline sends, like redis-py. A `NOSCRIPT` reply is settled by + /// [`Script::invocation`] alone, which loads the script and sends it again. pub(crate) fn evalsha(&self, keys: &[String], args: &[Vec]) -> Cmd { redis::cmd("EVALSHA") .arg(self.0.get_hash()) @@ -24,65 +29,10 @@ impl Script { } } -#[cfg(test)] -mod tests { - use std::time::Duration; - - use rstest::rstest; - - use super::Script; - use crate::{ - DEFAULT_TIMEOUT, Error, - test_support::{Server, Topology, flush_scripts}, - }; - - #[rstest] - #[tokio::test] - async fn script_keys_are_namespaced( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("a", DEFAULT_TIMEOUT).await; - let bare = server.client("", DEFAULT_TIMEOUT).await; - let read = Script::new("return redis.call('GET', KEYS[1])"); - bare.set("a:x", "1", Duration::from_secs(60)).await.unwrap(); - - let value: Option = client - .eval(&read, &["x"], Vec::::new()) - .await - .unwrap(); - - assert_eq!(value.as_deref(), Some("1")); - } - - #[rstest] - #[tokio::test] - async fn a_failing_script_is_a_server_error_and_keeps_the_connection( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); - client.rpush("list", "item").await.unwrap(); - - let failed = client.eval::(&incr, &["list"], 1).await; - - assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}"); - assert_eq!(client.eval::(&incr, &["counter"], 3).await.unwrap(), 3); - } - - #[rstest] - #[tokio::test] - async fn a_flushed_script_is_loaded_again( - #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, - ) { - let server = Server::start(topology).await; - let client = server.client("", DEFAULT_TIMEOUT).await; - let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); - client.eval::(&incr, &["counter"], 1).await.unwrap(); - - flush_scripts(&client).await; - - assert_eq!(client.eval::(&incr, &["counter"], 1).await.unwrap(), 2); - } +/// A script call in a pipeline, kept apart from plain commands because it is the one entry a +/// pipeline may send again. +pub(crate) struct ScriptCall { + pub(crate) script: Script, + pub(crate) keys: Vec, + pub(crate) args: Vec>, } diff --git a/litellm-rust/crates/redis-client/src/test_support.rs b/litellm-rust/crates/redis-client/src/test_support.rs deleted file mode 100644 index e503739a942..00000000000 --- a/litellm-rust/crates/redis-client/src/test_support.rs +++ /dev/null @@ -1,126 +0,0 @@ -use std::time::Duration; - -use litellm_circuit_breaker::Breaker; -use litellm_redis_testing::{RedisCluster, RedisServer}; -use litellm_secret_value::SecretValue; -use redis::cluster_routing::{Route, RoutingInfo, SingleNodeRoutingInfo, SlotAddr}; - -use crate::{ - ClusterEndpoint, ClusterNode, Endpoint, Namespace, RedisClient, RedisConfig, - connection::Connection, -}; - -#[derive(Clone, Copy, Debug)] -pub(crate) enum Topology { - Standalone, - Cluster, -} - -pub(crate) enum Server { - Standalone(RedisServer), - Cluster(RedisCluster), -} - -impl Server { - pub(crate) async fn start(topology: Topology) -> Self { - match topology { - Topology::Standalone => Self::Standalone(RedisServer::start().await.unwrap()), - Topology::Cluster => Self::Cluster(RedisCluster::start().await.unwrap()), - } - } - - pub(crate) fn endpoint(&self) -> Endpoint { - match self { - Self::Standalone(server) => Endpoint::Url(SecretValue::new(server.url().to_owned())), - Self::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { - nodes: cluster - .nodes() - .map(|(host, port)| ClusterNode { - host: host.to_owned(), - port, - }) - .collect(), - username: None, - password: None, - tls: false, - }), - } - } - - pub(crate) async fn client(&self, namespace: &str, timeout: Duration) -> RedisClient { - self.guarded_client(namespace, timeout, None).await - } - - pub(crate) async fn guarded_client( - &self, - namespace: &str, - timeout: Duration, - breaker: Option, - ) -> RedisClient { - RedisClient::new(RedisConfig { - endpoint: self.endpoint(), - namespace: Namespace::new(namespace), - timeout, - breaker, - }) - .unwrap() - } - - pub(crate) async fn restart(&self) { - match self { - Self::Standalone(server) => server.restart().await.unwrap(), - Self::Cluster(cluster) => cluster.restart().await.unwrap(), - } - } - - pub(crate) async fn stop(&self) { - match self { - Self::Standalone(server) => server.stop().await.unwrap(), - Self::Cluster(cluster) => cluster.stop().await.unwrap(), - } - } - - pub(crate) async fn resume(&self) { - match self { - Self::Standalone(server) => server.resume().await.unwrap(), - Self::Cluster(cluster) => cluster.resume().await.unwrap(), - } - } - - pub(crate) async fn freeze(&self) { - match self { - Self::Standalone(server) => server.freeze().await.unwrap(), - Self::Cluster(cluster) => cluster.freeze().await.unwrap(), - } - } - - pub(crate) async fn thaw(&self) { - match self { - Self::Standalone(server) => server.thaw().await.unwrap(), - Self::Cluster(cluster) => cluster.thaw().await.unwrap(), - } - } -} - -pub(crate) async fn client_id(client: &RedisClient, key: &str) -> i64 { - let command = redis::cmd("CLIENT").arg("ID").to_owned(); - match client.connection.clone() { - Connection::Standalone(mut connection) => { - command.query_async(&mut connection).await.unwrap() - } - Connection::Cluster(mut connection) => { - let node = Route::with_key(key, SlotAddr::Master); - let routing = RoutingInfo::SingleNode(SingleNodeRoutingInfo::SpecificNode(node)); - redis::from_redis_value(connection.route_command(command, routing).await.unwrap()) - .unwrap() - } - } -} - -pub(crate) async fn flush_scripts(client: &RedisClient) { - redis::cmd("SCRIPT") - .arg("FLUSH") - .query_async::<()>(&mut client.connection.clone()) - .await - .unwrap(); -} diff --git a/litellm-rust/crates/redis-client/tests/batch.rs b/litellm-rust/crates/redis-client/tests/batch.rs new file mode 100644 index 00000000000..45ca419951e --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/batch.rs @@ -0,0 +1,330 @@ +mod support; + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_redis_client::commands; +use litellm_redis_client::{ + Amount, CallEvent, Credentials, DEFAULT_TIMEOUT, Declared, Error, Expiry, Namespace, Operation, + RedisBatch, RedisClient, RedisConfig, RedisEvents, Script, Timeouts, Ttl, +}; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; +use support::{connect, endpoint}; + +const MINUTE: Option = Some(Expiry::KeepExisting(Duration::from_secs(60))); + +// RedisBatch pipelines a standalone flush and runs each op on its own on a cluster +fn expected_trips(topology: Topology, alone: &[Operation]) -> Vec { + match topology { + Topology::Standalone => vec![Operation::Pipeline { ops: alone.len() }], + Topology::Cluster => alone.to_vec(), + } +} + +#[derive(Default)] +struct RoundTrips(Mutex>); + +impl RoundTrips { + fn seen(&self) -> Vec { + self.0.lock().unwrap().clone() + } +} + +impl RedisEvents for RoundTrips { + fn call(&self, event: &CallEvent) { + self.0.lock().unwrap().push(event.operation); + } + + fn breaker(&self, _transition: &Transition) {} +} + +fn observed(server: &RedisDeployment) -> (RedisClient, Arc) { + let trips = Arc::new(RoundTrips::default()); + let client = RedisClient::new(RedisConfig { + target: endpoint(server), + credentials: Credentials::NONE, + namespace: Namespace::new("ns"), + timeouts: Timeouts::default(), + breaker: None, + events: Some(trips.clone()), + health_check: None, + keepalive: true, + }) + .unwrap(); + (client, trips) +} + +#[rstest] +#[tokio::test] +async fn declared_ops_wait_for_the_flush_and_share_one_round_trip( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = client.batch(); + let counter = batch.add(commands::incr_by_float("counter", Amount::Int(3), MINUTE)); + let total = batch.add(commands::incr_by_float("total", Amount::Float(1.5), MINUTE)); + let written = batch.add(commands::set("other", "v", Duration::from_secs(60))); + let read = batch.add(commands::get::("value")); + + let outside = connect(&server, "ns", DEFAULT_TIMEOUT); + assert_eq!( + outside.send(commands::get::("counter")).await.unwrap(), + None + ); + assert!(trips.seen().is_empty()); + + batch.flush().await; + + let expected = expected_trips( + topology, + &[ + Operation::Command("INCRBYFLOAT"), + Operation::Command("INCRBYFLOAT"), + Operation::Command("SET"), + Operation::Command("GET"), + ], + ); + let mut seen = trips.seen(); + seen.sort_by_key(|operation| format!("{operation:?}")); + let mut sorted = expected.clone(); + sorted.sort_by_key(|operation| format!("{operation:?}")); + assert_eq!(seen, sorted); + assert_eq!(counter.resolve().await.unwrap(), 3.0); + assert_eq!(total.resolve().await.unwrap(), 1.5); + written.resolve().await.unwrap(); + assert_eq!(read.resolve().await.unwrap(), None); + assert_eq!(trips.seen().len(), expected.len()); +} + +#[rstest] +#[tokio::test] +async fn resolving_before_a_flush_sends_everything_declared_so_far( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = client.batch(); + let first = batch.add(commands::incr_by_float("a", Amount::Int(1), MINUTE)); + let second = batch.add(commands::incr_by_float("b", Amount::Int(2), MINUTE)); + let increment = Operation::Command("INCRBYFLOAT"); + + assert_eq!(second.resolve().await.unwrap(), 2.0); + let late = batch.add(commands::incr_by_float("a", Amount::Int(10), MINUTE)); + assert_eq!(first.resolve().await.unwrap(), 1.0); + assert_eq!( + trips.seen(), + expected_trips(topology, &[increment, increment]) + ); + + assert_eq!(late.resolve().await.unwrap(), 11.0); + let both: Vec = expected_trips(topology, &[increment, increment]) + .into_iter() + .chain(expected_trips(topology, &[increment])) + .collect(); + assert_eq!(trips.seen(), both); + batch.flush().await; + assert_eq!(trips.seen().len(), both.len()); +} + +#[rstest] +#[tokio::test] +async fn a_pipelined_batch_is_one_pipeline_on_every_topology( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = RedisBatch::pipelined(client); + let declared = + ["a", "b", "c"].map(|key| batch.add(commands::incr_by_float(key, Amount::Int(2), MINUTE))); + + let totals = Declared::all(declared).await; + + assert_eq!( + totals.into_iter().map(Result::unwrap).collect::>(), + [2.0; 3] + ); + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 3 }]); +} + +#[rstest] +#[tokio::test] +async fn each_declared_op_gets_its_own_typed_result( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set( + "text", + "not a number", + Duration::from_secs(60), + )) + .await + .unwrap(); + let batch = client.batch(); + let rejected = batch.add(commands::incr_by_float("text", Amount::Int(1), MINUTE)); + let undecodable = batch.add(commands::get::("text")); + let missing = batch.add(commands::get::("missing")); + let script = batch.add(commands::eval::>( + &Script::new("return {KEYS[1], ARGV[1]}"), + &["k"], + "arg", + )); + let removed = batch.add(commands::del(&["text"])); + + batch.flush().await; + + assert!(matches!(rejected.resolve().await, Err(Error::Server(_)))); + assert!(matches!(undecodable.resolve().await, Err(Error::Decode(_)))); + assert_eq!(missing.resolve().await.unwrap(), None); + assert_eq!(script.resolve().await.unwrap(), ["k", "arg"]); + assert_eq!(removed.resolve().await.unwrap(), 1); +} + +#[rstest] +#[tokio::test] +async fn an_op_whose_flush_was_cancelled_resolves_as_transport( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_secs(5)); + let batch = client.batch(); + let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), MINUTE)); + + server.freeze().await.unwrap(); + let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await; + server.thaw().await.unwrap(); + + assert!(cancelled.is_err()); + assert!(matches!(counter.resolve().await, Err(Error::Transport(_)))); +} + +#[rstest] +#[tokio::test] +async fn a_settled_op_waits_for_someone_else_to_flush_and_sends_nothing_itself( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = observed(&server); + let batch = client.batch(); + let counter = tokio::spawn( + batch + .add(commands::incr_by_float("counter", Amount::Int(4), MINUTE)) + .settled(), + ); + + tokio::time::sleep(Duration::from_millis(200)).await; + assert!(!counter.is_finished()); + assert!(trips.seen().is_empty()); + + batch.flush().await; + + assert_eq!(counter.await.unwrap().unwrap(), 4.0); + assert_eq!( + trips.seen(), + expected_trips(topology, &[Operation::Command("INCRBYFLOAT")]) + ); +} + +#[rstest] +#[tokio::test] +async fn a_settled_op_whose_flush_was_cancelled_resolves_as_transport() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let client = connect(&server, "", Duration::from_secs(5)); + let batch = client.batch(); + let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), MINUTE)); + + server.freeze().await.unwrap(); + let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await; + server.thaw().await.unwrap(); + + assert!(cancelled.is_err()); + assert!(matches!(counter.settled().await, Err(Error::Transport(_)))); +} + +#[rstest] +#[tokio::test] +async fn clones_are_the_same_client_and_a_second_connection_is_not() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let twin = connect(&server, "", DEFAULT_TIMEOUT); + + assert!(client.same_client(&client.clone())); + assert!(!client.same_client(&twin)); +} + +#[rstest] +#[tokio::test] +async fn mget_reads_each_key_once_and_answers_every_key_in_order( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("a", "1", Duration::from_secs(60))) + .await + .unwrap(); + client + .send(commands::set("c", "3", Duration::from_secs(60))) + .await + .unwrap(); + server + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + let batch = client.batch(); + + let values = batch.add(commands::mget::(&["c", "missing", "a", "c"])); + batch.flush().await; + + assert_eq!( + values.resolve().await.unwrap(), + [ + Some("3".to_owned()), + None, + Some("1".to_owned()), + Some("3".to_owned()) + ] + ); + let stats = support::command_stats(&server).await; + assert!(support::calls(&stats, "mget") > 0, "{stats}"); + assert_eq!(support::calls(&stats, "get"), 0, "{stats}"); + assert_eq!(support::keys_read(&server).await, 3); +} + +// RedisBatch's _Increment refreshes the TTL in a pipeline, while its cluster path calls +// async_increment, which keeps a TTL the key already has +#[rstest] +#[case::pipelined(Topology::Standalone, false)] +#[case::alone_on_a_cluster(Topology::Cluster, true)] +#[tokio::test] +async fn a_batched_increment_treats_the_ttl_like_redis_batch( + #[case] topology: Topology, + #[case] keeps_the_ttl: bool, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("counter", "1", Duration::from_secs(30))) + .await + .unwrap(); + let batch = client.batch(); + + let counter = batch.add(commands::incr_by_float( + "counter", + Amount::Int(1), + Some(Expiry::KeepExisting(Duration::from_secs(600))), + )); + batch.flush().await; + + assert_eq!(counter.resolve().await.unwrap(), 2.0); + let Ttl::Expires(left) = client.send(commands::ttl("counter")).await.unwrap() else { + panic!("expected an expiring counter") + }; + assert_eq!(left <= Duration::from_secs(30), keeps_the_ttl, "{left:?}"); +} diff --git a/litellm-rust/crates/redis-client/tests/client.rs b/litellm-rust/crates/redis-client/tests/client.rs new file mode 100644 index 00000000000..f95023ca6db --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/client.rs @@ -0,0 +1,1039 @@ +mod support; + +use std::{ + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +use litellm_redis_testing::{RedisCluster, RedisDeployment, RedisServer, Topology, eventually}; +use rstest::rstest; + +use litellm_circuit_breaker::{Breaker, Config, State}; + +use litellm_redis_client::commands; +use litellm_redis_client::{ + Amount, CallEvent, Credentials, DEFAULT_TIMEOUT, Error, Expiry, FailureKind, Namespace, + Operation, Outcome, RedisClient, RedisConfig, RedisEvents, RedisUrl, Script, Standalone, + Target, Timeouts, Ttl, Value, +}; +use litellm_secret_value::SecretValue; +use support::{KEEP, connect, connect_guarded, connections, endpoint}; + +async fn persist(client: &RedisClient, key: &str) { + let set = Script::new("return redis.call('SET', KEYS[1], '0')"); + client + .send(commands::eval::<()>(&set, &[key], Vec::::new())) + .await + .unwrap(); +} + +#[derive(Clone, Copy, Debug)] +enum Counter { + IncrBy, + IncrByFloat, +} + +const MINUTE: Duration = Duration::from_secs(60); + +async fn write(client: &RedisClient, counter: Counter, key: &str, expiry: Expiry) { + match (counter, expiry) { + (Counter::IncrBy, Expiry::KeepExisting(ttl)) => { + client.send(commands::incr_by(key, 1, ttl)).await.unwrap(); + } + (Counter::IncrBy, Expiry::Refresh(_)) => unreachable!("increment_cache never refreshes"), + (Counter::IncrByFloat, expiry) => { + client + .send(commands::incr_by_float( + key, + Amount::Float(0.5), + Some(expiry), + )) + .await + .unwrap(); + } + } +} + +async fn persist_as(client: &RedisClient, key: &str) { + let source = "return redis.call('SET', KEYS[1], '1')"; + client + .send(commands::eval::( + &Script::new(source), + &[key], + Vec::::new(), + )) + .await + .unwrap(); +} + +fn remaining(ttl: Ttl) -> Duration { + match ttl { + Ttl::Expires(left) => left, + other => panic!("expected an expiring key, got {other:?}"), + } +} + +#[rstest] +#[case::increment_cache(Counter::IncrBy, Expiry::KeepExisting(MINUTE))] +#[case::async_increment(Counter::IncrByFloat, Expiry::KeepExisting(MINUTE))] +#[case::async_increment_refreshing(Counter::IncrByFloat, Expiry::Refresh(MINUTE))] +#[tokio::test] +async fn a_counter_write_expires_a_new_key( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[case] counter: Counter, + #[case] expiry: Expiry, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + write(&client, counter, "fresh", expiry).await; + + let left = remaining(client.send(commands::ttl("fresh")).await.unwrap()); + assert!(left > Duration::from_secs(50) && left <= MINUTE, "{left:?}"); +} + +#[rstest] +#[tokio::test] +async fn keep_existing_leaves_a_running_ttl_and_refresh_resets_it( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let short = Expiry::Refresh(Duration::from_secs(30)); + let long = Duration::from_secs(600); + write(&client, Counter::IncrByFloat, "kept", short).await; + write(&client, Counter::IncrByFloat, "refreshed", short).await; + write( + &client, + Counter::IncrBy, + "kept_int", + Expiry::KeepExisting(Duration::from_secs(30)), + ) + .await; + + write( + &client, + Counter::IncrByFloat, + "kept", + Expiry::KeepExisting(long), + ) + .await; + write( + &client, + Counter::IncrByFloat, + "refreshed", + Expiry::Refresh(long), + ) + .await; + write( + &client, + Counter::IncrBy, + "kept_int", + Expiry::KeepExisting(long), + ) + .await; + + assert!( + remaining(client.send(commands::ttl("kept")).await.unwrap()) <= Duration::from_secs(30) + ); + assert!( + remaining(client.send(commands::ttl("kept_int")).await.unwrap()) <= Duration::from_secs(30) + ); + assert!( + remaining(client.send(commands::ttl("refreshed")).await.unwrap()) + > Duration::from_secs(500) + ); +} + +#[rstest] +#[tokio::test] +async fn keep_existing_expires_a_key_that_had_no_ttl( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values(Counter::IncrBy, Counter::IncrByFloat)] counter: Counter, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "ns", DEFAULT_TIMEOUT); + persist_as(&client, "persistent").await; + + write(&client, counter, "persistent", KEEP).await; + + let left = remaining(client.send(commands::ttl("persistent")).await.unwrap()); + assert!(left > Duration::from_secs(50), "{left:?}"); +} + +#[rstest] +#[tokio::test] +async fn counters_send_the_commands_redis_cache_sends( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + server + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + + client + .send(commands::incr_by("int", 1, MINUTE)) + .await + .unwrap(); + client + .send(commands::incr_by_float("kept", Amount::Int(1), Some(KEEP))) + .await + .unwrap(); + client + .send(commands::incr_by_float( + "refreshed", + Amount::Float(0.5), + Some(Expiry::Refresh(MINUTE)), + )) + .await + .unwrap(); + client + .run( + &client + .pipeline() + .with(commands::incr_by_float( + "piped", + Amount::Float(0.25), + Some(Expiry::Refresh(MINUTE)), + )) + .with(commands::incr_by_float("bare", Amount::Float(0.25), None)), + ) + .await; + + let stats = support::command_stats(&server).await; + let calls = |command: &str| support::calls(&stats, command); + assert_eq!(calls("incrby"), 1); + assert_eq!(calls("incrbyfloat"), 4); + assert_eq!(calls("ttl"), 2); + assert_eq!(calls("expire"), 4); + assert_eq!( + calls("eval") + calls("evalsha") + calls("pttl") + calls("pexpire"), + 0 + ); +} + +#[derive(Default)] +struct Recorded(Mutex>); + +#[rstest] +#[case::beyond_double_precision(9_007_199_254_740_992_i64, 1_i64)] +#[case::negative_beyond_double_precision(-9_007_199_254_740_992_i64, -1_i64)] +#[case::maximum_integer(i64::MAX - 1, 1_i64)] +#[case::minimum_integer(i64::MIN + 1, -1_i64)] +#[tokio::test] +async fn integer_increments_preserve_every_bit( + #[case] initial: i64, + #[case] delta: i64, + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("integer", initial, Duration::from_secs(60))) + .await + .unwrap(); + let expected = initial + delta; + let result = client + .send(commands::incr_by("integer", delta, MINUTE)) + .await + .unwrap(); + assert_eq!(result, expected); + assert_eq!( + client.send(commands::get::("integer")).await.unwrap(), + Some(expected) + ); +} + +#[rstest] +#[tokio::test] +async fn outage_events_report_the_failure_kind( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values(false, true)] timeout: bool, + #[values(false, true)] pipeline: bool, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let recorded = Arc::new(Recorded::default()); + let client = RedisClient::new(RedisConfig { + target: endpoint(&server), + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_millis(300), None), + breaker: None, + events: Some(recorded.clone()), + credentials: Credentials::NONE, + health_check: None, + keepalive: true, + }) + .unwrap(); + client.ping().await.unwrap(); + if timeout { + server.freeze().await.unwrap(); + } else { + server.stop().await.unwrap(); + } + let result = if pipeline { + client + .run(&client.pipeline().with(commands::get::("missing"))) + .await + .remove(0) + .map(drop) + } else { + client + .send(commands::get::("missing")) + .await + .map(drop) + }; + if timeout { + server.thaw().await.unwrap(); + } else { + server.resume().await.unwrap(); + } + assert!(matches!(result, Err(Error::Transport(_)))); + let expected = if timeout { + FailureKind::Timeout + } else { + FailureKind::Transport + }; + assert_eq!( + recorded.0.lock().unwrap().last().unwrap().outcome, + Outcome::Failed(expected) + ); +} + +impl RedisEvents for Recorded { + fn call(&self, event: &CallEvent) { + self.0.lock().unwrap().push(*event); + } + + fn breaker(&self, _transition: &litellm_circuit_breaker::Transition) {} +} + +#[rstest] +#[tokio::test] +async fn every_call_reports_its_operation_and_outcome( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let recorded = Arc::new(Recorded::default()); + let client = RedisClient::new(RedisConfig { + target: endpoint(&server), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: Some(recorded.clone()), + health_check: None, + keepalive: true, + }) + .unwrap(); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + + let _ = client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await; + let _ = client.send(commands::get::("missing")).await; + let _ = client + .run( + &client + .pipeline() + .with(commands::get::("a")) + .with(commands::get::("b")), + ) + .await; + let failures = client + .run(&client.pipeline().with(commands::incr_by_float( + "list", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + ))) + .await; + assert!(matches!(failures[0], Err(Error::Server(_)))); + + let seen: Vec<(Operation, Outcome)> = recorded + .0 + .lock() + .unwrap() + .iter() + .map(|event| (event.operation, event.outcome)) + .collect(); + assert_eq!( + seen, + vec![ + (Operation::Command("SET"), Outcome::Succeeded), + ( + Operation::Command("INCRBY"), + Outcome::Failed(FailureKind::Server) + ), + (Operation::Command("GET"), Outcome::Succeeded), + (Operation::Pipeline { ops: 2 }, Outcome::Succeeded), + ( + Operation::Pipeline { ops: 1 }, + Outcome::Failed(FailureKind::Server) + ), + ] + ); +} + +#[rstest] +#[tokio::test] +async fn namespaces_isolate_keys_on_the_server( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let a = connect(&server, "a", DEFAULT_TIMEOUT); + let b = connect(&server, "b", DEFAULT_TIMEOUT); + let bare = connect(&server, "", DEFAULT_TIMEOUT); + + a.send(commands::set("x", "1", Duration::from_secs(60))) + .await + .unwrap(); + + assert_eq!( + a.send(commands::get::("x")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + assert_eq!( + a.send(commands::get::("a:x")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + assert_eq!( + bare.send(commands::get::("a:x")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + assert_eq!(b.send(commands::get::("x")).await.unwrap(), None); +} + +#[rstest] +#[tokio::test] +async fn a_rejected_command_keeps_the_connection( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + let before = connections(&server).await; + + let rejected = client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await; + + assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}"); + assert_eq!( + client + .send(commands::get::("list")) + .await + .unwrap() + .as_deref(), + Some("item") + ); + assert_eq!(connections(&server).await, before); +} + +#[rstest] +#[tokio::test] +async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("float", "15.0", Duration::from_secs(60))) + .await + .unwrap(); + + let decoded = client.send(commands::get::("float")).await; + + assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}"); + assert_eq!( + client.send(commands::get::("missing")).await.unwrap(), + None + ); +} + +#[rstest] +#[tokio::test] +async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + client + .send(commands::set("first", "1", Duration::from_secs(60))) + .await + .unwrap(); + client + .send(commands::set("second", "2", Duration::from_secs(60))) + .await + .unwrap(); + + server.freeze().await.unwrap(); + let timed_out = client.send(commands::get::("first")).await; + server.thaw().await.unwrap(); + + assert!( + matches!(timed_out, Err(Error::Transport(_))), + "{timed_out:?}" + ); + assert_eq!( + client + .send(commands::get::("second")) + .await + .unwrap() + .as_deref(), + Some("2") + ); +} + +#[rstest] +#[tokio::test] +async fn the_same_client_recovers_after_redis_restarts( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + + server.restart().await.unwrap(); + + let recovered = eventually(async || { + client + .send(commands::incr_by( + "after_restart", + 1, + Duration::from_secs(60), + )) + .await + .ok() + }) + .await; + assert_eq!(recovered, 1); +} + +#[rstest] +#[tokio::test] +async fn a_client_created_while_redis_is_down_works_once_redis_is_up( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server.stop().await.unwrap(); + + let client = connect(&server, "", Duration::from_millis(500)); + let started = Instant::now(); + let while_down = client.ping().await; + let failed_in = started.elapsed(); + server.resume().await.unwrap(); + + assert!(failed_in < Duration::from_secs(2), "{failed_in:?}"); + + assert!( + matches!(while_down, Err(Error::Transport(_))), + "{while_down:?}" + ); + let recovered = eventually(async || client.ping().await.ok()).await; + assert_eq!(recovered, ()); +} + +#[rstest] +#[tokio::test] +async fn empty_batches_succeed_without_asking_redis( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + assert_eq!( + client.send(commands::mget::(&[])).await.unwrap(), + Vec::>::new() + ); + assert_eq!(client.send(commands::del(&[])).await.unwrap(), 0); +} + +#[rstest] +#[tokio::test] +async fn mget_and_del_span_slots_in_input_order( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set("a", "1", Duration::from_secs(60))) + .await + .unwrap(); + client + .send(commands::set("c", "3", Duration::from_secs(60))) + .await + .unwrap(); + + let values = client + .send(commands::mget::(&["c", "b", "a"])) + .await + .unwrap(); + let deleted = client.send(commands::del(&["a", "b", "c"])).await.unwrap(); + + assert_eq!( + values, + vec![Some("3".to_owned()), None, Some("1".to_owned())] + ); + assert_eq!(deleted, 2); + assert_eq!( + client + .send(commands::mget::(&["a", "c"])) + .await + .unwrap(), + vec![None, None] + ); +} + +#[rstest] +#[tokio::test] +async fn set_nx_only_writes_a_missing_key( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + let first = client + .send(commands::set_nx("lock", "owner-1", Duration::from_secs(60))) + .await + .unwrap(); + let second = client + .send(commands::set_nx("lock", "owner-2", Duration::from_secs(60))) + .await + .unwrap(); + + assert!(first); + assert!(!second); + assert_eq!( + client + .send(commands::get::("lock")) + .await + .unwrap() + .as_deref(), + Some("owner-1") + ); +} + +#[rstest] +#[tokio::test] +async fn ttl_reports_missing_persistent_and_remaining_time( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + persist(&client, "persistent").await; + client + .send(commands::set("expiring", "1", Duration::from_secs(60))) + .await + .unwrap(); + client + .send(commands::incr_by_float( + "expired_later", + Amount::Float(0.5), + Some(KEEP), + )) + .await + .unwrap(); + assert!( + client + .send(commands::expire("expired_later", Duration::from_secs(30))) + .await + .unwrap() + ); + + assert_eq!( + client.send(commands::ttl("missing")).await.unwrap(), + Ttl::Missing + ); + assert_eq!( + client.send(commands::ttl("persistent")).await.unwrap(), + Ttl::Persistent + ); + assert!(matches!( + client.send(commands::ttl("expiring")).await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) + )); + assert!(matches!( + client.send(commands::ttl("expired_later")).await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30) + )); + assert!( + !client + .send(commands::expire("missing", Duration::from_secs(30))) + .await + .unwrap() + ); +} + +#[rstest] +#[tokio::test] +async fn incr_by_float_accumulates_fractions( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + client + .send(commands::incr_by_float( + "spend", + Amount::Float(0.25), + Some(KEEP), + )) + .await + .unwrap(); + + assert_eq!( + client + .send(commands::incr_by_float( + "spend", + Amount::Float(0.5), + Some(KEEP) + )) + .await + .unwrap(), + 0.75 + ); +} + +#[rstest] +#[tokio::test] +async fn concurrent_commands_on_one_cloned_client_all_apply( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + + let tasks: Vec<_> = (0..100) + .map(|_| { + let client = client.clone(); + tokio::spawn(async move { + for _ in 0..10 { + client + .send(commands::incr_by("shared", 1, Duration::from_secs(60))) + .await + .unwrap(); + } + }) + }) + .collect(); + for task in tasks { + task.await.unwrap(); + } + + assert_eq!( + client.send(commands::get::("shared")).await.unwrap(), + Some(1000) + ); +} + +fn breaker(failure_threshold: u32) -> Breaker { + Breaker::new( + Config { + failure_threshold, + recovery_timeout: Duration::from_millis(200), + timeout_min_duration: Duration::ZERO, + }, + Box::new(|_| {}), + ) +} + +#[rstest] +#[tokio::test] +async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded( + &server, + "", + Duration::from_millis(300), + Some(breaker.clone()), + ); + + server.freeze().await.unwrap(); + let first = client.send(commands::get::("k")).await; + let second = client.send(commands::get::("k")).await; + let started = Instant::now(); + let third = client.send(commands::get::("k")).await; + let rejected_in = started.elapsed(); + server.thaw().await.unwrap(); + + assert!(matches!(first, Err(Error::Transport(_))), "{first:?}"); + assert!(matches!(second, Err(Error::Transport(_))), "{second:?}"); + assert!(matches!(third, Err(Error::BreakerOpen)), "{third:?}"); + assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); + + let recovered = eventually(async || client.send(commands::get::("k")).await.ok()).await; + assert_eq!(recovered, None); + assert!( + matches!(breaker.state(), State::Closed { .. }), + "{:?}", + breaker.state() + ); +} + +#[rstest] +#[tokio::test] +async fn rejected_commands_never_open_the_breaker( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded(&server, "", DEFAULT_TIMEOUT, Some(breaker.clone())); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + + let rejected = [ + client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await, + client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await, + client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await, + ]; + + assert!( + matches!(rejected[2], Err(Error::Server(_))), + "{:?}", + rejected[2] + ); + assert!( + matches!(breaker.state(), State::Closed { hard: 0, .. }), + "{:?}", + breaker.state() + ); +} + +#[rstest] +#[tokio::test] +async fn an_open_breaker_fails_every_pipeline_op_without_sending( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect_guarded(&server, "", Duration::from_millis(300), Some(breaker(1))); + let pipeline = client + .pipeline() + .with(commands::incr_by_float( + "a", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::incr_by_float( + "c", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )); + + server.freeze().await.unwrap(); + let failed = client.run(&pipeline).await; + let rejected = client.run(&pipeline).await; + server.thaw().await.unwrap(); + + assert!( + matches!(failed[0], Err(Error::Transport(_))), + "{:?}", + failed[0] + ); + assert!( + matches!(rejected[0], Err(Error::BreakerOpen)), + "{:?}", + rejected[0] + ); + assert!( + matches!(rejected[1], Err(Error::BreakerOpen)), + "{:?}", + rejected[1] + ); +} + +#[rstest] +#[tokio::test] +async fn a_host_endpoint_connects_to_its_database() { + let server = RedisServer::start().await.unwrap(); + let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); + let connect = |target| { + RedisClient::new(RedisConfig { + target, + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + }; + let url = |url: String| { + Target::Url(RedisUrl { + url: SecretValue::new(url), + tls: None, + db: None, + username: None, + password: None, + }) + }; + let in_db_3 = connect(Target::Standalone(Standalone { + host: "127.0.0.1".to_owned(), + port, + tls: None, + db: 3, + })) + .unwrap(); + let in_db_0 = connect(url(server.url().to_owned())).unwrap(); + let via_url = connect(url(format!("{}/3", server.url()))).unwrap(); + + in_db_3 + .send(commands::set("k", "3", Duration::from_secs(60))) + .await + .unwrap(); + + assert_eq!( + in_db_0.send(commands::get::("k")).await.unwrap(), + None + ); + assert_eq!( + via_url + .send(commands::get::("k")) + .await + .unwrap() + .as_deref(), + Some("3") + ); +} + +async fn ping_calls(cluster: &RedisCluster) -> Vec { + cluster + .cli_outputs(&["INFO", "commandstats"]) + .await + .unwrap() + .iter() + .map(|info| { + info.lines() + .find_map(|line| line.strip_prefix("cmdstat_ping:calls=")) + .and_then(|rest| rest.split(',').next()) + .and_then(|calls| calls.trim().parse().ok()) + .unwrap_or(0) + }) + .collect() +} + +fn health_checked( + server: &RedisDeployment, + interval: Duration, + timeout: Duration, + breaker: Option, + events: Option>, +) -> RedisClient { + RedisClient::new(RedisConfig { + target: endpoint(server), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(timeout, None), + breaker, + events: events.map(|events| events as Arc), + health_check: Some(interval), + keepalive: true, + }) + .unwrap() +} + +// redis-py 5.3.1 redis/asyncio/connection.py check_health: a connection idle longer than +// health_check_interval gets a PING before the next command, a busy one does not +#[rstest] +#[tokio::test] +async fn an_idle_client_pings_every_node_before_the_next_command_like_redis_py() { + let server = RedisDeployment::start(Topology::Cluster).await.unwrap(); + let RedisDeployment::Cluster(cluster) = &server else { + unreachable!("started a cluster") + }; + let client = health_checked(&server, Duration::from_secs(3), DEFAULT_TIMEOUT, None, None); + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + let before = ping_calls(cluster).await; + + client + .send(commands::set("k", "2", Duration::from_secs(60))) + .await + .unwrap(); + let busy = ping_calls(cluster).await; + tokio::time::sleep(Duration::from_millis(3500)).await; + client + .send(commands::set("k", "3", Duration::from_secs(60))) + .await + .unwrap(); + let idle = ping_calls(cluster).await; + + assert_eq!(busy, before); + assert_eq!( + idle, + before.iter().map(|calls| calls + 1).collect::>(), + "before {before:?}" + ); +} + +#[rstest] +#[tokio::test] +async fn a_failed_health_check_is_neither_a_call_nor_a_breaker_failure() { + let server = RedisDeployment::start(Topology::Cluster).await.unwrap(); + let recorded = Arc::new(Recorded::default()); + let breaker = breaker(2); + let client = health_checked( + &server, + Duration::from_millis(200), + Duration::from_millis(300), + Some(breaker.clone()), + Some(recorded.clone()), + ); + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + + server.freeze().await.unwrap(); + tokio::time::sleep(Duration::from_millis(300)).await; + let failed = client.send(commands::get::("k")).await; + server.thaw().await.unwrap(); + + assert!(matches!(failed, Err(Error::Transport(_))), "{failed:?}"); + assert!( + matches!(breaker.state(), State::Closed { timeouts: 1, .. }), + "{:?}", + breaker.state() + ); + assert_eq!(recorded.0.lock().unwrap().len(), 2); +} + +#[rstest] +#[case::int(Amount::Int(1), "1")] +#[case::float_keeps_its_point(Amount::Float(1.0), "1.0")] +#[case::small_float(Amount::Float(1e-5), "1e-05")] +#[case::large_float(Amount::Float(1e16), "1e+16")] +#[case::fraction(Amount::Float(0.1 + 0.2), "0.30000000000000004")] +fn amounts_are_written_like_redis_py_encodes_them(#[case] amount: Amount, #[case] expected: &str) { + assert_eq!( + redis::ToRedisArgs::to_redis_args(&amount), + vec![expected.as_bytes().to_vec()] + ); +} diff --git a/litellm-rust/crates/redis-client/tests/config.rs b/litellm-rust/crates/redis-client/tests/config.rs new file mode 100644 index 00000000000..101c6529e94 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/config.rs @@ -0,0 +1,290 @@ +use rstest::rstest; + +use litellm_redis_client::{Credentials, RedisUrl, Sentinel, StaticAuth, Target, TlsOptions}; +use litellm_secret_value::SecretValue; + +#[rstest] +#[case::url(format!("{:?}", Target::Url(RedisUrl { url: SecretValue::new("redis://user:credential-value@cache:6379/0"), tls: None, db: None, username: None, password: None })))] +#[case::static_password(format!("{:?}", Credentials::Static { username: None, password: Some(SecretValue::new("credential-value")) }))] +#[case::sentinel_password(format!("{:?}", Target::Sentinel(Sentinel { + monitors: Vec::new(), + service_name: "primary".to_owned(), + monitor_auth: Some(StaticAuth { username: None, password: Some(SecretValue::new("credential-value")) }), + tls: None, + db: 0, +})))] +fn debug_output_never_shows_the_secret(#[case] debug: String) { + assert!(!debug.contains("credential-value"), "{debug}"); + assert!(debug.contains("[REDACTED]"), "{debug}"); +} + +#[derive(Clone, Copy, Debug)] +enum Ca { + Trusted, + TrustedData, + TrustedDirectory, + SystemRoots, +} + +fn directory_holding(ca: &std::path::Path) -> std::path::PathBuf { + let directory = std::env::temp_dir().join(format!("litellm-ca-dir-{}", std::process::id())); + std::fs::create_dir_all(&directory).unwrap(); + std::fs::copy(ca, directory.join("ca.pem")).unwrap(); + directory +} + +#[derive(Clone, Copy, Debug)] +enum ClientKey { + Absent, + Plain, + Encrypted { password: Option<&'static str> }, +} + +async fn tls_server(require_client_certificate: bool) -> litellm_redis_testing::RedisTlsServer { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } + litellm_redis_testing::RedisTlsServer::start(require_client_certificate) + .await + .unwrap() +} + +fn tls_client( + server: &litellm_redis_testing::RedisTlsServer, + host: &str, + tls: TlsOptions, +) -> Result { + litellm_redis_client::RedisClient::new(litellm_redis_client::RedisConfig { + target: Target::Standalone(litellm_redis_client::Standalone { + host: host.to_owned(), + port: server.port(), + tls: Some(tls), + db: 0, + }), + credentials: Credentials::NONE, + namespace: litellm_redis_client::Namespace::default(), + timeouts: litellm_redis_client::Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) +} + +fn tls_options( + files: &litellm_redis_testing::TlsFiles, + ca: Ca, + verify_certificate: bool, + check_hostname: bool, + client_key: ClientKey, +) -> TlsOptions { + let (keyfile, key_password) = match client_key { + ClientKey::Absent => (None, None), + ClientKey::Plain => (Some(files.client_key.clone()), None), + ClientKey::Encrypted { password } => ( + Some(files.encrypted_client_key.clone()), + password.map(SecretValue::new), + ), + }; + TlsOptions { + ca_certs: matches!(ca, Ca::Trusted).then(|| files.ca.clone()), + ca_data: matches!(ca, Ca::TrustedData).then(|| std::fs::read_to_string(&files.ca).unwrap()), + ca_path: matches!(ca, Ca::TrustedDirectory).then(|| directory_holding(&files.ca)), + certfile: keyfile.is_some().then(|| files.client_cert.clone()), + keyfile, + key_password, + verify_certificate, + check_hostname, + } +} + +async fn ping_over_tls( + require_client_certificate: bool, + host: &str, + ca: Ca, + verify_certificate: bool, + check_hostname: bool, + client_key: ClientKey, +) -> Result<(), litellm_redis_client::Error> { + let server = tls_server(require_client_certificate).await; + let tls = tls_options( + server.files(), + ca, + verify_certificate, + check_hostname, + client_key, + ); + tls_client(&server, host, tls).unwrap().ping().await +} + +#[rstest] +#[case::trusted_ca_without_hostname_check_like_redis_py( + "127.0.0.1", + Ca::Trusted, + true, + false, + true +)] +#[case::trusted_ca_as_pem_text_like_ssl_ca_data("127.0.0.1", Ca::TrustedData, true, false, true)] +#[case::trusted_ca_in_a_directory_like_ssl_ca_path( + "127.0.0.1", + Ca::TrustedDirectory, + true, + false, + true +)] +#[case::hostname_check_with_a_matching_name("localhost", Ca::Trusted, true, true, true)] +#[case::hostname_check_against_an_unlisted_address("127.0.0.1", Ca::Trusted, true, true, false)] +#[case::an_untrusted_ca_is_refused("127.0.0.1", Ca::SystemRoots, true, false, false)] +#[case::cert_reqs_none_skips_verification("127.0.0.1", Ca::SystemRoots, false, false, true)] +#[tokio::test] +async fn tls_options_decide_which_servers_are_trusted( + #[case] host: &str, + #[case] ca: Ca, + #[case] verify_certificate: bool, + #[case] check_hostname: bool, + #[case] connects: bool, +) { + let result = ping_over_tls( + false, + host, + ca, + verify_certificate, + check_hostname, + ClientKey::Absent, + ) + .await; + + assert_eq!(result.is_ok(), connects, "{result:?}"); + if let Err(error) = result { + assert!( + matches!(error, litellm_redis_client::Error::Transport(_)), + "{error:?}" + ); + } +} + +#[rstest] +#[case::with_the_client_certificate(ClientKey::Plain, true)] +#[case::with_an_encrypted_key_and_its_ssl_password( + ClientKey::Encrypted { password: Some(litellm_redis_testing::CLIENT_KEY_PASSWORD) }, + true +)] +#[case::without_a_client_certificate(ClientKey::Absent, false)] +#[tokio::test] +async fn mutual_tls_sends_the_configured_client_certificate( + #[case] client_key: ClientKey, + #[case] connects: bool, +) { + let result = ping_over_tls(true, "127.0.0.1", Ca::Trusted, true, false, client_key).await; + + assert_eq!(result.is_ok(), connects, "{result:?}"); +} + +#[rstest] +#[case::without_ssl_password(None, "ssl_password is required")] +#[case::with_the_wrong_ssl_password(Some("wrong-password"), "does not decrypt the key")] +#[tokio::test] +async fn an_encrypted_key_that_cannot_be_decrypted_fails_before_connecting( + #[case] password: Option<&'static str>, + #[case] reason: &str, +) { + let server = tls_server(true).await; + let tls = tls_options( + server.files(), + Ca::Trusted, + true, + false, + ClientKey::Encrypted { password }, + ); + + let error = tls_client(&server, "127.0.0.1", tls).err().unwrap(); + + let litellm_redis_client::ConnectError::TlsKey { + path, + reason: cause, + } = &error + else { + panic!("expected a TLS key error, got {error:?}"); + }; + assert_eq!( + *path, + server.files().encrypted_client_key.display().to_string() + ); + assert!(cause.to_string().contains(reason), "{cause}"); +} + +#[rstest] +fn a_missing_tls_file_names_the_path() { + let tls = TlsOptions { + ca_certs: Some("/nonexistent/ca.pem".into()), + ..TlsOptions::verified() + }; + + let error = litellm_redis_client::RedisClient::new(litellm_redis_client::RedisConfig { + target: Target::Standalone(litellm_redis_client::Standalone { + host: "127.0.0.1".to_owned(), + port: 6379, + tls: Some(tls), + db: 0, + }), + credentials: Credentials::NONE, + namespace: litellm_redis_client::Namespace::default(), + timeouts: litellm_redis_client::Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .err() + .unwrap(); + + assert!(error.to_string().contains("/nonexistent/ca.pem"), "{error}"); +} + +#[rstest] +fn a_legacy_encrypted_pem_key_is_refused_with_a_conversion_hint() { + let directory = std::env::temp_dir().join(format!("litellm-legacy-key-{}", std::process::id())); + std::fs::create_dir_all(&directory).unwrap(); + let keyfile = directory.join("client.key"); + std::fs::write( + &keyfile, + "-----BEGIN RSA PRIVATE KEY-----\nProc-Type: 4,ENCRYPTED\nDEK-Info: AES-256-CBC,00112233445566778899AABBCCDDEEFF\n\nAAAA\n-----END RSA PRIVATE KEY-----\n", + ) + .unwrap(); + let tls = TlsOptions { + certfile: Some(keyfile.clone()), + keyfile: Some(keyfile), + key_password: Some(SecretValue::new("legacy-password")), + ..TlsOptions::verified() + }; + + let error = litellm_redis_client::RedisClient::new(litellm_redis_client::RedisConfig { + target: Target::Standalone(litellm_redis_client::Standalone { + host: "127.0.0.1".to_owned(), + port: 6379, + tls: Some(tls), + db: 0, + }), + credentials: Credentials::NONE, + namespace: litellm_redis_client::Namespace::default(), + timeouts: litellm_redis_client::Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .err() + .unwrap(); + + assert!( + matches!( + &error, + litellm_redis_client::ConnectError::TlsKey { + reason: litellm_redis_client::KeyError::LegacyEncryption, + .. + } + ), + "{error:?}" + ); +} diff --git a/litellm-rust/crates/redis-client/tests/connection.rs b/litellm-rust/crates/redis-client/tests/connection.rs new file mode 100644 index 00000000000..4efa6038cf0 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/connection.rs @@ -0,0 +1,329 @@ +mod support; + +use std::{collections::BTreeSet, sync::Arc, time::Duration}; + +use litellm_redis_testing::{ + ClientConnection, RedisDeployment, RedisSentinel, RedisServer, SERVICE_NAME, Topology, + eventually, +}; +use rstest::rstest; + +use litellm_redis_client::commands; +use litellm_redis_client::{ + Credentials, CredentialsProvider, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel, + Standalone, Target, Timeouts, +}; +use litellm_secret_value::SecretValue; +use support::{Rotating, connect_as, connections, proxied}; + +#[rstest] +#[tokio::test] +async fn a_pushed_credential_reauthenticates_the_live_connection_and_a_reconnect_uses_the_latest( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = Arc::new(Rotating::new("alice", "first")); + let client = connect_as( + &server, + Credentials::Provider(CredentialsProvider(provider.clone())), + ); + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + let before = connections(&server).await; + + provider.rotate("bob", "second"); + let rotated = eventually(async || { + let now = connections(&server).await; + (users(&now) == BTreeSet::from(["bob"])).then_some(now) + }) + .await; + server + .cli_on_every_node(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + server + .cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"]) + .await + .unwrap(); + let after_reconnect = eventually(async || { + client + .send(commands::get::("k")) + .await + .ok() + .flatten() + }) + .await; + + let reconnected = connections(&server).await; + + assert_eq!(users(&before), BTreeSet::from(["alice"])); + assert_eq!(ids(&rotated), ids(&before)); + assert_eq!(after_reconnect, "1"); + assert!(!reconnected.is_empty()); + assert_eq!(users(&reconnected), BTreeSet::from(["bob"])); +} + +fn users(connections: &[ClientConnection]) -> BTreeSet<&str> { + connections + .iter() + .map(|connection| connection.user.as_str()) + .collect() +} + +fn ids(connections: &[ClientConnection]) -> Vec<(u16, u64)> { + connections + .iter() + .map(|connection| (connection.node, connection.id)) + .collect() +} + +#[rstest] +#[tokio::test] +async fn an_out_of_range_database_is_a_server_rejection() { + let server = RedisServer::start().await.unwrap(); + let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); + let client = RedisClient::new(RedisConfig { + target: Target::Standalone(Standalone { + host: "127.0.0.1".into(), + port, + tls: None, + db: 999, + }), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_secs(2), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + let result = client.ping().await; + assert!(matches!(result, Err(Error::Server(_))), "{result:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_rejected_credential_is_a_transport_failure() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + server + .cli_on_every_node(&["ACL", "SETUSER", "default", "on", ">right"]) + .await + .unwrap(); + let with_password = |password: &str| { + connect_as( + &server, + Credentials::Static { + username: None, + password: Some(SecretValue::new(password)), + }, + ) + }; + with_password("right").ping().await.unwrap(); + + let rejected = with_password("wrong").ping().await; + + assert!(matches!(rejected, Err(Error::Transport(_))), "{rejected:?}"); +} + +fn sentinel_client(sentinel: &RedisSentinel, credentials: Credentials) -> RedisClient { + let (host, port) = sentinel.sentinel(); + RedisClient::new(RedisConfig { + target: Target::Sentinel(Sentinel { + monitors: vec![Node { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + monitor_auth: None, + tls: None, + db: 0, + }), + credentials, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_secs(1), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap() +} + +#[rstest] +#[tokio::test] +async fn sentinel_writes_follow_the_master_through_a_failover() { + let sentinel = RedisSentinel::start().await.unwrap(); + let client = sentinel_client(&sentinel, Credentials::NONE); + client + .send(commands::set("before", "1", Duration::from_secs(600))) + .await + .unwrap(); + eventually(async || { + let replica = + redis::Client::open(format!("redis://127.0.0.1:{}", sentinel.replica_port())).ok()?; + let mut replica = replica.get_multiplexed_async_connection().await.ok()?; + let value: Option = redis::cmd("GET") + .arg("before") + .query_async(&mut replica) + .await + .ok()?; + value + }) + .await; + + let promoted = sentinel.failover().await.unwrap(); + eventually(async || { + client + .send(commands::set("after", "2", Duration::from_secs(600))) + .await + .ok() + }) + .await; + + assert_eq!(promoted, sentinel.replica_port()); + assert_eq!( + client + .send(commands::get::("before")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + let promoted = redis::Client::open(format!("redis://127.0.0.1:{promoted}")).unwrap(); + let mut promoted = promoted.get_multiplexed_async_connection().await.unwrap(); + let after: Option = redis::cmd("GET") + .arg("after") + .query_async(&mut promoted) + .await + .unwrap(); + assert_eq!(after.as_deref(), Some("2")); +} + +const DATA_PASSWORD: &str = "data-secret"; + +#[rstest] +#[tokio::test] +async fn sentinel_reaches_a_master_that_requires_a_password() { + let sentinel = RedisSentinel::start_with_password(DATA_PASSWORD) + .await + .unwrap(); + let client = sentinel_client( + &sentinel, + Credentials::Static { + username: None, + password: Some(SecretValue::new(DATA_PASSWORD)), + }, + ); + + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + + assert_eq!( + client + .send(commands::get::("k")) + .await + .unwrap() + .as_deref(), + Some("1") + ); +} + +#[rstest] +#[tokio::test] +async fn sentinel_rediscovers_with_the_credentials_the_provider_holds_now() { + let sentinel = RedisSentinel::start_with_password(DATA_PASSWORD) + .await + .unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = Arc::new(Rotating::new("alice", "first")); + let client = sentinel_client( + &sentinel, + Credentials::Provider(CredentialsProvider(provider.clone())), + ); + client + .send(commands::set("k", "1", Duration::from_secs(600))) + .await + .unwrap(); + + provider.rotate("bob", "second"); + sentinel + .cli_on_data_nodes(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + sentinel + .cli_on_data_nodes(&["CLIENT", "KILL", "USER", "bob"]) + .await + .unwrap(); + let after_rediscovery = eventually(async || { + client + .send(commands::get::("k")) + .await + .ok() + .flatten() + }) + .await; + + assert_eq!(after_rediscovery, "1"); +} + +#[rstest] +#[tokio::test] +async fn a_timed_out_command_replaces_a_connection_that_stopped_answering( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let proxies = proxied(&server).await; + let client = RedisClient::new(RedisConfig { + target: proxies.target.clone(), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_millis(500), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + client + .send(commands::set("k", "1", Duration::from_secs(600))) + .await + .unwrap(); + + proxies.blackhole_open_connections(); + let stalled = client.send(commands::get::("k")).await; + let next = client.send(commands::get::("k")).await; + + assert!( + matches!(&stalled, Err(Error::Transport(error)) if error.is_timeout()), + "{stalled:?}" + ); + assert_eq!(next.unwrap().as_deref(), Some("1")); +} diff --git a/litellm-rust/crates/redis-client/tests/pipeline.rs b/litellm-rust/crates/redis-client/tests/pipeline.rs new file mode 100644 index 00000000000..5aed70b42dd --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/pipeline.rs @@ -0,0 +1,362 @@ +mod support; + +use std::time::{Duration, Instant}; + +use litellm_redis_client::Value; +use litellm_redis_client::commands; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; + +use litellm_redis_client::{Amount, DEFAULT_TIMEOUT, Error, Expiry, Script}; +use support::connect; + +async fn begin_migration(cluster: &litellm_redis_testing::RedisCluster, key: &str) { + let (_, first) = cluster.nodes().next().unwrap(); + let slot = cluster + .cli_on_node(first, &["CLUSTER", "KEYSLOT", key]) + .await + .unwrap(); + let slot = slot.trim(); + let number: u16 = slot.parse().unwrap(); + let nodes = cluster + .cli_on_node(first, &["CLUSTER", "NODES"]) + .await + .unwrap(); + let owner = nodes.lines().find(|line| line.split_whitespace().skip(8).any(|range| { + let (start, end) = range.split_once('-').unwrap_or((range, range)); + matches!((start.parse::(), end.parse::()), (Ok(start), Ok(end)) if (start..=end).contains(&number)) + })).unwrap(); + let target = nodes.lines().find(|line| *line != owner).unwrap(); + let address = |line: &str| { + let value = line.split_whitespace().nth(1).unwrap(); + value + .split_once('@') + .unwrap() + .0 + .rsplit_once(':') + .unwrap() + .1 + .parse::() + .unwrap() + }; + let source_port = address(owner); + let target_port = address(target); + let source_id = owner.split_whitespace().next().unwrap(); + let target_id = target.split_whitespace().next().unwrap(); + assert_eq!( + cluster + .cli_on_node(source_port, &["GET", key]) + .await + .unwrap() + .trim(), + "0" + ); + assert_eq!( + cluster + .cli_on_node( + target_port, + &["CLUSTER", "SETSLOT", slot, "IMPORTING", source_id] + ) + .await + .unwrap() + .trim(), + "OK" + ); + assert_eq!( + cluster + .cli_on_node( + source_port, + &["CLUSTER", "SETSLOT", slot, "MIGRATING", target_id] + ) + .await + .unwrap() + .trim(), + "OK" + ); + let state = cluster + .cli_on_node(source_port, &["CLUSTER", "NODES"]) + .await + .unwrap(); + assert!( + state.contains(&format!("[{slot}->-{target_id}]")), + "{state}" + ); +} + +#[rstest] +#[tokio::test] +async fn a_retrying_neighbor_cannot_reapply_a_cluster_counter() { + let cluster = litellm_redis_testing::RedisCluster::start().await.unwrap(); + let server = RedisDeployment::Cluster(cluster); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let counter = "{migration}:counter"; + client + .send(commands::set(counter, 0, Duration::from_secs(60))) + .await + .unwrap(); + let RedisDeployment::Cluster(cluster) = &server else { + unreachable!() + }; + begin_migration(cluster, counter).await; + let both_keys = Script::new("return redis.call('MGET', KEYS[1], KEYS[2])"); + + let results = client + .run( + &client + .pipeline() + .with(commands::incr_by_float( + counter, + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::eval::( + &both_keys, + &[counter, "{migration}:missing"], + Vec::::new(), + )), + ) + .await; + + assert_eq!( + client.send(commands::get::(counter)).await.unwrap(), + Some(1) + ); + assert_eq!( + results[0].as_ref().unwrap(), + &Value::BulkString(b"1".to_vec()) + ); + assert!( + matches!(&results[1], Err(Error::Server(error)) if error.code() == Some("TRYAGAIN")), + "{:?}", + results[1] + ); +} + +#[rstest] +#[tokio::test] +async fn each_op_settles_on_its_own_in_input_order( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + client + .send(commands::set( + "text", + "not a number", + Duration::from_secs(60), + )) + .await + .unwrap(); + let pipeline = client + .pipeline() + .with(commands::incr_by_float( + "a", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::incr_by_float( + "text", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::incr_by_float( + "c", + Amount::Int(5), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::set("b", "2", Duration::from_secs(60))) + .with(commands::get::("b")) + .with(commands::incr_by_float( + "a", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::get::("missing")); + + let results = client.run(&pipeline).await; + + assert_eq!(results.len(), 7); + assert_eq!( + results[0].as_ref().unwrap(), + &Value::BulkString(b"1".to_vec()) + ); + assert!( + matches!(results[1], Err(Error::Server(_))), + "{:?}", + results[1] + ); + assert_eq!( + results[2].as_ref().unwrap(), + &Value::BulkString(b"5".to_vec()) + ); + assert_eq!(results[3].as_ref().unwrap(), &Value::Okay); + assert_eq!( + results[4].as_ref().unwrap(), + &Value::BulkString(b"2".to_vec()) + ); + assert_eq!( + results[5].as_ref().unwrap(), + &Value::BulkString(b"2".to_vec()) + ); + assert_eq!(results[6].as_ref().unwrap(), &Value::Nil); +} + +#[rstest] +#[case::cached(true, Value::BulkString(b"2".to_vec()))] +#[case::lost_by_the_server(false, Value::Nil)] +#[tokio::test] +async fn a_lost_script_reruns_after_the_rest_of_its_pipeline_like_redis_py( + #[case] cached: bool, + #[case] read_in_the_same_pipeline: Value, +) { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let client = connect(&server, "ns", DEFAULT_TIMEOUT); + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client + .send(commands::eval::(&incr, &["warmup"], 1)) + .await + .unwrap(); + if !cached { + server + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + } + let pipeline = client + .pipeline() + .with(commands::eval::(&incr, &["counter"], 2)) + .with(commands::get::("counter")); + + let results = client.run(&pipeline).await; + + assert_eq!(results[0].as_ref().unwrap(), &Value::Int(2)); + assert_eq!(results[1].as_ref().unwrap(), &read_in_the_same_pipeline); + assert_eq!( + client.send(commands::get::("counter")).await.unwrap(), + Some(2) + ); +} + +#[rstest] +#[tokio::test] +async fn an_unreachable_server_fails_every_op_as_transport( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + let pipeline = client + .pipeline() + .with(commands::incr_by_float( + "a", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::incr_by_float( + "c", + Amount::Int(1), + Some(Expiry::Refresh(Duration::from_secs(60))), + )); + + server.freeze().await.unwrap(); + let results = client.run(&pipeline).await; + server.thaw().await.unwrap(); + + assert_eq!(results.len(), 2); + assert!( + matches!(results[0], Err(Error::Transport(_))), + "{:?}", + results[0] + ); + assert!( + matches!(results[1], Err(Error::Transport(_))), + "{:?}", + results[1] + ); +} + +#[rstest] +#[tokio::test] +async fn an_empty_pipeline_returns_without_asking_redis( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_secs(2)); + + server.freeze().await.unwrap(); + let started = Instant::now(); + let results = client.run(&client.pipeline()).await; + let elapsed = started.elapsed(); + server.thaw().await.unwrap(); + + assert!(results.is_empty()); + assert!(elapsed < Duration::from_secs(1), "{elapsed:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_pipeline_sends_scripts_by_sha_and_reruns_only_the_ones_the_server_lost( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&deployment, "scripts", DEFAULT_TIMEOUT); + let count = Script::new("return redis.call('INCR', KEYS[1])"); + deployment + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + deployment + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + + let first = client + .run( + &client + .pipeline() + .with(commands::incr_by_float( + "{a}spend", + Amount::Float(0.5), + Some(Expiry::Refresh(Duration::from_secs(60))), + )) + .with(commands::eval::( + &count, + &["{a}"], + Vec::::new(), + )) + .with(commands::eval::( + &count, + &["b"], + Vec::::new(), + )), + ) + .await; + let second = client + .run(&client.pipeline().with(commands::eval::( + &count, + &["{a}"], + Vec::::new(), + ))) + .await; + + let replies: Vec = first + .into_iter() + .chain(second) + .map(Result::unwrap) + .collect(); + assert_eq!( + replies, + [ + Value::BulkString(b"0.5".to_vec()), + Value::Int(1), + Value::Int(1), + Value::Int(2) + ] + ); + assert_eq!( + client.send(commands::get::("{a}spend")).await.unwrap(), + Some(0.5) + ); + let stats = support::command_stats(&deployment).await; + assert!(support::calls(&stats, "evalsha") > 0, "{stats}"); + assert_eq!(support::calls(&stats, "eval"), 0, "{stats}"); +} diff --git a/litellm-rust/crates/redis-client/tests/pubsub.rs b/litellm-rust/crates/redis-client/tests/pubsub.rs new file mode 100644 index 00000000000..2ad7cb38a5a --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/pubsub.rs @@ -0,0 +1,189 @@ +mod support; + +use std::time::{Duration, Instant}; + +use futures_util::StreamExt; +use litellm_circuit_breaker::{Breaker, Config}; +use litellm_redis_client::commands; +use litellm_redis_client::{ + Credentials, DEFAULT_TIMEOUT, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel, + Target, Timeouts, +}; +use litellm_redis_testing::{RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually}; +use rstest::rstest; +use support::{connect, connect_guarded}; +use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; + +const CHANNEL: &str = "litellm_proxy.auth_cache_invalidation"; + +fn listen(client: &RedisClient, channel: &str) -> UnboundedReceiver> { + let (sender, receiver) = unbounded_channel(); + let mut messages = client.subscribe(channel); + tokio::spawn(async move { + while let Some(message) = messages.next().await { + if sender.send(message).is_err() { + return; + } + } + }); + receiver +} + +async fn delivered( + publisher: &RedisClient, + channel: &str, + received: &mut UnboundedReceiver>, + payload: &[u8], +) { + eventually(async || { + publisher + .send(commands::publish(channel, payload)) + .await + .ok()?; + let deadline = tokio::time::Instant::now() + Duration::from_millis(200); + while let Ok(Some(message)) = tokio::time::timeout_at(deadline, received.recv()).await { + if message == payload { + return Some(()); + } + } + None + }) + .await; +} + +#[rstest] +#[tokio::test] +async fn a_subscriber_receives_raw_bytes_on_the_channel_name_without_the_namespace( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let publisher = connect(&server, "ns", DEFAULT_TIMEOUT); + let plain = connect(&server, "", DEFAULT_TIMEOUT); + let mut received = listen(&plain, CHANNEL); + + delivered( + &publisher, + CHANNEL, + &mut received, + b"{\"cache_key\":\"k\"}\xff", + ) + .await; +} + +#[rstest] +#[tokio::test] +async fn a_subscriber_resubscribes_after_redis_restarts( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", Duration::from_millis(500)); + let mut received = listen(&client, CHANNEL); + delivered(&client, CHANNEL, &mut received, b"before").await; + + server.restart().await.unwrap(); + + delivered(&client, CHANNEL, &mut received, b"after").await; +} + +#[rstest] +#[tokio::test] +async fn a_sentinel_subscriber_receives_from_the_master() { + let server = RedisSentinel::start().await.unwrap(); + let (host, port) = server.sentinel(); + let client = RedisClient::new(RedisConfig { + target: Target::Sentinel(Sentinel { + monitors: vec![Node { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + monitor_auth: None, + tls: None, + db: 0, + }), + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + let mut received = listen(&client, CHANNEL); + + delivered(&client, CHANNEL, &mut received, b"payload").await; +} + +async fn publish_calls(server: &RedisDeployment) -> Vec { + let replies = match server { + RedisDeployment::Standalone(server) => { + vec![server.cli(&["INFO", "commandstats"]).await.unwrap()] + } + RedisDeployment::Cluster(cluster) => cluster + .cli_outputs(&["INFO", "commandstats"]) + .await + .unwrap(), + }; + replies + .iter() + .flat_map(|reply| reply.lines()) + .filter(|line| line.starts_with("cmdstat_publish:")) + .map(str::to_owned) + .collect() +} + +#[rstest] +#[tokio::test] +async fn an_open_breaker_rejects_a_publish_without_sending_it( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = Breaker::new( + Config { + failure_threshold: 2, + recovery_timeout: Duration::from_secs(60), + timeout_min_duration: Duration::ZERO, + }, + Box::new(|_| {}), + ); + let client = connect_guarded(&server, "", Duration::from_millis(300), Some(breaker)); + + server.freeze().await.unwrap(); + let opened = [client.ping().await, client.ping().await]; + let started = Instant::now(); + let rejected = client.send(commands::publish(CHANNEL, b"payload")).await; + let rejected_in = started.elapsed(); + server.thaw().await.unwrap(); + + assert!( + opened + .iter() + .all(|ping| matches!(ping, Err(Error::Transport(_)))), + "{opened:?}" + ); + assert!(matches!(rejected, Err(Error::BreakerOpen)), "{rejected:?}"); + assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); + assert_eq!(publish_calls(&server).await, Vec::::new()); +} + +// redis-py's PubSub subscribes on a RESP2 connection, so it never sends HELLO +#[rstest] +#[tokio::test] +async fn a_subscription_speaks_resp2_like_redis_py( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&["CONFIG", "RESETSTAT"]) + .await + .unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let mut received = listen(&client, CHANNEL); + + delivered(&client, CHANNEL, &mut received, b"evict").await; + + let stats = support::command_stats(&server).await; + assert!(support::calls(&stats, "subscribe") > 0, "{stats}"); + assert_eq!(support::calls(&stats, "hello"), 0, "{stats}"); +} diff --git a/litellm-rust/crates/redis-client/tests/script.rs b/litellm-rust/crates/redis-client/tests/script.rs new file mode 100644 index 00000000000..ea6d1c5ff78 --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/script.rs @@ -0,0 +1,133 @@ +mod support; + +use std::time::Duration; + +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; + +use litellm_redis_client::Script; +use litellm_redis_client::commands; +use litellm_redis_client::{DEFAULT_TIMEOUT, Error, Ttl}; +use support::connect; + +#[rstest] +#[tokio::test] +async fn script_keys_are_namespaced( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "a", DEFAULT_TIMEOUT); + let bare = connect(&server, "", DEFAULT_TIMEOUT); + let read = Script::new("return redis.call('GET', KEYS[1])"); + bare.send(commands::set("a:x", "1", Duration::from_secs(60))) + .await + .unwrap(); + + let value: Option = client + .send(commands::eval(&read, &["x"], Vec::::new())) + .await + .unwrap(); + + assert_eq!(value.as_deref(), Some("1")); +} + +#[rstest] +#[tokio::test] +async fn a_failing_script_is_a_server_error_and_keeps_the_connection( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + + let failed = client + .send(commands::eval::(&incr, &["list"], 1)) + .await; + + assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}"); + assert_eq!( + client + .send(commands::eval::(&incr, &["counter"], 3)) + .await + .unwrap(), + 3 + ); +} + +#[rstest] +#[tokio::test] +async fn a_flushed_script_is_loaded_again( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "", DEFAULT_TIMEOUT); + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client + .send(commands::eval::(&incr, &["counter"], 1)) + .await + .unwrap(); + + server + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + + assert_eq!( + client + .send(commands::eval::(&incr, &["counter"], 1)) + .await + .unwrap(), + 2 + ); +} + +const INCREMENT_WITH_FLOOR: &str = "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) \ +if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end \ +if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end \ +return count"; + +#[rstest] +#[tokio::test] +async fn pythons_increment_with_floor_script_floors_at_zero_and_expires_in_one_call( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[values(false, true)] pipelined: bool, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&server, "ns", DEFAULT_TIMEOUT); + let floor = Script::new(INCREMENT_WITH_FLOOR); + server + .cli_on_every_node(&["SCRIPT", "FLUSH"]) + .await + .unwrap(); + + let count = match pipelined { + false => client + .send(commands::eval::(&floor, &["slots"], (-5, 60))) + .await + .unwrap(), + true => { + let results = client + .run(&client.pipeline().with(commands::eval::( + &floor, + &["slots"], + (-5, 60), + ))) + .await; + redis::from_redis_value(results[0].as_ref().unwrap().clone()).unwrap() + } + }; + + assert_eq!(count, 0); + assert_eq!( + client.send(commands::get::("slots")).await.unwrap(), + Some(0) + ); + assert!(matches!( + client.send(commands::ttl("slots")).await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) + )); +} diff --git a/litellm-rust/crates/redis-client/tests/support/mod.rs b/litellm-rust/crates/redis-client/tests/support/mod.rs new file mode 100644 index 00000000000..9cf209dce4a --- /dev/null +++ b/litellm-rust/crates/redis-client/tests/support/mod.rs @@ -0,0 +1,206 @@ +#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset + +use std::{pin::Pin, time::Duration}; + +use futures_util::Stream; +use litellm_circuit_breaker::Breaker; +use litellm_redis_client::{ + BasicAuth, Cluster, Credentials, Expiry, Namespace, Node, RedisClient, RedisConfig, RedisError, + RedisUrl, Standalone, StreamingCredentialsProvider, Target, Timeouts, +}; +use litellm_redis_testing::{ClientConnection, RedisDeployment, TcpProxy}; +use litellm_secret_value::SecretValue; +use tokio::sync::watch; + +pub const KEEP: Expiry = Expiry::KeepExisting(Duration::from_secs(60)); + +pub fn endpoint(deployment: &RedisDeployment) -> Target { + match deployment { + RedisDeployment::Standalone(server) => Target::Url(RedisUrl { + url: SecretValue::new(server.url()), + tls: None, + db: None, + username: None, + password: None, + }), + RedisDeployment::Cluster(cluster) => Target::Cluster(Cluster { + nodes: cluster + .nodes() + .map(|(host, port)| Node { + host: host.to_owned(), + port, + }) + .collect(), + tls: None, + }), + } +} + +pub fn connect(deployment: &RedisDeployment, namespace: &str, timeout: Duration) -> RedisClient { + connect_guarded(deployment, namespace, timeout, None) +} + +pub fn connect_guarded( + deployment: &RedisDeployment, + namespace: &str, + timeout: Duration, + breaker: Option, +) -> RedisClient { + RedisClient::new(RedisConfig { + target: endpoint(deployment), + credentials: Credentials::NONE, + namespace: Namespace::new(namespace), + timeouts: Timeouts::new(timeout, None), + breaker, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap() +} + +pub fn connect_as(deployment: &RedisDeployment, credentials: Credentials) -> RedisClient { + RedisClient::new(RedisConfig { + target: endpoint(deployment), + credentials, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_secs(2), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap() +} + +pub async fn connections(deployment: &RedisDeployment) -> Vec { + deployment.clients().await.unwrap() +} + +/// A `TcpProxy` in front of every node, and a target that reaches the nodes only through them. A +/// cluster announces the proxy ports, so its slot map and redirects lead there too. +pub struct Proxied { + proxies: Vec, + pub target: Target, +} + +impl Proxied { + pub fn blackhole_open_connections(&self) { + self.proxies + .iter() + .for_each(TcpProxy::blackhole_open_connections); + } +} + +pub async fn proxied(deployment: &RedisDeployment) -> Proxied { + match deployment { + RedisDeployment::Standalone(server) => { + let proxy = TcpProxy::start(server.port()).await.unwrap(); + let target = Target::Standalone(Standalone { + host: "127.0.0.1".to_owned(), + port: proxy.port(), + tls: None, + db: 0, + }); + Proxied { + proxies: vec![proxy], + target, + } + } + RedisDeployment::Cluster(cluster) => { + let mut proxies = Vec::new(); + for (_, port) in cluster.nodes() { + proxies.push(TcpProxy::start(port).await.unwrap()); + } + let ports: [u16; 3] = std::array::from_fn(|node| proxies[node].port()); + cluster.announce_ports(ports).await.unwrap(); + let target = Target::Cluster(Cluster { + nodes: ports + .iter() + .map(|&port| Node { + host: "127.0.0.1".to_owned(), + port, + }) + .collect(), + tls: None, + }); + Proxied { proxies, target } + } + } +} + +/// `INFO
` from every node, concatenated. +async fn info(deployment: &RedisDeployment, section: &str) -> String { + match deployment { + RedisDeployment::Standalone(server) => server.cli(&["INFO", section]).await.unwrap(), + RedisDeployment::Cluster(cluster) => { + let mut stats = String::new(); + for (_, port) in cluster.nodes() { + stats.push_str(&cluster.cli_on_node(port, &["INFO", section]).await.unwrap()); + } + stats + } + } +} + +/// `INFO commandstats` from every node, so a test can count what was actually sent. +pub async fn command_stats(deployment: &RedisDeployment) -> String { + info(deployment, "commandstats").await +} + +/// Keys looked up on every node since the last `CONFIG RESETSTAT`, found or not. +pub async fn keys_read(deployment: &RedisDeployment) -> u64 { + let stats = info(deployment, "stats").await; + ["keyspace_hits:", "keyspace_misses:"] + .iter() + .flat_map(|field| { + stats + .lines() + .filter_map(move |line| line.trim().strip_prefix(field)?.parse::().ok()) + }) + .sum() +} + +/// How many times `command` reached a server, summed across the nodes in `stats`. +pub fn calls(stats: &str, command: &str) -> u64 { + let prefix = format!("cmdstat_{command}:calls="); + stats + .lines() + .filter_map(|line| line.trim().strip_prefix(&prefix)) + .filter_map(|rest| rest.split(',').next()?.parse::().ok()) + .sum() +} + +pub struct Rotating(watch::Sender); + +impl Rotating { + pub fn new(username: &str, password: &str) -> Self { + Self(watch::Sender::new(BasicAuth::new( + username.to_owned(), + password.to_owned(), + ))) + } + + pub fn rotate(&self, username: &str, password: &str) { + self.0 + .send_replace(BasicAuth::new(username.to_owned(), password.to_owned())); + } +} + +impl StreamingCredentialsProvider for Rotating { + fn subscribe( + &self, + ) -> Pin> + Send + 'static>> { + let receiver = self.0.subscribe(); + Box::pin(futures_util::stream::unfold( + (receiver, true), + |(mut receiver, first)| async move { + if !first { + receiver.changed().await.ok()?; + } + let credentials = receiver.borrow_and_update().clone(); + Some((Ok(credentials), (receiver, false))) + }, + )) + } +} diff --git a/litellm-rust/crates/redis-connect/AGENTS.md b/litellm-rust/crates/redis-connect/AGENTS.md new file mode 100644 index 00000000000..1b47804ddd8 --- /dev/null +++ b/litellm-rust/crates/redis-connect/AGENTS.md @@ -0,0 +1,10 @@ +- Turns a `litellm_redis_types::Resolution` into a `RedisClient`, the same way for every host. It depends on `litellm-redis-types`, never on `litellm-config`, so a host that resolves Redis some other way connects the same way. `redis-client` stays cloud-agnostic, so credential providers and cloud token sources are built here +- The host passes only what it alone can build in `HostServices`: the event sink, the HTTP client and the environment lookup. The breaker config comes from `Resolution::breaker`, which `config` reads from `REDIS_CIRCUIT_BREAKER_*` with Python's fallback, so a bad value keeps the default instead of failing connect. Each cloud source asks its auth crate for one explicit identity, the way the Python SDK builds it: `AwsAuthService::default_chain_credentials` (botocore's default chain), `CredentialSource::application_default` (`google.auth.default`) and `AzureIdentity` (the class `_build_azure_credential` picks). LiteLLM's provider settings such as `AWS_ROLE_NAME`, `AZURE_AD_TOKEN` or `VERTEXAI_CREDENTIALS` belong to Bedrock, Azure OpenAI and Vertex and never reach those constructors. `Auth::Provider` (a handle for a Python `credential_provider` object) fails at connect with an error naming the key, since no host can supply one yet +- Every expiring credential goes through the one refresh loop in `refresh.rs`. A cloud source only fetches. Never add a second loop or let a source sleep. The expiry deadline lives only in the published `Latest::Ready`, so expiring it leaves `Ready` and cannot fire again. Never keep a copy of the deadline in the loop +- AWS SDK identity settings must be in the process environment. `aws-config` cannot take the host lookup, so config-only overrides for profiles, role/web-identity settings, credentials, files and metadata sources fail before connecting rather than silently using another identity. Equal process/config values are allowed. Bedrock-only settings remain unrelated to Redis +- redis-rs subscribes once for the connection handshake and again for re-authentication, so a new connection sends a redundant initial AUTH. Unchanged refreshes do not trigger another AUTH. Do not skip a provider's first or alternating subscriptions, because concurrent connections and retries each need the current token. Removing the initial duplicate requires an upstream connection-level fix +- The loop starts on the first `subscribe()`, which redis-rs calls inside its runtime, so `connect` is safe off a runtime +- A token-only credential (Azure without `REDIS_USERNAME`, GCP) is sent as `AUTH default `, because redis-rs always sends a username with a credentials provider. On Redis 6+ that is the same as Python's `AUTH ` +- Tests for the refresh loop use a fake source and real Redis ACL users. Cloud sources are proven against the live service, unit tests only pin request shapes we own +- Rotation tests must retain the same server-side connection IDs across re-AUTH, then force a reconnect after removing the old user and verify the new identity still works +- Live cloud tests live in `tests/live/`. Set `LITELLM_REDIS_AWS_LIVE_CONFIG`, `LITELLM_REDIS_AZURE_LIVE_CONFIG` or `LITELLM_REDIS_GCP_LIVE_CONFIG` to a JSON file containing `redis` settings and `hold_seconds`. Set the waiting period from the deployment's token lifetime when checking expiry. Optional fixtures return early when unset, and configured tests run under plain `cargo test`. An unconfigured run never counts as live proof diff --git a/litellm-rust/crates/redis-connect/Cargo.toml b/litellm-rust/crates/redis-connect/Cargo.toml new file mode 100644 index 00000000000..72a453d36de --- /dev/null +++ b/litellm-rust/crates/redis-connect/Cargo.toml @@ -0,0 +1,39 @@ +[package] +name = "litellm-redis-connect" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-aws.workspace = true +litellm-auth-azure.workspace = true +litellm-auth-gcp.workspace = true +litellm-auth-types.workspace = true +litellm-error.workspace = true +litellm-secret-value.workspace = true +strum.workspace = true +litellm-circuit-breaker.workspace = true +litellm-redis-types.workspace = true +litellm-core-utils.workspace = true +litellm-http.workspace = true +rustls.workspace = true +litellm-redis-client.workspace = true +futures-util.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +veil.workspace = true +time.workspace = true +url.workspace = true +tokio = { workspace = true, features = ["sync", "time", "macros"] } +tracing.workspace = true + +[dev-dependencies] +litellm-redis-resolve.workspace = true +litellm-http = { workspace = true, features = ["test-support"] } +litellm-redis-testing.workspace = true +rustls.workspace = true +rstest.workspace = true +serde_json.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "test-util"] } diff --git a/litellm-rust/crates/redis-connect/src/connect.rs b/litellm-rust/crates/redis-connect/src/connect.rs new file mode 100644 index 00000000000..305aa181f2b --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/connect.rs @@ -0,0 +1,154 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_circuit_breaker::{Breaker, Config as BreakerConfig}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_client::{ + Credentials, CredentialsProvider, DEFAULT_TIMEOUT, Namespace, RedisClient, RedisConfig, + RedisEvents, Target, Timeouts, +}; +use litellm_redis_types::{Auth, HealthCheck, Resolution, StaticAuth}; + +use crate::{ + CredentialSource, Error, RefreshingCredentials, + sources::{AzureAd, ElastiCacheIam, GcpIam}, +}; + +pub type SharedEnvironment = Arc; + +pub struct HostServices { + pub events: Option>, + pub http: litellm_http::Client, + pub environment: SharedEnvironment, + pub credential_provider: Option>, +} + +// litellm/constants.py REDIS_CLUSTER_HEALTH_CHECK_INTERVAL, the default _redis.py sets for cluster clients only +const CLUSTER_HEALTH_CHECK_INTERVAL: Duration = Duration::from_secs(25); + +pub fn connect( + resolved: Resolution, + host: HostServices, + breaker: Option, +) -> Result { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } + tracing::info!(target = ?resolved.target, namespace = ?resolved.namespace, "redis client configured"); + let credentials = credentials(resolved.auth, &host)?; + let events = host.events; + let breaker = breaker.map(|config| { + let observer = events.clone(); + Breaker::new( + config, + Box::new(move |transition| { + if let Some(events) = &observer { + events.breaker(transition); + } + }), + ) + }); + let health_check = health_check(&resolved.target, resolved.health_check); + Ok(RedisClient::new(RedisConfig { + target: resolved.target, + credentials, + namespace: resolved.namespace.map(Namespace::new).unwrap_or_default(), + timeouts: Timeouts::new( + resolved.timeout.unwrap_or(DEFAULT_TIMEOUT), + resolved.connect_timeout, + ), + breaker, + events, + health_check, + keepalive: resolved.keepalive, + })?) +} + +fn health_check(target: &Target, configured: HealthCheck) -> Option { + match (target, configured) { + (_, HealthCheck::Every(interval)) => Some(interval), + (_, HealthCheck::Disabled) => None, + (Target::Cluster(_), HealthCheck::Unset) => Some(CLUSTER_HEALTH_CHECK_INTERVAL), + (_, HealthCheck::Unset) => None, + } +} + +fn credentials(auth: Auth, host: &HostServices) -> Result { + let refreshing = |provider: RefreshingCredentials| { + Credentials::Provider(CredentialsProvider(Arc::new(provider))) + }; + Ok(match auth { + Auth::Static(StaticAuth { username, password }) => { + Credentials::Static { username, password } + } + Auth::Provider(_) => match &host.credential_provider { + Some(source) => refreshing(RefreshingCredentials::new(Arc::clone(source))), + None => return Err(Error::NoHostCredentialProvider), + }, + Auth::AzureAd(auth) => refreshing(RefreshingCredentials::new(AzureAd::new( + auth, + host.environment.clone(), + ))), + Auth::AwsIam(auth) => refreshing(RefreshingCredentials::new(ElastiCacheIam::new( + auth, + &*host.environment, + )?)), + Auth::GcpIam { service_account } => refreshing(RefreshingCredentials::new(GcpIam::new( + host.http.clone(), + service_account, + host.environment.clone(), + ))), + }) +} + +#[cfg(test)] +mod tests { + use litellm_redis_client::{Cluster, Node, Standalone}; + use rstest::rstest; + + use super::*; + + fn cluster() -> Target { + Target::Cluster(Cluster { + nodes: vec![Node { + host: "node".to_owned(), + port: 7000, + }], + tls: None, + }) + } + + fn standalone() -> Target { + Target::Standalone(Standalone { + host: "redis".to_owned(), + port: 6379, + tls: None, + db: 0, + }) + } + + #[rstest] + #[case::an_unset_cluster_takes_the_python_default( + cluster(), + HealthCheck::Unset, + Some(CLUSTER_HEALTH_CHECK_INTERVAL) + )] + #[case::zero_keeps_a_cluster_off(cluster(), HealthCheck::Disabled, None)] + #[case::an_unset_standalone_stays_off(standalone(), HealthCheck::Unset, None)] + #[case::a_configured_interval_wins( + cluster(), + HealthCheck::Every(Duration::from_secs(5)), + Some(Duration::from_secs(5)) + )] + #[case::a_negative_interval_checks_every_command( + standalone(), + HealthCheck::Every(Duration::ZERO), + Some(Duration::ZERO) + )] + fn health_checks_follow_redis_py( + #[case] target: Target, + #[case] configured: HealthCheck, + #[case] expected: Option, + ) { + assert_eq!(health_check(&target, configured), expected); + } +} diff --git a/litellm-rust/crates/redis-connect/src/error.rs b/litellm-rust/crates/redis-connect/src/error.rs new file mode 100644 index 00000000000..98343a03fb8 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/error.rs @@ -0,0 +1,102 @@ +#[derive(veil::Redact, thiserror::Error)] +pub enum Error { + #[error(transparent)] + Connect(#[from] litellm_redis_client::ConnectError), + #[error( + "Redis setting credential_provider names a Python object, which this host cannot supply" + )] + NoHostCredentialProvider, + #[error( + "Redis credential setting {setting} must be set in the process environment, not only environment_variables" + )] + ProcessEnvironmentRequired { setting: &'static str }, + #[error("{stage} failed{}", http_status_suffix(*http_status))] + Credential { + stage: CredentialStage, + http_status: Option, + #[source] + #[redact(with = "[REDACTED]")] + cause: litellm_error::Source, + }, + #[error("credential acquisition exceeded {timeout:?}")] + CredentialTimeout { + timeout: std::time::Duration, + #[source] + cause: tokio::time::error::Elapsed, + }, + #[error("Redis credential has expired")] + ExpiredCredential, + #[error("no {provider} credential is available for Redis")] + MissingCredential { provider: &'static str }, +} + +#[derive(Clone, Copy, Debug, strum::Display, strum::VariantArray)] +pub enum CredentialStage { + #[strum(serialize = "caller identity")] + Identity, + #[strum(serialize = "Redis token signing")] + Signing, + #[strum(serialize = "Redis token exchange")] + Exchange, + #[strum(serialize = "Redis token response")] + Response, + #[strum(serialize = "Redis credential_provider")] + Provider, +} + +impl Error { + pub fn credential( + stage: CredentialStage, + ) -> impl FnOnce(E) -> Self { + move |cause| Self::Credential { + stage, + http_status: None, + cause: litellm_error::Source::new(cause), + } + } + + pub(crate) fn credential_response( + stage: CredentialStage, + status: u16, + ) -> impl FnOnce(litellm_http::Error) -> Self { + move |cause| Self::Credential { + stage, + http_status: Some(status), + cause: litellm_error::Source::new(cause), + } + } + + pub(crate) fn upstream_status(&self) -> Option { + match self { + Self::Credential { http_status, .. } => *http_status, + Self::Connect(_) + | Self::NoHostCredentialProvider + | Self::ProcessEnvironmentRequired { .. } + | Self::CredentialTimeout { .. } + | Self::ExpiredCredential + | Self::MissingCredential { .. } => None, + } + } + + pub(crate) fn kind(&self) -> &'static str { + match self { + Self::Connect(_) => "connection", + Self::NoHostCredentialProvider => "host-credential-provider", + Self::ProcessEnvironmentRequired { .. } => "process-environment-required", + Self::Credential { stage, .. } => match stage { + CredentialStage::Identity => "caller-identity", + CredentialStage::Signing => "token-signing", + CredentialStage::Exchange => "token-exchange", + CredentialStage::Response => "token-response", + CredentialStage::Provider => "credential-provider", + }, + Self::CredentialTimeout { .. } => "acquisition-timeout", + Self::ExpiredCredential => "expired-credential", + Self::MissingCredential { .. } => "missing-credential", + } + } +} + +fn http_status_suffix(status: Option) -> String { + status.map_or_else(String::new, |status| format!(" (HTTP {status})")) +} diff --git a/litellm-rust/crates/redis-connect/src/events.rs b/litellm-rust/crates/redis-connect/src/events.rs new file mode 100644 index 00000000000..a556b7d370b --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/events.rs @@ -0,0 +1,129 @@ +use std::{ + sync::Mutex, + time::{Duration, Instant}, +}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_client::{CallEvent, FailureKind, Outcome, RedisEvents}; + +const DEFAULT_TIMEOUT_LOG_INTERVAL: Duration = Duration::from_secs(5); + +pub struct TracingEvents { + throttle: TimeoutLogThrottle, +} + +impl TracingEvents { + pub fn new(lookup: &dyn Lookup) -> Self { + let interval = lookup + .get("REDIS_TIMEOUT_LOG_INTERVAL") + .and_then(|seconds| seconds.parse::().ok()) + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) + .unwrap_or(DEFAULT_TIMEOUT_LOG_INTERVAL); + Self { + throttle: TimeoutLogThrottle::new(interval), + } + } +} + +impl RedisEvents for TracingEvents { + fn call(&self, event: &CallEvent) { + let Outcome::Failed(kind) = event.outcome else { + return; + }; + match kind { + FailureKind::BreakerOpen => { + tracing::debug!(operation = ?event.operation, "redis call rejected by the open circuit breaker") + } + FailureKind::Timeout => match self.throttle.admit(Instant::now()) { + None => { + tracing::debug!(operation = ?event.operation, duration = ?event.duration, "redis call timed out") + } + Some(0) => { + tracing::error!(operation = ?event.operation, duration = ?event.duration, "redis call timed out") + } + Some(suppressed) => tracing::error!( + operation = ?event.operation, + duration = ?event.duration, + suppressed, + "redis call timed out, more redis timeouts since the previous line were logged at debug" + ), + }, + FailureKind::Transport | FailureKind::Server | FailureKind::Decode => { + tracing::error!(operation = ?event.operation, failure = ?kind, "redis call failed") + } + } + } + + fn breaker(&self, transition: &Transition) { + match transition.to { + State::Open { .. } => tracing::warn!( + generation = transition.generation, + "redis circuit breaker opened, fast-failing redis calls" + ), + State::HalfOpen => tracing::info!( + generation = transition.generation, + "redis circuit breaker half open, probing redis" + ), + State::Closed { .. } => tracing::info!( + generation = transition.generation, + "redis circuit breaker closed, redis recovered" + ), + } + } +} + +struct TimeoutLogThrottle { + interval: Duration, + state: Mutex<(Option, u64)>, +} + +impl TimeoutLogThrottle { + fn new(interval: Duration) -> Self { + Self { + interval, + state: Mutex::new((None, 0)), + } + } + + fn admit(&self, now: Instant) -> Option { + let mut state = self + .state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let (last_logged_at, suppressed) = &mut *state; + if last_logged_at.is_some_and(|last| now.duration_since(last) < self.interval) { + *suppressed += 1; + return None; + } + *last_logged_at = Some(now); + Some(std::mem::take(suppressed)) + } +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use rstest::rstest; + + use super::TimeoutLogThrottle; + + #[rstest] + fn admits_one_line_per_interval_and_counts_the_rest() { + let throttle = TimeoutLogThrottle::new(Duration::from_secs(5)); + let start = Instant::now(); + let at = |seconds: u64| start + Duration::from_secs(seconds); + + let admitted = [ + throttle.admit(at(0)), + throttle.admit(at(1)), + throttle.admit(at(4)), + throttle.admit(at(5)), + throttle.admit(at(6)), + throttle.admit(at(20)), + ]; + + assert_eq!(admitted, [Some(0), None, None, Some(2), None, Some(1)]); + } +} diff --git a/litellm-rust/crates/redis-connect/src/lib.rs b/litellm-rust/crates/redis-connect/src/lib.rs new file mode 100644 index 00000000000..b3d1d411eed --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/lib.rs @@ -0,0 +1,11 @@ +mod connect; +mod error; +mod events; +mod refresh; +mod sources; + +pub use connect::{HostServices, SharedEnvironment, connect}; +pub use error::{CredentialStage, Error}; +pub use events::TracingEvents; +pub(crate) use refresh::RefreshingCredentials; +pub use refresh::{CredentialSource, Expiring, SourceFuture}; diff --git a/litellm-rust/crates/redis-connect/src/refresh.rs b/litellm-rust/crates/redis-connect/src/refresh.rs new file mode 100644 index 00000000000..1e05d3d6289 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/refresh.rs @@ -0,0 +1,783 @@ +use std::{ + future::{Future, pending}, + pin::Pin, + sync::{Arc, Mutex}, + time::{Duration, SystemTime}, +}; + +use futures_util::{Stream, stream}; +use litellm_redis_client::{BasicAuth, RedisError, StreamingCredentialsProvider}; +use litellm_secret_value::SecretValue; +use tokio::{sync::watch, time::Instant}; + +use crate::Error; + +const TOKEN_ONLY_USERNAME: &str = "default"; +const UNKNOWN_LIFETIME: Duration = Duration::from_secs(300); +const MIN_REFRESH: Duration = Duration::from_secs(1); +const ACQUISITION_TIMEOUT: Duration = Duration::from_secs(30); +const FIRST_RETRY: Duration = Duration::from_secs(1); +const MAX_RETRY: Duration = Duration::from_secs(60); + +#[derive(Clone, PartialEq, Eq)] +pub struct Expiring { + pub username: Option, + pub password: SecretValue, + pub expires_at: Option, +} + +pub type SourceFuture<'a> = Pin> + Send + 'a>>; + +pub trait CredentialSource: Send + Sync + 'static { + fn fetch(&self) -> SourceFuture<'_>; +} + +impl CredentialSource for Arc { + fn fetch(&self) -> SourceFuture<'_> { + (**self).fetch() + } +} + +#[derive(Clone)] +enum Latest { + Pending, + Ready { + username: String, + password: SecretValue, + expires_at: Option, + }, + Unavailable(Arc), +} + +impl Latest { + fn is_same_token(&self, other: &Self) -> bool { + matches!( + (self, other), + ( + Self::Ready { username, password, .. }, + Self::Ready { username: other_username, password: other_password, .. }, + ) if username == other_username && password == other_password + ) + } + + fn expires_at(&self) -> Option { + match self { + Self::Ready { expires_at, .. } => *expires_at, + Self::Pending | Self::Unavailable(_) => None, + } + } +} + +type Start = Box; + +#[derive(Clone)] +pub struct RefreshingCredentials { + latest: watch::Receiver, + start: Arc>>, +} + +impl RefreshingCredentials { + pub fn new(source: impl CredentialSource) -> Self { + let (sender, latest) = watch::channel(Latest::Pending); + let start: Start = Box::new(move || drop(tokio::spawn(refresh(source, sender)))); + Self { + latest, + start: Arc::new(Mutex::new(Some(start))), + } + } + + fn start_once(&self) { + let start = self + .start + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take(); + if let Some(start) = start { + start(); + } + } +} + +impl StreamingCredentialsProvider for RefreshingCredentials { + fn subscribe( + &self, + ) -> Pin> + Send + 'static>> { + self.start_once(); + Box::pin(stream::unfold( + (self.latest.clone(), true), + |(mut receiver, first)| async move { + if !first { + receiver.changed().await.ok()?; + } + let latest = receiver + .wait_for(|latest| !matches!(latest, Latest::Pending)) + .await + .ok()? + .clone(); + let item = match latest { + Latest::Ready { + username, password, .. + } => Ok(BasicAuth::new(username, password.expose().to_owned())), + Latest::Unavailable(error) => Err(public_error(&error)), + Latest::Pending => unreachable!("wait_for skips Pending"), + }; + Some((item, (receiver, false))) + }, + )) + } +} + +fn public_error(error: &Error) -> RedisError { + RedisError::from(( + litellm_redis_client::RedisErrorKind::AuthenticationFailed, + "fetching Redis credentials failed", + error.to_string(), + )) +} + +struct Fresh { + latest: Latest, + refresh_in: Duration, +} + +fn fresh(credentials: Expiring, now: SystemTime) -> Result { + let remaining = credentials + .expires_at + .map(|expires_at| expires_at.duration_since(now).unwrap_or_default()); + if remaining == Some(Duration::ZERO) { + return Err(Error::ExpiredCredential); + } + Ok(Fresh { + latest: Latest::Ready { + username: credentials + .username + .unwrap_or_else(|| TOKEN_ONLY_USERNAME.to_owned()), + password: credentials.password, + expires_at: remaining.map(|remaining| Instant::now() + remaining), + }, + refresh_in: refresh_after(credentials.expires_at, now), + }) +} + +async fn until_expiry( + future: impl Future, + sender: &watch::Sender, +) -> Option { + tokio::pin!(future); + loop { + let expires_at = sender.borrow().expires_at(); + let expiration = async { + match expires_at { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => pending().await, + } + }; + tokio::select! { + biased; + () = sender.closed() => return None, + () = expiration => { + sender.send_replace(Latest::Unavailable(Arc::new(Error::ExpiredCredential))); + } + result = &mut future => return Some(result), + } + } +} + +async fn refresh(source: impl CredentialSource, sender: watch::Sender) { + let mut retry = FIRST_RETRY; + loop { + let acquisition = async { + tokio::time::timeout(ACQUISITION_TIMEOUT, source.fetch()) + .await + .map_err(|cause| Error::CredentialTimeout { + timeout: ACQUISITION_TIMEOUT, + cause, + })? + .and_then(|credentials| fresh(credentials, SystemTime::now())) + }; + let Some(result) = until_expiry(acquisition, &sender).await else { + return; + }; + let wait = match result { + Ok(Fresh { latest, refresh_in }) => { + retry = FIRST_RETRY; + sender.send_if_modified(|current| { + let changed = !current.is_same_token(&latest); + *current = latest; + changed + }); + refresh_in + } + Err(error) => { + tracing::warn!(failure = error.kind(), upstream_status = error.upstream_status(), retry_in = ?retry, "fetching Redis credentials failed"); + if !matches!(*sender.borrow(), Latest::Ready { .. }) { + sender.send_replace(Latest::Unavailable(Arc::new(error))); + } + let wait = retry; + retry = (retry * 2).min(MAX_RETRY); + wait + } + }; + if until_expiry(tokio::time::sleep(wait), &sender) + .await + .is_none() + { + return; + } + } +} + +fn refresh_after(expires_at: Option, now: SystemTime) -> Duration { + let remaining = expires_at + .map(|expires_at| expires_at.duration_since(now).unwrap_or_default()) + .unwrap_or(UNKNOWN_LIFETIME); + (remaining * 3 / 4).max(MIN_REFRESH) +} + +#[cfg(test)] +mod tests { + use litellm_redis_client::commands; + use std::{ + collections::BTreeSet, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::{Duration, SystemTime}, + }; + + use futures_util::StreamExt; + use litellm_redis_client::StreamingCredentialsProvider; + use litellm_secret_value::SecretValue; + use rstest::rstest; + + use super::{CredentialSource, Expiring, RefreshingCredentials, SourceFuture, refresh_after}; + use litellm_redis_client::{ + Cluster, Credentials, CredentialsProvider, Namespace, Node, RedisClient, RedisConfig, + RedisUrl, Sentinel, Target, Timeouts, + }; + use litellm_redis_testing::{ + ClientConnection, RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually, + }; + + #[rstest] + #[case::three_quarters_of_the_lifetime( + Some(Duration::from_secs(900)), + Duration::from_secs(675) + )] + #[case::already_expired(Some(Duration::ZERO), Duration::from_secs(1))] + #[case::unknown_expiry(None, Duration::from_secs(225))] + fn refreshes_well_before_expiry( + #[case] lifetime: Option, + #[case] expected: Duration, + ) { + let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000); + + assert_eq!( + refresh_after(lifetime.map(|lifetime| now + lifetime), now), + expected + ); + } + + struct Scripted { + replies: Mutex>>, + calls: Arc, + } + + impl CredentialSource for Scripted { + fn fetch(&self) -> SourceFuture<'_> { + self.calls.fetch_add(1, Ordering::SeqCst); + let reply = self.replies.lock().unwrap().remove(0); + Box::pin(async move { + reply.map_err(|reason| { + crate::Error::credential(crate::error::CredentialStage::Identity)( + std::io::Error::other(reason), + ) + }) + }) + } + } + + fn token(password: &str, lifetime: Duration) -> Result { + Ok(Expiring { + username: None, + password: SecretValue::new(password), + expires_at: Some(SystemTime::now() + lifetime), + }) + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn pushes_each_refresh_and_keeps_the_last_token_through_failures() { + let calls = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + token("first", Duration::from_secs(8)), + Err("token endpoint down"), + token("second", Duration::from_secs(3600)), + ]), + calls: calls.clone(), + }); + let mut stream = provider.subscribe(); + + let first = stream.next().await.unwrap().unwrap(); + let second = stream.next().await.unwrap().unwrap(); + let late = provider.subscribe().next().await.unwrap().unwrap(); + + assert_eq!((first.username(), first.password()), ("default", "first")); + assert_eq!(second.password(), "second"); + assert_eq!(late.password(), "second"); + assert_eq!(calls.load(Ordering::SeqCst), 3); + } + + #[rstest] + fn building_outside_a_runtime_fetches_nothing_until_a_connection_subscribes() { + let calls = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![token("first", Duration::from_secs(3600))]), + calls: calls.clone(), + }); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + let before = calls.load(Ordering::SeqCst); + let first = runtime.block_on(async { provider.subscribe().next().await.unwrap().unwrap() }); + + assert_eq!(before, 0); + assert_eq!(first.password(), "first"); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn a_failure_before_any_token_fails_the_connection_instead_of_hanging() { + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + Err("no identity"), + token("late", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + + assert!(stream.next().await.unwrap().is_err()); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "late"); + } + + struct Blocking { + first: Mutex>, + started: Arc, + cancelled: Arc, + } + + struct Cancellation(Arc); + + impl Drop for Cancellation { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + impl CredentialSource for Blocking { + fn fetch(&self) -> SourceFuture<'_> { + let first = self.first.lock().unwrap().take(); + Box::pin(async move { + if let Some(first) = first { + return Ok(first); + } + let _guard = Cancellation(self.cancelled.clone()); + self.started.notify_one(); + std::future::pending().await + }) + } + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn a_refresh_that_returns_the_same_token_is_not_pushed_again() { + let calls = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + token("same", Duration::from_secs(8)), + token("same", Duration::from_secs(3600)), + ]), + calls: calls.clone(), + }); + let mut stream = provider.subscribe(); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "same"); + + let pushed = tokio::time::timeout(Duration::from_secs(60), stream.next()).await; + + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert!(pushed.is_err(), "an unchanged token was pushed again"); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn expiration_is_published_while_refresh_is_stalled() { + let started = Arc::new(tokio::sync::Notify::new()); + let provider = RefreshingCredentials::new(Blocking { + first: Mutex::new(Some(token("first", Duration::from_secs(4)).unwrap())), + started: started.clone(), + cancelled: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "first"); + let start = tokio::time::Instant::now(); + started.notified().await; + let expired = stream.next().await.unwrap().unwrap_err(); + assert!(start.elapsed() <= Duration::from_secs(5)); + assert!(expired.to_string().contains("expired")); + let fresh = provider.subscribe().next().await.unwrap().unwrap_err(); + assert!(fresh.to_string().contains("expired")); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn expiration_is_published_once_and_later_failures_keep_their_cause() { + let replies = std::iter::once(token("first", Duration::from_secs(8))) + .chain(std::iter::repeat_n(Err("token endpoint down"), 16)) + .collect(); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(replies), + calls: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + assert_eq!(stream.next().await.unwrap().unwrap().password(), "first"); + + let failures: Vec = stream + .take(5) + .map(|item| item.unwrap_err().to_string()) + .collect() + .await; + + assert!(failures[0].contains("expired"), "{failures:?}"); + assert!( + failures[1..] + .iter() + .all(|failure| failure.contains("caller identity") && !failure.contains("expired")), + "{failures:?}" + ); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn initial_acquisition_has_a_bounded_wait() { + let provider = RefreshingCredentials::new(Blocking { + first: Mutex::new(None), + started: Arc::new(tokio::sync::Notify::new()), + cancelled: Arc::new(AtomicUsize::new(0)), + }); + let error = provider.subscribe().next().await.unwrap().unwrap_err(); + assert!(error.to_string().contains("exceeded")); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn dropping_all_subscribers_cancels_inflight_acquisition() { + let started = Arc::new(tokio::sync::Notify::new()); + let cancelled = Arc::new(AtomicUsize::new(0)); + let provider = RefreshingCredentials::new(Blocking { + first: Mutex::new(None), + started: started.clone(), + cancelled: cancelled.clone(), + }); + let stream = provider.subscribe(); + started.notified().await; + drop(stream); + drop(provider); + tokio::task::yield_now().await; + assert_eq!(cancelled.load(Ordering::SeqCst), 1); + } + + #[rstest] + #[tokio::test(start_paused = true)] + async fn an_expired_fetch_is_rejected_and_a_later_valid_token_recovers() { + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + Ok(Expiring { + username: None, + password: SecretValue::new("expired"), + expires_at: Some(SystemTime::UNIX_EPOCH), + }), + token("recovered", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let mut stream = provider.subscribe(); + assert!( + stream + .next() + .await + .unwrap() + .unwrap_err() + .to_string() + .contains("expired") + ); + assert_eq!( + stream.next().await.unwrap().unwrap().password(), + "recovered" + ); + } + + #[rstest] + fn credential_failure_keeps_its_typed_cause_but_hides_it_at_the_redis_boundary() { + let error = crate::Error::credential(crate::error::CredentialStage::Identity)( + std::io::Error::new(std::io::ErrorKind::PermissionDenied, "credential-value"), + ); + assert_eq!( + litellm_error::find::(&error) + .unwrap() + .kind(), + std::io::ErrorKind::PermissionDenied + ); + let debug = format!("{error:?}"); + assert!(!debug.contains("credential-value")); + assert!(debug.contains("[REDACTED]")); + let public = super::public_error(&error).to_string(); + assert!(public.contains("caller identity")); + assert!(!public.contains("credential-value")); + } + + #[rstest] + fn credential_http_failures_expose_status_without_exposing_the_response_body() { + let error = crate::Error::credential_response(crate::CredentialStage::Exchange, 503)( + litellm_http::Error::upstream(503, "credential-value".into()), + ); + let public = super::public_error(&error).to_string(); + assert!(public.contains("HTTP 503")); + assert!(!public.contains("credential-value")); + assert_eq!( + litellm_error::find::(&error) + .unwrap() + .upstream_response() + .unwrap() + .0, + 503 + ); + } + + fn endpoint(deployment: &RedisDeployment) -> Target { + match deployment { + RedisDeployment::Standalone(server) => Target::Url(RedisUrl { + url: SecretValue::new(server.url()), + tls: None, + db: None, + username: None, + password: None, + }), + RedisDeployment::Cluster(cluster) => Target::Cluster(Cluster { + nodes: cluster + .nodes() + .map(|(host, port)| Node { + host: host.to_owned(), + port, + }) + .collect(), + tls: None, + }), + } + } + + fn user(username: &str, password: &str, lifetime: Duration) -> Result { + Ok(Expiring { + username: Some(username.to_owned()), + password: SecretValue::new(password), + expires_at: Some(SystemTime::now() + lifetime), + }) + } + + fn connection_users(connections: &[ClientConnection]) -> BTreeSet<&str> { + connections + .iter() + .map(|connection| connection.user.as_str()) + .collect() + } + + fn connection_ids(connections: &[ClientConnection]) -> BTreeSet<(u16, u64)> { + connections + .iter() + .map(|connection| (connection.node, connection.id)) + .collect() + } + + #[rstest] + #[tokio::test] + async fn a_refreshed_token_reauthenticates_live_connections_before_the_old_one_expires( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + server + .cli_on_every_node(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + user("alice", "first", Duration::from_secs(2)), + user("bob", "second", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let client = RedisClient::new(RedisConfig { + target: endpoint(&server), + credentials: Credentials::Provider(CredentialsProvider(Arc::new(provider))), + namespace: Namespace::default(), + timeouts: Timeouts::default(), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + client + .send(commands::set("k", "1", Duration::from_secs(60))) + .await + .unwrap(); + let before = server.clients().await.unwrap(); + + let after = eventually(async || { + let connections = server.clients().await.ok()?; + (!connections.is_empty() && connection_users(&connections) == BTreeSet::from(["bob"])) + .then_some(connections) + }) + .await; + server + .cli_on_every_node(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + + assert!(!before.is_empty()); + assert_eq!(connection_users(&before), BTreeSet::from(["alice"])); + assert_eq!(connection_ids(&after), connection_ids(&before)); + assert_eq!( + client + .send(commands::get::("k")) + .await + .unwrap() + .as_deref(), + Some("1") + ); + server + .cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"]) + .await + .unwrap(); + let value = eventually(async || { + client + .send(commands::get::("k")) + .await + .ok() + .flatten() + }) + .await; + let reconnected = server.clients().await.unwrap(); + + assert_eq!(value, "1"); + assert!(!reconnected.is_empty()); + assert_eq!(connection_users(&reconnected), BTreeSet::from(["bob"])); + assert!(connection_ids(&reconnected).is_disjoint(&connection_ids(&after))); + } + + /// Cloud credentials rotate while Sentinel fails over: the refreshed user must be what the + /// client presents to the master Sentinel promotes, not the one it connected with first. + #[rstest] + #[tokio::test] + async fn a_refreshed_user_reaches_the_master_sentinel_promotes() { + let sentinel = RedisSentinel::start().await.unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "alice", "on", ">first", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + sentinel + .cli_on_data_nodes(&[ + "ACL", "SETUSER", "bob", "on", ">second", "~*", "&*", "+@all", + ]) + .await + .unwrap(); + let provider = RefreshingCredentials::new(Scripted { + replies: Mutex::new(vec![ + user("alice", "first", Duration::from_secs(2)), + user("bob", "second", Duration::from_secs(3600)), + ]), + calls: Arc::new(AtomicUsize::new(0)), + }); + let (host, port) = sentinel.sentinel(); + let client = RedisClient::new(RedisConfig { + target: Target::Sentinel(Sentinel { + monitors: vec![Node { + host: host.to_owned(), + port, + }], + service_name: SERVICE_NAME.to_owned(), + monitor_auth: None, + tls: None, + db: 0, + }), + credentials: Credentials::Provider(CredentialsProvider(Arc::new(provider))), + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_secs(1), None), + breaker: None, + events: None, + health_check: None, + keepalive: true, + }) + .unwrap(); + client + .send(commands::set("before", "1", Duration::from_secs(60))) + .await + .unwrap(); + let master = sentinel.master_port().await.unwrap(); + let app_users = |connections: &[ClientConnection]| -> BTreeSet { + connections + .iter() + .filter(|connection| connection.user != "default") + .map(|connection| connection.user.clone()) + .collect() + }; + eventually(async || { + let connections = sentinel.clients_on(master).await.ok()?; + (app_users(&connections) == BTreeSet::from(["bob".to_owned()])).then_some(()) + }) + .await; + sentinel + .cli_on_data_nodes(&["ACL", "DELUSER", "alice"]) + .await + .unwrap(); + + let promoted = sentinel.failover().await.unwrap(); + eventually(async || { + client + .send(commands::set("after", "2", Duration::from_secs(60))) + .await + .ok() + }) + .await; + let on_promoted = sentinel.clients_on(promoted).await.unwrap(); + + assert_eq!(promoted, sentinel.replica_port()); + assert_eq!( + app_users(&on_promoted), + BTreeSet::from(["bob".to_owned()]), + "{on_promoted:?}" + ); + assert_eq!( + client + .send(commands::get::("after")) + .await + .unwrap() + .as_deref(), + Some("2") + ); + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/aws.rs b/litellm-rust/crates/redis-connect/src/sources/aws.rs new file mode 100644 index 00000000000..1e87bdd1b12 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/aws.rs @@ -0,0 +1,188 @@ +use std::time::{Duration, SystemTime}; + +use litellm_auth_aws::{AwsAuthService, SigningScope, presign_url}; +use litellm_core_utils::settings::Lookup; +use litellm_redis_types::AwsIamAuth; + +use crate::Error; +use crate::error::CredentialStage; +use litellm_secret_value::SecretValue; + +use crate::{CredentialSource, Expiring, SourceFuture}; + +const SERVICE: &str = "elasticache"; +const TOKEN_LIFETIME: Duration = Duration::from_secs(900); + +const SDK_IDENTITY_ENVIRONMENT: &[&str] = &[ + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_SHARED_CREDENTIALS_FILE", + "AWS_CONFIG_FILE", + "AWS_WEB_IDENTITY_TOKEN_FILE", + "AWS_ROLE_ARN", + "AWS_ROLE_SESSION_NAME", + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", + "AWS_CONTAINER_CREDENTIALS_FULL_URI", + "AWS_CONTAINER_AUTHORIZATION_TOKEN", + "AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE", + "AWS_EC2_METADATA_DISABLED", + "AWS_EC2_METADATA_SERVICE_ENDPOINT", + "AWS_EC2_METADATA_SERVICE_ENDPOINT_MODE", + "AWS_ENDPOINT_URL", + "AWS_ENDPOINT_URL_STS", +]; + +pub(crate) struct ElastiCacheIam { + service: AwsAuthService, + auth: AwsIamAuth, +} + +impl ElastiCacheIam { + pub(crate) fn new(auth: AwsIamAuth, environment: &dyn Lookup) -> Result { + validate_environment(environment, &|name: &str| std::env::var(name).ok())?; + Ok(Self { + service: AwsAuthService::default(), + auth, + }) + } +} + +fn validate_environment(environment: &dyn Lookup, process: &dyn Lookup) -> Result<(), Error> { + match SDK_IDENTITY_ENVIRONMENT.iter().find(|name| { + environment.get(name).is_some_and(|value| { + Some(SecretValue::new(value)) != process.get(name).map(SecretValue::new) + }) + }) { + Some(setting) => Err(Error::ProcessEnvironmentRequired { setting }), + None => Ok(()), + } +} + +pub(crate) fn connect_url(auth: &AwsIamAuth) -> String { + let query = [ + Some(("Action", "connect")), + Some(("User", auth.user_name.as_str())), + auth.serverless + .then_some(("ResourceType", "ServerlessCache")), + ]; + let mut url = String::from("https://"); + url.push_str(&auth.cache_name.to_lowercase()); + url.push_str("/?"); + url.push_str( + &url::form_urlencoded::Serializer::new(String::new()) + .extend_pairs(query.into_iter().flatten()) + .finish(), + ); + url +} + +impl CredentialSource for ElastiCacheIam { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + let credentials = self + .service + .default_chain_credentials(&self.auth.region) + .await + .map_err(Error::credential(CredentialStage::Identity))?; + let now = SystemTime::now(); + let signed = presign_url( + &connect_url(&self.auth), + TOKEN_LIFETIME, + SigningScope { + region: &self.auth.region, + service: SERVICE, + credentials: &credentials, + time: now, + }, + ) + .map_err(Error::credential(CredentialStage::Signing))?; + let token_expiry = now + TOKEN_LIFETIME; + Ok(Expiring { + username: Some(self.auth.user_name.clone()), + password: SecretValue::new(signed.trim_start_matches("https://")), + expires_at: Some( + credentials + .expiry() + .map_or(token_expiry, |expiry| expiry.min(token_expiry)), + ), + }) + }) + } +} + +#[cfg(test)] +mod tests { + use litellm_redis_types::AwsIamAuth; + use rstest::rstest; + + use super::{connect_url, validate_environment}; + + #[rstest] + #[case::profile("AWS_PROFILE")] + #[case::role("AWS_ROLE_ARN")] + #[case::web_identity("AWS_WEB_IDENTITY_TOKEN_FILE")] + #[case::source_key("AWS_ACCESS_KEY_ID")] + #[case::source_secret("AWS_SECRET_ACCESS_KEY")] + #[case::session("AWS_SESSION_TOKEN")] + #[case::profile_path("AWS_SHARED_CREDENTIALS_FILE")] + #[case::container("AWS_CONTAINER_CREDENTIALS_FULL_URI")] + fn configured_identity_overrides_cannot_silently_use_an_ambient_identity( + #[case] setting: &'static str, + ) { + let configured = |name: &str| (name == setting).then(|| "credential-value".to_owned()); + let process = |name: &str| (name == setting).then(|| "ambient".to_owned()); + let error = validate_environment(&configured, &process).unwrap_err(); + + assert!(matches!( + error, + crate::Error::ProcessEnvironmentRequired { setting: rejected } if rejected == setting + )); + assert!(!error.to_string().contains("credential-value")); + assert!(!format!("{error:?}").contains("credential-value")); + } + + #[rstest] + #[case::process_environment("AWS_PROFILE", Some("profile"), Some("profile"))] + #[case::ambient_only("AWS_PROFILE", None, Some("profile"))] + #[case::unconfigured("AWS_PROFILE", None, None)] + #[case::unrelated_bedrock_setting("AWS_ROLE_NAME", Some("bedrock"), None)] + fn unchanged_or_unrelated_settings_preserve_the_sdk_chain( + #[case] setting: &str, + #[case] configured: Option<&str>, + #[case] ambient: Option<&str>, + ) { + let configured = |name: &str| { + (name == setting) + .then_some(configured) + .flatten() + .map(str::to_owned) + }; + let ambient = |name: &str| { + (name == setting) + .then_some(ambient) + .flatten() + .map(str::to_owned) + }; + assert!(validate_environment(&configured, &ambient).is_ok()); + } + + #[rstest] + #[case::provisioned(false, "https://cache-1/?Action=connect&User=iam+user%2F1")] + #[case::serverless( + true, + "https://cache-1/?Action=connect&User=iam+user%2F1&ResourceType=ServerlessCache" + )] + fn the_connect_url_matches_python_urlencode(#[case] serverless: bool, #[case] expected: &str) { + let auth = AwsIamAuth { + user_name: "iam user/1".to_owned(), + cache_name: "Cache-1".to_owned(), + region: "us-east-1".to_owned(), + serverless, + }; + + assert_eq!(connect_url(&auth), expected); + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/azure.rs b/litellm-rust/crates/redis-connect/src/sources/azure.rs new file mode 100644 index 00000000000..6dc72e6ce72 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/azure.rs @@ -0,0 +1,64 @@ +use litellm_auth_azure::{AzureAuthService, AzureIdentity}; +use litellm_auth_types::ResolvedCredential; +use litellm_redis_types::AzureAdAuth; + +use crate::Error; +use crate::error::CredentialStage; + +use crate::{CredentialSource, Expiring, SharedEnvironment, SourceFuture}; + +const AZURE_REDIS_SCOPE: &str = "https://redis.azure.com/.default"; + +pub(crate) struct AzureAd { + service: AzureAuthService, + identity: AzureIdentity, + username: Option, + lookup: SharedEnvironment, +} + +impl AzureAd { + pub(crate) fn new(auth: AzureAdAuth, lookup: SharedEnvironment) -> Self { + // _build_azure_credential in litellm/_redis.py picks the credential class the same way + let identity = match (auth.client_id, auth.tenant_id, auth.client_secret) { + (Some(client_id), Some(tenant_id), Some(client_secret)) => { + AzureIdentity::ClientSecret { + tenant_id, + client_id, + client_secret, + } + } + (Some(client_id), _, _) => AzureIdentity::ManagedIdentity { client_id }, + (None, _, _) => AzureIdentity::Default, + }; + Self { + service: AzureAuthService::default(), + identity, + username: auth.username, + lookup, + } + } +} + +impl CredentialSource for AzureAd { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + let credential = self + .service + .identity_token(&self.identity, AZURE_REDIS_SCOPE, &|name| { + self.lookup.get(name) + }) + .await + .map_err(Error::credential(CredentialStage::Identity))? + .ok_or(Error::MissingCredential { provider: "Azure" })?; + let (token, expires_at) = match credential { + ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on), + ResolvedCredential::Static(token) => (token, None), + }; + Ok(Expiring { + username: self.username.clone(), + password: token, + expires_at, + }) + }) + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/gcp.rs b/litellm-rust/crates/redis-connect/src/sources/gcp.rs new file mode 100644 index 00000000000..9da60fe9c08 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/gcp.rs @@ -0,0 +1,126 @@ +use std::time::SystemTime; + +use litellm_auth_gcp::{CredentialSource as GcpCredentialSource, GcpAuth}; +use serde::Deserialize; + +use crate::Error; +use crate::error::CredentialStage; +use litellm_secret_value::SecretValue; + +use crate::{CredentialSource, Expiring, SharedEnvironment, SourceFuture}; + +const IAM_CREDENTIALS: &str = "https://iamcredentials.googleapis.com/v1"; +const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; + +pub(crate) struct GcpIam { + http: litellm_http::Client, + identity: GcpAuth, + service_account: String, + lookup: SharedEnvironment, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct AccessToken { + access_token: String, + expire_time: String, +} + +impl GcpIam { + pub(crate) fn new( + http: litellm_http::Client, + service_account: String, + lookup: SharedEnvironment, + ) -> Self { + Self { + http, + identity: GcpAuth::default(), + service_account, + lookup, + } + } +} + +impl CredentialSource for GcpIam { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + let lookup = |name: &str| self.lookup.get(name); + let caller = self + .identity + .access_token(GcpCredentialSource::application_default(&lookup)) + .await + .map_err(Error::credential(CredentialStage::Identity))?; + let response = self + .http + .post(format!( + "{IAM_CREDENTIALS}/{}:generateAccessToken", + self.service_account + )) + .bearer_auth(caller.expose()) + .json(&serde_json::json!({ "scope": [CLOUD_PLATFORM_SCOPE] })) + .send() + .await + .map_err(Error::credential(CredentialStage::Exchange))?; + let status = response.status().as_u16(); + let reply: AccessToken = response + .error_for_status() + .map_err(litellm_http::Error::from) + .map_err(Error::credential_response( + CredentialStage::Exchange, + status, + ))? + .json() + .await + .map_err(litellm_http::Error::from) + .map_err(Error::credential(CredentialStage::Response))?; + let expires_at = time::OffsetDateTime::parse( + &reply.expire_time, + &time::format_description::well_known::Rfc3339, + ) + .map(SystemTime::from) + .map_err(Error::credential(CredentialStage::Response))?; + Ok(Expiring { + username: None, + password: SecretValue::new(reply.access_token), + expires_at: Some(expires_at), + }) + }) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::GcpIam; + use crate::CredentialSource; + + #[rstest] + #[tokio::test] + async fn the_caller_identity_comes_from_application_default_credentials_not_vertex_settings() { + let environment = |name: &str| { + match name { + "VERTEXAI_CREDENTIALS" => { + Some(r#"{"type":"service_account","project_id":"vertex"}"#) + } + "GOOGLE_APPLICATION_CREDENTIALS" => Some("/nonexistent/redis-adc.json"), + _ => None, + } + .map(str::to_owned) + }; + let source = GcpIam::new( + litellm_http::Client::for_test(), + "projects/-/serviceAccounts/redis@project.iam.gserviceaccount.com".to_owned(), + std::sync::Arc::new(environment), + ); + + let Err(error) = source.fetch().await else { + panic!("a missing credentials file must fail"); + }; + + assert_eq!( + litellm_error::find::(&error).map(std::io::Error::kind), + Some(std::io::ErrorKind::NotFound) + ); + } +} diff --git a/litellm-rust/crates/redis-connect/src/sources/mod.rs b/litellm-rust/crates/redis-connect/src/sources/mod.rs new file mode 100644 index 00000000000..988ad3e1d23 --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/sources/mod.rs @@ -0,0 +1,7 @@ +mod aws; +mod azure; +mod gcp; + +pub(crate) use aws::ElastiCacheIam; +pub(crate) use azure::AzureAd; +pub(crate) use gcp::GcpIam; diff --git a/litellm-rust/crates/redis-connect/tests/connect.rs b/litellm-rust/crates/redis-connect/tests/connect.rs new file mode 100644 index 00000000000..b1d541f4a5e --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/connect.rs @@ -0,0 +1,230 @@ +use litellm_redis_client::commands; +use std::{collections::BTreeMap, sync::Arc, time::Duration}; + +use litellm_redis_connect::{ + CredentialSource, Error, Expiring, HostServices, SourceFuture, connect, +}; +use litellm_redis_testing::{ + RedisDeployment, RedisSentinel, RedisServer, RedisTlsServer, SERVICE_NAME, Topology, eventually, +}; +use litellm_redis_types::{ProviderHandle, Resolution}; +use litellm_secret_value::SecretValue; +use rstest::rstest; + +// a host hands its credential_provider object over as a handle, so these tests do the same +fn resolved(kwargs: serde_json::Value) -> Resolution { + let provider = kwargs + .get("credential_provider") + .map(|_| ProviderHandle::new(1)); + let kwargs = serde_json::from_value::>(kwargs).unwrap(); + litellm_redis_resolve::interpret(&litellm_redis_resolve::kwargs(kwargs), provider, None) + .unwrap() +} + +fn services() -> HostServices { + HostServices { + events: None, + http: litellm_http::Client::for_test(), + environment: Arc::new(|_: &str| None), + credential_provider: None, + } +} + +struct HostPassword(&'static str); + +impl CredentialSource for HostPassword { + fn fetch(&self) -> SourceFuture<'_> { + Box::pin(async move { + Ok(Expiring { + username: None, + password: SecretValue::new(self.0), + expires_at: None, + }) + }) + } +} + +fn port(url: &str) -> u16 { + url.rsplit(':').next().unwrap().parse().unwrap() +} + +#[rstest] +#[case::trusted_ca(true, true)] +#[case::system_roots_only(false, false)] +#[tokio::test] +async fn python_tls_kwargs_reach_the_server(#[case] with_ca: bool, #[case] connects: bool) { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } + let server = RedisTlsServer::start(false).await.unwrap(); + let kwargs = serde_json::json!({ + "host": server.host(), + "port": server.port(), + "ssl": true, + "ssl_ca_certs": with_ca.then(|| server.files().ca.display().to_string()), + "socket_timeout": 2, + }); + + let pinged = connect(resolved(kwargs), services(), None) + .unwrap() + .ping() + .await; + + assert_eq!(pinged.is_ok(), connects, "{pinged:?}"); +} + +#[rstest] +#[case::right_password("right", true)] +#[case::wrong_password("wrong", false)] +#[tokio::test] +async fn static_credentials_authenticate( + #[case] password: &str, + #[case] connects: bool, + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + server + .cli_on_every_node(&["ACL", "SETUSER", "default", "on", ">right"]) + .await + .unwrap(); + let kwargs = match &server { + RedisDeployment::Standalone(server) => serde_json::json!({ + "host": "127.0.0.1", + "port": port(server.url()), + "password": password, + "socket_timeout": 2, + }), + RedisDeployment::Cluster(cluster) => serde_json::json!({ + "startup_nodes": cluster.nodes().map(|(host, port)| serde_json::json!({ + "host": host, + "port": port, + })).collect::>(), + "password": password, + "socket_timeout": 2, + }), + }; + + let pinged = connect(resolved(kwargs), services(), None) + .unwrap() + .ping() + .await; + + assert_eq!(pinged.is_ok(), connects, "{pinged:?}"); +} + +#[rstest] +#[tokio::test] +async fn python_sentinel_kwargs_follow_the_master_through_a_failover() { + let sentinel = RedisSentinel::start().await.unwrap(); + let (host, sentinel_port) = sentinel.sentinel(); + let kwargs = serde_json::json!({ + "sentinel_nodes": [[host, sentinel_port]], + "service_name": SERVICE_NAME, + "socket_timeout": 1, + }); + let client = connect(resolved(kwargs), services(), None).unwrap(); + client + .send(commands::set("before", "1", Duration::from_secs(600))) + .await + .unwrap(); + + let promoted = sentinel.failover().await.unwrap(); + eventually(async || { + client + .send(commands::set("after", "2", Duration::from_secs(600))) + .await + .ok() + }) + .await; + + assert_eq!(promoted, sentinel.replica_port()); + assert_eq!(sentinel.master_port().await.unwrap(), promoted); + assert_eq!( + client + .send(commands::get::("after")) + .await + .unwrap() + .as_deref(), + Some("2") + ); +} + +#[rstest] +fn a_python_credential_provider_without_a_host_source_fails_at_boot() { + let error = connect( + resolved(serde_json::json!({"host": "redis", "credential_provider": true})), + services(), + None, + ) + .err() + .unwrap(); + + assert!( + matches!(error, Error::NoHostCredentialProvider), + "{error:?}" + ); +} + +#[rstest] +#[case::profile("AWS_PROFILE")] +#[case::role("AWS_ROLE_ARN")] +#[case::web_identity("AWS_WEB_IDENTITY_TOKEN_FILE")] +fn config_only_aws_identities_fail_before_connecting(#[case] setting: &'static str) { + let configured = format!("{}-override", std::env::var(setting).unwrap_or_default()); + let host = HostServices { + environment: Arc::new(move |name: &str| (name == setting).then(|| configured.clone())), + ..services() + }; + let error = connect( + resolved(serde_json::json!({ + "host": "redis", "ssl": true, "aws_iam_auth": true, + "aws_iam_user_name": "user", "aws_iam_cache_name": "cache", + "aws_iam_region": "us-east-1", + })), + host, + None, + ) + .err() + .expect("must not select ambient credentials"); + + assert!(matches!( + error, + Error::ProcessEnvironmentRequired { setting: rejected } if rejected == setting + )); +} + +#[rstest] +#[case::right_password("right", true)] +#[case::wrong_password("wrong", false)] +#[tokio::test] +async fn a_python_credential_provider_authenticates_through_the_host_source( + #[case] password: &'static str, + #[case] connects: bool, +) { + let server = RedisServer::start().await.unwrap(); + server + .cli(&["ACL", "SETUSER", "default", "on", ">right"]) + .await + .unwrap(); + let kwargs = serde_json::json!({ + "host": "127.0.0.1", + "port": port(server.url()), + "password": "ignored", + "credential_provider": true, + "socket_timeout": 2, + }); + + let pinged = connect( + resolved(kwargs), + HostServices { + credential_provider: Some(Arc::new(HostPassword(password))), + ..services() + }, + None, + ) + .unwrap() + .ping() + .await; + + assert_eq!(pinged.is_ok(), connects, "{pinged:?}"); +} diff --git a/litellm-rust/crates/redis-connect/tests/live/aws.rs b/litellm-rust/crates/redis-connect/tests/live/aws.rs new file mode 100644 index 00000000000..da5ec323392 --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/aws.rs @@ -0,0 +1,17 @@ +use rstest::{fixture, rstest}; + +use super::{LiveConfig, config, roundtrip}; + +#[fixture] +fn live_config() -> Option { + config("LITELLM_REDIS_AWS_LIVE_CONFIG") +} + +#[rstest] +#[tokio::test] +async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option) { + let Some(config) = live_config else { + return; + }; + roundtrip(config).await; +} diff --git a/litellm-rust/crates/redis-connect/tests/live/azure.rs b/litellm-rust/crates/redis-connect/tests/live/azure.rs new file mode 100644 index 00000000000..a7f678edef4 --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/azure.rs @@ -0,0 +1,17 @@ +use rstest::{fixture, rstest}; + +use super::{LiveConfig, config, roundtrip}; + +#[fixture] +fn live_config() -> Option { + config("LITELLM_REDIS_AZURE_LIVE_CONFIG") +} + +#[rstest] +#[tokio::test] +async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option) { + let Some(config) = live_config else { + return; + }; + roundtrip(config).await; +} diff --git a/litellm-rust/crates/redis-connect/tests/live/gcp.rs b/litellm-rust/crates/redis-connect/tests/live/gcp.rs new file mode 100644 index 00000000000..0965a7a9e9a --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/gcp.rs @@ -0,0 +1,17 @@ +use rstest::{fixture, rstest}; + +use super::{LiveConfig, config, roundtrip}; + +#[fixture] +fn live_config() -> Option { + config("LITELLM_REDIS_GCP_LIVE_CONFIG") +} + +#[rstest] +#[tokio::test] +async fn credentials_authenticate_existing_and_fresh_connections(live_config: Option) { + let Some(config) = live_config else { + return; + }; + roundtrip(config).await; +} diff --git a/litellm-rust/crates/redis-connect/tests/live/main.rs b/litellm-rust/crates/redis-connect/tests/live/main.rs new file mode 100644 index 00000000000..90e7637530e --- /dev/null +++ b/litellm-rust/crates/redis-connect/tests/live/main.rs @@ -0,0 +1,111 @@ +mod aws; +mod azure; +mod gcp; + +use std::{collections::BTreeMap, path::PathBuf, sync::Arc, time::Duration}; + +use litellm_http::{ + ClientVariant, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, + media::PublicDnsResolver, +}; +use litellm_redis_client::commands; +use litellm_redis_connect::{HostServices, connect}; +use serde::Deserialize; + +#[derive(Deserialize)] +struct LiveConfig { + redis: BTreeMap, + hold_seconds: u64, +} + +fn config(variable: &str) -> Option { + let path = PathBuf::from(std::env::var_os(variable)?); + Some( + serde_json::from_slice(&std::fs::read(path).expect("read live Redis config")) + .expect("parse live Redis config"), + ) +} + +async fn roundtrip(config: LiveConfig) { + if rustls::crypto::CryptoProvider::get_default().is_none() { + let _ = rustls::crypto::ring::default_provider().install_default(); + } + let lookup = |name: &str| std::env::var(name).ok(); + let resolved = litellm_redis_resolve::interpret( + &litellm_redis_resolve::kwargs(config.redis), + None, + lookup("REDIS_USERNAME").as_deref(), + ) + .expect("resolve live Redis config"); + let http_settings = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&lookup)]); + let pool = HttpClientPool::new(Arc::new(PublicDnsResolver)); + let http = pool + .client( + &Resolution::from(&http_settings).config, + ClientVariant::Provider, + ) + .expect("build production HTTP client"); + let services = || HostServices { + events: None, + http: http.clone(), + environment: Arc::new(|name: &str| std::env::var(name).ok()), + credential_provider: None, + }; + let client = connect(resolved.clone(), services(), None).expect("build live Redis client"); + let marker = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let key = format!("litellm:redis-connect:live:{}:{marker}", std::process::id()); + let ttl = Duration::from_secs(config.hold_seconds + 120); + + client.ping().await.expect("authenticate and ping"); + client + .send(commands::set(&key, "before-refresh", ttl)) + .await + .expect("set before refresh"); + assert_eq!( + client + .send(commands::get::(&key)) + .await + .unwrap() + .as_deref(), + Some("before-refresh") + ); + + tokio::time::sleep(Duration::from_secs(config.hold_seconds)).await; + + assert_eq!( + client + .send(commands::get::(&key)) + .await + .unwrap() + .as_deref(), + Some("before-refresh") + ); + let fresh = connect(resolved, services(), None).expect("build a fresh live Redis client"); + assert_eq!( + fresh + .send(commands::get::(&key)) + .await + .unwrap() + .as_deref(), + Some("before-refresh") + ); + fresh + .send(commands::set(&key, "after-refresh", ttl)) + .await + .expect("set after refresh"); + assert_eq!( + client + .send(commands::get::(&key)) + .await + .unwrap() + .as_deref(), + Some("after-refresh") + ); + client + .send(commands::del(&[&key])) + .await + .expect("delete live test key"); +} diff --git a/litellm-rust/crates/redis-resolve/AGENTS.md b/litellm-rust/crates/redis-resolve/AGENTS.md new file mode 100644 index 00000000000..b7ce8278975 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/AGENTS.md @@ -0,0 +1,31 @@ +## Interpreting merged kwargs + +- The boundary is merged redis-py kwargs in, `litellm_redis_types::Resolution` out. `interpret` takes the dict a host already merged the way `_merge_redis_kwargs` in `litellm/_redis.py` does, with its `REDIS_*` values and `get_secret` fallbacks filled, and stands in for redis-py's constructors and `from_url`. It never reads the environment or a secret, so both hosts share it. Every `Entry` keeps the source a warning or error names, which `kwargs` sets to the kwarg name and the gateway merge sets to `REDIS_` or the fallback name +- The `Resolution` holds the target, the auth and the timeouts Rust models, and every setting the Rust client ignores as an `Ignored` with its source and consequence, which the host logs. A host passes its `credential_provider` object as a `ProviderHandle`, and a `credential_provider` key without one is `Error::ProviderWithoutHandle`. Python's Azure AD auth reads `REDIS_USERNAME` from `os.environ` instead of the merged kwargs, so it is a required argument of `interpret` that each host reads from its own process environment, the gateway through its `Lookup` and python-bridge through Python's `os.environ`. An empty one is unset, like `or None` +- `options.rs` reads the kwargs into the typed `Options` and collects the ignored ones. `endpoint.rs`, `auth.rs` and `tls.rs` pick the target, the credentials and TLS from `Options` alone +- Rust never introspects redis-py. Every Redis kwarg redis-py or `litellm/_redis.py` accepts is one row in `SETTINGS` in `src/keys.rs`, and that row is the whole answer for the kwarg: its `Handling` (`Read`, `Warn` with the consequence the warning names, or `Reject`), whether a url query carries it (`FromUrl`), whether `_get_redis_env_kwarg_mapping` reads it as `REDIS_`, and the names `_merge_redis_kwargs` falls back to. The last two columns are only read by the gateway merge, and they stay on the row so there is one table. A `Read` row is a field of `Options`, which an inline test pins. A kwarg with no row is not a Redis setting, so it is dropped without a warning +- `redis-py-kwargs.json` holds the Python facts each row answers: which redis-py versions name the kwarg in which signature, and whether Python reads its `REDIS_*` name under the locked redis-py. Never edit it by hand. `redis_py_kwargs.py` next to it regenerates it, and its docstring says how to run it. A kwarg litellm reads that no signature or `REDIS_*` name covers goes in its `LITELLM_ONLY`, and a new redis-py version in its `VERSIONS` and the `redis-compat` CI matrix together. `test_redis_py_kwargs_catalog_matches_the_installed_redis_py` in `tests/unit/test_redis.py` checks it against each version the `redis-compat` CI matrix installs, and `tests/catalog.rs` fails until `SETTINGS` has exactly one row per kwarg in it with the same `environment`. A redis-py bump that adds a kwarg therefore cannot land without a row deciding how Rust handles it +- The contract is `resolves_like_python_redis` in `tests/interpret.rs` and its mirror `test_redis_settings_resolve_like_the_rust_resolver_expects` in `tests/unit/test_redis.py`. Each case holds the merged dict, the sources the merge gave it, the `REDIS_USERNAME` the host passes, and the expected outcome written next to it. `Resolves` is the view both sides build. `Refuses` is a config Rust refuses and Python connects with, a disagreement on purpose. `Fails` is a config redis-py 5.3.1 already raises on, building the client or its first connection, so Rust failing changes no deployment that works today. Its Rust side names the class `Error::raised_building_the_python_client` answers, `None` when Python only raises once a command connects, and its Python mirror `_Fails` asserts the class and message redis-py raises. The Python side asserts that `_merge_redis_kwargs` returns that merged dict and reads the effective redis-py 5.3.1 connection objects without connecting, and `merges_like_python_redis` in `gateway/src/redis/surface_tests.rs` pins the same merged dicts from the Rust merge. A new rule is a new case in every file, never a generated fixture. Python and Rust pods share one Redis during the switch, so a change here is a change to a wire contract +- A url carries its own TLS. redis-py's url allowlist is the plain `Connection` signature, so every `ssl_*` setting outside the url is dropped, and a credential provider strips the url's userinfo. A `rediss://` url's `ssl_cert_reqs`, `ssl_ca_certs`, `ssl_ca_data`, `ssl_ca_path`, `ssl_certfile`, `ssl_keyfile`, `ssl_password` and `ssl_check_hostname` query parameters are read like `parse_url` reads them +- TLS options and `db` are only read for the endpoint that uses them, so a bad `ssl_cert_reqs` never fails a url target. A cluster with a non-zero `db`, and Sentinel with a `host` or `port`, are errors because redis-py raises building those clients +- An empty value (`""`, `[]`, `{}`) is the same as an unset one, like Python's `x or ...`. A null is unset too, except `ssl_cert_reqs`, where redis-py reads `None` as `CERT_NONE`, so a null skips certificate verification +- `ssl_cert_reqs` takes `none`, `optional` and `required` case-sensitively like redis-py's lookup, so `None` (what `REDIS_SSL_CERT_REQS=None` gives) is `Error::Invalid`, and an `ssl.VerifyMode` int, where 0 skips verification and 1 or 2 verifies. `optional` verifies, because rustls has no optional mode +- Precedence matches Python. The target is cluster, then URL, then Sentinel (`sentinel_nodes` plus `service_name`), then host. Auth is a `credential_provider` handle, then GCP, then Azure, then AWS, and the first one wins with a warning +- Numbers are read with `litellm_python_compat::number`, so whitespace and digit underscores parse like Python's `int()` and `float()`. A value they cannot parse (`port: abc`, `socket_timeout: slow`) is dropped with a warning and the default is used, like `_coerce_redis_kwargs_types`. A number that parses but is out of range (negative seconds, a port above 65535) fails with `Error::Invalid` +- Flags (`ssl`, `ssl_check_hostname`, `aws_iam_auth`, `aws_iam_serverless`) are read one way on every surface. `true`, `1` and `yes` are true and `false`, `0`, `no`, `off`, `f` and `n` are false, the text redis-py's `_str_to_bool` and the proxy's pydantic models agree on. Any other text is `Error::Invalid`, so a TLS setting is never guessed +- The resolver fails on a setting it cannot honor only when that setting is a security check. A `Warn` row is dropped and returned as an `Ignored` naming its source (`ssl_ciphers` or `REDIS_SSL_CIPHERS`) and the consequence on its row, because a Python deployment that worked must keep working when its stores move to Rust. Only a value that is malformed (`Error::Invalid` or `Error::UrlQuery`), a target that is missing or a `Reject` row that asks for its check fails. Promoting a `Warn` row to `Read` is then a pure improvement with no behavior cliff, while moving one to `Reject` can fail a config that connects today + - A `Warn` row with a `met` check is dropped silently when its value asks for nothing more than Rust does, so `ssl_min_version` of TLS 1.2 or lower never warns + - A `Reject` row is a security check rustls cannot perform that Python's locked async client already refuses with a `TypeError`, so failing closed breaks no proxy that runs today. `tests/catalog.rs` fails if the locked redis-py's async, cluster or connection signature accepts a `Reject` row, so a redis-py bump forces the decision again. A row's `Asked` says which values ask for the check, from the kwargs, `REDIS_*` or a url query, and those are `Error::Unenforceable` + - `ssl_validate_ocsp` and `ssl_validate_ocsp_stapled` ask when true, because rustls has no client-side OCSP. False asks for nothing and is dropped without a warning + - `ssl_include_verify_flags` and `ssl_exclude_verify_flags` ask when set, because rustls takes no X.509 verify flags such as CRL checking. An empty list asks for nothing + - `ssl_ca_data` and `ssl_ca_path` are read into `TlsOptions` and trusted next to `ssl_ca_certs`. `ssl_password` is read into `TlsOptions::key_password`, which `redis-client` uses to decrypt a PKCS#8 key like the sync client's `load_cert_chain(password=...)`. redis-py 5.3.1's async client refuses it, so this is a superset of the path the proxy uses, like `ssl_ca_path` + - `unix_socket_path` resolves to `Target::Unix`, which like redis-py's `UnixDomainSocketConnection` replaces host, port and TLS. Python still requires `host` next to it, so Rust does too + - A url's `?db=`, `?username=`, `?password=`, `?socket_timeout=`, `?socket_connect_timeout=` and `?health_check_interval=` are read like `parse_url` reads them, so `?db=` wins over the path, the userinfo wins over `?username=` and `?password=`, and the timeouts and health check win over the kwargs like `from_url`. A repeated parameter takes its first value like `parse_qs`. A number `int()` or `float()` cannot read is `Error::UrlQuery`, as `parse_url` raises `ValueError`. Any other query parameter, including one with no row, is ignored naming `url ?`, because redis-py turns each one into a client kwarg. Which ones a url carries is the `FromUrl` column + - `socket_keepalive` is a flag that defaults to true on every target, like `_get_redis_cluster_kwargs` defaults it for a cluster. Python leaves it off for a standalone or Sentinel client, so an unset value there is a difference on purpose, because without keepalive a load balancer drops an idle connection silently. False turns it off. `socket_keepalive_options` is warned, and `redis-client` probes after its own idle time instead + - `health_check_interval` is read as an `int` of seconds like `_coerce_redis_kwargs_types` into `HealthCheck`. Zero is `Disabled` and stays off on every target, a negative interval checks before every command like redis-py, and only `Unset` takes a default. `redis-connect` gives an unset cluster Python's 25s default and `redis-client` pings every node before the next command once the client sat idle that long, outside the breaker and the call events +- `Error::raised_building_the_python_client` is how python-bridge treats a config Rust cannot use. Where Python raises building its client, the bridge raises the same class from `NativeRedisClient.connect`. Where Python builds the client and only fails once a command connects, which `RedisCache` catches and logs on every call, the stores stay on memory and the bridge logs an error +- Cases where Rust and Python disagree on purpose are cases in both contract tests with a different expected outcome on each side + - `ssl_ca_path`: redis-py 5.3.1's async `Redis` has no such parameter and raises `TypeError` building the client, while Rust trusts the directory like `load_verify_locations(capath=...)` + - `on`, `t` and `y` for a flag are `Error::Invalid` on every surface. pydantic reads them as true for `coordination_redis` and `_str_to_bool` reads them as false everywhere else, so Rust refuses to pick one. Other unrecognized text is false in Python and an error here + - `ssl_cert_reqs` as an `ssl.VerifyMode` int: redis-py 5.3.1's sync client takes it, while its async `RedisSSLContext` leaves `cert_reqs` unset and raises `AttributeError` once a command connects. Rust reads it like the sync client + - Empty `sentinel_nodes` with no `REDIS_SENTINEL_NODES` is no target at all, so it fails as "either host or url must be specified". redis-py keeps the empty list and fails building the Sentinel client diff --git a/litellm-rust/crates/redis-resolve/Cargo.toml b/litellm-rust/crates/redis-resolve/Cargo.toml new file mode 100644 index 00000000000..b3d3fc01e44 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "litellm-redis-resolve" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-python-compat.workspace = true +litellm-redis-types.workspace = true +litellm-secret-value.workspace = true +serde_json.workspace = true +strum.workspace = true +thiserror.workspace = true +tracing.workspace = true +url.workspace = true + +[dev-dependencies] +rstest.workspace = true +serde.workspace = true diff --git a/litellm-rust/crates/redis-resolve/redis-py-kwargs.json b/litellm-rust/crates/redis-resolve/redis-py-kwargs.json new file mode 100644 index 00000000000..22c4495818b --- /dev/null +++ b/litellm-rust/crates/redis-resolve/redis-py-kwargs.json @@ -0,0 +1,118 @@ +{ + "meta": { + "description": "What redis-py and litellm/_redis.py accept for every Redis kwarg the Rust resolver can be handed. SETTINGS in src/keys.rs has exactly one row per arg here", + "redis_py_versions": [ + "5.3.1", + "6.4.0", + "7.4.1", + "8.0.1" + ], + "python": "3.12", + "method": "python redis_py_kwargs.py next to this file, see its docstring. It reads inspect.getfullargspec(inspect.unwrap(cls.__init__)) args and kwonlyargs, minus self", + "signatures": { + "sync": "redis.Redis.__init__. _get_redis_kwargs uses it as the config allowlist and the REDIS_ environment mapping", + "async": "redis.asyncio.Redis.__init__, which the proxy constructs. A kwarg in sync but not async raises TypeError when the client is built", + "connection": "redis.Connection or redis.asyncio.Connection MRO, the url allowlist from _get_redis_url_kwargs", + "cluster": "redis.asyncio.RedisCluster.__init__, the allowlist from _get_redis_cluster_kwargs" + }, + "environment": "true when _get_redis_env_kwarg_mapping reads the arg as REDIS_ under the locked redis-py", + "redis_py": "each redis-py version whose signatures name the arg, with which ones. An empty object means only litellm reads it", + "checked_by": "tests/unit/test_redis.py checks the installed redis-py's column and the REDIS_* names against this file, and tests/catalog.rs checks SETTINGS against it", + "redis_py_locked": "5.3.1" + }, + "args": { + "address_remap": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "auto_close_connection_pool": {"environment": false, "redis_py": {"5.3.1": ["async"], "6.4.0": ["async"], "7.4.1": ["async"], "8.0.1": ["async"]}}, + "aws_iam_auth": {"environment": true, "redis_py": {}}, + "aws_iam_cache_name": {"environment": true, "redis_py": {}}, + "aws_iam_region": {"environment": true, "redis_py": {}}, + "aws_iam_serverless": {"environment": true, "redis_py": {}}, + "aws_iam_user_name": {"environment": true, "redis_py": {}}, + "azure_client_id": {"environment": true, "redis_py": {}}, + "azure_client_secret": {"environment": true, "redis_py": {}}, + "azure_redis_ad_token": {"environment": true, "redis_py": {}}, + "azure_tenant_id": {"environment": true, "redis_py": {}}, + "cache": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "cache_config": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "charset": {"environment": true, "redis_py": {"5.3.1": ["sync"]}}, + "client_name": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "cluster_error_retry_attempts": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "command_packer": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "connection_class": {"environment": false, "redis_py": {}}, + "connection_error_retry_attempts": {"environment": false, "redis_py": {"5.3.1": ["cluster"]}}, + "connection_pool": {"environment": false, "redis_py": {"5.3.1": ["sync", "async"], "6.4.0": ["sync", "async"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "credential_provider": {"environment": false, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "db": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "decode_responses": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "driver_info": {"environment": false, "redis_py": {"7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "dynamic_startup_nodes": {"environment": false, "redis_py": {"6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "encoder_class": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "encoding": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "encoding_errors": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "errors": {"environment": true, "redis_py": {"5.3.1": ["sync"]}}, + "event_dispatcher": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "gcp_service_account": {"environment": true, "redis_py": {}}, + "gcp_ssl_ca_certs": {"environment": true, "redis_py": {}}, + "health_check_interval": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "host": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "legacy_responses": {"environment": false, "redis_py": {"8.0.1": ["sync", "async", "connection", "cluster"]}}, + "lib_name": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "lib_version": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "load_balancing_strategy": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "maint_notifications_config": {"environment": false, "redis_py": {"7.4.1": ["sync", "connection"], "8.0.1": ["sync", "connection"]}}, + "maint_notifications_pool_handler": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "maintenance_notification_hash": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "maintenance_state": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "max_connections": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "namespace": {"environment": false, "redis_py": {}}, + "orig_host_address": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "orig_socket_connect_timeout": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "orig_socket_timeout": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "oss_cluster_maint_notifications_handler": {"environment": false, "redis_py": {"7.4.1": ["sync", "connection"], "8.0.1": ["sync", "connection"]}}, + "parser": {"environment": false, "redis_py": {"7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "parser_class": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "password": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "path": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "policy_resolver": {"environment": false, "redis_py": {"7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "port": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "protocol": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "read_from_replicas": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "redis_connect_func": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection"], "6.4.0": ["sync", "async", "connection"], "7.4.1": ["sync", "async", "connection"], "8.0.1": ["sync", "async", "connection"]}}, + "reinitialize_steps": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "require_full_coverage": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "retry": {"environment": false, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "retry_on_error": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "retry_on_timeout": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection"], "6.4.0": ["sync", "async", "connection"], "7.4.1": ["sync", "async", "connection"], "8.0.1": ["sync", "async", "connection"]}}, + "sentinel_nodes": {"environment": false, "redis_py": {}}, + "sentinel_password": {"environment": false, "redis_py": {}}, + "service_name": {"environment": false, "redis_py": {}}, + "single_connection_client": {"environment": true, "redis_py": {"5.3.1": ["sync", "async"], "6.4.0": ["sync", "async"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "socket_connect_timeout": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_keepalive": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_keepalive_options": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_read_size": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_timeout": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}}, + "socket_type": {"environment": false, "redis_py": {"5.3.1": ["connection"], "6.4.0": ["connection"], "7.4.1": ["connection"], "8.0.1": ["connection"]}}, + "ssl": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ca_certs": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ca_data": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ca_path": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "ssl_cert_reqs": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_certfile": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_check_hostname": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ciphers": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_exclude_verify_flags": {"environment": false, "redis_py": {"7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_include_verify_flags": {"environment": false, "redis_py": {"7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_keyfile": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_min_version": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "cluster"], "6.4.0": ["sync", "async", "cluster"], "7.4.1": ["sync", "async", "cluster"], "8.0.1": ["sync", "async", "cluster"]}}, + "ssl_ocsp_context": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "ssl_ocsp_expected_cert": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "ssl_password": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "ssl_validate_ocsp": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "ssl_validate_ocsp_stapled": {"environment": true, "redis_py": {"5.3.1": ["sync"], "6.4.0": ["sync"], "7.4.1": ["sync"], "8.0.1": ["sync"]}}, + "startup_nodes": {"environment": false, "redis_py": {"5.3.1": ["cluster"], "6.4.0": ["cluster"], "7.4.1": ["cluster"], "8.0.1": ["cluster"]}}, + "unix_socket_path": {"environment": true, "redis_py": {"5.3.1": ["sync", "async"], "6.4.0": ["sync", "async"], "7.4.1": ["sync", "async"], "8.0.1": ["sync", "async"]}}, + "url": {"environment": true, "redis_py": {}}, + "username": {"environment": true, "redis_py": {"5.3.1": ["sync", "async", "connection", "cluster"], "6.4.0": ["sync", "async", "connection", "cluster"], "7.4.1": ["sync", "async", "connection", "cluster"], "8.0.1": ["sync", "async", "connection", "cluster"]}} + } +} diff --git a/litellm-rust/crates/redis-resolve/redis_py_kwargs.py b/litellm-rust/crates/redis-resolve/redis_py_kwargs.py new file mode 100644 index 00000000000..9459cfd3761 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/redis_py_kwargs.py @@ -0,0 +1,110 @@ +"""Regenerates redis-py-kwargs.json, the Python facts SETTINGS in src/keys.rs answers. + +Run it from the repository root in the project environment, which has the locked redis-py and litellm: + + uv run python litellm-rust/crates/redis-resolve/redis_py_kwargs.py + +It reads every other redis-py version's signatures in a throwaway environment through +`uv run --no-project --with redis==`, so nothing else is installed. Then update SETTINGS until +tests/catalog.rs passes. +""" + +import inspect +import json +import subprocess +import sys +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +VERSIONS: Final = ("5.3.1", "6.4.0", "7.4.1", "8.0.1") +PYTHON: Final = "3.12" +CATALOG: Final = Path(__file__).with_name("redis-py-kwargs.json") +# Keys litellm/_redis.py or its callers read that neither a redis-py signature nor the REDIS_* mapping names +LITELLM_ONLY: Final = frozenset( + {"connection_class", "namespace", "sentinel_nodes", "sentinel_password", "service_name"} +) + +META: Final = { + "description": "What redis-py and litellm/_redis.py accept for every Redis kwarg the Rust resolver can be handed. SETTINGS in src/keys.rs has exactly one row per arg here", + "redis_py_versions": list(VERSIONS), + "python": PYTHON, + "method": "python redis_py_kwargs.py next to this file, see its docstring. It reads inspect.getfullargspec(inspect.unwrap(cls.__init__)) args and kwonlyargs, minus self", + "signatures": { + "sync": "redis.Redis.__init__. _get_redis_kwargs uses it as the config allowlist and the REDIS_ environment mapping", + "async": "redis.asyncio.Redis.__init__, which the proxy constructs. A kwarg in sync but not async raises TypeError when the client is built", + "connection": "redis.Connection or redis.asyncio.Connection MRO, the url allowlist from _get_redis_url_kwargs", + "cluster": "redis.asyncio.RedisCluster.__init__, the allowlist from _get_redis_cluster_kwargs", + }, + "environment": "true when _get_redis_env_kwarg_mapping reads the arg as REDIS_ under the locked redis-py", + "redis_py": "each redis-py version whose signatures name the arg, with which ones. An empty object means only litellm reads it", + "checked_by": "tests/unit/test_redis.py checks the installed redis-py's column and the REDIS_* names against this file, and tests/catalog.rs checks SETTINGS against it", +} + + +def _named_args(*classes: type) -> frozenset[str]: + return frozenset( + name + for cls in classes + for spec in (inspect.getfullargspec(inspect.unwrap(cls.__init__)),) + for name in spec.args + spec.kwonlyargs + if name != "self" + ) + + +def arg_signatures() -> dict[str, list[str]]: + """Each kwarg the installed redis-py names, with the signatures that name it.""" + import redis + import redis.asyncio as async_redis + + connection_classes: Final = (*redis.Connection.__mro__, *async_redis.Connection.__mro__) + signatures: Final = { + "sync": _named_args(redis.Redis), + "async": _named_args(async_redis.Redis), + "connection": _named_args(*(cls for cls in connection_classes if cls is not object)), + "cluster": _named_args(async_redis.RedisCluster), + } + names: Final = sorted(frozenset().union(*signatures.values())) + return {name: [kind for kind, named in signatures.items() if name in named] for name in names} + + +def _arg_signatures_of(version: str) -> Mapping[str, list[str]]: + command: Final = ("uv", "run", "--no-project", "--python", PYTHON, "--with", f"redis=={version}") + script: Final = ("python", "-I", __file__, "--arg-signatures") + return json.loads(subprocess.run((*command, *script), check=True, capture_output=True, text=True).stdout) + + +def catalog() -> dict[str, object]: + import redis + + from litellm._redis import _get_redis_env_kwarg_mapping + + environment: Final = frozenset(_get_redis_env_kwarg_mapping().values()) + by_version: Final = {version: _arg_signatures_of(version) for version in VERSIONS} + names: Final = sorted(LITELLM_ONLY | environment | frozenset().union(*by_version.values())) + return { + "meta": {**META, "redis_py_locked": redis.__version__}, + "args": { + name: { + "environment": name in environment, + "redis_py": {version: args[name] for version, args in by_version.items() if name in args}, + } + for name in names + }, + } + + +# one line per arg, so a version bump reads as a diff of the args it touched +def render(catalog: Mapping[str, object]) -> str: + meta: Final = json.dumps(catalog["meta"], indent=2).replace("\n", "\n ") + args: Final = catalog["args"] + assert isinstance(args, dict) + rows: Final = ",\n".join(f" {json.dumps(name)}: {json.dumps(arg)}" for name, arg in args.items()) + return f'{{\n "meta": {meta},\n "args": {{\n{rows}\n }}\n}}\n' + + +if __name__ == "__main__": + if sys.argv[1:] == ["--arg-signatures"]: + sys.stdout.write(json.dumps(arg_signatures())) + else: + CATALOG.write_text(render(catalog())) diff --git a/litellm-rust/crates/redis-resolve/src/auth.rs b/litellm-rust/crates/redis-resolve/src/auth.rs new file mode 100644 index 00000000000..620f1a2e11b --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/auth.rs @@ -0,0 +1,95 @@ +use litellm_redis_types::{Auth, AwsIamAuth, AzureAdAuth, ProviderHandle, StaticAuth}; + +use crate::{ + Error, + options::{self, AwsOptions, Options, flag, secret, text}, +}; + +pub(super) fn redis_auth( + options: &Options, + provider: Option, + redis_username: Option<&str>, +) -> Result { + match (provider, &options.credential_provider) { + (Some(handle), _) => return Ok(Auth::Provider(handle)), + (None, Some(_)) => return Err(Error::ProviderWithoutHandle), + (None, None) => {} + } + let gcp = text(options.gcp.service_account.as_ref()); + let azure = text(options.azure.ad_token.as_ref()) + .is_some_and(|enabled| enabled.eq_ignore_ascii_case("true")); + let aws = flag(options.aws.auth.as_ref())?; + match (gcp, azure, aws) { + (Some(service_account), azure, aws) => { + if azure { + tracing::warn!( + "Both GCP IAM (gcp_service_account) and Azure AD (azure_redis_ad_token) are configured for Redis. Using GCP IAM. Remove one to avoid misconfiguration." + ); + } + if aws { + tracing::warn!( + "Both GCP IAM (gcp_service_account) and AWS ElastiCache IAM (aws_iam_auth) are configured for Redis. Using GCP IAM. Remove one to avoid misconfiguration." + ); + } + Ok(Auth::GcpIam { service_account }) + } + (None, true, aws) => { + if aws { + tracing::warn!( + "Both Azure AD (azure_redis_ad_token) and AWS ElastiCache IAM (aws_iam_auth) are configured for Redis. Using Azure AD. Remove one to avoid misconfiguration." + ); + } + Ok(Auth::AzureAd(AzureAdAuth { + client_id: text(options.azure.client_id.as_ref()), + tenant_id: text(options.azure.tenant_id.as_ref()), + client_secret: secret(options.azure.client_secret.as_ref()), + username: redis_username + .filter(|username| !username.is_empty()) + .map(str::to_owned), + })) + } + (None, false, true) => aws_iam(&options.aws, uses_tls(options)?).map(Auth::AwsIam), + (None, false, false) => Ok(Auth::Static(StaticAuth { + username: secret(options.username.as_ref()), + password: secret(options.password.as_ref()), + })), + } +} + +// _uses_tls: a url's scheme decides unless cluster nodes are set, otherwise the ssl kwarg does +fn uses_tls(options: &Options) -> Result { + match (&options.startup_nodes, text(options.url.as_ref())) { + (None, Some(url)) => Ok(url + .split_once("://") + .is_some_and(|(scheme, _)| scheme.eq_ignore_ascii_case("rediss"))), + _ => flag(options.ssl.as_ref()), + } +} + +fn aws_iam(aws: &AwsOptions, uses_tls: bool) -> Result { + if !uses_tls { + return Err(Error::AwsIamWithoutTls); + } + match ( + text(aws.user_name.as_ref()), + text(aws.cache_name.as_ref()), + text(aws.region.as_ref()), + ) { + (Some(user_name), Some(cache_name), Some(region)) => Ok(AwsIamAuth { + user_name, + cache_name, + region, + serverless: options::flag(aws.serverless.as_ref())?, + }), + (user_name, cache_name, region) => Err(Error::AwsIamIncomplete { + missing: [ + ("aws_iam_user_name", user_name.is_none()), + ("aws_iam_cache_name", cache_name.is_none()), + ("aws_iam_region", region.is_none()), + ] + .into_iter() + .filter_map(|(name, missing)| missing.then_some(name)) + .collect(), + }), + } +} diff --git a/litellm-rust/crates/redis-resolve/src/endpoint.rs b/litellm-rust/crates/redis-resolve/src/endpoint.rs new file mode 100644 index 00000000000..e3722b1b7cc --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/endpoint.rs @@ -0,0 +1,180 @@ +use litellm_redis_types::{ + Auth, Cluster, Node, RedisUrl, Sentinel, Standalone, StaticAuth, Target, TlsOptions, UnixSocket, +}; +use litellm_secret_value::SecretValue; + +use serde_json::Value; + +use crate::{ + Entry, Error, Expected, + options::{Options, as_text, flag, integer, invalid, secret, text}, + tls::{self, QuerySettings}, +}; + +/// The target, first match wins like `get_redis_client`: cluster, url, Sentinel, then host. A +/// url also returns the settings its query carries. +pub(super) fn target(options: &Options, auth: &Auth) -> Result<(Target, QuerySettings), Error> { + if let Some(nodes) = cluster_nodes(options.startup_nodes.as_ref())? { + reject_cluster_db(options.db.as_ref())?; + let cluster = Cluster { + nodes, + tls: tls_for(options, auth)?, + }; + return Ok((Target::Cluster(cluster), QuerySettings::default())); + } + if let Some(url) = text(options.url.as_ref()) { + let query = tls::url_query(&url)?; + let url = RedisUrl { + url: SecretValue::new(match auth { + Auth::Static(_) => url, + _ => without_userinfo(&url), + }), + tls: query.tls, + db: query.db, + username: query.username, + password: query.password, + }; + return Ok((Target::Url(url), query.settings)); + } + if let Some(sentinels) = options.sentinel.nodes.as_ref() { + let sentinel = sentinel(options, auth, sentinels)?; + return Ok((Target::Sentinel(sentinel), QuerySettings::default())); + } + let host = text(options.host.as_ref()).ok_or(Error::NoTarget)?; + let target = match text(options.unix_socket_path.as_ref()) { + Some(path) => Target::Unix(UnixSocket { + path: path.into(), + db: db(options)?, + }), + None => Target::Standalone(Standalone { + host, + port: integer(options.port.as_ref())?.unwrap_or(6379), + tls: tls_for(options, auth)?, + db: db(options)?, + }), + }; + Ok((target, QuerySettings::default())) +} + +fn tls_for(options: &Options, auth: &Auth) -> Result, Error> { + if !flag(options.ssl.as_ref())? { + return Ok(None); + } + let gcp_ca_certs = match auth { + Auth::GcpIam { .. } => options.gcp.ssl_ca_certs.as_ref(), + _ => None, + }; + tls::tls_options(&options.tls, gcp_ca_certs).map(Some) +} + +fn db(options: &Options) -> Result { + integer(options.db.as_ref()).map(|db| db.unwrap_or(0)) +} + +fn sentinel(options: &Options, auth: &Auth, sentinels: &Entry) -> Result { + let monitors = sentinel_nodes(&sentinels.value) + .ok_or_else(|| invalid(&sentinels.source, Expected::SentinelNodes))?; + let service_name = + text(options.sentinel.service_name.as_ref()).ok_or(Error::NoSentinelService)?; + // redis.asyncio.Sentinel passes the connection kwargs to each monitor's Redis(host=..., port=...), so a + // host or port setting raises TypeError when litellm builds the client + if let Some(entry) = options.host.as_ref().or(options.port.as_ref()) { + return Err(Error::SentinelWithHost { + setting: entry.source.clone(), + }); + } + Ok(Sentinel { + monitors, + service_name, + monitor_auth: monitor_auth(options, auth), + tls: tls_for(options, auth)?, + db: db(options)?, + }) +} + +// _sentinel_auth_kwargs gives the monitors the data node's connection kwargs with sentinel_password as the +// password, so a static username reaches them too +fn monitor_auth(options: &Options, auth: &Auth) -> Option { + let password = secret(options.sentinel.password.as_ref())?; + let username = match auth { + Auth::Static(StaticAuth { username, .. }) => username.clone(), + _ => None, + }; + Some(StaticAuth { + username, + password: Some(password), + }) +} + +fn cluster_nodes(entry: Option<&Entry>) -> Result>, Error> { + let Some(entry) = entry else { + return Ok(None); + }; + let nodes = nodes(&entry.value, |node| { + node_address( + node.get("host").and_then(as_text), + node.get("port").and_then(as_text), + ) + }) + .ok_or_else(|| invalid(&entry.source, Expected::StartupNodes))?; + Ok((!nodes.is_empty()).then_some(nodes)) +} + +fn sentinel_nodes(value: &Value) -> Option> { + let nodes = nodes(value, |pair| match pair { + Value::Array(pair) if pair.len() == 2 => node_address(as_text(&pair[0]), as_text(&pair[1])), + _ => None, + })?; + (!nodes.is_empty()).then_some(nodes) +} + +fn nodes(value: &Value, parse: impl Fn(&Value) -> Option) -> Option> { + match value { + Value::String(json) => serde_json::from_str::>(json) + .ok()? + .iter() + .map(parse) + .collect(), + Value::Array(nodes) => nodes.iter().map(parse).collect(), + _ => None, + } +} + +fn node_address(host: Option, port: Option) -> Option { + Some(Node { + host: host?, + port: port?.parse().ok()?, + }) +} + +// redis-py's RedisCluster raises unless db is 0 or None +fn reject_cluster_db(db: Option<&Entry>) -> Result<(), Error> { + let Some(entry) = db else { + return Ok(()); + }; + let is_zero = match &entry.value { + Value::String(_) => matches!(integer::(Some(entry)), Ok(None | Some(0))), + Value::Bool(flag) => !flag, + Value::Number(number) => number.as_f64() == Some(0.0), + _ => false, + }; + match is_zero { + true => Ok(()), + false => Err(Error::DbInCluster { + setting: entry.source.clone(), + }), + } +} + +// _get_redis_url_kwargs keeps only plain Connection kwargs, so redis-py drops every ssl_* setting for a url +// and a credential provider replaces the url's userinfo, which _url_without_userinfo strips +fn without_userinfo(url: &str) -> String { + let Some((scheme, rest)) = url.split_once("://") else { + return url.to_owned(); + }; + let (authority, tail) = rest.split_at(rest.find(['/', '?', '#']).unwrap_or(rest.len())); + let host = authority + .rsplit_once('@') + .map_or(authority, |(_, host)| host); + format!("{scheme}://{host}{tail}") +} diff --git a/litellm-rust/crates/redis-resolve/src/error.rs b/litellm-rust/crates/redis-resolve/src/error.rs new file mode 100644 index 00000000000..b13b65edfb5 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/error.rs @@ -0,0 +1,80 @@ +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid Redis setting {setting}: expected {expected}")] + Invalid { setting: String, expected: Expected }, + #[error("invalid Redis setting {setting}: redis-py's parse_url cannot read it as a number")] + UrlQuery { setting: String }, + #[error("either host or url must be specified for redis")] + NoTarget, + #[error("service_name is required with sentinel_nodes")] + NoSentinelService, + #[error("{setting} cannot be combined with sentinel_nodes, which redis-py rejects at startup")] + SentinelWithHost { setting: String }, + #[error("{setting} must be 0 in cluster mode")] + DbInCluster { setting: String }, + #[error("credential_provider is a host object, and the host passed no handle for it")] + ProviderWithoutHandle, + #[error( + "{setting} asks for {check}, which rustls cannot perform, so Redis is not connected without them" + )] + Unenforceable { + setting: String, + check: &'static str, + }, + #[error("AWS ElastiCache IAM Redis authentication requires TLS")] + AwsIamWithoutTls, + #[error("AWS ElastiCache IAM Redis authentication requires: {}", missing.join(", "))] + AwsIamIncomplete { missing: Vec<&'static str> }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display)] +pub enum Expected { + #[strum(serialize = "an integer")] + Integer, + #[strum(serialize = "a boolean")] + Boolean, + #[strum(serialize = "a non-negative number of seconds")] + Seconds, + #[strum(serialize = "none, optional, required, null or an ssl.VerifyMode value")] + CertRequirement, + #[strum(serialize = "a list of objects with host and port")] + StartupNodes, + #[strum(serialize = "a non-empty list of [host, port] pairs")] + SentinelNodes, +} + +/// The exception class redis-py 5.3.1 raises building its client. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum PythonException { + ValueError, + TypeError, + RedisClusterException, +} + +impl Error { + /// What `_get_redis_client_logic` and redis-py's async client constructors raise on the same kwargs, or `None` + /// when Python builds the client and fails only once a command opens a connection, or never fails. + pub fn raised_building_the_python_client(&self) -> Option { + match self { + Self::Invalid { expected, .. } => match expected { + Expected::StartupNodes | Expected::SentinelNodes => { + Some(PythonException::ValueError) + } + Expected::Integer + | Expected::Boolean + | Expected::Seconds + | Expected::CertRequirement => None, + }, + Self::UrlQuery { .. } + | Self::NoTarget + | Self::NoSentinelService + | Self::AwsIamWithoutTls + | Self::AwsIamIncomplete { .. } => Some(PythonException::ValueError), + Self::SentinelWithHost { .. } | Self::Unenforceable { .. } => { + Some(PythonException::TypeError) + } + Self::DbInCluster { .. } => Some(PythonException::RedisClusterException), + Self::ProviderWithoutHandle => None, + } + } +} diff --git a/litellm-rust/crates/redis-resolve/src/keys.rs b/litellm-rust/crates/redis-resolve/src/keys.rs new file mode 100644 index 00000000000..c4897701d6f --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/keys.rs @@ -0,0 +1,304 @@ +use litellm_redis_types::Ignored; +use serde_json::Value; + +/// How the Rust client honors a redis-py kwarg. +#[derive(Clone, Copy, Debug)] +pub enum Handling { + /// Modeled by the resolver into a field of `Options`. + Read, + /// Dropped and returned as an `Ignored` naming `consequence`, unless `met` says the value asks for nothing the + /// Rust client does not already do. + Warn { + consequence: &'static str, + met: Option bool>, + }, + /// A security check the Rust client cannot perform and Python's async client already refuses, so asking for it + /// fails instead of connecting without it. + Reject { check: &'static str, asked: Asked }, +} + +/// Which values of a `Reject` row ask for its check. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Asked { + /// A flag read as true. False asks for nothing. + WhenTrue, + /// Any present value, like a list of flags. An empty one asks for nothing. + WhenSet, +} + +/// Whether redis-py's `parse_url` reads the setting from a url's query and the Rust resolver does too. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FromUrl { + No, + Any, + Rediss, +} + +#[derive(Clone, Copy, Debug)] +pub struct Setting { + pub kwarg: &'static str, + pub handling: Handling, + pub from_url: FromUrl, + /// `_get_redis_env_kwarg_mapping` reads it as `REDIS_`. + pub environment: bool, + /// The names `_merge_redis_kwargs` reads with `get_secret` when the merged kwargs lack the setting. + pub fallback: &'static [&'static str], +} + +impl Setting { + const fn new(kwarg: &'static str, handling: Handling) -> Self { + Self { + kwarg, + handling, + from_url: FromUrl::No, + environment: false, + fallback: &[], + } + } + + const fn env(self) -> Self { + Self { + environment: true, + ..self + } + } + + const fn in_url(self, from_url: FromUrl) -> Self { + Self { from_url, ..self } + } + + const fn or_secrets(self, fallback: &'static [&'static str]) -> Self { + Self { fallback, ..self } + } + + pub fn environment_name(&self) -> Option { + self.environment + .then(|| format!("REDIS_{}", self.kwarg.to_ascii_uppercase())) + } + + pub fn is_read(&self) -> bool { + matches!(self.handling, Handling::Read) + } + + /// The warning a set value earns, or `None` when the setting is read, refused or already met. + pub(crate) fn warning(&self, value: &Value) -> Option<&'static str> { + match self.handling { + Handling::Warn { consequence, met } if !met.is_some_and(|met| met(value)) => { + Some(consequence) + } + _ => None, + } + } +} + +const fn read(kwarg: &'static str) -> Setting { + Setting::new(kwarg, Handling::Read) +} + +const fn warn(kwarg: &'static str, consequence: &'static str) -> Setting { + Setting::new( + kwarg, + Handling::Warn { + consequence, + met: None, + }, + ) +} + +const fn warn_unless( + kwarg: &'static str, + consequence: &'static str, + met: fn(&Value) -> bool, +) -> Setting { + Setting::new( + kwarg, + Handling::Warn { + consequence, + met: Some(met), + }, + ) +} + +const fn reject(kwarg: &'static str, check: &'static str, asked: Asked) -> Setting { + Setting::new(kwarg, Handling::Reject { check, asked }) +} + +const NO_EQUIVALENT: &str = "the Rust client has no equivalent"; +const OWN_STACK: &str = "the Rust client uses its own pool, retries and events"; +const OWN_CONNECTION: &str = "it tunes redis-py's own connection, and the Rust client has its own"; +const OWN_CLUSTER: &str = "the Rust cluster client discovers, routes and retries on its own"; +const MAINTENANCE: &str = "the Rust client does not follow server maintenance notifications"; +const RUSTLS: &str = "rustls negotiates TLS 1.2+ with its own ECDHE AEAD suites"; +const REVOCATION: &str = "OCSP revocation checks"; +const OCSP: &str = "it only configures ssl_validate_ocsp, which the Rust client refuses when set"; +const VERIFY_FLAGS: &str = "X.509 verify flags such as CRL checking"; + +/// Every Redis kwarg redis-py or `litellm/_redis.py` accepts, one row each. `redis-py-kwargs.json` records which +/// redis-py versions and `REDIS_*` names accept them, and `tests/catalog.rs` keeps the two in step. +pub const SETTINGS: &[Setting] = &[ + warn("address_remap", OWN_CLUSTER), + warn("auto_close_connection_pool", OWN_STACK), + read("aws_iam_auth").env(), + read("aws_iam_cache_name").env(), + read("aws_iam_region") + .env() + .or_secrets(&["AWS_REGION", "AWS_DEFAULT_REGION"]), + read("aws_iam_serverless").env(), + read("aws_iam_user_name").env(), + read("azure_client_id") + .env() + .or_secrets(&["AZURE_CLIENT_ID"]), + read("azure_client_secret") + .env() + .or_secrets(&["AZURE_CLIENT_SECRET"]), + read("azure_redis_ad_token") + .env() + .or_secrets(&["REDIS_AZURE_AD_TOKEN"]), + read("azure_tenant_id") + .env() + .or_secrets(&["AZURE_TENANT_ID"]), + warn("cache", NO_EQUIVALENT).env(), + warn("cache_config", NO_EQUIVALENT).env(), + warn("charset", NO_EQUIVALENT).env(), + warn("client_name", NO_EQUIVALENT).env(), + warn("cluster_error_retry_attempts", OWN_CLUSTER), + warn("command_packer", OWN_CONNECTION), + warn("connection_class", OWN_STACK), + warn("connection_error_retry_attempts", OWN_CLUSTER), + warn("connection_pool", OWN_STACK), + read("credential_provider"), + read("db").env().in_url(FromUrl::Any), + warn("decode_responses", NO_EQUIVALENT).env(), + warn("driver_info", NO_EQUIVALENT), + warn("dynamic_startup_nodes", OWN_CLUSTER), + warn("encoder_class", OWN_CONNECTION), + warn("encoding", NO_EQUIVALENT).env(), + warn("encoding_errors", NO_EQUIVALENT).env(), + warn("errors", NO_EQUIVALENT).env(), + warn("event_dispatcher", OWN_STACK).env(), + read("gcp_service_account") + .env() + .or_secrets(&["REDIS_GCP_SERVICE_ACCOUNT"]), + read("gcp_ssl_ca_certs") + .env() + .or_secrets(&["REDIS_GCP_SSL_CA_CERTS"]), + read("health_check_interval").env().in_url(FromUrl::Any), + read("host").env(), + warn("legacy_responses", NO_EQUIVALENT), + warn("lib_name", NO_EQUIVALENT).env(), + warn("lib_version", NO_EQUIVALENT).env(), + warn("load_balancing_strategy", OWN_CLUSTER), + warn("maint_notifications_config", MAINTENANCE), + warn("maint_notifications_pool_handler", MAINTENANCE), + warn("maintenance_notification_hash", MAINTENANCE), + warn("maintenance_state", MAINTENANCE), + warn("max_connections", "the Rust client has no connection cap").env(), + read("namespace"), + warn("orig_host_address", MAINTENANCE), + warn("orig_socket_connect_timeout", MAINTENANCE), + warn("orig_socket_timeout", MAINTENANCE), + warn("oss_cluster_maint_notifications_handler", MAINTENANCE), + warn("parser", OWN_CONNECTION), + warn("parser_class", OWN_CONNECTION), + read("password").env().in_url(FromUrl::Any), + warn("path", OWN_CLUSTER), + warn("policy_resolver", OWN_CLUSTER), + read("port").env(), + warn("protocol", NO_EQUIVALENT).env(), + warn("read_from_replicas", OWN_CLUSTER), + warn( + "redis_connect_func", + "the Python object is not callable from Rust, so static username and password are used", + ) + .env(), + warn("reinitialize_steps", OWN_CLUSTER), + warn("require_full_coverage", OWN_CLUSTER), + warn("retry", OWN_STACK), + warn("retry_on_error", NO_EQUIVALENT).env(), + warn("retry_on_timeout", NO_EQUIVALENT).env(), + read("sentinel_nodes").or_secrets(&["REDIS_SENTINEL_NODES"]), + read("sentinel_password").or_secrets(&["REDIS_SENTINEL_PASSWORD"]), + read("service_name").or_secrets(&["REDIS_SERVICE_NAME"]), + warn("single_connection_client", NO_EQUIVALENT).env(), + read("socket_connect_timeout").env().in_url(FromUrl::Any), + read("socket_keepalive").env(), + warn("socket_keepalive_options", NO_EQUIVALENT).env(), + warn("socket_read_size", OWN_CONNECTION), + read("socket_timeout").env().in_url(FromUrl::Any), + warn("socket_type", OWN_CONNECTION), + read("ssl").env(), + read("ssl_ca_certs").env().in_url(FromUrl::Rediss), + read("ssl_ca_data").env().in_url(FromUrl::Rediss), + read("ssl_ca_path").env().in_url(FromUrl::Rediss), + read("ssl_cert_reqs").env().in_url(FromUrl::Rediss), + read("ssl_certfile").env().in_url(FromUrl::Rediss), + read("ssl_check_hostname").env().in_url(FromUrl::Rediss), + warn("ssl_ciphers", RUSTLS).env(), + reject("ssl_exclude_verify_flags", VERIFY_FLAGS, Asked::WhenSet), + reject("ssl_include_verify_flags", VERIFY_FLAGS, Asked::WhenSet), + read("ssl_keyfile").env().in_url(FromUrl::Rediss), + warn_unless("ssl_min_version", RUSTLS, is_met_by_rustls).env(), + warn("ssl_ocsp_context", OCSP).env(), + warn("ssl_ocsp_expected_cert", OCSP).env(), + read("ssl_password").env().in_url(FromUrl::Rediss), + reject("ssl_validate_ocsp", REVOCATION, Asked::WhenTrue).env(), + reject("ssl_validate_ocsp_stapled", REVOCATION, Asked::WhenTrue).env(), + read("startup_nodes").or_secrets(&["REDIS_CLUSTER_NODES"]), + read("unix_socket_path").env(), + read("url").env(), + read("username").env().in_url(FromUrl::Any), +]; + +// rustls only speaks TLS 1.2 and 1.3, so a minimum of 1.2 or lower is already met. redis-py takes an +// ssl.TLSVersion, which reaches Rust as its name or its value +fn is_met_by_rustls(value: &Value) -> bool { + match value { + Value::Number(number) => matches!(number.as_i64(), Some(-2 | 769 | 770 | 771)), + Value::String(text) => matches!( + text.strip_prefix("TLSVersion.") + .unwrap_or(text) + .to_ascii_uppercase() + .as_str(), + "MINIMUM_SUPPORTED" | "TLSV1" | "TLSV1_1" | "TLSV1_2" + ), + _ => false, + } +} + +pub fn setting(kwarg: &str) -> Option<&'static Setting> { + SETTINGS.iter().find(|setting| setting.kwarg == kwarg) +} + +/// The check a `Reject` row stands for and which values ask for it. +pub(crate) fn refusal(kwarg: &str) -> Option<(&'static str, Asked)> { + match setting(kwarg)?.handling { + Handling::Reject { check, asked } => Some((check, asked)), + Handling::Read | Handling::Warn { .. } => None, + } +} + +/// A warned setting the kwargs set, with what dropping it means. A kwarg with no row is not a Redis setting, so it +/// is dropped without a warning. +pub(crate) fn ignored(source: &str, kwarg: &str, value: &Value) -> Option { + let consequence = setting(kwarg)?.warning(value)?; + Some(Ignored { + setting: source.to_owned(), + consequence, + }) +} + +/// A url query parameter the resolver does not read. redis-py turns every one into a client kwarg, so even one with +/// no row is reported. +pub(crate) fn ignored_query(source: &str, name: &str, value: &Value) -> Option { + let consequence = match setting(name) { + None => NO_EQUIVALENT, + Some(setting) if setting.is_read() => { + "the Rust client does not read it from this url's query" + } + Some(setting) => setting.warning(value)?, + }; + Some(Ignored { + setting: source.to_owned(), + consequence, + }) +} diff --git a/litellm-rust/crates/redis-resolve/src/lib.rs b/litellm-rust/crates/redis-resolve/src/lib.rs new file mode 100644 index 00000000000..b7aa5d18f64 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/lib.rs @@ -0,0 +1,95 @@ +use std::{collections::BTreeMap, time::Duration}; + +use litellm_redis_types::{Auth, HealthCheck, ProviderHandle, Resolution, StaticAuth, Target}; +use serde_json::Value; + +mod auth; +mod endpoint; +mod error; +mod keys; +mod options; +mod tls; + +pub use error::{Error, Expected, PythonException}; +pub use keys::{Asked, FromUrl, Handling, SETTINGS, Setting, setting}; +pub use options::is_present; + +use options::Options; + +/// One merged redis-py kwarg and the name a warning or error reports it under. +#[derive(Clone, Debug, PartialEq)] +pub struct Entry { + pub value: Value, + pub source: String, +} + +pub type Kwargs = BTreeMap; + +/// Kwargs whose every entry is reported under its own kwarg name. +pub fn kwargs(fields: impl IntoIterator) -> Kwargs { + fields + .into_iter() + .map(|(kwarg, value)| { + let entry = Entry { + value, + source: kwarg.clone(), + }; + (kwarg, entry) + }) + .collect() +} + +/// What redis-py's constructors make of kwargs a host already merged the way `_get_redis_client_logic` does. A +/// host passes its `credential_provider` object as `provider`, and `REDIS_USERNAME` as its process environment +/// holds it, which Python's Azure AD auth reads from `os.environ` instead of the merged kwargs. +pub fn interpret( + kwargs: &Kwargs, + provider: Option, + redis_username: Option<&str>, +) -> Result { + let (options, ignored) = Options::read(kwargs)?; + let auth = auth::redis_auth(&options, provider, redis_username)?; + let (target, query) = endpoint::target(&options, &auth)?; + let auth = match (&target, auth) { + (Target::Url(_), Auth::Static(StaticAuth { username, .. })) => Auth::Static(StaticAuth { + username, + password: None, + }), + (_, auth) => auth, + }; + let timeout = query.timeout.map_or_else( + || options::seconds(options.socket_timeout.as_ref()), + |timeout| Ok(Some(timeout)), + )?; + let connect_timeout = query.connect_timeout.map_or_else( + || options::seconds(options.socket_connect_timeout.as_ref()), + |timeout| Ok(Some(timeout)), + )?; + let health_check_interval = query.health_check_interval.map_or_else( + || options::integer::(options.health_check_interval.as_ref()), + |interval| Ok(Some(interval)), + )?; + Ok(Resolution { + target, + auth, + namespace: options::text(options.namespace.as_ref()), + timeout, + connect_timeout, + health_check: health_check(health_check_interval), + keepalive: options + .socket_keepalive + .as_ref() + .map_or(Ok(true), |entry| options::flag(Some(entry)))?, + ignored: ignored.into_iter().chain(query.ignored).collect(), + }) +} + +// redis-py checks health only when the interval is truthy, so zero turns it off and a negative interval checks +// before every command +fn health_check(seconds: Option) -> HealthCheck { + match seconds { + None => HealthCheck::Unset, + Some(0) => HealthCheck::Disabled, + Some(seconds) => HealthCheck::Every(Duration::from_secs(seconds.try_into().unwrap_or(0))), + } +} diff --git a/litellm-rust/crates/redis-resolve/src/options.rs b/litellm-rust/crates/redis-resolve/src/options.rs new file mode 100644 index 00000000000..f2539e46ca5 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/options.rs @@ -0,0 +1,281 @@ +use std::time::Duration; + +use litellm_python_compat::{BigInt, number}; +use litellm_redis_types::Ignored; +use litellm_secret_value::SecretValue; +use serde_json::Value; + +use crate::{ + Entry, Error, Expected, Kwargs, + keys::{self, Asked}, +}; + +/// Every setting the resolver reads. A setting is `None` when it is unset, null or empty, except that a null +/// `tls.cert_reqs` is kept because redis-py reads `ssl_cert_reqs=None` as `CERT_NONE`. +pub(super) struct Options { + pub(super) url: Option, + pub(super) host: Option, + pub(super) port: Option, + pub(super) db: Option, + pub(super) unix_socket_path: Option, + pub(super) startup_nodes: Option, + pub(super) sentinel: SentinelOptions, + pub(super) username: Option, + pub(super) password: Option, + pub(super) credential_provider: Option, + pub(super) gcp: GcpOptions, + pub(super) azure: AzureOptions, + pub(super) aws: AwsOptions, + pub(super) ssl: Option, + pub(super) tls: TlsInput, + pub(super) namespace: Option, + pub(super) socket_timeout: Option, + pub(super) socket_connect_timeout: Option, + pub(super) health_check_interval: Option, + pub(super) socket_keepalive: Option, +} + +pub(super) struct SentinelOptions { + pub(super) nodes: Option, + pub(super) service_name: Option, + pub(super) password: Option, +} + +pub(super) struct GcpOptions { + pub(super) service_account: Option, + pub(super) ssl_ca_certs: Option, +} + +pub(super) struct AzureOptions { + pub(super) ad_token: Option, + pub(super) client_id: Option, + pub(super) tenant_id: Option, + pub(super) client_secret: Option, +} + +pub(super) struct AwsOptions { + pub(super) auth: Option, + pub(super) user_name: Option, + pub(super) cache_name: Option, + pub(super) region: Option, + pub(super) serverless: Option, +} + +/// The TLS settings, read from the kwargs or from a `rediss://` url's query. +pub(super) struct TlsInput { + pub(super) cert_reqs: Option, + pub(super) ca_certs: Option, + pub(super) ca_data: Option, + pub(super) ca_path: Option, + pub(super) certfile: Option, + pub(super) keyfile: Option, + pub(super) key_password: Option, + pub(super) check_hostname: Option, +} + +/// The settings the kwargs set that the Rust client drops, once none of them asks for a check it cannot perform. +fn ignored(kwargs: &Kwargs) -> Result, Error> { + refuse_unenforceable(kwargs.iter())?; + Ok(kwargs + .iter() + .filter_map(|(key, entry)| keys::ignored(&entry.source, key, &entry.value)) + .collect()) +} + +impl Options { + /// The typed settings and the ones the Rust client drops. + pub(super) fn read(kwargs: &Kwargs) -> Result<(Self, Vec), Error> { + let ignored = ignored(kwargs)?; + let options = Self::from_lookup(|key| kwargs.get(key).cloned()); + Ok((options, ignored)) + } + + fn from_lookup(lookup: impl Fn(&'static str) -> Option) -> Self { + let get = |key| lookup(key).filter(|entry| is_present(&entry.value)); + Self { + url: get("url"), + host: get("host"), + port: get("port"), + db: get("db"), + unix_socket_path: get("unix_socket_path"), + startup_nodes: get("startup_nodes"), + sentinel: SentinelOptions { + nodes: get("sentinel_nodes"), + service_name: get("service_name"), + password: get("sentinel_password"), + }, + username: get("username"), + password: get("password"), + credential_provider: get("credential_provider"), + gcp: GcpOptions { + service_account: get("gcp_service_account"), + ssl_ca_certs: get("gcp_ssl_ca_certs"), + }, + azure: AzureOptions { + ad_token: get("azure_redis_ad_token"), + client_id: get("azure_client_id"), + tenant_id: get("azure_tenant_id"), + client_secret: get("azure_client_secret"), + }, + aws: AwsOptions { + auth: get("aws_iam_auth"), + user_name: get("aws_iam_user_name"), + cache_name: get("aws_iam_cache_name"), + region: get("aws_iam_region"), + serverless: get("aws_iam_serverless"), + }, + ssl: get("ssl"), + tls: TlsInput { + cert_reqs: lookup("ssl_cert_reqs") + .filter(|entry| entry.value.is_null() || is_present(&entry.value)), + ca_certs: get("ssl_ca_certs"), + ca_data: get("ssl_ca_data"), + ca_path: get("ssl_ca_path"), + certfile: get("ssl_certfile"), + keyfile: get("ssl_keyfile"), + key_password: get("ssl_password"), + check_hostname: get("ssl_check_hostname"), + }, + namespace: get("namespace"), + socket_timeout: get("socket_timeout"), + socket_connect_timeout: get("socket_connect_timeout"), + health_check_interval: get("health_check_interval"), + socket_keepalive: get("socket_keepalive"), + } + } +} + +/// A `Reject` row whose value asks for its check fails, because Rust cannot perform it. +pub(super) fn refuse_unenforceable<'a>( + entries: impl IntoIterator, +) -> Result<(), Error> { + entries + .into_iter() + .filter_map(|(key, entry)| Some((keys::refusal(key)?, entry))) + .try_for_each(|((check, asked), entry)| { + let asks = match asked { + Asked::WhenTrue => flag(Some(entry))?, + Asked::WhenSet => is_present(&entry.value), + }; + match asks { + true => Err(Error::Unenforceable { + setting: entry.source.clone(), + check, + }), + false => Ok(()), + } + }) +} + +// Python reads most of these with `x or fallback`, so an empty value counts as unset everywhere +pub fn is_present(value: &Value) -> bool { + match value { + Value::Null => false, + Value::String(text) => !text.is_empty(), + Value::Array(items) => !items.is_empty(), + Value::Object(entries) => !entries.is_empty(), + Value::Bool(_) | Value::Number(_) => true, + } +} + +pub(super) fn invalid(setting: &str, expected: Expected) -> Error { + Error::Invalid { + setting: setting.to_owned(), + expected, + } +} + +pub(super) fn as_text(value: &Value) -> Option { + match value { + Value::String(text) => Some(text.clone()), + Value::Number(number) => Some(number.to_string()), + Value::Bool(flag) => Some(flag.to_string()), + _ => None, + } +} + +pub(super) fn text(entry: Option<&Entry>) -> Option { + entry.and_then(|entry| as_text(&entry.value)) +} + +pub(super) fn secret(entry: Option<&Entry>) -> Option { + text(entry).map(SecretValue::new) +} + +/// A flag every Python surface reads the same way. `_str_to_bool` takes only `true`, `1` and `yes` +/// as true while the proxy's pydantic models also take `on`, `t` and `y`, so text they disagree on, +/// or text neither accepts, is an error instead of a guess. +pub(super) fn flag(entry: Option<&Entry>) -> Result { + let Some(entry) = entry else { + return Ok(false); + }; + match &entry.value { + Value::Bool(flag) => Ok(*flag), + Value::Number(number) => Ok(number.as_f64().is_some_and(|number| number != 0.0)), + Value::String(text) => match text.to_ascii_lowercase().as_str() { + "true" | "1" | "yes" => Ok(true), + "false" | "0" | "no" | "off" | "f" | "n" => Ok(false), + _ => Err(invalid(&entry.source, Expected::Boolean)), + }, + _ => Err(invalid(&entry.source, Expected::Boolean)), + } +} + +pub(super) fn integer>(entry: Option<&Entry>) -> Result, Error> { + coerced(entry, Expected::Integer, number::int, |number| { + T::try_from(number).ok() + }) +} + +pub(super) fn seconds(entry: Option<&Entry>) -> Result, Error> { + coerced(entry, Expected::Seconds, number::float, |seconds| { + Duration::try_from_secs_f64(seconds).ok() + }) +} + +fn coerced( + entry: Option<&Entry>, + expected: Expected, + parse: impl Fn(&str) -> Option, + convert: impl Fn(N) -> Option, +) -> Result, Error> { + let Some((entry, text)) = entry.and_then(|entry| Some((entry, as_text(&entry.value)?))) else { + return Ok(None); + }; + match parse(&text) { + Some(number) => convert(number) + .map(Some) + .ok_or_else(|| invalid(&entry.source, expected)), + None => { + tracing::warn!(setting = %entry.source, "Redis setting is not a number, so the default is used like redis-py"); + Ok(None) + } + } +} + +#[cfg(test)] +mod tests { + use std::{cell::RefCell, collections::BTreeSet}; + + use rstest::rstest; + + use super::Options; + use crate::SETTINGS; + + // a Read row that Options never asks for would be neither modeled nor passed through nor warned about + #[rstest] + fn options_ask_for_exactly_the_read_rows() { + let asked = RefCell::new(BTreeSet::new()); + let _ = Options::from_lookup(|kwarg| { + asked.borrow_mut().insert(kwarg); + None + }); + + let read: BTreeSet<&str> = SETTINGS + .iter() + .filter(|setting| setting.is_read()) + .map(|setting| setting.kwarg) + .collect(); + assert_eq!(asked.into_inner(), read); + } +} diff --git a/litellm-rust/crates/redis-resolve/src/tls.rs b/litellm-rust/crates/redis-resolve/src/tls.rs new file mode 100644 index 00000000000..a46a6ccc650 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/src/tls.rs @@ -0,0 +1,189 @@ +use std::time::Duration; + +use litellm_python_compat::number; +use litellm_redis_types::{Ignored, TlsOptions}; +use litellm_secret_value::SecretValue; + +use serde_json::Value; + +use crate::{ + Entry, Error, Expected, + keys::{self, FromUrl}, + options::{self, TlsInput, flag, invalid, secret, text}, +}; + +/// `gcp_ca_certs` replaces `ssl_ca_certs` when GCP IAM is on, as `_get_redis_client_logic` does. +pub(super) fn tls_options( + input: &TlsInput, + gcp_ca_certs: Option<&Entry>, +) -> Result { + Ok(TlsOptions { + ca_certs: text(gcp_ca_certs.or(input.ca_certs.as_ref())).map(Into::into), + ca_data: text(input.ca_data.as_ref()), + ca_path: text(input.ca_path.as_ref()).map(Into::into), + certfile: text(input.certfile.as_ref()).map(Into::into), + keyfile: text(input.keyfile.as_ref()).map(Into::into), + key_password: secret(input.key_password.as_ref()), + verify_certificate: verify_certificate(input.cert_reqs.as_ref())?, + check_hostname: flag(input.check_hostname.as_ref())?, + }) +} + +// RedisSSLContext looks text up case-sensitively and takes an ssl.VerifyMode as is +fn verify_certificate(cert_reqs: Option<&Entry>) -> Result { + let Some(entry) = cert_reqs else { + return Ok(true); + }; + match &entry.value { + // Explicit None intentionally means CERT_NONE, matching redis-py rather than the unset default. + Value::Null => Ok(false), + Value::String(text) if text == "none" => Ok(false), + Value::String(text) if text == "optional" || text == "required" => Ok(true), + Value::Number(mode) => match mode.as_i64() { + Some(0) => Ok(false), + Some(1 | 2) => Ok(true), + _ => Err(invalid(&entry.source, Expected::CertRequirement)), + }, + _ => Err(invalid(&entry.source, Expected::CertRequirement)), + } +} + +/// What redis-py's `parse_url` reads from a url's query on top of the url itself. +pub(super) struct UrlQuery { + pub(super) tls: Option, + pub(super) db: Option, + pub(super) username: Option, + pub(super) password: Option, + pub(super) settings: QuerySettings, +} + +/// The connection settings a url's query sets, which win over the kwargs because `from_url` applies them last. +#[derive(Default)] +pub(super) struct QuerySettings { + pub(super) timeout: Option, + pub(super) connect_timeout: Option, + pub(super) health_check_interval: Option, + pub(super) ignored: Vec, +} + +pub(super) fn url_query(url: &str) -> Result { + let Ok(url) = url::Url::parse(url) else { + return Ok(UrlQuery { + tls: None, + db: None, + username: None, + password: None, + settings: QuerySettings::default(), + }); + }; + let tls_scheme = url.scheme().eq_ignore_ascii_case("rediss"); + let pairs: Vec<(String, Entry)> = url + .query_pairs() + .filter(|(_, value)| !value.is_empty()) + .map(|(name, value)| { + ( + name.to_string(), + Entry { + value: Value::String(value.into_owned()), + source: format!("url ?{name}"), + }, + ) + }) + .collect(); + // parse_qs keeps every value of a repeated parameter and parse_url reads the first + let query: Vec<(String, Entry)> = pairs + .iter() + .enumerate() + .filter(|(index, (name, _))| !pairs[..*index].iter().any(|(earlier, _)| earlier == name)) + .map(|(_, pair)| pair.clone()) + .collect(); + options::refuse_unenforceable(query.iter().map(|(name, entry)| (name, entry)))?; + let (query, ignored): (Vec<_>, Vec<_>) = query.into_iter().partition(|(name, _)| { + keys::setting(name).is_some_and(|setting| match setting.from_url { + FromUrl::Any => true, + FromUrl::Rediss => tls_scheme, + FromUrl::No => false, + }) + }); + let get = |key: &str| { + query + .iter() + .find(|(name, _)| name == key) + .map(|(_, entry)| entry.clone()) + }; + let seconds = |key: &str| { + query_number( + get(key).as_ref(), + number::float, + Expected::Seconds, + |seconds| Duration::try_from_secs_f64(seconds).ok(), + ) + }; + let integer = |key: &str| { + query_number( + get(key).as_ref(), + number::int, + Expected::Integer, + |number| i64::try_from(number).ok(), + ) + }; + let db = integer("db")?; + let settings = QuerySettings { + timeout: seconds("socket_timeout")?, + connect_timeout: seconds("socket_connect_timeout")?, + health_check_interval: integer("health_check_interval")?, + ignored: ignored + .iter() + .filter_map(|(name, entry)| keys::ignored_query(&entry.source, name, &entry.value)) + .collect(), + }; + let tls = match tls_scheme { + true => Some(TlsOptions { + ca_certs: text(get("ssl_ca_certs").as_ref()).map(Into::into), + ca_data: text(get("ssl_ca_data").as_ref()), + ca_path: text(get("ssl_ca_path").as_ref()).map(Into::into), + certfile: text(get("ssl_certfile").as_ref()).map(Into::into), + keyfile: text(get("ssl_keyfile").as_ref()).map(Into::into), + key_password: secret(get("ssl_password").as_ref()), + verify_certificate: verify_certificate(get("ssl_cert_reqs").as_ref())?, + check_hostname: text(get("ssl_check_hostname").as_ref()) + .is_some_and(|text| url_bool(&text)), + }), + false => None, + }; + Ok(UrlQuery { + tls, + db, + username: secret(get("username").as_ref()), + password: secret(get("password").as_ref()), + settings, + }) +} + +// parse_url reads these with int() or float() and raises ValueError building the client on anything else +fn query_number( + entry: Option<&Entry>, + parse: impl Fn(&str) -> Option, + expected: Expected, + convert: impl Fn(N) -> Option, +) -> Result, Error> { + let Some(entry) = entry else { + return Ok(None); + }; + let number = text(Some(entry)) + .and_then(|text| parse(&text)) + .ok_or_else(|| Error::UrlQuery { + setting: entry.source.clone(), + })?; + convert(number) + .map(Some) + .ok_or_else(|| invalid(&entry.source, expected)) +} + +// redis-py's to_bool: only 0, F, FALSE, N and NO (any case) are false +fn url_bool(text: &str) -> bool { + !matches!( + text.to_ascii_uppercase().as_str(), + "" | "0" | "F" | "FALSE" | "N" | "NO" + ) +} diff --git a/litellm-rust/crates/redis-resolve/tests/catalog.rs b/litellm-rust/crates/redis-resolve/tests/catalog.rs new file mode 100644 index 00000000000..b26ad0eebfb --- /dev/null +++ b/litellm-rust/crates/redis-resolve/tests/catalog.rs @@ -0,0 +1,129 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use litellm_redis_resolve::{Asked, Error, Handling, SETTINGS, Setting, interpret, kwargs}; +use rstest::rstest; +use serde::Deserialize; +use serde_json::json; + +#[derive(Deserialize)] +struct Catalog { + meta: Meta, + args: BTreeMap, +} + +#[derive(Deserialize)] +struct Meta { + redis_py_locked: String, +} + +#[derive(Deserialize)] +struct Arg { + environment: bool, + redis_py: BTreeMap>, +} + +// the signatures of the clients the proxy builds, where an unknown kwarg raises TypeError +const PROXY_SIGNATURES: &[&str] = &["async", "cluster", "connection"]; + +fn rejected() -> impl Iterator { + SETTINGS + .iter() + .filter_map(|setting| match setting.handling { + Handling::Reject { asked, .. } => Some((setting, asked)), + Handling::Read | Handling::Warn { .. } => None, + }) +} + +fn catalog() -> Catalog { + serde_json::from_str(include_str!("../redis-py-kwargs.json")).unwrap() +} + +fn host_with(kwarg: &str, value: serde_json::Value) -> litellm_redis_resolve::Kwargs { + kwargs([ + ("host".to_owned(), json!("redis")), + (kwarg.to_owned(), value), + ]) +} + +#[rstest] +fn every_catalogued_kwarg_has_exactly_one_settings_row() { + let rows: Vec<&str> = SETTINGS.iter().map(|setting| setting.kwarg).collect(); + let unique: BTreeSet<&str> = rows.iter().copied().collect(); + assert_eq!(rows.len(), unique.len(), "a kwarg has two SETTINGS rows"); + + let catalog = catalog(); + let catalogued: BTreeSet<&str> = catalog.args.keys().map(String::as_str).collect(); + assert_eq!(unique, catalogued); +} + +#[rstest] +fn settings_read_redis_environment_names_exactly_where_python_does() { + let catalog = catalog(); + + let differing: Vec<&str> = SETTINGS + .iter() + .filter(|setting| { + catalog + .args + .get(setting.kwarg) + .is_some_and(|arg| arg.environment != setting.environment) + }) + .map(|setting| setting.kwarg) + .collect(); + assert!(differing.is_empty(), "{differing:?}"); +} + +#[rstest] +fn every_warned_setting_is_reported() { + let warned = SETTINGS + .iter() + .filter(|setting| matches!(setting.handling, Handling::Warn { met: None, .. })); + + for setting in warned { + let resolved = interpret(&host_with(setting.kwarg, json!("set")), None, None).unwrap(); + + let reported: Vec<&str> = resolved + .ignored + .iter() + .map(|ignored| ignored.setting.as_str()) + .collect(); + assert_eq!(reported, vec![setting.kwarg]); + } +} + +#[rstest] +fn every_rejected_setting_fails_when_asked_for_and_not_otherwise() { + for (setting, asked) in rejected() { + let (asking, not_asking) = match asked { + Asked::WhenTrue => (json!(true), json!(false)), + Asked::WhenSet => (json!(["VERIFY_CRL_CHECK_LEAF"]), json!([])), + }; + + let error = interpret(&host_with(setting.kwarg, asking), None, None).unwrap_err(); + assert!( + matches!(&error, Error::Unenforceable { setting: source, .. } if source == setting.kwarg), + "{error:?}" + ); + assert!(interpret(&host_with(setting.kwarg, not_asking), None, None).is_ok()); + } +} + +#[rstest] +fn python_already_refuses_every_rejected_setting() { + let catalog = catalog(); + + let accepted: Vec<&str> = rejected() + .map(|(setting, _)| setting.kwarg) + .filter(|kwarg| { + catalog.args[*kwarg] + .redis_py + .get(&catalog.meta.redis_py_locked) + .is_some_and(|signatures| { + signatures + .iter() + .any(|signature| PROXY_SIGNATURES.contains(&signature.as_str())) + }) + }) + .collect(); + assert!(accepted.is_empty(), "{accepted:?}"); +} diff --git a/litellm-rust/crates/redis-resolve/tests/interpret.rs b/litellm-rust/crates/redis-resolve/tests/interpret.rs new file mode 100644 index 00000000000..fd59ec2acb1 --- /dev/null +++ b/litellm-rust/crates/redis-resolve/tests/interpret.rs @@ -0,0 +1,790 @@ +use std::collections::BTreeMap; + +use litellm_redis_resolve::{Entry, Error, Expected, Kwargs, PythonException, interpret, kwargs}; +use litellm_redis_types::{ + Auth, HealthCheck, Node, ProviderHandle, Resolution, StaticAuth, Target, TlsOptions, +}; +use rstest::rstest; +use serde_json::json; + +use Outcome::{Fails, Refuses, Resolves}; +use PythonException::{RedisClusterException, TypeError, ValueError}; + +const PROVIDER: ProviderHandle = ProviderHandle::new(7); + +fn merged(fields: serde_json::Value, sources: &[(&str, &str)]) -> Kwargs { + let fields = serde_json::from_value::>(fields).unwrap(); + kwargs(fields) + .into_iter() + .map(|(kwarg, entry)| { + let source = sources + .iter() + .find(|(named, _)| *named == kwarg) + .map_or(entry.source, |(_, source)| (*source).to_owned()); + (kwarg, Entry { source, ..entry }) + }) + .collect() +} + +// a host hands its credential_provider object over as a handle, so the cases do the same +fn interpret_merged( + fields: serde_json::Value, + sources: &[(&str, &str)], + redis_username: Option<&str>, +) -> Result { + let has_provider = fields + .get("credential_provider") + .is_some_and(|provider| !provider.is_null()); + interpret( + &merged(fields, sources), + has_provider.then_some(PROVIDER), + redis_username, + ) +} + +fn interpret_kwargs(fields: serde_json::Value) -> Result { + interpret_merged(fields, &[], None) +} + +/// What a contract case expects of the resolver, next to what its Python mirror expects of redis-py. +enum Outcome { + Resolves(serde_json::Value), + /// Rust refuses a config Python connects with, a disagreement on purpose. + Refuses(&'static str), + /// redis-py 5.3.1 already raises on the config, building the client or its first connection, so failing + /// here breaks no deployment that works today. The class is what it raises building the client, and `None` + /// when it only raises once a command connects. + Fails(&'static str, Option), +} + +#[rstest] +fn a_provider_handle_beats_every_cloud_auth_and_static_credentials() { + let resolved = interpret_kwargs(json!({ + "host": "redis", "password": "pass", "credential_provider": "host-object", + "gcp_service_account": "account", "aws_iam_auth": true, "azure_redis_ad_token": "true", + })) + .unwrap(); + + assert_eq!(resolved.auth, Auth::Provider(PROVIDER)); +} + +#[rstest] +fn a_rediss_url_reads_ssl_password_from_its_query() { + let resolved = + interpret_kwargs(json!({"url": "rediss://cache:6380/0?ssl_password=query-key-secret"})) + .unwrap(); + + let Target::Url(parsed) = resolved.target else { + panic!("expected a url target, got {:?}", resolved.target); + }; + let key_password = parsed.tls.and_then(|tls| tls.key_password); + assert_eq!(expose(key_password.as_ref()), Some("query-key-secret")); + assert!(resolved.ignored.is_empty(), "{:?}", resolved.ignored); +} + +#[rstest] +#[case::read_on_a_plain_url( + "redis://cache:6379/1?db=3&username=query-user&password=query-secret", + Some(3), + Some("query-user"), + Some("query-secret") +)] +#[case::absent("rediss://cache:6380/1", None, None, None)] +fn a_url_query_carries_db_and_credentials_like_parse_url( + #[case] url: &str, + #[case] db: Option, + #[case] username: Option<&str>, + #[case] password: Option<&str>, +) { + let resolved = interpret_kwargs(json!({"url": url})).unwrap(); + + let Target::Url(parsed) = resolved.target else { + panic!("expected a url target, got {:?}", resolved.target); + }; + assert_eq!(parsed.db, db); + assert_eq!(expose(parsed.username.as_ref()), username); + assert_eq!(expose(parsed.password.as_ref()), password); + assert!(resolved.ignored.is_empty(), "{:?}", resolved.ignored); +} + +#[rstest] +fn a_url_query_db_that_is_not_an_integer_fails_like_parse_url() { + let error = interpret_kwargs(json!({"url": "redis://cache:6379?db=first"})).unwrap_err(); + + assert!(matches!(error, Error::UrlQuery { setting } if setting == "url ?db")); +} + +#[rstest] +#[case::tls_1_2_is_already_met(json!("TLSv1_2"), false)] +#[case::the_enum_name(json!("TLSVersion.TLSv1_2"), false)] +#[case::the_enum_value(json!(771), false)] +#[case::tls_1_3_is_not_enforced(json!("TLSv1_3"), true)] +fn ssl_min_version_warns_only_when_rustls_cannot_meet_it( + #[case] version: serde_json::Value, + #[case] warns: bool, +) { + let resolved = + interpret_kwargs(json!({"host": "redis", "ssl": true, "ssl_min_version": version})) + .unwrap(); + + assert_eq!( + !resolved.ignored.is_empty(), + warns, + "{:?}", + resolved.ignored + ); +} + +#[rstest] +#[case::unset(json!({"host": "redis"}), true)] +#[case::empty_like_unset(json!({"host": "redis", "socket_keepalive": ""}), true)] +#[case::enabled(json!({"host": "redis", "socket_keepalive": "yes"}), true)] +#[case::disabled(json!({"host": "redis", "socket_keepalive": false}), false)] +#[case::disabled_as_text(json!({"host": "redis", "socket_keepalive": "0"}), false)] +#[case::disabled_on_a_cluster(json!({"startup_nodes": [{"host": "node-a", "port": 7000}], "socket_keepalive": "false"}), false)] +fn keepalive_stays_on_unless_socket_keepalive_turns_it_off( + #[case] kwargs: serde_json::Value, + #[case] keepalive: bool, +) { + let resolved = interpret_kwargs(kwargs).unwrap(); + + assert_eq!(resolved.keepalive, keepalive); + assert!(resolved.ignored.is_empty(), "{:?}", resolved.ignored); +} + +#[rstest] +fn a_socket_keepalive_no_python_surface_agrees_on_is_invalid() { + let error = interpret_kwargs(json!({"host": "redis", "socket_keepalive": "on"})).unwrap_err(); + + assert!( + matches!(&error, Error::Invalid { setting, expected: Expected::Boolean } if setting == "socket_keepalive"), + "{error:?}" + ); +} + +fn tls_view(tls: Option<&TlsOptions>) -> serde_json::Value { + tls.map_or(serde_json::Value::Null, |tls| { + serde_json::json!({ + "ca_certs": tls.ca_certs, + "ca_data": tls.ca_data, + "ca_path": tls.ca_path, + "certfile": tls.certfile, + "keyfile": tls.keyfile, + "verify_certificate": tls.verify_certificate, + "check_hostname": tls.check_hostname, + }) + }) +} + +fn nodes_view(nodes: &[Node]) -> Vec { + nodes + .iter() + .map(|node| format!("{}:{}", node.host, node.port)) + .collect() +} + +fn expose(value: Option<&litellm_secret_value::SecretValue>) -> Option<&str> { + value.map(litellm_secret_value::SecretValue::expose) +} + +fn view(resolved: &Resolution) -> serde_json::Value { + let target = match &resolved.target { + Target::Unix(socket) => serde_json::json!({ + "mode": "unix", "path": socket.path, "tls": null, + }), + Target::Standalone(host) => serde_json::json!({ + "mode": "host", + "host": host.host, + "port": host.port, + "db": host.db, + "tls": tls_view(host.tls.as_ref()), + }), + Target::Url(url) => serde_json::json!({ + "mode": "url", + "url": url.url.expose(), + "tls": tls_view(url.tls.as_ref()), + }), + Target::Cluster(cluster) => serde_json::json!({ + "mode": "cluster", + "nodes": nodes_view(&cluster.nodes), + "tls": tls_view(cluster.tls.as_ref()), + }), + Target::Sentinel(sentinel) => serde_json::json!({ + "mode": "sentinel", + "nodes": nodes_view(&sentinel.monitors), + "service_name": sentinel.service_name, + "sentinel_password": expose(sentinel.monitor_auth.as_ref().and_then(|auth| auth.password.as_ref())), + "db": sentinel.db, + "tls": tls_view(sentinel.tls.as_ref()), + }), + }; + let auth = match &resolved.auth { + Auth::Static(StaticAuth { username, password }) => serde_json::json!({ + "kind": "static", + "username": expose(username.as_ref()), + "password": expose(password.as_ref()), + }), + Auth::Provider(_) => serde_json::json!({"kind": "host"}), + Auth::GcpIam { service_account } => { + serde_json::json!({"kind": "gcp_iam", "service_account": service_account}) + } + Auth::AzureAd(azure) => { + serde_json::json!({"kind": "azure_ad", "username": azure.username}) + } + Auth::AwsIam(aws) => serde_json::json!({ + "kind": "aws_iam", + "user_name": aws.user_name, + "cache_name": aws.cache_name, + "region": aws.region, + "serverless": aws.serverless, + }), + }; + let serde_json::Value::Object(target) = target else { + unreachable!("every target view is an object") + }; + let rest = serde_json::json!({ + "auth": auth, + "socket_timeout": resolved.timeout.map_or(5.0, |timeout| timeout.as_secs_f64()), + "socket_connect_timeout": resolved.connect_timeout.or(resolved.timeout).map_or(5.0, |timeout| timeout.as_secs_f64()), + "health_check_interval": match (&resolved.target, resolved.health_check) { + (_, HealthCheck::Every(interval)) => Some(interval.as_secs_f64()), + (Target::Cluster(_), HealthCheck::Unset) => Some(25.0), + (_, HealthCheck::Unset | HealthCheck::Disabled) => None, + }, + }); + let serde_json::Value::Object(rest) = rest else { + unreachable!("the view is an object") + }; + serde_json::Value::Object(target.into_iter().chain(rest).collect()) +} + +fn expected_view(fields: serde_json::Value) -> serde_json::Value { + let serde_json::Value::Object(fields) = fields else { + unreachable!("expected views are objects") + }; + let defaults = json!({ + "auth": {"kind": "static", "username": null, "password": null}, + "tls": null, + "socket_timeout": 5.0, + "socket_connect_timeout": 5.0, + "health_check_interval": null, + }); + let serde_json::Value::Object(defaults) = defaults else { + unreachable!("the defaults are an object") + }; + serde_json::Value::Object(defaults.into_iter().chain(fields).collect()) +} + +#[rstest] +#[case::host_defaults( + json!({"host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::empty_cluster_nodes_use_the_environment( + json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": 7000}]"}), + &[("startup_nodes", "REDIS_CLUSTER_NODES")], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), +)] +#[case::empty_sentinel_nodes_use_the_environment( + json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "service_name": "primary"}), + &[("sentinel_nodes", "REDIS_SENTINEL_NODES")], + None, + Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": null, "db": 0}))), +)] +#[case::empty_service_name_uses_the_environment( + json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}), + &[("service_name", "REDIS_SERVICE_NAME")], + None, + Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": null, "db": 0}))), +)] +#[case::plaintext_ignores_tls_requirements( + json!({"host": "redis", "ssl_cert_reqs": "strict"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::an_ambiguous_flag_is_refused( + json!({"host": "redis", "ssl": "on"}), + &[], + None, + Refuses("invalid Redis setting ssl: expected a boolean"), +)] +#[case::empty_sentinel_nodes_without_the_environment_name_no_target( + json!({"sentinel_nodes": [], "service_name": "primary"}), + &[], + None, + Fails("either host or url must be specified for redis", Some(ValueError)), +)] +#[case::ssl_ca_path_is_unsupported( + json!({"host": "redis", "ssl": true, "ssl_ca_path": "/ca"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": "/ca", "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::ssl_validate_ocsp_is_unsupported( + json!({"host": "redis", "ssl": true, "ssl_validate_ocsp": true}), + &[], + None, + Fails("ssl_validate_ocsp asks for OCSP revocation checks, which rustls cannot perform, so Redis is not connected without them", Some(TypeError)), +)] +#[case::ssl_include_verify_flags_is_unsupported( + json!({"host": "redis", "ssl": true, "ssl_include_verify_flags": ["VERIFY_CRL_CHECK_LEAF"]}), + &[], + None, + Fails("ssl_include_verify_flags asks for X.509 verify flags such as CRL checking, which rustls cannot perform, so Redis is not connected without them", Some(TypeError)), +)] +#[case::ssl_password_is_unsupported( + json!({"host": "redis", "ssl": true, "ssl_certfile": "/client.pem", "ssl_password": "key-secret"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": "/client.pem", "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::host_with_every_basic_setting( + json!({"db": 2, "host": "redis", "password": "pass", "port": 6380, "socket_timeout": 1.5, "ssl": true, "username": "user"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6380, "db": 2, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": "user", "password": "pass"}, "socket_timeout": 1.5, "socket_connect_timeout": 1.5}))), +)] +#[case::environment_host_port_and_password( + json!({"port": "6390", "ssl": true, "host": "environment", "password": "environment-secret"}), + &[("host", "REDIS_HOST"), ("password", "REDIS_PASSWORD"), ("port", "REDIS_PORT"), ("ssl", "REDIS_SSL")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "environment", "port": 6390, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": null, "password": "environment-secret"}}))), +)] +#[case::string_numbers_are_coerced( + json!({"db": "3", "host": "redis", "port": "6380", "socket_timeout": "2.5"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6380, "db": 3, "socket_timeout": 2.5, "socket_connect_timeout": 2.5}))), +)] +#[case::padded_and_underscored_numbers_parse_like_python_int_and_float( + json!({"host": "redis", "port": "6380\n", "db": " 2 ", "socket_timeout": "1_0.5\n"}), + &[], + None, + Resolves(expected_view(json!({"db": 2, "host": "redis", "mode": "host", "port": 6380, "socket_timeout": 10.5, "socket_connect_timeout": 10.5}))), +)] +#[case::padded_environment_numbers_parse_like_python_int( + json!({"port": "6_380", "host": "redis", "db": "3\n"}), + &[("db", "REDIS_DB"), ("host", "REDIS_HOST"), ("port", "REDIS_PORT")], + None, + Resolves(expected_view(json!({"db": 3, "host": "redis", "mode": "host", "port": 6380}))), +)] +#[case::socket_timeout_from_the_environment( + json!({"socket_timeout": "0.75", "host": "redis"}), + &[("socket_timeout", "REDIS_SOCKET_TIMEOUT")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "socket_timeout": 0.75, "socket_connect_timeout": 0.75}))), +)] +#[case::environment_url( + json!({"url": "redis://environment:6379/1"}), + &[("url", "REDIS_URL")], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://environment:6379/1"}))), +)] +#[case::environment_url_beats_the_environment_host( + json!({"host": "ignored", "url": "redis://from-url:6379"}), + &[("host", "REDIS_HOST"), ("url", "REDIS_URL")], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://from-url:6379"}))), +)] +#[case::an_explicit_host_drops_the_environment_url( + json!({"host": "configured"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "configured", "port": 6379, "db": 0}))), +)] +#[case::a_url_drops_host_db_and_password_and_keeps_the_username( + json!({"db": 4, "host": "dropped", "password": "dropped", "url": "rediss://cache:6380", "username": "user"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "static", "username": "user", "password": null}}))), +)] +#[case::a_url_ignores_tls_settings_outside_the_url( + json!({"ssl_ca_certs": "/ca.pem", "ssl_cert_reqs": "none", "ssl_check_hostname": true, "url": "rediss://cache:6380"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::startup_nodes_beat_a_url( + json!({"password": "cluster-secret", "startup_nodes": [{"host": "node-a", "port": 7000}], "url": "redis://single:6379"}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "auth": {"kind": "static", "username": null, "password": "cluster-secret"}, "health_check_interval": 25.0}))), +)] +#[case::environment_cluster_nodes( + json!({"startup_nodes": "[{\"host\": \"node-a\", \"port\": \"7000\"}, {\"host\": \"node-b\", \"port\": 7001}]"}), + &[("startup_nodes", "REDIS_CLUSTER_NODES")], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000", "node-b:7001"], "health_check_interval": 25.0}))), +)] +#[case::cluster_tls_options( + json!({"ssl": "true", "ssl_ca_certs": "/ca.pem", "ssl_check_hostname": true, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": true}, "health_check_interval": 25.0}))), +)] +#[case::sentinel_from_settings( + json!({"password": "data-secret", "sentinel_nodes": [["sentinel-a", 26379], ["sentinel-b", 26380]], "sentinel_password": "sentinel-secret", "service_name": "primary"}), + &[], + None, + Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379", "sentinel-b:26380"], "service_name": "primary", "sentinel_password": "sentinel-secret", "db": 0, "auth": {"kind": "static", "username": null, "password": "data-secret"}}))), +)] +#[case::sentinel_from_the_environment( + json!({"sentinel_nodes": "[[\"sentinel-a\", 26379]]", "sentinel_password": "sentinel-secret", "service_name": "primary"}), + &[("sentinel_nodes", "REDIS_SENTINEL_NODES"), ("sentinel_password", "REDIS_SENTINEL_PASSWORD"), ("service_name", "REDIS_SERVICE_NAME")], + None, + Resolves(expected_view(json!({"mode": "sentinel", "nodes": ["sentinel-a:26379"], "service_name": "primary", "sentinel_password": "sentinel-secret", "db": 0}))), +)] +#[case::a_host_breaks_sentinel( + json!({"host": "ignored", "sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}), + &[], + None, + Fails("host cannot be combined with sentinel_nodes, which redis-py rejects at startup", Some(TypeError)), +)] +#[case::a_url_beats_sentinels( + json!({"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary", "url": "redis://single:6379"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://single:6379"}))), +)] +#[case::an_explicit_credential_provider_beats_cloud_auth( + json!({"aws_iam_auth": true, "credential_provider": true, "gcp_service_account": "account", "host": "redis", "password": "dropped"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "host"}}))), +)] +#[case::gcp_beats_azure_and_aws( + json!({"gcp_service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", "aws_iam_auth": true, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}), + &[("gcp_service_account", "REDIS_GCP_SERVICE_ACCOUNT")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "gcp_iam", "service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com"}}))), +)] +#[case::gcp_ca_certs_replace_the_configured_ones_under_ssl( + json!({"gcp_service_account": "account", "gcp_ssl_ca_certs": "/gcp-ca.pem", "host": "redis", "ssl": true, "ssl_ca_certs": "/configured-ca.pem"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/gcp-ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "gcp_iam", "service_account": "account"}}))), +)] +#[case::azure_beats_aws_and_takes_the_environment_username( + json!({"aws_iam_auth": true, "azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": "dropped"}), + &[], + Some("environment-user"), + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": "environment-user"}}))), +)] +#[case::an_empty_environment_username_sends_only_the_token( + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis", "username": ""}), + &[("username", "REDIS_USERNAME")], + Some(""), + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": null}}))), +)] +#[case::azure_without_a_username_sends_only_the_token( + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": true, "azure_tenant_id": "tenant", "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "auth": {"kind": "azure_ad", "username": null}}))), +)] +#[case::azure_needs_the_literal_true( + json!({"azure_client_id": "client", "azure_client_secret": "secret", "azure_redis_ad_token": "1", "azure_tenant_id": "tenant", "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::aws_iam_over_tls_with_the_default_region( + json!({"aws_iam_auth": "yes", "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "password": "dropped", "ssl": true, "aws_iam_region": "us-east-1"}), + &[("aws_iam_region", "AWS_DEFAULT_REGION")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "cache-1", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "us-east-1", "serverless": false}}))), +)] +#[case::aws_region_beats_the_default_region( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true, "aws_iam_region": "eu-central-1"}), + &[("aws_iam_region", "AWS_REGION")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "cache-1", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "eu-central-1", "serverless": false}}))), +)] +#[case::aws_iam_serverless_over_a_rediss_url_drops_its_userinfo( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_region": "eu-west-1", "aws_iam_serverless": true, "aws_iam_user_name": "iam-user", "url": "rediss://user:pass@cache-1:6379"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache-1:6379", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}, "auth": {"kind": "aws_iam", "user_name": "iam-user", "cache_name": "cache-1", "region": "eu-west-1", "serverless": true}}))), +)] +#[case::aws_iam_without_tls_fails( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_region": "us-east-1", "aws_iam_user_name": "iam-user", "host": "cache-1"}), + &[], + None, + Fails("AWS ElastiCache IAM Redis authentication requires TLS", Some(ValueError)), +)] +#[case::aws_iam_without_a_region_fails( + json!({"aws_iam_auth": true, "aws_iam_cache_name": "cache-1", "aws_iam_user_name": "iam-user", "host": "cache-1", "ssl": true}), + &[], + None, + Fails("AWS ElastiCache IAM Redis authentication requires: aws_iam_region", Some(ValueError)), +)] +#[case::tls_options_from_settings( + json!({"host": "redis", "ssl": true, "ssl_ca_certs": "/ca.pem", "ssl_cert_reqs": "none", "ssl_check_hostname": true}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": true}}))), +)] +#[case::a_null_cert_requirement_skips_verification( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": null}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": false}}))), +)] +#[case::a_verify_mode_of_none_skips_verification( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": 0}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": false}}))), +)] +#[case::a_verify_mode_of_optional_verifies( + json!({"host": "redis", "ssl": true, "ssl_cert_reqs": 1}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::a_cert_requirement_spelled_like_python_none_fails( + json!({"ssl": true, "ssl_cert_reqs": "None", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")], + None, + Fails("invalid Redis setting REDIS_SSL_CERT_REQS: expected none, optional, required, null or an ssl.VerifyMode value", None), +)] +#[case::tls_options_from_the_environment( + json!({"ssl_ca_certs": "/env-ca.pem", "ssl": true, "ssl_cert_reqs": "optional", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_ca_certs", "REDIS_SSL_CA_CERTS"), ("ssl_cert_reqs", "REDIS_SSL_CERT_REQS")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": "/env-ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::connect_timeout_is_separate_from_the_socket_timeout( + json!({"host": "redis", "socket_connect_timeout": 0.25, "socket_timeout": 2.5}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "socket_timeout": 2.5, "socket_connect_timeout": 0.25}))), +)] +#[case::environment_references_resolve( + json!({"host": "referenced", "password": "referenced-secret"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "referenced", "port": 6379, "db": 0, "auth": {"kind": "static", "username": null, "password": "referenced-secret"}}))), +)] +#[case::no_target_fails( + json!({"port": 6379}), + &[], + None, + Fails("either host or url must be specified for redis", Some(ValueError)), +)] +#[case::client_certificates( + json!({"host": "redis", "ssl": true, "ssl_certfile": "/client.pem", "ssl_keyfile": "/client.key"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": "/client.pem", "keyfile": "/client.key", "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::a_rediss_url_takes_tls_options_from_its_query( + json!({"url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes", "tls": {"ca_certs": "/ca.pem", "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": false, "check_hostname": true}}))), +)] +#[case::a_url_query_carries_socket_timeout( + json!({"url": "redis://cache:6379/0?socket_timeout=2"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_timeout=2", "socket_timeout": 2.0, "socket_connect_timeout": 2.0}))), +)] +#[case::url_query_settings_win_over_the_kwargs( + json!({"health_check_interval": 4, "socket_connect_timeout": 7, "socket_timeout": 1, "url": "redis://cache:6379/0?socket_connect_timeout=0.5&health_check_interval=9&socket_timeout=3"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_connect_timeout=0.5&health_check_interval=9&socket_timeout=3", "socket_timeout": 3.0, "socket_connect_timeout": 0.5, "health_check_interval": 9.0}))), +)] +#[case::a_repeated_url_query_parameter_takes_its_first_value( + json!({"url": "redis://cache:6379/0?socket_timeout=2&socket_timeout=9"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "redis://cache:6379/0?socket_timeout=2&socket_timeout=9", "socket_timeout": 2.0, "socket_connect_timeout": 2.0}))), +)] +#[case::a_url_query_number_parse_url_cannot_read( + json!({"url": "redis://cache:6379/0?socket_timeout=slow"}), + &[], + None, + Fails("invalid Redis setting url ?socket_timeout: redis-py's parse_url cannot read it as a number", Some(ValueError)), +)] +#[case::tls_settings_outside_a_url_never_fail_it( + json!({"db": "99999999999999999999", "ssl_cert_reqs": "strict", "url": "rediss://cache:6380"}), + &[], + None, + Resolves(expected_view(json!({"mode": "url", "url": "rediss://cache:6380", "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::db_is_not_read_for_a_cluster( + json!({"db": 99_999_999_999_i64, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::cluster_db_true( + json!({"db": true, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::cluster_db_false( + json!({"db": false, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), +)] +#[case::cluster_db_fraction( + json!({"db": 0.5, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::cluster_db_float_zero( + json!({"db": 0.0, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), +)] +#[case::cluster_db_nonempty_list( + json!({"db": [1], "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::cluster_db_empty_list( + json!({"db": [], "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"], "health_check_interval": 25.0}))), +)] +#[case::ssl_ciphers( + json!({"host": "redis", "ssl": true, "ssl_ciphers": "HIGH"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::retry_on_timeout( + json!({"host": "redis", "retry_on_timeout": true}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::malformed_port( + json!({"host": "redis", "port": "abc"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::malformed_db( + json!({"db": "first", "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::socket_timeout_not_a_number( + json!({"host": "redis", "socket_timeout": "slow"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::malformed_port_from_the_environment( + json!({"port": "abc", "host": "redis"}), + &[("host", "REDIS_HOST"), ("port", "REDIS_PORT")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::protocol_3( + json!({"host": "redis", "protocol": 3}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::retry_on_error_from_the_environment( + json!({"retry_on_error": true, "host": "redis"}), + &[("retry_on_error", "REDIS_RETRY_ON_ERROR")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::unix_socket_path( + json!({"host": "redis", "unix_socket_path": "/tmp/redis.sock"}), + &[], + None, + Resolves(expected_view(json!({"mode": "unix", "path": "/tmp/redis.sock"}))), +)] +#[case::unix_socket_path_keeps_tls_settings_out_of_the_target( + json!({"host": "redis", "ssl": true, "ssl_ca_certs": "/ca.pem", "unix_socket_path": "/tmp/redis.sock"}), + &[], + None, + Resolves(expected_view(json!({"mode": "unix", "path": "/tmp/redis.sock"}))), +)] +#[case::ssl_ca_data( + json!({"host": "redis", "ssl": true, "ssl_ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----", "ca_path": null, "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::ssl_ca_path_from_the_environment( + json!({"ssl": true, "ssl_ca_path": "/etc/ssl/private-cas", "host": "redis"}), + &[("ssl", "REDIS_SSL"), ("ssl_ca_path", "REDIS_SSL_CA_PATH")], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "tls": {"ca_certs": null, "ca_data": null, "ca_path": "/etc/ssl/private-cas", "certfile": null, "keyfile": null, "verify_certificate": true, "check_hostname": false}}))), +)] +#[case::health_check_interval_is_whole_seconds( + json!({"health_check_interval": "7", "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0, "health_check_interval": 7.0}))), +)] +#[case::health_check_interval_zero_is_off( + json!({"health_check_interval": 0, "host": "redis"}), + &[], + None, + Resolves(expected_view(json!({"mode": "host", "host": "redis", "port": 6379, "db": 0}))), +)] +#[case::health_check_interval_zero_keeps_a_cluster_off( + json!({"health_check_interval": 0, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Resolves(expected_view(json!({"mode": "cluster", "nodes": ["node-a:7000"]}))), +)] +#[case::a_cluster_rejects_a_nonzero_db( + json!({"db": 3, "startup_nodes": [{"host": "node-a", "port": 7000}]}), + &[], + None, + Fails("db must be 0 in cluster mode", Some(RedisClusterException)), +)] +#[case::unix_socket_path_without_a_host_fails_like_python( + json!({"unix_socket_path": "/tmp/redis.sock"}), + &[], + None, + Fails("either host or url must be specified for redis", Some(ValueError)), +)] +fn resolves_like_python_redis( + #[case] merged: serde_json::Value, + #[case] sources: &[(&str, &str)], + #[case] redis_username: Option<&str>, + #[case] expected: Outcome, +) { + let outcome = interpret_merged(merged, sources, redis_username); + let rendered = outcome + .as_ref() + .map(view) + .map_err(|error| error.to_string()); + + match expected { + Resolves(expected) => assert_eq!(rendered, Ok(expected)), + Refuses(message) => assert_eq!(rendered, Err(message.to_owned())), + Fails(message, raised) => { + assert_eq!(rendered, Err(message.to_owned())); + assert_eq!( + outcome + .err() + .and_then(|error| error.raised_building_the_python_client()), + raised + ); + } + } +} diff --git a/litellm-rust/crates/redis-store/AGENTS.md b/litellm-rust/crates/redis-store/AGENTS.md new file mode 100644 index 00000000000..9aad49d1d71 --- /dev/null +++ b/litellm-rust/crates/redis-store/AGENTS.md @@ -0,0 +1,13 @@ +- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (Python-facing values with their `FromValue` / `IntoValue`). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`) +- An op declared on a `Send` returns `Pending`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on +- A store that sends a pipeline of its own uses `RedisBatch::pipelined`, never `Pipeline` and raw `Value`s +- Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec +- A family holds one `DualCache` per `DualCache` its Python held on `main`, and each `DualCache` method is the Python method it names, so a family op that was a `DualCache` call is one call here. Where `main`'s Python reached past `DualCache` to `in_memory_cache` or `redis_cache`, the family does the same through `memory()` and `redis()`, which keeps that reach visible. Caches a family's Python attached to one `RedisCache` share one `Arc`. Only a family that batch-reads builds its cache with `DualCache::batch_reading` from a `TierPolicy` the host passes in, and the Python host fills it from what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `Remote` (or none), never a cache. Redis-only families are plain functions over `&RedisClient` +- Every family works with no Redis, answering from memory. `attach` adds the Redis tier at most once and keeps memory, as `DualCache.attach_redis_cache` does, so a host builds the store at startup and never rebuilds it to connect +- Where `main` swallows a Redis failure and serves memory, the family does the same and logs. It never turns a failure into a miss that `main` would not +- Key suffixes and wire spellings derive from strum on the family's enums, and a type whose Python side is plain data implements `FromValue` / `IntoValue` over `litellm_python_compat::Value`, so the bridge never names a field or a spelling +- Wire contracts are hand-written cases with expected values, inlined as the same literals in this crate's test and in `tests/unit/redis/stores/`, checked here against real Redis and there against Python's own codec. Change both sides together +- `rate_limits.rs` keeps the v3 limiter's Lua byte for byte (`rate_limits/lua.rs`) so pods on either stack share one script SHA; the admission decisions stay in the Python limiter +- `RequestRedis` is one request's batches: a pre-call and a post-call `RedisBatch` per client, the post one flushed by `flush_post` or its deadline. Each op exists once and takes a `Send` (`Now`, `PreCall`, `PostCall`) picked by the caller, and the store calls `send.batch(&client)`. A store never encodes a phase in a method name. A post-call method exists only where a follow-up runs after the write settles (`token_increment_post_call`), as a thin wrapper over the shared op. `Send::Now` is an alone batch, so a script call outside a request goes out alone, like a redis-py `Script` outside a pipeline, so `NOSCRIPT` costs `SCRIPT LOAD` and one more `EVALSHA` the way it does on main. Only what main's Python puts on its request pipeline rides these batches. An op main sends on its own round trip, like a deployment's tpm and rpm increment, stays on its own round trip here +- Every store sends the commands main's `RedisCache` sends for that operation: `MGET` where Python reads several keys at once, `SET ... EX`, and `INCRBYFLOAT` with `EXPIRE` or `TTL`. A `MONITOR` diff against main is a parity check, and any difference is a defect +- Tests run against `litellm-redis-testing` on standalone and cluster. A test is single-topology only for outages diff --git a/litellm-rust/crates/redis-store/Cargo.toml b/litellm-rust/crates/redis-store/Cargo.toml new file mode 100644 index 00000000000..1dd61d76d3c --- /dev/null +++ b/litellm-rust/crates/redis-store/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "litellm-redis-store" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +futures-util.workspace = true +litellm-cache.workspace = true +litellm-cache-memory.workspace = true +litellm-python-compat.workspace = true +litellm-redis-client.workspace = true +strum.workspace = true +thiserror.workspace = true +tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } +tracing.workspace = true + +[dev-dependencies] +litellm-circuit-breaker.workspace = true +litellm-redis-testing.workspace = true +litellm-secret-value.workspace = true +litellm-tracing.workspace = true +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/redis-store/src/codec.rs b/litellm-rust/crates/redis-store/src/codec.rs new file mode 100644 index 00000000000..21cde7d0b81 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/codec.rs @@ -0,0 +1,74 @@ +use litellm_python_compat::{Value, json, literal, repr::to_str}; + +use crate::{Error, FromValue, IntoValue, Number}; + +#[derive(Debug, thiserror::Error)] +pub enum DecodeError { + #[error("Redis value is not UTF-8")] + NotUtf8, + #[error("Redis value is neither JSON nor a Python literal: {0}")] + Unparsed(#[source] litellm_python_compat::Error), + #[error("Redis value has the wrong shape: {0}")] + Shape(&'static str), +} + +/// `RedisCache._get_cache_logic`: `json.loads`, falling back to `ast.literal_eval`. +pub(crate) fn python_value(raw: &[u8]) -> Result { + let text = std::str::from_utf8(raw).map_err(|_| DecodeError::NotUtf8)?; + json::loads(text).or_else(|_| literal::literal_eval(text).map_err(DecodeError::Unparsed)) +} + +/// A cached value as `RedisCache._get_cache_logic` reads it. +pub(crate) fn decode(raw: &[u8]) -> Result { + V::from_value(&python_value(raw)?) +} + +/// `json.dumps(value)`, what `RedisCache.async_set_cache` writes. +pub(crate) fn dumps(value: V) -> Result { + Ok(json::dumps(&value.into_value())?) +} + +/// `str(value)`, what the sync `RedisCache.set_cache` writes. +pub(crate) fn str(value: V) -> String { + to_str(&value.into_value()) +} + +/// A number written by Python, `shape` naming what it was meant to be when it is not one. +pub(crate) fn number(raw: &[u8], shape: &'static str) -> Result { + Number::from_value(&python_value(raw)?).ok_or(DecodeError::Shape(shape)) +} + +#[cfg(test)] +mod tests { + use litellm_python_compat::{Value, json, repr::repr}; + use rstest::rstest; + + use super::{number, python_value}; + use crate::Number; + + // Expected values are the repr of main's RedisCache._get_cache_logic on these bytes + #[rstest] + #[case::nan(b"NaN", "nan")] + #[case::infinity(b"Infinity", "inf")] + #[case::negative_infinity(b"-Infinity", "-inf")] + #[case::nan_inside_a_document(br#"{"spend": NaN}"#, "{'spend': nan}")] + #[case::python_literal_fallback(b"{'a': (1, 2)}", "{'a': (1, 2)}")] + fn reads_what_python_reads(#[case] raw: &[u8], #[case] expected: &str) { + assert_eq!(repr(&python_value(raw).unwrap()), expected); + } + + #[rstest] + #[case::nan(f64::NAN)] + #[case::infinity(f64::INFINITY)] + #[case::negative_infinity(f64::NEG_INFINITY)] + fn a_non_finite_number_reads_back_as_written(#[case] value: f64) { + let written = json::dumps(&Value::Float(value)).unwrap(); + + let read = number(written.as_bytes(), "a number").unwrap(); + + assert_eq!( + repr(&read.to_value()), + repr(&Number::Float(value).to_value()) + ); + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs b/litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs new file mode 100644 index 00000000000..71dfda49233 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/batch_read.rs @@ -0,0 +1,242 @@ +use std::{ + collections::{HashMap, VecDeque}, + sync::Mutex, + time::Duration, +}; + +use litellm_redis_client::{Declared, Error as RedisError, commands}; + +use super::{DualCache, MemorySettings}; +use crate::{Error, FromValue, Send, codec, memory::lock, request::Misses}; + +/// A batch-reading family's `DualCache` settings, as the host reads them from `litellm` globals. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TierPolicy { + /// Python's `redis_batch_cache_expiry`: a key memory misses is read from Redis at most this often. + pub refresh_after: Duration, + /// Python's `default_max_redis_batch_cache_size`: how many keys the refresh throttle tracks. + pub tracked_reads: usize, + /// `litellm.default_in_memory_ttl`, which a value read from Redis keeps in memory. + pub default_in_memory_ttl: Option, +} + +impl TierPolicy { + pub(crate) fn memory(&self) -> MemorySettings { + MemorySettings { + entries: None, + ttl: self.default_in_memory_ttl, + } + } +} + +/// `DualCache.last_redis_batch_access_time` with its expiry, which only `async_batch_get_cache` reads. +pub(crate) struct BatchReads { + refresh_after: Duration, + recent: Mutex, +} + +impl BatchReads { + pub(super) fn new(policy: TierPolicy) -> Self { + Self { + refresh_after: policy.refresh_after, + recent: Mutex::new(RecentReads::new(policy.tracked_reads)), + } + } +} + +impl DualCache { + /// `DualCache.async_batch_get_cache`: memory first, then one `MGET` for the misses whose refresh + /// is due, written back to memory. A failed Redis read serves memory, as `RedisCache` swallows + /// it, and only a breaker refusal releases the refresh it reserved. + pub(crate) async fn batch_get(&self, keys: &[String]) -> Result>, Error> { + let declared = self.declare_batch_get(Send::Now, keys)?; + self.resolve_batch_get(declared).await + } + + /// The first half of `batch_get`: reserves the due refreshes and declares their `MGET` on the + /// batch `send` picks, without sending anything. + pub(crate) fn declare_batch_get( + &self, + send: Send<'_>, + keys: &[String], + ) -> Result { + let fetch = match self.redis.client() { + None => None, + Some(client) => { + let reservation = self.reserve(keys, &self.read_memory(keys)?); + (!reservation.keys.is_empty()).then(|| { + let keys: Vec<&str> = reservation.keys.iter().map(String::as_str).collect(); + let replies = send.batch(client).add(commands::mget::>(&keys)); + Fetch { + reservation, + replies, + misses: send.misses(client), + } + }) + } + }; + Ok(DeclaredRead { + keys: keys.to_vec(), + fetch, + }) + } + + /// The second half of `batch_get`. Memory is read again here, so a value written to memory + /// since the declaration wins over what Redis answered. The keys Redis lacks are noted as the + /// request's misses. An undecodable value serves memory for the whole read, where `main`'s + /// `RedisCache` returns the raw bytes and `DualCache` backfills memory with them. + pub(crate) async fn resolve_batch_get( + &self, + declared: DeclaredRead, + ) -> Result>, Error> { + let DeclaredRead { keys, fetch } = declared; + let Some(Fetch { + reservation, + replies, + misses, + }) = fetch + else { + return self.read_memory(&keys); + }; + let fetched = replies.resolve().await; + let local = self.read_memory(&keys)?; + let fetched = match fetched { + Ok(fetched) => fetched, + Err(RedisError::BreakerOpen) => { + self.release(reservation); + return Ok(local); + } + Err(error) => { + crate::failure_log::redis_failure!( + &error, + "Redis read failed, serving the memory tier" + ); + return Ok(local); + } + }; + if let Some(misses) = misses { + misses.note( + reservation + .keys + .iter() + .zip(&fetched) + .filter(|(_, raw)| raw.is_none()) + .map(|(key, _)| key.clone()), + ); + } + let decoded = match fetched + .iter() + .map(|raw| raw.as_deref().map(codec::decode).transpose()) + .collect::, _>>() + { + Ok(decoded) => decoded, + Err(error) => { + tracing::warn!(%error, "a Redis value did not decode, serving the memory tier"); + return Ok(local); + } + }; + let read = reservation + .keys + .into_iter() + .zip(decoded) + .collect::>(); + keys.iter() + .zip(local) + .map(|(key, local)| match local { + Some(value) => Ok(Some(value)), + None => self.backfill(key, read.get(key).cloned().flatten(), None), + }) + .collect() + } + + fn read_memory(&self, keys: &[String]) -> Result>, Error> { + keys.iter().map(|key| Ok(self.memory.get(key)?)).collect() + } + + fn reserve(&self, keys: &[String], local: &[Option]) -> Reservation { + let now = self.now(); + let refresh_after = self.reads.refresh_after; + let mut recent = lock(&self.reads.recent); + let due = keys + .iter() + .zip(local) + .filter(|(_, value)| value.is_none()) + .filter_map(|(key, _)| { + let previous = recent.get(key); + let due = previous.is_none_or(|read| now.saturating_sub(read) >= refresh_after); + due.then(|| { + recent.insert(key.clone(), now); + (key.clone(), previous) + }) + }) + .collect::>(); + let (keys, previous) = due.into_iter().unzip(); + Reservation { keys, previous } + } + + fn release(&self, reservation: Reservation) { + let mut recent = lock(&self.reads.recent); + for (key, previous) in reservation.keys.into_iter().zip(reservation.previous) { + match previous { + Some(read) => recent.insert(key, read), + None => recent.remove(&key), + } + } + } +} + +/// The `MGET` `declare_batch_get` put on a batch, with the refreshes it reserved. +pub(crate) struct DeclaredRead { + keys: Vec, + fetch: Option, +} + +struct Fetch { + reservation: Reservation, + replies: Declared>>>, + misses: Option, +} + +struct Reservation { + keys: Vec, + previous: Vec>, +} + +/// Python's `LimitedSizeOrderedDict`: insertion-ordered, and any write at capacity drops the +/// oldest entry first, even when the written key is already present. +struct RecentReads { + capacity: usize, + order: VecDeque, + times: HashMap, +} + +impl RecentReads { + fn new(capacity: usize) -> Self { + Self { + capacity, + order: VecDeque::new(), + times: HashMap::new(), + } + } + + fn get(&self, key: &str) -> Option { + self.times.get(key).copied() + } + + fn insert(&mut self, key: String, time: Duration) { + if self.times.len() >= self.capacity + && let Some(oldest) = self.order.pop_front() + { + self.times.remove(&oldest); + } + if self.times.insert(key.clone(), time).is_none() { + self.order.push_back(key); + } + } + + fn remove(&mut self, key: &str) { + if self.times.remove(key).is_some() { + self.order.retain(|tracked| tracked != key); + } + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/increment.rs b/litellm-rust/crates/redis-store/src/dual_cache/increment.rs new file mode 100644 index 00000000000..705bc71b2c1 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/increment.rs @@ -0,0 +1,129 @@ +use std::time::Duration; + +use futures_util::future::join_all; +use litellm_redis_client::{Declared, Expiry, RedisBatch, commands}; + +use super::DualCache; +use crate::{Error, Number, RequestRedis, Send, failure_log::redis_failure, memory::Counter}; + +const INCREMENT_FAILED: &str = + "Redis async_increment_cache failed, falling back to in-memory result"; +const PIPELINE_FAILED: &str = + "Redis async_increment_cache_pipeline failed, falling back to in-memory result"; + +/// A `RedisPipelineIncrementOperation`. +pub(crate) struct PipelineIncrement { + pub(crate) key: String, + pub(crate) amount: Number, + pub(crate) ttl: Option, +} + +impl DualCache { + /// `DualCache.async_increment_cache`: memory, then `INCRBYFLOAT` keeping an existing TTL, else + /// setting `ttl` or Redis's default. `None` when memory refuses the increment, which then never + /// reaches Redis, and the memory count when Redis fails. Both only log, as there. + pub(crate) async fn increment( + &self, + key: &str, + amount: Number, + ttl: Option, + ) -> Option { + let local = self + .memory + .increment(key, amount, ttl) + .map_err(|error| redis_failure!(&Error::from(error), INCREMENT_FAILED)) + .ok()?; + let Some(remote) = self.redis.get() else { + return Some(local); + }; + let expiry = Expiry::KeepExisting(ttl.unwrap_or(remote.default_ttl)); + let increment = commands::incr_by_float(key, amount.into(), Some(expiry)); + match remote.client.send(increment).await { + Ok(total) => Some(Number::Float(total)), + Err(error) => { + redis_failure!(&error, INCREMENT_FAILED); + Some(local) + } + } + } + + /// `DualCache.async_increment_cache_pipeline`: memory, then one pipeline of `INCRBYFLOAT` with + /// `EXPIRE` per increment. `None` when memory refuses an increment, which leaves the ones before + /// it applied and sends nothing, and the memory totals when Redis fails. + pub(crate) async fn increment_pipeline( + &self, + increments: &[PipelineIncrement], + ) -> Option> { + let local = increments + .iter() + .map(|increment| { + self.memory + .increment(&increment.key, increment.amount, increment.ttl) + }) + .collect::, _>>() + .map_err(|error| redis_failure!(&Error::from(error), PIPELINE_FAILED)) + .ok()?; + let Some(client) = self.redis.client() else { + return Some(local); + }; + let pipeline = RedisBatch::pipelined(client.clone()); + let totals = increments + .iter() + .map(|increment| { + pipeline.add(commands::incr_by_float( + &increment.key, + increment.amount.into(), + increment.ttl.map(Expiry::Refresh), + )) + }) + .collect::>(); + match Declared::all(totals) + .await + .into_iter() + .collect::, _>>() + { + Ok(totals) => Some(totals.into_iter().map(Number::Float).collect()), + Err(error) => { + redis_failure!(&error, PIPELINE_FAILED); + Some(local) + } + } + } + + /// `DualCache.async_increment_cache_pipeline_post_call`: memory now, stopping at the first + /// increment memory refuses and only logging it, then each increment on the request's post-call + /// batch, whose failure is logged once it settles. + pub(crate) fn increment_pipeline_post_call( + &self, + request: &RequestRedis, + increments: &[PipelineIncrement], + ) { + if let Err(error) = increments.iter().try_for_each(|increment| { + self.memory + .increment(&increment.key, increment.amount, increment.ttl) + .map(drop) + }) { + tracing::warn!(%error, "in-memory increment failed"); + } + let Some(batch) = self.redis.batch(Send::PostCall(request)) else { + return; + }; + let declared = increments + .iter() + .map(|increment| { + batch.add(commands::incr_by_float( + &increment.key, + increment.amount.into(), + increment.ttl.map(Expiry::KeepExisting), + )) + }) + .collect::>(); + request.spawn(async move { + join_all(declared.into_iter().map(Declared::settled)) + .await + .into_iter() + .filter_map(Result::err) + .for_each(|error| redis_failure!(&error, "deferred Redis increment failed")); + }); + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/mod.rs b/litellm-rust/crates/redis-store/src/dual_cache/mod.rs new file mode 100644 index 00000000000..b2c277a118e --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/mod.rs @@ -0,0 +1,279 @@ +mod batch_read; +mod increment; +mod read; +mod write; + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +pub use self::batch_read::TierPolicy; +pub(crate) use self::{batch_read::BatchReads, increment::PipelineIncrement}; +use crate::{ + memory::{Clock, Memory, lock}, + remote::RemoteSlot, +}; + +/// The memory tier's `InMemoryCache` size and the `default_in_memory_ttl` a write or a Redis +/// backfill takes when it names no TTL. `None` keeps `InMemoryCache`'s own default. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) struct MemorySettings { + pub(crate) entries: Option, + pub(crate) ttl: Option, +} + +/// Python's `DualCache`: a per-worker `InMemoryCache` in front of the Redis every pod shares. Each +/// method is the `DualCache` method it names. A family whose Python reached past `DualCache` to its +/// `in_memory_cache` or `redis_cache` uses `memory` and `redis` the same way. +pub(crate) struct DualCache { + memory: Memory, + redis: Arc, + default_in_memory_ttl: Mutex>, + now: Clock, + reads: R, +} + +impl DualCache { + /// `redis` may be shared with the family's other caches, as one `RedisCache` is attached to each. + pub(crate) fn new(redis: Arc, memory: MemorySettings, now: Clock) -> Self { + Self { + memory: Memory::new(memory.entries, &now), + redis, + default_in_memory_ttl: Mutex::new(memory.ttl), + now, + reads: (), + } + } + + /// A cache that batch-reads, with the `async_batch_get_cache` throttle `policy` configures. + pub(crate) fn batch_reading( + redis: Arc, + policy: TierPolicy, + now: Clock, + ) -> DualCache { + let Self { + memory, + redis, + default_in_memory_ttl, + now, + reads: (), + } = Self::new(redis, policy.memory(), now); + DualCache { + memory, + redis, + default_in_memory_ttl, + now, + reads: BatchReads::new(policy), + } + } +} + +impl DualCache { + pub(crate) fn memory(&self) -> &Memory { + &self.memory + } + + pub(crate) fn redis(&self) -> &RemoteSlot { + &self.redis + } + + pub(crate) fn now(&self) -> Duration { + (self.now)() + } + + pub(crate) fn default_in_memory_ttl(&self) -> Option { + *lock(&self.default_in_memory_ttl) + } + + /// `DualCache.update_cache_ttl` and `update_in_memory_max_size` together: both apply in place, + /// and the entries memory already holds stay. + pub(crate) fn reconfigure(&self, memory: MemorySettings) { + *lock(&self.default_in_memory_ttl) = memory.ttl; + self.memory.update_max_size(memory.entries); + } +} + +#[cfg(test)] +mod tests { + use std::{ + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, + }; + + use rstest::rstest; + use tokio::runtime::Handle; + + use super::{DualCache, MemorySettings, PipelineIncrement}; + use crate::{Clock, Number, RequestRedis, memory::Counter, remote::RemoteSlot}; + + #[derive(Clone, Debug, PartialEq)] + enum Entry { + Number(Number), + Text, + } + + impl From for Entry { + fn from(number: Number) -> Self { + Self::Number(number) + } + } + + impl Counter for Entry { + fn as_number(&self) -> Option { + match self { + Self::Number(number) => Some(*number), + Self::Text => None, + } + } + } + + fn clock() -> (Clock, Arc) { + let seconds = Arc::new(AtomicU64::new(1_000)); + let read = seconds.clone(); + let clock: Clock = Arc::new(move || Duration::from_secs(read.load(Ordering::SeqCst))); + (clock, seconds) + } + + fn without_redis( + default_in_memory_ttl: Option, + now: Clock, + ) -> DualCache { + DualCache::new( + Arc::new(RemoteSlot::new(None)), + MemorySettings { + entries: None, + ttl: default_in_memory_ttl, + }, + now, + ) + } + + fn increment(key: &str, amount: Number) -> PipelineIncrement { + PipelineIncrement { + key: key.to_owned(), + amount, + ttl: None, + } + } + + // InMemoryCache's own default TTL is 600 s, which an entry outlives here only when the cache + // configured none + #[rstest] + #[case::default_in_memory_ttl(Some(Duration::from_secs(5)), None)] + #[case::in_memory_cache_default(None, Some(Number::Int(1)))] + #[tokio::test] + async fn a_set_without_a_ttl_takes_default_in_memory_ttl( + #[case] default_in_memory_ttl: Option, + #[case] after_six_seconds: Option, + ) { + let (now, seconds) = clock(); + let cache = without_redis(default_in_memory_ttl, now); + + cache.set("k", Number::Int(1), None).await.unwrap(); + seconds.fetch_add(6, Ordering::SeqCst); + + assert_eq!(cache.memory().get("k").unwrap(), after_six_seconds); + } + + #[rstest] + #[case::absent_starts_at_int_zero(None, Number::Int(2), Number::Int(2))] + #[case::int_plus_int_stays_int(Some(Number::Int(1)), Number::Int(2), Number::Int(3))] + #[case::a_float_makes_a_float(Some(Number::Int(1)), Number::Float(0.5), Number::Float(1.5))] + fn a_local_increment_adds_as_python_adds( + #[case] stored: Option, + #[case] amount: Number, + #[case] expected: Number, + ) { + let (now, _) = clock(); + let cache = without_redis::(None, now); + if let Some(stored) = stored { + cache.memory().set("k", stored, None).unwrap(); + } + + let total = cache.memory().increment("k", amount, None).unwrap(); + + assert_eq!(total, expected); + assert_eq!(cache.memory().get("k").unwrap(), Some(expected)); + } + + #[test] + fn a_local_increment_of_a_non_number_fails_and_keeps_the_entry() { + let (now, _) = clock(); + let cache = without_redis::(None, now); + cache.memory().set("k", Entry::Text, None).unwrap(); + + let total = cache.memory().increment("k", Number::Int(1), None); + + assert!(total.is_err()); + assert_eq!(cache.memory().get("k").unwrap(), Some(Entry::Text)); + } + + // async_increment_cache and async_increment_cache_pipeline catch the in-memory failure with the + // Redis one, so the caller gets None and nothing is sent + #[tokio::test] + async fn an_increment_memory_refuses_answers_none() { + let (now, _) = clock(); + let cache = without_redis::(None, now); + cache.memory().set("text", Entry::Text, None).unwrap(); + + let single = cache.increment("text", Number::Int(1), None).await; + let pipelined = cache + .increment_pipeline(&[ + increment("before", Number::Int(1)), + increment("text", Number::Int(1)), + ]) + .await; + + assert_eq!(single, None); + assert_eq!(pipelined, None); + assert_eq!(cache.memory().get("text").unwrap(), Some(Entry::Text)); + } + + // update_cache_ttl and update_in_memory_max_size change the settings in place, so what memory + // already holds survives and the next write takes the new TTL + #[tokio::test] + async fn reconfigure_keeps_entries_and_applies_the_new_ttl() { + let (now, seconds) = clock(); + let cache = without_redis(None, now); + cache.set("kept", Number::Int(1), None).await.unwrap(); + + cache.reconfigure(MemorySettings { + entries: Some(10), + ttl: Some(Duration::from_secs(5)), + }); + cache.set("short", Number::Int(2), None).await.unwrap(); + seconds.fetch_add(6, Ordering::SeqCst); + + assert_eq!(cache.memory().get("kept").unwrap(), Some(Number::Int(1))); + assert_eq!(cache.memory().get("short").unwrap(), None); + } + + // async_increment_cache_pipeline_post_call: InMemoryCache.async_increment_pipeline raises at the + // first entry it cannot add to, so later increments never reach memory, and DualCache only logs + #[tokio::test] + async fn a_post_call_increment_stops_at_the_first_entry_memory_refuses() { + let (now, _) = clock(); + let cache = without_redis::(None, now); + let request = RequestRedis::new(Handle::current(), Duration::from_secs(60)); + cache.memory().set("text", Entry::Text, None).unwrap(); + + cache.increment_pipeline_post_call( + &request, + &[ + increment("before", Number::Int(1)), + increment("text", Number::Int(1)), + increment("after", Number::Int(1)), + ], + ); + + assert_eq!( + cache.memory().get("before").unwrap(), + Some(Entry::Number(Number::Int(1))) + ); + assert_eq!(cache.memory().get("after").unwrap(), None); + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/read.rs b/litellm-rust/crates/redis-store/src/dual_cache/read.rs new file mode 100644 index 00000000000..14156483c6f --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/read.rs @@ -0,0 +1,55 @@ +use std::time::Duration; + +use litellm_redis_client::commands; + +use super::DualCache; +use crate::{Error, FromValue, RequestRedis, codec, failure_log::redis_failure}; + +impl DualCache { + /// `DualCache.async_get_cache`: memory, else `None` for a key `request` already read as missing + /// from Redis, else a `GET` written back to memory for `ttl` or `default_in_memory_ttl`. A failed + /// or undecodable Redis read is a miss, as there. + pub(crate) async fn get( + &self, + request: Option<&RequestRedis>, + key: &str, + ttl: Option, + ) -> Result, Error> { + if let Some(value) = self.memory.get(key)? { + return Ok(Some(value)); + } + let Some(client) = self.redis.client() else { + return Ok(None); + }; + if request.is_some_and(|request| request.misses(client).contains(key)) { + return Ok(None); + } + let value = client + .send(commands::get::>(key)) + .await + .map_err(Error::from) + .and_then(|raw| Ok(raw.as_deref().map(codec::decode).transpose()?)) + .unwrap_or_else(|error| { + redis_failure!(&error, "LiteLLM Cache: exception in async_get_cache"); + None + }); + self.backfill(key, value, ttl) + } +} + +impl DualCache { + /// `_backfill_kwargs`: a value read from Redis stays in memory for `ttl` or + /// `default_in_memory_ttl`. + pub(super) fn backfill( + &self, + key: &str, + value: Option, + ttl: Option, + ) -> Result, Error> { + if let Some(value) = &value { + self.memory + .set(key, value.clone(), ttl.or(self.default_in_memory_ttl()))?; + } + Ok(value) + } +} diff --git a/litellm-rust/crates/redis-store/src/dual_cache/write.rs b/litellm-rust/crates/redis-store/src/dual_cache/write.rs new file mode 100644 index 00000000000..6fcc6f32741 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/dual_cache/write.rs @@ -0,0 +1,178 @@ +use std::time::Duration; + +use litellm_redis_client::{Declared, RedisBatch, commands}; + +use super::DualCache; +use crate::{ + Error, IntoValue, Pending, Send, codec, + failure_log::redis_failure, + request::{pending, ready}, +}; + +/// `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`, the chunk `async_delete_cache_keys` deletes at once. +const DELETE_CHUNK: usize = 1000; + +impl DualCache { + /// `DualCache.async_set_cache`: memory, then a `SET` of `json.dumps(value)`. A Redis failure, or + /// a value `json.dumps` refuses, only logs. + pub(crate) async fn set( + &self, + key: &str, + value: V, + ttl: Option, + ) -> Result<(), Error> { + self.set_now(key, value, ttl, codec::dumps).await + } + + /// `DualCache.set_cache`, whose sync `RedisCache.set_cache` writes `str(value)`. + pub(crate) async fn set_str( + &self, + key: &str, + value: V, + ttl: Option, + ) -> Result<(), Error> { + self.set_now(key, value, ttl, |value| Ok(codec::str(value))) + .await + } + + /// `DualCache._set_on_batch`: memory now, the `SET` on the batch `send` picks. Memory and Redis + /// both take `ttl`, else `default_in_memory_ttl`, and Redis falls back to its own default. + pub(crate) fn declare_set( + &self, + send: Send<'_>, + key: &str, + value: V, + ttl: Option, + ) -> Result, Error> { + self.declare_encoded(send, key, value, ttl, codec::dumps) + } + + /// `DualCache.async_set_cache_pipeline`: memory, then one pipeline of `SET`s. A value + /// `json.dumps` refuses aborts the whole pipeline, and any failure only logs, as `RedisCache` + /// does. + pub(crate) async fn set_pipeline( + &self, + entries: Vec<(String, V)>, + ttl: Option, + ) -> Result<(), Error> { + let ttl = ttl.or(self.default_in_memory_ttl()); + let encoded = self.redis.get().map(|_| { + entries + .iter() + .map(|(_, value)| codec::dumps(value.clone())) + .collect::, Error>>() + }); + let keys = entries + .into_iter() + .map(|(key, value)| { + self.memory.set(&key, value, ttl)?; + Ok(key) + }) + .collect::, Error>>()?; + let (Some(remote), Some(encoded)) = (self.redis.get(), encoded) else { + return Ok(()); + }; + let failure = match encoded { + Err(error) => Some(error), + Ok(values) => { + let pipeline = RedisBatch::pipelined(remote.client.clone()); + let ttl = ttl.unwrap_or(remote.default_ttl); + let written = keys + .iter() + .zip(values) + .map(|(key, value)| pipeline.add(commands::set(key, value, ttl))) + .collect::>(); + Declared::all(written) + .await + .into_iter() + .find_map(Result::err) + .map(Error::from) + } + }; + if let Some(error) = failure { + redis_failure!( + &error, + "LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS", + ); + } + Ok(()) + } + + async fn set_now( + &self, + key: &str, + value: V, + ttl: Option, + encode: impl FnOnce(V) -> Result, + ) -> Result<(), Error> { + if let Err(error) = self + .declare_encoded(Send::Now, key, value, ttl, encode)? + .await + { + redis_failure!(&error, "LiteLLM Cache: exception in async add_cache"); + } + Ok(()) + } + + fn declare_encoded( + &self, + send: Send<'_>, + key: &str, + value: V, + ttl: Option, + encode: impl FnOnce(V) -> Result, + ) -> Result, Error> { + let ttl = ttl.or(self.default_in_memory_ttl()); + let encoded = self.redis.get().map(|_| encode(value.clone())); + self.memory.set(key, value, ttl)?; + let (Some(remote), Some(encoded)) = (self.redis.get(), encoded) else { + return Ok(ready(())); + }; + let encoded = match encoded { + Ok(encoded) => encoded, + Err(error) => { + redis_failure!(&error, "a cache value did not encode, kept in memory only"); + return Ok(ready(())); + } + }; + if let Some(misses) = send.misses(&remote.client) { + misses.forget(key); + } + Ok(pending(send.batch(&remote.client).add(commands::set( + key, + encoded, + ttl.unwrap_or(remote.default_ttl), + )))) + } +} + +impl DualCache { + /// `DualCache.async_delete_cache_pre_call`: memory now, the `DEL` on the batch `send` picks, + /// after which the request reads the key as missing. + pub(crate) fn declare_delete(&self, send: Send<'_>, key: &str) -> Result, Error> { + self.memory.delete(key)?; + let Some(client) = self.redis.client() else { + return Ok(ready(())); + }; + if let Some(misses) = send.misses(client) { + misses.note([key.to_owned()]); + } + let deleted = send.batch(client).add(commands::del(&[key])); + Ok(Box::pin( + async move { Ok(deleted.resolve().await.map(drop)?) }, + )) + } + + /// `DualCache.async_delete_cache_keys`: memory, then one multi-key `DEL` per chunk, in order. + pub(crate) async fn delete_keys(&self, keys: &[String]) -> Result<(), Error> { + keys.iter().try_for_each(|key| self.memory.delete(key))?; + let Some(client) = self.redis.client() else { + return Ok(()); + }; + for chunk in keys.chunks(DELETE_CHUNK) { + let chunk: Vec<&str> = chunk.iter().map(String::as_str).collect(); + client.send(commands::del(&chunk)).await?; + } + Ok(()) + } +} diff --git a/litellm-rust/crates/redis-store/src/error.rs b/litellm-rust/crates/redis-store/src/error.rs new file mode 100644 index 00000000000..fdc722855bb --- /dev/null +++ b/litellm-rust/crates/redis-store/src/error.rs @@ -0,0 +1,63 @@ +use crate::DecodeError; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error(transparent)] + Redis(#[from] litellm_redis_client::Error), + #[error(transparent)] + Memory(#[from] litellm_cache::Error), + #[error(transparent)] + Decode(#[from] DecodeError), + #[error("the value is not JSON serializable: {0}")] + Encode(#[from] litellm_python_compat::Error), +} + +impl Error { + pub fn kind(&self) -> ErrorKind { + match self { + Self::Redis(error) => error.into(), + Self::Memory(_) => ErrorKind::Memory, + Self::Decode(_) => ErrorKind::Decode, + Self::Encode(_) => ErrorKind::Encode, + } + } +} + +/// The failure class a host reports with a store error, `timeout` and `breaker_open` being the ones a +/// caller may drop. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr)] +pub enum ErrorKind { + #[strum(serialize = "breaker_open")] + BreakerOpen, + #[strum(serialize = "timeout")] + Timeout, + #[strum(serialize = "transport")] + Transport, + #[strum(serialize = "rejected")] + Rejected, + #[strum(serialize = "decode")] + Decode, + #[strum(serialize = "encode")] + Encode, + #[strum(serialize = "memory")] + Memory, +} + +impl From<&Error> for ErrorKind { + fn from(error: &Error) -> Self { + error.kind() + } +} + +impl From<&litellm_redis_client::Error> for ErrorKind { + fn from(error: &litellm_redis_client::Error) -> Self { + use litellm_redis_client::Error; + match error { + Error::BreakerOpen => Self::BreakerOpen, + Error::Transport(error) if error.is_timeout() => Self::Timeout, + Error::Transport(_) => Self::Transport, + Error::Server(_) => Self::Rejected, + Error::Decode(_) => Self::Decode, + } + } +} diff --git a/litellm-rust/crates/redis-store/src/failure_log.rs b/litellm-rust/crates/redis-store/src/failure_log.rs new file mode 100644 index 00000000000..e6bf5187cfd --- /dev/null +++ b/litellm-rust/crates/redis-store/src/failure_log.rs @@ -0,0 +1,18 @@ +/// `log_redis_failure`: an open breaker is DEBUG, anything else WARNING. The event's `redis.failure` +/// field names the class, so the host can rate-limit timeout lines with the throttle redis-py shares. +/// A macro so the event carries its caller's target, file and line, as `stacklevel=2` does in Python. +macro_rules! redis_failure { + ($error:expr, $message:expr $(,)?) => {{ + let error = $error; + let kind = $crate::ErrorKind::from(error); + let failure: &'static str = kind.into(); + match kind { + $crate::ErrorKind::BreakerOpen => { + tracing::debug!(redis.failure = failure, "{}: {}", $message, error) + } + _ => tracing::warn!(redis.failure = failure, "{}: {}", $message, error), + } + }}; +} + +pub(crate) use redis_failure; diff --git a/litellm-rust/crates/redis-store/src/lib.rs b/litellm-rust/crates/redis-store/src/lib.rs new file mode 100644 index 00000000000..b7ee6e3e35e --- /dev/null +++ b/litellm-rust/crates/redis-store/src/lib.rs @@ -0,0 +1,30 @@ +#![expect( + dead_code, + reason = "the store families built on these tiers and codecs land in the next PR" +)] +#![cfg_attr( + not(test), + expect( + unused_imports, + reason = "the store families built on these tiers and codecs land in the next PR" + ) +)] + +mod codec; +mod dual_cache; +mod error; +mod failure_log; +mod memory; +mod number; +mod remote; +mod request; +mod value; + +pub use codec::DecodeError; +pub use dual_cache::TierPolicy; +pub use error::{Error, ErrorKind}; +pub use memory::{Clock, system_clock}; +pub use number::Number; +pub use remote::Remote; +pub use request::{Pending, RequestRedis, Send}; +pub use value::{FromValue, IntoValue}; diff --git a/litellm-rust/crates/redis-store/src/memory.rs b/litellm-rust/crates/redis-store/src/memory.rs new file mode 100644 index 00000000000..f89d5926760 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/memory.rs @@ -0,0 +1,122 @@ +use std::{ + sync::{Arc, Mutex, MutexGuard, PoisonError}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_cache::Error; +use litellm_cache_memory::InMemoryCache; + +use crate::Number; + +pub type Clock = Arc Duration + Send + Sync>; + +pub fn system_clock() -> Clock { + Arc::new(|| { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + }) +} + +/// A store's memory tier: Python's `InMemoryCache`, with its 200 entries and 600 s TTL unless the +/// store names others, plus one lock for every read-modify-write on it. +pub struct Memory { + cache: InMemoryCache, + updates: Mutex<()>, +} + +impl Memory { + pub fn new(entries: Option, now: &Clock) -> Self { + let now = now.clone(); + Self { + cache: InMemoryCache::with_clock(entries, None, move || now()), + updates: Mutex::new(()), + } + } + + pub fn get(&self, key: &str) -> Result, Error> { + self.cache.get_cache(key) + } + + pub fn set(&self, key: &str, value: V, ttl: Option) -> Result<(), Error> { + self.cache.set_cache(key, value, ttl).map(drop) + } + + pub fn delete(&self, key: &str) -> Result<(), Error> { + self.cache.delete_cache(key) + } + + pub fn flush(&self) -> Result<(), Error> { + self.cache.flush_cache() + } + + pub fn len(&self) -> Result { + self.cache.len() + } + + pub fn expires_at(&self, key: &str) -> Result, Error> { + self.cache.expires_at(key) + } + + pub fn entries(&self) -> Result, Error> { + self.cache.entries() + } + + /// `DualCache.update_in_memory_max_size`: `None` restores the 200-entry default. + pub fn update_max_size(&self, entries: Option) { + self.cache.update_max_size_in_memory(entries); + } + + /// Read-modify-write of one entry, atomic against every other update of this tier. + pub fn update( + &self, + key: &str, + ttl: Option, + update: impl FnOnce(Option) -> Result, + ) -> Result { + let _guard = self.guard(); + let value = update(self.get(key)?)?; + self.set(key, value.clone(), ttl)?; + Ok(value) + } + + /// Held across a read and a conditional write that `update` cannot express. + pub fn guard(&self) -> MutexGuard<'_, ()> { + lock(&self.updates) + } +} + +/// A value `InMemoryCache.increment_cache` can add to. +pub(crate) trait Counter: Clone + From { + fn as_number(&self) -> Option; +} + +impl Counter for Number { + fn as_number(&self) -> Option { + Some(*self) + } +} + +impl Memory { + /// `InMemoryCache.increment_cache`: `(current or 0) + amount`, refusing an entry that is not a + /// number and leaving it in place. + pub(crate) fn increment( + &self, + key: &str, + amount: Number, + ttl: Option, + ) -> Result { + let _guard = self.guard(); + let current = match self.get(key)? { + None => Number::Int(0), + Some(value) => value.as_number().ok_or(Error::InvalidEntry(None))?, + }; + let total = current + amount; + self.set(key, V::from(total), ttl)?; + Ok(total) + } +} + +pub(crate) fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(PoisonError::into_inner) +} diff --git a/litellm-rust/crates/redis-store/src/number.rs b/litellm-rust/crates/redis-store/src/number.rs new file mode 100644 index 00000000000..7f3e9f4d10b --- /dev/null +++ b/litellm-rust/crates/redis-store/src/number.rs @@ -0,0 +1,83 @@ +use std::{cmp::Ordering, ops::Add}; + +use litellm_python_compat::{BigInt, Value, json, repr}; + +/// A Python `int` or `float`, kept apart because `repr` and `json.dumps` write them differently. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Number { + Int(i64), + Float(f64), +} + +impl From for litellm_redis_client::Amount { + fn from(number: Number) -> Self { + match number { + Number::Int(value) => Self::Int(value), + Number::Float(value) => Self::Float(value), + } + } +} + +impl Number { + pub fn as_f64(self) -> f64 { + match self { + Self::Int(value) => value as f64, + Self::Float(value) => value, + } + } + + /// Python's `int(value)`, which truncates toward zero. + pub fn truncate(self) -> i64 { + match self { + Self::Int(value) => value, + Self::Float(value) => value.trunc() as i64, + } + } + + pub fn to_value(self) -> Value { + match self { + Self::Int(value) => Value::Int(BigInt::from(value)), + Self::Float(value) => Value::Float(value), + } + } + + /// `json.dumps`, which `RedisCache.async_set_cache` writes. + pub fn to_json(self) -> String { + let value = self.to_value(); + json::dumps(&value).unwrap_or_else(|_| repr::to_str(&value)) + } + + pub fn from_value(value: &Value) -> Option { + match value { + Value::Int(value) => i64::try_from(value).ok().map(Self::Int), + Value::Float(value) => Some(Self::Float(*value)), + _ => None, + } + } + + /// Python's `min` over numbers: the first of the smallest, keeping its type. + pub fn min(values: impl IntoIterator) -> Option { + values + .into_iter() + .fold(None, |smallest, value| match smallest { + Some(smallest) + if value.as_f64().partial_cmp(&smallest.as_f64()) != Some(Ordering::Less) => + { + Some(smallest) + } + _ => Some(value), + }) + } +} + +/// Python's `+`: two ints stay an int, anything with a float is a float. +impl Add for Number { + type Output = Self; + + fn add(self, other: Self) -> Self { + match (self, other) { + (Self::Int(left), Self::Int(right)) => Self::Int(left + right), + (left, right) => Self::Float(left.as_f64() + right.as_f64()), + } + } +} diff --git a/litellm-rust/crates/redis-store/src/remote.rs b/litellm-rust/crates/redis-store/src/remote.rs new file mode 100644 index 00000000000..fc06cae2764 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/remote.rs @@ -0,0 +1,40 @@ +use std::{sync::OnceLock, time::Duration}; + +use litellm_redis_client::{RedisBatch, RedisClient}; + +use crate::Send; + +/// The Redis a store shares with every pod, with the `RedisCache.default_ttl` a write takes when it +/// names none. +#[derive(Clone)] +pub struct Remote { + pub client: RedisClient, + pub default_ttl: Duration, +} + +/// A store's Redis, absent until the host attaches one, as `DualCache.attach_redis_cache` leaves +/// memory in place and only adds the second tier. It is attached at most once. +pub(crate) struct RemoteSlot(OnceLock); + +impl RemoteSlot { + pub(crate) fn new(remote: Option) -> Self { + Self(remote.map(OnceLock::from).unwrap_or_default()) + } + + /// `false` when a Redis was already attached, which keeps it. + pub(crate) fn attach(&self, remote: Remote) -> bool { + self.0.set(remote).is_ok() + } + + pub(crate) fn get(&self) -> Option<&Remote> { + self.0.get() + } + + pub(crate) fn client(&self) -> Option<&RedisClient> { + self.get().map(|remote| &remote.client) + } + + pub(crate) fn batch(&self, send: Send<'_>) -> Option { + Some(send.batch(self.client()?)) + } +} diff --git a/litellm-rust/crates/redis-store/src/request.rs b/litellm-rust/crates/redis-store/src/request.rs new file mode 100644 index 00000000000..79497a5d503 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/request.rs @@ -0,0 +1,219 @@ +use std::{ + collections::HashSet, + future::Future, + sync::{Arc, Mutex}, + time::Duration, +}; + +use futures_util::future::{BoxFuture, join_all}; +use litellm_redis_client::{CallTag, Declared, RedisBatch, RedisClient, current_tag, tagged}; +use tokio::{runtime::Handle, sync::oneshot, task::JoinHandle}; +use tracing::{Dispatch, Instrument, instrument::WithSubscriber}; + +use crate::memory::lock; + +/// One request's Redis batches: a pre-call batch for the reads it needs before answering and a +/// post-call batch for the writes nobody waits on, one of each per Redis client. The host that owns +/// the request builds one and passes it down; it knows nothing else about the request. +#[derive(Clone)] +pub struct RequestRedis(Arc); + +struct Shared { + runtime: Handle, + logs: Dispatch, + post_call_deadline: Duration, + deadline_tag: Option, + batches: Mutex>, + follow_ups: Mutex>>, + deadline: Mutex>>, +} + +#[derive(Clone)] +struct ClientBatches { + client: RedisClient, + pre: RedisBatch, + post: RedisBatch, + misses: Misses, +} + +/// `RedisBatch.read_as_missing`: keys this request's `MGET`s found absent on one client, or that it +/// deleted there, and has not set since, so a later `GET` of one is answered without a round trip. +#[derive(Clone, Default)] +pub(crate) struct Misses(Arc>>); + +impl Misses { + pub(crate) fn contains(&self, key: &str) -> bool { + lock(&self.0).contains(key) + } + + pub(crate) fn note(&self, keys: impl IntoIterator) { + lock(&self.0).extend(keys); + } + + pub(crate) fn forget(&self, key: &str) { + lock(&self.0).remove(key); + } +} + +impl RequestRedis { + /// `runtime` runs the post-call deadline, so declarations may come from any thread, including a + /// Python one that has no Tokio context. Work spawned for the request logs to the subscriber + /// current here, since a declaring thread may have none. + pub fn new(runtime: Handle, post_call_deadline: Duration) -> Self { + Self::with_deadline_tag(runtime, post_call_deadline, None) + } + + /// As `new`, with the deadline's flush of the post-call batch run inside `tagged(deadline_tag)`, + /// the scope an explicit `flush_post` runs in. + pub fn with_deadline_tag( + runtime: Handle, + post_call_deadline: Duration, + deadline_tag: Option, + ) -> Self { + Self(Arc::new(Shared { + runtime, + logs: tracing::dispatcher::get_default(Clone::clone), + post_call_deadline, + deadline_tag, + batches: Mutex::new(Vec::new()), + follow_ups: Mutex::new(Vec::new()), + deadline: Mutex::new(None), + })) + } + + pub fn pre(&self, client: &RedisClient) -> RedisBatch { + self.batches_for(client).pre + } + + pub(crate) fn misses(&self, client: &RedisClient) -> Misses { + self.batches_for(client).misses + } + + /// The post-call batch, which goes out on `flush_post` or `post_call_deadline` after the first + /// declaration, whichever comes first. + pub fn post(&self, client: &RedisClient) -> RedisBatch { + self.arm_deadline(); + self.batches_for(client).post + } + + /// Runs work that follows this request's post-call batch, such as mirroring settled totals + /// into memory, on the request's runtime, from any thread. `flush_post` waits for it, so the + /// shutdown drain does not drop a fallback write still in flight. + pub fn spawn(&self, task: impl Future + std::marker::Send + 'static) { + let follow_up = self + .0 + .runtime + .spawn(self.logged(within(current_tag(), task))); + lock(&self.0.follow_ups).push(follow_up); + } + + fn logged(&self, task: F) -> impl Future + use { + task.in_current_span().with_subscriber(self.0.logs.clone()) + } + + pub async fn flush_pre(&self) { + join_all(self.batches().iter().map(|batches| batches.pre.flush())).await; + } + + /// Sends every post-call op declared so far and waits for the follow-ups spawned on them; the + /// shutdown drain calls it too. + pub async fn flush_post(&self) { + lock(&self.0.deadline).take(); + let follow_ups = std::mem::take(&mut *lock(&self.0.follow_ups)); + join_all(self.batches().iter().map(|batches| batches.post.flush())).await; + join_all(follow_ups) + .await + .into_iter() + .filter_map(Result::err) + .for_each(|error| tracing::warn!(%error, "a post-call follow-up failed")); + } + + fn batches(&self) -> Vec { + lock(&self.0.batches).clone() + } + + fn batches_for(&self, client: &RedisClient) -> ClientBatches { + let mut batches = lock(&self.0.batches); + if let Some(existing) = batches + .iter() + .find(|batches| batches.client.same_client(client)) + { + return existing.clone(); + } + let created = ClientBatches { + client: client.clone(), + pre: client.batch(), + post: client.batch(), + misses: Misses::default(), + }; + batches.push(created.clone()); + created + } + + fn arm_deadline(&self) { + let mut deadline = lock(&self.0.deadline); + if deadline.is_some() { + return; + } + let (disarm, disarmed) = oneshot::channel(); + *deadline = Some(disarm); + let request = self.clone(); + let tag = self.0.deadline_tag.clone(); + self.0.runtime.spawn(self.logged(async move { + tokio::select! { + () = tokio::time::sleep(request.0.post_call_deadline) => within(tag, request.flush_post()).await, + _ = disarmed => {} + } + })); + } +} + +async fn within(tag: Option, task: F) -> F::Output { + match tag { + Some(tag) => tagged(tag, task).await, + None => task.await, + } +} + +/// Which batch an op rides, picked by the caller because only it knows when it needs the answer. +/// The store owns its client, so it calls `batch`, the one place a request and a client meet. +#[derive(Clone, Copy)] +pub enum Send<'a> { + /// On its own, sent at the first `resolve`. + Now, + /// The request's pre-call batch, sent at the first `resolve` or `flush_pre`. + PreCall(&'a RequestRedis), + /// The request's post-call batch, sent at `flush_post` or the post-call deadline. + PostCall(&'a RequestRedis), +} + +impl Send<'_> { + pub fn batch(&self, client: &RedisClient) -> RedisBatch { + match self { + Self::Now => RedisBatch::alone(client.clone()), + Self::PreCall(request) => request.pre(client), + Self::PostCall(request) => request.post(client), + } + } + + /// The request's misses on `client`; `None` for an op outside a request. + pub(crate) fn misses(&self, client: &RedisClient) -> Option { + match self { + Self::Now => None, + Self::PreCall(request) | Self::PostCall(request) => Some(request.misses(client)), + } + } +} + +/// A store op declared on the batch its `Send` picked: building it sends nothing, and awaiting it +/// flushes that batch if nobody has yet. +pub type Pending = BoxFuture<'static, Result>; + +pub(crate) fn pending(declared: Declared) -> Pending { + Box::pin(async move { Ok(declared.resolve().await?) }) +} + +/// An op answered without Redis, from memory. +pub(crate) fn ready(value: T) -> Pending { + Box::pin(std::future::ready(Ok(value))) +} diff --git a/litellm-rust/crates/redis-store/src/value.rs b/litellm-rust/crates/redis-store/src/value.rs new file mode 100644 index 00000000000..0e1b3854172 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/value.rs @@ -0,0 +1,79 @@ +use litellm_python_compat::Value; + +use crate::{DecodeError, Number}; + +/// A store type read from the plain Python data its host sends, so the host never names a field. +pub trait FromValue: Sized { + fn from_value(value: &Value) -> Result; +} + +/// A store type written as the plain Python data its host hands back. +pub trait IntoValue { + fn into_value(self) -> Value; +} + +impl FromValue for Value { + fn from_value(value: &Value) -> Result { + Ok(value.clone()) + } +} + +impl IntoValue for Value { + fn into_value(self) -> Value { + self + } +} + +impl FromValue for Number { + fn from_value(value: &Value) -> Result { + number(value, "a number") + } +} + +impl IntoValue for Number { + fn into_value(self) -> Value { + self.to_value() + } +} + +/// The fields of a Python tuple of exactly `N` items. +pub(crate) fn fields<'a, const N: usize>( + value: &'a Value, + shape: &'static str, +) -> Result<&'a [Value; N], DecodeError> { + match value { + Value::Tuple(items) | Value::List(items) => items + .as_slice() + .try_into() + .map_err(|_| DecodeError::Shape(shape)), + _ => Err(DecodeError::Shape(shape)), + } +} + +pub(crate) fn text(value: &Value, shape: &'static str) -> Result { + match value { + Value::Str(text) => Ok(text.clone()), + _ => Err(DecodeError::Shape(shape)), + } +} + +pub(crate) fn optional_text( + value: &Value, + shape: &'static str, +) -> Result, DecodeError> { + match value { + Value::None => Ok(None), + value => text(value, shape).map(Some), + } +} + +pub(crate) fn integer(value: &Value, shape: &'static str) -> Result { + match value { + Value::Int(value) => i64::try_from(value).map_err(|_| DecodeError::Shape(shape)), + _ => Err(DecodeError::Shape(shape)), + } +} + +pub(crate) fn number(value: &Value, shape: &'static str) -> Result { + Number::from_value(value).ok_or(DecodeError::Shape(shape)) +} diff --git a/litellm-rust/crates/redis-store/tests/request.rs b/litellm-rust/crates/redis-store/tests/request.rs new file mode 100644 index 00000000000..1ee437d4d5a --- /dev/null +++ b/litellm-rust/crates/redis-store/tests/request.rs @@ -0,0 +1,241 @@ +mod support; + +use std::{sync::Arc, time::Duration}; + +use litellm_redis_client::commands; +use litellm_redis_client::{Amount, CallTag, Expiry, Operation, tagged}; +use litellm_redis_store::RequestRedis; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; +use support::{connect, connect_observed, connect_tagged}; +use tokio::{runtime::Handle, sync::oneshot}; + +const TTL: Option = Some(Expiry::KeepExisting(Duration::from_secs(60))); + +// A batch is one pipeline on a standalone Redis and one call per increment on a cluster, like RedisBatch +fn increments(topology: Topology, count: usize) -> Vec { + match topology { + Topology::Standalone => vec![Operation::Pipeline { ops: count }], + Topology::Cluster => vec![Operation::Command("INCRBYFLOAT"); count], + } +} +const LONG_DEADLINE: Duration = Duration::from_secs(60); + +#[rstest] +#[tokio::test] +async fn pre_declarations_on_one_client_share_one_pipeline_and_another_client_gets_its_own( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = connect_observed(&deployment); + let (other, other_trips) = connect_observed(&deployment); + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + + let first = request + .pre(&client) + .add(commands::incr_by_float("a", Amount::Int(1), TTL)); + let second = + request + .pre(&client.clone()) + .add(commands::incr_by_float("b", Amount::Int(2), TTL)); + let third = request + .pre(&other) + .add(commands::incr_by_float("c", Amount::Int(3), TTL)); + request.flush_pre().await; + + assert_eq!(trips.seen(), increments(topology, 2)); + assert_eq!(other_trips.seen(), increments(topology, 1)); + assert_eq!( + [ + first.resolve().await.unwrap(), + second.resolve().await.unwrap(), + third.resolve().await.unwrap(), + ], + [1.0, 2.0, 3.0] + ); + assert_eq!(trips.seen(), increments(topology, 2)); +} + +#[rstest] +#[tokio::test] +async fn a_post_declaration_waits_for_flush_post_and_flush_pre_leaves_it_alone( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let (client, trips) = connect_observed(&deployment); + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + + let total = tokio::spawn( + request + .post(&client) + .add(commands::incr_by_float("spend", Amount::Float(2.5), TTL)) + .settled(), + ); + request.flush_pre().await; + tokio::time::sleep(Duration::from_millis(200)).await; + + assert!(!total.is_finished()); + assert!(trips.seen().is_empty()); + assert_eq!( + raw.send(commands::get::("spend")).await.unwrap(), + None + ); + + request.flush_post().await; + + assert_eq!(total.await.unwrap().unwrap(), 2.5); + assert_eq!(trips.seen(), increments(topology, 1)); +} + +#[rstest] +#[tokio::test] +async fn flush_post_returns_after_the_follow_ups_of_the_ops_it_sent( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let client = connect(&deployment); + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + let total = + request + .post(&client) + .add(commands::incr_by_float("spend", Amount::Float(2.5), TTL)); + let (mirror, mut mirrored) = oneshot::channel(); + request.spawn(async move { + mirror.send(total.settled().await.unwrap()).unwrap(); + }); + + request.flush_post().await; + + assert_eq!(mirrored.try_recv(), Ok(2.5)); +} + +#[rstest] +#[tokio::test] +async fn a_post_batch_nobody_flushes_goes_out_at_the_deadline( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = connect_observed(&deployment); + let request = RequestRedis::new(Handle::current(), Duration::from_millis(100)); + + let first = request + .post(&client) + .add(commands::incr_by_float("a", Amount::Int(1), TTL)); + let second = request + .post(&client) + .add(commands::incr_by_float("b", Amount::Int(2), TTL)); + let settled = tokio::time::timeout(Duration::from_secs(5), async { + (first.settled().await, second.settled().await) + }) + .await + .unwrap(); + + assert_eq!((settled.0.unwrap(), settled.1.unwrap()), (1.0, 2.0)); + assert_eq!(trips.seen(), increments(topology, 2)); +} + +#[rstest] +#[tokio::test] +async fn a_post_declaration_from_a_thread_without_a_runtime_goes_out_at_the_deadline( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = connect_observed(&deployment); + let request = RequestRedis::new(Handle::current(), Duration::from_millis(100)); + + let declared = std::thread::spawn({ + let request = request.clone(); + move || { + request + .post(&client) + .add(commands::incr_by_float("spend", Amount::Float(2.5), TTL)) + } + }) + .join() + .expect("a Python thread has no Tokio context, and declaring there must not panic"); + let total = tokio::time::timeout(Duration::from_secs(5), declared.settled()) + .await + .unwrap(); + + assert_eq!(total.unwrap(), 2.5); + assert_eq!(trips.seen(), increments(topology, 1)); +} + +#[rstest] +#[tokio::test] +async fn a_post_declaration_after_a_flush_arms_a_new_deadline( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, trips) = connect_observed(&deployment); + let request = RequestRedis::new(Handle::current(), Duration::from_millis(300)); + + let early = request + .post(&client) + .add(commands::incr_by_float("a", Amount::Int(1), TTL)); + request.flush_post().await; + tokio::time::sleep(Duration::from_millis(200)).await; + let late = request + .post(&client) + .add(commands::incr_by_float("a", Amount::Int(1), TTL)); + tokio::time::sleep(Duration::from_millis(200)).await; + let flushed_before_its_deadline = trips.seen().len() > 1; + let late = tokio::time::timeout(Duration::from_secs(5), late.settled()) + .await + .unwrap(); + + assert_eq!(early.settled().await.unwrap(), 1.0); + assert!(!flushed_before_its_deadline); + assert_eq!(late.unwrap(), 2.0); +} + +#[rstest] +#[tokio::test] +async fn the_deadline_flush_runs_under_the_deadline_tag( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, tags) = connect_tagged(&deployment); + let request = RequestRedis::with_deadline_tag( + Handle::current(), + Duration::from_millis(100), + Some(Arc::new("post_call_redis_batch") as CallTag), + ); + + let declared = tagged(Arc::new("declaring op"), async { + request + .post(&client) + .add(commands::incr_by_float("a", Amount::Int(1), TTL)) + }) + .await; + tokio::time::timeout(Duration::from_secs(5), declared.settled()) + .await + .unwrap() + .unwrap(); + + assert_eq!(tags.seen(), vec![Some("post_call_redis_batch")]); +} + +#[rstest] +#[tokio::test] +async fn a_follow_up_runs_under_the_tag_it_was_spawned_in( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let (client, tags) = connect_tagged(&deployment); + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + + tagged(Arc::new("async_increment"), async { + request.spawn(async move { + client + .send(commands::incr_by_float("a", Amount::Int(1), TTL)) + .await + .unwrap(); + }); + }) + .await; + request.flush_post().await; + + assert_eq!(tags.seen(), vec![Some("async_increment")]); +} diff --git a/litellm-rust/crates/redis-store/tests/support/mod.rs b/litellm-rust/crates/redis-store/tests/support/mod.rs new file mode 100644 index 00000000000..574f8e4e15c --- /dev/null +++ b/litellm-rust/crates/redis-store/tests/support/mod.rs @@ -0,0 +1,149 @@ +#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_circuit_breaker::{State, Transition}; +use litellm_redis_client::{ + CallEvent, Cluster, Credentials, Namespace, Node, Operation, RedisClient, RedisConfig, + RedisEvents, RedisUrl, Target, Timeouts, current_tag, +}; +use litellm_redis_store::Remote; +use litellm_redis_testing::RedisDeployment; + +pub fn connect(deployment: &RedisDeployment) -> RedisClient { + connect_with(deployment, None, None) +} + +/// `client` with `RedisCache`'s default TTL when `litellm.default_redis_ttl` is unset. +pub fn remote(client: RedisClient) -> Remote { + Remote { + client, + default_ttl: Duration::from_secs(60), + } +} + +/// Every round trip a client made, as the client reports it. +#[derive(Default)] +pub struct RoundTrips(Mutex>); + +impl RoundTrips { + pub fn seen(&self) -> Vec { + self.0.lock().unwrap().clone() + } +} + +impl RedisEvents for RoundTrips { + fn call(&self, event: &CallEvent) { + self.0.lock().unwrap().push(event.operation); + } + + fn breaker(&self, _transition: &Transition) {} +} + +pub fn connect_observed(deployment: &RedisDeployment) -> (RedisClient, Arc) { + let trips = Arc::new(RoundTrips::default()); + (connect_with(deployment, None, Some(trips.clone())), trips) +} + +/// The `&'static str` tag each call ran under, `None` for an untagged call. +#[derive(Default)] +pub struct CallTags(Mutex>>); + +impl CallTags { + pub fn seen(&self) -> Vec> { + self.0.lock().unwrap().clone() + } +} + +impl RedisEvents for CallTags { + fn call(&self, _event: &CallEvent) { + let tag = current_tag().and_then(|tag| tag.downcast_ref::<&'static str>().copied()); + self.0.lock().unwrap().push(tag); + } + + fn breaker(&self, _transition: &Transition) {} +} + +pub fn connect_tagged(deployment: &RedisDeployment) -> (RedisClient, Arc) { + let tags = Arc::new(CallTags::default()); + (connect_with(deployment, None, Some(tags.clone())), tags) +} + +pub fn connect_guarded( + deployment: &RedisDeployment, + breaker: litellm_circuit_breaker::Breaker, +) -> RedisClient { + connect_with(deployment, Some(breaker), None) +} + +fn connect_with( + deployment: &RedisDeployment, + breaker: Option, + events: Option>, +) -> RedisClient { + let endpoint = match deployment { + RedisDeployment::Standalone(server) => Target::Url(RedisUrl { + url: litellm_secret_value::SecretValue::new(server.url().to_owned()), + tls: None, + db: None, + username: None, + password: None, + }), + RedisDeployment::Cluster(cluster) => Target::Cluster(Cluster { + nodes: cluster + .nodes() + .map(|(host, port)| Node { + host: host.to_owned(), + port, + }) + .collect(), + tls: None, + }), + }; + RedisClient::new(RedisConfig { + target: endpoint, + credentials: Credentials::NONE, + namespace: Namespace::default(), + timeouts: Timeouts::new(Duration::from_millis(500), None), + breaker, + events, + health_check: None, + keepalive: true, + }) + .unwrap() +} + +/// A TTL Redis reports for a key written `whole_seconds` ago at most a few seconds back. +pub fn within(ttl: Duration, whole_seconds: u64) -> bool { + ttl <= Duration::from_secs(whole_seconds) + && ttl > Duration::from_secs(whole_seconds.saturating_sub(5)) +} + +/// How many times each command reached the servers since the last `CONFIG RESETSTAT`, so a test +/// can pin the commands a store sends against the ones main's Python sends. +pub async fn calls(deployment: &RedisDeployment, command: &str) -> u64 { + let stats = match deployment { + RedisDeployment::Standalone(server) => server.cli(&["INFO", "commandstats"]).await.unwrap(), + RedisDeployment::Cluster(cluster) => { + let mut stats = String::new(); + for (_, port) in cluster.nodes() { + stats.push_str( + &cluster + .cli_on_node(port, &["INFO", "commandstats"]) + .await + .unwrap(), + ); + } + stats + } + }; + let prefix = format!("cmdstat_{command}:calls="); + stats + .lines() + .filter_map(|line| line.trim().strip_prefix(&prefix)) + .filter_map(|rest| rest.split(',').next()?.parse::().ok()) + .sum() +} diff --git a/litellm-rust/crates/redis-testing/AGENTS.md b/litellm-rust/crates/redis-testing/AGENTS.md index fb7a4cd737a..8988be9b39b 100644 --- a/litellm-rust/crates/redis-testing/AGENTS.md +++ b/litellm-rust/crates/redis-testing/AGENTS.md @@ -4,3 +4,10 @@ - `RedisCluster` runs 3 masters in one container. Each node listens on the same port inside and outside and announces 127.0.0.1, so MOVED redirects work from the test process - Outages come from the container (`restart`, `freeze` / `thaw`), never `DEBUG`, which the official image disables - No env gate that makes a test pass without Redis. A test that can't start its container fails +- `RedisTlsServer` serves TLS with certificates `rcgen` generates per start, for `localhost`, optionally requiring a client certificate +- `RedisSentinel` runs a primary, a replica and one sentinel in one container, all announcing 127.0.0.1. `failover()` returns once the replica is master and the old master is a replica. `start_with_password` puts `requirepass` on the primary and the replica, which the sentinel monitors with, and the CLI helpers authenticate with it +- `TcpProxy` forwards one port from the test process. `blackhole_open_connections` stops forwarding on the connections open now while keeping their sockets open, which a container outage cannot do to one socket. `RedisCluster::announce_ports` makes the nodes announce proxy ports so a cluster client reaches them only through proxies +- `cli` / `cli_on_every_node` run `redis-cli -e` inside the container for setup such as ACL users, so tests don't need a second client. `-e` makes an error reply fail the call, and a test whose setup changes server state also asserts that state before relying on it +- Not `redis-test`'s server helpers, which need local `redis-server` and `openssl` binaries +- `RedisDeployment` runs either `Topology`, so a test can take `#[values(Topology::Standalone, Topology::Cluster)]`. `clients()` returns every client connection on every node from `CLIENT LIST`, leaving out the CLI's own +- Server fixtures only. This crate never depends on a Redis client crate, so the client crates can dev-depend on it without a cycle diff --git a/litellm-rust/crates/redis-testing/Cargo.toml b/litellm-rust/crates/redis-testing/Cargo.toml index c9bad227acc..a184a804e21 100644 --- a/litellm-rust/crates/redis-testing/Cargo.toml +++ b/litellm-rust/crates/redis-testing/Cargo.toml @@ -6,6 +6,9 @@ license.workspace = true repository.workspace = true [dependencies] +pkcs8 = { version = "0.10.2", features = ["encryption", "pem", "std"] } +rcgen = "0.14.10" +tempfile.workspace = true testcontainers-modules = { workspace = true, features = ["redis"] } thiserror.workspace = true -tokio = { workspace = true, features = ["io-util"] } +tokio = { workspace = true, features = ["io-util", "sync", "time"] } diff --git a/litellm-rust/crates/redis-testing/src/cluster.rs b/litellm-rust/crates/redis-testing/src/cluster.rs index e61f7832b1c..130495cefad 100644 --- a/litellm-rust/crates/redis-testing/src/cluster.rs +++ b/litellm-rust/crates/redis-testing/src/cluster.rs @@ -1,3 +1,5 @@ +use std::time::{Duration, Instant}; + use testcontainers_modules::testcontainers::{ ContainerAsync, GenericImage, ImageExt, core::{IntoContainerPort, WaitFor}, @@ -5,12 +7,14 @@ use testcontainers_modules::testcontainers::{ }; use crate::{ - Error, + ClientConnection, Error, + deployment::parse_client_list, ports::{HOST, start_on_free_ports, wait_for_pong}, standalone::REDIS_TAG, }; const READY: &str = "redis-testing cluster ready"; +const ANNOUNCE_TIMEOUT: Duration = Duration::from_secs(30); const START_CLUSTER: &str = r#" set -e @@ -70,6 +74,69 @@ impl RedisCluster { self.ports.iter().map(|&port| (HOST, port)) } + /// One redis-cli reply per node, in `nodes()` order. + pub async fn cli_outputs(&self, args: &[&str]) -> Result, Error> { + let mut replies = Vec::with_capacity(self.ports.len()); + for port in self.ports { + replies.push(crate::exec::run(&self.container, crate::exec::cli(port, args)).await?); + } + Ok(replies) + } + + pub async fn cli_on_every_node(&self, args: &[&str]) -> Result<(), Error> { + for port in self.ports { + crate::exec::run(&self.container, crate::exec::cli(port, args)).await?; + } + Ok(()) + } + + pub async fn cli_on_node(&self, port: u16, args: &[&str]) -> Result { + crate::exec::run(&self.container, crate::exec::cli(port, args)).await + } + + /// Makes each node tell clients it is at the port of `ports` in `nodes()` order, such as a + /// `TcpProxy` in front of it, and returns once every node lists every announced port, so slot + /// maps and redirects only lead a client to those ports. + pub async fn announce_ports(&self, ports: [u16; 3]) -> Result<(), Error> { + for (node, announced) in self.ports.into_iter().zip(ports) { + self.cli_on_node( + node, + &[ + "CONFIG", + "SET", + "cluster-announce-port", + &announced.to_string(), + ], + ) + .await?; + } + let deadline = Instant::now() + ANNOUNCE_TIMEOUT; + loop { + let views = self.cli_outputs(&["CLUSTER", "NODES"]).await?; + let learned = views + .iter() + .all(|view| ports.iter().all(|port| view.contains(&format!(":{port}@")))); + if learned { + return Ok(()); + } + if Instant::now() > deadline { + return Err(Error::Announce); + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + + pub async fn clients(&self) -> Result, Error> { + let mut clients = Vec::new(); + for port in self.ports { + let reply = + crate::exec::run(&self.container, crate::exec::cli(port, &["CLIENT", "LIST"])) + .await?; + clients.extend(parse_client_list(port, &reply)); + } + Ok(clients) + } + pub async fn restart(&self) -> Result<(), Error> { self.stop().await?; self.resume().await diff --git a/litellm-rust/crates/redis-testing/src/deployment.rs b/litellm-rust/crates/redis-testing/src/deployment.rs new file mode 100644 index 00000000000..5d090eae4ca --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/deployment.rs @@ -0,0 +1,113 @@ +use std::time::{Duration, Instant}; + +use crate::{Error, RedisCluster, RedisServer}; + +const EVENTUALLY_TIMEOUT: Duration = Duration::from_secs(30); + +#[derive(Clone, Copy, Debug)] +pub enum Topology { + Standalone, + Cluster, +} + +pub enum RedisDeployment { + Standalone(RedisServer), + Cluster(RedisCluster), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ClientConnection { + pub node: u16, + pub id: u64, + pub user: String, +} + +impl RedisDeployment { + pub async fn start(topology: Topology) -> Result { + Ok(match topology { + Topology::Standalone => Self::Standalone(RedisServer::start().await?), + Topology::Cluster => Self::Cluster(RedisCluster::start().await?), + }) + } + + pub async fn cli_on_every_node(&self, args: &[&str]) -> Result<(), Error> { + match self { + Self::Standalone(server) => server.cli(args).await.map(drop), + Self::Cluster(cluster) => cluster.cli_on_every_node(args).await, + } + } + + pub async fn clients(&self) -> Result, Error> { + match self { + Self::Standalone(server) => server.clients().await, + Self::Cluster(cluster) => cluster.clients().await, + } + } + + pub async fn restart(&self) -> Result<(), Error> { + match self { + Self::Standalone(server) => server.restart().await, + Self::Cluster(cluster) => cluster.restart().await, + } + } + + pub async fn stop(&self) -> Result<(), Error> { + match self { + Self::Standalone(server) => server.stop().await, + Self::Cluster(cluster) => cluster.stop().await, + } + } + + pub async fn resume(&self) -> Result<(), Error> { + match self { + Self::Standalone(server) => server.resume().await, + Self::Cluster(cluster) => cluster.resume().await, + } + } + + pub async fn freeze(&self) -> Result<(), Error> { + match self { + Self::Standalone(server) => server.freeze().await, + Self::Cluster(cluster) => cluster.freeze().await, + } + } + + pub async fn thaw(&self) -> Result<(), Error> { + match self { + Self::Standalone(server) => server.thaw().await, + Self::Cluster(cluster) => cluster.thaw().await, + } + } +} + +pub(crate) fn parse_client_list(node: u16, reply: &str) -> Vec { + reply + .lines() + .filter(|line| !line.split(' ').any(|field| field == "cmd=client|list")) + .filter_map(|line| { + let field = |name: &str| { + line.split(' ') + .find_map(|field| field.strip_prefix(name)?.strip_prefix('=')) + }; + Some(ClientConnection { + node, + id: field("id")?.parse().ok()?, + user: field("user")?.to_owned(), + }) + }) + .collect() +} + +pub async fn eventually(mut attempt: impl AsyncFnMut() -> Option) -> T { + let deadline = Instant::now() + EVENTUALLY_TIMEOUT; + loop { + if let Some(value) = attempt().await { + return value; + } + assert!( + Instant::now() < deadline, + "condition not met within {EVENTUALLY_TIMEOUT:?}" + ); + tokio::time::sleep(Duration::from_millis(100)).await; + } +} diff --git a/litellm-rust/crates/redis-testing/src/error.rs b/litellm-rust/crates/redis-testing/src/error.rs index f21cc959812..77d9ef6f8ea 100644 --- a/litellm-rust/crates/redis-testing/src/error.rs +++ b/litellm-rust/crates/redis-testing/src/error.rs @@ -8,4 +8,16 @@ pub enum Error { Unreachable { port: u16 }, #[error("controlling the Redis container")] Container(#[from] testcontainers_modules::testcontainers::TestcontainersError), + #[error("{command} exited with {code}")] + Command { command: String, code: i64 }, + #[error("generating test TLS certificates")] + Certificate(#[from] rcgen::Error), + #[error("encrypting the test client key")] + KeyEncryption(#[from] pkcs8::Error), + #[error("writing test TLS files")] + Files(#[source] io::Error), + #[error("Sentinel did not promote the replica")] + Failover, + #[error("the cluster nodes did not learn the announced ports")] + Announce, } diff --git a/litellm-rust/crates/redis-testing/src/exec.rs b/litellm-rust/crates/redis-testing/src/exec.rs new file mode 100644 index 00000000000..74a94c085ec --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/exec.rs @@ -0,0 +1,31 @@ +use testcontainers_modules::testcontainers::{ + ContainerAsync, Image, + core::{CmdWaitFor, ExecCommand}, +}; + +use crate::Error; + +pub(crate) async fn run( + container: &ContainerAsync, + command: Vec, +) -> Result { + let mut result = container + .exec(ExecCommand::new(command.clone()).with_cmd_ready_condition(CmdWaitFor::exit())) + .await?; + let stdout = String::from_utf8_lossy(&result.stdout_to_vec().await?).into_owned(); + match result.exit_code().await? { + Some(0) | None => Ok(stdout), + Some(code) => Err(Error::Command { + command: command.join(" "), + code, + }), + } +} + +pub(crate) fn cli(port: u16, args: &[&str]) -> Vec { + ["redis-cli", "-e", "-p", &port.to_string()] + .into_iter() + .chain(args.iter().copied()) + .map(str::to_owned) + .collect() +} diff --git a/litellm-rust/crates/redis-testing/src/lib.rs b/litellm-rust/crates/redis-testing/src/lib.rs index 6c2afd7f275..e8c892a2e7d 100644 --- a/litellm-rust/crates/redis-testing/src/lib.rs +++ b/litellm-rust/crates/redis-testing/src/lib.rs @@ -1,8 +1,17 @@ mod cluster; +mod deployment; mod error; +mod exec; mod ports; +mod proxy; +mod sentinel; mod standalone; +mod tls; pub use cluster::RedisCluster; +pub use deployment::{ClientConnection, RedisDeployment, Topology, eventually}; pub use error::Error; +pub use proxy::TcpProxy; +pub use sentinel::{RedisSentinel, SERVICE_NAME}; pub use standalone::RedisServer; +pub use tls::{CLIENT_KEY_PASSWORD, RedisTlsServer, SERVER_NAME, TlsFiles}; diff --git a/litellm-rust/crates/redis-testing/src/ports.rs b/litellm-rust/crates/redis-testing/src/ports.rs index e7b94988c10..eb0d5bd4e2e 100644 --- a/litellm-rust/crates/redis-testing/src/ports.rs +++ b/litellm-rust/crates/redis-testing/src/ports.rs @@ -42,6 +42,17 @@ pub(crate) async fn wait_for_pong(port: u16) -> Result<(), Error> { Err(Error::Unreachable { port }) } +pub(crate) async fn wait_for_listener(port: u16) -> Result<(), Error> { + let deadline = Instant::now() + READY_TIMEOUT; + while Instant::now() < deadline { + if TcpStream::connect((HOST, port)).await.is_ok() { + return Ok(()); + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + Err(Error::Unreachable { port }) +} + async fn ping(port: u16) -> io::Result { let mut stream = TcpStream::connect((HOST, port)).await?; stream.write_all(b"PING\r\n").await?; diff --git a/litellm-rust/crates/redis-testing/src/proxy.rs b/litellm-rust/crates/redis-testing/src/proxy.rs new file mode 100644 index 00000000000..91d8d9f5a08 --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/proxy.rs @@ -0,0 +1,65 @@ +use std::sync::Arc; + +use tokio::{ + net::{TcpListener, TcpStream}, + sync::watch, + task::JoinHandle, +}; + +use crate::{Error, ports::HOST}; + +/// A TCP proxy in front of one Redis port, for outages a container cannot stage on one socket. +pub struct TcpProxy { + port: u16, + cuts: Arc>, + accepting: JoinHandle<()>, +} + +impl TcpProxy { + pub async fn start(upstream: u16) -> Result { + let listener = TcpListener::bind((HOST, 0)).await.map_err(Error::Port)?; + let port = listener.local_addr().map_err(Error::Port)?.port(); + let cuts = Arc::new(watch::Sender::new(0)); + let accepting = tokio::spawn(accept(listener, upstream, cuts.clone())); + Ok(Self { + port, + cuts, + accepting, + }) + } + + pub fn port(&self) -> u16 { + self.port + } + + /// Stops forwarding on every connection open now while keeping both its sockets open, the way + /// a half-open TCP connection looks to the client: writes succeed and no reply ever comes. + /// Connections opened afterwards forward normally. + pub fn blackhole_open_connections(&self) { + self.cuts.send_modify(|cut| *cut += 1); + } +} + +impl Drop for TcpProxy { + fn drop(&mut self) { + self.accepting.abort(); + } +} + +async fn accept(listener: TcpListener, upstream: u16, cuts: Arc>) { + while let Ok((client, _)) = listener.accept().await { + tokio::spawn(forward(client, upstream, cuts.subscribe())); + } +} + +async fn forward(mut client: TcpStream, upstream: u16, mut cut: watch::Receiver) { + let Ok(mut server) = TcpStream::connect((HOST, upstream)).await else { + return; + }; + tokio::select! { + _ = tokio::io::copy_bidirectional(&mut client, &mut server) => {} + _ = cut.changed() => { + while cut.changed().await.is_ok() {} + } + } +} diff --git a/litellm-rust/crates/redis-testing/src/sentinel.rs b/litellm-rust/crates/redis-testing/src/sentinel.rs new file mode 100644 index 00000000000..cd93c72e0f6 --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/sentinel.rs @@ -0,0 +1,170 @@ +use std::time::{Duration, Instant}; + +use testcontainers_modules::testcontainers::{ + ContainerAsync, GenericImage, ImageExt, + core::{IntoContainerPort, WaitFor}, + runners::AsyncRunner, +}; + +use crate::{ + ClientConnection, Error, + deployment::parse_client_list, + exec::{cli, run}, + ports::{HOST, start_on_free_ports, wait_for_listener, wait_for_pong}, + standalone::REDIS_TAG, +}; + +pub const SERVICE_NAME: &str = "primary"; +const READY: &str = "redis-testing sentinel ready"; +const FAILOVER_TIMEOUT: Duration = Duration::from_secs(30); + +const START_SENTINEL: &str = r#" +set -e +primary=$1 replica=$2 sentinel=$3 password=$4 +common="--save '' --appendonly no --daemonize yes --replica-announce-ip 127.0.0.1" +if [ -n "$password" ]; then + common="$common --requirepass $password --masterauth $password" +fi +eval redis-server --port "$primary" $common +eval redis-server --port "$replica" --replicaof 127.0.0.1 "$primary" $common +until redis-cli ${password:+-a "$password" --no-auth-warning} -p "$replica" info replication \ + | grep -q master_link_status:up; do sleep 0.1; done +cat > /data/sentinel.conf <> /data/sentinel.conf +fi +redis-server /data/sentinel.conf --sentinel --daemonize yes +until redis-cli -p "$sentinel" sentinel replicas primary | grep -q "$replica"; do sleep 0.1; done +echo "redis-testing sentinel ready" +exec sleep infinity +"#; + +pub struct RedisSentinel { + container: ContainerAsync, + ports: [u16; 3], + password: Option, +} + +impl RedisSentinel { + pub async fn start() -> Result { + start_on_free_ports(|ports| Self::start_on(ports, None)).await + } + + /// The primary and the replica require `password` (`requirepass`, with `masterauth` between + /// them), and the sentinel monitors them with it. The sentinel itself needs none. + pub async fn start_with_password(password: &str) -> Result { + start_on_free_ports(|ports| Self::start_on(ports, Some(password))).await + } + + async fn start_on(ports: [u16; 3], password: Option<&str>) -> Result { + let image = GenericImage::new("redis", REDIS_TAG) + .with_wait_for(WaitFor::message_on_stdout(READY)) + .with_entrypoint("sh"); + let request = ports.iter().fold( + image.with_cmd( + ["-c", START_SENTINEL, "sh"] + .into_iter() + .map(str::to_owned) + .chain(ports.iter().map(u16::to_string)) + .chain([password.unwrap_or_default().to_owned()]), + ), + |request, &port| request.with_mapped_port(port, port.tcp()), + ); + let container = request.start().await?; + wait_for_listener(ports[0]).await?; + wait_for_listener(ports[1]).await?; + wait_for_pong(ports[2]).await?; + Ok(Self { + container, + ports, + password: password.map(str::to_owned), + }) + } + + fn data_cli(&self, port: u16, args: &[&str]) -> Vec { + let auth = self + .password + .as_deref() + .map(|password| ["-a", password, "--no-auth-warning"]); + let args: Vec<&str> = auth + .into_iter() + .flatten() + .chain(args.iter().copied()) + .collect(); + cli(port, &args) + } + + pub fn sentinel(&self) -> (&'static str, u16) { + (HOST, self.ports[2]) + } + + /// Runs one redis-cli command on the primary and the replica, which keep their own ACLs. + pub async fn cli_on_data_nodes(&self, args: &[&str]) -> Result<(), Error> { + for port in [self.ports[0], self.ports[1]] { + run(&self.container, self.data_cli(port, args)).await?; + } + Ok(()) + } + + pub async fn clients_on(&self, port: u16) -> Result, Error> { + Ok(parse_client_list( + port, + &run(&self.container, self.data_cli(port, &["CLIENT", "LIST"])).await?, + )) + } + + pub fn primary_port(&self) -> u16 { + self.ports[0] + } + + pub fn replica_port(&self) -> u16 { + self.ports[1] + } + + pub async fn master_port(&self) -> Result { + let reply = run( + &self.container, + cli( + self.ports[2], + &["sentinel", "get-master-addr-by-name", SERVICE_NAME], + ), + ) + .await?; + Ok(reply + .lines() + .nth(1) + .and_then(|port| port.trim().parse().ok()) + .unwrap_or_default()) + } + + pub async fn failover(&self) -> Result { + let before = self.master_port().await?; + run( + &self.container, + cli(self.ports[2], &["sentinel", "failover", SERVICE_NAME]), + ) + .await?; + let deadline = Instant::now() + FAILOVER_TIMEOUT; + loop { + let now = self.master_port().await?; + let promoted = now != before && self.role(now).await?.starts_with("master"); + if promoted && self.role(before).await?.starts_with("slave") { + return Ok(now); + } + if Instant::now() > deadline { + return Err(Error::Failover); + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + + async fn role(&self, port: u16) -> Result { + run(&self.container, self.data_cli(port, &["role"])).await + } +} diff --git a/litellm-rust/crates/redis-testing/src/standalone.rs b/litellm-rust/crates/redis-testing/src/standalone.rs index 8060ccbf190..6150aadff07 100644 --- a/litellm-rust/crates/redis-testing/src/standalone.rs +++ b/litellm-rust/crates/redis-testing/src/standalone.rs @@ -5,7 +5,8 @@ use testcontainers_modules::{ }; use crate::{ - Error, + ClientConnection, Error, + deployment::parse_client_list, ports::{HOST, start_on_free_ports, wait_for_pong}, }; @@ -14,6 +15,7 @@ pub(crate) const REDIS_TAG: &str = pub struct RedisServer { container: ContainerAsync, + port: u16, url: String, } @@ -31,6 +33,7 @@ impl RedisServer { wait_for_pong(host_port).await?; Ok(Self { container, + port: host_port, url: format!("redis://{HOST}:{host_port}"), }) } @@ -39,6 +42,21 @@ impl RedisServer { &self.url } + pub fn port(&self) -> u16 { + self.port + } + + pub async fn cli(&self, args: &[&str]) -> Result { + crate::exec::run(&self.container, crate::exec::cli(REDIS_PORT, args)).await + } + + pub async fn clients(&self) -> Result, Error> { + Ok(parse_client_list( + REDIS_PORT, + &self.cli(&["CLIENT", "LIST"]).await?, + )) + } + pub async fn restart(&self) -> Result<(), Error> { self.stop().await?; self.resume().await diff --git a/litellm-rust/crates/redis-testing/src/tls.rs b/litellm-rust/crates/redis-testing/src/tls.rs new file mode 100644 index 00000000000..d7f517175c9 --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/tls.rs @@ -0,0 +1,148 @@ +use std::path::PathBuf; + +use pkcs8::{LineEnding, PrivateKeyInfo, pkcs5::pbes2}; +use rcgen::{BasicConstraints, CertificateParams, CertifiedIssuer, DnType, IsCa, KeyPair}; +use tempfile::TempDir; +use testcontainers_modules::testcontainers::{ + ContainerAsync, GenericImage, ImageExt, + core::{IntoContainerPort, WaitFor}, + runners::AsyncRunner, +}; + +use crate::{ + Error, + ports::{HOST, start_on_free_ports, wait_for_listener}, + standalone::REDIS_TAG, +}; + +pub const SERVER_NAME: &str = "localhost"; +pub const CLIENT_KEY_PASSWORD: &str = "redis-testing-key-password"; + +pub struct TlsFiles { + _directory: TempDir, + pub ca: PathBuf, + pub client_cert: PathBuf, + pub client_key: PathBuf, + /// `client_key` as a PKCS#8 `ENCRYPTED PRIVATE KEY` under `CLIENT_KEY_PASSWORD` + pub encrypted_client_key: PathBuf, +} + +struct Generated { + files: TlsFiles, + ca_pem: String, + server_cert: String, + server_key: String, +} + +pub struct RedisTlsServer { + _container: ContainerAsync, + port: u16, + files: TlsFiles, +} + +impl RedisTlsServer { + pub async fn start(require_client_certificate: bool) -> Result { + start_on_free_ports(|[port]| Self::start_on(port, require_client_certificate)).await + } + + async fn start_on(port: u16, require_client_certificate: bool) -> Result { + let generated = generate()?; + let auth_clients = match require_client_certificate { + true => "yes", + false => "no", + }; + let container = GenericImage::new("redis", REDIS_TAG) + .with_wait_for(WaitFor::message_on_stdout("Ready to accept connections")) + .with_copy_to("/tls/ca.crt", generated.ca_pem.into_bytes()) + .with_copy_to("/tls/server.crt", generated.server_cert.into_bytes()) + .with_copy_to("/tls/server.key", generated.server_key.into_bytes()) + .with_cmd([ + "redis-server", + "--port", + "0", + "--tls-port", + &port.to_string(), + "--tls-cert-file", + "/tls/server.crt", + "--tls-key-file", + "/tls/server.key", + "--tls-ca-cert-file", + "/tls/ca.crt", + "--tls-auth-clients", + auth_clients, + ]) + .with_mapped_port(port, port.tcp()) + .start() + .await?; + wait_for_listener(port).await?; + Ok(Self { + _container: container, + port, + files: generated.files, + }) + } + + pub fn host(&self) -> &'static str { + HOST + } + + pub fn port(&self) -> u16 { + self.port + } + + pub fn files(&self) -> &TlsFiles { + &self.files + } +} + +fn generate() -> Result { + let named = |names: Vec, common_name: &str| -> Result { + let mut params = CertificateParams::new(names)?; + params + .distinguished_name + .push(DnType::CommonName, common_name); + Ok(params) + }; + let mut ca_params = named(Vec::new(), "redis-testing CA")?; + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + let ca = CertifiedIssuer::self_signed(ca_params, KeyPair::generate()?)?; + let server_key = KeyPair::generate()?; + let server = + named(vec![SERVER_NAME.to_owned()], "redis-testing server")?.signed_by(&server_key, &ca)?; + let client_key = KeyPair::generate()?; + let client = + named(vec!["client".to_owned()], "redis-testing client")?.signed_by(&client_key, &ca)?; + let directory = tempfile::tempdir().map_err(Error::Files)?; + let write = |name: &str, contents: String| -> Result { + let path = directory.path().join(name); + std::fs::write(&path, contents).map_err(Error::Files)?; + Ok(path) + }; + let ca_path = write("ca.pem", ca.pem())?; + let client_cert = write("client.pem", client.pem())?; + let client_key_path = write("client.key", client_key.serialize_pem())?; + let encrypted_client_key = write("client-encrypted.key", encrypt(&client_key)?)?; + Ok(Generated { + ca_pem: ca.pem(), + server_cert: server.pem(), + server_key: server_key.serialize_pem(), + files: TlsFiles { + _directory: directory, + ca: ca_path, + client_cert, + client_key: client_key_path, + encrypted_client_key, + }, + }) +} + +// The parameters `openssl pkcs8 -topk8` uses by default: PBES2 with PBKDF2-HMAC-SHA256 and AES-256-CBC +fn encrypt(key: &KeyPair) -> Result { + let der = key.serialize_der(); + let params = pbes2::Parameters::pbkdf2_sha256_aes256cbc(2048, b"redis-testing-salt", &[7; 16])?; + let encrypted = PrivateKeyInfo::try_from(der.as_slice())? + .encrypt_with_params(params, CLIENT_KEY_PASSWORD)?; + Ok(encrypted + .to_pem("ENCRYPTED PRIVATE KEY", LineEnding::LF)? + .to_string()) +} diff --git a/litellm-rust/crates/redis-types/AGENTS.md b/litellm-rust/crates/redis-types/AGENTS.md new file mode 100644 index 00000000000..bfce4e28e63 --- /dev/null +++ b/litellm-rust/crates/redis-types/AGENTS.md @@ -0,0 +1,2 @@ +- `Resolution` and the target, node, TLS and auth types inside it: what `redis-resolve` interprets merged kwargs into, `redis-connect` turns into a client and `redis-client` connects with, so none of them converts another's copy. +- Data only. It depends on `secret-value`, `serde_json` and `veil` and nothing else, so `redis-resolve` never pulls in redis-rs or tokio. Anything that needs redis-rs, like building a `ConnectionInfo` or reading TLS files, lives in `redis-client` diff --git a/litellm-rust/crates/redis-types/Cargo.toml b/litellm-rust/crates/redis-types/Cargo.toml new file mode 100644 index 00000000000..229325a27bd --- /dev/null +++ b/litellm-rust/crates/redis-types/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "litellm-redis-types" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-secret-value.workspace = true diff --git a/litellm-rust/crates/redis-types/src/lib.rs b/litellm-rust/crates/redis-types/src/lib.rs new file mode 100644 index 00000000000..0ea69edfa6a --- /dev/null +++ b/litellm-rust/crates/redis-types/src/lib.rs @@ -0,0 +1,175 @@ +use std::{path::PathBuf, time::Duration}; + +use litellm_secret_value::SecretValue; + +/// One Redis connection with every setting, environment value and secret already applied. Each +/// client library builds its client from this one value, so they cannot read a setting differently. +#[derive(Clone, Debug, PartialEq)] +pub struct Resolution { + pub target: Target, + pub auth: Auth, + pub namespace: Option, + pub timeout: Option, + pub connect_timeout: Option, + pub health_check: HealthCheck, + /// redis-py's `socket_keepalive`, on unless set false, so an idle connection a load balancer + /// dropped is noticed. + pub keepalive: bool, + pub ignored: Vec, +} + +/// redis-py's `health_check_interval`, which `_coerce_redis_kwargs_types` reads as an `int` of +/// seconds. Unset and zero differ, because a client may give an unset interval its own default. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum HealthCheck { + Unset, + Disabled, + Every(Duration), +} + +/// A setting the Rust client accepts without honoring, so the host can log what ignoring it means. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Ignored { + /// The key, the environment variable or the url query parameter that set it. + pub setting: String, + pub consequence: &'static str, +} + +#[derive(Clone, Debug, PartialEq)] +pub enum Target { + Standalone(Standalone), + /// Like redis-py's `UnixDomainSocketConnection`, which replaces host, port and TLS. + Unix(UnixSocket), + Url(RedisUrl), + Cluster(Cluster), + Sentinel(Sentinel), +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Standalone { + pub host: String, + pub port: u16, + pub tls: Option, + pub db: i64, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct UnixSocket { + pub path: PathBuf, + pub db: i64, +} + +/// A url kept verbatim, plus what redis-py's `parse_url` reads from its query that the url's own +/// grammar does not carry. +#[derive(Clone, Debug, PartialEq)] +pub struct RedisUrl { + pub url: SecretValue, + /// Set only for a `rediss://` url, from its query parameters. + pub tls: Option, + /// `?db=`, which wins over the path. + pub db: Option, + /// `?username=` and `?password=`, which the url's userinfo wins over. + pub username: Option, + pub password: Option, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Cluster { + pub nodes: Vec, + pub tls: Option, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Node { + pub host: String, + pub port: u16, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Sentinel { + pub monitors: Vec, + pub service_name: String, + /// How to authenticate to the monitors, which are separate servers from the data node. + pub monitor_auth: Option, + pub tls: Option, + pub db: i64, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct TlsOptions { + pub ca_certs: Option, + /// PEM text trusted alongside `ca_certs`, redis-py's `ssl_ca_data` + pub ca_data: Option, + /// A directory whose files are PEM certificates to trust, redis-py's `ssl_ca_path` + pub ca_path: Option, + pub certfile: Option, + pub keyfile: Option, + /// Unlocks an encrypted `keyfile`, redis-py's `ssl_password` + pub key_password: Option, + pub verify_certificate: bool, + pub check_hostname: bool, +} + +impl TlsOptions { + pub fn verified() -> Self { + Self { + ca_certs: None, + ca_data: None, + ca_path: None, + certfile: None, + keyfile: None, + key_password: None, + verify_certificate: true, + check_hostname: false, + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub enum Auth { + Static(StaticAuth), + /// Credentials come from a host object, like a redis-py `credential_provider`. + Provider(ProviderHandle), + GcpIam { + service_account: String, + }, + AzureAd(AzureAdAuth), + AwsIam(AwsIamAuth), +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct StaticAuth { + pub username: Option, + pub password: Option, +} + +/// Names one credential provider the host holds. Two resolutions with the same handle use the +/// same provider. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct ProviderHandle(u64); + +impl ProviderHandle { + pub const fn new(id: u64) -> Self { + Self(id) + } + + pub const fn id(self) -> u64 { + self.0 + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct AzureAdAuth { + pub client_id: Option, + pub tenant_id: Option, + pub client_secret: Option, + pub username: Option, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct AwsIamAuth { + pub user_name: String, + pub cache_name: String, + pub region: String, + pub serverless: bool, +} diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index f2f7869876d..e659e6c1ad1 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -12,7 +12,12 @@ from contextlib import contextmanager, suppress from contextvars import ContextVar, Token from datetime import datetime, timezone from functools import wraps -from typing import Final, ParamSpec, TypeVar, cast +from typing import TYPE_CHECKING, Final, ParamSpec, TypeVar, cast + +if TYPE_CHECKING: + from opentelemetry.trace import Span +else: + Span: Final = object _P = ParamSpec("_P") _R = TypeVar("_R") @@ -39,6 +44,8 @@ _service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", de _mirrored_call: Final[ContextVar[bool]] = ContextVar("mirrored_call", default=False) +_request_parent_otel_span: Final[ContextVar[Span | None]] = ContextVar("request_parent_otel_span", default=None) + @contextmanager def internal_sub_call(parent_request_id: str | None = None) -> Generator[None]: @@ -154,6 +161,16 @@ def in_mirrored_call() -> bool: return _mirrored_call.get() +def set_request_parent_otel_span(span: Span) -> None: + """Anchor the proxy request's OTEL span for datastore service spans that no call site passes a parent to. + Never reset: it dies with the request task and rides the tasks it spawns, like ``request.state``.""" + _request_parent_otel_span.set(span) + + +def current_request_parent_otel_span() -> Span | None: + return _request_parent_otel_span.get() + + @contextmanager def outside_mirrored_call() -> Generator[None]: """Logging a mirrored call: the calls its callbacks start are not mirrors themselves.""" diff --git a/litellm/_redis.py b/litellm/_redis.py index 96777887c74..f9066d03270 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -13,6 +13,7 @@ import json # s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation import os from collections.abc import Callable, Mapping +from functools import reduce from types import MappingProxyType from typing import Final from urllib.parse import urlsplit, urlunsplit @@ -300,9 +301,7 @@ def _uses_tls(redis_kwargs: Mapping[str, object]) -> bool: def _build_elasticache_iam_provider(redis_kwargs: Mapping[str, object]) -> ElastiCacheIAMCredentialProvider: user_name: Final = redis_kwargs.get("aws_iam_user_name") cache_name: Final = redis_kwargs.get("aws_iam_cache_name") - region: Final = ( - redis_kwargs.get("aws_iam_region") or get_secret_str("AWS_REGION") or get_secret_str("AWS_DEFAULT_REGION") - ) + region: Final = redis_kwargs.get("aws_iam_region") required_settings: Final = ( ("aws_iam_user_name", user_name), ("aws_iam_cache_name", cache_name), @@ -510,36 +509,80 @@ def _url_without_userinfo(url: str) -> str: return urlunsplit((parts.scheme, netloc, parts.path, parts.query, parts.fragment)) +def _setting_text(redis_kwargs: Mapping[str, object], setting: str) -> str | None: + value: Final = redis_kwargs.get(setting) + return value if isinstance(value, str) and value else None + + +def _or_secrets(value: object, read: Callable[[str], object], *names: str) -> object: + return reduce(lambda found, name: found or read(name), names, value) + + +def _merge_redis_kwargs( + env_overrides: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: _get_redis_client_logic swaps settings for client objects in place + overrides: Final = { + k: get_secret(v.replace("os.environ/", "")) if isinstance(v, str) and v.startswith("os.environ/") else v + for k, v in env_overrides.items() + } + # An explicitly configured connection target outranks REDIS_URL from the + # environment. Without this, the url branch strips the caller's + # host/port/password and silently connects to whatever REDIS_URL names. + caller_named_a_target: Final = any( + overrides.get(key) is not None for key in ("host", "startup_nodes", "sentinel_nodes") + ) + drops_environment_url: Final = caller_named_a_target and overrides.get("url") is None + layered: Final = { + **{k: v for k, v in _redis_kwargs_from_environment().items() if not (k == "url" and drops_environment_url)}, + **overrides, + } + + def fallback(setting: str, read: Callable[[str], object], *names: str) -> tuple[str, object]: + return setting, _or_secrets(layered.get(setting), read, *names) + + targets: Final = ( + fallback("startup_nodes", get_secret, "REDIS_CLUSTER_NODES"), + fallback("sentinel_nodes", get_secret, "REDIS_SENTINEL_NODES"), + fallback("sentinel_password", get_secret_str, "REDIS_SENTINEL_PASSWORD"), + fallback("service_name", get_secret, "REDIS_SERVICE_NAME"), + ) + if layered.get("credential_provider") is not None: + return {**layered, **{k: v for k, v in targets if v is not None}} + + identities: Final = dict( + ( + fallback("gcp_service_account", get_secret_str, "REDIS_GCP_SERVICE_ACCOUNT"), + fallback("gcp_ssl_ca_certs", get_secret_str, "REDIS_GCP_SSL_CA_CERTS"), + fallback("azure_redis_ad_token", get_secret, "REDIS_AZURE_AD_TOKEN"), + ) + ) + gcp_enabled: Final = identities["gcp_service_account"] is not None + azure_enabled: Final = str(identities["azure_redis_ad_token"]).lower() == "true" + azure: Final = ( + ( + fallback("azure_client_id", get_secret_str, "AZURE_CLIENT_ID"), + fallback("azure_tenant_id", get_secret_str, "AZURE_TENANT_ID"), + fallback("azure_client_secret", get_secret_str, "AZURE_CLIENT_SECRET"), + ) + if azure_enabled and not gcp_enabled + else () + ) + aws: Final = ( + (fallback("aws_iam_region", get_secret_str, "AWS_REGION", "AWS_DEFAULT_REGION"),) + if _coerces_to_true(layered.get("aws_iam_auth")) and not gcp_enabled and not azure_enabled + else () + ) + filled: Final = (*targets, *identities.items(), *azure, *aws) + return {**layered, **{k: v for k, v in filled if v is not None}} + + def _get_redis_client_logic(**env_overrides): """ Common functionality across sync + async redis client implementations """ - ### check if "os.environ/" passed in - for k, v in env_overrides.items(): - if isinstance(v, str) and v.startswith("os.environ/"): - v = v.replace("os.environ/", "") - value = get_secret(v) - env_overrides[k] = value + redis_kwargs: Final = _merge_redis_kwargs(env_overrides) - environment_kwargs: Final = _redis_kwargs_from_environment() - - # An explicitly configured connection target outranks REDIS_URL from the - # environment. Without this, the url branch below strips the caller's - # host/port/password and silently connects to whatever REDIS_URL names. - caller_named_a_target: Final = any( - env_overrides.get(key) is not None for key in ("host", "startup_nodes", "sentinel_nodes") - ) - if caller_named_a_target and env_overrides.get("url") is None: - environment_kwargs.pop("url", None) - - redis_kwargs: Final = { - **environment_kwargs, - **env_overrides, - } - - _startup_nodes: Final[str | list | None] = redis_kwargs.get("startup_nodes", None) or get_secret( - "REDIS_CLUSTER_NODES" - ) + _startup_nodes: Final = redis_kwargs.get("startup_nodes") or None # If startup_nodes resolved to None (not set by kwarg or env), remove the key # entirely so callers can rely on key presence as a reliable cluster-mode signal. @@ -548,31 +591,14 @@ def _get_redis_client_logic(**env_overrides): elif _startup_nodes is None: redis_kwargs.pop("startup_nodes", None) - _sentinel_nodes: Final[str | list | None] = redis_kwargs.get("sentinel_nodes", None) or get_secret( - "REDIS_SENTINEL_NODES" - ) - - if _sentinel_nodes is not None and isinstance(_sentinel_nodes, str): + _sentinel_nodes: Final = redis_kwargs.get("sentinel_nodes") + if isinstance(_sentinel_nodes, str): redis_kwargs["sentinel_nodes"] = json.loads(_sentinel_nodes) - _sentinel_password: Final[str | None] = redis_kwargs.get("sentinel_password", None) or get_secret_str( - "REDIS_SENTINEL_PASSWORD" - ) - - if _sentinel_password is not None: - redis_kwargs["sentinel_password"] = _sentinel_password - - _service_name: Final[str | None] = redis_kwargs.get("service_name", None) or get_secret("REDIS_SERVICE_NAME") - - if _service_name is not None: - redis_kwargs["service_name"] = _service_name - if redis_kwargs.get("credential_provider") is None: # Handle GCP IAM authentication - _gcp_service_account: Final = redis_kwargs.get("gcp_service_account") or get_secret_str( - "REDIS_GCP_SERVICE_ACCOUNT" - ) - _gcp_ssl_ca_certs: Final = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS") + _gcp_service_account: Final = _setting_text(redis_kwargs, "gcp_service_account") + _gcp_ssl_ca_certs: Final = _setting_text(redis_kwargs, "gcp_ssl_ca_certs") if _gcp_service_account is not None: verbose_logger.debug("Setting up GCP IAM authentication for Redis with service account.") @@ -587,7 +613,7 @@ def _get_redis_client_logic(**env_overrides): redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs # Handle Azure AD authentication (after GCP IAM block) - _azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN") + _azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or None _azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true" _aws_iam_enabled: Final = _coerces_to_true(redis_kwargs.get("aws_iam_auth")) @@ -599,11 +625,9 @@ def _get_redis_client_logic(**env_overrides): ) if _azure_ad_enabled and _gcp_service_account is None: - _azure_client_id: Final = redis_kwargs.get("azure_client_id") or get_secret_str("AZURE_CLIENT_ID") - _azure_tenant_id: Final = redis_kwargs.get("azure_tenant_id") or get_secret_str("AZURE_TENANT_ID") - _azure_client_secret: Final = redis_kwargs.get("azure_client_secret") or get_secret_str( - "AZURE_CLIENT_SECRET" - ) + _azure_client_id: Final = _setting_text(redis_kwargs, "azure_client_id") + _azure_tenant_id: Final = _setting_text(redis_kwargs, "azure_tenant_id") + _azure_client_secret: Final = _setting_text(redis_kwargs, "azure_client_secret") verbose_logger.debug("Setting up Azure AD authentication for Redis.") redis_kwargs["redis_connect_func"] = create_azure_ad_redis_connect_func( @@ -649,8 +673,9 @@ def _get_redis_client_logic(**env_overrides): redis_kwargs.pop("redis_connect_func", None) redis_kwargs.pop("username", None) redis_kwargs.pop("password", None) - if redis_kwargs.get("url") is not None: - redis_kwargs["url"] = _url_without_userinfo(redis_kwargs["url"]) + _url: Final = redis_kwargs.get("url") + if isinstance(_url, str): + redis_kwargs["url"] = _url_without_userinfo(_url) if "url" in redis_kwargs and redis_kwargs["url"] is not None: # Only strip host/port/db/password when not routing to a cluster. diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index aa458c9c926..fc9aa999cac 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -19,8 +19,9 @@ from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence from contextvars import ContextVar, Token from dataclasses import dataclass, field from datetime import timedelta +from enum import Enum, auto from types import MappingProxyType, TracebackType -from typing import Final, Generic, Protocol, TypeVar +from typing import TYPE_CHECKING, Final, Generic, Literal, Protocol, TypeAlias, TypeVar from litellm._internal_context import ( REDIS_FAMILIES_METADATA_KEY, @@ -38,7 +39,11 @@ from litellm.caching.redis_cache import ( from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.types.services import ServiceTypes +if TYPE_CHECKING: + from litellm.rust_bridge._native import NativeRequestRedis, NativeSend + _T = TypeVar("_T") +_T_co: Final = TypeVar("_T_co", covariant=True) _ScriptArg = str | bytes | int | float SettledHook = Callable[[asyncio.Future[_T]], Awaitable[None] | None] POST_CALL_FLUSH_DEADLINE_SECONDS: Final = 1.0 @@ -443,6 +448,61 @@ def _backend_key(redis_cache: RedisCache) -> object: return (type(redis_cache), redis_cache.namespace, settings) +class NativePending(Protocol[_T_co]): + def resolve(self) -> Awaitable[_T_co]: ... + + +class NativeBatchResult(Generic[_T]): + """Awaitable handle for a read declared on the request's native pre-call batch; awaiting it runs the + native flush hooks, flushes that batch, then answers from what it read.""" + + __slots__ = ("_pending", "_request") + + def __init__(self, request: RequestRedisBatches, pending: NativePending[_T]) -> None: + self._request: Final = request + self._pending: Final = pending + + def __await__(self) -> Generator[object, None, _T]: + return self._wait().__await__() + + async def _wait(self) -> _T: + await self._request.flush_native_pre() + return await self._pending.resolve() + + +NativeRequestFactory: TypeAlias = Callable[[float], "NativeRequestRedis"] + + +def _native_request(post_call_deadline: float) -> NativeRequestRedis: + from litellm.rust_bridge._native import NativeRequestRedis + + return NativeRequestRedis(post_call_deadline) + + +class Batch(Enum): + """Which of the request's batches a store op rides: none, so it goes out on its own when awaited, the pre-call + batch of reads the request needs before answering, or the post-call batch of writes nobody waits on. Outside a + request every op goes out on its own.""" + + NONE = auto() + PRE_CALL = auto() + POST_CALL = auto() + + +RequestPhase: TypeAlias = Literal[Batch.PRE_CALL, Batch.POST_CALL] +NativeSendFactory: TypeAlias = Callable[[RequestPhase, "NativeRequestRedis"], "NativeSend"] + + +def _native_send(phase: RequestPhase, request: NativeRequestRedis) -> NativeSend: + from litellm.rust_bridge._native import NativeSend + + match phase: + case Batch.PRE_CALL: + return NativeSend.PreCall(request) + case Batch.POST_CALL: + return NativeSend.PostCall(request) + + _open_post_call: Final[weakref.WeakSet[RequestRedisBatches]] = weakref.WeakSet() """Requests whose post-call batch still holds declared ops, so a shutdown can send them before Redis goes away.""" @@ -460,15 +520,30 @@ class RequestRedisBatches: "_batches", "_deadline", "_deadline_flush", + "_native", + "_native_factory", + "_native_flush_hooks", + "_native_misses", + "_native_send_factory", "_post_call", "post_call_deadline", "prefetched", ) - def __init__(self, post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS) -> None: + def __init__( + self, + post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS, + native_factory: NativeRequestFactory = _native_request, + native_send_factory: NativeSendFactory = _native_send, + ) -> None: self._batches: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend self._post_call: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend self.post_call_deadline: Final = post_call_deadline + self._native_factory: Final = native_factory + self._native_send_factory: Final = native_send_factory + self._native: NativeRequestRedis | None = None + self._native_flush_hooks: Final[list[Callable[[], None]]] = [] # mutable-ok: append-only registry + self._native_misses: frozenset[str] = frozenset() self._deadline: asyncio.TimerHandle | None = None self._deadline_flush: asyncio.Task[None] | None = None # Reads declared early for a consumer that runs later in the request, keyed by consumer name. @@ -495,20 +570,72 @@ class RequestRedisBatches: _open_post_call.add(self) return batch + def native(self) -> NativeRequestRedis: + """The request's native batches, built on first use so requests that touch no native store pay nothing.""" + if self._native is None: + self._native = self._native_factory(self.post_call_deadline) + return self._native + + def native_post_call(self) -> NativeRequestRedis: + """``native`` for a post-call declaration, so a shutdown drain also flushes it.""" + _open_post_call.add(self) + return self.native() + + def native_send(self, phase: RequestPhase) -> NativeSend: + """The ``NativeSend`` for ``phase``; a post-call one registers the request with the shutdown drain.""" + match phase: + case Batch.PRE_CALL: + return self._native_send_factory(phase, self.native()) + case Batch.POST_CALL: + return self._native_send_factory(phase, self.native_post_call()) + + def native_result(self, pending: NativePending[_T]) -> NativeBatchResult[_T]: + return NativeBatchResult(self, pending) + + def add_pre_call_flush_hook(self, hook: Callable[[], None]) -> None: + """Called before every pre-call flush so late readers can declare into the same trip. The hook runs with this + request active, so a ``Batch.PRE_CALL`` op it declares rides this request whoever triggered the flush.""" + self._native_flush_hooks.append(hook) + + def note_native_misses(self, keys: Sequence[str]) -> None: + """Keys a native pre-call read found absent (or a native DEL removed), like ``RedisBatch.read_as_missing``.""" + self._native_misses = self._native_misses | frozenset(keys) + + def forget_native_miss(self, key: str) -> None: + self._native_misses = self._native_misses - frozenset((key,)) + + def native_read_as_missing(self, key: str) -> bool: + return key in self._native_misses + + async def flush_native_pre(self) -> None: + token: Final = _active_request_batches.set(self) + try: + for hook in self._native_flush_hooks: + hook() + finally: + _active_request_batches.reset(token) + if self._native is not None: + await self._native.flush_pre() + def _flush_on_deadline(self) -> None: self._deadline = None self._deadline_flush = asyncio.ensure_future(self.flush_post_call()) async def flush_all(self) -> None: - """Send whatever is still declared (write-backs nobody awaits) before the request scope closes.""" - await asyncio.gather(*(batch.flush() for batch in self._batches.values() if batch.pending)) + """Batch whatever is still declared (write-backs nobody awaits) before the request scope closes.""" + await asyncio.gather( + *(batch.flush() for batch in self._batches.values() if batch.pending), self.flush_native_pre() + ) async def flush_post_call(self) -> None: """One pipeline per backend for the post-call writes; the deadline is disarmed since this is that flush.""" if self._deadline is not None: self._deadline.cancel() self._deadline = None - await asyncio.gather(*(batch.flush() for batch in self._post_call.values() if batch.pending)) + await asyncio.gather( + *(batch.flush() for batch in self._post_call.values() if batch.pending), + *(() if self._native is None else (self._native.flush_post(),)), + ) if not any(batch.pending for batch in self._post_call.values()): _open_post_call.discard(self) @@ -557,17 +684,28 @@ async def drain_post_call_redis_batches() -> None: class request_redis_batch_scope: """Redis reads declared inside share one pipeline per backend; nested scopes join the outer one.""" - __slots__ = ("_post_call_deadline", "_token") + __slots__ = ("_native_factory", "_native_send_factory", "_post_call_deadline", "_token") - def __init__(self, post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS) -> None: + def __init__( + self, + post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS, + native_factory: NativeRequestFactory = _native_request, + native_send_factory: NativeSendFactory = _native_send, + ) -> None: self._token: Token[RequestRedisBatches | None] | None = None self._post_call_deadline: Final = post_call_deadline + self._native_factory: Final = native_factory + self._native_send_factory: Final = native_send_factory def __enter__(self) -> RequestRedisBatches: outer: Final = _active_request_batches.get() if outer is not None: return outer - batches: Final = RequestRedisBatches(post_call_deadline=self._post_call_deadline) + batches: Final = RequestRedisBatches( + post_call_deadline=self._post_call_deadline, + native_factory=self._native_factory, + native_send_factory=self._native_send_factory, + ) self._token = _active_request_batches.set(batches) return batches diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index afe6f8a1eab..fab613ab8a7 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -42,6 +42,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.coroutine_checker import coroutine_checker +from litellm.redis.store_settings import RedisStoreSettings from litellm.types.caching import ( RedisPipelineIncrementOperation, RedisPipelineLpopOperation, @@ -472,15 +473,21 @@ _LIVE_BREAKERS: Final[weakref.WeakSet[RedisCircuitBreaker]] = weakref.WeakSet() def redis_breaker_state_counts() -> Mapping[str, int]: - """How many live Redis circuit breakers sit in each state, for the Prometheus services logger.""" - return MappingProxyType(dict(Counter(breaker.state for breaker in _LIVE_BREAKERS))) + """How many live Redis circuit breakers sit in each state, for the Prometheus services logger. + + Counts the Python breakers of every live ``RedisCache`` and the breaker of every native Redis + client the router stores built, since each opens and closes on its own.""" + from litellm.rust_bridge._native import native_redis_breaker_states + + python_states: Final = Counter(breaker.state for breaker in _LIVE_BREAKERS) + return MappingProxyType(dict(python_states + Counter(native_redis_breaker_states()))) class RedisCircuitBreakerOpenError(Exception): pass -class _RedisTimeoutLogThrottle: +class RedisTimeoutLogThrottle: """Admits one Redis timeout log line per interval and counts the timeouts it suppressed in between.""" def __init__(self, interval: float, clock: Callable[[], float] = time.monotonic) -> None: @@ -503,7 +510,7 @@ class _RedisTimeoutLogThrottle: return suppressed -_redis_timeout_log_throttle: Final = _RedisTimeoutLogThrottle(REDIS_TIMEOUT_LOG_INTERVAL) +redis_timeout_log_throttle: Final = RedisTimeoutLogThrottle(REDIS_TIMEOUT_LOG_INTERVAL) def log_redis_failure( @@ -516,7 +523,7 @@ def log_redis_failure( if not is_redis_timeout_failure(exc): logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2) return - suppressed: Final = _redis_timeout_log_throttle.admit() + suppressed: Final = redis_timeout_log_throttle.admit() if suppressed is None: logger.debug("%s: %s", message, exc, stacklevel=2) return @@ -639,6 +646,8 @@ def _redis_circuit_breaker_guard_sync(method: Callable[..., _RedisCallResult]) - class RedisCache(BaseCache): # if users don't provider one, use the default litellm cache + store_settings: RedisStoreSettings | None = None + def __init__( self, host=None, @@ -705,6 +714,11 @@ class RedisCache(BaseCache): super().__init__(default_ttl=int(litellm.default_redis_ttl)) else: super().__init__() # defaults to 60s + self.store_settings = RedisStoreSettings( + redis_kwargs=MappingProxyType(dict(redis_kwargs)), + namespace=namespace, + default_ttl=self.default_ttl, + ) def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]: return call() diff --git a/litellm/redis/__init__.py b/litellm/redis/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/redis/store_settings.py b/litellm/redis/store_settings.py new file mode 100644 index 00000000000..dc3cf7d1391 --- /dev/null +++ b/litellm/redis/store_settings.py @@ -0,0 +1,11 @@ +from collections.abc import Mapping +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class RedisStoreSettings: + """What `connect` takes from a `RedisCache`, captured once when the cache is built.""" + + redis_kwargs: Mapping[str, object] + namespace: str | None + default_ttl: int diff --git a/litellm/redis/stores/__init__.py b/litellm/redis/stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/redis/stores/_route.py b/litellm/redis/stores/_route.py new file mode 100644 index 00000000000..5397c1be5d1 --- /dev/null +++ b/litellm/redis/stores/_route.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Iterable +from dataclasses import dataclass +from typing import Final, TypeVar + +from litellm._internal_context import current_service_target, service_target +from litellm.caching.redis_batch import Batch, NativePending, RequestRedisBatches, active_request_redis_batches +from litellm.rust_bridge._native import NativeRequestRedis, NativeSend + +_T: Final = TypeVar("_T") + + +@dataclass(frozen=True, slots=True) +class Route: + """The batch a native store op rides, and how its reply is awaited. A pre-call read awaits through + ``native_result``, so the pre-call flush hooks run first and late readers join the same trip. The reply is + reported under the service target the op was declared with, however late it is awaited.""" + + send: NativeSend + request: RequestRedisBatches | None + + def result(self, pending: NativePending[_T]) -> Awaitable[_T]: + if self.request is None: + return pending.resolve() + return _resolved_under(current_service_target(), self.request.native_result(pending)) + + +async def _resolved_under(target: str | None, reply: Awaitable[_T]) -> _T: + with service_target(target): + return await reply + + +async def as_tuple(reply: Awaitable[Iterable[_T]]) -> tuple[_T, ...]: + return tuple(await reply) + + +def route(batch: Batch) -> Route: + request: Final = active_request_redis_batches() + match batch: + case Batch.NONE: + return Route(NativeSend.Now(), None) + case Batch.PRE_CALL | Batch.POST_CALL: + return Route(NativeSend.Now(), None) if request is None else Route(request.native_send(batch), request) + + +def post_call_request() -> NativeRequestRedis | None: + """The current request's native batches for a post-call write nobody awaits, None outside a request.""" + request: Final = active_request_redis_batches() + return None if request is None else request.native_post_call() diff --git a/litellm/redis/stores/connection.py b/litellm/redis/stores/connection.py new file mode 100644 index 00000000000..7111f0df576 --- /dev/null +++ b/litellm/redis/stores/connection.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final + +from pydantic import TypeAdapter + +from litellm.constants import ( + REDIS_CIRCUIT_BREAKER_ENABLED, + REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, + REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT, + REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, +) +from litellm.rust_bridge._native import NativeRedisClient + +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + +_KWARGS: Final = TypeAdapter(dict[str, object]) + + +@dataclass(frozen=True, slots=True) +class RedisConnection: + """A Redis connection the stores share. Pass one to every store that should ride the same client.""" + + client: NativeRedisClient + + +def _breaker() -> tuple[int, float, float] | None: + if not REDIS_CIRCUIT_BREAKER_ENABLED: + return None + return ( + REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, + float(REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT), + REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, + ) + + +def connect(redis_cache: RedisCache | None) -> RedisConnection | None: + settings: Final = None if redis_cache is None else redis_cache.store_settings + if settings is None: + return None + from litellm._redis import ( # redis-py, which importing litellm must not require + _merge_redis_kwargs, # pyright: ignore[reportPrivateUsage] # Reuse the shared internal kwargs normalizer + ) + + client: Final = NativeRedisClient.connect( + _merge_redis_kwargs(_KWARGS.validate_python(settings.redis_kwargs)), + default_ttl_seconds=settings.default_ttl, + namespace=settings.namespace, + breaker=_breaker(), + ) + return None if client is None else RedisConnection(client) + + +def client_of(connection: RedisConnection | None) -> NativeRedisClient | None: + return None if connection is None else connection.client diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index dc170116089..76728b2e0a7 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -4,17 +4,22 @@ from typing import Generic, Literal, Never, NotRequired, TypedDict, TypeVar, fin import httpx from pydantic import JsonValue +from typing_extensions import disjoint_base from litellm.rust_bridge.execution import Execution from litellm.rust_bridge.public_call import NativeCall +_T_co = TypeVar("_T_co", covariant=True) + class RustBridgeDeclined(Exception): ... class RustUpstreamError(Exception): ... class ForkedAfterNativeRuntimeStarted(RuntimeError): ... class ProcessReservedForForking(RuntimeError): ... +class NativeRedisStoreError(Exception): ... def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... +def native_redis_breaker_states() -> list[str]: ... @final class NativeClickHouseSpendConfig: @@ -26,6 +31,50 @@ class NativeClickHouseSpendStorage: def ensure_schema(self) -> Future[None]: ... def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Future[None]: ... +@final +class NativeRedisClient: + def __new__(cls, _uninstantiable: Never, /) -> Never: ... + @staticmethod + def connect( + redis_kwargs: Mapping[str, object], + *, + default_ttl_seconds: float, + namespace: str | None = None, + breaker: tuple[int, float, float] | None = None, + ) -> NativeRedisClient | None: ... + @property + def default_ttl_seconds(self) -> float: ... + +@final +class NativeRequestRedis: + def __new__(cls, post_call_deadline_seconds: float) -> NativeRequestRedis: ... + def flush_pre(self) -> Future[None]: ... + def flush_post(self) -> Future[None]: ... + +@final +class NativePending(Generic[_T_co]): + def resolve(self) -> Future[_T_co]: ... + +@disjoint_base +class NativeSend: + @final + class Now(NativeSend): + __match_args__: tuple[()] = () + + def __new__(cls) -> NativeSend.Now: ... + + @final + class PreCall(NativeSend): + __match_args__: tuple[Literal["_0"]] = ("_0",) + + def __new__(cls, _0: NativeRequestRedis) -> NativeSend.PreCall: ... + + @final + class PostCall(NativeSend): + __match_args__: tuple[Literal["_0"]] = ("_0",) + + def __new__(cls, _0: NativeRequestRedis) -> NativeSend.PostCall: ... + @final class NativeDatabase: @staticmethod @@ -216,6 +265,11 @@ __all__ = [ "NativeClickHouseSpendStorage", "NativeDatabase", "NativeDiagnosticProcessor", + "NativePending", + "NativeRedisClient", + "NativeRedisStoreError", + "NativeRequestRedis", + "NativeSend", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RouterCall", @@ -235,6 +289,7 @@ __all__ = [ "embedding", "gil_stats", "messages", + "native_redis_breaker_states", "ocr", "ocr_health_check_document", "ocr_passthrough_response", diff --git a/litellm/rust_bridge/host/AGENTS.md b/litellm/rust_bridge/host/AGENTS.md index 9b751a1ada2..81463065acd 100644 --- a/litellm/rust_bridge/host/AGENTS.md +++ b/litellm/rust_bridge/host/AGENTS.md @@ -2,7 +2,7 @@ Python state Rust needs but does not own, behind one module per question. Rust imports a module here by path and calls a function by name -- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds) +- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds), `redis_events` (reporting native Redis calls and breaker transitions to `ServiceLogging`), `redis_credentials` (what a caller's redis-py `CredentialProvider` authenticates with) - A module is the only place Rust reads that Python state. Rust never does `py.import("litellm").getattr(...)` on a global; it asks here - Each module reuses the Python implementation the Python route runs (`AnthropicModelInfo`, `Logging`). Do not re-derive those rules here or in Rust; one resolver, two callers - The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green diff --git a/litellm/rust_bridge/host/redis_credentials.py b/litellm/rust_bridge/host/redis_credentials.py new file mode 100644 index 00000000000..ad6449e325e --- /dev/null +++ b/litellm/rust_bridge/host/redis_credentials.py @@ -0,0 +1,19 @@ +"""Credentials from a redis-py ``CredentialProvider`` an SDK caller passed in ``cache_kwargs``, asked the way +redis-py's async connection asks it.""" + +from __future__ import annotations + +import asyncio +from typing import Protocol + + +class CredentialProvider(Protocol): + async def get_credentials_async(self) -> tuple[str] | tuple[str, str]: ... + + +def credentials(provider: CredentialProvider) -> tuple[str | None, str]: + match asyncio.run(provider.get_credentials_async()): + case (password,): + return None, password + case (username, password): + return username, password diff --git a/litellm/rust_bridge/host/redis_events.py b/litellm/rust_bridge/host/redis_events.py new file mode 100644 index 00000000000..30116a1e15b --- /dev/null +++ b/litellm/rust_bridge/host/redis_events.py @@ -0,0 +1,124 @@ +"""Redis call and circuit breaker events from the native Redis client, reported the way ``RedisCache`` reports them.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable, Coroutine +from functools import lru_cache +from typing import Final, Literal, TypeAlias + +from redis import exceptions as redis_exceptions +from typing_extensions import assert_never + +from litellm._internal_context import current_request_parent_otel_span +from litellm._logging import verbose_logger +from litellm._service_logger import ServiceLogging +from litellm.types.services import ServiceTypes + +CallFailure: TypeAlias = Literal["timeout", "transport", "server", "decode"] +Failure: TypeAlias = CallFailure | Literal["breaker_open"] +BreakerState: TypeAlias = Literal["closed", "open", "half_open"] +HealthFailure: TypeAlias = Literal["timeout", "connectivity"] +Hook: TypeAlias = Callable[[], Coroutine[object, object, None]] + + +@lru_cache(maxsize=1) +def _service_logging() -> ServiceLogging: + return ServiceLogging() + + +def _redis_error(failure: CallFailure) -> redis_exceptions.RedisError: + """The redis-py exception ``RedisCache`` would have caught, so ``error_class`` keeps its labels.""" + match failure: + case "timeout": + return redis_exceptions.TimeoutError(failure) + case "transport": + return redis_exceptions.ConnectionError(failure) + case "server": + return redis_exceptions.ResponseError(failure) + case "decode": + return redis_exceptions.InvalidResponse(failure) + case _: + assert_never(failure) + + +def _health_failure(failure: CallFailure) -> HealthFailure | None: + match failure: + case "timeout": + return "timeout" + case "transport": + return "connectivity" + case "server" | "decode": + return None + case _: + assert_never(failure) + + +async def _guarded(hook: Hook) -> None: + try: + await hook() + except Exception as e: # noqa: BLE001 # Service-event hooks must not fail Redis operations + verbose_logger.exception("Error emitting a Redis service event - %s", e) + + +def _spawn(hook: Hook) -> None: + asyncio.get_running_loop().create_task(_guarded(hook)) + + +def _emit(event_loop: asyncio.AbstractEventLoop | None, hook: Hook) -> None: + """Schedule ``hook`` as a task on ``event_loop`` in the current context, as ``RedisCache`` did with + ``asyncio.create_task``; with no open loop, run it here the way ``ServiceLogging``'s sync hooks do.""" + if event_loop is None or event_loop.is_closed(): + asyncio.run(_guarded(hook)) + return + event_loop.call_soon_threadsafe(_spawn, hook) + + +def record_redis_call( + event_loop: asyncio.AbstractEventLoop | None, + call_type: str, + start_time: float, + end_time: float, + failure: Failure | None, +) -> None: + logger: Final = _service_logging() + duration: Final = end_time - start_time + parent_otel_span: Final = current_request_parent_otel_span() + match failure: + case "breaker_open": + return + case None: + _emit( + event_loop, + lambda: logger.async_service_success_hook( + service=ServiceTypes.REDIS, + call_type=call_type, + duration=duration, + parent_otel_span=parent_otel_span, + start_time=start_time, + end_time=end_time, + ), + ) + case "timeout" | "transport" | "server" | "decode": + error: Final = _redis_error(failure) + _emit( + event_loop, + lambda: logger.async_service_failure_hook( + service=ServiceTypes.REDIS, + duration=duration, + error=error, + call_type=call_type, + parent_otel_span=parent_otel_span, + start_time=start_time, + end_time=end_time, + ), + ) + health: Final = _health_failure(failure) + if health is not None: + logger.redis_breaker_hook(transitions_to=None, failure_class=health) + case _: + assert_never(failure) + + +def record_redis_breaker(state: BreakerState) -> None: + _service_logging().redis_breaker_hook(transitions_to=state, failure_class=None) diff --git a/litellm/rust_bridge/logger.py b/litellm/rust_bridge/logger.py index 544e7195848..8795625bc49 100644 --- a/litellm/rust_bridge/logger.py +++ b/litellm/rust_bridge/logger.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from collections.abc import Mapping from typing import Final @@ -27,6 +28,21 @@ def enabled(level: int) -> bool: return verbose_logger.isEnabledFor(level) +def _redis_timeout_line(level: int, message: str) -> tuple[int, str]: + """``log_redis_failure``'s timeout rule, on the throttle redis-py's timeouts count against too.""" + from litellm.caching.redis_cache import redis_timeout_log_throttle + + suppressed: Final = redis_timeout_log_throttle.admit() + if suppressed is None: + return logging.DEBUG, message + if suppressed == 0: + return level, message + return ( + level, + f"{message} ({suppressed} more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", + ) + + def emit( level: int, message: str, @@ -35,6 +51,22 @@ def emit( target: str, fields: Mapping[str, JsonValue], correlation: tuple[str, str], +) -> None: + if fields.get("redis.failure") == "timeout": + throttled_level, throttled_message = _redis_timeout_line(level, message) + _emit(throttled_level, throttled_message, pathname, lineno, target, fields, correlation) + return + _emit(level, message, pathname, lineno, target, fields, correlation) + + +def _emit( + level: int, + message: str, + pathname: str, + lineno: int, + target: str, + fields: Mapping[str, JsonValue], + correlation: tuple[str, str], ) -> None: if not enabled(level): return diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index 0911d5ef888..e9a075af507 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -21,7 +21,7 @@ from litellm.caching.caching import Cache, CacheMode, response_cache_phase from litellm.caching.caching_handler import _PENDING_CACHE_WRITES from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.redis_cache import RedisCache, _RedisTimeoutLogThrottle +from litellm.caching.redis_cache import RedisCache, RedisTimeoutLogThrottle from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.types.caching import EMBEDDING_CACHE_FORMAT_VERSION, LiteLLMCacheType, SemanticCacheScope from litellm.types.utils import Embedding, EmbeddingResponse, Usage @@ -94,9 +94,9 @@ def test_cache_key_debug_log_does_not_include_prompt_material(caplog): ], ) def test_add_cache_timeout_only_joins_redis_throttle_for_redis_backends(backend, expected_level, caplog, monkeypatch): - throttle = _RedisTimeoutLogThrottle(interval=5.0, clock=MagicMock(return_value=1_000.0)) + throttle = RedisTimeoutLogThrottle(interval=5.0, clock=MagicMock(return_value=1_000.0)) assert throttle.admit() == 0 - monkeypatch.setattr(redis_cache_module, "_redis_timeout_log_throttle", throttle) + monkeypatch.setattr(redis_cache_module, "redis_timeout_log_throttle", throttle) cache = Cache(type=LiteLLMCacheType.LOCAL) backend.set_cache.side_effect = TimeoutError("lit7520 backend timed out") diff --git a/tests/unit/caching/test_dual_cache.py b/tests/unit/caching/test_dual_cache.py index 5ebc3bf7336..0a9eee6630a 100644 --- a/tests/unit/caching/test_dual_cache.py +++ b/tests/unit/caching/test_dual_cache.py @@ -728,11 +728,11 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo from redis.exceptions import TimeoutError as RedisTimeoutError from litellm.caching import redis_cache as redis_cache_module - from litellm.caching.redis_cache import _RedisTimeoutLogThrottle + from litellm.caching.redis_cache import RedisTimeoutLogThrottle clock = MagicMock(return_value=1_000.0) monkeypatch.setattr( - redis_cache_module, "_redis_timeout_log_throttle", _RedisTimeoutLogThrottle(interval=5.0, clock=clock) + redis_cache_module, "redis_timeout_log_throttle", RedisTimeoutLogThrottle(interval=5.0, clock=clock) ) class _TimingOutRedis: diff --git a/tests/unit/caching/test_redis_batch.py b/tests/unit/caching/test_redis_batch.py index cd270035416..ea8ad0ebcea 100644 --- a/tests/unit/caching/test_redis_batch.py +++ b/tests/unit/caching/test_redis_batch.py @@ -6,8 +6,9 @@ import asyncio import hashlib import json from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass, field from datetime import timedelta -from typing import Any +from typing import Any, Generic, TypeVar import pytest from redis.exceptions import NoScriptError @@ -26,8 +27,10 @@ from litellm.caching.redis_cache import ( _get_call_stack_info, # pyright: ignore[reportPrivateUsage] # the chain the service hook reports ) from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.redis.store_settings import RedisStoreSettings SCRIPT = "return redis.call('GET', KEYS[1])" +_Answer = TypeVar("_Answer") SHA = hashlib.sha1(SCRIPT.encode()).hexdigest() # noqa: S324 @@ -93,6 +96,8 @@ class FakeRedisCache(RedisCache): def __init__(self, client: FakeClient, namespace: str | None = None) -> None: # super().__init__ needs a server self.client = client self.namespace = namespace + self.redis_kwargs = {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2, "db": id(self) % 1_000_000} + self.store_settings = RedisStoreSettings(redis_kwargs=self.redis_kwargs, namespace=namespace, default_ttl=60) self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30) self.service_logger_obj = ServiceLogging() self.default_ttl = None @@ -147,6 +152,63 @@ def replies(command: tuple[Any, ...]) -> Any: raise AssertionError(command) +@dataclass +class FakeNativeRequest: + """Stands in for ``NativeRequestRedis``: records each flush in ``events`` next to what the test's pendings log.""" + + post_call_deadline: float + events: list[str] = field(default_factory=list) + + async def flush_pre(self) -> None: + self.events.append("flush_pre") + + async def flush_post(self) -> None: + self.events.append("flush_post") + + +@dataclass +class FakeNatives: + built: list[FakeNativeRequest] = field(default_factory=list) + + def __call__(self, post_call_deadline: float) -> FakeNativeRequest: + self.built.append(FakeNativeRequest(post_call_deadline)) + return self.built[-1] + + +@dataclass +class FakeNativePending(Generic[_Answer]): + """A read declared on a ``FakeNativeRequest``; it only answers once that request's pre-call batch was flushed.""" + + request: FakeNativeRequest + answer: Callable[[], _Answer] + + async def resolve(self) -> _Answer: + if "flush_pre" not in self.request.events: + raise ConnectionError("resolved before its batch was flushed") + self.request.events.append("resolve") + return self.answer() + + +@dataclass +class ReadyPending(Generic[_Answer]): + """An op declared with ``NativeSend.Now()``: it answers whenever it is resolved.""" + + answer: Callable[[], _Answer] + + async def resolve(self) -> _Answer: + return self.answer() + + +def fake_send(_phase: object, request: FakeNativeRequest) -> FakeNativeRequest: + """A ``native_send_factory`` for ``FakeNatives``: the fake request stands in for its ``NativeSend``.""" + return request + + +def fake_pending(send: object, answer: Callable[[], _Answer]) -> FakeNativePending[_Answer] | ReadyPending[_Answer]: + """What a fake handle declares for ``send``: on the fake request's batch, or on its own for ``Now()``.""" + return FakeNativePending(send, answer) if isinstance(send, FakeNativeRequest) else ReadyPending(answer) + + def make(fail: Exception | None = None, namespace: str | None = None) -> tuple[FakeRedisCache, FakeClient]: client = FakeClient(replies, fail) return FakeRedisCache(client, namespace), client diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 2bdc7b9549b..0fbeb05e154 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -1198,11 +1198,11 @@ def test_timeouts_during_a_blip_log_once_per_interval_not_once_per_call(sync_bat from redis.exceptions import TimeoutError as RedisTimeoutError from litellm.caching import redis_cache as redis_cache_module - from litellm.caching.redis_cache import _RedisTimeoutLogThrottle + from litellm.caching.redis_cache import RedisTimeoutLogThrottle clock = MagicMock(return_value=1_000.0) monkeypatch.setattr( - redis_cache_module, "_redis_timeout_log_throttle", _RedisTimeoutLogThrottle(interval=5.0, clock=clock) + redis_cache_module, "redis_timeout_log_throttle", RedisTimeoutLogThrottle(interval=5.0, clock=clock) ) sync_batch_redis_cache.redis_client.get.side_effect = RedisTimeoutError("Timeout reading from 127.0.0.1:6379") sync_batch_redis_cache.redis_client.mget.side_effect = RedisTimeoutError("Timeout reading from 127.0.0.1:6379") @@ -1272,12 +1272,12 @@ async def test_write_path_timeouts_inside_the_interval_stay_at_debug(call_method from redis.exceptions import TimeoutError as RedisTimeoutError from litellm.caching import redis_cache as redis_cache_module - from litellm.caching.redis_cache import _RedisTimeoutLogThrottle + from litellm.caching.redis_cache import RedisTimeoutLogThrottle clock = MagicMock(return_value=1_000.0) - throttle = _RedisTimeoutLogThrottle(interval=5.0, clock=clock) + throttle = RedisTimeoutLogThrottle(interval=5.0, clock=clock) assert throttle.admit() == 0 - monkeypatch.setattr(redis_cache_module, "_redis_timeout_log_throttle", throttle) + monkeypatch.setattr(redis_cache_module, "redis_timeout_log_throttle", throttle) timeout = RedisTimeoutError("Timeout reading from 127.0.0.1:6379") client = MagicMock() diff --git a/tests/unit/redis/__init__.py b/tests/unit/redis/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/redis/stores/__init__.py b/tests/unit/redis/stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/redis/stores/test_connection.py b/tests/unit/redis/stores/test_connection.py new file mode 100644 index 00000000000..6fe25403592 --- /dev/null +++ b/tests/unit/redis/stores/test_connection.py @@ -0,0 +1,74 @@ +import gc +import uuid +from concurrent.futures import ThreadPoolExecutor +from typing import Final + +import pytest + +from litellm.caching.redis_cache import RedisCache +from litellm.redis.stores.connection import connect +from litellm.rust_bridge._native import NativeRedisClient, native_redis_breaker_states + + +def test_concurrent_first_connects_share_one_client_without_deadlocking() -> None: + before: Final = len(native_redis_breaker_states()) + namespace: Final = f"concurrent-connect-{uuid.uuid4()}" + pool: Final = ThreadPoolExecutor(max_workers=8) + futures: Final = tuple( + pool.submit( + NativeRedisClient.connect, + {"host": "127.0.0.1", "port": 1}, + default_ttl_seconds=60, + namespace=namespace, + breaker=(5, 60.0, 5.0), + ) + for _ in range(8) + ) + try: + clients: Final = tuple(future.result(timeout=30) for future in futures) + finally: + pool.shutdown(wait=False) + + assert all(client is not None for client in clients) + assert len(native_redis_breaker_states()) == before + 1 + + +def _breaker_states_while_connected(namespace: str) -> int: + client: Final = NativeRedisClient.connect( + {"host": "127.0.0.1", "port": 1}, default_ttl_seconds=60, namespace=namespace, breaker=(5, 60.0, 5.0) + ) + assert client is not None + return len(native_redis_breaker_states()) + + +def test_a_discarded_client_leaves_the_process_registry() -> None: + gc.collect() + before: Final = len(native_redis_breaker_states()) + + while_connected: Final = _breaker_states_while_connected(f"discarded-{uuid.uuid4()}") + gc.collect() + + assert while_connected == before + 1 + assert len(native_redis_breaker_states()) == before + + +def test_a_redis_cache_captures_what_connect_needs_at_construction() -> None: + cache: Final = RedisCache(host="127.0.0.1", port=1, namespace="ns", socket_timeout=0.2) + + assert cache.store_settings is not None + assert dict(cache.store_settings.redis_kwargs) == {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2} + assert cache.store_settings.namespace == "ns" + assert cache.store_settings.default_ttl == cache.default_ttl + + +class _UnbuiltRedisCache(RedisCache): + def __init__(self) -> None: + pass + + +def test_a_redis_cache_stand_in_that_skips_the_constructor_connects_nothing() -> None: + assert connect(_UnbuiltRedisCache()) is None + + +def test_no_redis_cache_connects_nothing() -> None: + assert connect(None) is None diff --git a/tests/unit/rust_bridge/host/test_redis_credentials.py b/tests/unit/rust_bridge/host/test_redis_credentials.py new file mode 100644 index 00000000000..4a12c8807ba --- /dev/null +++ b/tests/unit/rust_bridge/host/test_redis_credentials.py @@ -0,0 +1,38 @@ +import asyncio +from typing import Final + +import pytest +from redis.credentials import CredentialProvider + +from litellm.rust_bridge.host.redis_credentials import credentials + + +class _SyncOnly(CredentialProvider): + def __init__(self, answer: tuple[str] | tuple[str, str]) -> None: + self.answer: Final = answer + + def get_credentials(self) -> tuple[str] | tuple[str, str]: + return self.answer + + +class _AsyncOnly(CredentialProvider): + def get_credentials(self) -> tuple[str, str]: + raise AssertionError("redis-py's async connection asks get_credentials_async") + + async def get_credentials_async(self) -> tuple[str, str]: + await asyncio.sleep(0) + return "async-user", "async-password" + + +@pytest.mark.parametrize( + ("provider", "expected"), + [ + (_SyncOnly(("token",)), (None, "token")), + (_SyncOnly(("user", "password")), ("user", "password")), + (_AsyncOnly(), ("async-user", "async-password")), + ], +) +def test_credentials_are_what_redis_py_would_authenticate_with( + provider: CredentialProvider, expected: tuple[str | None, str] +) -> None: + assert credentials(provider) == expected diff --git a/tests/unit/rust_bridge/host/test_redis_events.py b/tests/unit/rust_bridge/host/test_redis_events.py new file mode 100644 index 00000000000..85f856113c5 --- /dev/null +++ b/tests/unit/rust_bridge/host/test_redis_events.py @@ -0,0 +1,139 @@ +import asyncio +import queue +import threading +from collections.abc import Iterator +from typing import Final + +import pytest +from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor, TracerProvider +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from prometheus_client import REGISTRY +from redis import exceptions as redis_exceptions + +import litellm +from litellm.integrations.otel import OpenTelemetryV2Config +from litellm.integrations.otel.logger import OpenTelemetryV2 +from litellm.integrations.otel.plumbing import providers +from litellm.rust_bridge.host import redis_events + + +def _sample(name: str, labels: dict[str, str]) -> float: + return REGISTRY.get_sample_value(name, labels) or 0.0 + + +@pytest.fixture +def prometheus(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + redis_events._service_logging.cache_clear() + yield + redis_events._service_logging.cache_clear() + + +def _failed_requests(error: type[redis_exceptions.RedisError]) -> tuple[str, dict[str, str]]: + return ( + "litellm_redis_failed_requests_total", + {"redis": "redis", "error_class": error.__name__, "function_name": "async_increment"}, + ) + + +@pytest.mark.parametrize( + ("failure", "error", "failure_class"), + [ + ("timeout", redis_exceptions.TimeoutError, "timeout"), + ("transport", redis_exceptions.ConnectionError, "connectivity"), + ], +) +def test_a_health_failure_counts_as_a_failed_request_of_its_redis_py_class_and_a_breaker_failure( + prometheus: None, failure: redis_events.Failure, error: type[redis_exceptions.RedisError], failure_class: str +): + breaker: Final = ("litellm_redis_circuit_breaker_failures_total", {"failure_class": failure_class}) + before: Final = (_sample(*_failed_requests(error)), _sample(*breaker)) + + redis_events.record_redis_call(None, "async_increment", 10.0, 10.25, failure) + + assert (_sample(*_failed_requests(error)), _sample(*breaker)) == (before[0] + 1, before[1] + 1) + + +@pytest.mark.parametrize( + ("failure", "error"), + [("server", redis_exceptions.ResponseError), ("decode", redis_exceptions.InvalidResponse)], +) +def test_a_rejected_command_is_a_failed_request_of_its_redis_py_class_but_not_a_breaker_failure( + prometheus: None, failure: redis_events.Failure, error: type[redis_exceptions.RedisError] +): + breaker: Final = { + failure_class: _sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": failure_class}) + for failure_class in ("timeout", "connectivity") + } + before: Final = _sample(*_failed_requests(error)) + + redis_events.record_redis_call(None, "async_increment", 10.0, 10.25, failure) + + assert _sample(*_failed_requests(error)) == before + 1 + assert { + failure_class: _sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": failure_class}) + for failure_class in ("timeout", "connectivity") + } == breaker + + +def test_a_success_is_counted_and_a_breaker_rejection_is_not(prometheus: None): + total: Final = ("litellm_redis_total_requests_total", {"redis": "redis"}) + before: Final = _sample(*total) + + redis_events.record_redis_call(None, "async_get_cache", 10.0, 10.01, None) + redis_events.record_redis_call(None, "async_get_cache", 10.0, 10.0, "breaker_open") + + assert _sample(*total) == before + 1 + + +def test_a_call_whose_loop_has_closed_is_still_counted(prometheus: None): + total: Final = ("litellm_redis_total_requests_total", {"redis": "redis"}) + closed: Final = asyncio.new_event_loop() + closed.close() + before: Final = _sample(*total) + + redis_events.record_redis_call(closed, "async_get_cache", 10.0, 10.01, None) + + assert _sample(*total) == before + 1 + + +@pytest.mark.parametrize("state", ["open", "half_open", "closed"]) +def test_breaker_transitions_are_counted_by_state(prometheus: None, state: redis_events.BreakerState): + labels: Final = {"state": state} + before: Final = _sample("litellm_redis_circuit_breaker_transitions_total", labels) + + redis_events.record_redis_breaker(state) + + assert _sample("litellm_redis_circuit_breaker_transitions_total", labels) == before + 1 + + +@pytest.fixture +def otel(monkeypatch: pytest.MonkeyPatch) -> Iterator[tuple[TracerProvider, InMemorySpanExporter]]: + config: Final = OpenTelemetryV2Config(exporter="in_memory") + exporter: Final = InMemorySpanExporter() + provider: Final = providers.build_tracer_provider(config, exporter=exporter) + monkeypatch.setattr(litellm, "service_callback", [OpenTelemetryV2(config=config, tracer_provider=provider)]) + yield provider, exporter + + +class _EndedSpans(SpanProcessor): + def __init__(self) -> None: + self.ended: Final = queue.SimpleQueue[tuple[ReadableSpan, int]]() + + def on_end(self, span: ReadableSpan) -> None: + self.ended.put((span, threading.get_ident())) + + +async def test_a_call_reported_from_a_worker_thread_is_traced_on_the_callers_loop_with_its_times( + otel: tuple[TracerProvider, InMemorySpanExporter], +): + provider, _exporter = otel + ended: Final = _EndedSpans() + provider.add_span_processor(ended) + + await asyncio.to_thread( + redis_events.record_redis_call, asyncio.get_running_loop(), "async_get_cache", 10.0, 10.25, None + ) + + span, thread = await asyncio.to_thread(ended.ended.get, timeout=5) + assert (span.start_time, span.end_time, thread) == (10_000_000_000, 10_250_000_000, threading.get_ident()) diff --git a/tests/unit/rust_bridge/test_logger.py b/tests/unit/rust_bridge/test_logger.py index 4745aca9d50..df3aedf4cb9 100644 --- a/tests/unit/rust_bridge/test_logger.py +++ b/tests/unit/rust_bridge/test_logger.py @@ -69,6 +69,32 @@ def test_native_logging_observes_level_changes(caplog: pytest.LogCaptureFixture) assert [record.getMessage() for record in caplog.records] == ["visible"] +def test_native_redis_timeouts_share_the_redis_py_timeout_throttle( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching import redis_cache as redis_cache_module + from litellm.caching.redis_cache import RedisTimeoutLogThrottle, log_redis_failure + + monkeypatch.setattr( + redis_cache_module, "redis_timeout_log_throttle", RedisTimeoutLogThrottle(interval=5.0, clock=lambda: 1_000.0) + ) + store: Final = "litellm_redis_store::rate_limits" + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + logger.emit(logging.WARNING, "native timeout", "mod.rs", 1, store, {"redis.failure": "timeout"}, ("", "")) + log_redis_failure(verbose_logger, logging.WARNING, "redis-py timeout", RedisTimeoutError("timed out")) + logger.emit(logging.WARNING, "native timeout again", "mod.rs", 2, store, {"redis.failure": "timeout"}, ("", "")) + logger.emit(logging.WARNING, "native refused", "mod.rs", 3, store, {"redis.failure": "transport"}, ("", "")) + + assert [(record.getMessage(), record.levelno) for record in caplog.records] == [ + ("native timeout", logging.WARNING), + ("redis-py timeout: timed out", logging.DEBUG), + ("native timeout again", logging.DEBUG), + ("native refused", logging.WARNING), + ] + + @pytest.mark.parametrize( "text", ( diff --git a/tests/unit/test_redis.py b/tests/unit/test_redis.py index bfe3b55497d..9a5af89207b 100644 --- a/tests/unit/test_redis.py +++ b/tests/unit/test_redis.py @@ -1,13 +1,20 @@ +import importlib.util import inspect import json -from types import SimpleNamespace +import os +import ssl +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType, SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest import redis import redis.asyncio as async_redis +from redis.asyncio.connection import AbstractConnection from redis.credentials import CredentialProvider +from redis.exceptions import RedisClusterException, RedisError import litellm from litellm._redis import ( @@ -19,6 +26,7 @@ from litellm._redis import ( _get_redis_env_kwarg_mapping, _get_redis_kwargs, _get_redis_url_kwargs, + _merge_redis_kwargs, _pretty_print_redis_config, _uses_tls, get_redis_async_client, @@ -1960,3 +1968,1413 @@ def test_async_sentinel_keeps_the_credential_provider_off_the_monitors(markers, master_kwargs = mock_sentinel_cls.return_value.master_for.call_args[1] assert isinstance(master_kwargs["credential_provider"], provider_cls) assert "password" not in master_kwargs + + +def _tls_view(connection: AbstractConnection) -> dict[str, object] | None: + if not isinstance(connection, async_redis.SSLConnection): + return None + context: Final = connection.ssl_context + return { + "ca_certs": context.ca_certs, + "ca_data": context.ca_data, + "ca_path": None, + "certfile": context.certfile, + "keyfile": context.keyfile, + "verify_certificate": context.cert_reqs != ssl.CERT_NONE, + "check_hostname": context.check_hostname, + } + + +def _auth_view(kwargs: dict[str, object]) -> dict[str, object]: + provider: Final = kwargs.get("credential_provider") + match provider: + case None: + return {"kind": "static", "username": kwargs.get("username"), "password": kwargs.get("password")} + case GCPIAMCredentialProvider(): + return {"kind": "gcp_iam", "service_account": provider._gcp_service_account} + case AzureADCredentialProvider(): + return {"kind": "azure_ad", "username": provider._username} + case ElastiCacheIAMCredentialProvider(): + return { + "kind": "aws_iam", + "user_name": provider._user_name, + "cache_name": provider._cache_name, + "region": provider._region, + "serverless": provider._is_serverless, + } + case _: + return {"kind": "host"} + + +def _connection_view(connection: AbstractConnection, kwargs: dict[str, object]) -> dict[str, object]: + return { + "tls": _tls_view(connection), + "auth": _auth_view(kwargs), + "socket_timeout": connection.socket_timeout, + "socket_connect_timeout": connection.socket_connect_timeout, + "health_check_interval": float(connection.health_check_interval) or None, + } + + +def _resolution_view(kwargs: dict[str, object]) -> dict[str, object]: + url: Final = _credential_provider_auth_kwargs(_get_redis_client_logic(**kwargs)).get("url") + client: Final = get_redis_async_client(**kwargs) + if isinstance(client, async_redis.RedisCluster): + nodes: Final = list(client.nodes_manager.startup_nodes.values()) + return { + "mode": "cluster", + "nodes": [f"{node.host}:{node.port}" for node in nodes], + **_connection_view(nodes[0].acquire_connection(), client.connection_kwargs), + } + pool: Final = client.connection_pool + connection: Final = pool.make_connection() + if isinstance(pool, async_redis.SentinelConnectionPool): + sentinels: Final = [sentinel.connection_pool.connection_kwargs for sentinel in pool.sentinel_manager.sentinels] + return { + "mode": "sentinel", + "nodes": [f"{sentinel['host']}:{sentinel['port']}" for sentinel in sentinels], + "service_name": pool.service_name, + "sentinel_password": sentinels[0].get("password"), + "db": connection.db, + **_connection_view(connection, pool.connection_kwargs), + } + match connection: + case async_redis.UnixDomainSocketConnection(): + target = {"mode": "unix", "path": connection.path} + case _ if url is not None: + target = {"mode": "url", "url": url} + case _: + target = {"mode": "host", "host": connection.host, "port": connection.port, "db": connection.db} + return {**target, **_connection_view(connection, pool.connection_kwargs)} + + +@dataclass(frozen=True, slots=True) +class _Fails: + """redis-py raises building the client or its first connection, so Rust failing too is no behavior change""" + + type: str + message: str + + +def _expected_view(fields: dict[str, object]) -> dict[str, object]: + return { + "auth": {"kind": "static", "username": None, "password": None}, + "tls": None, + "socket_timeout": 5.0, + "socket_connect_timeout": 5.0, + "health_check_interval": None, + **fields, + } + + +@pytest.mark.skipif( + redis.VERSION != (5, 3, 1), + reason="the Rust interpreter in litellm-rust/crates/redis-resolve pins redis-py 5.3.1, per its AGENTS.md", +) +@pytest.mark.parametrize( + ("kwargs", "environment", "merged", "expected"), + [ + pytest.param( + {"host": "redis"}, + {}, + {"host": "redis"}, + _expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}), + id="host_defaults", + ), + pytest.param( + {"startup_nodes": []}, + {"REDIS_CLUSTER_NODES": '[{"host": "node-a", "port": 7000}]'}, + {"startup_nodes": '[{"host": "node-a", "port": 7000}]'}, + _expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"]}), + id="empty_cluster_nodes_use_the_environment", + ), + pytest.param( + {"sentinel_nodes": [], "service_name": "primary"}, + {"REDIS_SENTINEL_NODES": '[["sentinel-a", 26379]]'}, + {"sentinel_nodes": '[["sentinel-a", 26379]]', "service_name": "primary"}, + _expected_view( + { + "db": 0, + "mode": "sentinel", + "nodes": ["sentinel-a:26379"], + "sentinel_password": None, + "service_name": "primary", + } + ), + id="empty_sentinel_nodes_use_the_environment", + ), + pytest.param( + {"sentinel_nodes": [["sentinel-a", 26379]], "service_name": ""}, + {"REDIS_SERVICE_NAME": "primary"}, + {"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}, + _expected_view( + { + "db": 0, + "mode": "sentinel", + "nodes": ["sentinel-a:26379"], + "sentinel_password": None, + "service_name": "primary", + } + ), + id="empty_service_name_uses_the_environment", + ), + pytest.param( + {"host": "redis", "ssl_cert_reqs": "strict"}, + {}, + {"host": "redis", "ssl_cert_reqs": "strict"}, + _expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}), + id="plaintext_ignores_tls_requirements", + ), + pytest.param( + {"host": "redis", "ssl": "on"}, + {}, + {"host": "redis", "ssl": "on"}, + _expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}), + id="an_ambiguous_flag_is_refused", + ), + pytest.param( + {"sentinel_nodes": [], "service_name": "primary"}, + {}, + {"sentinel_nodes": [], "service_name": "primary"}, + _Fails("ValueError", "Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel."), + id="empty_sentinel_nodes_without_the_environment_name_no_target", + ), + pytest.param( + {"host": "redis", "ssl": True, "ssl_ca_path": "/ca"}, + {}, + {"host": "redis", "ssl": True, "ssl_ca_path": "/ca"}, + { + "rejected": { + "message": "Redis.__init__() got an unexpected keyword argument 'ssl_ca_path'", + "type": "TypeError", + } + }, + id="ssl_ca_path_is_unsupported", + ), + pytest.param( + {"host": "redis", "ssl": True, "ssl_certfile": "/client.pem", "ssl_password": "key-secret"}, + {}, + {"host": "redis", "ssl": True, "ssl_certfile": "/client.pem", "ssl_password": "key-secret"}, + { + "rejected": { + "message": "Redis.__init__() got an unexpected keyword argument 'ssl_password'", + "type": "TypeError", + } + }, + id="ssl_password_is_unsupported", + ), + pytest.param( + {"host": "redis", "ssl": True, "ssl_validate_ocsp": True}, + {}, + {"host": "redis", "ssl": True, "ssl_validate_ocsp": True}, + _Fails("TypeError", "Redis.__init__() got an unexpected keyword argument 'ssl_validate_ocsp'"), + id="ssl_validate_ocsp_is_unsupported", + ), + pytest.param( + {"host": "redis", "ssl": True, "ssl_include_verify_flags": ["VERIFY_CRL_CHECK_LEAF"]}, + {}, + {"host": "redis", "ssl": True, "ssl_include_verify_flags": ["VERIFY_CRL_CHECK_LEAF"]}, + _Fails("TypeError", "Redis.__init__() got an unexpected keyword argument 'ssl_include_verify_flags'"), + id="ssl_include_verify_flags_is_unsupported", + ), + pytest.param( + { + "db": 2, + "host": "redis", + "password": "pass", + "port": 6380, + "socket_timeout": 1.5, + "ssl": True, + "username": "user", + }, + {}, + { + "db": 2, + "host": "redis", + "password": "pass", + "port": 6380, + "socket_timeout": 1.5, + "ssl": True, + "username": "user", + }, + _expected_view( + { + "auth": {"kind": "static", "password": "pass", "username": "user"}, + "db": 2, + "host": "redis", + "mode": "host", + "port": 6380, + "socket_connect_timeout": 1.5, + "socket_timeout": 1.5, + "tls": { + "ca_certs": None, + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": True, + }, + } + ), + id="host_with_every_basic_setting", + ), + pytest.param( + {}, + { + "REDIS_HOST": "environment", + "REDIS_PASSWORD": "environment-secret", + "REDIS_PORT": "6390", + "REDIS_SSL": "True", + }, + {"port": "6390", "ssl": True, "host": "environment", "password": "environment-secret"}, + _expected_view( + { + "auth": {"kind": "static", "password": "environment-secret", "username": None}, + "db": 0, + "host": "environment", + "mode": "host", + "port": 6390, + "tls": { + "ca_certs": None, + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": True, + }, + } + ), + id="environment_host_port_and_password", + ), + pytest.param( + {"db": "3", "host": "redis", "port": "6380", "socket_timeout": "2.5"}, + {}, + {"db": "3", "host": "redis", "port": "6380", "socket_timeout": "2.5"}, + _expected_view( + { + "db": 3, + "host": "redis", + "mode": "host", + "port": 6380, + "socket_connect_timeout": 2.5, + "socket_timeout": 2.5, + } + ), + id="string_numbers_are_coerced", + ), + pytest.param( + {"host": "redis", "port": "6380\n", "db": " 2 ", "socket_timeout": "1_0.5\n"}, + {}, + {"host": "redis", "port": "6380\n", "db": " 2 ", "socket_timeout": "1_0.5\n"}, + _expected_view( + { + "db": 2, + "host": "redis", + "mode": "host", + "port": 6380, + "socket_timeout": 10.5, + "socket_connect_timeout": 10.5, + } + ), + id="padded_and_underscored_numbers_parse_like_python_int_and_float", + ), + pytest.param( + {}, + {"REDIS_HOST": "redis", "REDIS_PORT": "6_380", "REDIS_DB": "3\n"}, + {"port": "6_380", "host": "redis", "db": "3\n"}, + _expected_view({"db": 3, "host": "redis", "mode": "host", "port": 6380}), + id="padded_environment_numbers_parse_like_python_int", + ), + pytest.param( + {"host": "redis"}, + {"REDIS_SOCKET_TIMEOUT": "0.75"}, + {"socket_timeout": "0.75", "host": "redis"}, + _expected_view( + { + "db": 0, + "host": "redis", + "mode": "host", + "port": 6379, + "socket_connect_timeout": 0.75, + "socket_timeout": 0.75, + } + ), + id="socket_timeout_from_the_environment", + ), + pytest.param( + {}, + {"REDIS_URL": "redis://environment:6379/1"}, + {"url": "redis://environment:6379/1"}, + _expected_view({"mode": "url", "url": "redis://environment:6379/1"}), + id="environment_url", + ), + pytest.param( + {}, + {"REDIS_HOST": "ignored", "REDIS_URL": "redis://from-url:6379"}, + {"host": "ignored", "url": "redis://from-url:6379"}, + _expected_view({"mode": "url", "url": "redis://from-url:6379"}), + id="environment_url_beats_the_environment_host", + ), + pytest.param( + {"host": "configured"}, + {"REDIS_URL": "redis://environment:6379"}, + {"host": "configured"}, + _expected_view({"db": 0, "host": "configured", "mode": "host", "port": 6379}), + id="an_explicit_host_drops_the_environment_url", + ), + pytest.param( + {"db": 4, "host": "dropped", "password": "dropped", "url": "rediss://cache:6380", "username": "user"}, + {}, + {"db": 4, "host": "dropped", "password": "dropped", "url": "rediss://cache:6380", "username": "user"}, + _expected_view( + { + "auth": {"kind": "static", "password": None, "username": "user"}, + "mode": "url", + "tls": { + "ca_certs": None, + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": True, + }, + "url": "rediss://cache:6380", + } + ), + id="a_url_drops_host_db_and_password_and_keeps_the_username", + ), + pytest.param( + { + "ssl_ca_certs": "/ca.pem", + "ssl_cert_reqs": "none", + "ssl_check_hostname": True, + "url": "rediss://cache:6380", + }, + {}, + { + "ssl_ca_certs": "/ca.pem", + "ssl_cert_reqs": "none", + "ssl_check_hostname": True, + "url": "rediss://cache:6380", + }, + _expected_view( + { + "mode": "url", + "tls": { + "ca_certs": None, + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": True, + }, + "url": "rediss://cache:6380", + } + ), + id="a_url_ignores_tls_settings_outside_the_url", + ), + pytest.param( + { + "password": "cluster-secret", + "startup_nodes": [{"host": "node-a", "port": 7000}], + "url": "redis://single:6379", + }, + {}, + { + "password": "cluster-secret", + "startup_nodes": [{"host": "node-a", "port": 7000}], + "url": "redis://single:6379", + }, + _expected_view( + { + "auth": {"kind": "static", "password": "cluster-secret", "username": None}, + "health_check_interval": 25.0, + "mode": "cluster", + "nodes": ["node-a:7000"], + } + ), + id="startup_nodes_beat_a_url", + ), + pytest.param( + {}, + {"REDIS_CLUSTER_NODES": '[{"host": "node-a", "port": "7000"}, {"host": "node-b", "port": 7001}]'}, + {"startup_nodes": '[{"host": "node-a", "port": "7000"}, {"host": "node-b", "port": 7001}]'}, + _expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000", "node-b:7001"]}), + id="environment_cluster_nodes", + ), + pytest.param( + { + "ssl": "true", + "ssl_ca_certs": "/ca.pem", + "ssl_check_hostname": True, + "startup_nodes": [{"host": "node-a", "port": 7000}], + }, + {}, + { + "ssl": "true", + "ssl_ca_certs": "/ca.pem", + "ssl_check_hostname": True, + "startup_nodes": [{"host": "node-a", "port": 7000}], + }, + _expected_view( + { + "health_check_interval": 25.0, + "mode": "cluster", + "nodes": ["node-a:7000"], + "tls": { + "ca_certs": "/ca.pem", + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": True, + "keyfile": None, + "verify_certificate": True, + }, + } + ), + id="cluster_tls_options", + ), + pytest.param( + { + "password": "data-secret", + "sentinel_nodes": [["sentinel-a", 26379], ["sentinel-b", 26380]], + "sentinel_password": "sentinel-secret", + "service_name": "primary", + }, + {}, + { + "password": "data-secret", + "sentinel_nodes": [["sentinel-a", 26379], ["sentinel-b", 26380]], + "sentinel_password": "sentinel-secret", + "service_name": "primary", + }, + _expected_view( + { + "auth": {"kind": "static", "password": "data-secret", "username": None}, + "db": 0, + "mode": "sentinel", + "nodes": ["sentinel-a:26379", "sentinel-b:26380"], + "sentinel_password": "sentinel-secret", + "service_name": "primary", + } + ), + id="sentinel_from_settings", + ), + pytest.param( + {}, + { + "REDIS_SENTINEL_NODES": '[["sentinel-a", 26379]]', + "REDIS_SENTINEL_PASSWORD": "sentinel-secret", + "REDIS_SERVICE_NAME": "primary", + }, + { + "sentinel_nodes": '[["sentinel-a", 26379]]', + "sentinel_password": "sentinel-secret", + "service_name": "primary", + }, + _expected_view( + { + "db": 0, + "mode": "sentinel", + "nodes": ["sentinel-a:26379"], + "sentinel_password": "sentinel-secret", + "service_name": "primary", + } + ), + id="sentinel_from_the_environment", + ), + pytest.param( + {"host": "ignored", "sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}, + {}, + {"host": "ignored", "sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary"}, + _Fails("TypeError", "redis.asyncio.client.Redis() got multiple values for keyword argument 'host'"), + id="a_host_breaks_sentinel", + ), + pytest.param( + {"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary", "url": "redis://single:6379"}, + {}, + {"sentinel_nodes": [["sentinel-a", 26379]], "service_name": "primary", "url": "redis://single:6379"}, + _expected_view({"mode": "url", "url": "redis://single:6379"}), + id="a_url_beats_sentinels", + ), + pytest.param( + { + "aws_iam_auth": True, + "credential_provider": True, + "gcp_service_account": "account", + "host": "redis", + "password": "dropped", + }, + {"REDIS_AZURE_AD_TOKEN": "true"}, + { + "aws_iam_auth": True, + "credential_provider": True, + "gcp_service_account": "account", + "host": "redis", + "password": "dropped", + }, + _expected_view({"auth": {"kind": "host"}, "db": 0, "host": "redis", "mode": "host", "port": 6379}), + id="an_explicit_credential_provider_beats_cloud_auth", + ), + pytest.param( + {"aws_iam_auth": True, "azure_redis_ad_token": "True", "host": "redis", "password": "dropped"}, + {"REDIS_GCP_SERVICE_ACCOUNT": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com"}, + { + "gcp_service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", + "aws_iam_auth": True, + "azure_redis_ad_token": "True", + "host": "redis", + "password": "dropped", + }, + _expected_view( + { + "auth": { + "kind": "gcp_iam", + "service_account": "projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com", + }, + "db": 0, + "host": "redis", + "mode": "host", + "port": 6379, + } + ), + id="gcp_beats_azure_and_aws", + ), + pytest.param( + { + "gcp_service_account": "account", + "gcp_ssl_ca_certs": "/gcp-ca.pem", + "host": "redis", + "ssl": True, + "ssl_ca_certs": "/configured-ca.pem", + }, + {}, + { + "gcp_service_account": "account", + "gcp_ssl_ca_certs": "/gcp-ca.pem", + "host": "redis", + "ssl": True, + "ssl_ca_certs": "/configured-ca.pem", + }, + _expected_view( + { + "auth": {"kind": "gcp_iam", "service_account": "account"}, + "db": 0, + "host": "redis", + "mode": "host", + "port": 6379, + "tls": { + "ca_certs": "/gcp-ca.pem", + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": True, + }, + } + ), + id="gcp_ca_certs_replace_the_configured_ones_under_ssl", + ), + pytest.param( + { + "aws_iam_auth": True, + "azure_client_id": "client", + "azure_client_secret": "secret", + "azure_redis_ad_token": True, + "azure_tenant_id": "tenant", + "host": "redis", + "username": "dropped", + }, + {"REDIS_USERNAME": "environment-user"}, + { + "username": "dropped", + "aws_iam_auth": True, + "azure_client_id": "client", + "azure_client_secret": "secret", + "azure_redis_ad_token": True, + "azure_tenant_id": "tenant", + "host": "redis", + }, + _expected_view( + { + "auth": {"kind": "azure_ad", "username": "environment-user"}, + "db": 0, + "host": "redis", + "mode": "host", + "port": 6379, + } + ), + id="azure_beats_aws_and_takes_the_environment_username", + ), + pytest.param( + { + "azure_client_id": "client", + "azure_client_secret": "secret", + "azure_redis_ad_token": True, + "azure_tenant_id": "tenant", + "host": "redis", + }, + {}, + { + "azure_client_id": "client", + "azure_client_secret": "secret", + "azure_redis_ad_token": True, + "azure_tenant_id": "tenant", + "host": "redis", + }, + _expected_view( + {"auth": {"kind": "azure_ad", "username": None}, "db": 0, "host": "redis", "mode": "host", "port": 6379} + ), + id="azure_without_a_username_sends_only_the_token", + ), + pytest.param( + { + "azure_client_id": "client", + "azure_client_secret": "secret", + "azure_redis_ad_token": True, + "azure_tenant_id": "tenant", + "host": "redis", + }, + {"REDIS_USERNAME": ""}, + { + "username": "", + "azure_client_id": "client", + "azure_client_secret": "secret", + "azure_redis_ad_token": True, + "azure_tenant_id": "tenant", + "host": "redis", + }, + _expected_view( + {"auth": {"kind": "azure_ad", "username": None}, "db": 0, "host": "redis", "mode": "host", "port": 6379} + ), + id="an_empty_environment_username_sends_only_the_token", + ), + pytest.param( + { + "azure_client_id": "client", + "azure_client_secret": "secret", + "azure_redis_ad_token": "1", + "azure_tenant_id": "tenant", + "host": "redis", + }, + {}, + { + "azure_client_id": "client", + "azure_client_secret": "secret", + "azure_redis_ad_token": "1", + "azure_tenant_id": "tenant", + "host": "redis", + }, + _expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}), + id="azure_needs_the_literal_true", + ), + pytest.param( + { + "aws_iam_auth": "yes", + "aws_iam_cache_name": "cache-1", + "aws_iam_user_name": "iam-user", + "host": "cache-1", + "password": "dropped", + "ssl": True, + }, + {"AWS_DEFAULT_REGION": "us-east-1"}, + { + "aws_iam_auth": "yes", + "aws_iam_cache_name": "cache-1", + "aws_iam_user_name": "iam-user", + "host": "cache-1", + "password": "dropped", + "ssl": True, + "aws_iam_region": "us-east-1", + }, + _expected_view( + { + "auth": { + "cache_name": "cache-1", + "kind": "aws_iam", + "region": "us-east-1", + "serverless": False, + "user_name": "iam-user", + }, + "db": 0, + "host": "cache-1", + "mode": "host", + "port": 6379, + "tls": { + "ca_certs": None, + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": True, + }, + } + ), + id="aws_iam_over_tls_with_the_default_region", + ), + pytest.param( + { + "aws_iam_auth": True, + "aws_iam_cache_name": "cache-1", + "aws_iam_user_name": "iam-user", + "host": "cache-1", + "ssl": True, + }, + {"AWS_DEFAULT_REGION": "us-east-1", "AWS_REGION": "eu-central-1"}, + { + "aws_iam_auth": True, + "aws_iam_cache_name": "cache-1", + "aws_iam_user_name": "iam-user", + "host": "cache-1", + "ssl": True, + "aws_iam_region": "eu-central-1", + }, + _expected_view( + { + "auth": { + "cache_name": "cache-1", + "kind": "aws_iam", + "region": "eu-central-1", + "serverless": False, + "user_name": "iam-user", + }, + "db": 0, + "host": "cache-1", + "mode": "host", + "port": 6379, + "tls": { + "ca_certs": None, + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": True, + }, + } + ), + id="aws_region_beats_the_default_region", + ), + pytest.param( + { + "aws_iam_auth": True, + "aws_iam_cache_name": "cache-1", + "aws_iam_region": "eu-west-1", + "aws_iam_serverless": True, + "aws_iam_user_name": "iam-user", + "url": "rediss://user:pass@cache-1:6379", + }, + {}, + { + "aws_iam_auth": True, + "aws_iam_cache_name": "cache-1", + "aws_iam_region": "eu-west-1", + "aws_iam_serverless": True, + "aws_iam_user_name": "iam-user", + "url": "rediss://user:pass@cache-1:6379", + }, + _expected_view( + { + "auth": { + "cache_name": "cache-1", + "kind": "aws_iam", + "region": "eu-west-1", + "serverless": True, + "user_name": "iam-user", + }, + "mode": "url", + "tls": { + "ca_certs": None, + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": True, + }, + "url": "rediss://cache-1:6379", + } + ), + id="aws_iam_serverless_over_a_rediss_url_drops_its_userinfo", + ), + pytest.param( + { + "aws_iam_auth": True, + "aws_iam_cache_name": "cache-1", + "aws_iam_region": "us-east-1", + "aws_iam_user_name": "iam-user", + "host": "cache-1", + }, + {}, + { + "aws_iam_auth": True, + "aws_iam_cache_name": "cache-1", + "aws_iam_region": "us-east-1", + "aws_iam_user_name": "iam-user", + "host": "cache-1", + }, + _Fails("ValueError", "AWS ElastiCache IAM Redis authentication requires TLS"), + id="aws_iam_without_tls_fails", + ), + pytest.param( + { + "aws_iam_auth": True, + "aws_iam_cache_name": "cache-1", + "aws_iam_user_name": "iam-user", + "host": "cache-1", + "ssl": True, + }, + {}, + { + "aws_iam_auth": True, + "aws_iam_cache_name": "cache-1", + "aws_iam_user_name": "iam-user", + "host": "cache-1", + "ssl": True, + }, + _Fails("ValueError", "AWS ElastiCache IAM Redis authentication requires: aws_iam_region"), + id="aws_iam_without_a_region_fails", + ), + pytest.param( + { + "host": "redis", + "ssl": True, + "ssl_ca_certs": "/ca.pem", + "ssl_cert_reqs": "none", + "ssl_check_hostname": True, + }, + {}, + { + "host": "redis", + "ssl": True, + "ssl_ca_certs": "/ca.pem", + "ssl_cert_reqs": "none", + "ssl_check_hostname": True, + }, + _expected_view( + { + "db": 0, + "host": "redis", + "mode": "host", + "port": 6379, + "tls": { + "ca_certs": "/ca.pem", + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": True, + "keyfile": None, + "verify_certificate": False, + }, + } + ), + id="tls_options_from_settings", + ), + pytest.param( + {"host": "redis", "ssl": True, "ssl_cert_reqs": None}, + {"REDIS_SSL_CERT_REQS": "required"}, + {"host": "redis", "ssl": True, "ssl_cert_reqs": None}, + _expected_view( + { + "db": 0, + "host": "redis", + "mode": "host", + "port": 6379, + "tls": { + "ca_certs": None, + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": False, + }, + } + ), + id="a_null_cert_requirement_skips_verification", + ), + pytest.param( + {"host": "redis", "ssl": True, "ssl_cert_reqs": ssl.CERT_NONE}, + {}, + {"host": "redis", "ssl": True, "ssl_cert_reqs": ssl.CERT_NONE}, + { + "rejected": { + "message": "'RedisSSLContext' object has no attribute 'cert_reqs'", + "type": "AttributeError", + } + }, + id="a_verify_mode_of_none_skips_verification", + ), + pytest.param( + {"host": "redis", "ssl": True, "ssl_cert_reqs": ssl.CERT_OPTIONAL}, + {}, + {"host": "redis", "ssl": True, "ssl_cert_reqs": ssl.CERT_OPTIONAL}, + { + "rejected": { + "message": "'RedisSSLContext' object has no attribute 'cert_reqs'", + "type": "AttributeError", + } + }, + id="a_verify_mode_of_optional_verifies", + ), + pytest.param( + {"host": "redis"}, + {"REDIS_SSL": "true", "REDIS_SSL_CERT_REQS": "None"}, + {"ssl": True, "ssl_cert_reqs": "None", "host": "redis"}, + _Fails("RedisError", "Invalid SSL Certificate Requirements Flag: None"), + id="a_cert_requirement_spelled_like_python_none_fails", + ), + pytest.param( + {"host": "redis"}, + {"REDIS_SSL": "true", "REDIS_SSL_CA_CERTS": "/env-ca.pem", "REDIS_SSL_CERT_REQS": "optional"}, + {"ssl_ca_certs": "/env-ca.pem", "ssl": True, "ssl_cert_reqs": "optional", "host": "redis"}, + _expected_view( + { + "db": 0, + "host": "redis", + "mode": "host", + "port": 6379, + "tls": { + "ca_certs": "/env-ca.pem", + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": True, + }, + } + ), + id="tls_options_from_the_environment", + ), + pytest.param( + {"host": "redis", "socket_connect_timeout": 0.25, "socket_timeout": 2.5}, + {}, + {"host": "redis", "socket_connect_timeout": 0.25, "socket_timeout": 2.5}, + _expected_view( + { + "db": 0, + "host": "redis", + "mode": "host", + "port": 6379, + "socket_connect_timeout": 0.25, + "socket_timeout": 2.5, + } + ), + id="connect_timeout_is_separate_from_the_socket_timeout", + ), + pytest.param( + {"host": "os.environ/COORD_HOST", "password": "os.environ/COORD_PASSWORD"}, + {"COORD_HOST": "referenced", "COORD_PASSWORD": "referenced-secret"}, + {"host": "referenced", "password": "referenced-secret"}, + _expected_view( + { + "auth": {"kind": "static", "password": "referenced-secret", "username": None}, + "db": 0, + "host": "referenced", + "mode": "host", + "port": 6379, + } + ), + id="environment_references_resolve", + ), + pytest.param( + {"port": 6379}, + {}, + {"port": 6379}, + _Fails("ValueError", "Either 'host' or 'url' must be specified for redis."), + id="no_target_fails", + ), + pytest.param( + {"host": "redis", "ssl": True, "ssl_certfile": "/client.pem", "ssl_keyfile": "/client.key"}, + {}, + {"host": "redis", "ssl": True, "ssl_certfile": "/client.pem", "ssl_keyfile": "/client.key"}, + _expected_view( + { + "db": 0, + "host": "redis", + "mode": "host", + "port": 6379, + "tls": { + "ca_certs": None, + "ca_data": None, + "ca_path": None, + "certfile": "/client.pem", + "check_hostname": False, + "keyfile": "/client.key", + "verify_certificate": True, + }, + } + ), + id="client_certificates", + ), + pytest.param( + {"url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes"}, + {}, + {"url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes"}, + _expected_view( + { + "mode": "url", + "tls": { + "ca_certs": "/ca.pem", + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": True, + "keyfile": None, + "verify_certificate": False, + }, + "url": "rediss://cache:6380/0?ssl_cert_reqs=none&ssl_ca_certs=/ca.pem&ssl_check_hostname=yes", + } + ), + id="a_rediss_url_takes_tls_options_from_its_query", + ), + pytest.param( + {"url": "redis://cache:6379/0?socket_timeout=2"}, + {}, + {"url": "redis://cache:6379/0?socket_timeout=2"}, + _expected_view( + { + "mode": "url", + "socket_connect_timeout": 2.0, + "socket_timeout": 2.0, + "url": "redis://cache:6379/0?socket_timeout=2", + } + ), + id="a_url_query_carries_socket_timeout", + ), + pytest.param( + { + "health_check_interval": 4, + "socket_connect_timeout": 7, + "socket_timeout": 1, + "url": "redis://cache:6379/0?socket_connect_timeout=0.5&health_check_interval=9&socket_timeout=3", + }, + {}, + { + "health_check_interval": 4, + "socket_connect_timeout": 7, + "socket_timeout": 1, + "url": "redis://cache:6379/0?socket_connect_timeout=0.5&health_check_interval=9&socket_timeout=3", + }, + _expected_view( + { + "health_check_interval": 9.0, + "mode": "url", + "socket_connect_timeout": 0.5, + "socket_timeout": 3.0, + "url": "redis://cache:6379/0?socket_connect_timeout=0.5&health_check_interval=9&socket_timeout=3", + } + ), + id="url_query_settings_win_over_the_kwargs", + ), + pytest.param( + {"url": "redis://cache:6379/0?socket_timeout=2&socket_timeout=9"}, + {}, + {"url": "redis://cache:6379/0?socket_timeout=2&socket_timeout=9"}, + _expected_view( + { + "mode": "url", + "socket_connect_timeout": 2.0, + "socket_timeout": 2.0, + "url": "redis://cache:6379/0?socket_timeout=2&socket_timeout=9", + } + ), + id="a_repeated_url_query_parameter_takes_its_first_value", + ), + pytest.param( + {"url": "redis://cache:6379/0?socket_timeout=slow"}, + {}, + {"url": "redis://cache:6379/0?socket_timeout=slow"}, + _Fails("ValueError", "Invalid value for 'socket_timeout' in connection URL."), + id="a_url_query_number_parse_url_cannot_read", + ), + pytest.param( + {"db": "99999999999999999999", "ssl_cert_reqs": "strict", "url": "rediss://cache:6380"}, + {}, + {"db": "99999999999999999999", "ssl_cert_reqs": "strict", "url": "rediss://cache:6380"}, + _expected_view( + { + "mode": "url", + "tls": { + "ca_certs": None, + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": True, + }, + "url": "rediss://cache:6380", + } + ), + id="tls_settings_outside_a_url_never_fail_it", + ), + pytest.param( + {"db": 99999999999, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + {}, + {"db": 99999999999, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + _Fails("RedisClusterException", "Argument 'db' must be 0 or None in cluster mode"), + id="db_is_not_read_for_a_cluster", + ), + pytest.param( + {"db": True, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + {}, + {"db": True, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + _Fails("RedisClusterException", "Argument 'db' must be 0 or None in cluster mode"), + id="cluster_db_true", + ), + pytest.param( + {"db": False, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + {}, + {"db": False, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + _expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"]}), + id="cluster_db_false", + ), + pytest.param( + {"db": 0.5, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + {}, + {"db": 0.5, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + _Fails("RedisClusterException", "Argument 'db' must be 0 or None in cluster mode"), + id="cluster_db_fraction", + ), + pytest.param( + {"db": 0.0, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + {}, + {"db": 0.0, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + _expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"]}), + id="cluster_db_float_zero", + ), + pytest.param( + {"db": [1], "startup_nodes": [{"host": "node-a", "port": 7000}]}, + {}, + {"db": [1], "startup_nodes": [{"host": "node-a", "port": 7000}]}, + _Fails("RedisClusterException", "Argument 'db' must be 0 or None in cluster mode"), + id="cluster_db_nonempty_list", + ), + pytest.param( + {"db": [], "startup_nodes": [{"host": "node-a", "port": 7000}]}, + {}, + {"db": [], "startup_nodes": [{"host": "node-a", "port": 7000}]}, + _expected_view({"health_check_interval": 25.0, "mode": "cluster", "nodes": ["node-a:7000"]}), + id="cluster_db_empty_list", + ), + pytest.param( + {"host": "redis", "ssl": True, "ssl_ciphers": "HIGH"}, + {}, + {"host": "redis", "ssl": True, "ssl_ciphers": "HIGH"}, + _expected_view( + { + "db": 0, + "host": "redis", + "mode": "host", + "port": 6379, + "tls": { + "ca_certs": None, + "ca_data": None, + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": True, + }, + } + ), + id="ssl_ciphers", + ), + pytest.param( + {"host": "redis", "retry_on_timeout": True}, + {}, + {"host": "redis", "retry_on_timeout": True}, + _expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}), + id="retry_on_timeout", + ), + pytest.param( + {"host": "redis", "port": "abc"}, + {}, + {"host": "redis", "port": "abc"}, + _expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}), + id="malformed_port", + ), + pytest.param( + {"db": "first", "host": "redis"}, + {}, + {"db": "first", "host": "redis"}, + _expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}), + id="malformed_db", + ), + pytest.param( + {"host": "redis", "socket_timeout": "slow"}, + {}, + {"host": "redis", "socket_timeout": "slow"}, + _expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}), + id="socket_timeout_not_a_number", + ), + pytest.param( + {}, + {"REDIS_HOST": "redis", "REDIS_PORT": "abc"}, + {"port": "abc", "host": "redis"}, + _expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}), + id="malformed_port_from_the_environment", + ), + pytest.param( + {"host": "redis", "protocol": 3}, + {}, + {"host": "redis", "protocol": 3}, + _expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}), + id="protocol_3", + ), + pytest.param( + {"host": "redis"}, + {"REDIS_RETRY_ON_ERROR": "true"}, + {"retry_on_error": True, "host": "redis"}, + {"rejected": {"message": "'bool' object is not iterable", "type": "TypeError"}}, + id="retry_on_error_from_the_environment", + ), + pytest.param( + {"host": "redis", "unix_socket_path": "/tmp/redis.sock"}, + {}, + {"host": "redis", "unix_socket_path": "/tmp/redis.sock"}, + _expected_view({"mode": "unix", "path": "/tmp/redis.sock"}), + id="unix_socket_path", + ), + pytest.param( + {"host": "redis", "ssl": True, "ssl_ca_certs": "/ca.pem", "unix_socket_path": "/tmp/redis.sock"}, + {}, + {"host": "redis", "ssl": True, "ssl_ca_certs": "/ca.pem", "unix_socket_path": "/tmp/redis.sock"}, + _expected_view({"mode": "unix", "path": "/tmp/redis.sock"}), + id="unix_socket_path_keeps_tls_settings_out_of_the_target", + ), + pytest.param( + { + "host": "redis", + "ssl": True, + "ssl_ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----", + }, + {}, + { + "host": "redis", + "ssl": True, + "ssl_ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----", + }, + _expected_view( + { + "db": 0, + "host": "redis", + "mode": "host", + "port": 6379, + "tls": { + "ca_certs": None, + "ca_data": "-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----", + "ca_path": None, + "certfile": None, + "check_hostname": False, + "keyfile": None, + "verify_certificate": True, + }, + } + ), + id="ssl_ca_data", + ), + pytest.param( + {"host": "redis"}, + {"REDIS_SSL": "true", "REDIS_SSL_CA_PATH": "/etc/ssl/private-cas"}, + {"ssl": True, "ssl_ca_path": "/etc/ssl/private-cas", "host": "redis"}, + { + "rejected": { + "message": "Redis.__init__() got an unexpected keyword argument 'ssl_ca_path'", + "type": "TypeError", + } + }, + id="ssl_ca_path_from_the_environment", + ), + pytest.param( + {"health_check_interval": "7", "host": "redis"}, + {}, + {"health_check_interval": "7", "host": "redis"}, + _expected_view({"db": 0, "health_check_interval": 7.0, "host": "redis", "mode": "host", "port": 6379}), + id="health_check_interval_is_whole_seconds", + ), + pytest.param( + {"health_check_interval": 0, "host": "redis"}, + {}, + {"health_check_interval": 0, "host": "redis"}, + _expected_view({"db": 0, "host": "redis", "mode": "host", "port": 6379}), + id="health_check_interval_zero_is_off", + ), + pytest.param( + {"health_check_interval": 0, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + {}, + {"health_check_interval": 0, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + _expected_view({"mode": "cluster", "nodes": ["node-a:7000"]}), + id="health_check_interval_zero_keeps_a_cluster_off", + ), + pytest.param( + {"db": 3, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + {}, + {"db": 3, "startup_nodes": [{"host": "node-a", "port": 7000}]}, + _Fails("RedisClusterException", "Argument 'db' must be 0 or None in cluster mode"), + id="a_cluster_rejects_a_nonzero_db", + ), + pytest.param( + {"unix_socket_path": "/tmp/redis.sock"}, + {}, + {"unix_socket_path": "/tmp/redis.sock"}, + _Fails("ValueError", "Either 'host' or 'url' must be specified for redis."), + id="unix_socket_path_without_a_host_fails_like_python", + ), + ], +) +def test_redis_settings_resolve_like_the_rust_resolver_expects(monkeypatch, kwargs, environment, merged, expected): + for name in [name for name in os.environ if name.startswith(("REDIS_", "AZURE_", "AWS_"))]: + monkeypatch.delenv(name) + for name, value in environment.items(): + monkeypatch.setenv(name, value) + + assert _merge_redis_kwargs(kwargs) == merged + try: + outcome = _resolution_view(kwargs) + except (ValueError, TypeError, AttributeError, RedisError, RedisClusterException) as error: + raised: Final = {"type": type(error).__name__, "message": str(error)} + outcome = _Fails(**raised) if isinstance(expected, _Fails) else {"rejected": raised} + + assert outcome == expected + + +_REDIS_RESOLVE: Final = Path(__file__).parents[2] / "litellm-rust/crates/redis-resolve" +_REDIS_PY_KWARGS: Final = json.loads((_REDIS_RESOLVE / "redis-py-kwargs.json").read_text()) + + +def _catalog_script() -> ModuleType: + spec: Final = importlib.util.spec_from_file_location("redis_py_kwargs", _REDIS_RESOLVE / "redis_py_kwargs.py") + assert spec is not None and spec.loader is not None + module: Final = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_redis_py_kwargs_catalog_matches_the_installed_redis_py(): + assert redis.__version__ in _REDIS_PY_KWARGS["meta"]["redis_py_versions"] + catalogued: Final = { + name: arg["redis_py"][redis.__version__] + for name, arg in _REDIS_PY_KWARGS["args"].items() + if redis.__version__ in arg["redis_py"] + } + + assert catalogued == _catalog_script().arg_signatures() + + +@pytest.mark.skipif( + redis.__version__ != _REDIS_PY_KWARGS["meta"]["redis_py_locked"], + reason="the catalog records REDIS_* names for the locked redis-py only", +) +def test_redis_py_kwargs_catalog_names_the_environment_kwargs_python_reads(): + catalogued: Final = frozenset(name for name, arg in _REDIS_PY_KWARGS["args"].items() if arg["environment"]) + + assert catalogued == frozenset(_get_redis_env_kwarg_mapping().values())