feat: add frozen count, apply_cached, update_from_result to CompressionCache

Add three methods and supporting helpers for token headroom mode:
- compute_frozen_count: counts consecutive stable messages from start
- apply_cached: swaps cached compressions into tool results (immutable)
- update_from_result: learns new compressions from original/compressed pairs

Supports both Anthropic and OpenAI tool result formats.

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

headroom/cache/compression_cache.py

@@ -6,11 +6,15 @@ Maps original content hashes to their compressed versions.
 
 from __future__ import annotations
 
+import copy
 import hashlib
 import json
+import logging
 from collections import OrderedDict
 from dataclasses import dataclass
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class _CacheEntry:
@@ -20,6 +24,54 @@ class _CacheEntry:
     tokens_saved: int
 
 
+def _is_tool_result_message(msg: dict) -> bool:
+    """Check if a message is a tool result in Anthropic or OpenAI format."""
+    # OpenAI format: role="tool"
+    if msg.get("role") == "tool":
+        return True
+    # Anthropic format: role="user" with content list containing tool_result blocks
+    content = msg.get("content")
+    if isinstance(content, list):
+        return any(
+            isinstance(block, dict) and block.get("type") == "tool_result" for block in content
+        )
+    return False
+
+
+def _extract_tool_result_content(msg: dict) -> str | None:
+    """Extract text content from a tool result message (both formats)."""
+    # OpenAI format
+    if msg.get("role") == "tool":
+        content = msg.get("content")
+        return content if isinstance(content, str) else None
+    # Anthropic format
+    content = msg.get("content")
+    if isinstance(content, list):
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "tool_result":
+                inner = block.get("content")
+                if isinstance(inner, str):
+                    return inner
+    return None
+
+
+def _swap_tool_result_content(msg: dict, new_content: str) -> dict:
+    """Deep copy msg and replace tool result content with new_content."""
+    new_msg = copy.deepcopy(msg)
+    # OpenAI format
+    if new_msg.get("role") == "tool":
+        new_msg["content"] = new_content
+        return new_msg
+    # Anthropic format
+    content = new_msg.get("content")
+    if isinstance(content, list):
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "tool_result":
+                block["content"] = new_content
+                break
+    return new_msg
+
+
 class CompressionCache:
     """Content-addressed cache mapping content hashes to compressed versions.
 
@@ -84,3 +136,71 @@ class CompressionCache:
         else:
             raw = content
         return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]
+
+    def compute_frozen_count(self, messages: list[dict]) -> int:
+        """Count consecutive stable messages from the start.
+
+        A message is stable if it is a plain user/assistant/system message,
+        an assistant message with tool_use blocks, or a tool_result whose
+        content hash is already in the cache. The first unstable tool_result
+        (cache miss) stops the count.
+        """
+        count = 0
+        for msg in messages:
+            if _is_tool_result_message(msg):
+                content = _extract_tool_result_content(msg)
+                if content is not None:
+                    h = self.content_hash(content)
+                    if self.get_compressed(h) is None:
+                        break
+                else:
+                    # tool_result with non-string content; treat as unstable
+                    break
+            # Regular user/assistant/system messages and assistant+tool_use
+            # are always stable — fall through.
+            count += 1
+        return count
+
+    def apply_cached(self, messages: list[dict]) -> list[dict]:
+        """Return a new list with cached compressions swapped into tool results.
+
+        Never mutates the input list or any message dict within it.
+        Output always has the same length as input.
+        """
+        result: list[dict] = []
+        for msg in messages:
+            if _is_tool_result_message(msg):
+                content = _extract_tool_result_content(msg)
+                if content is not None:
+                    h = self.content_hash(content)
+                    compressed = self.get_compressed(h)
+                    if compressed is not None:
+                        result.append(_swap_tool_result_content(msg, compressed))
+                        continue
+            result.append(msg)
+        return result
+
+    def update_from_result(self, originals: list[dict], compressed: list[dict]) -> None:
+        """Cache new compressions by comparing original and compressed messages.
+
+        Index-aligned: for each position, if both are tool results and the
+        content differs, store the mapping original_hash -> compressed_content.
+        """
+        if len(originals) != len(compressed):
+            logger.warning(
+                "update_from_result: length mismatch (originals=%d, compressed=%d), skipping",
+                len(originals),
+                len(compressed),
+            )
+            return
+
+        for orig, comp in zip(originals, compressed):
+            orig_content = _extract_tool_result_content(orig)
+            comp_content = _extract_tool_result_content(comp)
+            if orig_content is None or comp_content is None:
+                continue
+            if orig_content == comp_content:
+                continue
+            h = self.content_hash(orig_content)
+            tokens_saved = len(orig_content) - len(comp_content)
+            self.store_compressed(h, comp_content, tokens_saved=max(tokens_saved, 0))

tests/test_compression_cache.py

@@ -107,3 +107,177 @@ class TestCompressionCache:
     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_always_stable(self, cache: CompressionCache) -> None:
+        messages = [
+            {"role": "user", "content": "hello"},
+            {"role": "assistant", "content": "hi there"},
+            {"role": "user", "content": "how are you"},
+        ]
+        assert cache.compute_frozen_count(messages) == 3
+
+    def test_tool_result_with_cache_hit_is_stable(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}],
+            },
+        ]
+        assert cache.compute_frozen_count(messages) == 3
+
+    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
+
+
+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