diff --git a/litellm-rust/crates/gateway/src/redis/mod.rs b/litellm-rust/crates/gateway/src/redis/mod.rs index 3a091dd7bed..5bf8ea96250 100644 --- a/litellm-rust/crates/gateway/src/redis/mod.rs +++ b/litellm-rust/crates/gateway/src/redis/mod.rs @@ -87,6 +87,7 @@ pub(crate) async fn connect( events: Some(events.clone()), http: resources.http.clone(), environment: resources.lookup.clone(), + credential_provider: None, }, breaker, )?) 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 f709a688bcb..c4f56e01ed7 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, sync::{Arc, LazyLock, Mutex, PoisonError}, @@ -6,22 +8,28 @@ use std::{ use litellm_circuit_breaker::Config as BreakerConfig; use litellm_core_utils::settings::ProcessEnvironment; -use litellm_redis_client::{RedisClient, RedisEvents}; -use litellm_redis_connect::{HostServices, SharedEnvironment}; +use litellm_host_python::PythonContext; +use litellm_redis_client::{RedisClient, RedisEvents, WeakRedisClient}; +use litellm_redis_connect::{CredentialSource, HostServices, SharedEnvironment}; use litellm_redis_resolve::Kwargs; -use litellm_redis_types::Resolution; +use litellm_redis_types::{ProviderHandle, Resolution}; use pyo3::{ prelude::*, sync::MutexExt, types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple}, }; +use self::credentials::PythonCredentialSource; use crate::redis::events::{ServiceLoggingEvents, state_label}; /// 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. -pub(crate) struct RedisSpec(Kwargs); +/// 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; @@ -36,12 +44,20 @@ impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec { 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(litellm_redis_resolve::kwargs(fields.into_iter().chain(objects)))) + Ok(Self { + kwargs: litellm_redis_resolve::kwargs(fields.into_iter().chain(objects)), + credential_provider, + }) } } @@ -76,14 +92,23 @@ fn is_plain(value: &Bound<'_, PyAny>) -> bool { || 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: RedisClient, + client: WeakRedisClient, } 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 +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. 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`. @@ -94,8 +119,16 @@ fn shared_client( 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 resolved = match logger.scope(|| litellm_redis_resolve::interpret(&spec.0, None)) { + let resolved = match logger.scope(|| litellm_redis_resolve::interpret(&kwargs, handle)) { Ok(resolution) => { logger.scope(|| { for ignored in &resolution.ignored { @@ -119,12 +152,23 @@ fn shared_client( 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 = CLIENTS - .lock_py_attached(py) - .unwrap_or_else(PoisonError::into_inner); - if let Some(shared) = clients.iter().find(|shared| shared.resolved == resolved) { - return Ok(Some(shared.client.clone())); + 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(); @@ -143,7 +187,7 @@ fn shared_client( }; clients.push(Shared { resolved, - client: client.clone(), + client: client.downgrade(), }); Ok(Some(client)) } @@ -178,12 +222,10 @@ impl NativeRedisClient { /// gauge counts them next to the Python `RedisCache` breakers. #[pyfunction] pub(crate) fn native_redis_breaker_states(py: Python<'_>) -> Vec<&'static str> { - let clients = CLIENTS - .lock_py_attached(py) - .unwrap_or_else(PoisonError::into_inner); - clients + live_clients(py) .iter() - .filter_map(|shared| shared.client.breaker_state()) + .filter_map(|shared| shared.client.upgrade()) + .filter_map(|client| client.breaker_state()) .map(|state| state_label(&state)) .collect() } diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index c63d2121081..d615f616d19 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -1,6 +1,6 @@ use std::{ future::Future, - sync::{Arc, Mutex, PoisonError}, + sync::{Arc, Mutex, PoisonError, Weak}, time::{Duration, Instant}, }; @@ -14,15 +14,28 @@ use crate::{ 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: Arc>, + last_used: Mutex, } #[derive(Debug, PartialEq, Eq)] @@ -35,7 +48,7 @@ pub enum Ttl { impl RedisClient { pub fn new(config: RedisConfig) -> Result { let connection = Connection::new(&config.target, &config.credentials, config.timeouts)?; - Ok(Self { + Ok(Self(Arc::new(Shared { subscriber: Subscriber::new( &config.target, &config.credentials, @@ -47,13 +60,17 @@ impl RedisClient { breaker: config.breaker, events: config.events, health_check: config.health_check, - last_used: Arc::new(Mutex::new(Instant::now())), - }) + 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.breaker.as_ref().map(Breaker::state) + self.0.breaker.as_ref().map(Breaker::state) } pub async fn ping(&self) -> Result<(), Error> { @@ -119,7 +136,7 @@ impl RedisClient { pub async fn incr_by(&self, key: &str, delta: i64, ttl: Duration) -> Result { let key = self.key(key); self.guarded(Operation::Command("INCRBY"), async { - let mut connection = self.connection.clone(); + let mut connection = self.0.connection.clone(); let value: i64 = commands::incr_by(key.clone(), delta) .query_async(&mut connection) .await?; @@ -159,7 +176,7 @@ impl RedisClient { ) -> Result { let key = self.key(key); self.guarded(Operation::Command("INCRBYFLOAT"), async { - let mut connection = self.connection.clone(); + let mut connection = self.0.connection.clone(); let mut pipeline = redis::pipe(); pipeline.add_command(commands::incr_by_float(key.clone(), amount)); match expiry { @@ -198,7 +215,7 @@ impl RedisClient { self.guarded(Operation::Script, async { script .invocation(&keys, &args) - .invoke_async(&mut self.connection.clone()) + .invoke_async(&mut self.0.connection.clone()) .await }) .await @@ -213,19 +230,19 @@ impl RedisClient { /// 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.subscriber.clone().messages(channel.to_owned()) + 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 fn same_client(&self, other: &Self) -> bool { - Arc::ptr_eq(&self.last_used, &other.last_used) + Arc::ptr_eq(&self.0, &other.0) } pub fn is_cluster(&self) -> bool { - matches!(self.connection, Connection::Cluster(_)) + matches!(self.0.connection, Connection::Cluster(_)) } pub fn batch(&self) -> RedisBatch { @@ -251,7 +268,7 @@ impl RedisClient { } }; self.check_health().await; - let results = pipeline.run(&self.connection).await; + let results = pipeline.run(&self.0.connection).await; let transport = results.iter().find_map(|result| { result .as_ref() @@ -268,7 +285,7 @@ impl RedisClient { } fn key(&self, key: &str) -> String { - self.namespace.apply(key) + self.0.namespace.apply(key) } fn keys(&self, keys: &[&str]) -> Vec { @@ -281,7 +298,7 @@ impl RedisClient { command: &Cmd, ) -> Result { self.guarded(Operation::Command(name), async { - command.query_async(&mut self.connection.clone()).await + command.query_async(&mut self.0.connection.clone()).await }) .await } @@ -297,7 +314,7 @@ impl RedisClient { self.check_health().await; let result = call .await - .map_err(|error| Error::classify(error, self.connection.deployment())); + .map_err(|error| Error::classify(error, self.0.connection.deployment())); settle(admission, result.as_ref().err()); result } @@ -309,10 +326,11 @@ impl RedisClient { 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.events { + if let Some(events) = &self.0.events { events.call(&CallEvent { operation, duration: started.elapsed(), @@ -325,10 +343,11 @@ impl RedisClient { /// sending the next command on it, and reconnects when the ping fails. 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.health_check else { + let Some(interval) = self.0.health_check else { return; }; let idle = self + .0 .last_used .lock() .unwrap_or_else(PoisonError::into_inner) @@ -336,13 +355,14 @@ impl RedisClient { if idle < interval { return; } - if let Err(error) = self.connection.ping_every_node().await { + if let Err(error) = self.0.connection.ping_every_node().await { tracing::debug!(%error, "idle Redis connection failed its health check, the next command reconnects"); } } fn admit(&self) -> Result, Error> { - self.breaker + self.0 + .breaker .as_ref() .map(|breaker| { breaker diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs index a5ed803dca8..6064a81892d 100644 --- a/litellm-rust/crates/redis-client/src/lib.rs +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -11,7 +11,7 @@ mod pubsub; mod script; pub use batch::{Declared, RedisBatch}; -pub use client::{RedisClient, Ttl}; +pub use client::{RedisClient, Ttl, WeakRedisClient}; pub use commands::{Amount, Expiry}; pub use config::{Credentials, CredentialsProvider, DEFAULT_TIMEOUT, RedisConfig, Timeouts}; pub use error::{ConnectError, Error, KeyError}; diff --git a/litellm-rust/crates/redis-connect/src/connect.rs b/litellm-rust/crates/redis-connect/src/connect.rs index 758cb4acd82..755f1b3af10 100644 --- a/litellm-rust/crates/redis-connect/src/connect.rs +++ b/litellm-rust/crates/redis-connect/src/connect.rs @@ -9,7 +9,7 @@ use litellm_redis_client::{ use litellm_redis_types::{Auth, HealthCheck, Resolution, StaticAuth}; use crate::{ - Error, RefreshingCredentials, + CredentialSource, Error, RefreshingCredentials, sources::{AzureAd, ElastiCacheIam, GcpIam}, }; @@ -19,6 +19,7 @@ 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 @@ -78,7 +79,10 @@ fn credentials(auth: Auth, host: &HostServices) -> Result { Auth::Static(StaticAuth { username, password }) => { Credentials::Static { username, password } } - Auth::Provider(_) => return Err(Error::NoHostCredentialProvider), + 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(), diff --git a/litellm-rust/crates/redis-connect/src/error.rs b/litellm-rust/crates/redis-connect/src/error.rs index b9e3f1ad5b5..98343a03fb8 100644 --- a/litellm-rust/crates/redis-connect/src/error.rs +++ b/litellm-rust/crates/redis-connect/src/error.rs @@ -40,10 +40,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 { @@ -86,6 +88,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 08d7af91fb7..b3d1d411eed 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, 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 6c0cf3027ce..7d50e947d74 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 5f834687cc3..30754a1912c 100644 --- a/litellm-rust/crates/redis-connect/tests/connect.rs +++ b/litellm-rust/crates/redis-connect/tests/connect.rs @@ -1,10 +1,13 @@ use std::{collections::BTreeMap, sync::Arc, time::Duration}; -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::{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 @@ -21,6 +24,21 @@ fn services() -> 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, + }) + }) } } @@ -126,7 +144,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})), @@ -170,3 +188,39 @@ 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() + }, + None, + ) + .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 e0e3f9ae39e..9c658c64001 100644 --- a/litellm-rust/crates/redis-connect/tests/live/main.rs +++ b/litellm-rust/crates/redis-connect/tests/live/main.rs @@ -45,6 +45,7 @@ async fn roundtrip(config: LiveConfig) { 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() 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/redis/stores/test_connection.py b/tests/unit/redis/stores/test_connection.py index 9eb7707b18d..83cbeddaaaa 100644 --- a/tests/unit/redis/stores/test_connection.py +++ b/tests/unit/redis/stores/test_connection.py @@ -1,3 +1,4 @@ +import gc import uuid from concurrent.futures import ThreadPoolExecutor from typing import Final @@ -26,3 +27,21 @@ def test_concurrent_first_connects_share_one_client_without_deadlocking() -> Non 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}, namespace=namespace) + assert client is not None + return len(native_redis_breaker_states()) + + +def test_a_discarded_client_leaves_the_process_registry(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true") + 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 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