From 6423c83a2442438cf5ce596e1a9dcfd123f819b3 Mon Sep 17 00:00:00 2001 From: emme1t <149944796+emme1t@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:19:58 -0700 Subject: [PATCH 1/5] fix(cache): preserve live entries when updating a full cache --- litellm/caching/in_memory_cache.py | 6 +++--- .../caching/test_in_memory_cache.py | 21 +++++++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) 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. From 04a76b8e154f5c2be951ab38f949b8eddfcde804 Mon Sep 17 00:00:00 2001 From: emme1t <149944796+emme1t@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:43:24 -0700 Subject: [PATCH 2/5] fix(cache): reserve capacity after resize eviction --- litellm/caching/in_memory_cache.py | 3 ++ .../caching/test_in_memory_cache.py | 42 +++++++++++++++++-- 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index ce698f589ed..d8501508134 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -167,6 +167,9 @@ class InMemoryCache(BaseCache): if not self.check_value_size(value): return + if key not in self.cache_dict and len(self.cache_dict) >= self.max_size_in_memory: + self.evict_cache() + self.cache_dict[key] = value if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl if "ttl" in kwargs and kwargs["ttl"] is not None: diff --git a/tests/test_litellm/caching/test_in_memory_cache.py b/tests/test_litellm/caching/test_in_memory_cache.py index 53e37c04c3e..054aacce0b0 100644 --- a/tests/test_litellm/caching/test_in_memory_cache.py +++ b/tests/test_litellm/caching/test_in_memory_cache.py @@ -4,7 +4,7 @@ import threading import time from concurrent.futures import ThreadPoolExecutor from typing import Final, Literal -from unittest.mock import MagicMock, patch +from unittest.mock import patch import httpx import pytest @@ -13,6 +13,7 @@ from fastapi.testclient import TestClient from unittest.mock import AsyncMock +from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache @@ -22,6 +23,17 @@ class _SlowInt(int): return _SlowInt(int(self) + value) +class _FakeClock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + def test_increment_cache_is_atomic_under_thread_concurrency(): cache = InMemoryCache() seed = 1000 @@ -153,7 +165,7 @@ def test_in_memory_cache_max_size_with_ttl(): @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) + clock: Final = _FakeClock() 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) @@ -163,14 +175,36 @@ def test_updating_full_cache_preserves_entries_and_expiration(updated_key: Liter 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 + clock.advance(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 + clock.advance(20.0) assert cache.get_cache("second") is None +@pytest.mark.parametrize( + ("updated_key", "expected_values"), + [ + ("first", ("updated", None, "original-c")), + ("third", (None, "original-b", "updated")), + ], +) +def test_updating_cache_after_capacity_reduction_respects_limit( + updated_key: Literal["first", "third"], expected_values: tuple[str | None, str | None, str | None] +) -> None: + cache: Final = InMemoryCache(max_size_in_memory=3, clock=_FakeClock()) + cache.set_cache("first", "original-a", ttl=10) + cache.set_cache("second", "original-b", ttl=20) + cache.set_cache("third", "original-c", ttl=30) + DualCache(in_memory_cache=cache).update_in_memory_max_size(2) + + cache.set_cache(updated_key, "updated", ttl=60) + + assert len(cache.cache_dict) == 2 + assert tuple(cache.get_cache(key) for key in ("first", "second", "third")) == expected_values + + def test_in_memory_cache_expired_items_evicted_first(): """ Test that expired items are evicted before non-expired items when cache is full. From 0002b538c1c3a47233c41ae1eab7131d8b371828 Mon Sep 17 00:00:00 2001 From: Bao Yansong Date: Sat, 10 Oct 2026 04:55:26 +0000 Subject: [PATCH 3/5] fix(cache): type in-memory cache entries for lint gate --- litellm/caching/in_memory_cache.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index d8501508134..4c6955d0f73 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -47,7 +47,7 @@ class InMemoryCache(BaseCache): self.max_size_per_item = max_size_per_item or MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB # 1MB = 1024KB # in-memory cache - self.cache_dict: dict = {} + self.cache_dict: dict[str, Any] = {} self.ttl_dict: dict = {} self.expiration_heap: list[tuple[float, str]] = [] self._increment_lock = threading.Lock() From f9eb141e779d5bb69c699065b2c9869244f282d9 Mon Sep 17 00:00:00 2001 From: Bao Yansong Date: Sat, 10 Oct 2026 05:06:14 +0000 Subject: [PATCH 4/5] fix(cache): preserve generic cache key typing --- litellm/caching/in_memory_cache.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 4c6955d0f73..76bf941dcd7 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -47,7 +47,7 @@ class InMemoryCache(BaseCache): self.max_size_per_item = max_size_per_item or MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB # 1MB = 1024KB # in-memory cache - self.cache_dict: dict[str, Any] = {} + self.cache_dict: dict[object, Any] = {} self.ttl_dict: dict = {} self.expiration_heap: list[tuple[float, str]] = [] self._increment_lock = threading.Lock() From 538d72cb9e6f2b4070c1212b37d117d633414f1f Mon Sep 17 00:00:00 2001 From: emme1t <149944796+emme1t@users.noreply.github.com> Date: Wed, 21 Oct 2026 03:00:29 -0700 Subject: [PATCH 5/5] fix: align cache consumers with the cache key type --- litellm/caching/llm_caching_handler.py | 2 +- .../specialty_caches/dynamic_logging_cache.py | 2 +- litellm/proxy/agent_endpoints/databricks_oauth.py | 4 ++-- litellm/proxy/common_utils/debug_utils.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/caching/llm_caching_handler.py b/litellm/caching/llm_caching_handler.py index a89e43b78b4..e7e7d9e6a90 100644 --- a/litellm/caching/llm_caching_handler.py +++ b/litellm/caching/llm_caching_handler.py @@ -37,7 +37,7 @@ class LLMClientCache(InMemoryCache): ) self.evicted_client_closer = evicted_client_closer or default_evicted_client_closer - def _remove_key(self, key: str) -> None: + def _remove_key(self, key: object) -> None: evicted: Final[object] = self.cache_dict.get(key) super()._remove_key(key) self.evicted_client_closer.schedule(evicted) diff --git a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py b/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py index 9361fc8dec4..f7a321e24eb 100644 --- a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py +++ b/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py @@ -27,7 +27,7 @@ class LangfuseInMemoryCache(InMemoryCache): (https://github.com/BerriAI/litellm/issues/11169). """ - def _remove_key(self, key: str) -> None: + def _remove_key(self, key: object) -> None: from litellm.integrations.langfuse.langfuse import LangFuseLogger evicted: Final = self.cache_dict.pop(key, None) diff --git a/litellm/proxy/agent_endpoints/databricks_oauth.py b/litellm/proxy/agent_endpoints/databricks_oauth.py index 7425c240d10..ac4fc021660 100644 --- a/litellm/proxy/agent_endpoints/databricks_oauth.py +++ b/litellm/proxy/agent_endpoints/databricks_oauth.py @@ -212,12 +212,12 @@ class DatabricksAppOAuthTokenCache(InMemoryCache): def __init__(self) -> None: super().__init__(default_ttl=_DEFAULT_TTL_SECONDS) - self._locks: dict[str, asyncio.Lock] = {} + self._locks: dict[object, asyncio.Lock] = {} def _get_lock(self, cache_key: str) -> asyncio.Lock: return self._locks.setdefault(cache_key, asyncio.Lock()) - def _remove_key(self, key: str) -> None: + def _remove_key(self, key: object) -> None: # Drop the per-key lock alongside the cached token so ``_locks`` stays # bounded by the live key set rather than growing for every key ever seen. super()._remove_key(key) diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 33633150711..e10598673c1 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -130,7 +130,7 @@ if os.environ.get("LITELLM_PROFILE", "false").lower() == "true": return {"top_50_memory_usage": result} -def _prompt_caching_pins(llm_router: "Router | None") -> tuple[Mapping[str, object], Mapping[str, float]]: +def _prompt_caching_pins(llm_router: "Router | None") -> tuple[Mapping[object, object], Mapping[str, float]]: if llm_router is None: return {}, {} pin_cache: Final = llm_router.prompt_caching_pin_cache.in_memory_cache