Python: fix Purview cache accounting when replacing entries

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)

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)