fix(cache): preserve live entries when updating a full cache

litellm/caching/in_memory_cache.py

@@ -103,7 +103,7 @@ class InMemoryCache(BaseCache):
         self.cache_dict.pop(key, None)
         self.ttl_dict.pop(key, None)
 
-    def evict_cache(self):
+    def evict_cache(self, *, reserve_space: bool = True) -> None:
         """
         Eviction policy:
         1. First, remove expired items from ttl_dict and cache_dict
@@ -134,7 +134,7 @@ class InMemoryCache(BaseCache):
                 break
 
         # Step 2: Evict if cache is still full
-        while len(self.cache_dict) >= self.max_size_in_memory:
+        while len(self.cache_dict) > self.max_size_in_memory - int(reserve_space):
             expiration_time, key = heapq.heappop(self.expiration_heap)
             # Skip if key was removed or updated
             if self.ttl_dict.get(key) == expiration_time:
@@ -163,7 +163,7 @@ class InMemoryCache(BaseCache):
         # Always prune expired/outdated heap roots before inserting.
         # This keeps expiration_heap bounded even when the live cache stays
         # below max_size_in_memory and keys are reinserted after TTL expiry.
-        self.evict_cache()
+        self.evict_cache(reserve_space=key not in self.cache_dict)
         if not self.check_value_size(value):
             return
 

tests/test_litellm/caching/test_in_memory_cache.py

@@ -3,6 +3,7 @@ import json
 import threading
 import time
 from concurrent.futures import ThreadPoolExecutor
+from typing import Final, Literal
 from unittest.mock import MagicMock, patch
 
 import httpx
@@ -150,6 +151,26 @@ def test_in_memory_cache_max_size_with_ttl():
     assert "key_3" in in_memory_cache.cache_dict
 
 
+@pytest.mark.parametrize("updated_key", ["first", "second"])
+def test_updating_full_cache_preserves_entries_and_expiration(updated_key: Literal["first", "second"]) -> None:
+    clock: Final = MagicMock(return_value=0.0)
+    cache: Final = InMemoryCache(max_size_in_memory=2, clock=clock)
+    cache.set_cache("first", "original-a", ttl=10)
+    cache.set_cache("second", "original-b", ttl=30)
+
+    cache.set_cache(updated_key, "updated", ttl=60)
+
+    assert cache.get_cache("first") == ("updated" if updated_key == "first" else "original-a")
+    assert cache.get_cache("second") == ("updated" if updated_key == "second" else "original-b")
+
+    clock.return_value = 11.0
+    assert cache.get_cache("first") is None
+    assert cache.get_cache("second") == ("updated" if updated_key == "second" else "original-b")
+
+    clock.return_value = 31.0
+    assert cache.get_cache("second") is None
+
+
 def test_in_memory_cache_expired_items_evicted_first():
     """
     Test that expired items are evicted before non-expired items when cache is full.