mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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:
parent
a639540959
commit
c2fc4d3753
33 changed files with 315 additions and 1381 deletions
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
265
headroom/cache/compression_store.py
vendored
265
headroom/cache/compression_store.py
vendored
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -65,7 +65,6 @@ const retrieve = createHeadroomRetrieveTool({
|
|||
|
||||
const result = await retrieve.execute({
|
||||
hash: "0123456789abcdef01234567",
|
||||
query: "needle",
|
||||
});
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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", "")
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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")),
|
||||
|
|
|
|||
|
|
@ -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"]}}}'
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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, ...):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue