diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 56c9147e066..ce698f589ed 100644 --- a/litellm/caching/in_memory_cache.py +++ b/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 diff --git a/tests/test_litellm/caching/test_in_memory_cache.py b/tests/test_litellm/caching/test_in_memory_cache.py index 40ad4f0c6f0..53e37c04c3e 100644 --- a/tests/test_litellm/caching/test_in_memory_cache.py +++ b/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.