From efcafa7f120341187991bd2655cdec3257459b3b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 12:30:17 -0700 Subject: [PATCH] fix(rust): invalidate changed Redis pool settings --- .../crates/python-bridge/src/cache/facade.rs | 69 ++++++++++++++++++- tests/test_litellm_rust/test_cache.py | 4 ++ 2 files changed, 70 insertions(+), 3 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 83508356263..f2f86c14b37 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -25,9 +25,17 @@ struct ObjectGuard { config: Vec, } +struct RedisPoolGuard { + reference: Py, + connection_class: Py, + connection_kwargs: Py, + max_connections: usize, +} + pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, + redis_pool: Option, } impl ObjectGuard { @@ -129,6 +137,45 @@ impl ObjectGuard { } } +impl RedisPoolGuard { + fn capture(backend: &Bound<'_, PyAny>) -> PyResult { + let pool = backend + .getattr("redis_client")? + .getattr("connection_pool")?; + Ok(Self { + reference: pool.clone().unbind(), + connection_class: pool.getattr("connection_class")?.unbind(), + connection_kwargs: pool + .getattr("connection_kwargs")? + .call_method0("copy")? + .unbind(), + max_connections: pool.getattr("max_connections")?.extract::()?, + }) + } + + fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { + let pool = backend + .getattr("redis_client")? + .getattr("connection_pool")?; + Ok(self.reference.bind(py).is(&pool) + && self + .connection_class + .bind(py) + .is(&pool.getattr("connection_class")?) + && self.max_connections == pool.getattr("max_connections")?.extract::()? + && self + .connection_kwargs + .bind(py) + .eq(pool.getattr("connection_kwargs")?)?) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.reference)?; + visit.call(&self.connection_class)?; + visit.call(&self.connection_kwargs) + } +} + impl FacadeGuard { pub(super) fn capture( py: Python<'_>, @@ -190,17 +237,33 @@ impl FacadeGuard { "redis_flush_size", ], )?, + redis_pool: (kind == "redis") + .then(|| RedisPoolGuard::capture(&backend)) + .transpose()?, }) } fn matches(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult { - Ok(self.outer.matches(py, facade)? - && self.backend.matches(py, &facade.getattr("cache")?)?) + if !self.outer.matches(py, facade)? { + return Ok(false); + } + let backend = facade.getattr("cache")?; + if !self.backend.matches(py, &backend)? { + return Ok(false); + } + match &self.redis_pool { + Some(guard) => guard.matches(py, &backend), + None => Ok(true), + } } pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { self.outer.traverse(&visit)?; - self.backend.traverse(&visit) + self.backend.traverse(&visit)?; + if let Some(guard) = &self.redis_pool { + guard.traverse(&visit)?; + } + Ok(()) } } diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 796a0ec36ac..c35cb1a20fb 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -381,6 +381,10 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}): assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + pool: Final = facade.cache.redis_client.connection_pool + with rebound(pool, "connection_kwargs", {**pool.connection_kwargs, "db": 1}): + assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + await binding.async_store(request("first"), {"value": 1}) assert client.get("first") is None await binding.async_store(request("second"), {"value": 2})