diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 3b62c80b125..1908ff44a81 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -154,7 +154,7 @@ impl InMemoryCache { while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() { if state.expirations.get(&key).copied() != Some(expiration) { state.expiration_heap.pop(); - } else if expiration < now { + } else if expiration <= now { state.expiration_heap.pop(); Self::remove(state, &key); } else { diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index 423b8c959f2..aaf82641db7 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -45,6 +45,24 @@ fn default_explicit_and_override_ttls_follow_python_rules(clock: Arc) ); } +#[rstest] +fn write_at_expiry_boundary_refreshes_ttl(clock: Arc) { + let cache = cache(clock.clone(), 4); + cache + .set_cache("key", "first".into(), Some(Duration::from_secs(10))) + .unwrap(); + clock.store(110, Ordering::SeqCst); + cache + .set_cache("key", "second".into(), Some(Duration::from_secs(10))) + .unwrap(); + assert_eq!( + cache.expires_at("key").unwrap(), + Some(Duration::from_secs(120)) + ); + clock.store(115, Ordering::SeqCst); + assert_eq!(cache.get_cache("key").unwrap(), Some("second".into())); +} + #[rstest] fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc) { let cache = cache(clock, 2);