"""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 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"