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. -->
This commit is contained in:
Tejas Chopra 2026-06-28 10:32:43 -07:00 committed by GitHub
parent a639540959
commit c2fc4d3753
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
33 changed files with 315 additions and 1381 deletions

View file

@ -252,13 +252,14 @@ def test_needle_looks_exactly_like_hay() -> AdversarialResult:
tool_name="user_search",
)
# Try to find user 456 via search
search_results = store.search(hash_key, "user_id 456")
# Recover the target user via CCR retrieval (hash-only → full content)
entry_for_search = store.retrieve(hash_key)
search_results = json.loads(entry_for_search.original_content) if entry_for_search else []
found_target = any(item.get("user_id") == target_id for item in search_results)
if found_target:
result.actual_behavior = "Found target user via CCR search"
result.actual_behavior = "Found target user via CCR retrieval"
result.passed = True
else:
# Try full retrieval as fallback
@ -718,12 +719,16 @@ def test_extremely_long_strings() -> AdversarialResult:
def test_query_injection_in_search() -> AdversarialResult:
"""
ATTACK: Malicious search query.
ATTACK: Malicious input on the retrieval surface.
Retrieval is hash-only (no query/search parameter), so the only
attacker-controlled input is the hash. A malicious string must never
crash the store or return another entry's data — it must be a clean miss.
"""
result = AdversarialResult(
name="Search Query Injection",
name="Retrieval Hash Injection",
category="injection",
expected_behavior="Should sanitize search queries",
expected_behavior="Malicious hash input is a safe cache miss, never a crash",
severity="high",
)
@ -732,35 +737,36 @@ def test_query_injection_in_search() -> AdversarialResult:
items = [{"id": i, "data": f"item {i}"} for i in range(100)]
hash_key = store.store(
store.store(
original=json.dumps(items),
compressed=json.dumps(items[:10]),
original_item_count=100,
compressed_item_count=10,
)
# Various injection attempts
malicious_queries = [
# Various injection attempts, now aimed at the hash (the only input)
malicious_hashes = [
"'; DROP TABLE items; --",
"<script>alert('xss')</script>",
"{{7*7}}", # Template injection
"${7*7}", # Expression injection
"\\x00\\x01\\x02", # Null bytes
"*" * 10000, # Long query
"*" * 10000, # Long input
".*", # Regex wildcard
"(a]", # Invalid regex
]
failures = []
for query in malicious_queries:
for bad_hash in malicious_hashes:
try:
store.search(hash_key, query)
# If it returns without error, it handled the injection
entry = store.retrieve(bad_hash)
if entry is not None:
failures.append(f"{bad_hash[:20]}: unexpected hit")
except Exception as e:
failures.append(f"{query[:20]}: {type(e).__name__}")
failures.append(f"{bad_hash[:20]}: {type(e).__name__}")
if not failures:
result.actual_behavior = "All malicious queries handled safely"
result.actual_behavior = "All malicious hashes handled safely (clean miss)"
result.passed = True
else:
result.actual_behavior = f"Failures: {failures}"
@ -1261,7 +1267,7 @@ def test_catastrophic_regex_in_search() -> AdversarialResult:
result = AdversarialResult(
name="Regex Catastrophic Backtracking",
category="extreme",
expected_behavior="Should not hang on malicious search patterns",
expected_behavior="Should not hang on malicious hash input",
severity="critical",
)
@ -1270,7 +1276,7 @@ def test_catastrophic_regex_in_search() -> AdversarialResult:
items = [{"id": i, "content": "a" * 50 + "b"} for i in range(100)]
hash_key = store.store(
store.store(
original=json.dumps(items),
compressed=json.dumps(items[:10]),
original_item_count=100,
@ -1278,7 +1284,9 @@ def test_catastrophic_regex_in_search() -> AdversarialResult:
tool_name="regex_test",
)
# These patterns could cause catastrophic backtracking in naive regex
# Retrieval is hash-only, so the only attacker input is the hash. These
# patterns could cause catastrophic backtracking in a naive matcher;
# the hash lookup must not hang on any of them.
evil_patterns = [
"(a+)+$",
"(a|aa)+$",
@ -1290,23 +1298,23 @@ def test_catastrophic_regex_in_search() -> AdversarialResult:
import signal
def timeout_handler(signum, frame):
raise TimeoutError("Search took too long")
raise TimeoutError("Retrieval took too long")
# Set 2 second timeout
old_handler = signal.signal(signal.SIGALRM, timeout_handler)
signal.alarm(2)
for pattern in evil_patterns:
# BM25 search doesn't use regex, so should be safe
store.search(hash_key, pattern)
# Hash lookup is a plain dict/store get — no regex, so it is safe
store.retrieve(pattern)
signal.alarm(0)
signal.signal(signal.SIGALRM, old_handler)
result.actual_behavior = "Search completed without hanging"
result.actual_behavior = "Retrieval completed without hanging"
result.passed = True
except TimeoutError:
result.actual_behavior = "Search hung on regex-like pattern"
result.actual_behavior = "Retrieval hung on regex-like input"
result.passed = False
except Exception as e:
result.actual_behavior = f"Error: {type(e).__name__}: {e}"
@ -1533,7 +1541,6 @@ def test_concurrent_reset_during_operation() -> AdversarialResult:
tool_name="reset_test",
)
store.retrieve(hash_key)
store.search(hash_key, "test")
operations_completed[0] += 1
except Exception as e:
errors.append(f"Op error: {type(e).__name__}: {e}")

View file

@ -11,12 +11,10 @@ does not cause any regression in agent behavior. Specifically:
- Anomalies and outliers
2. RETRIEVAL ACCURACY: When retrieval is needed, correct items are returned
- Full retrieval returns original content
- Search retrieval finds relevant items
- Retrieval is by hash and always returns the full original content
3. FEEDBACK LEARNING: System learns from retrieval patterns
- High retrieval rate triggers less aggressive compression
- Common queries improve future compression
Usage:
python benchmarks/ccr_regression_benchmark.py
@ -73,6 +71,23 @@ class RegressionResult:
failures: list[str] = field(default_factory=list)
def _ccr_retrieve_items(store: Any, hash_key: str) -> list[dict[str, Any]]:
"""Full CCR retrieval (hash-only) → parsed original items.
Retrieval is by hash and always returns the complete original content,
so any "needle" present at compression time is guaranteed to survive the
round-trip. Returns the parsed list, or [] on a miss / non-list payload.
"""
entry = store.retrieve(hash_key)
if not entry:
return []
try:
data = json.loads(entry.original_content)
except (json.JSONDecodeError, TypeError):
return []
return data if isinstance(data, list) else []
# =============================================================================
# TEST 1: Needle in Haystack - Error Retention
# =============================================================================
@ -204,7 +219,7 @@ def test_uuid_retrieval() -> RegressionResult:
)
# Search for the specific UUID
search_results = store.search(hash_key, target_uuid)
search_results = _ccr_retrieve_items(store, hash_key)
result.latency_ms = (time.perf_counter() - start) * 1000
# Check if target UUID was found
@ -476,10 +491,12 @@ def test_feedback_learning() -> RegressionResult:
def test_search_accuracy() -> RegressionResult:
"""
Test that BM25 search within cached content finds relevant items.
Test that hash-keyed retrieval returns the full original content (the
needle is always present in the losslessly-retrieved superset).
"""
result = RegressionResult(
name="Search Accuracy", description="Verify BM25 search finds relevant items in cache"
name="Retrieval Accuracy",
description="Verify hash retrieval returns the full original content from cache",
)
reset_compression_store()
@ -535,7 +552,7 @@ def test_search_accuracy() -> RegressionResult:
start = time.perf_counter()
# Search for authentication errors
search_results = store.search(hash_key, "authentication failed token")
search_results = _ccr_retrieve_items(store, hash_key)
result.latency_ms = (time.perf_counter() - start) * 1000
@ -640,8 +657,8 @@ def test_ccr_end_to_end() -> RegressionResult:
feedback.record_compression("alert_search", 500, 20)
# Step 4: Retrieve and search
critical_results = store.search(hash_key, "critical system overload P0")
error_results = store.search(hash_key, "Error position P1")
critical_results = _ccr_retrieve_items(store, hash_key)
error_results = _ccr_retrieve_items(store, hash_key)
# Step 5: Process feedback
store.process_pending_feedback()

View file

@ -55,8 +55,7 @@ Headroom injects a `headroom_retrieve` tool into the LLM's available tools:
"name": "headroom_retrieve",
"description": "Retrieve original uncompressed data from Headroom cache",
"parameters": {
"hash": "The hash key from the compression marker",
"query": "Optional: search within the cached data"
"hash": "The hash key from the compression marker"
}
}
```
@ -93,22 +92,6 @@ Turn 5: User asks "What about the auth middleware?"
-> LLM finds auth_middleware.py in the full list
```
## BM25 search within compressed data
The LLM does not have to retrieve everything. It can search within compressed data using the optional `query` parameter:
```json
{
"name": "headroom_retrieve",
"parameters": {
"hash": "abc123",
"query": "authentication errors"
}
}
```
This runs a BM25 search over the cached items, returning only the relevant subset instead of the full original payload.
## Retrieving originals
CCR works automatically through the proxy, but you can also retrieve cached data programmatically:

View file

@ -24,11 +24,8 @@ Usage:
tool_name="search_api",
)
# Retrieve later
# Retrieve later (by hash; always returns the full original content)
entry = store.retrieve(hash_key)
# Or search within
results = store.search(hash_key, "user query")
"""
from __future__ import annotations
@ -45,8 +42,6 @@ from contextvars import ContextVar
from dataclasses import dataclass, field, replace
from typing import TYPE_CHECKING, Any
from ..relevance.bm25 import BM25Scorer
if TYPE_CHECKING:
from ..memory.tracker import ComponentStats
from .backends import CompressionStoreBackend
@ -190,7 +185,7 @@ class RetrievalEvent:
total_items: int
tool_name: str | None
timestamp: float
retrieval_type: str # "full" or "search"
retrieval_type: str # always "full" (retrieval is by hash)
tool_signature_hash: str | None = None # For TOIN correlation
@ -206,7 +201,7 @@ class CompressionStore:
- Thread-safe for concurrent access
- TTL-based expiration (default 300 seconds, env-configurable)
- LRU-style eviction when capacity is reached
- Built-in BM25 search for filtering
- Hash-keyed retrieval that always returns the full original content
"""
def __init__(
@ -250,9 +245,6 @@ class CompressionStore:
# Threshold for triggering heap rebuild (when 50% are stale)
self._heap_rebuild_threshold = 0.5
# BM25 scorer for search
self._scorer = BM25Scorer()
@property
def default_ttl_seconds(self) -> int:
"""Default TTL applied to new entries when callers do not override it."""
@ -487,74 +479,6 @@ class CompressionStore:
"ttl": entry.ttl,
}
def search(
self,
hash_key: str,
query: str,
max_results: int = 20,
score_threshold: float = 0.3,
) -> list[dict[str, Any]]:
"""Search within cached content using BM25.
Args:
hash_key: Hash key of cached content.
query: Search query.
max_results: Maximum number of results to return.
score_threshold: Minimum BM25 score to include.
Returns:
List of matching items from original content.
"""
# Get entry without logging (we'll log the search separately)
entry = self._get_entry_for_search(hash_key, query)
if entry is None:
return []
items = self._search_items_from_original(entry.original_content)
if not items:
return []
# Score each item using BM25
item_strs = [json.dumps(item, default=str) for item in items]
scores = self._scorer.score_batch(item_strs, query)
# Filter and sort by score
scored_items = [
(items[i], scores[i].score)
for i in range(len(items))
if scores[i].score >= score_threshold
]
scored_items.sort(key=lambda x: x[1], reverse=True)
results = [item for item, _ in scored_items[:max_results]]
# Log retrieval event
if self._enable_feedback:
with self._lock:
self._log_retrieval(
hash_key=hash_key,
query=query,
items_retrieved=len(results),
total_items=len(items),
tool_name=entry.tool_name,
retrieval_type="search",
tool_signature_hash=entry.tool_signature_hash,
)
# Process feedback immediately to ensure TOIN learns in real-time
self.process_pending_feedback()
self._log_retrieval_payload(
hash_key=hash_key,
query=query,
retrieval_type="search",
payload=json.dumps(results, ensure_ascii=False),
items_retrieved=len(results),
total_items=len(items),
entry=entry,
)
return results
def _log_retrieval_payload(
self,
*,
@ -588,189 +512,6 @@ class CompressionStore:
json.dumps(event, ensure_ascii=False, separators=(",", ":")),
)
def _search_items_from_original(self, original_content: str) -> list[Any]:
"""Normalize cached originals into searchable items.
CCR producers store different shapes:
- SmartCrusher/search-style paths usually store JSON arrays.
- Kompress stores the original plain text.
- Some callers store JSON objects or scalar JSON values.
Search should work for all of them. Preserve the legacy JSON-array
result shape, but fall back to structured text chunks for everything
else so `headroom_retrieve(hash, query=...)` can find plain-text
originals.
"""
try:
parsed = json.loads(original_content)
except json.JSONDecodeError:
return self._plain_text_search_items(original_content)
if isinstance(parsed, list):
return parsed
if isinstance(parsed, dict):
return self._json_object_search_items(parsed)
if isinstance(parsed, str):
return self._plain_text_search_items(parsed)
if parsed is None:
return []
return [{"type": "json_scalar", "value": parsed}]
def _json_object_search_items(self, value: dict[str, Any]) -> list[dict[str, Any]]:
"""Return searchable leaf records for a JSON object."""
items: list[dict[str, Any]] = []
def walk(node: Any, path: str) -> None:
if isinstance(node, dict):
for key, child in node.items():
child_path = f"{path}.{key}" if path else str(key)
walk(child, child_path)
return
if isinstance(node, list):
for idx, child in enumerate(node):
walk(child, f"{path}[{idx}]")
return
if node is None:
return
items.append({"type": "json_leaf", "path": path, "value": node})
walk(value, "")
if items:
return items
return [{"type": "json_object", "value": value}]
def _plain_text_search_items(self, text: str) -> list[dict[str, Any]]:
"""Chunk arbitrary text into searchable records.
Line-aware chunks work well for logs/source. Word-window chunks handle
Kompress originals, which are often long single-line text blobs.
"""
if not text or not text.strip():
return []
normalized = text.replace("\r\n", "\n").replace("\r", "\n")
lines = normalized.split("\n")
if len(lines) > 1:
return self._line_text_search_items(lines)
words = normalized.split()
if not words:
return []
max_words = 350
overlap_words = 50
if len(words) <= max_words:
return [
{
"type": "text",
"text": normalized,
"chunk_index": 0,
"word_start": 1,
"word_end": len(words),
}
]
items: list[dict[str, Any]] = []
start = 0
chunk_index = 0
step = max_words - overlap_words
while start < len(words):
end = min(len(words), start + max_words)
items.append(
{
"type": "text",
"text": " ".join(words[start:end]),
"chunk_index": chunk_index,
"word_start": start + 1,
"word_end": end,
}
)
if end == len(words):
break
start += step
chunk_index += 1
return items
@staticmethod
def _line_text_search_items(lines: list[str]) -> list[dict[str, Any]]:
max_chars = 2000
items: list[dict[str, Any]] = []
current: list[str] = []
line_start = 1
char_count = 0
for idx, line in enumerate(lines, start=1):
line_len = len(line) + 1
if current and char_count + line_len > max_chars:
items.append(
{
"type": "text",
"text": "\n".join(current),
"chunk_index": len(items),
"line_start": line_start,
"line_end": idx - 1,
}
)
current = []
line_start = idx
char_count = 0
current.append(line)
char_count += line_len
if current:
items.append(
{
"type": "text",
"text": "\n".join(current),
"chunk_index": len(items),
"line_start": line_start,
"line_end": len(lines),
}
)
return items
def _get_entry_for_search(
self,
hash_key: str,
query: str | None = None,
) -> CompressionEntry | None:
"""Get entry without logging retrieval (used by search to avoid double-logging).
CRITICAL FIX #4: Returns a copy of the entry to prevent race conditions.
The caller may use the entry after we release the lock, and another thread
could modify or evict the original entry.
Args:
hash_key: Hash key returned by store().
query: Optional query for access tracking.
Returns:
CompressionEntry copy if found and not expired, None otherwise.
"""
with self._lock:
entry = self._backend.get(hash_key)
if entry is None:
return None
if entry.is_expired():
self._backend.delete(hash_key)
# CRITICAL FIX: Track stale heap entry
self._stale_heap_entries += 1
return None
# Track access but don't log retrieval event (search will log separately)
entry.record_access(query)
# Update the backend with the modified entry
self._backend.set(hash_key, entry)
# CRITICAL FIX #4: Return a copy to prevent race conditions
# The entry contains mutable fields (search_queries list) that could be
# modified by other threads after we release the lock
return replace(entry, search_queries=list(entry.search_queries))
def exists(self, hash_key: str, clean_expired: bool = False) -> bool:
"""Check if a hash key exists and is not expired.

View file

@ -20,7 +20,6 @@ Example:
from __future__ import annotations
import json
import logging
import re
import time
@ -66,8 +65,6 @@ class ExpansionRecommendation:
hash_key: str
reason: str
relevance_score: float
expand_full: bool = True # True = expand all, False = search only
search_query: str | None = None
@dataclass
@ -256,18 +253,11 @@ class ContextTracker:
relevance *= age_factor
if relevance >= self.config.relevance_threshold:
# Determine if full expansion or search
expand_full, search_query = self._determine_expansion_type(
query, context, relevance
)
recommendations.append(
ExpansionRecommendation(
hash_key=hash_key,
reason=self._generate_reason(query, context, relevance),
relevance_score=relevance,
expand_full=expand_full,
search_query=search_query,
)
)
@ -436,39 +426,6 @@ class ContextTracker:
return [w for w in words if w not in stop_words and len(w) >= 2]
def _determine_expansion_type(
self,
query: str,
context: CompressedContext,
relevance: float,
) -> tuple[bool, str | None]:
"""Determine whether to do full expansion or search.
Returns:
Tuple of (expand_full, search_query)
"""
# High relevance + small original count = full expansion
if relevance > 0.6 or context.original_item_count <= 50:
return True, None
# Extract specific search terms from query
keywords = self._extract_keywords(query.lower())
# Filter to most specific keywords (longer, less common)
specific_keywords = [
k
for k in keywords
if len(k) >= 4 and k not in {"file", "code", "show", "find", "list", "what"}
]
if specific_keywords:
# Use top keywords as search query
search_query = " ".join(specific_keywords[:3])
return False, search_query
# Default to full expansion if we can't form a good search
return True, None
def _generate_reason(
self,
query: str,
@ -509,39 +466,23 @@ class ContextTracker:
for rec in recommendations:
try:
if rec.expand_full:
entry = store.retrieve(rec.hash_key)
if entry:
results.append(
{
"hash": rec.hash_key,
"type": "full",
"content": entry.original_content,
"item_count": entry.original_item_count,
"reason": rec.reason,
}
)
logger.info(
f"CCR Tracker: Proactively expanded {rec.hash_key} "
f"({entry.original_item_count} items)"
)
else:
search_results = store.search(rec.hash_key, rec.search_query or "")
if search_results:
results.append(
{
"hash": rec.hash_key,
"type": "search",
"query": rec.search_query,
"content": search_results,
"item_count": len(search_results),
"reason": rec.reason,
}
)
logger.info(
f"CCR Tracker: Proactive search in {rec.hash_key} "
f"for '{rec.search_query}' ({len(search_results)} results)"
)
# Retrieval is by hash: proactive expansion always restores the
# full original content (no partial/search expansion).
entry = store.retrieve(rec.hash_key)
if entry:
results.append(
{
"hash": rec.hash_key,
"type": "full",
"content": entry.original_content,
"item_count": entry.original_item_count,
"reason": rec.reason,
}
)
logger.info(
f"CCR Tracker: Proactively expanded {rec.hash_key} "
f"({entry.original_item_count} items)"
)
except Exception as e:
logger.warning(f"CCR Tracker: Failed to expand {rec.hash_key}: {e}")
@ -577,15 +518,9 @@ class ContextTracker:
parts = [header]
for exp in expansions:
if exp["type"] == "full":
parts.append(f"\n--- Expanded from earlier ({exp['reason']}) ---")
parts.append(exp["content"])
else:
parts.append(f"\n--- Search results for '{exp['query']}' ({exp['reason']}) ---")
if isinstance(exp["content"], list):
parts.append(json.dumps(exp["content"], indent=2))
else:
parts.append(str(exp["content"]))
# Expansions are always full (retrieval is by hash).
parts.append(f"\n--- Expanded from earlier ({exp['reason']}) ---")
parts.append(exp["content"])
parts.append("[End Proactive Expansion]")
body = "\n".join(parts)

View file

@ -415,56 +415,29 @@ class HeadroomMCPServer:
async def _retrieve_content(
self,
hash_key: str,
query: str | None,
) -> dict[str, Any]:
"""Retrieve content. Checks local store first, then proxy."""
"""Retrieve content by hash. Checks local store first, then proxy.
Retrieval is by hash and always returns the full original content.
"""
# Check local store first
store = self._get_local_store()
if query:
results = store.search(hash_key, query)
if results:
self._stats.record_retrieval(hash_key)
return {
"hash": hash_key,
"source": "local",
"query": query,
"results": results,
"count": len(results),
}
# The query matched no items above the relevance floor, but the
# entry itself may still be present and unexpired. An empty search
# is not the same as a missing/expired hash, so fall back to the
# full content rather than reporting it as not found.
entry = store.retrieve(hash_key)
if entry:
self._stats.record_retrieval(hash_key)
return {
"hash": hash_key,
"source": "local",
"query": query,
"results": [],
"count": 0,
"original_content": entry.original_content,
"note": "Entry exists but no item matched the query above "
"the relevance threshold; returning the full content.",
}
else:
entry = store.retrieve(hash_key)
if entry:
self._stats.record_retrieval(hash_key)
return {
"hash": hash_key,
"source": "local",
"original_content": entry.original_content,
"original_item_count": entry.original_item_count,
"compressed_item_count": entry.compressed_item_count,
"retrieval_count": entry.retrieval_count,
}
entry = store.retrieve(hash_key)
if entry:
self._stats.record_retrieval(hash_key)
return {
"hash": hash_key,
"source": "local",
"original_content": entry.original_content,
"original_item_count": entry.original_item_count,
"compressed_item_count": entry.compressed_item_count,
"retrieval_count": entry.retrieval_count,
}
# Fall back to proxy if available
if self.check_proxy and HTTPX_AVAILABLE:
try:
result = await self._retrieve_via_proxy(hash_key, query)
result = await self._retrieve_via_proxy(hash_key)
if "error" not in result:
result["source"] = "proxy"
self._stats.record_retrieval(hash_key)
@ -485,16 +458,13 @@ class HeadroomMCPServer:
async def _retrieve_via_proxy(
self,
hash_key: str,
query: str | None,
) -> dict[str, Any]:
"""Retrieve content via proxy's HTTP endpoint."""
"""Retrieve full content by hash via proxy's HTTP endpoint."""
if self._http_client is None:
self._http_client = httpx.AsyncClient(timeout=15.0)
url = f"{self.proxy_url}/v1/retrieve"
payload: dict[str, str] = {"hash": hash_key}
if query:
payload["query"] = query
response = await self._http_client.post(url, json=payload)
@ -549,13 +519,6 @@ class HeadroomMCPServer:
"type": "string",
"description": "Hash key from compression (e.g., 'abc123' from hash=abc123)",
},
"query": {
"type": "string",
"description": (
"Optional search query to filter results. "
"If provided, returns only items matching the query."
),
},
},
"required": ["hash"],
},
@ -723,17 +686,11 @@ class HeadroomMCPServer:
)
]
query = arguments.get("query")
logger.info("event=mcp_retrieve_started hash=%s", hash_key)
result = await self._retrieve_content(hash_key)
logger.info(
"event=mcp_retrieve_started hash=%s query=%s",
"event=mcp_retrieve_completed hash=%s result=%s",
hash_key,
json.dumps(query, ensure_ascii=False, default=str),
)
result = await self._retrieve_content(hash_key, query)
logger.info(
"event=mcp_retrieve_completed hash=%s query=%s result=%s",
hash_key,
json.dumps(query, ensure_ascii=False, default=str),
json.dumps(result, ensure_ascii=False, default=str),
)

View file

@ -31,7 +31,6 @@ class CCRToolCall:
tool_call_id: str
hash_key: str
query: str | None = None
@dataclass
@ -42,7 +41,6 @@ class CCRToolResult:
content: str
success: bool
items_retrieved: int = 0
was_search: bool = False
@dataclass
@ -166,7 +164,7 @@ class CCRResponseHandler:
other_calls = []
for tc in all_tool_calls:
hash_key, query = parse_tool_call(tc, provider)
hash_key = parse_tool_call(tc, provider)
if hash_key is not None:
# This is a CCR tool call - extract tool_call_id based on provider
@ -181,7 +179,6 @@ class CCRResponseHandler:
CCRToolCall(
tool_call_id=tool_call_id,
hash_key=hash_key,
query=query,
)
)
else:
@ -224,15 +221,14 @@ class CCRResponseHandler:
success=False,
)
if ccr_call.query:
# Search within compressed content
results = store.search(ccr_call.hash_key, ccr_call.query)
# Retrieval is by hash: always return the full original content.
entry = store.retrieve(ccr_call.hash_key)
if entry:
content = json.dumps(
{
"hash": ccr_call.hash_key,
"query": ccr_call.query,
"results": results,
"count": len(results),
"original_content": entry.original_content,
"original_item_count": entry.original_item_count,
},
indent=2,
)
@ -240,48 +236,28 @@ class CCRResponseHandler:
tool_call_id=ccr_call.tool_call_id,
content=content,
success=True,
items_retrieved=len(results),
was_search=True,
items_retrieved=entry.original_item_count,
)
else:
# Full retrieval
entry = store.retrieve(ccr_call.hash_key)
if entry:
content = json.dumps(
{
"hash": ccr_call.hash_key,
"original_content": entry.original_content,
"original_item_count": entry.original_item_count,
},
indent=2,
)
return CCRToolResult(
tool_call_id=ccr_call.tool_call_id,
content=content,
success=True,
items_retrieved=entry.original_item_count,
was_search=False,
)
else:
miss_status = (
get_status(ccr_call.hash_key, clean_expired=True)
if callable(get_status)
else {"hash": ccr_call.hash_key, "status": "missing"}
)
content = json.dumps(
{
"error": format_retrieval_miss_detail(miss_status),
"hash": ccr_call.hash_key,
"status": miss_status["status"],
"ttl_seconds": miss_status.get("ttl_seconds"),
},
indent=2,
)
return CCRToolResult(
tool_call_id=ccr_call.tool_call_id,
content=content,
success=False,
)
miss_status = (
get_status(ccr_call.hash_key, clean_expired=True)
if callable(get_status)
else {"hash": ccr_call.hash_key, "status": "missing"}
)
content = json.dumps(
{
"error": format_retrieval_miss_detail(miss_status),
"hash": ccr_call.hash_key,
"status": miss_status["status"],
"ttl_seconds": miss_status.get("ttl_seconds"),
},
indent=2,
)
return CCRToolResult(
tool_call_id=ccr_call.tool_call_id,
content=content,
success=False,
)
except Exception as e:
logger.error(f"CCR retrieval failed for {ccr_call.hash_key}: {e}")
@ -485,10 +461,8 @@ class CCRResponseHandler:
# Log retrieval stats
total_items = sum(r.items_retrieved for r in results)
searches = sum(1 for r in results if r.was_search)
logger.debug(
f"CCR: Retrieved {total_items} items "
f"({searches} searches, {len(results) - searches} full)"
f"CCR: Retrieved {total_items} items across {len(results)} full retrieval(s)"
)
# Build continuation messages

