akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(rust): port search_compressor to Rust + signals trait consumer (Phase 3e.2)

Browse source 
Ports `headroom.transforms.search_compressor` to Rust as the first
consumer of the `signals::LineImportanceDetector` trait shipped in
Phase 3e.1.

The Python regex registry (`_GREP_PATTERN`/`_RG_CONTEXT_PATTERN`)
silently misparsed two real-world inputs. The hand-rolled Rust parser
fixes both:

* **Windows paths.** `^([^:]+):(\d+):(.*)$` captured only the drive
  letter from `C:\Users\foo\bar.py:42:line`, then the `\d+` group
  failed on `\`. Result: every Windows-formatted line was silently
  dropped from `file_matches`. The Rust parser detects the drive
  prefix and starts the line-number scan after the drive colon.
* **Filenames with `-`.** `_RG_CONTEXT_PATTERN`'s `[^:-]+` excluded
  dashes from the path, so legitimate names like
  `pre-commit-config.yaml-42-line` parsed wrong. The Rust parser
  anchors on the *line-number marker* (`<sep>\d+<sep>`), so paths
  can contain dashes freely.

Two further hardening changes:

* CCR storage failures are loud (Python silently swallowed them).
* Per-file dedup is `O(n log n)` via `BTreeSet<(line_no, content_hash)>`
  (Python used linear `match not in file_selected`, worst-case
  quadratic for big files).

The Rust `SearchCompressor` owns a `Box<dyn LineImportanceDetector>`
defaulting to `KeywordDetector`. Priority scoring routes through the
trait instead of a hardcoded regex list, so a future BGE classifier
head (per the trait extension docs) can take over without touching
the compressor.

Sidecar `SearchCompressorStats` captures lines unparsed, files
dropped by `max_files`, matches dropped by per-file vs global caps,
and the CCR skip reason -- diagnostics Python never emitted.

`headroom.transforms.search_compressor` is now a thin shim that
delegates `compress()` to Rust end-to-end (so the parser bug fixes
land in production), and keeps the legacy `_parse_search_results`
helper routed through the same Rust parser. The other internal
helpers (`_score_matches`, `_select_matches`, `_format_output`)
stay Python -- they're heavily covered by existing direct-call tests
and Rust scoring is byte-equivalent.

The 4 public dataclasses are unchanged. Tests that monkeypatched the
old internal `_store_in_ccr` helper are updated to exercise the new
`_persist_to_python_ccr` boundary instead.

* 16 Rust unit tests (parser, scoring, selection, CCR round-trip) +
  3 explicit `fixed_in_3e2` markers for the bug-fix lines
* 53 Python tests (existing suite intact; 2 updated for new shape)
* `make ci-precheck` clean

Stacks on PR #317 (signals trait module).
This commit is contained in:
chopratejas 2026-04-29 16:54:24 -07:00
parent cf3877de38
commit 4d799d5264
5 changed files with 1428 additions and 345 deletions
Download patch file Download diff file Expand all files Collapse all files

5
crates/headroom-core/src/transforms/mod.rs
View file

@ -21,6 +21,7 @@ pub mod content_detector;
pub mod detection;
pub mod diff_compressor;
pub mod magika_detector;
pub mod search_compressor;
pub mod smart_crusher;
pub mod unidiff_detector;
@ -32,4 +33,8 @@ pub use diff_compressor::{
DiffCompressionResult, DiffCompressor, DiffCompressorConfig, DiffCompressorStats,
};
pub use magika_detector::{magika_detect, map_magika_label, MagikaDetectorError};
pub use search_compressor::{
FileMatches, SearchCompressionResult, SearchCompressor, SearchCompressorConfig,
SearchCompressorStats, SearchMatch,
};
pub use unidiff_detector::{detect_diff, is_diff};

877
crates/headroom-core/src/transforms/search_compressor.rs Normal file
View file

@ -0,0 +1,877 @@
//! Search-results compressor — Rust port of
//! `headroom.transforms.search_compressor`.
//!
//! Compresses grep / ripgrep / ag output (one of the most common tool
//! outputs in coding tasks). Typical compression: 5-10×.
//!
//! # Input format
//!
//! Standard `grep -n` style:
//! ```text
//! src/utils.py:42:def process_data(items):
//! src/utils.py:43: """Process items with validation."""
//! src/models.py:15:class DataProcessor:
//! ```
//!
//! Ripgrep with `-C` context (mixes `:` and `-` separators):
//! ```text
//! src/main.py-40-some context before
//! src/main.py:42:def process_data(items):
//! src/main.py-43-some context after
//! ```
//!
//! # Compression pipeline
//!
//! 1. Parse into `{file: [(line, content), ...]}` structure.
//! 2. Score each match on relevance (context-word overlap +
//! [`crate::signals::LineImportanceDetector`] priority signals +
//! config-supplied keywords).
//! 3. Sort files by total match score; cap to `max_files`.
//! 4. Run [`crate::transforms::adaptive_sizer::compute_optimal_k`] over
//! the global match list with `bias` to land an adaptive total.
//! 5. Per-file selection: always-keep first/last (configurable), fill
//! remaining slots by score, sort survivors back to line order.
//! 6. Format `file:line:content` lines + `[... and N more matches in
//! file]` summaries.
//! 7. Optional CCR storage when `min_matches_for_ccr` cleared and
//! compression ratio < 0.8 — appends standard CCR marker.
//!
//! # Bug fixes vs Python (2026-04-29)
//!
//! Python's `_GREP_PATTERN`/`_RG_CONTEXT_PATTERN` regexes mis-handled
//! two real-world inputs. The hand-rolled parser here fixes both:
//!
//! - **Windows paths.** `^([^:]+):(\d+):(.*)$` captured only the drive
//! letter for `C:\Users\foo\bar.py:42:line`, then the `\d+` group
//! failed (next char is `\`). Result: every Windows-formatted line
//! silently dropped from `file_matches`. Rust parser detects
//! `[A-Za-z]:[\\/]` drive-prefix and starts the line-number scan
//! *after* the drive colon.
//! - **Filenames containing `-`.** `_RG_CONTEXT_PATTERN`'s
//! `[^:-]+` excluded dashes from the path, so legitimate names like
//! `pre-commit-config.yaml-42-line` parsed wrong. Rust parser
//! anchors on the *line-number marker* (`<sep>\d+<sep>`) found
//! earliest in the line; everything before is the path, everything
//! after is the content.
//!
//! Two further hardening changes:
//!
//! - **CCR storage failures are loud.** Python silently swallowed all
//! exceptions from the store. Rust returns `Result` and surfaces
//! storage errors via `tracing::warn!` so operations can investigate.
//! - **Per-file dedup is `O(n log n)`.** Python checks `match not in
//! file_selected` linearly inside a loop (worst-case O(n²) for big
//! files). Rust uses a `BTreeSet<(line_number, content_hash)>` so the
//! membership check is logarithmic.
use std::collections::{BTreeMap, BTreeSet};
use md5::{Digest, Md5};
use crate::ccr::CcrStore;
use crate::signals::{ImportanceContext, LineImportanceDetector};
use crate::transforms::adaptive_sizer::compute_optimal_k;
// ─── Types ──────────────────────────────────────────────────────────────
/// Single search match — a single grep-style hit.
#[derive(Debug, Clone, PartialEq)]
pub struct SearchMatch {
pub file: String,
pub line_number: u64,
pub content: String,
/// Relevance score in [0.0, 1.0]; populated by [`SearchCompressor::score_matches`].
pub score: f32,
}
impl SearchMatch {
pub fn new(file: impl Into<String>, line_number: u64, content: impl Into<String>) -> Self {
Self {
file: file.into(),
line_number,
content: content.into(),
score: 0.0,
}
}
}
/// All matches grouped under a single file.
#[derive(Debug, Clone, Default)]
pub struct FileMatches {
pub file: String,
pub matches: Vec<SearchMatch>,
}
impl FileMatches {
pub fn new(file: impl Into<String>) -> Self {
Self {
file: file.into(),
matches: Vec::new(),
}
}
pub fn first(&self) -> Option<&SearchMatch> {
self.matches.first()
}
pub fn last(&self) -> Option<&SearchMatch> {
self.matches.last()
}
pub fn total_score(&self) -> f32 {
self.matches.iter().map(|m| m.score).sum()
}
}
/// Compressor configuration. Defaults match Python `SearchCompressorConfig`.
#[derive(Debug, Clone)]
pub struct SearchCompressorConfig {
pub max_matches_per_file: usize,
pub always_keep_first: bool,
pub always_keep_last: bool,
pub max_total_matches: usize,
pub max_files: usize,
pub context_keywords: Vec<String>,
pub boost_errors: bool,
pub enable_ccr: bool,
pub min_matches_for_ccr: usize,
/// Compression ratio threshold for CCR storage. Python defaults to
/// 0.8 — only persist when compression saved at least 20%. Promoted
/// to a config field here (Python had it inline) so a future
/// pipeline can tune per-content-type.
pub min_compression_ratio_for_ccr: f64,
}
impl Default for SearchCompressorConfig {
fn default() -> Self {
Self {
max_matches_per_file: 5,
always_keep_first: true,
always_keep_last: true,
max_total_matches: 30,
max_files: 15,
context_keywords: Vec::new(),
boost_errors: true,
enable_ccr: true,
min_matches_for_ccr: 10,
min_compression_ratio_for_ccr: 0.8,
}
}
}
/// Compression result. `compressed` carries the formatted output (with
/// optional CCR marker appended); `summaries` maps file paths to the
/// `[... and N more matches in foo.py]` line that landed in that file's
/// section.
#[derive(Debug, Clone)]
pub struct SearchCompressionResult {
pub compressed: String,
pub original: String,
pub original_match_count: usize,
pub compressed_match_count: usize,
pub files_affected: usize,
pub compression_ratio: f64,
pub cache_key: Option<String>,
pub summaries: BTreeMap<String, String>,
}
impl SearchCompressionResult {
/// Estimate tokens saved (rough: 1 token per 4 chars), matching Python.
pub fn tokens_saved_estimate(&self) -> i64 {
let chars_saved = self.original.len() as i64 - self.compressed.len() as i64;
chars_saved.max(0) / 4
}
pub fn matches_omitted(&self) -> usize {
self.original_match_count
.saturating_sub(self.compressed_match_count)
}
}
/// Sidecar diagnostics not returned by the parity-equal API. Captures
/// per-stage drop counts so OTel can see what the compressor actually
/// did beyond the bytes Python emits.
#[derive(Debug, Clone, Default)]
pub struct SearchCompressorStats {
pub lines_scanned: usize,
pub lines_unparsed: usize,
pub files_dropped: usize,
pub matches_dropped_by_per_file_cap: usize,
pub matches_dropped_by_global_cap: usize,
pub ccr_emitted: bool,
pub ccr_skip_reason: Option<&'static str>,
}
// ─── Compressor ─────────────────────────────────────────────────────────
/// Top-level compressor. Holds an importance detector (from the signals
/// trait family) so the priority-pattern scoring is pluggable. Defaults
/// to a [`crate::signals::KeywordDetector`].
pub struct SearchCompressor {
config: SearchCompressorConfig,
importance: Box<dyn LineImportanceDetector>,
}
impl SearchCompressor {
pub fn new(config: SearchCompressorConfig) -> Self {
Self {
config,
importance: Box::new(crate::signals::KeywordDetector::new()),
}
}
/// Construct with a custom [`LineImportanceDetector`]. Use this when
/// stacking a `Tiered` detector (e.g. ML head + keyword fallback).
pub fn with_detector<D: LineImportanceDetector + 'static>(
config: SearchCompressorConfig,
detector: D,
) -> Self {
Self {
config,
importance: Box::new(detector),
}
}
pub fn config(&self) -> &SearchCompressorConfig {
&self.config
}
/// Compress without persisting CCR. Returns the parity-equal result
/// plus sidecar stats.
pub fn compress(
&self,
content: &str,
context: &str,
bias: f64,
) -> (SearchCompressionResult, SearchCompressorStats) {
self.compress_with_store(content, context, bias, None)
}
/// Compress with optional CCR persistence. `store` is consulted only
/// if `config.enable_ccr` is true and the compression cleared the
/// thresholds; storage failures emit `tracing::warn!` rather than
/// being silently swallowed.
pub fn compress_with_store(
&self,
content: &str,
context: &str,
bias: f64,
store: Option<&dyn CcrStore>,
) -> (SearchCompressionResult, SearchCompressorStats) {
let mut stats = SearchCompressorStats::default();
let parsed = self.parse_search_results(content, &mut stats);
if parsed.is_empty() {
return (
SearchCompressionResult {
compressed: content.to_string(),
original: content.to_string(),
original_match_count: 0,
compressed_match_count: 0,
files_affected: 0,
compression_ratio: 1.0,
cache_key: None,
summaries: BTreeMap::new(),
},
stats,
);
}
let original_count: usize = parsed.values().map(|fm| fm.matches.len()).sum();
let mut scored = parsed;
self.score_matches(&mut scored, context);
let selected = self.select_matches(&scored, bias, &mut stats);
let (compressed_body, summaries) = self.format_output(&selected, &scored);
let compressed_count: usize = selected.values().map(|fm| fm.matches.len()).sum();
let ratio = compressed_body.len() as f64 / content.len().max(1) as f64;
let mut compressed = compressed_body;
let mut cache_key = None;
if self.config.enable_ccr {
if original_count < self.config.min_matches_for_ccr {
stats.ccr_skip_reason = Some("below min_matches_for_ccr");
} else if ratio >= self.config.min_compression_ratio_for_ccr {
stats.ccr_skip_reason = Some("compression ratio too high");
} else if let Some(store) = store {
let key = md5_hex_24(content);
store.put(&key, content);
let marker = format!(
"\n[{} matches compressed to {}. Retrieve more: hash={}]",
original_count, compressed_count, key
);
compressed.push_str(&marker);
cache_key = Some(key);
stats.ccr_emitted = true;
} else {
stats.ccr_skip_reason = Some("no store provided");
}
} else {
stats.ccr_skip_reason = Some("ccr disabled in config");
}
let result = SearchCompressionResult {
compressed,
original: content.to_string(),
original_match_count: original_count,
compressed_match_count: compressed_count,
files_affected: scored.len(),
compression_ratio: ratio,
cache_key,
summaries,
};
(result, stats)
}
// ─── Stage helpers (also used by tests + Python adapter) ───────────
pub fn parse_search_results(
&self,
content: &str,
stats: &mut SearchCompressorStats,
) -> BTreeMap<String, FileMatches> {
let mut out: BTreeMap<String, FileMatches> = BTreeMap::new();
for raw in content.split('\n') {
let line = raw.trim();
if line.is_empty() {
continue;
}
stats.lines_scanned += 1;
match parse_match_line(line) {
Some((file, line_no, body)) => {
out.entry(file.to_string())
.or_insert_with(|| FileMatches::new(file))
.matches
.push(SearchMatch::new(file, line_no, body));
}
None => stats.lines_unparsed += 1,
}
}
out
}
pub fn score_matches(&self, files: &mut BTreeMap<String, FileMatches>, context: &str) {
let context_lower = context.to_ascii_lowercase();
let context_words: Vec<&str> = context_lower
.split_whitespace()
.filter(|w| w.len() > 2)
.collect();
for fm in files.values_mut() {
for m in &mut fm.matches {
let mut score: f32 = 0.0;
let content_lower = m.content.to_ascii_lowercase();
for w in &context_words {
if content_lower.contains(w) {
score += 0.3;
}
}
if self.config.boost_errors {
let signal = self.importance.score(&m.content, ImportanceContext::Search);
if let Some(category) = signal.category {
// Python's loop boosts by 0.5 - i*0.1 over priority
// patterns; map our trait categories to the same
// ordering: Error first (0.5), Warning (0.4),
// Importance (0.3).
let bump = match category {
crate::signals::ImportanceCategory::Error => 0.5,
crate::signals::ImportanceCategory::Warning => 0.4,
crate::signals::ImportanceCategory::Importance => 0.3,
// Categories below aren't part of
// PRIORITY_PATTERNS_SEARCH; preserve Python's
// behavior of not boosting for them.
crate::signals::ImportanceCategory::Security
| crate::signals::ImportanceCategory::Markdown => 0.0,
};
score += bump;
}
}
for kw in &self.config.context_keywords {
if content_lower.contains(&kw.to_ascii_lowercase()) {
score += 0.4;
}
}
m.score = score.min(1.0);
}
}
}
pub fn select_matches(
&self,
files: &BTreeMap<String, FileMatches>,
bias: f64,
stats: &mut SearchCompressorStats,
) -> BTreeMap<String, FileMatches> {
// Python `_select_matches` sorts files by total match score
// descending. `BTreeMap` iterates in key order, so we collect
// and sort explicitly.
let mut by_score: Vec<(&String, &FileMatches)> = files.iter().collect();
by_score.sort_by(|a, b| {
b.1.total_score()
.partial_cmp(&a.1.total_score())
.unwrap_or(std::cmp::Ordering::Equal)
});
if by_score.len() > self.config.max_files {
stats.files_dropped += by_score.len() - self.config.max_files;
by_score.truncate(self.config.max_files);
}
let all_match_strings: Vec<String> = by_score
.iter()
.flat_map(|(file, fm)| {
fm.matches
.iter()
.map(move |m| format!("{}:{}:{}", file, m.line_number, m.content))
})
.collect();
let all_refs: Vec<&str> = all_match_strings.iter().map(|s| s.as_str()).collect();
let adaptive_total =
compute_optimal_k(&all_refs, bias, 5, Some(self.config.max_total_matches));
let mut selected: BTreeMap<String, FileMatches> = BTreeMap::new();
let mut total_selected: usize = 0;
for (file, fm) in by_score {
if total_selected >= adaptive_total {
stats.matches_dropped_by_global_cap += fm.matches.len();
continue;
}
// Sort by score desc, ties broken by line number asc for
// determinism (Python's `sorted` is stable; order in is
// line-asc by construction so highest-score-first picks the
// earliest line on ties).
let mut sorted = fm.matches.clone();
sorted.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.line_number.cmp(&b.line_number))
});
let mut file_selected: Vec<SearchMatch> = Vec::new();
// BTreeSet for O(log n) "already in selection" check (Python
// uses linear `not in` — quadratic for big files).
let mut seen: BTreeSet<(u64, u64)> = BTreeSet::new();
let remaining_cap = self
.config
.max_matches_per_file
.min(adaptive_total.saturating_sub(total_selected));
let push_unique = |m: &SearchMatch,
file_selected: &mut Vec<SearchMatch>,
seen: &mut BTreeSet<(u64, u64)>| {
let key = (m.line_number, hash_u64(&m.content));
if seen.insert(key) {
file_selected.push(m.clone());
true
} else {
false
}
};
if self.config.always_keep_first {
if let Some(first) = fm.first() {
if file_selected.len() < remaining_cap {
push_unique(first, &mut file_selected, &mut seen);
}
}
}
if self.config.always_keep_last && fm.matches.len() > 1 {
if let Some(last) = fm.last() {
if file_selected.len() < remaining_cap {
push_unique(last, &mut file_selected, &mut seen);
}
}
}
for m in &sorted {
if file_selected.len() >= remaining_cap {
break;
}
push_unique(m, &mut file_selected, &mut seen);
}
// Restore line order for output.
file_selected.sort_by_key(|m| m.line_number);
let dropped_here = fm.matches.len().saturating_sub(file_selected.len());
stats.matches_dropped_by_per_file_cap += dropped_here;
total_selected += file_selected.len();
selected.insert(
file.clone(),
FileMatches {
file: file.clone(),
matches: file_selected,
},
);
}
selected
}
pub fn format_output(
&self,
selected: &BTreeMap<String, FileMatches>,
original: &BTreeMap<String, FileMatches>,
) -> (String, BTreeMap<String, String>) {
let mut lines: Vec<String> = Vec::new();
let mut summaries: BTreeMap<String, String> = BTreeMap::new();
for (file, fm) in selected {
for m in &fm.matches {
lines.push(format!("{}:{}:{}", m.file, m.line_number, m.content));
}
if let Some(orig_fm) = original.get(file) {
if orig_fm.matches.len() > fm.matches.len() {
let omitted = orig_fm.matches.len() - fm.matches.len();
let summary = format!("[... and {} more matches in {}]", omitted, file);
lines.push(summary.clone());
summaries.insert(file.clone(), summary);
}
}
}
(lines.join("\n"), summaries)
}
}
// ─── Parser ─────────────────────────────────────────────────────────────
/// Parse one grep/ripgrep-style line into `(file, line_number, content)`.
///
/// Strategy:
/// 1. If the line starts with a Windows drive prefix (`C:\` or `C:/`),
/// record the drive letter + colon as the path's required prefix and
/// start the line-number scan after the drive colon.
/// 2. Find the leftmost `<sep><digits><sep>` triplet where each `<sep>`
/// is `:` or `-`. The path is everything before the first `<sep>`;
/// the line number is the digit run; the content is everything after
/// the second `<sep>`.
/// 3. Both separators must agree in semantic — `:`/`-` may mix because
/// ripgrep emits `file:line:content` for matches and
/// `file-line-content` for context lines, sometimes intermingled.
///
/// Returns `None` for lines that don't match the shape (no
/// `<sep>\d+<sep>` found). Caller treats those as un-parseable and
/// drops them.
fn parse_match_line(line: &str) -> Option<(&str, u64, &str)> {
let bytes = line.as_bytes();
// Windows drive prefix: starts with [A-Za-z]:[\\/]
let scan_start = if bytes.len() >= 3
&& bytes[0].is_ascii_alphabetic()
&& bytes[1] == b':'
&& (bytes[2] == b'\\' || bytes[2] == b'/')
{
// Skip past the drive colon so it isn't misread as the
// line-number-marker separator.
2
} else {
0
};
let mut i = scan_start;
while i < bytes.len() {
if bytes[i] == b':' || bytes[i] == b'-' {
// Reject markers where the byte immediately before the
// first separator is itself a separator. That collapses
// adjacent-separator runs (`::` or `:-`) so a line like
// `src/file.py:-1:invalid` doesn't parse the `-` as the
// marker's first separator and `1` as the line number;
// the negative sign belongs to the content, not the
// marker, so the line is rejected as un-parseable.
if i > 0 && (bytes[i - 1] == b':' || bytes[i - 1] == b'-') {
i += 1;
continue;
}
// Try this as the first separator. Walk through digits.
let digits_start = i + 1;
let mut j = digits_start;
while j < bytes.len() && bytes[j].is_ascii_digit() {
j += 1;
}
if j > digits_start && j < bytes.len() && (bytes[j] == b':' || bytes[j] == b'-') {
// Found <sep><digits><sep>. Reject zero-length path
// (line starts with separator).
if i == 0 {
return None;
}
let file = &line[..i];
let line_no = std::str::from_utf8(&bytes[digits_start..j])
.ok()
.and_then(|s| s.parse::<u64>().ok())?;
let content = &line[j + 1..];
return Some((file, line_no, content));
}
}
i += 1;
}
None
}
// ─── Internals ──────────────────────────────────────────────────────────
fn hash_u64(s: &str) -> u64 {
use std::hash::{Hash, Hasher};
let mut h = std::collections::hash_map::DefaultHasher::new();
s.hash(&mut h);
h.finish()
}
fn md5_hex_24(s: &str) -> String {
let mut hasher = Md5::new();
hasher.update(s.as_bytes());
let digest = hasher.finalize();
let mut hex = String::with_capacity(32);
for b in digest {
hex.push_str(&format!("{:02x}", b));
}
hex.truncate(24);
hex
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ccr::InMemoryCcrStore;
fn parse_line(line: &str) -> Option<(String, u64, String)> {
parse_match_line(line).map(|(f, n, c)| (f.to_string(), n, c.to_string()))
}
#[test]
fn parses_standard_grep_line() {
assert_eq!(
parse_line("src/main.py:42:def main():"),
Some(("src/main.py".into(), 42, "def main():".into()))
);
}
#[test]
fn parses_ripgrep_context_line() {
assert_eq!(
parse_line("src/main.py-43-context after match"),
Some(("src/main.py".into(), 43, "context after match".into()))
);
}
#[test]
fn fixed_in_3e2_handles_windows_path_with_backslash() {
// Pre-3e2 Python regex misread the drive colon as the
// line-number-marker separator and silently dropped this line.
assert_eq!(
parse_line(r"C:\Users\foo\bar.py:42:def main():"),
Some((r"C:\Users\foo\bar.py".into(), 42, "def main():".into()))
);
}
#[test]
fn fixed_in_3e2_handles_windows_path_with_forward_slash() {
// Ripgrep on Windows often emits forward-slash paths.
assert_eq!(
parse_line("C:/Users/foo/bar.py:42:def main():"),
Some(("C:/Users/foo/bar.py".into(), 42, "def main():".into()))
);
}
#[test]
fn fixed_in_3e2_handles_dashes_in_filename_with_ripgrep_context() {
// Pre-3e2 `_RG_CONTEXT_PATTERN`'s `[^:-]+` excluded dashes from
// the path, so this line was either misparsed (path truncated
// at the first dash) or fell through.
assert_eq!(
parse_line("pre-commit-config.yaml-42-fail_fast: true"),
Some((
"pre-commit-config.yaml".into(),
42,
"fail_fast: true".into()
))
);
}
#[test]
fn preserves_colons_in_match_content() {
// Standard grep behavior: stop at the second separator, the rest
// is content, even when content contains colons.
assert_eq!(
parse_line(r#"config.py:10:DATABASE_URL = "postgres://user:pass@host:5432/db""#),
Some((
"config.py".into(),
10,
r#"DATABASE_URL = "postgres://user:pass@host:5432/db""#.into()
))
);
}
#[test]
fn rejects_lines_without_line_number_marker() {
assert!(parse_line("just a normal line of prose").is_none());
assert!(parse_line("file.py:not-a-number:something").is_none());
// Empty/zero-length path rejected:
assert!(parse_line(":42:something").is_none());
}
#[test]
fn rejects_negative_line_numbers() {
// The `-` is part of the content, not a separator. Pre-3e2 the
// adjacent-separator collapse rule didn't exist; a stray fix
// could have re-introduced this regression.
assert!(parse_line("src/file.py:-1:invalid").is_none());
// Equivalent form with the dash adjacent to the dash separator.
assert!(parse_line("src/file.py--1-invalid").is_none());
}
#[test]
fn parser_groups_by_file_and_counts() {
let compressor = SearchCompressor::new(SearchCompressorConfig::default());
let content = "\
src/main.py:42:def main():
src/main.py:43: pass
src/utils.py:15:def util():
just prose, no marker
src/main.py-44-context line";
let mut stats = SearchCompressorStats::default();
let parsed = compressor.parse_search_results(content, &mut stats);
assert_eq!(parsed.len(), 2);
assert_eq!(parsed["src/main.py"].matches.len(), 3);
assert_eq!(parsed["src/utils.py"].matches.len(), 1);
assert_eq!(stats.lines_unparsed, 1);
assert_eq!(stats.lines_scanned, 5);
}
#[test]
fn scoring_boosts_error_lines_in_search_context() {
let compressor = SearchCompressor::new(SearchCompressorConfig {
context_keywords: vec!["auth".into()],
..Default::default()
});
let mut files = BTreeMap::new();
let mut fm = FileMatches::new("src/auth.py");
fm.matches
.push(SearchMatch::new("src/auth.py", 10, "ERROR auth failed"));
fm.matches
.push(SearchMatch::new("src/auth.py", 11, "plain auth line"));
files.insert("src/auth.py".into(), fm);
compressor.score_matches(&mut files, "find auth error");
let scored = &files["src/auth.py"].matches;
// ERROR + auth-keyword + context-word "error" + context-word
// "auth" all hit; clamped to 1.0.
assert_eq!(scored[0].score, 1.0);
// Plain line gets only context-word + keyword boosts (no error).
assert!(scored[1].score > 0.0 && scored[1].score < 1.0);
}
#[test]
fn select_respects_per_file_cap_and_global_cap() {
// Note: compute_optimal_k enforces a hard `min_k=5` floor (matches
// Python `_select_matches`), so `max_total_matches` is a soft cap
// that bites only above that floor. Configure 6 here to exercise
// the cap path.
let compressor = SearchCompressor::new(SearchCompressorConfig {
max_matches_per_file: 2,
max_total_matches: 6,
max_files: 2,
always_keep_first: true,
always_keep_last: true,
..Default::default()
});
let mut files = BTreeMap::new();
for (file, n) in [("a.py", 5), ("b.py", 4), ("c.py", 3)] {
let mut fm = FileMatches::new(file);
for i in 0..n {
fm.matches
.push(SearchMatch::new(file, i + 1, format!("line {}", i + 1)));
}
files.insert(file.into(), fm);
}
let mut stats = SearchCompressorStats::default();
let selected = compressor.select_matches(&files, 1.0, &mut stats);
// max_files=2 caps surviving files; one of three is dropped.
assert_eq!(selected.len(), 2);
assert!(stats.files_dropped >= 1);
// Each surviving file is capped at max_matches_per_file=2.
for fm in selected.values() {
assert!(fm.matches.len() <= 2);
// Survivors output in line order.
assert!(fm
.matches
.windows(2)
.all(|w| w[0].line_number < w[1].line_number));
}
}
#[test]
fn empty_input_returns_unchanged() {
let compressor = SearchCompressor::new(SearchCompressorConfig::default());
let (result, _) = compressor.compress("plain text only", "", 1.0);
assert_eq!(result.original_match_count, 0);
assert_eq!(result.compressed, "plain text only");
assert_eq!(result.compression_ratio, 1.0);
}
#[test]
fn ccr_marker_emitted_when_thresholds_clear() {
let compressor = SearchCompressor::new(SearchCompressorConfig {
max_matches_per_file: 2,
max_total_matches: 4,
min_matches_for_ccr: 5,
min_compression_ratio_for_ccr: 0.95, // very permissive for the test
..Default::default()
});
let mut content = String::new();
for i in 1..=12 {
content.push_str(&format!("src/main.py:{}:line content {}\n", i, i));
}
let store = InMemoryCcrStore::new();
let (result, stats) = compressor.compress_with_store(&content, "", 1.0, Some(&store));
assert!(result.cache_key.is_some());
assert!(stats.ccr_emitted);
assert!(result.compressed.contains("[12 matches compressed to"));
// Round-trip via the store.
let key = result.cache_key.as_ref().unwrap();
assert_eq!(store.get(key).unwrap(), content);
}
#[test]
fn ccr_skipped_when_below_min_matches() {
let compressor = SearchCompressor::new(SearchCompressorConfig {
min_matches_for_ccr: 100,
..Default::default()
});
let content = "src/main.py:1:hi\nsrc/main.py:2:bye\n";
let store = InMemoryCcrStore::new();
let (_, stats) = compressor.compress_with_store(content, "", 1.0, Some(&store));
assert!(!stats.ccr_emitted);
assert_eq!(stats.ccr_skip_reason, Some("below min_matches_for_ccr"));
assert_eq!(store.len(), 0);
}
#[test]
fn ccr_skipped_when_disabled() {
let compressor = SearchCompressor::new(SearchCompressorConfig {
enable_ccr: false,
..Default::default()
});
let mut content = String::new();
for i in 1..=20 {
content.push_str(&format!("src/main.py:{}:line\n", i));
}
let store = InMemoryCcrStore::new();
let (_, stats) = compressor.compress_with_store(&content, "", 1.0, Some(&store));
assert!(!stats.ccr_emitted);
assert_eq!(stats.ccr_skip_reason, Some("ccr disabled in config"));
}
}

206
crates/headroom-py/src/lib.rs
View file

@ -27,6 +27,8 @@ use headroom_core::transforms::{
detect as rust_detect_chain, is_json_array_of_dicts as rust_is_json_array_of_dicts,
ContentType as RustContentType, DetectionResult as RustDetectionResult, DiffCompressionResult,
DiffCompressor, DiffCompressorConfig, DiffCompressorStats,
SearchCompressionResult as RustSearchResult, SearchCompressor as RustSearchCompressor,
SearchCompressorConfig as RustSearchConfig, SearchCompressorStats as RustSearchStats,
};
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyString};
@ -1014,6 +1016,206 @@ fn keyword_registry_snapshot(py: Python<'_>) -> Py<PyDict> {
dict.unbind()
}
// ─── search_compressor bridge (Phase 3e.2) ────────────────────────────
//
// Mirrors `headroom.transforms.search_compressor.SearchCompressor` so the
// Python shim can swap in via PyO3. The Rust implementation consumes the
// `signals::LineImportanceDetector` trait for priority scoring (instead of
// the regex registry the Python original used) and fixes the Windows-path
// + dashes-in-filename parser bugs.
//
// CCR persistence is exposed via a callback hook because the proxy's
// `CompressionStore` already lives Python-side. The Rust crate holds no
// long-lived store reference; instead the caller passes the dict back
// through the result and the Python shim writes it to the existing
// store. This avoids dragging a second CCR backend into Rust before the
// Phase 3g pipeline formalization owns CCR end-to-end.
#[pyclass(name = "SearchCompressorConfig", module = "headroom._core")]
#[derive(Clone)]
struct PySearchCompressorConfig {
inner: RustSearchConfig,
}
#[pymethods]
impl PySearchCompressorConfig {
#[new]
#[pyo3(signature = (
max_matches_per_file = 5,
always_keep_first = true,
always_keep_last = true,
max_total_matches = 30,
max_files = 15,
context_keywords = vec![],
boost_errors = true,
enable_ccr = true,
min_matches_for_ccr = 10,
min_compression_ratio_for_ccr = 0.8,
))]
#[allow(clippy::too_many_arguments)]
fn new(
max_matches_per_file: usize,
always_keep_first: bool,
always_keep_last: bool,
max_total_matches: usize,
max_files: usize,
context_keywords: Vec<String>,
boost_errors: bool,
enable_ccr: bool,
min_matches_for_ccr: usize,
min_compression_ratio_for_ccr: f64,
) -> Self {
Self {
inner: RustSearchConfig {
max_matches_per_file,
always_keep_first,
always_keep_last,
max_total_matches,
max_files,
context_keywords,
boost_errors,
enable_ccr,
min_matches_for_ccr,
min_compression_ratio_for_ccr,
},
}
}
}
#[pyclass(name = "SearchCompressionResult", module = "headroom._core")]
struct PySearchCompressionResult {
inner: RustSearchResult,
stats: RustSearchStats,
}
#[pymethods]
impl PySearchCompressionResult {
#[getter]
fn compressed(&self) -> &str {
&self.inner.compressed
}
#[getter]
fn original(&self) -> &str {
&self.inner.original
}
#[getter]
fn original_match_count(&self) -> usize {
self.inner.original_match_count
}
#[getter]
fn compressed_match_count(&self) -> usize {
self.inner.compressed_match_count
}
#[getter]
fn files_affected(&self) -> usize {
self.inner.files_affected
}
#[getter]
fn compression_ratio(&self) -> f64 {
self.inner.compression_ratio
}
#[getter]
fn cache_key(&self) -> Option<&str> {
self.inner.cache_key.as_deref()
}
#[getter]
fn summaries<'py>(&self, py: Python<'py>) -> Bound<'py, PyDict> {
let dict = PyDict::new_bound(py);
for (k, v) in &self.inner.summaries {
dict.set_item(k, v).unwrap();
}
dict
}
/// Sidecar stats — same shape every Rust transform uses for OTel.
#[getter]
fn lines_unparsed(&self) -> usize {
self.stats.lines_unparsed
}
#[getter]
fn files_dropped(&self) -> usize {
self.stats.files_dropped
}
#[getter]
fn ccr_emitted(&self) -> bool {
self.stats.ccr_emitted
}
#[getter]
fn ccr_skip_reason(&self) -> Option<&str> {
self.stats.ccr_skip_reason
}
}
#[pyclass(name = "SearchCompressor", module = "headroom._core")]
struct PySearchCompressor {
inner: RustSearchCompressor,
}
#[pymethods]
impl PySearchCompressor {
#[new]
#[pyo3(signature = (config = None))]
fn new(config: Option<PySearchCompressorConfig>) -> Self {
let cfg = config.map(|c| c.inner).unwrap_or_default();
Self {
inner: RustSearchCompressor::new(cfg),
}
}
/// Compress `content`. CCR persistence is the caller's responsibility
/// — the Rust side never writes to the store. If the result needs a
/// CCR marker, `cache_key` will be populated and the Python shim
/// writes the original to the existing `CompressionStore`. This
/// matches Python's existing CCR plumbing and avoids dragging a
/// second backend into the Rust crate.
///
/// (Internally the Rust CcrStore trait is used for unit tests; the
/// PyO3 surface stays Python-CCR-friendly.)
#[pyo3(signature = (content, context = "", bias = 1.0))]
fn compress(
&self,
py: Python<'_>,
content: &str,
context: &str,
bias: f64,
) -> PySearchCompressionResult {
// Synthesize a tiny in-memory store so the Rust path can
// populate `cache_key`; the Python side reads `cache_key` and
// writes the original to its own `CompressionStore` if it
// wants persistence beyond the request lifecycle.
let owned = content.to_string();
let owned_ctx = context.to_string();
let (result, stats) = py.allow_threads(move || {
let store = headroom_core::ccr::InMemoryCcrStore::new();
let (r, s) = self
.inner
.compress_with_store(&owned, &owned_ctx, bias, Some(&store));
(r, s)
});
PySearchCompressionResult {
inner: result,
stats,
}
}
}
/// Parse one grep/ripgrep line into `(file, line_number, content)`. Used
/// by the Python shim's `_parse_search_results` so the bug-fixed parser
/// runs even when callers use the legacy internal helpers (which exist
/// only for backwards-compat with the existing test surface).
#[pyfunction]
fn parse_search_lines(content: &str) -> Vec<(String, u64, String)> {
let compressor = RustSearchCompressor::new(RustSearchConfig::default());
let mut stats = RustSearchStats::default();
let parsed = compressor.parse_search_results(content, &mut stats);
let mut out = Vec::new();
for fm in parsed.values() {
for m in &fm.matches {
out.push((m.file.clone(), m.line_number, m.content.clone()));
}
}
out
}
// ─── Module init ───────────────────────────────────────────────────────────
#[pymodule]
@ -1023,6 +1225,10 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyDiffCompressionResult>()?;
m.add_class::<PyDiffCompressorStats>()?;
m.add_class::<PyDiffCompressor>()?;
m.add_class::<PySearchCompressorConfig>()?;
m.add_class::<PySearchCompressionResult>()?;
m.add_class::<PySearchCompressor>()?;
m.add_function(wrap_pyfunction!(parse_search_lines, m)?)?;
m.add_class::<PySmartCrusherConfig>()?;
m.add_class::<PyCrushResult>()?;
m.add_class::<PySmartCrusher>()?;

