diff --git a/python/packages/purview/agent_framework_purview/_cache.py b/python/packages/purview/agent_framework_purview/_cache.py index df5d03b97ee..3c54c7382ab 100644 --- a/python/packages/purview/agent_framework_purview/_cache.py +++ b/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] diff --git a/python/packages/purview/tests/purview/test_cache.py b/python/packages/purview/tests/purview/test_cache.py index 1892842d424..881deec706f 100644 --- a/python/packages/purview/tests/purview/test_cache.py +++ b/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()