View file

@ -55,14 +55,6 @@ def create_ccr_tool_definition(
"type": "string",
"description": "Hash key from the compression marker (e.g., 'abc123' from hash=abc123)",
},
"query": {
"type": "string",
"description": (
"Optional search query to filter results. "
"If provided, only returns items matching the query. "
"If omitted, returns all original items."
),
},
},
"required": ["hash"],
},
@ -88,14 +80,6 @@ def create_ccr_tool_definition(
"type": "string",
"description": "Hash key from the compression marker (e.g., 'abc123' from hash=abc123)",
},
"query": {
"type": "string",
"description": (
"Optional search query to filter results. "
"If provided, only returns items matching the query. "
"If omitted, returns all original items."
),
},
},
"required": ["hash"],
},
@ -116,10 +100,6 @@ def create_ccr_tool_definition(
"type": "string",
"description": "Hash key from the compression marker",
},
"query": {
"type": "string",
"description": "Optional search query to filter results",
},
},
"required": ["hash"],
},
@ -155,8 +135,7 @@ Some tool outputs have been compressed to reduce context size. If you need
the full uncompressed data, you can retrieve it using the `{CCR_TOOL_NAME}` tool.
**How to retrieve:**
- Call `{CCR_TOOL_NAME}(hash="<hash>")` to get all original items
- Call `{CCR_TOOL_NAME}(hash="<hash>", query="search terms")` to search within
- Call `{CCR_TOOL_NAME}(hash="<hash>")` to get the full original content back
**Available hashes:** {hash_list}
@ -465,15 +444,15 @@ class CCRToolInjector:
def parse_tool_call(
tool_call: dict[str, Any],
provider: str = "anthropic",
) -> tuple[str | None, str | None]:
"""Parse a CCR tool call to extract hash and query.
) -> str | None:
"""Parse a CCR tool call to extract the content hash.
Args:
tool_call: The tool call object from the LLM response.
provider: The provider type for format detection.
Returns:
Tuple of (hash, query) or (None, None) if not a CCR tool call.
The hash key, or None if this is not a (valid) CCR tool call.
"""
# Get tool name and input data based on provider format
if provider == "anthropic":
@ -499,19 +478,19 @@ def parse_tool_call(
input_data = tool_call.get("input", tool_call.get("args", {}))
if name != CCR_TOOL_NAME:
return None, None
return None
hash_key = input_data.get("hash")
query = input_data.get("query")
if hash_key is None:
return None
# Validate hash format. SmartCrusher emits 12-hex-char hashes while legacy
# bracket markers / the compression_store use 24-hex-char hashes; accept
# either real length and reject anything else as malformed.
if hash_key is not None:
if not isinstance(hash_key, str) or len(hash_key) not in (12, 24):
return None, None
# Validate hex characters only
if not all(c in "0123456789abcdef" for c in hash_key.lower()):
return None, None
if not isinstance(hash_key, str) or len(hash_key) not in (12, 24):
return None
# Validate hex characters only
if not all(c in "0123456789abcdef" for c in hash_key.lower()):
return None
return hash_key, query
return hash_key

View file

@ -488,7 +488,7 @@ class CCRConfig:
1. COMPRESS: SmartCrusher compresses array from 1000 to 20 items
2. CACHE: Original 1000 items stored in CompressionStore
3. INJECT: Marker added to tell LLM how to retrieve more
4. RETRIEVE: If LLM needs more, it calls headroom_retrieve(hash, query)
4. RETRIEVE: If LLM needs more, it calls headroom_retrieve(hash) to get the full original back
Benefits:
- Zero-risk compression: worst case = LLM retrieves what it needs

View file

@ -588,24 +588,17 @@ class StreamingMixin:
input_data = block.get("input", {})
hash_key = input_data.get("hash")
query = input_data.get("query")
if not hash_key:
continue
logger.info(
f"[{request_id}] CCR Feedback: Recording retrieval "
f"hash={hash_key[:8]}... query={query!r}"
)
logger.info(f"[{request_id}] CCR Feedback: Recording retrieval hash={hash_key[:8]}...")
# Call store.retrieve()/search() for the side effect of triggering
# the feedback chain: _log_retrieval -> process_pending_feedback
# Call store.retrieve() for the side effect of triggering the
# feedback chain: _log_retrieval -> process_pending_feedback
# -> toin.record_retrieval(). We discard the returned content.
try:
if query:
store.search(hash_key, query)
else:
store.retrieve(hash_key, query=None)
store.retrieve(hash_key)
except Exception as e:
logger.debug(f"[{request_id}] CCR Feedback recording failed: {e}")
@ -670,19 +663,15 @@ class StreamingMixin:
if not isinstance(input_data, dict):
continue
hash_key = input_data.get("hash")
query = input_data.get("query")
if not hash_key:
continue
logger.info(
f"[{request_id}] CCR Feedback (openai stream): Recording retrieval "
f"hash={hash_key[:8]}... query={query!r}"
f"hash={hash_key[:8]}..."
)
try:
if query:
store.search(hash_key, query)
else:
store.retrieve(hash_key, query=None)
store.retrieve(hash_key)
except Exception as e:
logger.debug(f"[{request_id}] CCR Feedback (openai stream) failed: {e}")

View file

@ -3406,15 +3406,12 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
Request body:
hash (str): Hash key from compression marker (required)
query (str): Optional search query to filter results
Response:
Full retrieval: {"hash": "...", "original_content": "...", ...}
Search: {"hash": "...", "query": "...", "results": [...], "count": N}
{"hash": "...", "original_content": "...", ...}
"""
data = await request.json()
hash_key = data.get("hash")
query = data.get("query")
if not hash_key:
raise HTTPException(status_code=400, detail="hash required")
@ -3428,36 +3425,24 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
detail=format_retrieval_miss_detail(entry_status),
)
if query:
# Search within cached content. The get_entry_status check above
# (clean_expired=True) already guaranteed availability or raised
# 404, so no second exists()/status backend read is needed here.
results = store.search(hash_key, query)
# Retrieval is by hash: always return the full original content.
entry = store.retrieve(hash_key)
if entry:
return {
"hash": hash_key,
"query": query,
"results": results,
"count": len(results),
"original_content": entry.original_content,
"original_tokens": entry.original_tokens,
"original_item_count": entry.original_item_count,
"compressed_item_count": entry.compressed_item_count,
"tool_name": entry.tool_name,
"retrieval_count": entry.retrieval_count,
}
else:
# Return full original content
entry = store.retrieve(hash_key)
if entry:
return {
"hash": hash_key,
"original_content": entry.original_content,
"original_tokens": entry.original_tokens,
"original_item_count": entry.original_item_count,
"compressed_item_count": entry.compressed_item_count,
"tool_name": entry.tool_name,
"retrieval_count": entry.retrieval_count,
}
raise HTTPException(
status_code=404,
detail=format_retrieval_miss_detail(
store.get_entry_status(hash_key, clean_expired=True)
),
)
raise HTTPException(
status_code=404,
detail=format_retrieval_miss_detail(
store.get_entry_status(hash_key, clean_expired=True)
),
)
@app.get("/v1/retrieve/stats", dependencies=[Depends(_require_loopback)])
async def ccr_stats():
@ -3739,7 +3724,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
)
@app.get("/v1/retrieve/{hash_key}", dependencies=[Depends(_require_loopback)])
async def ccr_retrieve_get(hash_key: str, query: str | None = None):
async def ccr_retrieve_get(hash_key: str):
"""GET version of CCR retrieve for easier testing."""
store = get_compression_store()
entry_status = store.get_entry_status(hash_key, clean_expired=True)
@ -3750,32 +3735,24 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
detail=format_retrieval_miss_detail(entry_status),
)
if query:
results = store.search(hash_key, query)
# Retrieval is by hash: always return the full original content.
entry = store.retrieve(hash_key)
if entry:
return {
"hash": hash_key,
"query": query,
"results": results,
"count": len(results),
"original_content": entry.original_content,
"original_tokens": entry.original_tokens,
"original_item_count": entry.original_item_count,
"compressed_item_count": entry.compressed_item_count,
"tool_name": entry.tool_name,
"retrieval_count": entry.retrieval_count,
}
else:
entry = store.retrieve(hash_key)
if entry:
return {
"hash": hash_key,
"original_content": entry.original_content,
"original_tokens": entry.original_tokens,
"original_item_count": entry.original_item_count,
"compressed_item_count": entry.compressed_item_count,
"tool_name": entry.tool_name,
"retrieval_count": entry.retrieval_count,
}
raise HTTPException(
status_code=404,
detail=format_retrieval_miss_detail(
store.get_entry_status(hash_key, clean_expired=True)
),
)
raise HTTPException(
status_code=404,
detail=format_retrieval_miss_detail(
store.get_entry_status(hash_key, clean_expired=True)
),
)
# CCR Tool Call Handler - for agent frameworks to call when LLM uses headroom_retrieve
@app.post("/v1/retrieve/tool_call", dependencies=[Depends(_require_loopback)])
@ -3791,7 +3768,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
"tool_call": {
"id": "toolu_123",
"name": "headroom_retrieve",
"input": {"hash": "abc123", "query": "optional search"}
"input": {"hash": "abc123"}
},
"provider": "anthropic"
}
@ -3820,7 +3797,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
provider = data.get("provider", "anthropic")
# Parse the tool call
hash_key, query = parse_tool_call(tool_call, provider)
hash_key = parse_tool_call(tool_call, provider)
if hash_key is None:
raise HTTPException(
@ -3838,15 +3815,8 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
"status": entry_status["status"],
"ttl_seconds": entry_status.get("ttl_seconds", entry_status["default_ttl_seconds"]),
}
elif query:
results = store.search(hash_key, query)
retrieval_data = {
"hash": hash_key,
"query": query,
"results": results,
"count": len(results),
}
else:
# Retrieval is by hash: always return the full original content.
entry = store.retrieve(hash_key)
if entry:
retrieval_data = {

View file

@ -55,7 +55,7 @@ HEADROOM_EXCLUDE_TOOLS=read_file,headroom_retrieve
## Behavior
- Accepts the bare hash or the whole marker — `<<ccr:abc123,base64,4.5KB>>`, `ccr:abc123`, and `hash=abc123` are all normalized to `abc123`.
- Optional `query` parameter filters very large results via the proxy's BM25 search.
- Retrieval is by hash and always returns the full original content.
- Clear, actionable errors: expired hash (TTL) and proxy-unreachable cases both tell the model to re-run the original command instead of retrying blindly.
## Requirements

View file

@ -27,10 +27,9 @@ HEADROOM_RETRIEVE_SCHEMA = {
"to cat/read them. When you see one in a tool result or in "
"conversation history, call this tool with the hash (the hex string "
"after 'hash=' or 'ccr:') to read the full original content instead "
"of guessing or re-running the command. For very large results, pass "
"the optional 'query' to filter to the relevant parts (BM25 search). "
"Content expires after a TTL — if expired, re-run the original "
"command instead."
"of guessing or re-running the command. Retrieval is by hash and "
"always returns the complete original content. Content expires after "
"a TTL — if expired, re-run the original command instead."
),
"parameters": {
"type": "object",
@ -39,10 +38,6 @@ HEADROOM_RETRIEVE_SCHEMA = {
"type": "string",
"description": "Hash from the compression marker, e.g. 'abc123' from '[... hash=abc123]' or '<<ccr:abc123>>'",
},
"query": {
"type": "string",
"description": "Optional search query to filter large results to relevant items",
},
},
"required": ["hash"],
},
@ -61,9 +56,6 @@ def _handle_headroom_retrieve(args: dict, **kw) -> str:
)
payload: dict = {"hash": hash_key}
query = str(args.get("query") or "").strip()
if query:
payload["query"] = query
try:
resp = httpx.post(f"{_PROXY_URL}/v1/retrieve", json=payload, timeout=15)

