diff --git a/headroom/cache/compression_cache.py b/headroom/cache/compression_cache.py index a8cfc5f53..30f2ae103 100644 --- a/headroom/cache/compression_cache.py +++ b/headroom/cache/compression_cache.py @@ -10,6 +10,7 @@ import copy import hashlib import json import logging +import time from collections import OrderedDict from dataclasses import dataclass @@ -82,6 +83,8 @@ class CompressionCache: def __init__(self, max_entries: int = 10000) -> None: self.max_entries = max_entries self._cache: OrderedDict[str, _CacheEntry] = OrderedDict() + self._stable_hashes: set[str] = set() + self._first_seen: dict[str, float] = {} self._hits: int = 0 self._misses: int = 0 self._total_tokens_saved: int = 0 @@ -115,10 +118,51 @@ class CompressionCache: _, evicted = self._cache.popitem(last=False) self._total_tokens_saved -= evicted.tokens_saved + def mark_stable(self, content_hash: str) -> None: + """Mark a content hash as stable (unchanged, not compressed). + + Used for tool_results that the content router excluded or skipped. + These messages appear verbatim every turn, so they are prefix-stable + even though no compressed version exists in the cache. + """ + self._stable_hashes.add(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.""" + for msg in messages[:up_to]: + 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)) + + def should_defer_compression( + self, + content_hash: str, + ttl_seconds: float = 300.0, + batch_window: float = 30.0, + ) -> bool: + """Whether to defer compressing this content to avoid mid-TTL busts. + + Returns True if the content was first seen recently enough that + compressing it now would bust the cached prefix with no TTL benefit. + Returns False near the TTL boundary (within batch_window of expiry), + meaning we should compress now and accept one bust. + """ + now = time.time() + first_seen = self._first_seen.get(content_hash) + if first_seen is None: + self._first_seen[content_hash] = now + return True # First time seeing this — defer + age = now - first_seen + if age >= ttl_seconds - batch_window: + return False # Near TTL boundary — compress now (batch window) + return True # Still within TTL — defer to preserve cache + def get_stats(self) -> dict: """Return cache statistics.""" return { "entries": len(self._cache), + "stable_hashes": len(self._stable_hashes), "hits": self._hits, "misses": self._misses, "tokens_saved": self._total_tokens_saved, @@ -151,7 +195,7 @@ class CompressionCache: content = _extract_tool_result_content(msg) if content is not None: h = self.content_hash(content) - if h not in self._cache: + if h not in self._cache and h not in self._stable_hashes: break else: # tool_result with non-string content; treat as unstable @@ -200,6 +244,8 @@ class CompressionCache: if orig_content is None or comp_content is None: continue if orig_content == comp_content: + # Content unchanged — mark as stable for frozen count walk + self._stable_hashes.add(self.content_hash(orig_content)) continue h = self.content_hash(orig_content) tokens_saved = len(orig_content) // 4 - len(comp_content) // 4 diff --git a/headroom/cache/prefix_tracker.py b/headroom/cache/prefix_tracker.py index 56453bb3d..8e0fc0129 100644 --- a/headroom/cache/prefix_tracker.py +++ b/headroom/cache/prefix_tracker.py @@ -18,6 +18,7 @@ from __future__ import annotations import copy import hashlib +import json import logging import time from dataclasses import dataclass @@ -228,19 +229,45 @@ class PrefixCacheTracker: @staticmethod def _estimate_message_tokens(messages: list[dict[str, Any]]) -> list[int]: - """Rough token count per message (chars / 3.5).""" + """Rough token count per message (chars / 3.5). + + Counts text, tool_result content, and tool_use input fields + for accurate Anthropic-format estimation. + """ counts = [] for msg in messages: content = msg.get("content", "") if isinstance(content, str): chars = len(content) elif isinstance(content, list): - chars = sum( - len(str(block.get("text", ""))) for block in content if isinstance(block, dict) - ) + chars = 0 + for block in content: + if not isinstance(block, dict): + continue + block_type = block.get("type", "") + if block_type == "text": + chars += len(block.get("text", "")) + elif block_type == "tool_result": + inner = block.get("content", "") + if isinstance(inner, str): + chars += len(inner) + elif isinstance(inner, list): + chars += sum( + len(b.get("text", "")) for b in inner if isinstance(b, dict) + ) + elif block_type == "tool_use": + inp = block.get("input") + if isinstance(inp, str): + chars += len(inp) + elif isinstance(inp, dict): + chars += len(json.dumps(inp, separators=(",", ":"))) + else: + text = block.get("text", "") + if text: + chars += len(text) else: chars = 0 - # Add overhead for role, tool_use blocks, etc. + # Add overhead for role, block structure, etc. chars += 20 counts.append(max(1, int(chars / 3.5))) return counts diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index ad59da3ee..6f32729f7 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -458,6 +458,35 @@ class AnthropicHandlerMixin: tokenizer = get_tokenizer(model) original_tokens = tokenizer.count_messages(messages) + # Enterprise Security: scan request before compression + _security_ctx = None + if self.security: + try: + messages, _security_ctx = self.security.scan_request( + messages, + { + "provider": "anthropic", + "model": model, + "request_id": str(request_id), + "user_id": headers.get("x-api-key", "")[:16], + }, + ) + except Exception as e: + if hasattr(e, "reason"): + from fastapi.responses import JSONResponse as _JSONResp + + return _JSONResp( + status_code=403, + content={ + "type": "error", + "error": { + "type": "security_block", + "message": str(e), + }, + }, + ) + logger.warning(f"[{request_id}] Security scan error: {e}") + # Hook: pre_compress — let hooks modify messages before compression if self.config.hooks and not is_cache_mode(self.config.mode): @@ -538,6 +567,41 @@ class AnthropicHandlerMixin: # Safety: never freeze beyond provider-confirmed cached prefix. cache_frozen_count = comp_cache.compute_frozen_count(messages) frozen_message_count = min(frozen_message_count, cache_frozen_count) + # Record all tool_results in the verified frozen prefix as stable + comp_cache.mark_stable_from_messages(messages, frozen_message_count) + + # TTL-aware deferral: extend freeze to cover messages whose + # first-time compression would bust the cache mid-TTL window. + # This batches first-time compressions near the 5-min TTL + # boundary, trading one big bust for many small ones. + ttl_frozen = frozen_message_count + from headroom.cache.compression_cache import ( + _extract_tool_result_content, + _is_tool_result_message, + ) + + for idx in range(frozen_message_count, len(messages)): + msg = messages[idx] + if _is_tool_result_message(msg): + tr_content = _extract_tool_result_content(msg) + if tr_content is not None: + h = comp_cache.content_hash(tr_content) + # Already compressed or stable — keep going + if h in comp_cache._cache or h in comp_cache._stable_hashes: + ttl_frozen = idx + 1 + elif comp_cache.should_defer_compression(h): + # New content within TTL — defer and freeze + comp_cache.mark_stable(h) + ttl_frozen = idx + 1 + else: + break # TTL expired — compress this turn + else: + break + else: + # Non-tool_result messages are stable (user/assistant text) + ttl_frozen = idx + 1 + + frozen_message_count = ttl_frozen result = await asyncio.wait_for( asyncio.to_thread( @@ -966,6 +1030,7 @@ class AnthropicHandlerMixin: memory_user_id=memory_user_id, pipeline_timing=pipeline_timing, prefix_tracker=prefix_tracker, + original_messages=original_client_messages, ) else: response = await self._retry_request("POST", url, headers, body) @@ -1358,6 +1423,23 @@ class AnthropicHandlerMixin: if _compression_failed: response_headers["x-headroom-compression-failed"] = "true" + # Enterprise Security: scan response + de-anonymize + if self.security and _security_ctx and resp_json: + try: + resp_json = self.security.scan_response(resp_json, _security_ctx) + response = httpx.Response( + status_code=200, + content=json.dumps(resp_json).encode(), + headers=response_headers, + ) + return Response( + content=response.content, + status_code=response.status_code, + headers=response_headers, + ) + except Exception as sec_err: + logger.warning(f"[{request_id}] Security response scan error: {sec_err}") + return Response( content=response.content, status_code=response.status_code, diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index 80c6656c8..5480d8fd6 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -450,6 +450,7 @@ class StreamingMixin: memory_user_id: str | None = None, pipeline_timing: dict[str, float] | None = None, prefix_tracker: Any | None = None, + original_messages: list[dict] | None = None, ) -> StreamingResponse: """Stream response with metrics tracking and memory tool handling. @@ -551,9 +552,13 @@ class StreamingMixin: # real-time clients (LangGraph, LangChain, etc.) yield chunk - if memory_enabled: - # Also buffer for post-stream memory processing - buffered_chunks.append(chunk) + # Buffer SSE data for memory processing and/or prefix tracker + _track_sse = memory_enabled or ( + prefix_tracker is not None and provider == "anthropic" + ) + if _track_sse: + if memory_enabled: + buffered_chunks.append(chunk) full_sse_data += chunk_str if len(full_sse_data) > MAX_SSE_BUFFER_SIZE: logger.warning( @@ -696,10 +701,31 @@ class StreamingMixin: # Update prefix cache tracker for next turn (streaming path) if prefix_tracker is not None: + import copy as _copy + + forwarded_messages = body.get("messages", []) + next_forwarded = _copy.deepcopy(forwarded_messages) + next_original = _copy.deepcopy(original_messages or forwarded_messages) + + # Reconstruct assistant response from SSE data so the + # prefix tracker accounts for it in the cached prefix + if full_sse_data and provider == "anthropic": + _parsed = ( + parsed_response + if parsed_response is not None + else self._parse_sse_to_response(full_sse_data, provider) + ) + if _parsed: + asst_msg = self._assistant_message_from_response_json(_parsed) + if asst_msg is not None: + next_forwarded.append(_copy.deepcopy(asst_msg)) + next_original.append(_copy.deepcopy(asst_msg)) + prefix_tracker.update_from_response( cache_read_tokens=cache_read_tokens, cache_write_tokens=cache_write_tokens, - messages=body.get("messages", []), + messages=next_forwarded, + original_messages=next_original, ) if self.cost_tracker: diff --git a/tests/test_cache/test_prefix_tracker.py b/tests/test_cache/test_prefix_tracker.py index 2eb4e63f1..4a79356ea 100644 --- a/tests/test_cache/test_prefix_tracker.py +++ b/tests/test_cache/test_prefix_tracker.py @@ -212,6 +212,66 @@ class TestPrefixCacheTracker: assert len(counts) == 1 assert counts[0] > 100 + def test_estimate_tool_result_content(self): + """Token estimation should count tool_result content field.""" + tool_content = "x" * 3500 # ~1000 tokens + messages = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t1", + "content": tool_content, + } + ], + }, + ] + counts = PrefixCacheTracker._estimate_message_tokens(messages) + assert len(counts) == 1 + # Should be ~1000 tokens, definitely > 100 + assert counts[0] > 100 + + def test_estimate_tool_use_input(self): + """Token estimation should count tool_use input field.""" + messages = [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "t1", + "name": "Read", + "input": {"file_path": "/very/long/path/" + "x" * 700}, + } + ], + }, + ] + counts = PrefixCacheTracker._estimate_message_tokens(messages) + assert len(counts) == 1 + # Should count the serialized input dict + assert counts[0] > 50 + + def test_estimate_tool_result_nested_blocks(self): + """Token estimation should handle nested content blocks in tool_result.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t1", + "content": [ + {"type": "text", "text": "A" * 3500}, + ], + } + ], + }, + ] + counts = PrefixCacheTracker._estimate_message_tokens(messages) + assert len(counts) == 1 + assert counts[0] > 100 + def test_session_ttl_expiry(self): """Tracker should report as expired after TTL.""" config = PrefixFreezeConfig(session_ttl_seconds=1) diff --git a/tests/test_compression_cache.py b/tests/test_compression_cache.py index 1e7f5c45d..6e9f14caf 100644 --- a/tests/test_compression_cache.py +++ b/tests/test_compression_cache.py @@ -172,6 +172,111 @@ class TestCompressionCacheFrozenCount: ] 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, this would stop at msg[1] → frozen=1. + # With stable hash, the walk continues past msg[1] → frozen=3. + assert cache.compute_frozen_count(messages) == 3 + + 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 should now walk past this tool_result + 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) == 3 + + 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 be deferred.""" + h = CompressionCache.content_hash("brand new content") + assert cache.should_defer_compression(h, ttl_seconds=300, batch_window=30) is True + + 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: