test: add integration tests for token_headroom mode

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

tests/test_token_headroom_mode.py

@@ -0,0 +1,307 @@
+"""Integration tests for token_headroom mode.
+
+Tests the CompressionCache working across simulated multi-turn conversations,
+verifying the critical invariants: no message injection, correct frozen counts,
+proper handling of both Anthropic and OpenAI formats, and correct behavior
+when Claude Code drops messages.
+"""
+
+import copy
+
+from headroom.cache.compression_cache import CompressionCache
+
+
+def _make_user_msg(text: str) -> dict:
+    return {"role": "user", "content": text}
+
+
+def _make_assistant_msg(text: str) -> dict:
+    return {"role": "assistant", "content": text}
+
+
+def _make_tool_use_msg(tool_id: str, name: str) -> dict:
+    return {
+        "role": "assistant",
+        "content": [{"type": "tool_use", "id": tool_id, "name": name, "input": {}}],
+    }
+
+
+def _make_tool_result_msg(tool_id: str, content: str) -> dict:
+    """Anthropic-format tool result."""
+    return {
+        "role": "user",
+        "content": [{"type": "tool_result", "tool_use_id": tool_id, "content": content}],
+    }
+
+
+def _make_openai_tool_msg(tool_call_id: str, content: str) -> dict:
+    """OpenAI-format tool result."""
+    return {"role": "tool", "tool_call_id": tool_call_id, "content": content}
+
+
+def _large_code_content(n: int = 200) -> str:
+    """Generate realistic Python code content."""
+    parts = ["import os\nimport sys\nfrom typing import List, Dict\n\n"]
+    for i in range(n // 10):
+        parts.append(
+            f"def function_{i}(arg: str) -> str:\n"
+            f'    """Docstring for function {i}."""\n'
+            f"    result = arg.strip()\n"
+            f"    for j in range({i}):\n"
+            f"        result += str(j)\n"
+            f"    return result\n\n"
+        )
+    return "".join(parts)
+
+
+class TestMultiTurnCompression:
+    """Simulate multi-turn conversations to verify compression cascade."""
+
+    def test_first_turn_nothing_cached(self):
+        """On first turn, no cache hits, frozen count is minimal."""
+        cache = CompressionCache()
+        messages = [
+            _make_user_msg("Read file.py"),
+            _make_tool_use_msg("t1", "Read"),
+            _make_tool_result_msg("t1", _large_code_content(100)),
+        ]
+        frozen = cache.compute_frozen_count(messages)
+        # user (stable) + tool_use (stable) + tool_result (miss) → 2
+        assert frozen == 2
+
+    def test_second_turn_cache_hits(self):
+        """After caching, same content gets cache hits."""
+        cache = CompressionCache()
+        code = _large_code_content(100)
+        compressed = "# compressed version"
+
+        # Simulate first turn: pipeline compressed the code
+        h = CompressionCache.content_hash(code)
+        cache.store_compressed(h, compressed, tokens_saved=500)
+
+        # Second turn: same messages
+        messages = [
+            _make_user_msg("Read file.py"),
+            _make_tool_use_msg("t1", "Read"),
+            _make_tool_result_msg("t1", code),
+            _make_user_msg("now edit it"),
+        ]
+        frozen = cache.compute_frozen_count(messages)
+        # All 4 messages stable (user, tool_use, tool_result cached, user)
+        assert frozen == 4
+
+        # apply_cached should swap the content
+        result = cache.apply_cached(messages)
+        tool_result = result[2]["content"][0]
+        assert tool_result["content"] == compressed
+
+    def test_multi_turn_waterfall(self):
+        """Messages age out progressively across turns."""
+        cache = CompressionCache()
+
+        # Build conversation with 3 read results
+        code_a = "code A " * 200
+        code_b = "code B " * 200
+        code_c = "code C " * 200
+
+        messages = [
+            _make_user_msg("read A"),
+            _make_tool_result_msg("t1", code_a),
+            _make_user_msg("read B"),
+            _make_tool_result_msg("t2", code_b),
+            _make_user_msg("read C"),
+            _make_tool_result_msg("t3", code_c),
+        ]
+
+        # Turn 1: nothing cached
+        frozen = cache.compute_frozen_count(messages)
+        assert frozen == 1  # only first user msg
+
+        # Simulate pipeline compressing A and B (not C — in protection window)
+        cache.store_compressed(CompressionCache.content_hash(code_a), "ca", tokens_saved=100)
+        cache.store_compressed(CompressionCache.content_hash(code_b), "cb", tokens_saved=100)
+
+        # Turn 2: A and B cached, C still uncached
+        frozen = cache.compute_frozen_count(messages)
+        # user(stable) + tool_result_A(cached) + user(stable) + tool_result_B(cached) + user(stable) + tool_result_C(miss)
+        assert frozen == 5
+
+        # Now cache C too
+        cache.store_compressed(CompressionCache.content_hash(code_c), "cc", tokens_saved=100)
+
+        # Turn 3: all cached
+        frozen = cache.compute_frozen_count(messages)
+        assert frozen == 6  # all stable
+
+
+class TestNoMessageInjection:
+    """Critical invariant: proxy never adds messages."""
+
+    def test_output_length_equals_input(self):
+        cache = CompressionCache()
+        messages = [
+            _make_user_msg("hello"),
+            _make_tool_result_msg("t1", _large_code_content(50)),
+            _make_user_msg("bye"),
+        ]
+        result = cache.apply_cached(messages)
+        assert len(result) == len(messages)
+
+    def test_orphan_cache_entries_not_injected(self):
+        """Cache entries with no matching message are NOT injected."""
+        cache = CompressionCache()
+        cache.store_compressed("orphan_hash_1", "orphan content 1", tokens_saved=100)
+        cache.store_compressed("orphan_hash_2", "orphan content 2", tokens_saved=200)
+
+        messages = [_make_user_msg("hello")]
+        result = cache.apply_cached(messages)
+        assert len(result) == 1
+        assert result[0]["content"] == "hello"
+
+    def test_input_not_mutated(self):
+        """apply_cached must NOT mutate the input list or messages."""
+        cache = CompressionCache()
+        code = "original code content"
+        h = CompressionCache.content_hash(code)
+        cache.store_compressed(h, "compressed", tokens_saved=50)
+
+        messages = [
+            _make_tool_result_msg("t1", code),
+        ]
+        original = copy.deepcopy(messages)
+        cache.apply_cached(messages)
+        assert messages == original
+
+
+class TestClaudeCodeDropsMessages:
+    """When Claude Code drops messages via its own context management."""
+
+    def test_dropped_messages_not_readded(self):
+        cache = CompressionCache()
+        content_a = "content A " * 100
+        content_b = "content B " * 100
+        cache.store_compressed(CompressionCache.content_hash(content_a), "ca", tokens_saved=100)
+        cache.store_compressed(CompressionCache.content_hash(content_b), "cb", tokens_saved=100)
+
+        # CC dropped the message with content_b
+        messages = [
+            _make_user_msg("hello"),
+            _make_tool_result_msg("t1", content_a),
+            _make_user_msg("continue"),
+        ]
+        result = cache.apply_cached(messages)
+        assert len(result) == 3  # NOT 4
+
+    def test_frozen_count_breaks_at_gap(self):
+        """Dropped cached message creates a gap that stops frozen count."""
+        cache = CompressionCache()
+        content_a = "content A " * 100
+        content_c = "content C " * 100
+        cache.store_compressed(CompressionCache.content_hash(content_a), "ca", tokens_saved=100)
+        cache.store_compressed(CompressionCache.content_hash(content_c), "cc", tokens_saved=100)
+
+        # CC dropped content_b, content_c is still here but preceded by uncached gap
+        messages = [
+            _make_tool_result_msg("t1", content_a),
+            _make_tool_result_msg("t2", "UNCACHED content_b replacement"),
+            _make_tool_result_msg("t3", content_c),
+        ]
+        frozen = cache.compute_frozen_count(messages)
+        # t1 (cached, stable), t2 (NOT cached, stop)
+        assert frozen == 1
+
+
+class TestOpenAIFormat:
+    """Verify OpenAI-format tool messages work correctly."""
+
+    def test_openai_tool_result_cached(self):
+        cache = CompressionCache()
+        content = "large openai output " * 100
+        h = CompressionCache.content_hash(content)
+        cache.store_compressed(h, "compressed openai output", tokens_saved=300)
+
+        messages = [
+            _make_user_msg("run command"),
+            {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [
+                    {
+                        "id": "tc1",
+                        "type": "function",
+                        "function": {"name": "bash", "arguments": "{}"},
+                    }
+                ],
+            },
+            _make_openai_tool_msg("tc1", content),
+        ]
+
+        result = cache.apply_cached(messages)
+        assert len(result) == 3
+        assert result[2]["content"] == "compressed openai output"
+
+    def test_openai_frozen_count(self):
+        cache = CompressionCache()
+        content = "openai tool output " * 100
+        h = CompressionCache.content_hash(content)
+        cache.store_compressed(h, "compressed", tokens_saved=200)
+
+        messages = [
+            _make_user_msg("hello"),
+            _make_openai_tool_msg("tc1", content),
+            _make_openai_tool_msg("tc2", "uncached content"),
+        ]
+        frozen = cache.compute_frozen_count(messages)
+        # user (stable), tool tc1 (cached), tool tc2 (miss)
+        assert frozen == 2
+
+
+class TestUpdateFromResult:
+    """Verify update_from_result correctly caches compression results."""
+
+    def test_caches_compressed_anthropic(self):
+        cache = CompressionCache()
+        original_content = "long original " * 100
+        compressed_content = "short compressed"
+
+        originals = [
+            _make_user_msg("hello"),
+            _make_tool_result_msg("t1", original_content),
+        ]
+        compressed = [
+            _make_user_msg("hello"),
+            _make_tool_result_msg("t1", compressed_content),
+        ]
+
+        cache.update_from_result(originals, compressed)
+
+        h = CompressionCache.content_hash(original_content)
+        assert cache.get_compressed(h) == compressed_content
+
+    def test_caches_compressed_openai(self):
+        cache = CompressionCache()
+        original_content = "long openai output " * 100
+        compressed_content = "short compressed"
+
+        originals = [_make_openai_tool_msg("tc1", original_content)]
+        compressed = [_make_openai_tool_msg("tc1", compressed_content)]
+
+        cache.update_from_result(originals, compressed)
+
+        h = CompressionCache.content_hash(original_content)
+        assert cache.get_compressed(h) == compressed_content
+
+    def test_length_mismatch_no_crash(self):
+        """If pipeline somehow changes message count, don't crash."""
+        cache = CompressionCache()
+        originals = [_make_user_msg("a"), _make_user_msg("b")]
+        compressed = [_make_user_msg("a")]  # shorter
+        # Should not raise, just log warning
+        cache.update_from_result(originals, compressed)
+        assert cache.get_stats()["entries"] == 0
+
+    def test_unchanged_content_not_cached(self):
+        cache = CompressionCache()
+        msg = _make_user_msg("same content")
+        cache.update_from_result([msg], [msg])
+        assert cache.get_stats()["entries"] == 0