View file

@ -21,7 +21,7 @@ export function createHeadroomRetrieveTool(config: RetrieveToolConfig) {
"Retrieve original uncompressed content from Headroom's compression store. " +
"Use when compressed context mentions a hash and you need the full details. " +
"Pass the hash from the compression marker (24 hex characters). " +
"Optionally pass a query to search within the original content.",
"Retrieval is by hash and always returns the full original content.",
parameters: {
type: "object" as const,
properties: {
@ -29,15 +29,11 @@ export function createHeadroomRetrieveTool(config: RetrieveToolConfig) {
type: "string",
description: "The 24-character hex hash from the compression marker",
},
query: {
type: "string",
description: "Optional search query to filter results within the original content",
},
},
required: ["hash"],
},
execute: async (args: { hash: string; query?: string }): Promise<string> => {
const { hash, query } = args;
execute: async (args: { hash: string }): Promise<string> => {
const { hash } = args;
// Validate hash format
if (!/^[a-f0-9]{24}$/i.test(hash)) {
@ -47,9 +43,7 @@ export function createHeadroomRetrieveTool(config: RetrieveToolConfig) {
}
try {
const url = query
? `${proxyOrigin}/v1/retrieve/${hash}?query=${encodeURIComponent(query)}`
: `${proxyOrigin}/v1/retrieve/${hash}`;
const url = `${proxyOrigin}/v1/retrieve/${hash}`;
const resp = await fetch(url, {
signal: AbortSignal.timeout(10_000),

View file

@ -65,7 +65,6 @@ const retrieve = createHeadroomRetrieveTool({
const result = await retrieve.execute({
hash: "0123456789abcdef01234567",
query: "needle",
});
```

View file

@ -55,13 +55,13 @@ describe("HeadroomPlugin", () => {
proxyUrl: "http://127.0.0.1:8787",
});
const result = await plugin.tool?.headroom_retrieve.execute(
{ hash: "0123456789abcdef01234567", query: "needle" },
{ hash: "0123456789abcdef01234567" },
{} as never,
);
expect(result).toBe("original content");
expect(fetchMock).toHaveBeenCalledWith(
"http://127.0.0.1:8787/v1/retrieve/0123456789abcdef01234567?query=needle",
"http://127.0.0.1:8787/v1/retrieve/0123456789abcdef01234567",
expect.any(Object),
);
});

View file

@ -45,7 +45,6 @@ export const HeadroomPlugin: Plugin = async (input, options = {}) => {
hash: z
.string()
.regex(/^[a-f0-9]{24}$/i, "Expected 24-character hex hash"),
query: z.string().optional(),
},
async execute(args) {
return retrieveTool.execute(args);

View file

@ -24,7 +24,7 @@ export function createHeadroomRetrieveTool(config: RetrieveToolConfig) {
"Retrieve original uncompressed content from Headroom's compression store. " +
"Use when compressed context mentions a hash and you need the full details. " +
"Pass the hash from the compression marker (24 hex characters). " +
"Optionally pass a query to search within the original content.",
"Retrieval is by hash and always returns the full original content.",
parameters: {
type: "object" as const,
properties: {
@ -32,15 +32,11 @@ export function createHeadroomRetrieveTool(config: RetrieveToolConfig) {
type: "string",
description: "The 24-character hex hash from the compression marker",
},
query: {
type: "string",
description: "Optional search query to filter results within the original content",
},
},
required: ["hash"],
},
execute: async (args: { hash: string; query?: string }): Promise<string> => {
const { hash, query } = args;
execute: async (args: { hash: string }): Promise<string> => {
const { hash } = args;
if (!/^[a-f0-9]{24}$/i.test(hash)) {
return JSON.stringify({
@ -49,9 +45,7 @@ export function createHeadroomRetrieveTool(config: RetrieveToolConfig) {
}
try {
const url = query
? `${origin}/v1/retrieve/${hash}?query=${encodeURIComponent(query)}`
: `${origin}/v1/retrieve/${hash}`;
const url = `${origin}/v1/retrieve/${hash}`;
const resp = await fetch(url, {
signal: AbortSignal.timeout(10_000),

View file

@ -109,33 +109,6 @@ class TestCompressionStore:
assert store.exists(hashes[3])
assert store.exists(hashes[4])
def test_search_with_bm25(self):
"""Search within cached content using BM25."""
store = CompressionStore()
items = [
{"id": 1, "content": "Python programming language"},
{"id": 2, "content": "JavaScript web development"},
{"id": 3, "content": "Python data science pandas"},
{"id": 4, "content": "Java enterprise applications"},
{"id": 5, "content": "Python machine learning tensorflow"},
]
hash_key = store.store(
original=json.dumps(items),
compressed=json.dumps(items[:2]),
original_item_count=5,
compressed_item_count=2,
)
# Search for Python items
results = store.search(hash_key, "Python programming")
assert len(results) >= 1
# Should prioritize Python items
result_ids = [r["id"] for r in results]
assert 1 in result_ids # "Python programming language"
def test_retrieval_tracking(self):
"""Retrieval events are tracked for feedback."""
store = CompressionStore(enable_feedback=True)
@ -149,14 +122,12 @@ class TestCompressionStore:
# Retrieve multiple times
store.retrieve(hash_key)
store.retrieve(hash_key, query="test query")
store.search(hash_key, "another query")
events = store.get_retrieval_events(limit=10)
assert len(events) >= 2
# Check event details
# Check event details — retrieval is always full
assert any(e.retrieval_type == "full" for e in events)
assert any(e.retrieval_type == "search" for e in events)
def test_access_tracking_on_entry(self):
"""Entry tracks access count and queries."""
@ -274,20 +245,18 @@ class TestCCRFeedbackLoop:
# Simulate retrievals
store.retrieve(hash_key)
store.search(hash_key, "specific query")
store.search(hash_key, "another query")
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
# Check event types — retrieval is always full
full_events = [e for e in events if e.retrieval_type == "full"]
search_events = [e for e in events if e.retrieval_type == "search"]
assert len(full_events) >= 1
assert len(search_events) >= 2
assert len(full_events) >= 3
def test_tool_name_in_events(self):
"""Tool name is preserved in retrieval events."""
@ -341,59 +310,6 @@ class TestCCREdgeCases:
yield
reset_compression_store()
def test_search_expired_entry(self):
"""Search on expired entry returns empty."""
store = CompressionStore(default_ttl=1)
hash_key = store.store(
original=json.dumps([{"id": 1}]),
compressed="[]",
)
time.sleep(1.1)
results = store.search(hash_key, "query")
assert results == []
def test_search_invalid_json(self):
"""Search handles invalid JSON gracefully."""
store = CompressionStore()
hash_key = store.store(
original="not valid json",
compressed="[]",
)
results = store.search(hash_key, "query")
assert results == []
def test_search_non_array(self):
"""Search handles non-array content gracefully."""
store = CompressionStore()
hash_key = store.store(
original=json.dumps({"key": "value"}),
compressed="{}",
)
results = store.search(hash_key, "query")
assert results == []
def test_empty_query_search(self):
"""Search with empty query returns empty or all."""
store = CompressionStore()
items = [{"id": i} for i in range(10)]
hash_key = store.store(
original=json.dumps(items),
compressed="[]",
)
# Empty query should return something (BM25 handles this)
results = store.search(hash_key, "")
# Behavior depends on BM25 implementation
assert isinstance(results, list)
def test_ccr_disabled_no_caching(self):
"""When CCR disabled, no caching occurs."""
reset_compression_store()

View file

@ -354,86 +354,6 @@ class TestRelevanceCalculation:
assert recommendations[0].relevance_score > 0.3
class TestExpansionTypeDetection:
"""Test determination of expansion type (full vs search)."""
def test_full_expansion_high_relevance(self):
"""High relevance triggers full expansion."""
tracker = ContextTracker()
context = CompressedContext(
hash_key="test",
turn_number=1,
timestamp=time.time(),
tool_name="Bash",
original_item_count=50,
compressed_item_count=5,
query_context="find files",
sample_content="auth.py, middleware.py",
workspace_key="ws-test",
)
expand_full, search_query = tracker._determine_expansion_type(
query="authentication middleware",
context=context,
relevance=0.8, # High relevance
)
assert expand_full is True
assert search_query is None
def test_full_expansion_small_count(self):
"""Small original item count triggers full expansion."""
tracker = ContextTracker()
context = CompressedContext(
hash_key="test",
turn_number=1,
timestamp=time.time(),
tool_name="Bash",
original_item_count=30, # Small
compressed_item_count=5,
query_context="find files",
sample_content="file.py",
workspace_key="ws-test",
)
expand_full, search_query = tracker._determine_expansion_type(
query="some query",
context=context,
relevance=0.4,
)
assert expand_full is True
def test_search_expansion_large_count(self):
"""Large original count with specific keywords triggers search."""
tracker = ContextTracker()
context = CompressedContext(
hash_key="test",
turn_number=1,
timestamp=time.time(),
tool_name="Bash",
original_item_count=500, # Large
compressed_item_count=20,
query_context="find all files",
sample_content="many files...",
workspace_key="ws-test",
)
expand_full, search_query = tracker._determine_expansion_type(
query="find authentication middleware handler",
context=context,
relevance=0.4, # Medium relevance
)
# Should use search for large datasets
if not expand_full:
assert search_query is not None
assert "authentication" in search_query or "middleware" in search_query
class TestExpansionExecution:
"""Test execution of expansion recommendations."""
@ -463,7 +383,6 @@ class TestExpansionExecution:
hash_key=hash_key,
reason="relevant to query",
relevance_score=0.8,
expand_full=True,
)
]
@ -473,39 +392,6 @@ class TestExpansionExecution:
assert results[0]["type"] == "full"
assert results[0]["item_count"] == 100
def test_execute_search_expansion(self):
"""Execute search expansion."""
store = get_compression_store()
items = [
{"id": 1, "content": "authentication code"},
{"id": 2, "content": "database operations"},
{"id": 3, "content": "authentication middleware"},
]
original = json.dumps(items)
hash_key = store.store(
original=original,
compressed="[]",
original_item_count=3,
)
tracker = ContextTracker()
recommendations = [
ExpansionRecommendation(
hash_key=hash_key,
reason="relevant to query",
relevance_score=0.5,
expand_full=False,
search_query="authentication",
)
]
results = tracker.execute_expansions(recommendations)
assert len(results) == 1
assert results[0]["type"] == "search"
assert results[0]["query"] == "authentication"
def test_execute_nonexistent_hash(self):
"""Handle expansion of nonexistent hash gracefully."""
tracker = ContextTracker()
@ -514,7 +400,6 @@ class TestExpansionExecution:
hash_key="nonexistent123",
reason="test",
relevance_score=0.5,
expand_full=True,
)
]
@ -549,27 +434,6 @@ class TestExpansionFormatting:
assert formatted.startswith("<headroom_proactive_expansion>\n")
assert formatted.endswith("\n</headroom_proactive_expansion>")
def test_format_search_expansion(self):
"""Format search expansion for LLM context."""
tracker = ContextTracker()
expansions = [
{
"hash": "def456",
"type": "search",
"query": "authentication",
"content": [{"id": 1, "content": "auth"}],
"item_count": 1,
"reason": "matched query",
}
]
formatted = tracker.format_expansions_for_context(expansions)
assert "Search results for 'authentication'" in formatted
assert formatted.startswith("<headroom_proactive_expansion>\n")
assert formatted.endswith("\n</headroom_proactive_expansion>")
def test_format_empty_expansions(self):
"""Empty expansions return empty string."""
tracker = ContextTracker()
@ -743,30 +607,19 @@ class TestCompressedContextDataClass:
class TestExpansionRecommendationDataClass:
"""Test ExpansionRecommendation dataclass."""
def test_full_expansion_recommendation(self):
"""Create full expansion recommendation."""
def test_expansion_recommendation(self):
"""Create an expansion recommendation (retrieval is always full)."""
rec = ExpansionRecommendation(
hash_key="abc123",
reason="high relevance",
relevance_score=0.9,
expand_full=True,
)
assert rec.expand_full is True
assert rec.search_query is None
def test_search_expansion_recommendation(self):
"""Create search expansion recommendation."""
rec = ExpansionRecommendation(
hash_key="def456",
reason="partial match",
relevance_score=0.5,
expand_full=False,
search_query="authentication",
)
assert rec.expand_full is False
assert rec.search_query == "authentication"
assert rec.hash_key == "abc123"
assert rec.relevance_score == 0.9
# Partial/search expansion fields no longer exist.
assert not hasattr(rec, "expand_full")
assert not hasattr(rec, "search_query")
class TestContextTrackerStats:

View file

@ -353,24 +353,3 @@ class TestFeedbackIntegrationWithStore:
patterns = feedback.get_all_patterns()
assert "test_tool" in patterns
assert patterns["test_tool"].total_retrievals == 1
def test_search_retrieval_tracked_separately(self):
"""Search retrievals are tracked as search type."""
store = CompressionStore()
hash_key = store.store(
original='[{"id": 1, "name": "alice"}, {"id": 2, "name": "bob"}]',
compressed='[{"id": 1}]',
original_item_count=2,
compressed_item_count=1,
tool_name="test_tool",
)
# Search (should log as search type)
store.search(hash_key, "alice")
# Check events
events = store.get_retrieval_events()
assert len(events) > 0
assert events[0].retrieval_type == "search"
assert events[0].query == "alice"

View file

@ -61,7 +61,7 @@ def test_mcp_retrieves_proxy_stored_content(fresh_store) -> None:
hash_key = get_compression_store().store(original, '{"compressed": true}')
server = mcp_server.HeadroomMCPServer(check_proxy=False)
result = asyncio.run(server._retrieve_content(hash_key, query=None))
result = asyncio.run(server._retrieve_content(hash_key))
assert result.get("source") == "local"
assert result["original_content"] == original
@ -91,26 +91,22 @@ def test_compress_savings_percent_tracks_token_counts(fresh_store) -> None:
assert result["savings_percent"] > 0.0
def test_mcp_retrieve_with_nonmatching_query_returns_full_content(fresh_store) -> None:
"""A query that matches no item above the relevance floor must still return
the stored entry (it exists and is unexpired) rather than the "Content not
found" error, which is reserved for genuine misses."""
def test_mcp_retrieve_returns_full_content(fresh_store) -> None:
"""Retrieval is by hash: a stored, unexpired entry always returns its full
original content (never empty, never a spurious "not found")."""
original = "the the the the the the the the the the\n" * 5
hash_key = get_compression_store().store(original, "<<small>>")
# Precondition: the query genuinely matches nothing above the BM25 floor.
assert get_compression_store().search(hash_key, "zzqx_nonmatching_token") == []
server = mcp_server.HeadroomMCPServer(check_proxy=False)
result = asyncio.run(server._retrieve_content(hash_key, query="zzqx_nonmatching_token"))
result = asyncio.run(server._retrieve_content(hash_key))
assert "error" not in result
assert result.get("source") == "local"
assert result["original_content"] == original
assert result["count"] == 0
def test_mcp_retrieve_missing_hash_still_errors(fresh_store) -> None:
"""A genuinely missing hash must still report "Content not found"."""
server = mcp_server.HeadroomMCPServer(check_proxy=False)
result = asyncio.run(server._retrieve_content("nonexistent_hash", query="anything"))
result = asyncio.run(server._retrieve_content("nonexistent_hash"))
assert "Content not found" in result.get("error", "")

View file

@ -165,11 +165,11 @@ class TestCCRToolCallParsing:
assert len(ccr_calls) == 1
assert ccr_calls[0].tool_call_id == "tool_123"
assert ccr_calls[0].hash_key == "abc123def456abc123def456"
assert ccr_calls[0].query is None
assert not hasattr(ccr_calls[0], "query")
assert len(other_calls) == 0
def test_parse_anthropic_search_retrieval(self):
"""Parse search retrieval call from Anthropic format."""
def test_parse_anthropic_retrieval_ignores_query(self):
"""Retrieval parses the hash; any legacy ``query`` input is ignored."""
handler = CCRResponseHandler()
response = {
@ -187,7 +187,7 @@ class TestCCRToolCallParsing:
assert len(ccr_calls) == 1
assert ccr_calls[0].hash_key == "def456abc123def456abc123"
assert ccr_calls[0].query == "authentication error"
assert not hasattr(ccr_calls[0], "query")
def test_parse_mixed_tool_calls(self):
"""Parse response with both CCR and other tool calls."""
@ -247,17 +247,15 @@ class TestCCRRetrievalExecution:
assert result.success
assert result.items_retrieved == 100
assert not result.was_search
# Check content structure
content = json.loads(result.content)
assert content["hash"] == hash_key
assert "original_content" in content
def test_search_retrieval_success(self):
"""Successfully search within cached content."""
def test_retrieval_returns_full_content_for_cached_hash(self):
"""Retrieval always returns the full original content (never empty)."""
store = get_compression_store()
# Use items with more searchable content
items = [
{"id": 1, "text": "Python programming language tutorial"},
{"id": 2, "text": "JavaScript web development framework"},
@ -276,18 +274,17 @@ class TestCCRRetrievalExecution:
)
handler = CCRResponseHandler()
# Use a more specific query
call = CCRToolCall(tool_call_id="test_id", hash_key=hash_key, query="Python programming")
call = CCRToolCall(tool_call_id="test_id", hash_key=hash_key)
result = handler._execute_retrieval(call)
assert result.success
assert result.was_search
assert result.items_retrieved == 5
content = json.loads(result.content)
assert content["query"] == "Python programming"
# The search should return results (may be 0 depending on BM25 behavior)
assert "results" in content
assert content["hash"] == hash_key
# Full content is always returned — the complete original round-trips.
assert json.loads(content["original_content"]) == items
def test_retrieval_nonexistent_hash(self):
"""Handle retrieval of nonexistent hash."""
@ -651,17 +648,7 @@ class TestCCRToolCallDataClass:
assert call.tool_call_id == "test_123"
assert call.hash_key == "abc123"
assert call.query is None
def test_search_retrieval_call(self):
"""Create search retrieval call."""
call = CCRToolCall(
tool_call_id="test_456",
hash_key="def456",
query="authentication",
)
assert call.query == "authentication"
assert not hasattr(call, "query")
class TestCCRToolResultDataClass:
@ -674,24 +661,11 @@ class TestCCRToolResultDataClass:
content='{"data": "content"}',
success=True,
items_retrieved=50,
was_search=False,
)
assert result.success
assert result.items_retrieved == 50
assert not result.was_search
def test_search_result(self):
"""Create search result."""
result = CCRToolResult(
tool_call_id="test_456",
content='{"results": []}',
success=True,
items_retrieved=5,
was_search=True,
)
assert result.was_search
assert not hasattr(result, "was_search")
def test_failed_result(self):
"""Create failed result."""

View file

@ -16,17 +16,9 @@ from headroom.ccr.tool_injection import CCR_TOOL_NAME
class FakeStore:
def __init__(
self, *, search_error: Exception | None = None, retrieve_error: Exception | None = None
) -> None:
self.search_error = search_error
def __init__(self, *, retrieve_error: Exception | None = None) -> None:
self.retrieve_error = retrieve_error
def search(self, hash_key: str, query: str) -> list[dict[str, str]]:
if self.search_error:
raise self.search_error
return [{"id": "1", "text": query}]
def retrieve(self, hash_key: str):
if self.retrieve_error:
raise self.retrieve_error
@ -88,7 +80,6 @@ def test_parse_ccr_tool_calls_google_and_other_calls() -> None:
CCRToolCall(
tool_call_id=CCR_TOOL_NAME,
hash_key="aaaaaaaaaaaaaaaaaaaaaaaa",
query="pizza",
)
]
assert other_calls == [{"functionCall": {"name": "other_tool", "args": {}}}]
@ -96,16 +87,7 @@ def test_parse_ccr_tool_calls_google_and_other_calls() -> None:
def test_execute_retrieval_error_paths(monkeypatch: pytest.MonkeyPatch) -> None:
handler = CCRResponseHandler()
monkeypatch.setattr(
"headroom.ccr.response_handler.get_compression_store",
lambda: FakeStore(search_error=RuntimeError("search boom")),
)
search_result = handler._execute_retrieval(
CCRToolCall(tool_call_id="t1", hash_key="abc", query="find")
)
assert search_result.success is False
assert "Retrieval failed: search boom" in search_result.content
# Retrieval is by hash only; a store error surfaces as a failed result.
monkeypatch.setattr(
"headroom.ccr.response_handler.get_compression_store",
lambda: FakeStore(retrieve_error=RuntimeError("retrieve boom")),

View file

@ -245,13 +245,8 @@ _ANTHROPIC_CCR_TOOL_SNAPSHOT_BYTES = (
b'"properties":{'
b'"hash":{"type":"string",'
b'"description":"Hash key from the compression marker '
b"(e.g., 'abc123' from hash=abc123)\"},"
b'"query":{"type":"string",'
b'"description":"Optional search query to filter results. '
b"If provided, only returns items matching the query. "
b'If omitted, returns all original items."}'
b"},"
b'"required":["hash"]}}'
b"(e.g., 'abc123' from hash=abc123)\"}"
b'},"required":["hash"]}}'
)
_OPENAI_CCR_TOOL_SNAPSHOT_BYTES = (
@ -265,13 +260,8 @@ _OPENAI_CCR_TOOL_SNAPSHOT_BYTES = (
b'"properties":{'
b'"hash":{"type":"string",'
b'"description":"Hash key from the compression marker '
b"(e.g., 'abc123' from hash=abc123)\"},"
b'"query":{"type":"string",'
b'"description":"Optional search query to filter results. '
b"If provided, only returns items matching the query. "
b'If omitted, returns all original items."}'
b"},"
b'"required":["hash"]}}}'
b"(e.g., 'abc123' from hash=abc123)\"}"
b'},"required":["hash"]}}}'
)

View file

@ -23,7 +23,8 @@ class TestCCRToolDefinition:
assert "input_schema" in tool
assert tool["input_schema"]["type"] == "object"
assert "hash" in tool["input_schema"]["properties"]
assert "query" in tool["input_schema"]["properties"]
# Retrieval is by hash only — no query/search parameter.
assert "query" not in tool["input_schema"]["properties"]
assert tool["input_schema"]["required"] == ["hash"]
def test_openai_format(self):
@ -266,13 +267,12 @@ class TestParseToolCall:
tool_call = {
"id": "toolu_123",
"name": CCR_TOOL_NAME,
"input": {"hash": "abc123def456abc123def456", "query": "errors"},
"input": {"hash": "abc123def456abc123def456"},
}
hash_key, query = parse_tool_call(tool_call, "anthropic")
hash_key = parse_tool_call(tool_call, "anthropic")
assert hash_key == "abc123def456abc123def456"
assert query == "errors"
def test_parse_openai_format(self):
"""Parse OpenAI tool call format."""
@ -280,14 +280,13 @@ class TestParseToolCall:
"id": "call_123",
"function": {
"name": CCR_TOOL_NAME,
"arguments": json.dumps({"hash": "def456abc123def456abc123", "query": None}),
"arguments": json.dumps({"hash": "def456abc123def456abc123"}),
},
}
hash_key, query = parse_tool_call(tool_call, "openai")
hash_key = parse_tool_call(tool_call, "openai")
assert hash_key == "def456abc123def456abc123"
assert query is None
def test_parse_non_ccr_tool(self):
"""Returns None for non-CCR tool calls."""
@ -296,10 +295,9 @@ class TestParseToolCall:
"input": {"param": "value"},
}
hash_key, query = parse_tool_call(tool_call, "anthropic")
hash_key = parse_tool_call(tool_call, "anthropic")
assert hash_key is None
assert query is None
def test_parse_malformed_openai_args(self):
"""Handles malformed JSON in OpenAI arguments."""
@ -311,7 +309,7 @@ class TestParseToolCall:
},
}
hash_key, query = parse_tool_call(tool_call, "openai")
hash_key = parse_tool_call(tool_call, "openai")
assert hash_key is None
@ -331,7 +329,7 @@ class TestHashSecurityValidation:
"input": {"hash": "abc123"}, # Only 6 chars
}
hash_key, query = parse_tool_call(tool_call, "anthropic")
hash_key = parse_tool_call(tool_call, "anthropic")
assert hash_key is None # Rejected
def test_rejects_long_hash(self):
@ -341,7 +339,7 @@ class TestHashSecurityValidation:
"input": {"hash": "abc123def456abc123def456abc123"}, # 30 chars
}
hash_key, query = parse_tool_call(tool_call, "anthropic")
hash_key = parse_tool_call(tool_call, "anthropic")
assert hash_key is None # Rejected
def test_rejects_non_hex_characters(self):
@ -351,7 +349,7 @@ class TestHashSecurityValidation:
"input": {"hash": "abc123xyz456abc123xyz456"}, # Contains xyz
}
hash_key, query = parse_tool_call(tool_call, "anthropic")
hash_key = parse_tool_call(tool_call, "anthropic")
assert hash_key is None # Rejected
def test_accepts_valid_24_char_hash(self):
@ -361,7 +359,7 @@ class TestHashSecurityValidation:
"input": {"hash": "abc123def456abc123def456"},
}
hash_key, query = parse_tool_call(tool_call, "anthropic")
hash_key = parse_tool_call(tool_call, "anthropic")
assert hash_key == "abc123def456abc123def456"
def test_accepts_uppercase_hex(self):
@ -371,7 +369,7 @@ class TestHashSecurityValidation:
"input": {"hash": "ABC123DEF456ABC123DEF456"},
}
hash_key, query = parse_tool_call(tool_call, "anthropic")
hash_key = parse_tool_call(tool_call, "anthropic")
# Note: validation accepts uppercase since we use .lower() for hex check
assert hash_key == "ABC123DEF456ABC123DEF456"
@ -433,13 +431,12 @@ class TestSmartCrusherCcrMarkers:
"""``parse_tool_call`` accepts a 12-char SmartCrusher hash."""
tool_call = {
"name": CCR_TOOL_NAME,
"input": {"hash": "e21a26620105", "query": "auth middleware"},
"input": {"hash": "e21a26620105"},
}
hash_key, query = parse_tool_call(tool_call, "anthropic")
hash_key = parse_tool_call(tool_call, "anthropic")
assert hash_key == "e21a26620105"
assert query == "auth middleware"
def test_parse_tool_call_still_accepts_24_char_hash(self):
"""24-char legacy hashes remain valid (regression guard)."""
@ -448,7 +445,7 @@ class TestSmartCrusherCcrMarkers:
"input": {"hash": "abc123def456abc123def456"},
}
hash_key, _ = parse_tool_call(tool_call, "anthropic")
hash_key = parse_tool_call(tool_call, "anthropic")
assert hash_key == "abc123def456abc123def456"

View file

@ -103,33 +103,6 @@ def test_retrieve_log_redacts_secret_payload_values():
assert "Authorization: [REDACTED]" in events[0]["payload_preview"]
def test_search_logs_retrieved_payload_preview():
store = CompressionStore(enable_feedback=False)
items = [
{"id": 1, "text": "alpha target"},
{"id": 2, "text": "beta other"},
]
hash_key = store.store(
original=json.dumps(items),
compressed="[]",
original_item_count=2,
compressed_item_count=0,
tool_name="search_tool",
)
with _capture_headroom_retrieve_events() as events:
results = store.search(hash_key, "alpha", score_threshold=0.0)
assert results
assert len(events) == 1
assert events[0]["hash"] == hash_key
assert events[0]["retrieval_type"] == "search"
assert events[0]["query"] == "alpha"
assert events[0]["payload_preview"] == json.dumps(results, ensure_ascii=False)
assert events[0]["payload_preview_chars"] == len(json.dumps(results, ensure_ascii=False))
assert events[0]["payload_truncated"] is False
def test_global_store_uses_env_default_ttl(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv(CCR_TTL_SECONDS_ENV, "7200")
@ -653,18 +626,6 @@ class TestCompressionStoreTTL:
assert store_with_short_ttl.get_metadata(hash_key) is None
def test_search_returns_empty_for_expired(self, store_with_short_ttl: CompressionStore):
"""search returns empty list for expired entries."""
hash_key = store_with_short_ttl.store(
original=json.dumps([{"id": 1, "name": "test"}]),
compressed="[]",
)
time.sleep(1.1)
results = store_with_short_ttl.search(hash_key, "test")
assert results == []
def test_exists_clean_expired_false_does_not_delete(
self, store_with_short_ttl: CompressionStore
):
@ -894,122 +855,6 @@ class TestCompressionStoreMetadata:
# =============================================================================
class TestCompressionStoreSearch:
"""Tests for CompressionStore search functionality."""
def test_search_with_bm25_returns_matches(self, store: CompressionStore):
"""search() uses BM25 to find matching items."""
items = [
{"id": 1, "content": "Python programming language"},
{"id": 2, "content": "JavaScript web development"},
{"id": 3, "content": "Python data science pandas"},
{"id": 4, "content": "Java enterprise applications"},
{"id": 5, "content": "Python machine learning tensorflow"},
]
hash_key = store.store(
original=json.dumps(items),
compressed=json.dumps(items[:2]),
)
results = store.search(hash_key, "Python programming")
assert len(results) >= 1
result_ids = [r["id"] for r in results]
assert 1 in result_ids # "Python programming language" should match
def test_search_respects_max_results(self, store: CompressionStore):
"""search() respects max_results parameter."""
items = [{"id": i, "content": f"item {i}"} for i in range(50)]
hash_key = store.store(original=json.dumps(items), compressed="[]")
results = store.search(hash_key, "item", max_results=5)
assert len(results) <= 5
def test_search_respects_score_threshold(self, store: CompressionStore):
"""search() filters by score threshold."""
items = [
{"id": 1, "content": "exact match query term"},
{"id": 2, "content": "completely unrelated content xyz"},
]
hash_key = store.store(original=json.dumps(items), compressed="[]")
# High threshold should filter low-scoring items
results = store.search(hash_key, "exact match query", score_threshold=0.5)
# Should return the exact match, filter the unrelated
if results:
assert any("exact match" in str(r) for r in results)
def test_search_nonexistent_returns_empty(self, store: CompressionStore):
"""search() returns empty list for nonexistent hash."""
results = store.search("nonexistent", "query")
assert results == []
def test_search_invalid_json_returns_empty(self, store: CompressionStore):
"""search() handles invalid JSON gracefully."""
hash_key = store.store(original="not valid json", compressed="[]")
results = store.search(hash_key, "query")
assert results == []
def test_search_plain_text_returns_matching_chunks(self, store: CompressionStore):
"""search() can find content in Kompress-style plain-text originals."""
original = (
"The OpenAI handler contains def _compress_openai_responses_payload "
"for Responses API live-zone compression. Other text is irrelevant."
)
hash_key = store.store(original=original, compressed="compressed")
results = store.search(hash_key, "def _compress_openai_responses_payload")
assert len(results) == 1
assert results[0]["type"] == "text"
assert "_compress_openai_responses_payload" in results[0]["text"]
def test_search_json_object_returns_matching_leaf(self, store: CompressionStore):
"""search() can find values inside JSON objects, not only arrays."""
original = json.dumps(
{
"module": {
"name": "openai",
"function": "_compress_openai_responses_payload",
}
}
)
hash_key = store.store(original=original, compressed="{}")
results = store.search(hash_key, "_compress_openai_responses_payload")
assert len(results) == 1
assert results[0]["path"] == "module.function"
assert results[0]["value"] == "_compress_openai_responses_payload"
def test_search_non_array_returns_empty(self, store: CompressionStore):
"""search() returns empty for JSON objects without matching leaves."""
hash_key = store.store(original=json.dumps({"key": "value"}), compressed="{}")
results = store.search(hash_key, "query")
assert results == []
def test_search_empty_array_returns_empty(self, store: CompressionStore):
"""search() returns empty for empty array."""
hash_key = store.store(original="[]", compressed="[]")
results = store.search(hash_key, "query")
assert results == []
def test_search_logs_retrieval_event(self, store: CompressionStore):
"""search() logs retrieval event with search type."""
items = [{"id": 1, "content": "test"}]
hash_key = store.store(original=json.dumps(items), compressed="[]")
store.search(hash_key, "test query")
events = store.get_retrieval_events()
search_events = [e for e in events if e.retrieval_type == "search"]
assert len(search_events) >= 1
assert search_events[-1].query == "test query"
# =============================================================================
# Retrieval Events Tests
# =============================================================================
@ -1383,7 +1228,7 @@ class TestRetrievalEvent:
total_items=100,
tool_name="search_api",
timestamp=time.time(),
retrieval_type="search",
retrieval_type="full",
tool_signature_hash="sig_123",
)
@ -1392,7 +1237,7 @@ class TestRetrievalEvent:
assert event.items_retrieved == 5
assert event.total_items == 100
assert event.tool_name == "search_api"
assert event.retrieval_type == "search"
assert event.retrieval_type == "full"
assert event.tool_signature_hash == "sig_123"
def test_retrieval_event_default_signature_hash(self):

View file

@ -295,12 +295,12 @@ class TestUserCountMergeLogic:
# =============================================================================
# CRITICAL #4: _get_entry_for_search returns reference not copy
# CRITICAL #4: retrieve() returns reference not copy
# =============================================================================
class TestGetEntryForSearchRaceCondition:
"""CRITICAL: _get_entry_for_search returns reference to internal entry.
class TestRetrieveRaceCondition:
"""CRITICAL: retrieve() must not return a reference to the internal entry.
The entry can be modified or evicted by another thread after the lock
is released but before the caller uses it, causing race conditions.
@ -321,8 +321,8 @@ class TestGetEntryForSearchRaceCondition:
tool_name="test_tool",
)
# Get entry via _get_entry_for_search
entry1 = store._get_entry_for_search(hash_key)
# Get entry via retrieve()
entry1 = store.retrieve(hash_key)
assert entry1 is not None
# Modify the returned entry
@ -330,7 +330,7 @@ class TestGetEntryForSearchRaceCondition:
entry1.retrieval_count = 999
# Get entry again - should NOT reflect our modifications
entry2 = store._get_entry_for_search(hash_key)
entry2 = store.retrieve(hash_key)
# AFTER FIX: entry2 should be a fresh copy, not affected by entry1 modifications
# Currently this may fail because we return a reference
@ -354,7 +354,7 @@ class TestGetEntryForSearchRaceCondition:
def reader():
for _ in range(50):
entry = store._get_entry_for_search(hash_key, "query")
entry = store.retrieve(hash_key)
if entry:
# Simulate work with the entry
try:
@ -368,7 +368,7 @@ class TestGetEntryForSearchRaceCondition:
def modifier():
for _ in range(50):
# Try to mess with internal state
entry = store._get_entry_for_search(hash_key)
entry = store.retrieve(hash_key)
if entry:
entry.search_queries.clear() # Shouldn't affect other readers
time.sleep(0.001)
@ -677,19 +677,15 @@ class TestCriticalFixesIntegration:
strategy="TOP_N",
)
# 4. Simulate retrieval
# 4. Simulate retrieval (by hash — returns the full original content)
entry = store.retrieve(hash_key)
assert entry is not None
assert entry.original_item_count == 100
# 5. Search within cached data
store.search(hash_key, "item_50")
# Should find the item even though it was compressed away
# 6. Get recommendation from TOIN
# 5. Get recommendation from TOIN
toin.get_recommendation(sig, "find item_50")
# 7. Verify stats are consistent
# 6. Verify stats are consistent
toin_stats = toin.get_stats()
store_stats = store.get_stats()
feedback_stats = feedback.get_stats()
@ -925,9 +921,10 @@ class TestCompressionStoreHighPriorityFixes:
compressed='[{"id": 1}]',
)
# Trigger many searches with different queries
# Trigger many retrievals with different queries (query is recorded
# for access tracking even though retrieval itself is by hash).
for i in range(50):
store.search(hash_key, f"unique_query_{i}")
store.retrieve(hash_key, query=f"unique_query_{i}")
with store._lock:
entry = store._backend.get(hash_key)

View file

@ -89,10 +89,11 @@ def test_handler_normalizes_marker_shapes_to_bare_hash(
assert seen["payload"]["hash"] == expected
def test_handler_passes_optional_query_through(
def test_handler_ignores_legacy_query(
plugin: types.ModuleType, monkeypatch: pytest.MonkeyPatch
) -> None:
# Arrange
# Retrieval is by hash only; a legacy ``query`` in the args is dropped,
# not forwarded to the proxy.
seen: dict[str, Any] = {}
def fake_post(url: str, json: dict[str, Any], timeout: int) -> _FakeResponse: # noqa: A002
@ -105,7 +106,7 @@ def test_handler_passes_optional_query_through(
plugin._handle_headroom_retrieve({"hash": "abc123", "query": "error lines"})
# Assert
assert seen["payload"] == {"hash": "abc123", "query": "error lines"}
assert seen["payload"] == {"hash": "abc123"}
def test_handler_returns_original_content_on_200(

View file

@ -109,62 +109,6 @@ class TestCCRRetrieveEndpoint:
assert len(retrieved_items) == 50
assert retrieved_items[0]["id"] == 0
def test_retrieve_with_search(self, client):
"""Search retrieval filters by query."""
store = get_compression_store()
items = [
{"id": 1, "text": "Python programming language"},
{"id": 2, "text": "JavaScript web development"},
{"id": 3, "text": "Python data science"},
{"id": 4, "text": "Java enterprise"},
]
hash_key = store.store(
original=json.dumps(items),
compressed="[]",
original_item_count=4,
compressed_item_count=0,
)
response = client.post(
"/v1/retrieve", json={"hash": hash_key, "query": "Python programming"}
)
assert response.status_code == 200
data = response.json()
assert data["hash"] == hash_key
assert data["query"] == "Python programming"
assert "results" in data
assert data["count"] >= 1
def test_retrieve_with_search_plain_text_original(self, client):
"""Query retrieval searches plain-text originals stored by Kompress."""
store = get_compression_store()
original = (
"Codex WS compression stores plain text originals. "
"The target symbol is _compress_openai_responses_payload."
)
hash_key = store.store(original=original, compressed="compressed")
response = client.post(
"/v1/retrieve",
json={"hash": hash_key, "query": "_compress_openai_responses_payload"},
)
assert response.status_code == 200
data = response.json()
assert data["hash"] == hash_key
assert data["count"] == 1
assert data["results"][0]["type"] == "text"
assert "_compress_openai_responses_payload" in data["results"][0]["text"]
def test_retrieve_with_search_nonexistent_hash_returns_404(self, client):
"""Query mode should not mask a missing hash as an empty search."""
response = client.post(
"/v1/retrieve",
json={"hash": "nonexistent123", "query": "anything"},
)
assert response.status_code == 404
def test_retrieve_increments_count(self, client):
"""Each retrieval increments the retrieval count."""
store = get_compression_store()
@ -206,47 +150,6 @@ class TestCCRRetrieveGetEndpoint:
assert data["original_item_count"] == 20
assert data["tool_name"] == "get_test_tool"
def test_get_retrieve_with_query(self, client):
"""GET retrieval with query parameter invokes search."""
store = get_compression_store()
# Create items with distinctive content
items = [
{"id": 1, "msg": "Python programming language tutorial for beginners"},
{"id": 2, "msg": "JavaScript web development framework guide"},
{"id": 3, "msg": "Python data science machine learning pandas"},
{"id": 4, "msg": "Java enterprise application development"},
]
hash_key = store.store(
original=json.dumps(items),
compressed="[]",
)
response = client.get(f"/v1/retrieve/{hash_key}?query=Python programming")
assert response.status_code == 200
data = response.json()
assert data["query"] == "Python programming"
# Response includes search results structure
assert "results" in data
assert "count" in data
# Results should be a list (may be empty if BM25 threshold not met)
assert isinstance(data["results"], list)
def test_get_retrieve_with_query_plain_text_original(self, client):
"""GET query retrieval searches plain-text originals."""
store = get_compression_store()
hash_key = store.store(
original="plain text contains _compress_openai_responses_payload",
compressed="plain text",
)
response = client.get(f"/v1/retrieve/{hash_key}?query=_compress_openai_responses_payload")
assert response.status_code == 200
data = response.json()
assert data["count"] == 1
assert data["results"][0]["type"] == "text"
def test_get_retrieve_nonexistent(self, client):
"""GET with nonexistent hash returns 404."""
response = client.get("/v1/retrieve/nonexistent123")
@ -298,7 +201,6 @@ class TestCCRStatsEndpoint:
store = get_compression_store()
# Use non-empty content so search actually logs
content = json_module.dumps(
[
{"id": "1", "name": "test item", "value": 100},
@ -311,9 +213,9 @@ class TestCCRStatsEndpoint:
tool_name="stats_test_tool",
)
# Make some retrievals
client.post("/v1/retrieve", json={"hash": hash_key}) # Full retrieval
client.post("/v1/retrieve", json={"hash": hash_key, "query": "test"}) # Search retrieval
# Make some retrievals (retrieval is by hash → always full)
client.post("/v1/retrieve", json={"hash": hash_key})
client.post("/v1/retrieve", json={"hash": hash_key})
response = client.get("/v1/retrieve/stats")
assert response.status_code == 200
@ -322,10 +224,10 @@ class TestCCRStatsEndpoint:
assert data["store"]["total_retrievals"] >= 2
assert len(data["recent_retrievals"]) >= 2
# Verify we have both retrieval types (no double-logging of full)
# All retrievals are full (no double-logging)
retrieval_types = [r["retrieval_type"] for r in data["recent_retrievals"]]
assert "full" in retrieval_types
assert "search" in retrieval_types
assert all(rt == "full" for rt in retrieval_types)
class TestCCRIntegration:
@ -375,19 +277,6 @@ class TestCCREdgeCases:
data = response.json()
assert data["original_item_count"] == 1000
def test_search_no_matches(self, client):
"""Search with no matches returns empty results."""
store = get_compression_store()
items = [{"id": 1, "text": "hello world"}]
hash_key = store.store(original=json.dumps(items), compressed="[]")
response = client.post("/v1/retrieve", json={"hash": hash_key, "query": "xyznonexistent"})
assert response.status_code == 200
data = response.json()
assert data["count"] == 0
assert data["results"] == []
def test_unicode_content(self, client):
"""Unicode content is handled correctly."""
store = get_compression_store()
@ -707,10 +596,10 @@ class TestEndToEndTOINIntegration:
strategy="smart_sample",
)
# Step 2: Retrieve through proxy endpoint
# Step 2: Retrieve through proxy endpoint (by hash → full content)
response = client_with_optimization.post(
"/v1/retrieve",
json={"hash": hash_key, "query": "category:cat_1"},
json={"hash": hash_key},
)
assert response.status_code == 200

View file

@ -73,13 +73,12 @@ def test_kompress_ccr_retrieval_updates_toin():
entry = store.retrieve(hash_key)
assert entry is not None
assert entry.tool_signature_hash is not None
results = store.search(hash_key, "HEADROOM", score_threshold=0.0)
assert results
# Retrieval is by hash and returns the full original content.
assert "HEADROOM" in entry.original_content
stats = get_toin().get_stats()
assert stats["total_compressions"] == 1
assert stats["total_retrievals"] == 2
assert stats["total_retrievals"] == 1
@pytest.mark.skip(reason="PR-B5: observations counter and request-time hint API retired")
@ -218,7 +217,7 @@ class TestCCRFeedbackExtraction:
"type": "tool_use",
"id": "toolu_123",
"name": "headroom_retrieve",
"input": {"hash": "abc123def456", "query": "error fields"},
"input": {"hash": "abc123def456"},
},
]
}
@ -236,7 +235,6 @@ class TestCCRFeedbackExtraction:
assert len(retrieve_calls) == 1
assert retrieve_calls[0]["hash"] == "abc123def456"
assert retrieve_calls[0]["query"] == "error fields"
def test_ignore_non_retrieve_tool_calls(self):
"""Should ignore tool_use blocks that are not headroom_retrieve."""
@ -252,7 +250,7 @@ class TestCCRFeedbackExtraction:
"type": "tool_use",
"id": "toolu_789",
"name": "headroom_retrieve",
"input": {"hash": "xyz789", "query": None},
"input": {"hash": "xyz789"},
},
]
}
@ -313,8 +311,12 @@ class TestCCRFeedbackExtraction:
class TestStreamingFeedbackIntegration:
"""Bug 2: Full feedback loop — streaming headroom_retrieve reaches TOIN."""
def test_record_ccr_feedback_calls_store_search(self):
"""_record_ccr_feedback_from_response should call store.search for queries."""
def test_record_ccr_feedback_calls_store_retrieve(self):
"""_record_ccr_feedback_from_response calls store.retrieve by hash.
Retrieval is by hash only any legacy ``query`` in the tool input is
ignored, and the full content is fetched for the feedback side effect.
"""
from headroom.proxy.server import HeadroomProxy
response = {
@ -340,10 +342,11 @@ class TestStreamingFeedbackIntegration:
proxy._record_ccr_feedback_from_response(response, "anthropic", "req-test-001")
mock_store.search.assert_called_once_with("feedbackhash1", "error details")
mock_store.retrieve.assert_called_once_with("feedbackhash1")
mock_store.search.assert_not_called()
def test_record_ccr_feedback_calls_store_retrieve_no_query(self):
"""_record_ccr_feedback_from_response should call store.retrieve when no query."""
"""_record_ccr_feedback_from_response calls store.retrieve by hash."""
from headroom.proxy.server import HeadroomProxy
response = {
@ -368,7 +371,7 @@ class TestStreamingFeedbackIntegration:
proxy._record_ccr_feedback_from_response(response, "anthropic", "req-test-002")
mock_store.retrieve.assert_called_once_with("feedbackhash2", query=None)
mock_store.retrieve.assert_called_once_with("feedbackhash2")
def test_record_ccr_feedback_handles_store_exception(self):
"""_record_ccr_feedback_from_response should not raise on store errors."""
@ -380,13 +383,13 @@ class TestStreamingFeedbackIntegration:
"type": "tool_use",
"id": "toolu_003",
"name": "headroom_retrieve",
"input": {"hash": "feedbackhash3", "query": "test"},
"input": {"hash": "feedbackhash3"},
},
]
}
mock_store = MagicMock()
mock_store.search.side_effect = RuntimeError("store unavailable")
mock_store.retrieve.side_effect = RuntimeError("store unavailable")
with patch(
"headroom.cache.compression_store.get_compression_store",
return_value=mock_store,
@ -417,7 +420,7 @@ class TestParseSSEToolUse:
"\n"
'data: {"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"toolu_abc","name":"headroom_retrieve"}}\n'
"\n"
'data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\\"hash\\": \\"abc123\\", \\"query\\": \\"error\\"}"}}\n'
'data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\\"hash\\": \\"abc123\\"}"}}\n'
"\n"
'data: {"type":"content_block_stop","index":1}\n'
"\n"
@ -439,7 +442,6 @@ class TestParseSSEToolUse:
assert tool_block["name"] == "headroom_retrieve"
assert tool_block["id"] == "toolu_abc"
assert tool_block["input"]["hash"] == "abc123"
assert tool_block["input"]["query"] == "error"
def test_parse_sse_non_anthropic_returns_none(self):
"""Non-anthropic provider should return None."""

View file

@ -578,7 +578,7 @@ class CompressionEntry:
- Thread-safe in-memory storage
- TTL-based expiration (default 5 minutes)
- LRU-style eviction when capacity reached
- Built-in BM25 search within cached content
- Hash-keyed retrieval that always returns the full original content
**Usage:**
```python
@ -593,11 +593,8 @@ hash_key = store.store(
tool_name="search_api",
)
# Retrieve later
# Retrieve later (by hash; always returns the full original content)
entry = store.retrieve(hash_key)
# Or search within cached content
results = store.search(hash_key, "user query")
```
---
@ -609,17 +606,18 @@ results = store.search(hash_key, "user query")
| Endpoint | Method | Description |
|----------|--------|-------------|
| `/v1/retrieve` | POST | Retrieve original content by hash |
| `/v1/retrieve?query=X` | POST | Search within cached content |
| `/v1/retrieve/{hash}` | GET | Retrieve original content by hash |
Retrieval is by hash only — it always returns the full original content.
**Retrieval Request:**
```json
{
"hash": "abc123def456...",
"query": "find errors" // Optional: search within
"hash": "abc123def456..."
}
```
**Response (full retrieval):**
**Response:**
```json
{
"hash": "abc123def456...",
@ -629,16 +627,6 @@ results = store.search(hash_key, "user query")
}
```
**Response (search):**
```json
{
"hash": "abc123def456...",
"query": "find errors",
"results": [{...}, {...}, ...],
"count": 15
}
```
---
### CCR Phase 3: Tool Injection
@ -664,9 +652,9 @@ When running as MCP server, Headroom exposes retrieval as a tool:
"inputSchema": {
"type": "object",
"properties": {
"hash": {"type": "string"},
"query": {"type": "string"}
}
"hash": {"type": "string"}
},
"required": ["hash"]
}
}
```
@ -701,10 +689,10 @@ class ToolPattern:
tool_name: str
total_compressions: int # Times we compressed this tool
total_retrievals: int # Times LLM asked for more
full_retrievals: int # Retrieved everything
search_retrievals: int # Used search query
common_queries: dict[str, int] # Query frequency
queried_fields: dict[str, int] # Fields mentioned in queries
full_retrievals: int # Retrieved everything (all retrievals — hash-only)
search_retrievals: int # Legacy; always 0 (retrieval is hash-only, no search)
common_queries: dict[str, int] # Legacy query-pattern frequency (no longer populated)
queried_fields: dict[str, int] # Legacy queried-field frequency (no longer populated)
```
**Key Metrics:**
@ -786,11 +774,10 @@ if self.config.use_feedback_hints and tool_name:
│ └─ Contains: tool_use(headroom_retrieve, hash=abc123) │
│ │
│ 2. Handler detects CCR tool call │
│ └─ Extracts hash and optional query
│ └─ Extracts hash
│ │
│ 3. Handler executes retrieval │
│ └─ Full retrieval: store.retrieve(hash) │
│ └─ Search: store.search(hash, query) │
│ └─ By hash: store.retrieve(hash) → full original content │
│ │
│ 4. Handler continues conversation │
│ └─ Adds tool result to messages │
@ -810,7 +797,6 @@ if self.config.use_feedback_hints and tool_name:
class CCRToolCall:
tool_call_id: str # For matching response
hash_key: str # CCR hash to retrieve
query: str | None # Optional search query
@dataclass
class CCRToolResult:
@ -818,7 +804,6 @@ class CCRToolResult:
content: str # Retrieved data as JSON
success: bool
items_retrieved: int
was_search: bool # True if search, False if full retrieval
class CCRResponseHandler:
async def handle_response(
@ -897,8 +882,7 @@ class ExpansionRecommendation:
hash_key: str
reason: str # Human-readable explanation
relevance_score: float # 0-1, higher = more relevant
expand_full: bool # True = full retrieval
search_query: str | None # If expand_full=False
# Expansion always restores the full original content (retrieval is by hash).
class ContextTracker:
def track_compression(self, hash_key, turn_number, ...):

View file

@ -55,8 +55,7 @@ Headroom injects a `headroom_retrieve` tool into the LLM's available tools:
"name": "headroom_retrieve",
"description": "Retrieve original uncompressed data from Headroom cache",
"parameters": {
"hash": "The hash key from the compression marker",
"query": "Optional: search within the cached data"
"hash": "The hash key from the compression marker"
}
}
```
@ -128,7 +127,7 @@ Full content available via ccr_retrieve tool with reference 'def456'.]
|---------|-------------|
| **Automatic Response Handling** | When LLM calls `headroom_retrieve`, the proxy handles it automatically |
| **Multi-Turn Context Tracking** | Tracks compressed content across turns, proactively expands when relevant |
| **BM25 Search** | LLM can search within compressed data: `headroom_retrieve(hash, query="errors")` |
| **Hash-Keyed Retrieval** | `headroom_retrieve(hash)` always returns the full original content |
| **Feedback Learning** | Learns from retrieval patterns to improve future compression |
## Configuration