Python: Enforce the Purview cache capacity for oversized values

python/packages/purview/agent_framework_purview/_cache.py

@@ -139,6 +139,10 @@ async def set(self, key: str, value: Any, ttl_seconds: int | None = None) -> Non
         expiry = time.time() + ttl
         size = self._estimate_size(value)
 
+        if size > self._max_size_bytes:
+            await self.remove(key)
+            return
+
         # Remove old entry if exists
         if key in self._cache:
             old_size = self._cache[key][2]

python/packages/purview/tests/purview/test_cache.py

@@ -94,6 +94,45 @@ async def test_cache_size_limit_eviction(self) -> None:
         result1 = await cache.get("key1")
         assert result1 is None
 
+    async def test_cache_oversized_entry_preserves_other_entries(self) -> None:
+        """Skip a value exceeding the total capacity while retaining cached peers."""
+        cache = InMemoryCacheProvider(max_size_bytes=20)
+        await cache.set("existing", "value")
+
+        await cache.set("oversized", "x" * 20)
+
+        assert await cache.get("oversized") is None
+        assert await cache.get("existing") == "value"
+        assert cache._current_size_bytes == 7
+
+    async def test_cache_oversized_replacement_invalidates_previous_value(self) -> None:
+        """Remove the superseded value when its replacement exceeds capacity."""
+        cache = InMemoryCacheProvider(max_size_bytes=20)
+        await cache.set("updated", "old", ttl_seconds=10)
+        await cache.set("other", "peer", ttl_seconds=20)
+
+        await cache.set("updated", "x" * 20)
+
+        assert await cache.get("updated") is None
+        assert await cache.get("other") == "peer"
+        assert cache._current_size_bytes == 6
+
+        await cache.set("next", "n" * 12)
+        assert await cache.get("next") == "n" * 12
+        assert await cache.get("other") == "peer"
+        assert cache._current_size_bytes == 20
+
+    async def test_cache_entry_filling_capacity_is_cached(self) -> None:
+        """Allow a value whose estimated size exactly matches the capacity."""
+        cache = InMemoryCacheProvider(max_size_bytes=20)
+        await cache.set("existing", "value")
+
+        await cache.set("full", "x" * 18)
+
+        assert await cache.get("full") == "x" * 18
+        assert await cache.get("existing") is None
+        assert cache._current_size_bytes == 20
+
     async def test_estimate_size_with_pydantic_model(self) -> None:
         """Test size estimation with Pydantic models."""
         cache = InMemoryCacheProvider()