mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
0ae948c151
## 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.
799 lines
30 KiB
Python
799 lines
30 KiB
Python
"""Tests for CompressionCache with LRU eviction."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from headroom.cache.compression_cache import CompressionCache
|
|
|
|
|
|
@pytest.fixture
|
|
def cache() -> CompressionCache:
|
|
return CompressionCache()
|
|
|
|
|
|
@pytest.fixture
|
|
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")
|
|
assert cache.get_compressed(h) is None
|
|
|
|
def test_store_and_retrieve(self, cache: CompressionCache) -> None:
|
|
content = "hello world this is a long message"
|
|
h = CompressionCache.content_hash(content)
|
|
cache.store_compressed(h, "hello world...compressed", tokens_saved=15)
|
|
assert cache.get_compressed(h) == "hello world...compressed"
|
|
|
|
def test_different_content_different_hash(self) -> None:
|
|
h1 = CompressionCache.content_hash("content A")
|
|
h2 = CompressionCache.content_hash("content B")
|
|
assert h1 != h2
|
|
|
|
def test_overwrite_same_hash(self, cache: CompressionCache) -> None:
|
|
h = CompressionCache.content_hash("some content")
|
|
cache.store_compressed(h, "v1", tokens_saved=10)
|
|
cache.store_compressed(h, "v2", tokens_saved=20)
|
|
assert cache.get_compressed(h) == "v2"
|
|
|
|
def test_stats_tracking(self, cache: CompressionCache) -> None:
|
|
h = CompressionCache.content_hash("content")
|
|
cache.store_compressed(h, "compressed", tokens_saved=5)
|
|
|
|
# One hit
|
|
cache.get_compressed(h)
|
|
# One miss
|
|
cache.get_compressed("nonexistent")
|
|
|
|
stats = cache.get_stats()
|
|
assert stats["hits"] == 1
|
|
assert stats["misses"] == 1
|
|
assert stats["entries"] == 1
|
|
assert stats["tokens_saved"] == 5
|
|
|
|
def test_eviction_at_max_entries(self, small_cache: CompressionCache) -> None:
|
|
h1 = CompressionCache.content_hash("a")
|
|
h2 = CompressionCache.content_hash("b")
|
|
h3 = CompressionCache.content_hash("c")
|
|
h4 = CompressionCache.content_hash("d")
|
|
|
|
small_cache.store_compressed(h1, "ca", tokens_saved=1)
|
|
small_cache.store_compressed(h2, "cb", tokens_saved=1)
|
|
small_cache.store_compressed(h3, "cc", tokens_saved=1)
|
|
|
|
# Adding a 4th should evict the oldest (h1)
|
|
small_cache.store_compressed(h4, "cd", tokens_saved=1)
|
|
|
|
assert small_cache.get_compressed(h1) is None
|
|
assert small_cache.get_compressed(h2) == "cb"
|
|
assert small_cache.get_compressed(h4) == "cd"
|
|
|
|
def test_access_refreshes_lru(self, small_cache: CompressionCache) -> None:
|
|
h1 = CompressionCache.content_hash("a")
|
|
h2 = CompressionCache.content_hash("b")
|
|
h3 = CompressionCache.content_hash("c")
|
|
h4 = CompressionCache.content_hash("d")
|
|
|
|
small_cache.store_compressed(h1, "ca", tokens_saved=1)
|
|
small_cache.store_compressed(h2, "cb", tokens_saved=1)
|
|
small_cache.store_compressed(h3, "cc", tokens_saved=1)
|
|
|
|
# Access h1 to refresh it
|
|
small_cache.get_compressed(h1)
|
|
|
|
# Adding h4 should evict h2 (oldest untouched), not h1
|
|
small_cache.store_compressed(h4, "cd", tokens_saved=1)
|
|
|
|
assert small_cache.get_compressed(h1) == "ca"
|
|
assert small_cache.get_compressed(h2) is None
|
|
assert small_cache.get_compressed(h4) == "cd"
|
|
|
|
def test_content_hash_list_content(self) -> None:
|
|
"""content_hash handles Anthropic-format list content."""
|
|
list_content = [
|
|
{"type": "text", "text": "hello"},
|
|
{"type": "text", "text": "world"},
|
|
]
|
|
h = CompressionCache.content_hash(list_content)
|
|
assert isinstance(h, str)
|
|
assert len(h) == 16
|
|
|
|
# Same content produces same hash
|
|
assert CompressionCache.content_hash(list_content) == h
|
|
|
|
def test_content_hash_string_length(self) -> None:
|
|
h = CompressionCache.content_hash("test")
|
|
assert len(h) == 16
|
|
|
|
|
|
class TestCompressionCacheFrozenCount:
|
|
def test_empty_cache_returns_zero(self, cache: CompressionCache) -> None:
|
|
assert cache.compute_frozen_count([]) == 0
|
|
|
|
def test_user_assistant_stable_with_live_zone_cap(self, cache: CompressionCache) -> None:
|
|
"""Plain user/assistant turns are individually stable, but the
|
|
trailing message is reserved as the live zone — the new turn
|
|
cannot be in any provider prefix cache. See docstring on
|
|
``CompressionCache.compute_frozen_count``."""
|
|
messages = [
|
|
{"role": "user", "content": "hello"},
|
|
{"role": "assistant", "content": "hi there"},
|
|
{"role": "user", "content": "how are you"},
|
|
]
|
|
# 3 messages structurally stable; cap clamps to len-1 = 2.
|
|
assert cache.compute_frozen_count(messages) == 2
|
|
|
|
def test_tool_result_with_cache_hit_capped_at_live_zone(self, cache: CompressionCache) -> None:
|
|
tool_content = "tool output data"
|
|
h = CompressionCache.content_hash(tool_content)
|
|
cache.store_compressed(h, "compressed tool output", tokens_saved=5)
|
|
|
|
messages = [
|
|
{"role": "user", "content": "do something"},
|
|
{
|
|
"role": "assistant",
|
|
"content": [{"type": "tool_use", "id": "t1", "name": "my_tool", "input": {}}],
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": tool_content}],
|
|
},
|
|
]
|
|
# All 3 stable; cap clamps to len-1 = 2 (trailing tool_result is
|
|
# the live zone).
|
|
assert cache.compute_frozen_count(messages) == 2
|
|
|
|
def test_tool_result_cache_miss_stops_frozen(self, cache: CompressionCache) -> None:
|
|
messages = [
|
|
{"role": "user", "content": "hello"},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "tool_result", "tool_use_id": "t1", "content": "uncached stuff"}
|
|
],
|
|
},
|
|
{"role": "user", "content": "follow up"},
|
|
]
|
|
assert cache.compute_frozen_count(messages) == 1
|
|
|
|
def test_frozen_count_with_dropped_messages(self, cache: CompressionCache) -> None:
|
|
cached_content = "cached tool output"
|
|
h = CompressionCache.content_hash(cached_content)
|
|
cache.store_compressed(h, "compressed", tokens_saved=3)
|
|
|
|
messages = [
|
|
{"role": "user", "content": "start"},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "tool_result", "tool_use_id": "t1", "content": cached_content}
|
|
],
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "t2", "content": "not cached"}],
|
|
},
|
|
]
|
|
assert cache.compute_frozen_count(messages) == 2
|
|
|
|
def test_stable_hash_allows_frozen_count_past_uncached_tool_result(
|
|
self, cache: CompressionCache
|
|
) -> None:
|
|
"""Tool_results marked stable should not stop the frozen count walk."""
|
|
tool_content = "excluded Read output — big file contents"
|
|
h = CompressionCache.content_hash(tool_content)
|
|
cache.mark_stable(h)
|
|
|
|
messages = [
|
|
{"role": "user", "content": "hello"},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": tool_content}],
|
|
},
|
|
{"role": "user", "content": "follow up"},
|
|
]
|
|
# Without mark_stable, the walk would stop at msg[1] → frozen=1.
|
|
# With stable hash, the walk continues past msg[1]; structural
|
|
# count = 3, then capped at len-1 = 2 (live-zone reservation).
|
|
assert cache.compute_frozen_count(messages) == 2
|
|
|
|
def test_update_from_result_identical_content_marks_stable(
|
|
self, cache: CompressionCache
|
|
) -> None:
|
|
"""When orig == compressed, update_from_result marks the hash as stable."""
|
|
tool_content = "unchanged tool output"
|
|
originals = [
|
|
{"role": "user", "content": "hi"},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": tool_content}],
|
|
},
|
|
]
|
|
# Compressed is identical to originals (no compression happened)
|
|
compressed = [
|
|
{"role": "user", "content": "hi"},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": tool_content}],
|
|
},
|
|
]
|
|
cache.update_from_result(originals, compressed)
|
|
|
|
h = CompressionCache.content_hash(tool_content)
|
|
assert h in cache._stable_hashes
|
|
|
|
# Frozen count walks past this tool_result (its hash is stable),
|
|
# but the trailing message is still reserved as live zone.
|
|
messages = [
|
|
{"role": "user", "content": "hello"},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": tool_content}],
|
|
},
|
|
{"role": "user", "content": "more stuff"},
|
|
]
|
|
assert cache.compute_frozen_count(messages) == 2
|
|
|
|
def test_mark_stable_from_messages(self, cache: CompressionCache) -> None:
|
|
"""mark_stable_from_messages records hashes for tool_results."""
|
|
content_a = "tool output A"
|
|
content_b = "tool output B"
|
|
messages = [
|
|
{"role": "user", "content": "hi"},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": content_a}],
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "t2", "content": content_b}],
|
|
},
|
|
]
|
|
# Mark first 2 messages (msg[0] + msg[1])
|
|
cache.mark_stable_from_messages(messages, 2)
|
|
|
|
ha = CompressionCache.content_hash(content_a)
|
|
hb = CompressionCache.content_hash(content_b)
|
|
assert ha in cache._stable_hashes
|
|
assert hb not in cache._stable_hashes # msg[2] not included
|
|
|
|
def test_should_defer_compression_new_content(self, cache: CompressionCache) -> None:
|
|
"""First-time content should NOT be deferred — there is no
|
|
prefix-cache entry to preserve, so compression carries no bust
|
|
cost. Issue #327: prior behavior deferred first-sight, which
|
|
marked every fresh tool_result as stable and disabled
|
|
compression for typical Claude Code workloads.
|
|
"""
|
|
h = CompressionCache.content_hash("brand new content")
|
|
assert cache.should_defer_compression(h, ttl_seconds=300, batch_window=30) is False
|
|
# Subsequent sightings within TTL should defer (batch window).
|
|
assert cache.should_defer_compression(h, ttl_seconds=300, batch_window=30) is True
|
|
|
|
def test_should_defer_compression_records_first_seen(self, cache: CompressionCache) -> None:
|
|
"""First-sight call must record the timestamp so subsequent
|
|
in-window calls can defer. Without this the deferral pathway
|
|
for genuinely-repeated content stops working."""
|
|
h = CompressionCache.content_hash("seen-twice content")
|
|
cache.should_defer_compression(h) # first sight
|
|
assert h in cache._first_seen
|
|
|
|
def test_should_defer_compression_near_ttl(self, cache: CompressionCache) -> None:
|
|
"""Content near TTL boundary should NOT be deferred."""
|
|
import time
|
|
|
|
h = CompressionCache.content_hash("old content")
|
|
# Backdate first_seen to simulate age near TTL
|
|
cache._first_seen[h] = time.time() - 280 # 280s old, TTL=300, window=30
|
|
assert cache.should_defer_compression(h, ttl_seconds=300, batch_window=30) is False
|
|
|
|
|
|
class TestCompressionCacheApplyAndUpdate:
|
|
def test_apply_cached_swaps_tool_results(self, cache: CompressionCache) -> None:
|
|
original_content = "big tool output"
|
|
h = CompressionCache.content_hash(original_content)
|
|
cache.store_compressed(h, "small output", tokens_saved=5)
|
|
|
|
messages = [
|
|
{"role": "user", "content": "hi"},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "tool_result", "tool_use_id": "t1", "content": original_content}
|
|
],
|
|
},
|
|
]
|
|
result = cache.apply_cached(messages)
|
|
assert result[1]["content"][0]["content"] == "small output"
|
|
|
|
def test_apply_cached_preserves_uncached_messages(self, cache: CompressionCache) -> None:
|
|
messages = [
|
|
{"role": "user", "content": "hello"},
|
|
{"role": "assistant", "content": "world"},
|
|
]
|
|
result = cache.apply_cached(messages)
|
|
assert result[0] is messages[0]
|
|
assert result[1] is messages[1]
|
|
|
|
def test_apply_cached_never_adds_messages(self, cache: CompressionCache) -> None:
|
|
# Store something in cache that doesn't correspond to any message
|
|
cache.store_compressed("orphan_hash", "orphan_value", tokens_saved=1)
|
|
|
|
messages = [
|
|
{"role": "user", "content": "hello"},
|
|
{"role": "assistant", "content": "hi"},
|
|
]
|
|
result = cache.apply_cached(messages)
|
|
assert len(result) == len(messages)
|
|
|
|
def test_update_from_result_caches_changes(self, cache: CompressionCache) -> None:
|
|
originals = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "tool_result", "tool_use_id": "t1", "content": "original output"}
|
|
],
|
|
},
|
|
]
|
|
compressed = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "tool_result", "tool_use_id": "t1", "content": "compressed output"}
|
|
],
|
|
},
|
|
]
|
|
cache.update_from_result(originals, compressed)
|
|
|
|
h = CompressionCache.content_hash("original output")
|
|
assert cache.get_compressed(h) == "compressed output"
|
|
|
|
def test_update_from_result_ignores_unchanged(self, cache: CompressionCache) -> None:
|
|
originals = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "tool_result", "tool_use_id": "t1", "content": "same content"}
|
|
],
|
|
},
|
|
]
|
|
compressed = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "tool_result", "tool_use_id": "t1", "content": "same content"}
|
|
],
|
|
},
|
|
]
|
|
cache.update_from_result(originals, compressed)
|
|
h = CompressionCache.content_hash("same content")
|
|
assert cache.get_compressed(h) is None
|
|
|
|
def test_apply_does_not_modify_original_messages(self, cache: CompressionCache) -> None:
|
|
original_content = "big tool output"
|
|
h = CompressionCache.content_hash(original_content)
|
|
cache.store_compressed(h, "small output", tokens_saved=5)
|
|
|
|
msg = {
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": original_content}],
|
|
}
|
|
messages = [msg]
|
|
cache.apply_cached(messages)
|
|
|
|
# Original must be untouched
|
|
assert msg["content"][0]["content"] == original_content
|
|
|
|
def test_openai_format_tool_result(self, cache: CompressionCache) -> None:
|
|
original_content = "openai tool output"
|
|
h = CompressionCache.content_hash(original_content)
|
|
cache.store_compressed(h, "compressed openai", tokens_saved=4)
|
|
|
|
messages = [
|
|
{"role": "tool", "tool_call_id": "tc1", "content": original_content},
|
|
]
|
|
result = cache.apply_cached(messages)
|
|
assert result[0]["content"] == "compressed openai"
|
|
# Original untouched
|
|
assert messages[0]["content"] == original_content
|
|
|
|
|
|
# ─── C1 (audit follow-up): concurrency regression suite ────────────────────
|
|
#
|
|
# CompressionCache must be safe under multi-threaded mutation. The proxy is
|
|
# async and dispatches multiple concurrent requests per `session_id` into
|
|
# `asyncio.to_thread` workers — a single CompressionCache instance therefore
|
|
# sees concurrent calls to `store_compressed` / `get_compressed` /
|
|
# `mark_stable_from_messages` / `apply_cached` / `update_from_result`.
|
|
# These tests provoke the race conditions that motivated adding `_lock`.
|
|
|
|
|
|
class TestCompressionCacheConcurrency:
|
|
"""Threading regression suite for the audit-followup lock."""
|
|
|
|
def test_concurrent_store_does_not_corrupt_total_tokens_saved(self) -> None:
|
|
"""Many threads each store_compressed with tokens_saved=N; the
|
|
bookkeeping field must equal SUM(N) when threads finish. Pre-lock
|
|
this races (read-modify-write of `_total_tokens_saved`)."""
|
|
import threading
|
|
|
|
cache = CompressionCache(max_entries=1_000_000)
|
|
n_threads = 32
|
|
per_thread = 100
|
|
per_thread_tokens = 7
|
|
|
|
def worker(tid: int) -> None:
|
|
for i in range(per_thread):
|
|
h = CompressionCache.content_hash(f"thread-{tid}-item-{i}")
|
|
cache.store_compressed(h, f"comp-{tid}-{i}", tokens_saved=per_thread_tokens)
|
|
|
|
threads = [threading.Thread(target=worker, args=(t,)) for t in range(n_threads)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
expected = n_threads * per_thread * per_thread_tokens
|
|
stats = cache.get_stats()
|
|
assert stats["entries"] == n_threads * per_thread
|
|
# The expected token count is exact only because each (thread, item)
|
|
# produces a unique hash → no overwrite path. Pre-lock this would be
|
|
# < expected due to lost updates.
|
|
assert stats["tokens_saved"] == expected
|
|
|
|
def test_concurrent_apply_cached_with_concurrent_store_does_not_raise(self) -> None:
|
|
"""`apply_cached` iterates `_cache` (via `get_compressed`); if a
|
|
concurrent `store_compressed` mutates the OrderedDict during the
|
|
iteration, pre-lock you'd get `RuntimeError: OrderedDict mutated
|
|
during iteration`. Locks make this a single critical section."""
|
|
import threading
|
|
|
|
cache = CompressionCache()
|
|
|
|
# Pre-populate so apply_cached has work to do.
|
|
for i in range(50):
|
|
h = CompressionCache.content_hash(f"seed-{i}")
|
|
cache.store_compressed(h, f"comp-{i}", tokens_saved=1)
|
|
|
|
msgs = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "tool_result",
|
|
"tool_use_id": f"t{i}",
|
|
"content": f"seed-{i}",
|
|
}
|
|
],
|
|
}
|
|
for i in range(50)
|
|
]
|
|
|
|
stop = threading.Event()
|
|
errors: list[Exception] = []
|
|
|
|
def reader() -> None:
|
|
try:
|
|
while not stop.is_set():
|
|
_ = cache.apply_cached(msgs)
|
|
except Exception as e: # pragma: no cover
|
|
errors.append(e)
|
|
|
|
def writer() -> None:
|
|
try:
|
|
for i in range(500):
|
|
h = CompressionCache.content_hash(f"writer-{i}")
|
|
cache.store_compressed(h, f"w-{i}", tokens_saved=1)
|
|
except Exception as e: # pragma: no cover
|
|
errors.append(e)
|
|
|
|
readers = [threading.Thread(target=reader) for _ in range(4)]
|
|
writers = [threading.Thread(target=writer) for _ in range(4)]
|
|
for t in readers + writers:
|
|
t.start()
|
|
for t in writers:
|
|
t.join()
|
|
stop.set()
|
|
for t in readers:
|
|
t.join()
|
|
|
|
assert errors == [], f"Concurrent ops raised: {errors}"
|
|
|
|
def test_concurrent_update_from_result_no_partial_state(self) -> None:
|
|
"""update_from_result must be all-or-nothing per call. With many
|
|
threads calling update_from_result in parallel on the same cache,
|
|
the final state must reflect every call's full effect (no partial
|
|
writes)."""
|
|
import threading
|
|
|
|
cache = CompressionCache()
|
|
|
|
n_threads = 16
|
|
per_thread_calls = 20
|
|
|
|
def worker(tid: int) -> None:
|
|
for i in range(per_thread_calls):
|
|
orig_text = f"orig-{tid}-{i}-" + "X" * 200
|
|
comp_text = f"comp-{tid}-{i}-" + "X" * 50
|
|
originals = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "tool_result",
|
|
"tool_use_id": f"t-{tid}-{i}",
|
|
"content": orig_text,
|
|
}
|
|
],
|
|
}
|
|
]
|
|
compressed = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "tool_result",
|
|
"tool_use_id": f"t-{tid}-{i}",
|
|
"content": comp_text,
|
|
}
|
|
],
|
|
}
|
|
]
|
|
cache.update_from_result(originals, compressed)
|
|
|
|
threads = [threading.Thread(target=worker, args=(t,)) for t in range(n_threads)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
stats = cache.get_stats()
|
|
# Each (tid, i) is a unique hash → cache entries == n_threads * per_thread_calls.
|
|
assert stats["entries"] == n_threads * per_thread_calls
|
|
assert stats["tokens_saved"] > 0
|
|
|
|
def test_concurrent_hits_misses_consistent(self) -> None:
|
|
"""Under concurrent reads + writes, hits+misses must be bounded by
|
|
total lookups (hits ≤ entries, misses ≥ 0 at all moments)."""
|
|
import random
|
|
import threading
|
|
|
|
cache = CompressionCache(max_entries=1_000_000)
|
|
n_threads = 16
|
|
per_thread = 50
|
|
|
|
# Pre-populate so reads have something to hit
|
|
for i in range(per_thread):
|
|
h = CompressionCache.content_hash(f"hit-{i}")
|
|
cache.store_compressed(h, f"comp-{i}", tokens_saved=3)
|
|
|
|
errors: list[Exception] = []
|
|
barrier = threading.Barrier(n_threads)
|
|
|
|
def worker(tid: int) -> None:
|
|
try:
|
|
barrier.wait()
|
|
for i in range(per_thread):
|
|
if random.random() < 0.6:
|
|
# Read path
|
|
_ = cache.get_compressed(
|
|
CompressionCache.content_hash(
|
|
f"hit-{random.randint(0, per_thread - 1)}"
|
|
)
|
|
)
|
|
else:
|
|
# Write path
|
|
h = CompressionCache.content_hash(f"write-{tid}-{i}")
|
|
cache.store_compressed(h, f"w-{tid}-{i}", tokens_saved=1)
|
|
except Exception as e: # pragma: no cover
|
|
errors.append(e)
|
|
|
|
threads = [threading.Thread(target=worker, args=(t,)) for t in range(n_threads)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert errors == [], f"Concurrent reads+writes raised: {errors}"
|
|
stats = cache.get_stats()
|
|
# hits + misses should be non-negative (sanity)
|
|
assert stats["hits"] >= 0
|
|
assert stats["misses"] >= 0
|
|
assert stats["entries"] > 0
|
|
|
|
def test_concurrent_stable_hash_ops_no_race(self) -> None:
|
|
"""Concurrent mark_stable_from_messages + compute_frozen_count must
|
|
not race — stable_hashes must remain self-consistent."""
|
|
import threading
|
|
|
|
cache = CompressionCache()
|
|
n_threads = 12
|
|
per_thread = 30
|
|
|
|
# Each thread has its own content; produce tool_result messages
|
|
# and mark them stable, then verify frozen count.
|
|
errors: list[Exception] = []
|
|
barrier = threading.Barrier(n_threads)
|
|
|
|
def worker(tid: int) -> None:
|
|
try:
|
|
barrier.wait()
|
|
for i in range(per_thread):
|
|
content = f"stable-content-{tid}-{i}"
|
|
h = CompressionCache.content_hash(content)
|
|
# Also store to make it appear cached
|
|
cache.store_compressed(h, f"comp-{tid}-{i}", tokens_saved=2)
|
|
# Mark stable
|
|
cache.mark_stable(h)
|
|
except Exception as e: # pragma: no cover
|
|
errors.append(e)
|
|
|
|
threads = [threading.Thread(target=worker, args=(t,)) for t in range(n_threads)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert errors == [], f"Concurrent stable-hash ops raised: {errors}"
|
|
stats = cache.get_stats()
|
|
# All entries should be recorded; stable_hashes should match entries
|
|
# (every store_compressed was followed by mark_stable in our test)
|
|
assert stats["entries"] == n_threads * per_thread
|
|
|
|
|
|
def test_get_compression_cache_returns_same_instance_under_contention() -> None:
|
|
"""`HeadroomProxy._get_compression_cache(session_id)` must return the
|
|
SAME `CompressionCache` instance for concurrent calls with the same
|
|
session_id. Pre-lock, two concurrent calls could both see "not in dict"
|
|
and each create a new instance, splitting the cache state across them.
|
|
"""
|
|
import threading
|
|
|
|
pytest.importorskip("fastapi")
|
|
from headroom.proxy.server import ProxyConfig, create_app
|
|
|
|
config = ProxyConfig(
|
|
optimize=False,
|
|
cache_enabled=False,
|
|
rate_limit_enabled=False,
|
|
cost_tracking_enabled=False,
|
|
log_requests=False,
|
|
ccr_inject_tool=False,
|
|
ccr_handle_responses=False,
|
|
ccr_context_tracking=False,
|
|
image_optimize=False,
|
|
)
|
|
app = create_app(config)
|
|
proxy = app.state.proxy
|
|
|
|
n_threads = 32
|
|
results: list[CompressionCache] = []
|
|
results_lock = threading.Lock()
|
|
|
|
def worker() -> None:
|
|
c = proxy._get_compression_cache("shared-session-id")
|
|
with results_lock:
|
|
results.append(c)
|
|
|
|
threads = [threading.Thread(target=worker) for _ in range(n_threads)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert len(results) == n_threads
|
|
first = results[0]
|
|
for c in results[1:]:
|
|
assert c is first, "Concurrent _get_compression_cache returned different instances"
|