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:
Joseph Benno
2026-08-11 12:52:27 -04:00
committed by GitHub
parent 739fdef423
commit 0ae948c151
2 changed files with 151 additions and 6 deletions
+41 -6
View File
@@ -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
+110
View File
@@ -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")