fix(cache): bound compression cache bookkeeping
## Description `CompressionCache.max_entries` bounded the main compression cache, but not `_stable_hashes` or `_first_seen`. A long-lived session could therefore retain every unique tool-result hash even while `_cache` stayed empty. This change applies the same bounded retention to both side tables. It also cleans up expired first-seen entries and resets the timing window when compression occurs near the TTL boundary. Fixes #2874 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Store stable hashes and first-seen timestamps in ordered mappings. - Evict oldest entries when either side table exceeds `max_entries`. - Keep all bookkeeping under the existing reentrant lock. - Reset first-seen timing after compression near the TTL boundary. - Add tests covering size limits, TTL behavior, frozen-prefix safety, and concurrency. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text uv run ruff format --check . Passed uv run ruff check . All checks passed! uv run mypy headroom Success: no issues found in 515 source files uv run pytest Passed ``` Focused cache tests on macOS 26.5.2 arm64 with Python 3.11.14: ```text uv run pytest tests/test_compression_cache.py::TestCompressionCacheRetention -v 5 passed in 0.30s uv run pytest tests/test_compression_cache.py -q 38 passed in 5.76s ``` After the final formatting-only commit, the cache test file was also run on Linux with Python 3.12.13: ```text 37 passed, 1 skipped in 32.70s ``` ## Real Behavior Proof - Environment: Linux 6.18 x86_64, Python 3.12.13, `CompressionCache(max_entries=100)`. - Exact command / steps: Created a `CompressionCache(max_entries=100)`, generated 20,000 unique content hashes, and passed each hash through `mark_stable()` and `should_defer_compression()`. Store sizes were sampled after 100, 1,000, 5,000, and 20,000 results. - Observed result: `_cache=0`, `_stable_hashes=100`, and `_first_seen=100` at every sample after reaching the configured limit. At 20,000 results, traced memory was approximately 0.03 MB current and 0.04 MB peak. Before the fix, the same workload retained all 20,000 hashes and timestamps. - Not tested: A live multi-hour proxy/provider session. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented the code where retention behavior is not obvious - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove the fix is effective - [x] New and existing unit tests pass locally - [x] I did **not** edit `CHANGELOG.md` ## Screenshots N/A — internal cache bookkeeping change. ## Additional Notes No changes to dependencies, public APIs, or configuration. No user-facing behavior changes.
This commit is contained in:
Vendored
+41
-6
@@ -135,8 +135,10 @@ class CompressionCache:
|
||||
# `compute_frozen_count` (bounded above by the `min` clamp at
|
||||
# `proxy/handlers/anthropic.py`) and `update_from_result`'s
|
||||
# "unchanged content" tracking.
|
||||
self._stable_hashes: set[str] = set()
|
||||
self._first_seen: dict[str, float] = {}
|
||||
# Ordered mappings preserve set/dict-style membership while allowing
|
||||
# deterministic oldest-first eviction.
|
||||
self._stable_hashes: OrderedDict[str, None] = OrderedDict()
|
||||
self._first_seen: OrderedDict[str, float] = OrderedDict()
|
||||
self._hits: int = 0
|
||||
self._misses: int = 0
|
||||
self._total_tokens_saved: int = 0
|
||||
@@ -172,6 +174,34 @@ class CompressionCache:
|
||||
_, evicted = self._cache.popitem(last=False)
|
||||
self._total_tokens_saved -= evicted.tokens_saved
|
||||
|
||||
def _mark_stable_locked(self, content_hash: str) -> None:
|
||||
"""Record a stable hash while bounding retained bookkeeping."""
|
||||
self._stable_hashes[content_hash] = None
|
||||
self._stable_hashes.move_to_end(content_hash)
|
||||
|
||||
while len(self._stable_hashes) > self.max_entries:
|
||||
self._stable_hashes.popitem(last=False)
|
||||
|
||||
def _record_first_seen_locked(self, content_hash: str, seen_at: float) -> None:
|
||||
"""Record a first-seen timestamp while bounding retained bookkeeping."""
|
||||
self._first_seen[content_hash] = seen_at
|
||||
self._first_seen.move_to_end(content_hash)
|
||||
|
||||
while len(self._first_seen) > self.max_entries:
|
||||
self._first_seen.popitem(last=False)
|
||||
|
||||
def _prune_expired_first_seen_locked(
|
||||
self,
|
||||
now: float,
|
||||
ttl_seconds: float,
|
||||
) -> None:
|
||||
"""Remove first-seen entries whose cache timing window has expired."""
|
||||
while self._first_seen:
|
||||
_, oldest_seen_at = next(iter(self._first_seen.items()))
|
||||
if now - oldest_seen_at < ttl_seconds:
|
||||
break
|
||||
self._first_seen.popitem(last=False)
|
||||
|
||||
def mark_stable(self, content_hash: str) -> None:
|
||||
"""Mark a content hash as stable (unchanged, not compressed).
|
||||
|
||||
@@ -180,7 +210,7 @@ class CompressionCache:
|
||||
even though no compressed version exists in the cache.
|
||||
"""
|
||||
with self._lock:
|
||||
self._stable_hashes.add(content_hash)
|
||||
self._mark_stable_locked(content_hash)
|
||||
|
||||
def mark_stable_from_messages(self, messages: list[dict], up_to: int) -> None:
|
||||
"""Mark all tool_result hashes in messages[:up_to] as stable."""
|
||||
@@ -189,7 +219,7 @@ class CompressionCache:
|
||||
if _is_tool_result_message(msg):
|
||||
content = _extract_tool_result_content(msg)
|
||||
if content is not None:
|
||||
self._stable_hashes.add(self.content_hash(content))
|
||||
self._mark_stable_locked(self.content_hash(content))
|
||||
|
||||
def should_defer_compression(
|
||||
self,
|
||||
@@ -216,13 +246,18 @@ class CompressionCache:
|
||||
"""
|
||||
with self._lock:
|
||||
now = time.time()
|
||||
self._prune_expired_first_seen_locked(now, ttl_seconds)
|
||||
|
||||
first_seen = self._first_seen.get(content_hash)
|
||||
if first_seen is None:
|
||||
self._first_seen[content_hash] = now
|
||||
self._record_first_seen_locked(content_hash, now)
|
||||
return False # First time — compress now (no cache entry to preserve)
|
||||
|
||||
age = now - first_seen
|
||||
if age >= ttl_seconds - batch_window:
|
||||
self._record_first_seen_locked(content_hash, now)
|
||||
return False # Near TTL boundary — compress now (batch window)
|
||||
|
||||
return True # Seen recently within TTL — defer to preserve cache
|
||||
|
||||
def get_stats(self) -> dict:
|
||||
@@ -335,7 +370,7 @@ class CompressionCache:
|
||||
continue
|
||||
if orig_content == comp_content:
|
||||
# Content unchanged — mark as stable for frozen count walk
|
||||
self._stable_hashes.add(self.content_hash(orig_content))
|
||||
self._mark_stable_locked(self.content_hash(orig_content))
|
||||
continue
|
||||
h = self.content_hash(orig_content)
|
||||
tokens_saved = len(orig_content) // 4 - len(comp_content) // 4
|
||||
|
||||
@@ -17,6 +17,116 @@ def small_cache() -> CompressionCache:
|
||||
return CompressionCache(max_entries=3)
|
||||
|
||||
|
||||
class TestCompressionCacheRetention:
|
||||
def test_stable_hashes_are_bounded(self) -> None:
|
||||
cache = CompressionCache(max_entries=3)
|
||||
hashes = [CompressionCache.content_hash(f"stable-{index}") for index in range(4)]
|
||||
|
||||
for content_hash in hashes:
|
||||
cache.mark_stable(content_hash)
|
||||
|
||||
assert len(cache._stable_hashes) == 3
|
||||
assert hashes[0] not in cache._stable_hashes
|
||||
assert hashes[-1] in cache._stable_hashes
|
||||
|
||||
def test_first_seen_is_bounded(self) -> None:
|
||||
cache = CompressionCache(max_entries=3)
|
||||
hashes = [CompressionCache.content_hash(f"first-seen-{index}") for index in range(4)]
|
||||
|
||||
for content_hash in hashes:
|
||||
cache.should_defer_compression(content_hash)
|
||||
|
||||
assert len(cache._first_seen) == 3
|
||||
assert hashes[0] not in cache._first_seen
|
||||
assert hashes[-1] in cache._first_seen
|
||||
|
||||
def test_expired_first_seen_starts_new_window(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
cache = CompressionCache(max_entries=3)
|
||||
content_hash = CompressionCache.content_hash("repeated content")
|
||||
timestamps = iter([1_000.0, 1_271.0, 1_272.0])
|
||||
|
||||
monkeypatch.setattr(
|
||||
"headroom.cache.compression_cache.time.time",
|
||||
lambda: next(timestamps),
|
||||
)
|
||||
|
||||
assert (
|
||||
cache.should_defer_compression(
|
||||
content_hash,
|
||||
ttl_seconds=300,
|
||||
batch_window=30,
|
||||
)
|
||||
is False
|
||||
)
|
||||
assert (
|
||||
cache.should_defer_compression(
|
||||
content_hash,
|
||||
ttl_seconds=300,
|
||||
batch_window=30,
|
||||
)
|
||||
is False
|
||||
)
|
||||
assert cache._first_seen[content_hash] == 1_271.0
|
||||
assert (
|
||||
cache.should_defer_compression(
|
||||
content_hash,
|
||||
ttl_seconds=300,
|
||||
batch_window=30,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
def test_evicted_stable_hash_does_not_extend_frozen_prefix(self) -> None:
|
||||
cache = CompressionCache(max_entries=1)
|
||||
old_content = "old stable tool output"
|
||||
new_content = "new stable tool output"
|
||||
|
||||
cache.mark_stable(CompressionCache.content_hash(old_content))
|
||||
cache.mark_stable(CompressionCache.content_hash(new_content))
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": "start"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "tool-1",
|
||||
"content": old_content,
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "follow up"},
|
||||
]
|
||||
|
||||
assert cache.compute_frozen_count(messages) == 1
|
||||
|
||||
def test_concurrent_bookkeeping_stays_bounded(self) -> None:
|
||||
import threading
|
||||
|
||||
cache = CompressionCache(max_entries=50)
|
||||
errors: list[Exception] = []
|
||||
|
||||
def worker(thread_id: int) -> None:
|
||||
try:
|
||||
for index in range(100):
|
||||
content_hash = CompressionCache.content_hash(f"thread-{thread_id}-{index}")
|
||||
cache.mark_stable(content_hash)
|
||||
cache.should_defer_compression(content_hash)
|
||||
except Exception as exc: # pragma: no cover
|
||||
errors.append(exc)
|
||||
|
||||
threads = [threading.Thread(target=worker, args=(index,)) for index in range(8)]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
assert errors == []
|
||||
assert len(cache._stable_hashes) <= cache.max_entries
|
||||
assert len(cache._first_seen) <= cache.max_entries
|
||||
|
||||
|
||||
class TestCompressionCache:
|
||||
def test_cache_miss_returns_none(self, cache: CompressionCache) -> None:
|
||||
h = CompressionCache.content_hash("some content")
|
||||
|
||||
Reference in New Issue
Block a user