headroom/tests/test_ccr.py
Tejas Chopra c2fc4d3753
fix(ccr): make headroom_retrieve a hash-only full-content lookup (#1532)
The optional `query` parameter on headroom_retrieve routed retrieval
through CompressionStore.search(), which BM25-scored the items inside a
single cached blob and dropped everything below a 0.3 relevance floor.
On small per-blob corpora with conversational queries this returned an
empty result the large majority of the time, so the LLM saw "nothing
found" for content that was actually present — pushing users to turn
compression off entirely.

Retrieval is fundamentally a hash lookup (this already matches the Rust
proxy's CCR store, which is put/get only — "no BM25 search"). Remove the
query/search path end to end and always return the full original
content:

Core (Python proxy):
- tool schemas (anthropic/openai/google) drop the `query` property
- parse_tool_call returns the hash (str | None) instead of (hash, query)
- response handler, proxy POST/GET/tool-call handlers, the MCP retrieve
tool, and the streaming feedback recorders retrieve by hash only
- proactive context-tracker expansion always restores full content
- delete CompressionStore.search() and its BM25 machinery (the bm25
module stays — it is still used by relevance/)
- CCRToolCall.query, CCRToolResult.was_search, and
ExpansionRecommendation.expand_full/search_query are removed

Plugins (advertised a now-defunct query param to the LLM):
- hermes (Python), openclaw + opencode (TypeScript) retrieve tools drop
`query` from their schemas, signatures, request URLs, and tests

Benchmarks/docs:
- ccr_regression + adversarial benchmarks switch from store.search() to
full hash retrieval (search input-injection tests repurposed to the
hash, the only remaining input surface)
- wiki/ARCHITECTURE.md, wiki/ccr.md, docs/content/docs/ccr.mdx,
config.py and store docstrings updated to describe hash-only retrieval

Tests updated to assert full-content retrieval and guard the removed
surface; the full CCR/proxy/store/TOIN suite passes. ruff + mypy clean.

## Description

<!-- Briefly explain the change and why it is needed. -->

Closes #

## Type of Change

- [ ] 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

- 

## Testing

<!-- Check what you actually ran, then paste the real command output
below. -->

- [ ] Unit tests pass (`pytest`)
- [ ] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`)
- [ ] New tests added for new functionality
- [ ] Manual testing performed

### Test Output

```text
# Paste relevant command output or artifact links here
```

## Real Behavior Proof

- Environment:
- Exact command / steps:
- Observed result:
- Not tested:

## Review Readiness

- [ ] I have performed a self-review
- [ ] This PR is ready for human review

## Checklist

- [ ] My code follows the project's style guidelines
- [ ] I have performed a self-review of my code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable

## Screenshots (if applicable)

Add screenshots to help explain your changes.

## Additional Notes

<!-- Mention any N/A checklist items, tradeoffs, follow-ups, or
maintainer context. -->
2026-06-28 10:32:43 -07:00

361 lines
11 KiB
Python

"""Tests for Compress-Cache-Retrieve (CCR) architecture.
These tests verify that:
1. CompressionStore correctly caches compressed content
2. SmartCrusher integrates with CompressionStore
3. Retrieval works correctly (full and search)
4. Feedback tracking works
5. TTL expiration works
"""
import json
import time
import pytest
from headroom.cache.compression_store import (
CompressionStore,
get_compression_store,
reset_compression_store,
)
from headroom.config import CCRConfig
from headroom.transforms.smart_crusher import (
SmartCrusherConfig,
smart_crush_tool_output,
)
class TestCompressionStore:
"""Test CompressionStore functionality."""
@pytest.fixture(autouse=True)
def reset_store(self):
"""Reset global store before each test."""
reset_compression_store()
yield
reset_compression_store()
def test_store_and_retrieve(self):
"""Basic store and retrieve flow."""
store = CompressionStore()
original = json.dumps([{"id": i} for i in range(100)])
compressed = json.dumps([{"id": i} for i in range(10)])
hash_key = store.store(
original=original,
compressed=compressed,
original_tokens=1000,
compressed_tokens=100,
original_item_count=100,
compressed_item_count=10,
)
assert (
len(hash_key) == 24
) # SHA256 truncated to 24 chars (96 bits for collision resistance)
entry = store.retrieve(hash_key)
assert entry is not None
assert entry.original_content == original
assert entry.compressed_content == compressed
assert entry.original_tokens == 1000
assert entry.compressed_tokens == 100
def test_retrieve_nonexistent(self):
"""Retrieve returns None for nonexistent hash."""
store = CompressionStore()
entry = store.retrieve("nonexistent1234")
assert entry is None
def test_ttl_expiration(self):
"""Entries expire after TTL."""
store = CompressionStore(default_ttl=1) # 1 second TTL
hash_key = store.store(
original="[1,2,3]",
compressed="[1]",
ttl=1,
)
# Should exist immediately
assert store.exists(hash_key)
# Wait for expiration
time.sleep(1.1)
# Should be expired
assert not store.exists(hash_key)
entry = store.retrieve(hash_key)
assert entry is None
def test_eviction_at_capacity(self):
"""Oldest entries evicted when at capacity."""
store = CompressionStore(max_entries=3)
hashes = []
for i in range(5):
h = store.store(
original=f"original_{i}",
compressed=f"compressed_{i}",
)
hashes.append(h)
time.sleep(0.01) # Ensure different timestamps
# Only last 3 should exist
assert not store.exists(hashes[0])
assert not store.exists(hashes[1])
assert store.exists(hashes[2])
assert store.exists(hashes[3])
assert store.exists(hashes[4])
def test_retrieval_tracking(self):
"""Retrieval events are tracked for feedback."""
store = CompressionStore(enable_feedback=True)
hash_key = store.store(
original="[1,2,3]",
compressed="[1]",
tool_name="test_tool",
)
# Retrieve multiple times
store.retrieve(hash_key)
store.retrieve(hash_key, query="test query")
events = store.get_retrieval_events(limit=10)
assert len(events) >= 2
# Check event details — retrieval is always full
assert any(e.retrieval_type == "full" for e in events)
def test_access_tracking_on_entry(self):
"""Entry tracks access count and queries."""
store = CompressionStore()
hash_key = store.store(
original=json.dumps([{"id": i} for i in range(10)]),
compressed="[]",
)
# Access multiple times with queries
store.retrieve(hash_key, query="first query")
store.retrieve(hash_key, query="second query")
store.retrieve(hash_key, query="first query") # Duplicate
entry = store.retrieve(hash_key)
assert entry.retrieval_count >= 3
assert "first query" in entry.search_queries
assert "second query" in entry.search_queries
def test_stats(self):
"""Store statistics are accurate."""
store = CompressionStore()
store.store(
original="x" * 100,
compressed="x" * 10,
original_tokens=100,
compressed_tokens=10,
)
store.store(
original="y" * 200,
compressed="y" * 20,
original_tokens=200,
compressed_tokens=20,
)
stats = store.get_stats()
assert stats["entry_count"] == 2
assert stats["total_original_tokens"] == 300
assert stats["total_compressed_tokens"] == 30
def test_global_store_singleton(self):
"""Global store uses singleton pattern."""
reset_compression_store()
store1 = get_compression_store()
store2 = get_compression_store()
assert store1 is store2
def test_thread_safety(self):
"""Store is thread-safe."""
import threading
store = CompressionStore()
hashes = []
lock = threading.Lock()
def store_item(i):
h = store.store(
original=f"original_{i}",
compressed=f"compressed_{i}",
)
with lock:
hashes.append(h)
threads = [threading.Thread(target=store_item, args=(i,)) for i in range(10)]
for t in threads:
t.start()
for t in threads:
t.join()
assert len(hashes) == 10
for h in hashes:
assert store.exists(h)
class TestCCRConfig:
"""Test CCR configuration options."""
def test_default_config(self):
"""Default CCR config values."""
config = CCRConfig()
assert config.enabled is True
assert config.store_max_entries == 1000
assert config.store_ttl_seconds == 1800 # session-scale (was 300)
assert config.inject_retrieval_marker is True
assert config.feedback_enabled is True
assert config.min_items_to_cache == 20
class TestCCRFeedbackLoop:
"""Test CCR feedback tracking for learning."""
@pytest.fixture(autouse=True)
def reset_store(self):
"""Reset global store before each test."""
reset_compression_store()
yield
reset_compression_store()
def test_retrieval_events_logged(self):
"""Retrieval events are logged for feedback."""
store = CompressionStore(enable_feedback=True)
items = [{"id": i, "data": f"item_{i}"} for i in range(50)]
hash_key = store.store(
original=json.dumps(items),
compressed=json.dumps(items[:10]),
original_item_count=50,
compressed_item_count=10,
tool_name="search_api",
)
# Simulate retrievals
store.retrieve(hash_key)
store.retrieve(hash_key)
store.retrieve(hash_key)
events = store.get_retrieval_events(limit=10)
# Should have logged all retrievals
assert len(events) >= 3
# Check event types — retrieval is always full
full_events = [e for e in events if e.retrieval_type == "full"]
assert len(full_events) >= 3
def test_tool_name_in_events(self):
"""Tool name is preserved in retrieval events."""
store = CompressionStore(enable_feedback=True)
hash_key = store.store(
original="[1,2,3]",
compressed="[1]",
tool_name="github_search",
)
store.retrieve(hash_key)
events = store.get_retrieval_events(tool_name="github_search")
assert len(events) >= 1
assert all(e.tool_name == "github_search" for e in events)
def test_event_filtering_by_tool(self):
"""Events can be filtered by tool name."""
store = CompressionStore(enable_feedback=True)
hash1 = store.store(
original="[1]",
compressed="[1]",
tool_name="tool_a",
)
hash2 = store.store(
original="[2]",
compressed="[2]",
tool_name="tool_b",
)
store.retrieve(hash1)
store.retrieve(hash1)
store.retrieve(hash2)
tool_a_events = store.get_retrieval_events(tool_name="tool_a")
tool_b_events = store.get_retrieval_events(tool_name="tool_b")
assert len(tool_a_events) == 2
assert len(tool_b_events) == 1
class TestCCREdgeCases:
"""Test edge cases and error handling."""
@pytest.fixture(autouse=True)
def reset_store(self):
"""Reset global store before each test."""
reset_compression_store()
yield
reset_compression_store()
def test_ccr_disabled_no_caching(self):
"""When CCR disabled, no caching occurs."""
reset_compression_store()
items = [{"id": i, "score": 100 - i} for i in range(100)]
content = json.dumps(items)
config = SmartCrusherConfig(max_items_after_crush=15)
ccr_config = CCRConfig(enabled=False) # Disabled
smart_crush_tool_output(content, config, ccr_config)
store = get_compression_store()
stats = store.get_stats()
assert stats["entry_count"] == 0
def test_concurrent_store_and_retrieve(self):
"""Concurrent operations don't corrupt data."""
import threading
store = CompressionStore()
errors = []
def store_and_retrieve(i):
try:
items = [{"id": j, "batch": i} for j in range(10)]
hash_key = store.store(
original=json.dumps(items),
compressed="[]",
tool_name=f"tool_{i}",
)
# Immediately retrieve
entry = store.retrieve(hash_key)
if entry is None:
errors.append(f"Entry {i} not found after store")
elif f'"batch": {i}' not in entry.original_content:
errors.append(f"Entry {i} has wrong content")
except Exception as e:
errors.append(str(e))
threads = [threading.Thread(target=store_and_retrieve, args=(i,)) for i in range(20)]
for t in threads:
t.start()
for t in threads:
t.join()
assert errors == [], f"Errors during concurrent operations: {errors}"