From 839a5ba429b61bed65ea6ceb4c0c7a44da9603e6 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 2 Nov 2026 00:18:01 +0000 Subject: [PATCH] feat(redis): authenticate host credential providers through Rust --- litellm-rust/crates/gateway/src/redis.rs | 1 + .../src/redis/client/credentials.rs | 47 ++++++++++++++ .../python-bridge/src/redis/client/mod.rs | 65 ++++++++++++++++--- .../crates/redis-connect/src/connect.rs | 8 ++- .../crates/redis-connect/src/error.rs | 5 +- litellm-rust/crates/redis-connect/src/lib.rs | 3 +- .../crates/redis-connect/src/refresh.rs | 6 ++ .../crates/redis-connect/tests/connect.rs | 59 ++++++++++++++++- .../crates/redis-connect/tests/live/main.rs | 1 + litellm/rust_bridge/host/AGENTS.md | 2 +- litellm/rust_bridge/host/redis_credentials.py | 19 ++++++ .../host/test_redis_credentials.py | 38 +++++++++++ 12 files changed, 236 insertions(+), 18 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/redis/client/credentials.rs create mode 100644 litellm/rust_bridge/host/redis_credentials.py create mode 100644 tests/unit/rust_bridge/host/test_redis_credentials.py diff --git a/litellm-rust/crates/gateway/src/redis.rs b/litellm-rust/crates/gateway/src/redis.rs index e7b4f6ea119..e4db0df5063 100644 --- a/litellm-rust/crates/gateway/src/redis.rs +++ b/litellm-rust/crates/gateway/src/redis.rs @@ -82,6 +82,7 @@ pub(crate) async fn connect( breaker, http: resources.http.clone(), environment: resources.lookup.clone(), + credential_provider: None, }, )?) }; 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 index 9c4618af1a5..a0cbf694df4 100644 --- a/litellm-rust/crates/python-bridge/src/redis/client/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/client/mod.rs @@ -1,3 +1,5 @@ +mod credentials; + use std::{ collections::{BTreeMap, BTreeSet}, sync::{Arc, LazyLock, Mutex, PoisonError}, @@ -5,8 +7,9 @@ use std::{ use litellm_config::{ConfiguredRedis, Resolution, SecretSnapshot, Value}; use litellm_core_utils::settings::ProcessEnvironment; +use litellm_host_python::PythonContext; use litellm_redis_client::{RedisClient, RedisEvents}; -use litellm_redis_connect::{HostServices, SharedEnvironment}; +use litellm_redis_connect::{CredentialSource, HostServices, SharedEnvironment}; use litellm_redis_types::ResolvedRedis; use pyo3::{ prelude::*, @@ -14,29 +17,42 @@ use pyo3::{ types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple}, }; +use self::credentials::PythonCredentialSource; use crate::redis::events::{ServiceLoggingEvents, state_label}; /// The `redis_kwargs` dict `RedisCache` hands redis-py, read the way `litellm/_redis.py` reads it. /// A `None` stays an explicit null that hides its `REDIS_*` value, as in redis-py. A value that is not -/// plain data, such as a `credential_provider` object, reaches the resolver by its type name, so the -/// resolver or the connector names the setting it cannot honor. -pub(crate) struct PythonRedisSettings(ConfiguredRedis); +/// plain data reaches the resolver by its type name, so the resolver or the connector names the +/// setting it cannot honor. A `credential_provider` object is also kept, so the connector can ask it +/// for credentials. +pub(crate) struct PythonRedisSettings { + settings: ConfiguredRedis, + credential_provider: Option>, +} impl<'a, 'py> FromPyObject<'a, 'py> for PythonRedisSettings { type Error = PyErr; fn extract(kwargs: Borrowed<'a, 'py, PyAny>) -> PyResult { let py = kwargs.py(); + let kwargs = kwargs.cast::()?; let plain = PyDict::new(py); - for (name, value) in kwargs.cast::()?.iter() { + for (name, value) in kwargs.iter() { if value.is_none() || is_plain(&value) { plain.set_item(&name, value)?; } else { plain.set_item(&name, value.get_type().name()?)?; } } + let credential_provider = kwargs + .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())?; - Ok(Self(ConfiguredRedis::redis_cache_kwargs(fields))) + Ok(Self { + settings: ConfiguredRedis::redis_cache_kwargs(fields), + credential_provider, + }) } } @@ -52,12 +68,24 @@ fn is_plain(value: &Bound<'_, PyAny>) -> bool { struct Shared { resolved: ResolvedRedis, + credential_provider: Option>, client: RedisClient, } +impl Shared { + fn serves(&self, resolved: &ResolvedRedis, credential_provider: Option<&Py>) -> bool { + let same_provider = match (&self.credential_provider, credential_provider) { + (None, None) => true, + (Some(shared), Some(provider)) => shared.is(provider), + _ => false, + }; + same_provider && self.resolved == *resolved + } +} + static CLIENTS: LazyLock>> = LazyLock::new(|| Mutex::new(Vec::new())); -/// The one client for a resolved Redis config, so every family in this process shares its pool +/// The one client for a resolved Redis config and credential provider, so every family in this process shares its pool /// and its breaker. Like redis-py, which never connects at construction, a config the resolver or /// the connector cannot honor leaves the family on memory with a warning rather than failing the /// `Router`. @@ -67,12 +95,16 @@ fn shared_client( namespace: Option, ) -> PyResult> { let environment: SharedEnvironment = Arc::new(ProcessEnvironment); - let secrets = match secret_snapshot(py, kwargs.0.secret_names())? { + let PythonRedisSettings { + settings, + credential_provider, + } = kwargs; + let secrets = match secret_snapshot(py, settings.secret_names())? { Ok(secrets) => secrets, Err(error) => return Ok(memory_only(py, &error)), }; let logger = crate::logger::capture(py); - let resolved = match logger.scope(|| kwargs.0.resolve(&secrets, &*environment)) { + let resolved = match logger.scope(|| settings.resolve(&secrets, &*environment)) { Ok(Resolution { redis, ignored }) => { logger.scope(|| { for ignored in ignored { @@ -98,11 +130,23 @@ fn shared_client( breaker, 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 = CLIENTS .lock_py_attached(py) .unwrap_or_else(PoisonError::into_inner); - if let Some(shared) = clients.iter().find(|shared| shared.resolved == resolved) { + if let Some(shared) = clients + .iter() + .find(|shared| shared.serves(&resolved, credential_provider.as_ref())) + { return Ok(Some(shared.client.clone())); } let connected = { @@ -118,6 +162,7 @@ fn shared_client( }; clients.push(Shared { resolved, + credential_provider, client: client.clone(), }); Ok(Some(client)) diff --git a/litellm-rust/crates/redis-connect/src/connect.rs b/litellm-rust/crates/redis-connect/src/connect.rs index b8611d309c2..624c5018ef0 100644 --- a/litellm-rust/crates/redis-connect/src/connect.rs +++ b/litellm-rust/crates/redis-connect/src/connect.rs @@ -12,7 +12,7 @@ use litellm_redis_client::{ use litellm_redis_types::{RedisAuth, ResolvedRedis}; use crate::{ - Error, RefreshingCredentials, + CredentialSource, Error, RefreshingCredentials, sources::{AzureAd, ElastiCacheIam, GcpIam}, }; @@ -23,6 +23,7 @@ pub struct HostServices { pub breaker: 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 @@ -73,7 +74,10 @@ fn credentials(auth: RedisAuth, host: &HostServices) -> Result Credentials::Static { username, password }, - RedisAuth::Host => return Err(Error::NoHostCredentialProvider), + RedisAuth::Host => match &host.credential_provider { + Some(source) => refreshing(RefreshingCredentials::new(Arc::clone(source))), + None => return Err(Error::NoHostCredentialProvider), + }, RedisAuth::AzureAd(auth) => refreshing(RefreshingCredentials::new(AzureAd::new( auth, host.environment.clone(), diff --git a/litellm-rust/crates/redis-connect/src/error.rs b/litellm-rust/crates/redis-connect/src/error.rs index 07f5178cbd2..8a85ee6be98 100644 --- a/litellm-rust/crates/redis-connect/src/error.rs +++ b/litellm-rust/crates/redis-connect/src/error.rs @@ -42,10 +42,12 @@ pub enum CredentialStage { Exchange, #[strum(serialize = "Redis token response")] Response, + #[strum(serialize = "Redis credential_provider")] + Provider, } impl Error { - pub(crate) fn credential( + pub fn credential( stage: CredentialStage, ) -> impl FnOnce(E) -> Self { move |cause| Self::Credential { @@ -88,6 +90,7 @@ impl Error { CredentialStage::Signing => "token-signing", CredentialStage::Exchange => "token-exchange", CredentialStage::Response => "token-response", + CredentialStage::Provider => "credential-provider", }, Self::CredentialTimeout { .. } => "acquisition-timeout", Self::ExpiredCredential => "expired-credential", diff --git a/litellm-rust/crates/redis-connect/src/lib.rs b/litellm-rust/crates/redis-connect/src/lib.rs index 4e82af2db3b..044177a74c2 100644 --- a/litellm-rust/crates/redis-connect/src/lib.rs +++ b/litellm-rust/crates/redis-connect/src/lib.rs @@ -7,4 +7,5 @@ mod sources; pub use connect::{HostServices, SharedEnvironment, breaker_config, connect}; pub use error::{CredentialStage, Error}; pub use events::TracingEvents; -pub(crate) use refresh::{CredentialSource, Expiring, RefreshingCredentials, SourceFuture}; +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 index 852ba412b52..9423dcdbae7 100644 --- a/litellm-rust/crates/redis-connect/src/refresh.rs +++ b/litellm-rust/crates/redis-connect/src/refresh.rs @@ -32,6 +32,12 @@ 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, diff --git a/litellm-rust/crates/redis-connect/tests/connect.rs b/litellm-rust/crates/redis-connect/tests/connect.rs index db2f9fafa4b..cc30ae5b2c9 100644 --- a/litellm-rust/crates/redis-connect/tests/connect.rs +++ b/litellm-rust/crates/redis-connect/tests/connect.rs @@ -1,11 +1,14 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use litellm_config::{RedisSettings, SecretSnapshot}; -use litellm_redis_connect::{Error, HostServices, connect}; +use litellm_redis_connect::{ + CredentialSource, Error, Expiring, HostServices, SourceFuture, connect, +}; use litellm_redis_testing::{ - RedisDeployment, RedisSentinel, RedisTlsServer, SERVICE_NAME, Topology, eventually, + RedisDeployment, RedisSentinel, RedisServer, RedisTlsServer, SERVICE_NAME, Topology, eventually, }; use litellm_redis_types::ResolvedRedis; +use litellm_secret_value::SecretValue; use rstest::rstest; fn resolved(kwargs: serde_json::Value, environment: &[(&str, &str)]) -> ResolvedRedis { @@ -28,6 +31,21 @@ fn services() -> HostServices { breaker: 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, + }) + }) } } @@ -135,7 +153,7 @@ async fn python_sentinel_kwargs_follow_the_master_through_a_failover() { } #[rstest] -fn a_python_credential_provider_fails_at_boot() { +fn a_python_credential_provider_without_a_host_source_fails_at_boot() { let error = connect( resolved( serde_json::json!({"host": "redis", "credential_provider": true}), @@ -181,3 +199,38 @@ fn config_only_aws_identities_fail_before_connecting(#[case] setting: &'static s 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() + }, + ) + .unwrap() + .ping() + .await; + + assert_eq!(pinged.is_ok(), connects, "{pinged:?}"); +} diff --git a/litellm-rust/crates/redis-connect/tests/live/main.rs b/litellm-rust/crates/redis-connect/tests/live/main.rs index bd142c0d005..80138e21706 100644 --- a/litellm-rust/crates/redis-connect/tests/live/main.rs +++ b/litellm-rust/crates/redis-connect/tests/live/main.rs @@ -51,6 +51,7 @@ async fn roundtrip(config: LiveConfig) { breaker: None, http: http.clone(), environment: Arc::new(|name: &str| std::env::var(name).ok()), + credential_provider: None, }; let client = connect(resolved.clone(), services()).expect("build live Redis client"); let marker = std::time::SystemTime::now() diff --git a/litellm/rust_bridge/host/AGENTS.md b/litellm/rust_bridge/host/AGENTS.md index d4cb25cb4fd..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), `redis_events` (reporting native Redis calls and breaker transitions to `ServiceLogging`) +- 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/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