From e3769b18d5d1b72ff09cd924c6d1b06429de8538 Mon Sep 17 00:00:00 2001 From: suhasagg Date: Mon, 5 Oct 2026 16:55:52 +0530 Subject: [PATCH] fix(purview): prevent cache size double accounting --- .../purview/agent_framework_purview/_cache.py | 9 +++++---- .../packages/purview/tests/purview/test_cache.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/python/packages/purview/agent_framework_purview/_cache.py b/python/packages/purview/agent_framework_purview/_cache.py index df5d03b97ee..39321082e8c 100644 --- a/python/packages/purview/agent_framework_purview/_cache.py +++ b/python/packages/purview/agent_framework_purview/_cache.py @@ -139,10 +139,11 @@ async def set(self, key: str, value: Any, ttl_seconds: int | None = None) -> Non expiry = time.time() + ttl size = self._estimate_size(value) - # Remove old entry if exists - if key in self._cache: - old_size = self._cache[key][2] - self._current_size_bytes -= old_size + # Remove old entry if it exists before eviction so size accounting + # reflects only entries that are still present in the cache. + old_entry = self._cache.pop(key, None) + if old_entry is not None: + self._current_size_bytes -= old_entry[2] # Evict if needed self._evict_if_needed(size) diff --git a/python/packages/purview/tests/purview/test_cache.py b/python/packages/purview/tests/purview/test_cache.py index 1892842d424..170b54c367e 100644 --- a/python/packages/purview/tests/purview/test_cache.py +++ b/python/packages/purview/tests/purview/test_cache.py @@ -149,6 +149,20 @@ async def test_cache_multiple_updates(self) -> None: assert cache._current_size_bytes != initial_size + async def test_cache_update_larger_entry_at_capacity_maintains_size_accounting(self) -> None: + """Test growing an existing entry at capacity maintains correct size accounting.""" + cache = InMemoryCacheProvider(max_size_bytes=200) + + await cache.set("key1", "a" * 50, ttl_seconds=10) + await cache.set("key2", "b" * 50, ttl_seconds=20) + await cache.set("key3", "c" * 50, ttl_seconds=30) + + await cache.set("key1", "d" * 100, ttl_seconds=40) + + assert await cache.get("key1") == "d" * 100 + assert cache._current_size_bytes == sum(entry[2] for entry in cache._cache.values()) + assert cache._current_size_bytes <= cache._max_size_bytes + async def test_eviction_with_stale_heap_entries(self) -> None: """Test that eviction correctly handles stale heap entries.""" cache = InMemoryCacheProvider(max_size_bytes=500)