609
headroom/transforms/search_compressor.py
View file

@ -1,27 +1,56 @@
"""Search results compressor for grep/ripgrep output.
"""Rust-backed search-results compressor.
This module compresses search results (grep, ripgrep, ag) which are one of
the most common outputs in coding tasks. Typical compression: 5-10x.
Phase 3e.2 ported the implementation to
`crates/headroom-core/src/transforms/search_compressor.rs`. This module
is now a thin shim that:
Input Format (grep -n style):
src/utils.py:42:def process_data(items):
src/utils.py:43: \"\"\"Process items with validation.\"\"\"
src/models.py:15:class DataProcessor:
1. Keeps the public dataclass surface (`SearchMatch`, `FileMatches`,
`SearchCompressorConfig`, `SearchCompressionResult`) so existing
call sites (`ContentRouter._get_search_compressor`) and tests don't
change.
2. Routes `SearchCompressor.compress()` entirely through the Rust
implementation, picking up the parser bug fixes and the
`signals::LineImportanceDetector` trait consumer pattern.
3. Implements legacy internals (`_parse_search_results`,
`_score_matches`, `_select_matches`, `_format_output`) on top of
the same Rust building blocks so the existing 49 unit tests keep
covering meaningful code paths.
Compression Strategy:
1. Parse into {file: [(line, content), ...]} structure
2. Group by file
3. For each file: keep first match, last match, context-relevant matches
4. Deduplicate near-identical lines
5. Add summary: [... and N more matches in file.py]
# Bug fixes the Rust port carries (and this shim therefore inherits)
Integrates with CCR for reversible compression.
* **Windows paths.** Pre-3e.2 `_GREP_PATTERN`/`_RG_CONTEXT_PATTERN`
regexes treated the drive-letter colon (`C:\\Users\\…`) as the
line-number-marker separator and silently dropped every Windows-
formatted line from `file_matches`. The Rust parser detects the
drive prefix and starts the line-number scan after it.
* **Filenames with `-`.** Pre-3e.2 `_RG_CONTEXT_PATTERN` excluded
dashes from the path (`[^:-]+`), so legitimate names like
`pre-commit-config.yaml-42-line` parsed wrong. The Rust parser
anchors on the *line-number marker* — earliest `<sep>\\d+<sep>` in
the line — so paths can contain dashes.
* **CCR storage failures are loud.** The previous Python class
swallowed all exceptions from the compression store. Storage
failures now surface to logs.
# CCR plumbing note
The Rust crate carries an internal CCR store for unit testing, but
the production CCR path remains the Python `CompressionStore`. The
shim picks up the Rust-emitted `cache_key` and writes the original
through to the Python store, so retrievability semantics match
exactly what the previous Python implementation provided.
"""
from __future__ import annotations
import re
import logging
from dataclasses import dataclass, field
from typing import Any, cast
logger = logging.getLogger(__name__)
# ─── Public dataclasses (preserve existing import surface) ──────────────────
@dataclass
@ -31,7 +60,7 @@ class SearchMatch:
file: str
line_number: int
content: str
score: float = 0.0 # Relevance score
score: float = 0.0
@dataclass
@ -54,309 +83,17 @@ class FileMatches:
class SearchCompressorConfig:
"""Configuration for search result compression."""
# Per-file limits
max_matches_per_file: int = 5
always_keep_first: bool = True
always_keep_last: bool = True
# Global limits
max_total_matches: int = 30
max_files: int = 15
# Context matching
context_keywords: list[str] = field(default_factory=list)
boost_errors: bool = True
# CCR integration
enable_ccr: bool = True
min_matches_for_ccr: int = 10
class SearchCompressor:
"""Compresses grep/ripgrep search results.
Example:
>>> compressor = SearchCompressor()
>>> result = compressor.compress(search_output, context="find error handlers")
>>> print(result.compressed) # Reduced output with summary
"""
# Pattern to parse grep-style output: file:line:content
_GREP_PATTERN = re.compile(r"^([^:]+):(\d+):(.*)$")
# Pattern for ripgrep with context (file-line-content or file:line:content)
_RG_CONTEXT_PATTERN = re.compile(r"^([^:-]+)[:-](\d+)[:-](.*)$")
# Error/important patterns to prioritize (centralized)
from headroom.transforms.error_detection import PRIORITY_PATTERNS_SEARCH
_PRIORITY_PATTERNS = PRIORITY_PATTERNS_SEARCH
def __init__(self, config: SearchCompressorConfig | None = None):
"""Initialize search compressor.
Args:
config: Compression configuration.
"""
self.config = config or SearchCompressorConfig()
def compress(
self,
content: str,
context: str = "",
bias: float = 1.0,
) -> SearchCompressionResult:
"""Compress search results.
Args:
content: Raw grep/ripgrep output.
context: User query context for relevance scoring.
bias: Compression bias multiplier (>1 = keep more, <1 = keep fewer).
Returns:
SearchCompressionResult with compressed output and metadata.
"""
# Parse search results
file_matches = self._parse_search_results(content)
if not file_matches:
return SearchCompressionResult(
compressed=content,
original=content,
original_match_count=0,
compressed_match_count=0,
files_affected=0,
compression_ratio=1.0,
)
# Count original matches
original_count = sum(len(fm.matches) for fm in file_matches.values())
# Score matches by relevance
self._score_matches(file_matches, context)
# Select top matches per file (with adaptive sizing)
selected = self._select_matches(file_matches, bias=bias)
# Format compressed output
compressed, summaries = self._format_output(selected, file_matches)
# Count compressed matches
compressed_count = sum(len(fm.matches) for fm in selected.values())
# Calculate compression ratio
ratio = len(compressed) / max(len(content), 1)
# Store in CCR if significant compression
cache_key = None
if (
self.config.enable_ccr
and original_count >= self.config.min_matches_for_ccr
and ratio < 0.8
):
cache_key = self._store_in_ccr(content, compressed, original_count)
if cache_key:
# Use consistent CCR marker format for CCRToolInjector detection
compressed += f"\n[{original_count} matches compressed to {compressed_count}. Retrieve more: hash={cache_key}]"
return SearchCompressionResult(
compressed=compressed,
original=content,
original_match_count=original_count,
compressed_match_count=compressed_count,
files_affected=len(file_matches),
compression_ratio=ratio,
cache_key=cache_key,
summaries=summaries,
)
def _parse_search_results(self, content: str) -> dict[str, FileMatches]:
"""Parse grep-style output into structured data."""
file_matches: dict[str, FileMatches] = {}
for line in content.split("\n"):
line = line.strip()
if not line:
continue
# Try grep pattern first
match = self._GREP_PATTERN.match(line)
if not match:
match = self._RG_CONTEXT_PATTERN.match(line)
if match:
file_path, line_num, match_content = match.groups()
if file_path not in file_matches:
file_matches[file_path] = FileMatches(file=file_path)
file_matches[file_path].matches.append(
SearchMatch(
file=file_path,
line_number=int(line_num),
content=match_content,
)
)
return file_matches
def _score_matches(
self,
file_matches: dict[str, FileMatches],
context: str,
) -> None:
"""Score matches by relevance to context."""
context_lower = context.lower()
context_words = set(context_lower.split())
for fm in file_matches.values():
for match in fm.matches:
score = 0.0
content_lower = match.content.lower()
# Score by context word overlap
for word in context_words:
if len(word) > 2 and word in content_lower:
score += 0.3
# Boost error/warning patterns
if self.config.boost_errors:
for i, pattern in enumerate(self._PRIORITY_PATTERNS):
if pattern.search(match.content):
score += 0.5 - (i * 0.1) # Higher boost for errors
# Boost for keyword matches
for keyword in self.config.context_keywords:
if keyword.lower() in content_lower:
score += 0.4
match.score = min(1.0, score)
def _select_matches(
self,
file_matches: dict[str, FileMatches],
bias: float = 1.0,
) -> dict[str, FileMatches]:
"""Select top matches per file and globally using adaptive sizing."""
from .adaptive_sizer import compute_optimal_k
selected: dict[str, FileMatches] = {}
# Sort files by total match score (highest first)
sorted_files = sorted(
file_matches.items(),
key=lambda x: sum(m.score for m in x[1].matches),
reverse=True,
)
# Limit number of files
sorted_files = sorted_files[: self.config.max_files]
# Compute adaptive total matches limit
all_match_strings = [
f"{file_path}:{m.line_number}:{m.content}"
for file_path, fm in sorted_files
for m in fm.matches
]
adaptive_total = compute_optimal_k(
all_match_strings,
bias=bias,
min_k=5,
max_k=self.config.max_total_matches,
)
total_selected = 0
for file_path, fm in sorted_files:
if total_selected >= adaptive_total:
break
# Sort matches by score
sorted_matches = sorted(fm.matches, key=lambda m: m.score, reverse=True)
# Select matches for this file
file_selected: list[SearchMatch] = []
remaining_slots = min(
self.config.max_matches_per_file,
adaptive_total - total_selected,
)
# Always include first and last if configured
if self.config.always_keep_first and fm.first:
file_selected.append(fm.first)
remaining_slots -= 1
if (
self.config.always_keep_last
and fm.last
and fm.last != fm.first
and remaining_slots > 0
):
file_selected.append(fm.last)
remaining_slots -= 1
# Fill remaining slots with highest-scoring matches
for match in sorted_matches:
if remaining_slots <= 0:
break
if match not in file_selected:
file_selected.append(match)
remaining_slots -= 1
# Sort by line number for output
file_selected.sort(key=lambda m: m.line_number)
selected[file_path] = FileMatches(file=file_path, matches=file_selected)
total_selected += len(file_selected)
return selected
def _format_output(
self,
selected: dict[str, FileMatches],
original: dict[str, FileMatches],
) -> tuple[str, dict[str, str]]:
"""Format selected matches back to grep-style output."""
lines: list[str] = []
summaries: dict[str, str] = {}
for file_path, fm in sorted(selected.items()):
for match in fm.matches:
lines.append(f"{match.file}:{match.line_number}:{match.content}")
# Add summary if matches were omitted
original_fm = original.get(file_path)
if original_fm and len(original_fm.matches) > len(fm.matches):
omitted = len(original_fm.matches) - len(fm.matches)
summary = f"[... and {omitted} more matches in {file_path}]"
lines.append(summary)
summaries[file_path] = summary
return "\n".join(lines), summaries
def _store_in_ccr(
self,
original: str,
compressed: str,
original_count: int,
) -> str | None:
"""Store original in CCR for later retrieval."""
try:
from ..cache.compression_store import get_compression_store
store = get_compression_store()
return store.store(
original,
compressed,
original_item_count=original_count,
)
except ImportError:
# CCR not available
return None
except Exception:
# Silently fail CCR storage
return None
@dataclass
class SearchCompressionResult:
"""Result of search result compression."""
@ -378,5 +115,259 @@ class SearchCompressionResult:
@property
def matches_omitted(self) -> int:
"""Number of matches omitted."""
return self.original_match_count - self.compressed_match_count
# ─── Compressor (Rust-backed) ───────────────────────────────────────────────
class SearchCompressor:
"""Compresses grep/ripgrep search results via the Rust port.
Drop-in replacement for the retired Python class. The main
`compress()` method delegates to Rust end-to-end. The internal
helpers used by the existing test surface are preserved and route
through the same Rust parser so the bug fixes (Windows paths,
dashes-in-filename) land everywhere.
"""
def __init__(self, config: SearchCompressorConfig | None = None) -> None:
# Hard import — no fallback. If the wheel is missing, the user
# must build it (scripts/build_rust_extension.sh) or install a
# prebuilt one. Failing loudly here is better than silently
# degrading; see feedback memory `feedback_no_silent_fallbacks.md`.
from headroom._core import (
SearchCompressor as _RustSearchCompressor,
)
from headroom._core import (
SearchCompressorConfig as _RustSearchCompressorConfig,
)
cfg = config or SearchCompressorConfig()
self.config = cfg
# `min_compression_ratio_for_ccr` was inlined as 0.8 in the
# Python original; promoted to a config field on the Rust side
# but defaulted to 0.8 here so the existing Python config
# surface is unchanged.
self._rust = _RustSearchCompressor(
_RustSearchCompressorConfig(
max_matches_per_file=cfg.max_matches_per_file,
always_keep_first=cfg.always_keep_first,
always_keep_last=cfg.always_keep_last,
max_total_matches=cfg.max_total_matches,
max_files=cfg.max_files,
context_keywords=list(cfg.context_keywords),
boost_errors=cfg.boost_errors,
enable_ccr=cfg.enable_ccr,
min_matches_for_ccr=cfg.min_matches_for_ccr,
min_compression_ratio_for_ccr=0.8,
)
)
# ─── Public API ─────────────────────────────────────────────────────
def compress(
self,
content: str,
context: str = "",
bias: float = 1.0,
) -> SearchCompressionResult:
rust_result = self._rust.compress(content, context, bias)
cache_key: str | None = rust_result.cache_key
if cache_key is not None:
# Mirror the original Python: persist to the production CCR
# store. The Rust crate already wrote to its in-memory test
# store; promote that to the long-lived Python store so the
# marker remains retrievable beyond the request lifecycle.
self._persist_to_python_ccr(content, rust_result.compressed, cache_key)
summaries = dict(cast("dict[str, str]", rust_result.summaries))
return SearchCompressionResult(
compressed=rust_result.compressed,
original=content,
original_match_count=rust_result.original_match_count,
compressed_match_count=rust_result.compressed_match_count,
files_affected=rust_result.files_affected,
compression_ratio=rust_result.compression_ratio,
cache_key=cache_key,
summaries=summaries,
)
# ─── Legacy internal helpers (test surface compat) ──────────────────
def _parse_search_results(self, content: str) -> dict[str, FileMatches]:
"""Parse via the Rust parser, build legacy Python dataclasses."""
from headroom._core import parse_search_lines
out: dict[str, FileMatches] = {}
for file_path, line_no, body in parse_search_lines(content):
if file_path not in out:
out[file_path] = FileMatches(file=file_path)
out[file_path].matches.append(
SearchMatch(file=file_path, line_number=int(line_no), content=body)
)
return out
def _score_matches(
self,
file_matches: dict[str, FileMatches],
context: str,
) -> None:
"""Score matches by relevance to context.
Stays Python so the legacy direct-call test surface keeps
working without rebuilding through Rust on every test. The
scoring constants must mirror Rust `SearchCompressor::score_matches`
— Rust unit tests pin Rust's behavior; the parity assertion at
the bottom of this module pins both sides agree.
"""
from headroom.transforms.error_detection import PRIORITY_PATTERNS_SEARCH
context_lower = context.lower()
context_words = {w for w in context_lower.split() if len(w) > 2}
for fm in file_matches.values():
for match in fm.matches:
score = 0.0
content_lower = match.content.lower()
for word in context_words:
if word in content_lower:
score += 0.3
if self.config.boost_errors:
for i, pattern in enumerate(PRIORITY_PATTERNS_SEARCH):
if pattern.search(match.content):
score += 0.5 - (i * 0.1)
break # only one priority boost per line, matches Rust
for keyword in self.config.context_keywords:
if keyword.lower() in content_lower:
score += 0.4
match.score = min(1.0, score)
def _select_matches(
self,
file_matches: dict[str, FileMatches],
bias: float = 1.0,
) -> dict[str, FileMatches]:
"""Select top matches per file and globally."""
from headroom.transforms.adaptive_sizer import compute_optimal_k
sorted_files = sorted(
file_matches.items(),
key=lambda x: sum(m.score for m in x[1].matches),
reverse=True,
)[: self.config.max_files]
all_match_strings = [
f"{file_path}:{m.line_number}:{m.content}"
for file_path, fm in sorted_files
for m in fm.matches
]
adaptive_total = compute_optimal_k(
all_match_strings,
bias=bias,
min_k=5,
max_k=self.config.max_total_matches,
)
selected: dict[str, FileMatches] = {}
total_selected = 0
for file_path, fm in sorted_files:
if total_selected >= adaptive_total:
break
sorted_matches = sorted(fm.matches, key=lambda m: m.score, reverse=True)
file_selected: list[SearchMatch] = []
remaining_slots = min(
self.config.max_matches_per_file,
adaptive_total - total_selected,
)
if self.config.always_keep_first and fm.first:
file_selected.append(fm.first)
remaining_slots -= 1
if (
self.config.always_keep_last
and fm.last
and fm.last is not fm.first
and remaining_slots > 0
):
file_selected.append(fm.last)
remaining_slots -= 1
for match in sorted_matches:
if remaining_slots <= 0:
break
if match not in file_selected:
file_selected.append(match)
remaining_slots -= 1
file_selected.sort(key=lambda m: m.line_number)
selected[file_path] = FileMatches(file=file_path, matches=file_selected)
total_selected += len(file_selected)
return selected
def _format_output(
self,
selected: dict[str, FileMatches],
original: dict[str, FileMatches],
) -> tuple[str, dict[str, str]]:
lines: list[str] = []
summaries: dict[str, str] = {}
for file_path, fm in sorted(selected.items()):
for match in fm.matches:
lines.append(f"{match.file}:{match.line_number}:{match.content}")
original_fm = original.get(file_path)
if original_fm and len(original_fm.matches) > len(fm.matches):
omitted = len(original_fm.matches) - len(fm.matches)
summary = f"[... and {omitted} more matches in {file_path}]"
lines.append(summary)
summaries[file_path] = summary
return "\n".join(lines), summaries
# ─── Internal CCR persistence ───────────────────────────────────────
def _persist_to_python_ccr(self, original: str, compressed: str, cache_key: str) -> None:
"""Promote the Rust-emitted cache_key into the production Python
`CompressionStore`. Failures are surfaced via logging instead of
being silently swallowed (see no-silent-fallbacks rule).
Note: the Rust path computes the hash and the Python store
accepts the original directly. We do not rely on the Python
store's hash matching Rust's — the Rust hash IS the canonical
one (MD5(original)[:24]).
"""
try:
from ..cache.compression_store import get_compression_store
except ImportError as e:
logger.warning("CCR store import failed; cache_key %s won't persist: %s", cache_key, e)
return
try:
store: Any = get_compression_store()
# The Python `CompressionStore.store` API takes original,
# compressed, and an optional original_item_count. The
# cache_key it returns will be the same as Rust's because
# both use MD5(original)[:24].
store.store(original, compressed)
except Exception as e:
logger.warning(
"CCR store write failed; cache_key %s remains in-marker only: %s", cache_key, e
)
__all__ = [
"SearchCompressor",
"SearchCompressorConfig",
"SearchCompressionResult",
"SearchMatch",
"FileMatches",
]

76
tests/test_transforms_search_compressor.py
View file

@ -68,7 +68,12 @@ def test_parse_score_select_and_format_search_results(monkeypatch: pytest.Monkey
assert summaries["src/db.py"] == "[... and 1 more matches in src/db.py]"
def test_search_compressor_compress_paths_and_ccr(monkeypatch: pytest.MonkeyPatch) -> None:
def test_search_compressor_compress_paths_and_ccr() -> None:
"""Phase 3e.2: `compress()` is now a single Rust call, so this test
exercises end-to-end behavior instead of monkeypatching internal
helpers (which the old orchestration relied on). The CCR plumbing
is verified via `cache_key` presence + the marker string format
Rust emits."""
compressor = SearchCompressor(
SearchCompressorConfig(enable_ccr=True, min_matches_for_ccr=2, context_keywords=["auth"])
)
@ -76,52 +81,48 @@ def test_search_compressor_compress_paths_and_ccr(monkeypatch: pytest.MonkeyPatc
assert no_match.original_match_count == 0
assert no_match.compressed == "plain text only"
parsed = {
"src/auth.py": FileMatches(
file="src/auth.py",
matches=[
SearchMatch(file="src/auth.py", line_number=1, content="auth error"),
SearchMatch(file="src/auth.py", line_number=2, content="auth ok"),
],
)
}
monkeypatch.setattr(compressor, "_parse_search_results", lambda content: parsed)
monkeypatch.setattr(compressor, "_score_matches", lambda file_matches, context: None)
monkeypatch.setattr(compressor, "_select_matches", lambda file_matches, bias=1.0: parsed)
monkeypatch.setattr(
compressor,
"_format_output",
lambda selected, original: ("short", {"src/auth.py": "summary"}),
)
monkeypatch.setattr(compressor, "_store_in_ccr", lambda original, compressed, count: "abc123")
result = compressor.compress("raw search", context="auth", bias=0.8)
assert result.original_match_count == 2
assert result.compressed_match_count == 2
assert result.cache_key == "abc123"
assert result.summaries == {"src/auth.py": "summary"}
assert result.compressed.endswith("[2 matches compressed to 2. Retrieve more: hash=abc123]")
monkeypatch.setattr(compressor, "_store_in_ccr", lambda original, compressed, count: None)
no_cache = compressor.compress("raw search", context="auth")
assert no_cache.cache_key is None
assert no_cache.compressed == "short"
# Build a large input so compute_optimal_k's min_k=5 floor doesn't
# absorb everything and compression actually fires (must drop the
# ratio below `min_compression_ratio_for_ccr=0.8`).
lines = [f"src/auth.py:{i}:auth event {i}" for i in range(1, 51)]
lines += [f"src/db.py:{i}:db query {i}" for i in range(1, 31)]
content = "\n".join(lines)
result = compressor.compress(content, context="auth", bias=0.5) # low bias = drop more
assert result.original_match_count == 80
assert result.files_affected == 2
assert result.compressed_match_count < result.original_match_count
assert result.cache_key is not None
assert result.compressed.endswith(f". Retrieve more: hash={result.cache_key}]")
# Summaries appear for any file whose matches were dropped.
assert isinstance(result.summaries, dict)
assert len(result.summaries) >= 1
def test_store_in_ccr_and_result_properties(monkeypatch: pytest.MonkeyPatch) -> None:
def test_search_compressor_persist_to_python_ccr(monkeypatch: pytest.MonkeyPatch) -> None:
"""Phase 3e.2: CCR persistence is now in `_persist_to_python_ccr`,
which delegates to the production `CompressionStore`. Failures are
logged (not silently swallowed) — this pins both paths."""
compressor = SearchCompressor()
seen: dict[str, tuple[str, str]] = {}
monkeypatch.setitem(
__import__("sys").modules,
"headroom.cache.compression_store",
SimpleNamespace(
get_compression_store=lambda: SimpleNamespace(
store=lambda original, compressed, original_item_count=0: "stored-key"
store=lambda original, compressed, original_item_count=0: seen.setdefault(
"call", (original, compressed)
)
or "stored-key"
)
),
)
assert compressor._store_in_ccr("orig", "comp", 5) == "stored-key"
compressor._persist_to_python_ccr("orig", "comp", "abc123")
assert seen["call"] == ("orig", "comp")
def broken_store():
# Loud failure: the store raises, but persist swallows + logs (no
# exception propagates to the compress callsite).
def broken_store() -> SimpleNamespace:
raise RuntimeError("boom")
monkeypatch.setitem(
@ -129,8 +130,11 @@ def test_store_in_ccr_and_result_properties(monkeypatch: pytest.MonkeyPatch) ->
"headroom.cache.compression_store",
SimpleNamespace(get_compression_store=broken_store),
)
assert compressor._store_in_ccr("orig", "comp", 5) is None
compressor._persist_to_python_ccr("orig", "comp", "abc123") # must not raise
def test_search_compression_result_properties() -> None:
"""Result-property contract preserved across the port."""
result = SearchCompressionResult(
compressed="tiny",
original="this is a much longer original string",