From b7be3814f1d38375bc27901272bbe919e6b35940 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Tue, 16 Jun 2026 20:21:13 -0700 Subject: [PATCH] =?UTF-8?q?feat:=20compression=20extraction=20=E2=80=94=20?= =?UTF-8?q?Rust=20knob=20exposure,=20CCR=20hardening,=20traffic=20audits?= =?UTF-8?q?=20(#818)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description A data-driven push for better compression savings without accuracy loss, in four parts: expose and tune the Rust compressor knobs, harden the CCR retrieval store, add traffic-audit tooling that sizes opportunities from real transcripts, and introduce **read maturation** — a new, live-validated mechanism that compresses Read outputs *before* they ever enter the provider prefix cache. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [x] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made ### 1. Rust compressor extraction - Expose `lossless_min_savings_ratio` end-to-end and lower the default 0.30 → 0.15 (lockstep across Rust, PyO3, and both Python config classes) so the lossless Table/CSV compaction path wins more often. - Expose the `CompactConfig` heuristics (core-field fraction, heterogeneity ratio, flatten cap, bucket bounds) through PyO3 + Python. - `SearchCompressor` grouped-by-file output (`rg --heading` style — path once per file instead of per match). Library default off; the proxy enables it in token mode. - Complete `factor_out_constants`: constant fields now emit once in a `_constant_fields` sentinel with slim rows (defensive per-item value match; default off). - `ContentRouter` accepts a SmartCrusher config override and the search-grouping knob. ### 2. CCR store hardening - Session-scale TTL: 300s → 1800s (CCRConfig, CompressionEntry, CompressionStore, Rust `DEFAULT_TTL` — lockstep). - **SQLite is the default CCR backend** (`~/.headroom/ccr_store.db`, WAL): survives proxy restarts and is shared across workers. `HEADROOM_CCR_BACKEND=memory` opts out. - Multi-worker safety: `busy_timeout`, and corruption detection narrowed so transient `SQLITE_BUSY` errors can never trigger database deletion. - Data-at-rest hygiene: `chmod 600` on db + sidecars, expired rows swept at open. - Retrieval-miss messages are actionable (re-read the file / re-run the command). ### 3. Traffic audit tooling (measure before tuning) - `headroom audit-reads`: sizes Read opportunities from local Claude Code transcripts (read share, stale %, line-number overhead, context residency, cache-death windows). - `--simulate-maturation`: Mechanism B risk sizing (re-read rates, never-touched-again share, quiesce coverage, at-risk edits). - `--codex`: shell-read classifier for Codex transcripts (rtk-wrapper aware, workdir resolution). - Findings that shaped this PR (81 sessions): Reads are 67% of tool bytes; median Read lingers 118 turns (~13x lifetime cost); a prototyped repeat-Read dedup measured 0.1% and was **removed** rather than shipped as dead code. ### 4. Read maturation (Mechanism B) — experimental, default OFF - Activity-based: a fresh large Read is held **out** of the provider cache (trailing breakpoint relocated before it), stays verbatim while its file is active, and matures into a CCR-backed marker once the file is quiet for `quiesce_turns` (default 5; `max_hold_turns` bounds busy files). - Only the final compressed form ever enters the cache — **no cached byte is ever mutated**; matured markers replay byte-identically. - Wired into the Anthropic handler behind `--read-maturation` / `HEADROOM_READ_MATURATION=1`; session state rides on the prefix tracker; advisory (can never fail a request). - Live-validated against the Anthropic API: held content excluded from cache_creation; after maturation the prior cached prefix still served — the no-bust invariant holds end-to-end. ### 5. Rebase / CI fixups (this update) - Rebased onto latest `main` (was 28 commits behind): picks up `ci: pass CODECOV_TOKEN to coverage uploads (#968)`, which is what was turning the 4 test shards red — the tests themselves passed (1528) but the post-test codecov upload exited non-zero on a protected branch. - Resolved the duplicate `lossless_min_savings_ratio` that two independent main/branch additions left in `SmartCrusherConfig` and the Rust-config kwarg (import-time `SyntaxError` + mypy `no-redef`). - Aligned CCR tests with the new defaults (SQLite backend, 1800s TTL) across `test_ccr`, `test_adapter_hooks`, `test_compression_store`, `test_proxy_ccr`, and the lossy row-drop bridge test. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ python -m pytest tests/test_proxy_ccr.py tests/test_ccr.py tests/test_compression_store.py tests/test_adapter_hooks.py tests/test_ccr_row_drop_store_bridge.py -q 170 passed, 4 warnings in 42.49s $ python -m pytest tests/test_audit_reads.py tests/test_audit_codex.py tests/test_read_maturation.py tests/test_transforms_content_router.py tests/test_smart_crusher_toin_attachment.py -q 83 passed $ mypy headroom/ Success: no issues found in 365 source files $ python -m compileall headroom/ -q COMPILE-OK # CI (run 27488990477, pre-rebase head): all 4 shards ran to completion — # "1528 passed, 120 skipped, 4922 deselected" # The red shards were the codecov upload step, not test failures; fixed by # the #968 rebase above. ``` ## Real Behavior Proof - Environment: macOS (darwin), Python 3.12 venv; branch `feat/compression-extraction` rebased onto `origin/main` (head 7cb0f43b); GitHub Actions CI run 27488990477 for the test shards - Exact command / steps: rebased onto latest main (clean, 13 commits replayed, 0 conflicts); ran the pytest suites and mypy above locally; inspected CI shard logs to confirm the failure was the codecov upload, not the test phase - Observed result: 253 targeted tests pass locally; mypy clean on 365 files; CI test phase reports `1528 passed, 120 skipped`; the only red step (codecov `upload-coverage` → "Token required because branch is protected") is resolved by the rebased-in #968 CODECOV_TOKEN fix - Not tested: the read-maturation live-API no-bust validation (`tests/test_live/`) was not re-run in this rebase pass (requires provider keys); it was validated when the feature first landed, and no maturation code changed in the rebase — only CCR-default test assertions and the duplicate-field resolution ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Additional Notes CHANGELOG is generated by release-please from the conventional commits, so the CHANGELOG box is intentionally left unchecked. "Manual testing performed" is unchecked deliberately — see `Real Behavior Proof` → `Not tested` for the exact boundary (the live-API maturation validation was not re-run in this rebase pass). ### Follow-ups (tracked, not in this PR) - Mechanism B provider extensions: OpenAI-family wiring (no breakpoint hold — bounded near-tail bust) and the Codex runtime read-detector (the audit classifier is the prototype). - Pilot enablement playbook: run `audit-reads --simulate-maturation` on target traffic → pick `quiesce_turns` → enable via env → watch cache hit rate + `read_maturation:N` transform tags. --- crates/headroom-core/src/ccr/mod.rs | 7 +- .../src/transforms/search_compressor.rs | 31 +- .../smart_crusher/compaction/mod.rs | 10 + .../src/transforms/smart_crusher/config.rs | 40 +- .../src/transforms/smart_crusher/crusher.rs | 130 ++++++- crates/headroom-parity/src/lib.rs | 4 + crates/headroom-py/src/lib.rs | 44 ++- headroom/audit/__init__.py | 17 + headroom/audit/codex.py | 217 +++++++++++ headroom/audit/maturation.py | 201 ++++++++++ headroom/audit/reads.py | 344 ++++++++++++++++++ headroom/cache/backends/__init__.py | 26 +- headroom/cache/backends/sqlite.py | 275 ++++++++++++++ headroom/cache/compression_store.py | 49 ++- headroom/ccr/mcp_server.py | 7 +- headroom/cli/__init__.py | 1 + headroom/cli/audit.py | 107 ++++++ headroom/cli/main.py | 1 + headroom/config.py | 28 +- headroom/proxy/server.py | 11 +- headroom/transforms/content_router.py | 20 +- headroom/transforms/read_lifecycle.py | 17 +- headroom/transforms/search_compressor.py | 6 + headroom/transforms/smart_crusher.py | 38 +- tests/test_adapter_hooks.py | 10 +- tests/test_audit_codex.py | 135 +++++++ tests/test_audit_reads.py | 195 ++++++++++ tests/test_ccr.py | 2 +- tests/test_ccr_row_drop_store_bridge.py | 8 +- tests/test_ccr_sqlite_backend.py | 229 ++++++++++++ tests/test_compression_store.py | 6 +- tests/test_compressor_config_exposure.py | 143 ++++++++ tests/test_live/__init__.py | 0 tests/test_live/conftest.py | 33 ++ tests/test_live/test_live_anthropic.py | 149 ++++++++ tests/test_live/test_live_openai.py | 109 ++++++ tests/test_proxy_ccr.py | 4 +- tests/test_proxy_scalability.py | 28 +- 38 files changed, 2606 insertions(+), 76 deletions(-) create mode 100644 headroom/audit/__init__.py create mode 100644 headroom/audit/codex.py create mode 100644 headroom/audit/maturation.py create mode 100644 headroom/audit/reads.py create mode 100644 headroom/cache/backends/sqlite.py create mode 100644 headroom/cli/audit.py create mode 100644 tests/test_audit_codex.py create mode 100644 tests/test_audit_reads.py create mode 100644 tests/test_ccr_sqlite_backend.py create mode 100644 tests/test_compressor_config_exposure.py create mode 100644 tests/test_live/__init__.py create mode 100644 tests/test_live/conftest.py create mode 100644 tests/test_live/test_live_anthropic.py create mode 100644 tests/test_live/test_live_openai.py diff --git a/crates/headroom-core/src/ccr/mod.rs b/crates/headroom-core/src/ccr/mod.rs index 38d207a37..2dd3ef28a 100644 --- a/crates/headroom-core/src/ccr/mod.rs +++ b/crates/headroom-core/src/ccr/mod.rs @@ -59,8 +59,11 @@ pub trait CcrStore: Send + Sync { /// Default capacity — matches Python's `CompressionStore` default. pub const DEFAULT_CAPACITY: usize = 1000; -/// Default TTL — 5 minutes, matching Python. -pub const DEFAULT_TTL: Duration = Duration::from_secs(300); +/// Default TTL — 30 minutes, matching Python +/// (`CCRConfig.store_ttl_seconds`). Session-scale: agentic sessions +/// routinely outlive the old 5-minute default, and an expired entry +/// silently converts "lossless with retrieval" into "lossy". +pub const DEFAULT_TTL: Duration = Duration::from_secs(1800); /// Compute the canonical CCR key for `payload`. BLAKE3 → first 24 hex /// chars (96 bits — collision-resistant for the bounded LRU population diff --git a/crates/headroom-core/src/transforms/search_compressor.rs b/crates/headroom-core/src/transforms/search_compressor.rs index a5dd02f78..e6f79d93d 100644 --- a/crates/headroom-core/src/transforms/search_compressor.rs +++ b/crates/headroom-core/src/transforms/search_compressor.rs @@ -140,6 +140,14 @@ pub struct SearchCompressorConfig { /// 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, + /// Group output by file (`rg --heading` style): emit each file path + /// once as a header line, then `line:content` rows beneath it, with + /// a blank line between file groups. Eliminates per-match path + /// repetition — the dominant remaining token waste on large result + /// sets (a 70-char path repeated 15× is ~250 wasted tokens). + /// Default `false` (classic `file:line:content`) for parity; the + /// proxy enables it in token mode. + pub group_by_file: bool, } impl Default for SearchCompressorConfig { @@ -155,6 +163,7 @@ impl Default for SearchCompressorConfig { enable_ccr: true, min_matches_for_ccr: 10, min_compression_ratio_for_ccr: 0.8, + group_by_file: false, } } } @@ -527,15 +536,31 @@ impl SearchCompressor { ) -> (String, BTreeMap) { let mut lines: Vec = Vec::new(); let mut summaries: BTreeMap = BTreeMap::new(); + let grouped = self.config.group_by_file; for (file, fm) in selected { - for m in &fm.matches { - lines.push(format!("{}:{}:{}", m.file, m.line_number, m.content)); + if grouped { + // `rg --heading` style: path once, then line:content rows. + if !lines.is_empty() { + lines.push(String::new()); + } + lines.push(file.clone()); + for m in &fm.matches { + lines.push(format!("{}:{}", m.line_number, m.content)); + } + } else { + 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); + let summary = if grouped { + format!("[... and {} more matches]", omitted) + } else { + format!("[... and {} more matches in {}]", omitted, file) + }; lines.push(summary.clone()); summaries.insert(file.clone(), summary); } diff --git a/crates/headroom-core/src/transforms/smart_crusher/compaction/mod.rs b/crates/headroom-core/src/transforms/smart_crusher/compaction/mod.rs index 891975aea..b3ce01e8c 100644 --- a/crates/headroom-core/src/transforms/smart_crusher/compaction/mod.rs +++ b/crates/headroom-core/src/transforms/smart_crusher/compaction/mod.rs @@ -59,6 +59,16 @@ impl CompactionStage { } } + /// CSV+schema formatter with an explicit config. Used by + /// `SmartCrusher::new` to honor the compaction heuristics carried + /// on `SmartCrusherConfig` instead of pinning `CompactConfig::default()`. + pub fn csv_schema(config: CompactConfig) -> Self { + Self { + config, + formatter: Box::new(CsvSchemaFormatter::new()), + } + } + /// JSON formatter, default config — useful for debugging or for /// downstream consumers that want structured rather than CSV-shaped /// output. diff --git a/crates/headroom-core/src/transforms/smart_crusher/config.rs b/crates/headroom-core/src/transforms/smart_crusher/config.rs index 0e0b7dfd5..8096b7624 100644 --- a/crates/headroom-core/src/transforms/smart_crusher/config.rs +++ b/crates/headroom-core/src/transforms/smart_crusher/config.rs @@ -53,7 +53,10 @@ pub struct SmartCrusherConfig { /// path to be chosen over lossy. Computed as /// `1 - len(rendered) / len(input)`. If lossless saves less than /// this fraction, `crush_array` falls through to the lossy path - /// (with CCR-Dropped retrieval markers). Default `0.30`. + /// (with CCR-Dropped retrieval markers). Default `0.15` — kept in + /// lockstep with the Python `SmartCrusherConfig` dataclass default + /// (lowered from 0.30 so cleanly tabular input takes the lossless + /// path more often; lossless needs no CCR retrieval round-trip). /// /// **Override semantics.** OSS users can tune this via the config /// directly. Enterprise plug-ins replace the entire decision via @@ -77,6 +80,27 @@ pub struct SmartCrusherConfig { /// still emit always; they have no Python equivalent and no /// production caller has asked for them to be suppressed. pub enable_ccr_marker: bool, + /// Compaction heuristic: a field is "core" if it appears in at + /// least this fraction of rows. Mirrors + /// `CompactConfig::core_field_fraction`. Default 0.8. + pub compaction_core_field_fraction: f64, + /// Compaction heuristic: when fewer than this fraction of all + /// observed keys are core, treat the array as heterogeneous and + /// look for a discriminator. Mirrors + /// `CompactConfig::heterogeneous_core_ratio`. Default 0.6. + pub compaction_heterogeneous_core_ratio: f64, + /// Compaction heuristic: cap on inner-key count for + /// nested-uniform flattening. Mirrors + /// `CompactConfig::max_flatten_inner_keys`. Default 6. + pub compaction_max_flatten_inner_keys: usize, + /// Compaction heuristic: minimum bucket count before a candidate + /// discriminator is "useful". Mirrors `CompactConfig::min_buckets`. + /// Default 2. + pub compaction_min_buckets: usize, + /// Compaction heuristic: maximum bucket count — too many buckets + /// means the discriminator is too granular (e.g. an ID column). + /// Mirrors `CompactConfig::max_buckets`. Default 8. + pub compaction_max_buckets: usize, } impl Default for SmartCrusherConfig { @@ -101,8 +125,13 @@ impl Default for SmartCrusherConfig { first_fraction: 0.3, last_fraction: 0.15, relevance_threshold: 0.3, - lossless_min_savings_ratio: 0.30, + lossless_min_savings_ratio: 0.15, enable_ccr_marker: true, + compaction_core_field_fraction: 0.8, + compaction_heterogeneous_core_ratio: 0.6, + compaction_max_flatten_inner_keys: 6, + compaction_min_buckets: 2, + compaction_max_buckets: 8, } } } @@ -133,7 +162,12 @@ mod tests { assert_eq!(c.first_fraction, 0.3); assert_eq!(c.last_fraction, 0.15); assert_eq!(c.relevance_threshold, 0.3); - assert_eq!(c.lossless_min_savings_ratio, 0.30); + assert_eq!(c.lossless_min_savings_ratio, 0.15); assert!(c.enable_ccr_marker); + assert_eq!(c.compaction_core_field_fraction, 0.8); + assert_eq!(c.compaction_heterogeneous_core_ratio, 0.6); + assert_eq!(c.compaction_max_flatten_inner_keys, 6); + assert_eq!(c.compaction_min_buckets, 2); + assert_eq!(c.compaction_max_buckets, 8); } } diff --git a/crates/headroom-core/src/transforms/smart_crusher/crusher.rs b/crates/headroom-core/src/transforms/smart_crusher/crusher.rs index 9ef951f54..456271b67 100644 --- a/crates/headroom-core/src/transforms/smart_crusher/crusher.rs +++ b/crates/headroom-core/src/transforms/smart_crusher/crusher.rs @@ -42,7 +42,7 @@ use super::builder::SmartCrusherBuilder; use super::classifier::{classify_array, ArrayType}; use super::compaction::{ classify_cell, emit_opaque_ccr_marker, try_parse_json_container, CellClass, ClassifyConfig, - Compaction, CompactionStage, + CompactConfig, Compaction, CompactionStage, }; use super::config::SmartCrusherConfig; use super::crushers::{compute_k_split, crush_number_array, crush_object, crush_string_array}; @@ -154,9 +154,20 @@ impl SmartCrusher { /// CCR cache, not to nowhere — same semantics as Python's /// SmartCrusher with CCR enabled. pub fn new(config: SmartCrusherConfig) -> Self { + // Carry the compaction heuristics from the crusher config into + // the compaction stage; everything not exposed on + // SmartCrusherConfig keeps its CompactConfig default. + let compact_cfg = CompactConfig { + core_field_fraction: config.compaction_core_field_fraction, + heterogeneous_core_ratio: config.compaction_heterogeneous_core_ratio, + max_flatten_inner_keys: config.compaction_max_flatten_inner_keys, + min_buckets: config.compaction_min_buckets, + max_buckets: config.compaction_max_buckets, + ..CompactConfig::default() + }; SmartCrusherBuilder::new(config) .with_default_oss_setup() - .with_default_compaction() + .with_compaction(CompactionStage::csv_schema(compact_cfg)) .with_default_ccr_store() .build() } @@ -263,16 +274,49 @@ impl SmartCrusher { /// kept-items list in original-array order. Mirrors Python's /// `_execute_plan` (line 3617-3633). /// - /// Schema-preserving: each kept item is cloned unchanged. No - /// summary objects, generated fields, or wrapper metadata. + /// Schema-preserving by default: each kept item is cloned unchanged. + /// No summary objects, generated fields, or wrapper metadata. + /// + /// When `factor_out_constants` is enabled (default off), fields the + /// analyzer found constant across ALL items are stripped from each + /// kept object and emitted once in a leading + /// `{"_constant_fields": {...}}` sentinel — same output-shape + /// convention as the `_ccr_dropped` sentinel. Stripping is + /// defensive: a key is only removed from an item when its value + /// equals the recorded constant, so a drifted item keeps its own + /// value. The CCR store always holds the full unfactored original. pub fn execute_plan(&self, plan: &CompressionPlan, items: &[Value]) -> Vec { let mut indices = plan.keep_indices.clone(); indices.sort_unstable(); - indices + let mut kept: Vec = indices .into_iter() .filter(|&idx| idx < items.len()) .map(|idx| items[idx].clone()) - .collect() + .collect(); + + if self.config.factor_out_constants && !plan.constant_fields.is_empty() && kept.len() >= 2 { + let mut any_stripped = false; + for item in kept.iter_mut() { + if let Value::Object(map) = item { + for (key, constant) in &plan.constant_fields { + if map.get(key) == Some(constant) { + map.remove(key); + any_stripped = true; + } + } + } + } + if any_stripped { + let mut sentinel = serde_json::Map::new(); + sentinel.insert( + "_constant_fields".to_string(), + Value::Object(plan.constant_fields.clone().into_iter().collect()), + ); + kept.insert(0, Value::Object(sentinel)); + } + } + + kept } /// Top-level entry point. Mirrors Python `SmartCrusher.crush` @@ -1023,6 +1067,80 @@ mod tests { assert_eq!(result.len(), 2); } + #[test] + fn execute_plan_factors_constants_when_enabled() { + let cfg = SmartCrusherConfig { + factor_out_constants: true, + ..Default::default() + }; + let c = SmartCrusher::new(cfg); + let items: Vec = (0..4) + .map(|i| json!({"id": i, "region": "us-west-2", "status": "ok"})) + .collect(); + let mut constant_fields = std::collections::BTreeMap::new(); + constant_fields.insert("region".to_string(), json!("us-west-2")); + constant_fields.insert("status".to_string(), json!("ok")); + let plan = CompressionPlan { + keep_indices: vec![0, 1, 2], + constant_fields, + ..CompressionPlan::default() + }; + let result = c.execute_plan(&plan, &items); + // Sentinel first, then 3 slim items. + assert_eq!(result.len(), 4); + assert_eq!(result[0]["_constant_fields"]["region"], "us-west-2"); + assert_eq!(result[0]["_constant_fields"]["status"], "ok"); + for item in &result[1..] { + assert!(item.get("region").is_none()); + assert!(item.get("status").is_none()); + assert!(item.get("id").is_some()); + } + } + + #[test] + fn execute_plan_keeps_drifted_values_when_factoring() { + // Defensive strip: an item whose value differs from the recorded + // constant keeps its own value. + let cfg = SmartCrusherConfig { + factor_out_constants: true, + ..Default::default() + }; + let c = SmartCrusher::new(cfg); + let items = vec![ + json!({"id": 0, "status": "ok"}), + json!({"id": 1, "status": "FAILED"}), + ]; + let mut constant_fields = std::collections::BTreeMap::new(); + constant_fields.insert("status".to_string(), json!("ok")); + let plan = CompressionPlan { + keep_indices: vec![0, 1], + constant_fields, + ..CompressionPlan::default() + }; + let result = c.execute_plan(&plan, &items); + assert_eq!(result.len(), 3); + assert!(result[1].get("status").is_none()); // matched → stripped + assert_eq!(result[2]["status"], "FAILED"); // drifted → kept + } + + #[test] + fn execute_plan_default_off_leaves_items_unchanged() { + // factor_out_constants defaults to false: schema preserved even + // when the plan carries constant_fields. + let c = crusher(); + let items: Vec = (0..3).map(|i| json!({"id": i, "k": "v"})).collect(); + let mut constant_fields = std::collections::BTreeMap::new(); + constant_fields.insert("k".to_string(), json!("v")); + let plan = CompressionPlan { + keep_indices: vec![0, 1, 2], + constant_fields, + ..CompressionPlan::default() + }; + let result = c.execute_plan(&plan, &items); + assert_eq!(result.len(), 3); + assert_eq!(result[0]["k"], "v"); + } + // ---------- crush_array ---------- #[test] diff --git a/crates/headroom-parity/src/lib.rs b/crates/headroom-parity/src/lib.rs index bf0b33e97..37310e6ff 100644 --- a/crates/headroom-parity/src/lib.rs +++ b/crates/headroom-parity/src/lib.rs @@ -404,6 +404,10 @@ impl TransformComparator for SmartCrusherComparator { .get("enable_ccr_marker") .and_then(|v| v.as_bool()) .unwrap_or(defaults.enable_ccr_marker), + // Compaction heuristics are moot here: this comparator uses + // `without_compaction` (fixtures were recorded against the + // lossy-only path). Take the defaults wholesale. + ..defaults }; // Use without_compaction so the legacy fixtures (recorded diff --git a/crates/headroom-py/src/lib.rs b/crates/headroom-py/src/lib.rs index 99f08d62b..249356d86 100644 --- a/crates/headroom-py/src/lib.rs +++ b/crates/headroom-py/src/lib.rs @@ -469,8 +469,13 @@ impl PySmartCrusherConfig { first_fraction = 0.3, last_fraction = 0.15, relevance_threshold = 0.3, - lossless_min_savings_ratio = 0.30, + lossless_min_savings_ratio = 0.15, enable_ccr_marker = true, + compaction_core_field_fraction = 0.8, + compaction_heterogeneous_core_ratio = 0.6, + compaction_max_flatten_inner_keys = 6, + compaction_min_buckets = 2, + compaction_max_buckets = 8, ))] #[allow(clippy::too_many_arguments)] fn new( @@ -492,6 +497,11 @@ impl PySmartCrusherConfig { relevance_threshold: f64, lossless_min_savings_ratio: f64, enable_ccr_marker: bool, + compaction_core_field_fraction: f64, + compaction_heterogeneous_core_ratio: f64, + compaction_max_flatten_inner_keys: usize, + compaction_min_buckets: usize, + compaction_max_buckets: usize, ) -> Self { Self { inner: RustSmartCrusherConfig { @@ -513,6 +523,11 @@ impl PySmartCrusherConfig { relevance_threshold, lossless_min_savings_ratio, enable_ccr_marker, + compaction_core_field_fraction, + compaction_heterogeneous_core_ratio, + compaction_max_flatten_inner_keys, + compaction_min_buckets, + compaction_max_buckets, }, } } @@ -585,6 +600,30 @@ impl PySmartCrusherConfig { fn enable_ccr_marker(&self) -> bool { self.inner.enable_ccr_marker } + #[getter] + fn lossless_min_savings_ratio(&self) -> f64 { + self.inner.lossless_min_savings_ratio + } + #[getter] + fn compaction_core_field_fraction(&self) -> f64 { + self.inner.compaction_core_field_fraction + } + #[getter] + fn compaction_heterogeneous_core_ratio(&self) -> f64 { + self.inner.compaction_heterogeneous_core_ratio + } + #[getter] + fn compaction_max_flatten_inner_keys(&self) -> usize { + self.inner.compaction_max_flatten_inner_keys + } + #[getter] + fn compaction_min_buckets(&self) -> usize { + self.inner.compaction_min_buckets + } + #[getter] + fn compaction_max_buckets(&self) -> usize { + self.inner.compaction_max_buckets + } fn __repr__(&self) -> String { format!( @@ -1080,6 +1119,7 @@ impl PySearchCompressorConfig { enable_ccr = true, min_matches_for_ccr = 10, min_compression_ratio_for_ccr = 0.8, + group_by_file = false, ))] #[allow(clippy::too_many_arguments)] fn new( @@ -1093,6 +1133,7 @@ impl PySearchCompressorConfig { enable_ccr: bool, min_matches_for_ccr: usize, min_compression_ratio_for_ccr: f64, + group_by_file: bool, ) -> Self { Self { inner: RustSearchConfig { @@ -1106,6 +1147,7 @@ impl PySearchCompressorConfig { enable_ccr, min_matches_for_ccr, min_compression_ratio_for_ccr, + group_by_file, }, } } diff --git a/headroom/audit/__init__.py b/headroom/audit/__init__.py new file mode 100644 index 000000000..2ef7578cb --- /dev/null +++ b/headroom/audit/__init__.py @@ -0,0 +1,17 @@ +"""Offline traffic audits — measure opportunity sizes before tuning defaults.""" + +from .codex import CodexAuditReport, audit_codex, render_codex_text +from .maturation import MaturationSimReport, render_sim_text, simulate_maturation +from .reads import ReadAuditReport, audit_reads, render_text + +__all__ = [ + "CodexAuditReport", + "MaturationSimReport", + "ReadAuditReport", + "audit_codex", + "audit_reads", + "render_codex_text", + "render_sim_text", + "render_text", + "simulate_maturation", +] diff --git a/headroom/audit/codex.py b/headroom/audit/codex.py new file mode 100644 index 000000000..d887d5376 --- /dev/null +++ b/headroom/audit/codex.py @@ -0,0 +1,217 @@ +"""Codex transcript audit — read-pattern analysis for shell-based clients. + +Codex has no structured Read tool: it reads files through shell commands +(``cat``, ``sed -n 'a,bp'``, ``head``/``tail``, ``nl``) — frequently +wrapped by rtk (``rtk read ``, ``rtk proxy ``). This module +classifies ``exec_command`` calls in Codex session transcripts +(``~/.codex/sessions/**/*.jsonl``) and measures the read pattern so the +read-maturation mechanism can be sized for Codex workloads. + +Findings on the development corpus (2026-06-10, 144 sessions, 50MB of +tool output): reads are 51.9% of output bytes, 66% of reads are partial +slices, 55% of reads target an already-read path, and hot files are read +hundreds of times per corpus — the "slice grinder" profile. 78% of read +outputs clear the 2KB maturation floor. +""" + +from __future__ import annotations + +import json +import re +import shlex +from collections import Counter +from dataclasses import asdict, dataclass, field +from pathlib import Path + +# Programs whose output is file content. "read" is rtk's read command. +_READ_PROGS = frozenset({"cat", "sed", "head", "tail", "nl", "bat", "more", "read"}) +_SEARCH_PROGS = frozenset({"rg", "grep", "ugrep", "ag", "fd", "find"}) +_BUILD_PROGS = frozenset({"python", "python3", "pytest", "cargo", "npm", "make", "uv", "ruff"}) +_RANGE_RE = re.compile(r"^\d+([,:-]\d+)?p?$") + +MATURE_FLOOR = 2048 # ReadMaturationConfig.min_size_bytes + + +@dataclass +class CodexAuditReport: + """Aggregated Codex read-pattern results.""" + + sessions: int = 0 + exec_calls: int = 0 + calls_by_category: dict[str, int] = field(default_factory=dict) + bytes_by_category: dict[str, int] = field(default_factory=dict) + total_output_bytes: int = 0 + read_calls: int = 0 + read_bytes: int = 0 + reads_partial: int = 0 + rereads_same_path: int = 0 + distinct_files_read: int = 0 + reads_over_floor: int = 0 + read_size_p50: int = 0 + read_size_p90: int = 0 + top_reread_files: list[tuple[str, int]] = field(default_factory=list) + + def to_dict(self) -> dict: + return asdict(self) + + +def strip_wrappers(cmd: str) -> str: + """Peel rtk wrappers: ``rtk `` and ``rtk proxy ``.""" + c = cmd.strip() + while True: + if c.startswith("rtk "): + c = c[4:].strip() + continue + if c.startswith("proxy "): + c = c[6:].strip() + continue + return c + + +def classify_command(cmd: str, workdir: str = "") -> tuple[str, str | None, bool]: + """Classify a shell command: (category, file_path|None, is_partial). + + Categories: read, search, git, edit, build/test, compound, other. + For reads, the path is resolved against ``workdir`` when relative. + """ + c = strip_wrappers(cmd) + if "apply_patch" in c: + return "edit", None, False + try: + toks = shlex.split(c) + except ValueError: + toks = c.split() + if not toks: + return "other", None, False + prog = toks[0].rsplit("/", 1)[-1] + + if prog in _READ_PROGS: + candidates = [ + t + for t in toks[1:] + if not t.startswith("-") and ("/" in t or "." in t.rsplit("/", 1)[-1]) + ] + # Range tokens like 1,200p (sed) are not paths. + candidates = [t for t in candidates if not _RANGE_RE.match(t.strip("'\""))] + fpath = candidates[0] if candidates else None + if fpath and workdir and not fpath.startswith("/"): + fpath = f"{workdir.rstrip('/')}/{fpath}" + partial = ( + prog in ("sed", "head", "tail") + or any(_RANGE_RE.match(t.strip("'\"")) for t in toks[1:]) + or "--lines" in c + ) + return "read", fpath, partial + if prog in _SEARCH_PROGS: + return "search", None, False + if prog == "git": + return "git", None, False + if prog in _BUILD_PROGS: + return "build/test", None, False + if "&&" in cmd or "|" in cmd: + for part in re.split(r"&&|\|", cmd): + cat, fpath, partial = classify_command(part, workdir) + if cat == "read": + return cat, fpath, partial + return "compound", None, False + return "other", None, False + + +def _output_text(payload: dict) -> str: + out = payload.get("output", "") + if isinstance(out, dict): + out = out.get("output", "") or str(out) + return str(out) + + +def audit_codex(root: Path) -> CodexAuditReport: + """Audit all Codex ``*.jsonl`` transcripts under ``root``.""" + r = CodexAuditReport() + calls: Counter[str] = Counter() + cat_bytes: Counter[str] = Counter() + read_sizes: list[int] = [] + per_file_reads: Counter[str] = Counter() + + for path in sorted(root.glob("**/*.jsonl")): + pending: dict[str, str] = {} + seen_paths: set[str] = set() + saw_lines = False + try: + with path.open(errors="replace") as f: + for raw in f: + try: + line = json.loads(raw) + except json.JSONDecodeError: + continue + saw_lines = True + pl = line.get("payload") or {} + t = pl.get("type") + if t == "function_call" and pl.get("name") == "exec_command": + try: + args = json.loads(pl.get("arguments", "{}")) + except (json.JSONDecodeError, TypeError): + args = {} + cat, fpath, partial = classify_command( + args.get("cmd", ""), args.get("workdir", "") + ) + r.exec_calls += 1 + calls[cat] += 1 + pending[pl.get("call_id", "")] = cat + if cat == "read": + r.read_calls += 1 + r.reads_partial += partial + if fpath: + if fpath in seen_paths: + r.rereads_same_path += 1 + seen_paths.add(fpath) + per_file_reads[fpath] += 1 + elif t == "function_call_output": + size = len(_output_text(pl).encode("utf-8", errors="replace")) + r.total_output_bytes += size + cat = pending.get(pl.get("call_id", ""), "untracked") + cat_bytes[cat] += size + if cat == "read": + r.read_bytes += size + read_sizes.append(size) + if size >= MATURE_FLOOR: + r.reads_over_floor += 1 + except OSError: + continue + if saw_lines: + r.sessions += 1 + + r.calls_by_category = dict(calls.most_common()) + r.bytes_by_category = dict(cat_bytes.most_common()) + r.distinct_files_read = len(per_file_reads) + if read_sizes: + rs = sorted(read_sizes) + r.read_size_p50 = rs[len(rs) // 2] + r.read_size_p90 = rs[int(len(rs) * 0.9)] + r.top_reread_files = [ + (f.rsplit("/", 1)[-1], n) for f, n in per_file_reads.most_common(5) if n > 1 + ] + return r + + +def render_codex_text(r: CodexAuditReport) -> str: + """Human-readable Codex audit summary.""" + total = r.total_output_bytes or 1 + out: list[str] = [] + out.append("── codex read-pattern audit ──") + out.append(f" sessions: {r.sessions}, exec_command calls: {r.exec_calls}") + out.append(f" output bytes by category ({total / 1e6:.1f}MB total):") + for cat, b in r.bytes_by_category.items(): + out.append(f" {cat:<12} {b / 1e6:>6.2f}MB {100 * b / total:.1f}%") + rc = r.read_calls or 1 + out.append( + f" reads: {r.read_calls} ({100 * r.reads_partial / rc:.0f}% partial slices); " + f"re-reads of same path: {r.rereads_same_path} ({100 * r.rereads_same_path / rc:.0f}%)" + ) + out.append( + f" distinct files read: {r.distinct_files_read}; read size p50={r.read_size_p50}B " + f"p90={r.read_size_p90}B; ≥{MATURE_FLOOR}B: {r.reads_over_floor} " + f"({100 * r.reads_over_floor / rc:.0f}%)" + ) + if r.top_reread_files: + out.append(f" most re-read files: {r.top_reread_files}") + return "\n".join(out) diff --git a/headroom/audit/maturation.py b/headroom/audit/maturation.py new file mode 100644 index 000000000..ccf4be20d --- /dev/null +++ b/headroom/audit/maturation.py @@ -0,0 +1,201 @@ +"""Simulate read maturation (Mechanism B) against local transcripts. + +Answers, from real traffic, the questions that size Mechanism B's risk +and tune its `quiesce_turns` policy: + +- How often is the same file re-read at all (and how often as a partial + range)? Partial re-reads happening *despite* full content in context + are evidence the model's natural recovery is already "go back to disk". +- What share of maturation-eligible reads is never touched again + (pure savings, zero recovery events)? +- For files that are touched again, how long until the next touch + (the quiesce-window coverage table)? +- Activity-based at-risk edits: edits landing on a file that had been + quiet longer than N turns — the moments a matured read would force a + re-read under the activity policy. + +Findings on the development corpus (2026-06-10, 81 sessions): 35.5% of +reads are re-reads (95% of those partial); 60.7% of big reads are never +touched again; next-touch p50 is 4 turns — hence quiesce_turns=5. +""" + +from __future__ import annotations + +import json +from collections import defaultdict +from dataclasses import asdict, dataclass, field +from pathlib import Path + +MATURE_FLOOR = 2048 # ReadMaturationConfig.min_size_bytes +QUIESCE_CANDIDATES = [1, 2, 3, 5, 10, 25] +_MUTATING = ("Edit", "Write", "MultiEdit", "NotebookEdit") + + +@dataclass +class MaturationSimReport: + """Aggregated simulation results.""" + + read_calls: int = 0 + rereads_any: int = 0 + rereads_partial: int = 0 + big_reads: int = 0 + big_read_bytes: int = 0 + never_touched_again: int = 0 + next_touch_p50: int = 0 + next_touch_p90: int = 0 + next_touch_p95: int = 0 + # quiesce N -> % of touched-again reads whose next touch is within N + next_touch_within: dict[int, float] = field(default_factory=dict) + edits_with_prior_read: int = 0 + edits_without_prior_read: int = 0 + # quiesce N -> edits whose file was quiet > N turns when edited + # (the matured-read moments under the activity policy) + at_risk_edits: dict[int, int] = field(default_factory=dict) + + def to_dict(self) -> dict: + return asdict(self) + + +def _block_text(content: object) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + return "".join( + b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text" + ) + return "" + + +def simulate_maturation(root: Path) -> MaturationSimReport: + """Run the maturation simulation over ``root/**/*.jsonl``.""" + r = MaturationSimReport() + next_touch_gaps: list[int] = [] + prev_touch_gaps: list[int] = [] # edit -> previous touch of same file + at_risk = dict.fromkeys(QUIESCE_CANDIDATES, 0) + + for path in sorted(root.glob("**/*.jsonl")): + tool_meta: dict[str, tuple[str, dict]] = {} + timeline: dict[str, list[tuple[int, str, int]]] = defaultdict(list) + session_reads: list[tuple[str, int, int]] = [] + seen_files: set[str] = set() + a_idx = 0 + try: + with path.open(errors="replace") as f: + for raw in f: + try: + line = json.loads(raw) + except json.JSONDecodeError: + continue + msg = line.get("message") or {} + role, content = msg.get("role"), msg.get("content") + if role == "assistant" and isinstance(content, list): + a_idx += 1 + for b in content: + if isinstance(b, dict) and b.get("type") == "tool_use": + name = b.get("name", "") + inp = b.get("input") or {} + tool_meta[b.get("id", "")] = (name, inp) + fp = inp.get("file_path") or inp.get("path") or "" + if name in _MUTATING and fp: + timeline[fp].append((a_idx, "edit", 0)) + if role == "user" and isinstance(content, list): + for b in content: + if not (isinstance(b, dict) and b.get("type") == "tool_result"): + continue + name, inp = tool_meta.get(b.get("tool_use_id", ""), ("", {})) + if name != "Read": + continue + fp = inp.get("file_path") or inp.get("path") or "" + if not fp: + continue + text = _block_text(b.get("content")) + size = len(text.encode("utf-8", errors="replace")) + partial = inp.get("offset") is not None or inp.get("limit") is not None + r.read_calls += 1 + if fp in seen_files: + r.rereads_any += 1 + if partial: + r.rereads_partial += 1 + seen_files.add(fp) + timeline[fp].append((a_idx, "read", size)) + session_reads.append((fp, a_idx, size)) + except OSError: + continue + + for ops in timeline.values(): + ops.sort(key=lambda t: t[0]) + for turn, kind, _size in ops: + if kind != "edit": + continue + prev = [t for t, _, _ in ops if t < turn] + had_read = any(k == "read" and t <= turn for t, k, _ in ops) + if not had_read: + r.edits_without_prior_read += 1 + continue + r.edits_with_prior_read += 1 + if prev: + gap = turn - max(prev) + prev_touch_gaps.append(gap) + for n in QUIESCE_CANDIDATES: + if gap > n: + at_risk[n] += 1 + + for fp, rturn, size in session_reads: + if size < MATURE_FLOOR: + continue + r.big_reads += 1 + r.big_read_bytes += size + later = [t for t, _, _ in timeline[fp] if t > rturn] + if later: + next_touch_gaps.append(min(later) - rturn) + else: + r.never_touched_again += 1 + + def pct(xs: list[int], p: float) -> int: + return sorted(xs)[int(len(xs) * p)] if xs else 0 + + r.next_touch_p50 = pct(next_touch_gaps, 0.5) + r.next_touch_p90 = pct(next_touch_gaps, 0.9) + r.next_touch_p95 = pct(next_touch_gaps, 0.95) + if next_touch_gaps: + r.next_touch_within = { + n: round(100 * sum(1 for g in next_touch_gaps if g <= n) / len(next_touch_gaps), 1) + for n in QUIESCE_CANDIDATES + } + r.at_risk_edits = at_risk + return r + + +def render_sim_text(r: MaturationSimReport) -> str: + """Human-readable simulation summary.""" + out: list[str] = [] + out.append("── maturation simulation (Mechanism B) ──") + out.append( + f" re-reads: {r.rereads_any}/{r.read_calls} reads target an already-read file " + f"({100 * r.rereads_any / max(r.read_calls, 1):.1f}%); " + f"{r.rereads_partial} of those are partial ranges" + ) + out.append( + f" big reads (≥{MATURE_FLOOR}B): {r.big_reads} " + f"({r.big_read_bytes / 1e6:.1f}MB); never touched again: " + f"{r.never_touched_again} ({100 * r.never_touched_again / max(r.big_reads, 1):.1f}%) " + f"← pure savings" + ) + out.append( + f" next-touch gap for the rest (turns): p50={r.next_touch_p50} " + f"p90={r.next_touch_p90} p95={r.next_touch_p95}" + ) + if r.next_touch_within: + for n, share in r.next_touch_within.items(): + out.append(f" next touch within {n:>2} turn(s): {share:.1f}%") + total_edits = r.edits_with_prior_read + r.edits_without_prior_read + out.append( + f" edits: {r.edits_with_prior_read} with a prior read of the file, " + f"{r.edits_without_prior_read} without" + ) + out.append(" activity-based at-risk edits (file quiet > N turns when edited):") + for n, count in r.at_risk_edits.items(): + out.append( + f" quiesce {n:>2}: {count:>5} edits ({100 * count / max(total_edits, 1):.1f}%)" + ) + return "\n".join(out) diff --git a/headroom/audit/reads.py b/headroom/audit/reads.py new file mode 100644 index 000000000..932401d0a --- /dev/null +++ b/headroom/audit/reads.py @@ -0,0 +1,344 @@ +"""Read-opportunity audit over local Claude Code transcripts. + +Measures, from REAL session data, the addressable bytes for each Read +compression mechanism — so defaults are set from traffic, not theory. +Run it on a deployment's transcripts before tuning anything: + + headroom audit-reads + headroom audit-reads --path /path/to/projects --format json + +Read-only: streams ``/**/*.jsonl`` (Claude Code session transcripts) +and never modifies anything. + +What it sizes, per mechanism: + +- **identical repeat** — a later Read byte-identical to an earlier Read of + the same file. (A dedup mechanism for this was prototyped and removed: + it measured 0.1% of Read bytes on real traffic. If this number is + material on YOUR traffic, the implementation lives in git history — + see the feat/compression-extraction branch.) +- **subset containment** — a later partial Read contained in an earlier + full Read of the same file. +- **write-readback** — a Read whose content echoes a prior Write input. +- **stale** — Reads of files later edited (read_lifecycle's stale class; + mostly freeze-blocked in production, unlockable at cache-death). +- **line-number scaffolding** — `cat -n` prefix bytes inside Read output. +- **context residency** — how many assistant turns each Read stays in + context (the multiplier on its prefix-cache read cost; the case for + compress-before-cache-entry). +- **cache-death windows** — inter-message gaps exceeding the provider + cache TTL (free recompression moments). +""" + +from __future__ import annotations + +import hashlib +import json +import re +from collections import defaultdict +from dataclasses import asdict, dataclass, field +from datetime import datetime +from pathlib import Path + +MIN_SIZE = 512 # matches ReadLifecycleConfig.min_size_bytes +_LINE_NUM_RE = re.compile(r"^\s*\d+\t", re.M) + +_LOCK_GENERATED = ( + "package-lock.json", + "yarn.lock", + "pnpm-lock.yaml", + "cargo.lock", + "go.sum", + "poetry.lock", + "uv.lock", + "gemfile.lock", + "composer.lock", +) +_SOURCE_EXT = ( + ".py", + ".ts", + ".tsx", + ".js", + ".jsx", + ".rs", + ".go", + ".java", + ".c", + ".cpp", + ".h", + ".rb", + ".swift", + ".kt", + ".scala", + ".sh", + ".zsh", +) +_DATA_EXT = (".json", ".jsonl", ".csv", ".yaml", ".yml", ".toml", ".xml") +_DOC_EXT = (".md", ".rst", ".txt") + +_MUTATING_TOOLS = ("Edit", "Write", "MultiEdit", "NotebookEdit") + + +@dataclass +class ReadAuditReport: + """Aggregated audit results. All byte figures are UTF-8 bytes of + tool_result content; tokens ≈ bytes / 4.""" + + sessions: int = 0 + files_skipped: int = 0 + tool_bytes: dict[str, int] = field(default_factory=dict) + read_calls: int = 0 + read_bytes: int = 0 + read_calls_small: int = 0 + dedup_identical_calls: int = 0 + dedup_identical_bytes: int = 0 + subset_calls: int = 0 + subset_bytes: int = 0 + write_readback_calls: int = 0 + write_readback_bytes: int = 0 + stale_calls: int = 0 + stale_bytes: int = 0 + linenum_overhead_bytes: int = 0 + class_bytes: dict[str, int] = field(default_factory=dict) + residency_median: int = 0 + residency_p90: int = 0 + residency_mean: float = 0.0 + gaps_over_5m: int = 0 + gaps_over_1h: int = 0 + sessions_with_gap: int = 0 + reads_per_file_max_median: int = 0 + reads_per_file_max: int = 0 + + def to_json(self) -> str: + return json.dumps(asdict(self), indent=2, sort_keys=True) + + +def _classify_path(p: str) -> str: + low = p.lower() + name = low.rsplit("/", 1)[-1] + if name in _LOCK_GENERATED or "/node_modules/" in low or "/dist/" in low or ".min." in name: + return "lock/generated/vendored" + if name.endswith(_SOURCE_EXT): + return "source code" + if name.endswith(_DOC_EXT): + return "docs/text" + if name.endswith(_DATA_EXT): + return "data/config" + return "other" + + +def _block_text(content: object) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + return "".join( + b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text" + ) + return "" + + +def _parse_ts(line: dict) -> float | None: + ts = line.get("timestamp") + if not ts: + return None + try: + return datetime.fromisoformat(str(ts).replace("Z", "+00:00")).timestamp() + except ValueError: + return None + + +class _Agg: + def __init__(self) -> None: + self.report = ReadAuditReport() + self.tool_bytes: dict[str, int] = defaultdict(int) + self.class_bytes: dict[str, int] = defaultdict(int) + self.residency: list[int] = [] + self.reads_per_file_max: list[int] = [] + + +def _audit_session(path: Path, agg: _Agg) -> None: + r = agg.report + tool_meta: dict[str, tuple[str, dict]] = {} + file_reads: dict[str, list[tuple[str, str]]] = defaultdict(list) + file_writes: dict[str, list[str]] = defaultdict(list) + read_events: list[tuple[str, int, int, bool]] = [] # (file, size, at, deduped) + edit_files_at: list[tuple[int, str]] = [] + assistant_idx = 0 + prev_ts: float | None = None + had_gap = False + + with path.open(errors="replace") as f: + for raw in f: + try: + line = json.loads(raw) + except json.JSONDecodeError: + continue + msg = line.get("message") or {} + role = msg.get("role") + content = msg.get("content") + + ts = _parse_ts(line) + if ts is not None and prev_ts is not None: + gap = ts - prev_ts + if gap > 3600: + r.gaps_over_1h += 1 + r.gaps_over_5m += 1 + had_gap = True + elif gap > 300: + r.gaps_over_5m += 1 + had_gap = True + if ts is not None: + prev_ts = ts + + if role == "assistant" and isinstance(content, list): + assistant_idx += 1 + for b in content: + if isinstance(b, dict) and b.get("type") == "tool_use": + name = b.get("name", "") + inp = b.get("input") or {} + tool_meta[b.get("id", "")] = (name, inp) + fp = inp.get("file_path") or inp.get("path") or "" + if name in _MUTATING_TOOLS and fp: + edit_files_at.append((assistant_idx, fp)) + if name == "Write": + file_writes[fp].append(str(inp.get("content", ""))) + + if role == "user" and isinstance(content, list): + for b in content: + if not (isinstance(b, dict) and b.get("type") == "tool_result"): + continue + tid = b.get("tool_use_id", "") + name, inp = tool_meta.get(tid, ("", {})) + text = _block_text(b.get("content")) + size = len(text.encode("utf-8", errors="replace")) + agg.tool_bytes[name or "unknown"] += size + if name != "Read": + continue + + r.read_calls += 1 + r.read_bytes += size + fp = inp.get("file_path") or inp.get("path") or "" + is_partial = inp.get("offset") is not None or inp.get("limit") is not None + if size < MIN_SIZE: + r.read_calls_small += 1 + agg.class_bytes[_classify_path(fp)] += size + r.linenum_overhead_bytes += sum( + len(m.group(0)) for m in _LINE_NUM_RE.finditer(text) + ) + + h = hashlib.sha256(text.encode("utf-8", errors="replace")).hexdigest() + deduped = False + if size >= MIN_SIZE and fp: + prior = file_reads[fp] + if any(ph == h for ph, _ in prior): + r.dedup_identical_bytes += size + r.dedup_identical_calls += 1 + deduped = True + elif ( + is_partial + and text + and any(text in pc for _, pc in prior if len(pc) > len(text)) + ): + r.subset_bytes += size + r.subset_calls += 1 + elif any( + text.strip() and w.strip() and text.strip() in w + for w in file_writes.get(fp, []) + ): + r.write_readback_bytes += size + r.write_readback_calls += 1 + if fp: + file_reads[fp].append((h, text)) + read_events.append((fp, size, assistant_idx, deduped)) + + for fp, size, at, deduped in read_events: + if size >= MIN_SIZE and fp and not deduped: + if any(idx > at and ef == fp for idx, ef in edit_files_at): + r.stale_bytes += size + r.stale_calls += 1 + agg.residency.append(max(0, assistant_idx - at)) + + per_file: dict[str, int] = defaultdict(int) + for fp, _, _, _ in read_events: + if fp: + per_file[fp] += 1 + if per_file: + agg.reads_per_file_max.append(max(per_file.values())) + if had_gap: + r.sessions_with_gap += 1 + r.sessions += 1 + + +def audit_reads(root: Path) -> ReadAuditReport: + """Audit all ``*.jsonl`` transcripts under ``root``.""" + agg = _Agg() + for p in sorted(root.glob("**/*.jsonl")): + try: + _audit_session(p, agg) + except OSError: + agg.report.files_skipped += 1 + + r = agg.report + r.tool_bytes = dict(sorted(agg.tool_bytes.items(), key=lambda kv: -kv[1])) + r.class_bytes = dict(sorted(agg.class_bytes.items(), key=lambda kv: -kv[1])) + if agg.residency: + rt = sorted(agg.residency) + r.residency_median = rt[len(rt) // 2] + r.residency_p90 = rt[int(len(rt) * 0.9)] + r.residency_mean = sum(rt) / len(rt) + if agg.reads_per_file_max: + m = sorted(agg.reads_per_file_max) + r.reads_per_file_max_median = m[len(m) // 2] + r.reads_per_file_max = m[-1] + return r + + +def _fmt(b: int) -> str: + if b > 1_000_000: + return f"{b / 1_000_000:.1f}MB (~{b // 4000}K tok)" + return f"{b / 1000:.0f}KB (~{b // 4000}K tok)" + + +def render_text(r: ReadAuditReport) -> str: + """Render the report as the human-readable summary.""" + total_tool = sum(r.tool_bytes.values()) or 1 + rb = r.read_bytes or 1 + out: list[str] = [] + out.append(f"sessions analyzed: {r.sessions}") + if r.files_skipped: + out.append(f"files skipped (unreadable): {r.files_skipped}") + out.append("\n── tool_result bytes by tool ──") + for name, b in list(r.tool_bytes.items())[:10]: + out.append(f" {name or '?':<24} {_fmt(b):<28} {100 * b / total_tool:.1f}%") + out.append("\n── Read opportunity sizing (share of Read bytes) ──") + out.append(f" Read calls: {r.read_calls} ({r.read_calls_small} below {MIN_SIZE}B floor)") + out.append( + f" Read bytes: {_fmt(r.read_bytes)} = {100 * r.read_bytes / total_tool:.1f}% of all tool bytes" + ) + rows = [ + ("identical repeat", r.dedup_identical_calls, r.dedup_identical_bytes), + ("subset containment", r.subset_calls, r.subset_bytes), + ("write-readback", r.write_readback_calls, r.write_readback_bytes), + ("stale (edit after read)", r.stale_calls, r.stale_bytes), + ] + for label, calls, b in rows: + out.append(f" {label:<32} {calls:>5} calls {_fmt(b):<28} {100 * b / rb:.1f}%") + out.append( + f" {'line-number scaffolding':<32} {'':>11} {_fmt(r.linenum_overhead_bytes):<28} " + f"{100 * r.linenum_overhead_bytes / rb:.1f}%" + ) + out.append("\n── Read bytes by file class ──") + for cls, b in r.class_bytes.items(): + out.append(f" {cls:<24} {_fmt(b):<28} {100 * b / rb:.1f}%") + out.append("\n── context residency (assistant turns after each Read) ──") + out.append(f" median {r.residency_median}, p90 {r.residency_p90}, mean {r.residency_mean:.0f}") + out.append("\n── cache-death windows ──") + out.append( + f" gaps >5min: {r.gaps_over_5m} ({r.gaps_over_1h} of them >1h); " + f"sessions with ≥1 gap: {r.sessions_with_gap}/{r.sessions}" + ) + out.append( + f" max reads of one file per session: median {r.reads_per_file_max_median}, " + f"max {r.reads_per_file_max}" + ) + return "\n".join(out) diff --git a/headroom/cache/backends/__init__.py b/headroom/cache/backends/__init__.py index bdfd46b6b..8dd19ac3d 100644 --- a/headroom/cache/backends/__init__.py +++ b/headroom/cache/backends/__init__.py @@ -1,19 +1,25 @@ """Storage backends for CompressionStore. This module provides pluggable storage backends for CCR (Compress-Cache-Retrieve). -The default is in-memory storage, but alternative backends can be implemented for: -- Persistence (MongoDB, Redis, etc.) -- Distributed caching -- Custom storage solutions +Backend selection depends on how the store is constructed: +- ``get_compression_store()`` (the proxy path) defaults to SQLite + (restart-safe, shared across workers); ``HEADROOM_CCR_BACKEND=memory`` + forces in-memory, and other backends (Redis, MongoDB via entry points) + can be selected by env. +- ``CompressionStore()`` constructed directly defaults to **in-memory** + unless a backend is passed explicitly. Usage: - from headroom.cache.backends import InMemoryBackend, CompressionStoreBackend - from headroom.cache.compression_store import CompressionStore + from headroom.cache.backends import SQLiteBackend, CompressionStoreBackend + from headroom.cache.compression_store import CompressionStore, get_compression_store - # Use default in-memory backend - store = CompressionStore() + # Env-driven default (SQLite at ~/.headroom/ccr_store.db) + store = get_compression_store() - # Use custom backend + # Direct construction defaults to in-memory; pass a backend for persistence + store = CompressionStore(backend=SQLiteBackend()) + + # Use a custom backend class MyBackend: # Implement CompressionStoreBackend protocol ... @@ -22,8 +28,10 @@ Usage: from .base import CompressionStoreBackend from .memory import InMemoryBackend +from .sqlite import SQLiteBackend __all__ = [ "CompressionStoreBackend", "InMemoryBackend", + "SQLiteBackend", ] diff --git a/headroom/cache/backends/sqlite.py b/headroom/cache/backends/sqlite.py new file mode 100644 index 000000000..641ae369b --- /dev/null +++ b/headroom/cache/backends/sqlite.py @@ -0,0 +1,275 @@ +"""SQLite storage backend for CompressionStore. + +Default backend for the CCR store. Two properties the in-memory backend +cannot provide, both load-bearing for the no-accuracy-loss guarantee: + +- **Restart survival.** A proxy restart no longer destroys every + retrievable original mid-session. With the session-scale 30-minute + TTL, entries are expected to outlive any single process. +- **Multi-worker sharing.** The database file (WAL mode) is shared + across worker processes, so a `headroom_retrieve` call served by a + different worker than the one that compressed still finds the entry. + This closes the largest of the documented multi-worker gaps. + +Set ``HEADROOM_CCR_BACKEND=memory`` to opt back into the in-memory +backend, or ``HEADROOM_CCR_SQLITE_PATH`` to relocate the database file +(default ``~/.headroom/ccr_store.db``). +""" + +from __future__ import annotations + +import json +import logging +import os +import sqlite3 +import threading +import time +from dataclasses import asdict, fields +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from ..compression_store import CompressionEntry + +logger = logging.getLogger(__name__) + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS ccr_entries ( + hash TEXT PRIMARY KEY, + entry_json TEXT NOT NULL, + created_at REAL NOT NULL, + ttl INTEGER NOT NULL +); +CREATE INDEX IF NOT EXISTS idx_ccr_expiry ON ccr_entries (created_at); +""" + +# Purge expired rows at most this often (seconds). Purging is hygiene, +# not correctness — CompressionStore checks TTL on every get(). +_PURGE_INTERVAL = 60.0 + + +def default_db_path() -> Path: + """Resolve the database path (env override or ~/.headroom/).""" + env = os.environ.get("HEADROOM_CCR_SQLITE_PATH", "").strip() + if env: + return Path(env).expanduser() + return Path.home() / ".headroom" / "ccr_store.db" + + +class SQLiteBackend: + """Thread-safe SQLite storage backend (WAL mode). + + Entries are serialized as one JSON blob per row; ``created_at`` and + ``ttl`` are duplicated into columns so expired rows can be purged + with one DELETE. TTL *enforcement* on reads stays in + CompressionStore, matching the backend protocol contract. + + Deserialization is field-filtered: unknown keys in stored JSON are + dropped (forward-compatible with newer versions that add fields). + Missing keys load cleanly only when the corresponding + ``CompressionEntry`` field has a default; a blob missing a required + field (one without a default) raises ``TypeError`` on construction. + """ + + def __init__(self, db_path: str | Path | None = None) -> None: + self._path = Path(db_path).expanduser() if db_path else default_db_path() + self._path.parent.mkdir(parents=True, exist_ok=True) + self._lock = threading.Lock() + self._last_purge = 0.0 + self._conn = self._open() + + def _open(self) -> sqlite3.Connection: + conn = sqlite3.connect(self._path, check_same_thread=False) + # Wait for competing writers instead of failing with SQLITE_BUSY — + # multiple proxy workers share this file, and writes are frequent + # but tiny, so contention resolves in milliseconds. + conn.execute("PRAGMA busy_timeout=5000") + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA synchronous=NORMAL") + conn.executescript(_SCHEMA) + # Startup hygiene: expired rows are only purged opportunistically + # on writes, so a quiet store could otherwise hold expired + # originals (which may contain sensitive tool output) on disk + # indefinitely. Sweep them on every open. + conn.execute( + "DELETE FROM ccr_entries WHERE created_at + ttl < ?", + (time.time(),), + ) + conn.commit() + # Originals can contain sensitive tool output (file contents, + # command output) — keep the database private to the user. + for suffix in ("", "-wal", "-shm"): + p = Path(str(self._path) + suffix) + if p.exists(): + try: + p.chmod(0o600) + except OSError: + pass + return conn + + @staticmethod + def _is_corruption(error: Exception) -> bool: + """Only genuine file corruption justifies recreating the database. + + ``sqlite3.OperationalError`` (a DatabaseError subclass) also covers + transient conditions like ``database is locked`` under multi-worker + write contention — misclassifying those as corruption would delete + live data while sibling workers still hold handles to the unlinked + inode (split-brain). Match the corruption messages explicitly. + """ + msg = str(error).lower() + return "malformed" in msg or "not a database" in msg + + def _handle_db_error(self, error: sqlite3.DatabaseError, op: str) -> None: + """Corruption → recreate (loud). Anything else (busy/locked/io) → + log and treat the operation as a miss; never destroy data over a + transient error.""" + if not self._is_corruption(error): + logger.warning("CCR SQLite %s failed (transient, no reset): %s", op, error) + return + logger.warning( + "CCR SQLite store at %s is corrupt (%s); recreating. " + "Previously stored originals are lost — affected retrieval " + "markers will miss until their content is re-compressed.", + self._path, + error, + ) + try: + self._conn.close() + except Exception: # noqa: BLE001 - best-effort close on corrupt handle + pass + self._path.unlink(missing_ok=True) + self._conn = self._open() + + def _entry_from_json(self, raw: str) -> CompressionEntry | None: + from ..compression_store import CompressionEntry + + try: + data = json.loads(raw) + except (json.JSONDecodeError, TypeError): + return None + known = {f.name for f in fields(CompressionEntry)} + return CompressionEntry(**{k: v for k, v in data.items() if k in known}) + + def _maybe_purge(self) -> None: + """Delete expired rows; called opportunistically under the lock.""" + now = time.time() + if now - self._last_purge < _PURGE_INTERVAL: + return + self._last_purge = now + self._conn.execute( + "DELETE FROM ccr_entries WHERE created_at + ttl < ?", + (now,), + ) + self._conn.commit() + + def get(self, hash_key: str) -> CompressionEntry | None: + with self._lock: + try: + row = self._conn.execute( + "SELECT entry_json FROM ccr_entries WHERE hash = ?", + (hash_key,), + ).fetchone() + except sqlite3.DatabaseError as e: + self._handle_db_error(e, "get") + return None + if row is None: + return None + return self._entry_from_json(row[0]) + + def set(self, hash_key: str, entry: CompressionEntry) -> None: + payload = json.dumps(asdict(entry), ensure_ascii=False) + with self._lock: + try: + self._conn.execute( + "INSERT OR REPLACE INTO ccr_entries " + "(hash, entry_json, created_at, ttl) VALUES (?, ?, ?, ?)", + (hash_key, payload, entry.created_at, entry.ttl), + ) + self._conn.commit() + self._maybe_purge() + except sqlite3.DatabaseError as e: + self._handle_db_error(e, "set") + + def delete(self, hash_key: str) -> bool: + with self._lock: + try: + cur = self._conn.execute( + "DELETE FROM ccr_entries WHERE hash = ?", + (hash_key,), + ) + self._conn.commit() + return cur.rowcount > 0 + except sqlite3.DatabaseError as e: + self._handle_db_error(e, "op") + return False + + def exists(self, hash_key: str) -> bool: + with self._lock: + try: + row = self._conn.execute( + "SELECT 1 FROM ccr_entries WHERE hash = ?", + (hash_key,), + ).fetchone() + except sqlite3.DatabaseError as e: + self._handle_db_error(e, "op") + return False + return row is not None + + def clear(self) -> None: + with self._lock: + try: + self._conn.execute("DELETE FROM ccr_entries") + self._conn.commit() + except sqlite3.DatabaseError as e: + self._handle_db_error(e, "op") + + def count(self) -> int: + with self._lock: + try: + row = self._conn.execute("SELECT COUNT(*) FROM ccr_entries").fetchone() + except sqlite3.DatabaseError as e: + self._handle_db_error(e, "op") + return 0 + return int(row[0]) + + def keys(self) -> list[str]: + with self._lock: + try: + rows = self._conn.execute("SELECT hash FROM ccr_entries").fetchall() + except sqlite3.DatabaseError as e: + self._handle_db_error(e, "op") + return [] + return [r[0] for r in rows] + + def items(self) -> list[tuple[str, CompressionEntry]]: + with self._lock: + try: + rows = self._conn.execute("SELECT hash, entry_json FROM ccr_entries").fetchall() + except sqlite3.DatabaseError as e: + self._handle_db_error(e, "op") + return [] + out: list[tuple[str, CompressionEntry]] = [] + for hash_key, raw in rows: + entry = self._entry_from_json(raw) + if entry is not None: + out.append((hash_key, entry)) + return out + + def get_stats(self) -> dict[str, Any]: + with self._lock: + try: + count_row = self._conn.execute("SELECT COUNT(*) FROM ccr_entries").fetchone() + except sqlite3.DatabaseError as e: + self._handle_db_error(e, "op") + count_row = (0,) + try: + bytes_used = self._path.stat().st_size + except OSError: + bytes_used = 0 + return { + "backend_type": "sqlite", + "entry_count": int(count_row[0]), + "bytes_used": bytes_used, + "db_path": str(self._path), + } diff --git a/headroom/cache/compression_store.py b/headroom/cache/compression_store.py index b6eb94660..2f2fa0244 100644 --- a/headroom/cache/compression_store.py +++ b/headroom/cache/compression_store.py @@ -53,7 +53,7 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -DEFAULT_CCR_TTL_SECONDS = 300 +DEFAULT_CCR_TTL_SECONDS = 1800 # session-scale; override via HEADROOM_CCR_TTL_SECONDS CCR_TTL_SECONDS_ENV = "HEADROOM_CCR_TTL_SECONDS" _RETRIEVAL_LOG_PREVIEW_CHARS = 4096 @@ -125,6 +125,19 @@ def _payload_for_retrieval_log(payload: str) -> dict[str, Any]: } +# Single source of truth for the retrieval-miss message. Actionable by +# design: the model still has the marker in context (Read markers carry +# the file path), so tell it how to recover instead of just reporting +# the miss. +CCR_MISS_MESSAGE = ( + "Entry not found or expired. To recover: if the compression marker " + "references a file Read, re-read that file (the path is in the " + "marker; disk is the source of truth). If it was command output, " + "re-run the command. Entries expire after the store TTL " + "(default 30 minutes; configurable via HEADROOM_CCR_TTL_SECONDS)." +) + + @dataclass class CompressionEntry: """A cached compression entry with metadata for retrieval and feedback.""" @@ -207,10 +220,13 @@ class CompressionStore: Args: max_entries: Maximum number of entries to store. - default_ttl: Default TTL in seconds. + default_ttl: Default TTL in seconds (default 30 minutes — session scale). enable_feedback: Whether to track retrieval events. - backend: Storage backend to use. Defaults to InMemoryBackend. - Custom backends can be passed for persistence (MongoDB, Redis). + backend: Storage backend to use. Defaults to InMemoryBackend + when constructed directly; `get_compression_store()` + defaults to SQLiteBackend for restart/multi-worker + safety. Custom backends can be passed for + persistence (MongoDB, Redis). """ # Import here to avoid circular imports from .backends import InMemoryBackend @@ -1189,12 +1205,29 @@ def clear_request_compression_store() -> None: def _create_default_ccr_backend() -> CompressionStoreBackend | None: """Create a CCR backend from env (e.g. HEADROOM_CCR_BACKEND=redis). - Loads adapters via setuptools entry point 'headroom.ccr_backend'. - Returns None to use default InMemoryBackend. + Default (env unset or "sqlite"): SQLiteBackend at + ~/.headroom/ccr_store.db — restart-safe and shared across worker + processes, which the session-scale 30-minute TTL assumes. + "memory" opts back into the in-process dict. Other values load + adapters via setuptools entry point 'headroom.ccr_backend'. + Returns None to use InMemoryBackend. """ backend_type = (os.environ.get("HEADROOM_CCR_BACKEND") or "").strip().lower() - if not backend_type or backend_type == "memory": + if backend_type == "memory": return None + if not backend_type or backend_type == "sqlite": + try: + from .backends.sqlite import SQLiteBackend + + return SQLiteBackend() + except Exception as e: + logger.warning( + "Failed to initialize SQLite CCR backend (%s); " + "falling back to in-memory store. Retrieval will not " + "survive proxy restarts.", + e, + ) + return None try: from importlib.metadata import entry_points @@ -1233,7 +1266,7 @@ def get_compression_store( Args: max_entries: Maximum entries (only used on first call for global store). default_ttl: Default TTL (only used on first call for global store). - When omitted, HEADROOM_CCR_TTL_SECONDS overrides the 300-second default. + When omitted, HEADROOM_CCR_TTL_SECONDS overrides the 1800-second default. backend: Custom storage backend (only used on first call for global store). Defaults to InMemoryBackend if not provided; env backend used if backend is None. diff --git a/headroom/ccr/mcp_server.py b/headroom/ccr/mcp_server.py index 74425b020..650e375bc 100644 --- a/headroom/ccr/mcp_server.py +++ b/headroom/ccr/mcp_server.py @@ -451,8 +451,11 @@ class HeadroomMCPServer: return { "error": "Content not found. It may have expired or the hash may be incorrect.", "hash": hash_key, - "hint": "Content compressed via headroom_compress is stored for the session. " - "Content compressed by the proxy uses the configured CCR TTL.", + "hint": "To recover: if the compression marker references a file Read, " + "re-read that file (the path is in the marker; disk is the source of " + "truth). If it was command output, re-run the command. Content " + "compressed via headroom_compress is stored for the session; content " + "compressed by the proxy uses the configured CCR TTL.", } async def _retrieve_via_proxy( diff --git a/headroom/cli/__init__.py b/headroom/cli/__init__.py index 9ac5bd81a..626ea1d3d 100644 --- a/headroom/cli/__init__.py +++ b/headroom/cli/__init__.py @@ -13,6 +13,7 @@ survives that kind of sys.modules mutation. """ from . import ( # noqa: F401 + audit, capture, copilot_auth, evals, diff --git a/headroom/cli/audit.py b/headroom/cli/audit.py new file mode 100644 index 000000000..e1cfa5998 --- /dev/null +++ b/headroom/cli/audit.py @@ -0,0 +1,107 @@ +"""Traffic audit CLI commands.""" + +from pathlib import Path + +import click + +from .main import main + + +@main.command(name="audit-reads") +@click.option( + "--path", + "root", + type=click.Path(exists=True, file_okay=False, path_type=Path), + default=None, + help="Transcript directory to audit (default: ~/.claude/projects)", +) +@click.option( + "--format", + "output_format", + type=click.Choice(["text", "json"]), + default="text", + help="Output format (default: text)", +) +@click.option( + "--simulate-maturation", + is_flag=True, + help="Also simulate Mechanism B (read maturation): re-read rates, " + "never-touched-again share, quiesce-window coverage, at-risk edits.", +) +@click.option( + "--codex", + "codex_mode", + is_flag=True, + help="Audit Codex transcripts instead (shell-based reads: cat/sed/rtk read). " + "Default path becomes ~/.codex/sessions.", +) +def audit_reads_cmd( + root: Path | None, output_format: str, simulate_maturation: bool, codex_mode: bool +) -> None: + """Audit Read-tool traffic for compression opportunities. + + Streams local Claude Code transcripts (read-only) and sizes the + addressable bytes for each Read mechanism: identical repeats, subset + containment, write-readback, stale reads, line-number scaffolding, + context residency, and cache-death windows. + + \b + Run this on a deployment's transcripts BEFORE tuning compression + defaults — opportunity sizes vary heavily by workload. + + \b + Examples: + headroom audit-reads + headroom audit-reads --path /var/transcripts --format json + headroom audit-reads --codex + """ + import json as _json + + from headroom.audit.reads import audit_reads, render_text + + if codex_mode: + if root is None: + root = Path.home() / ".codex" / "sessions" + if not root.exists(): + raise click.ClickException( + f"{root} does not exist — pass --path to the Codex sessions directory" + ) + from headroom.audit.codex import audit_codex, render_codex_text + + codex_report = audit_codex(root) + if codex_report.sessions == 0: + raise click.ClickException(f"no *.jsonl transcripts found under {root}") + if output_format == "json": + click.echo(_json.dumps(codex_report.to_dict(), indent=2, sort_keys=True)) + else: + click.echo(render_codex_text(codex_report)) + return + + if root is None: + root = Path.home() / ".claude" / "projects" + if not root.exists(): + raise click.ClickException( + f"{root} does not exist — pass --path to the transcript directory" + ) + + report = audit_reads(root) + if report.sessions == 0: + raise click.ClickException(f"no *.jsonl transcripts found under {root}") + + sim = None + if simulate_maturation: + from headroom.audit.maturation import render_sim_text + from headroom.audit.maturation import simulate_maturation as run_sim + + sim = run_sim(root) + + if output_format == "json": + data = _json.loads(report.to_json()) + if sim is not None: + data["maturation_simulation"] = sim.to_dict() + click.echo(_json.dumps(data, indent=2, sort_keys=True)) + else: + click.echo(render_text(report)) + if sim is not None: + click.echo() + click.echo(render_sim_text(sim)) diff --git a/headroom/cli/main.py b/headroom/cli/main.py index b726e7018..472028357 100644 --- a/headroom/cli/main.py +++ b/headroom/cli/main.py @@ -37,6 +37,7 @@ def _register_commands() -> None: """Register all subcommand groups.""" from . import ( agent_savings, # noqa: F401 + audit, # noqa: F401 capture, # noqa: F401 copilot_auth, # noqa: F401 evals, # noqa: F401 diff --git a/headroom/config.py b/headroom/config.py index 2dd07f8d2..beb49d3de 100644 --- a/headroom/config.py +++ b/headroom/config.py @@ -356,10 +356,20 @@ class SmartCrusherConfig: first_fraction: float = 0.3 # 30% of K from start of array last_fraction: float = 0.15 # 15% of K from end of array - # Lossless compaction only replaces the original when it saves at - # least this byte fraction vs the (minified) input. Mirrors the - # Rust default. - lossless_min_savings_ratio: float = 0.30 + # Lossless-first dispatch: minimum byte-savings ratio for the lossless + # Table/CSV compaction path to win over the lossy path. Must stay in + # lockstep with the Rust default (smart_crusher config.rs) and the + # transforms-level dataclass. + lossless_min_savings_ratio: float = 0.15 + + # Compaction heuristics (mirror Rust CompactConfig). A field is "core" + # if present in at least this fraction of rows; arrays whose key sets + # are mostly non-core are bucketed by a discriminator instead. + compaction_core_field_fraction: float = 0.8 + compaction_heterogeneous_core_ratio: float = 0.6 + compaction_max_flatten_inner_keys: int = 6 + compaction_min_buckets: int = 2 + compaction_max_buckets: int = 8 @dataclass @@ -410,14 +420,20 @@ class CCRConfig: - Network effect: retrieval patterns improve compression for all users GOTCHAS: - - Cache has TTL (default 300 seconds) - retrieval fails after expiration + - Cache has TTL (default 30 min) - retrieval fails after expiration - Memory usage: ~1KB per cached entry - Only works with array compression (not string truncation) """ enabled: bool = True # Enable CCR (cache + retrieval markers) store_max_entries: int = 1000 # Max entries in compression store - store_ttl_seconds: int = 300 # Cache TTL in seconds + # Session-scale TTL. The original 5-minute default predates agentic + # sessions that routinely run 30+ minutes; an expired entry silently + # converts "lossless with retrieval" into "lossy", so the TTL is the + # weakest link in the no-accuracy-loss guarantee. Kept in lockstep + # with Rust DEFAULT_TTL (crates/headroom-core/src/ccr/mod.rs) and + # DEFAULT_CCR_TTL_SECONDS (cache/compression_store.py). + store_ttl_seconds: int = 1800 # Cache TTL (30 minutes) inject_retrieval_marker: bool = True # Add retrieval hint to compressed output feedback_enabled: bool = True # Track retrieval events for learning min_items_to_cache: int = 20 # Only cache if original had >= N items diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 07b343c38..e5e82797b 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -628,9 +628,12 @@ class HeadroomProxy( # ContentRouter, so merge rather than assign. if config.exclude_tools: router_config.exclude_tools = set(DEFAULT_EXCLUDE_TOOLS) | config.exclude_tools - # Token mode: allow compression of older excluded-tool results. + # Token mode: allow compression of older excluded-tool results, + # and emit search results grouped by file (path once per file + # instead of repeated on every match line). if is_token_mode(config.mode): router_config.protect_recent_reads_fraction = 0.3 + router_config.search_group_by_file = True # `--compress-user-messages` flips the router's default skip rule. # Off by default for prefix-cache safety; enabled for workloads where # user-message content dominates input (OpenAI/Azure chat with pasted @@ -2737,7 +2740,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: profile_kwargs = proxy_pipeline_kwargs(config) target_ratio = profile_kwargs.get("target_ratio", config.target_ratio) target_savings_percent = None - if isinstance(target_ratio, (int, float)): + if isinstance(target_ratio, int | float): target_savings_percent = round(max(0.0, min(1.0, 1.0 - float(target_ratio))) * 100, 1) return { "savings_profile": config.savings_profile, @@ -2976,7 +2979,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: ) if query: - # Search within cached content + # 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) return { "hash": hash_key, diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index d8b961738..7e203b42e 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -636,6 +636,16 @@ class ContentRouterConfig: # Set to None to use DEFAULT_TOOL_PROFILES from config tool_profiles: dict[str, Any] | None = None + # SmartCrusher configuration override. None → transforms-level + # SmartCrusherConfig() defaults. Lets deployments tune the lossless + # dispatch threshold and compaction heuristics without constructing + # the crusher themselves. + smart_crusher: Any | None = None + + # Group search-compressor output by file (`rg --heading` style). + # Default False; the proxy enables it in token mode. + search_group_by_file: bool = False + # Patterns for detecting mixed content _CODE_FENCE_PATTERN = re.compile(r"^```(\w*)\s*$", re.MULTILINE) @@ -1696,7 +1706,9 @@ class ContentRouter(Transform): enabled=self.config.ccr_enabled, inject_retrieval_marker=self.config.ccr_inject_marker, ) - crusher_config = SmartCrusherConfig() + # Full config override (smart_crusher) wins as the base; + # the per-field knobs from savings profiles still apply on top. + crusher_config = self.config.smart_crusher or SmartCrusherConfig() if self.config.smart_crusher_max_items_after_crush is not None: crusher_config.max_items_after_crush = ( self.config.smart_crusher_max_items_after_crush @@ -1714,9 +1726,11 @@ class ContentRouter(Transform): """Get SearchCompressor (lazy load).""" if self._search_compressor is None: try: - from .search_compressor import SearchCompressor + from .search_compressor import SearchCompressor, SearchCompressorConfig - self._search_compressor = SearchCompressor() + self._search_compressor = SearchCompressor( + SearchCompressorConfig(group_by_file=self.config.search_group_by_file) + ) except ImportError: logger.debug("SearchCompressor not available") return self._search_compressor diff --git a/headroom/transforms/read_lifecycle.py b/headroom/transforms/read_lifecycle.py index a46538e8a..c5ea9296b 100644 --- a/headroom/transforms/read_lifecycle.py +++ b/headroom/transforms/read_lifecycle.py @@ -11,6 +11,12 @@ Real-world data shows 75% of Read output bytes fall into these two categories: - 67% stale (file edited after Read) - 12% superseded (file re-Read later) - Only 20% are fresh (untouched) + +NOTE: a first-sight repeat-Read dedup mechanism (DEDUP_REPEAT) was +prototyped here and removed — `headroom audit-reads` measured +byte-identical repeats at 0.1% of Read bytes on real traffic. If a +deployment's audit shows otherwise, the implementation is in git history +on the feat/compression-extraction branch. """ from __future__ import annotations @@ -374,7 +380,7 @@ class ReadLifecycleManager: ccr_hashes: list[str] = [] bytes_before = 0 bytes_after = 0 - counts = {ReadState.FRESH: 0, ReadState.STALE: 0, ReadState.SUPERSEDED: 0} + counts = dict.fromkeys(ReadState, 0) for c in classifications: counts[c.state] += 1 @@ -486,14 +492,19 @@ class ReadLifecycleManager: file_display = classification.file_path or "unknown" + # NOTE: the literal phrase "Retrieve original: hash=" is load-bearing — + # the compression-pinning checks in ContentRouter and the + # marker-preserving regex in compression_units.py match on it. if classification.state == ReadState.STALE: marker = ( - f"[Read content stale: {file_display} was modified after this read. " + f"[Read content stale: {file_display} was modified after this read — " + f"re-read the file for current content. " f"Retrieve original: hash={ccr_hash}]" ) else: # SUPERSEDED marker = ( - f"[Read content superseded: {file_display} was re-read later. " + f"[Read content superseded: {file_display} was re-read later — " + f"re-read the file if needed. " f"Retrieve original: hash={ccr_hash}]" ) diff --git a/headroom/transforms/search_compressor.py b/headroom/transforms/search_compressor.py index 7eb51971e..d9316c23b 100644 --- a/headroom/transforms/search_compressor.py +++ b/headroom/transforms/search_compressor.py @@ -92,6 +92,11 @@ class SearchCompressorConfig: boost_errors: bool = True enable_ccr: bool = True min_matches_for_ccr: int = 10 + # Group output by file (`rg --heading` style): path emitted once per + # file, then `line:content` rows. Removes per-match path repetition. + # Default False (classic `file:line:content`); the proxy enables it + # in token mode. + group_by_file: bool = False @dataclass @@ -161,6 +166,7 @@ class SearchCompressor: enable_ccr=cfg.enable_ccr, min_matches_for_ccr=cfg.min_matches_for_ccr, min_compression_ratio_for_ccr=0.8, + group_by_file=getattr(cfg, "group_by_file", False), ) ) diff --git a/headroom/transforms/smart_crusher.py b/headroom/transforms/smart_crusher.py index a0c59797d..fe283e177 100644 --- a/headroom/transforms/smart_crusher.py +++ b/headroom/transforms/smart_crusher.py @@ -179,11 +179,27 @@ class SmartCrusherConfig: dedup_identical_items: bool = True first_fraction: float = 0.3 last_fraction: float = 0.15 - # Lossless compaction only replaces the original when it saves at - # least this byte fraction vs the (minified) input. Mirrors the Rust - # default; mainly lowered in tests and KV experiments — KV repeats - # field names per row, so it clears the gate less often than CSV. - lossless_min_savings_ratio: float = 0.30 + # Minimum byte-savings ratio for the lossless Table/CSV compaction + # path to win over the lossy path (0.15, matching the Rust default — + # the two must stay in lockstep, see config.rs). Lossless output + # needs no CCR retrieval round-trip when the model wants more rows, + # so it gets a lower bar than the lossy path. Mainly raised in tests + # and KV experiments — KV repeats field names per row, so it clears + # the gate less often than CSV. + lossless_min_savings_ratio: float = 0.15 + + # Compaction heuristics (mirror Rust CompactConfig; see + # crates/headroom-core/src/transforms/smart_crusher/compaction/compactor.rs). + # A field is "core" if present in at least this fraction of rows. + compaction_core_field_fraction: float = 0.8 + # Below this fraction of core keys, treat the array as heterogeneous + # and look for a discriminator to bucket by. + compaction_heterogeneous_core_ratio: float = 0.6 + # Cap on inner-key count for nested-uniform flattening. + compaction_max_flatten_inner_keys: int = 6 + # Bucket-count bounds for discriminator usefulness. + compaction_min_buckets: int = 2 + compaction_max_buckets: int = 8 # ─── Rust-backed SmartCrusher ───────────────────────────────────────────── @@ -316,11 +332,21 @@ class SmartCrusher(Transform): dedup_identical_items=cfg.dedup_identical_items, first_fraction=cfg.first_fraction, last_fraction=cfg.last_fraction, - lossless_min_savings_ratio=cfg.lossless_min_savings_ratio, relevance_threshold=0.3, enable_ccr_marker=( self._ccr_config.enabled and self._ccr_config.inject_retrieval_marker ), + # getattr fallbacks: callers may pass the structurally-similar + # `headroom.config.SmartCrusherConfig` (MCP server, SDK) or a + # pre-existing config object that predates these fields. + lossless_min_savings_ratio=getattr(cfg, "lossless_min_savings_ratio", 0.15), + compaction_core_field_fraction=getattr(cfg, "compaction_core_field_fraction", 0.8), + compaction_heterogeneous_core_ratio=getattr( + cfg, "compaction_heterogeneous_core_ratio", 0.6 + ), + compaction_max_flatten_inner_keys=getattr(cfg, "compaction_max_flatten_inner_keys", 6), + compaction_min_buckets=getattr(cfg, "compaction_min_buckets", 2), + compaction_max_buckets=getattr(cfg, "compaction_max_buckets", 8), ) # Default: lossless-first compaction (PR4). Lossless wins for # cleanly tabular input where it saves ≥ 30% bytes; otherwise diff --git a/tests/test_adapter_hooks.py b/tests/test_adapter_hooks.py index efb61ecf8..f0a039987 100644 --- a/tests/test_adapter_hooks.py +++ b/tests/test_adapter_hooks.py @@ -385,12 +385,16 @@ class TestCCRContextVarScoping: class TestCCREntryPointLoading: """Verify _create_default_ccr_backend() env-based loading.""" - def test_no_env_returns_none(self, monkeypatch): - """No HEADROOM_CCR_BACKEND → returns None (use InMemoryBackend).""" + def test_no_env_returns_sqlite(self, monkeypatch, tmp_path): + """No HEADROOM_CCR_BACKEND → SQLiteBackend (the persistent default; + restart survival + cross-worker sharing for the 30-min TTL).""" monkeypatch.delenv("HEADROOM_CCR_BACKEND", raising=False) + monkeypatch.setenv("HEADROOM_CCR_SQLITE_PATH", str(tmp_path / "ccr.db")) from headroom.cache.compression_store import _create_default_ccr_backend - assert _create_default_ccr_backend() is None + backend = _create_default_ccr_backend() + assert backend is not None + assert backend.get_stats()["backend_type"] == "sqlite" def test_memory_env_returns_none(self, monkeypatch): """HEADROOM_CCR_BACKEND=memory → returns None (use default).""" diff --git a/tests/test_audit_codex.py b/tests/test_audit_codex.py new file mode 100644 index 000000000..66d036e08 --- /dev/null +++ b/tests/test_audit_codex.py @@ -0,0 +1,135 @@ +"""Tests for the Codex read-pattern audit (headroom.audit.codex).""" + +from __future__ import annotations + +import json + +import pytest + +from headroom.audit.codex import audit_codex, classify_command, render_codex_text, strip_wrappers + + +class TestClassifier: + def test_strip_wrappers(self): + assert strip_wrappers("rtk cat foo.py") == "cat foo.py" + assert strip_wrappers("rtk proxy sed -n '1,20p' foo.py") == "sed -n '1,20p' foo.py" + assert strip_wrappers("git status") == "git status" + + @pytest.mark.parametrize( + ("cmd", "category", "partial"), + [ + ("cat src/foo.py", "read", False), + ("sed -n '1,200p' src/foo.py", "read", True), + ("rtk read src/foo.py --lines 10-50", "read", True), + ("head -50 src/foo.py", "read", True), + ("nl headroom/config.py", "read", False), + ("rg -n 'def apply' headroom/", "search", False), + ("rtk grep -n pattern .", "search", False), + ("git diff HEAD~1", "git", False), + ("apply_patch <<'EOF'\n*** Begin Patch\nEOF", "edit", False), + ("pytest tests/ -x -q", "build/test", False), + ("echo hello", "other", False), + ], + ) + def test_categories(self, cmd, category, partial): + cat, _path, is_partial = classify_command(cmd) + assert cat == category + if category == "read": + assert is_partial == partial + + def test_path_extraction_and_workdir(self): + _, path, _ = classify_command("cat src/foo.py", workdir="/repo") + assert path == "/repo/src/foo.py" + _, path, _ = classify_command("cat /abs/foo.py", workdir="/repo") + assert path == "/abs/foo.py" + + def test_sed_range_not_mistaken_for_path(self): + _, path, _ = classify_command("sed -n '5,30p' headroom/config.py") + assert path == "headroom/config.py" + + def test_compound_command_with_read(self): + cat, path, _ = classify_command("cat foo.py | grep def", workdir="/r") + assert cat == "read" + assert path == "/r/foo.py" + + +def _call(call_id: str, cmd: str, workdir: str = "/repo") -> str: + return json.dumps( + { + "payload": { + "type": "function_call", + "name": "exec_command", + "call_id": call_id, + "arguments": json.dumps({"cmd": cmd, "workdir": workdir}), + } + } + ) + + +def _output(call_id: str, text: str) -> str: + return json.dumps( + {"payload": {"type": "function_call_output", "call_id": call_id, "output": text}} + ) + + +@pytest.fixture +def codex_dir(tmp_path): + content = "line\n" * 600 # 3000B — over the maturation floor + lines = [ + _call("c1", "cat src/foo.py"), + _output("c1", content), + _call("c2", "sed -n '1,100p' src/foo.py"), # partial re-read, same path + _output("c2", content[:500]), + _call("c3", "rg -n 'def ' src/"), + _output("c3", "src/foo.py:1:def x():"), + _call("c4", "rtk read src/bar.py --lines 1-50"), + _output("c4", "bar content " * 10), + ] + sessions = tmp_path / "sessions" / "2026" / "06" + sessions.mkdir(parents=True) + (sessions / "rollout-1.jsonl").write_text("\n".join(lines)) + return tmp_path / "sessions" + + +class TestAuditCodex: + def test_metrics(self, codex_dir): + r = audit_codex(codex_dir) + assert r.sessions == 1 + assert r.exec_calls == 4 + assert r.read_calls == 3 # c1, c2, c4 + assert r.reads_partial == 2 # c2 (sed range), c4 (--lines) + assert r.rereads_same_path == 1 # c2 re-reads foo.py + assert r.distinct_files_read == 2 + assert r.reads_over_floor == 1 # c1 (3000B) + assert r.calls_by_category["search"] == 1 + assert r.bytes_by_category["read"] > r.bytes_by_category["search"] + + def test_render_runs(self, codex_dir): + out = render_codex_text(audit_codex(codex_dir)) + assert "codex read-pattern audit" in out + assert "partial slices" in out + + def test_empty_dir(self, tmp_path): + assert audit_codex(tmp_path).sessions == 0 + + +class TestCli: + def test_cli_codex_mode(self, codex_dir): + from click.testing import CliRunner + + from headroom.cli.main import main + + runner = CliRunner() + res = runner.invoke(main, ["audit-reads", "--codex", "--path", str(codex_dir)]) + assert res.exit_code == 0, res.output + assert "codex read-pattern audit" in res.output + + res = runner.invoke( + main, ["audit-reads", "--codex", "--path", str(codex_dir), "--format", "json"] + ) + assert res.exit_code == 0 + assert json.loads(res.output)["read_calls"] == 3 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_audit_reads.py b/tests/test_audit_reads.py new file mode 100644 index 000000000..813c42906 --- /dev/null +++ b/tests/test_audit_reads.py @@ -0,0 +1,195 @@ +"""Tests for the audit-reads traffic audit (headroom.audit.reads).""" + +from __future__ import annotations + +import json + +import pytest + +from headroom.audit.reads import audit_reads, render_text + +CONTENT = " 1\tdef foo():\n 2\t return 42\n" * 30 # >512B + + +def _line(role: str, content, ts: str = "2026-06-09T10:00:00Z") -> str: + return json.dumps({"message": {"role": role, "content": content}, "timestamp": ts}) + + +def _tool_use(tc_id: str, name: str, inp: dict) -> dict: + return {"type": "tool_use", "id": tc_id, "name": name, "input": inp} + + +def _tool_result(tc_id: str, text: str) -> dict: + return {"type": "tool_result", "tool_use_id": tc_id, "content": text} + + +@pytest.fixture +def transcript_dir(tmp_path): + """Synthetic session: read foo.py twice (identical), partial read + contained in the full read, edit foo.py, then a >5min gap.""" + lines = [ + _line("user", "look at foo.py", "2026-06-09T10:00:00Z"), + _line( + "assistant", + [_tool_use("r1", "Read", {"file_path": "/x/foo.py"})], + "2026-06-09T10:00:01Z", + ), + _line("user", [_tool_result("r1", CONTENT)], "2026-06-09T10:00:02Z"), + _line( + "assistant", + [_tool_use("r2", "Read", {"file_path": "/x/foo.py"})], + "2026-06-09T10:00:03Z", + ), + _line("user", [_tool_result("r2", CONTENT)], "2026-06-09T10:00:04Z"), + _line( + "assistant", + [_tool_use("r3", "Read", {"file_path": "/x/foo.py", "offset": 1, "limit": 2})], + "2026-06-09T10:00:05Z", + ), + # Partial read: a strict substring of the earlier full read. + _line("user", [_tool_result("r3", CONTENT[: len(CONTENT) // 2])], "2026-06-09T10:00:06Z"), + _line( + "assistant", + [_tool_use("e1", "Edit", {"file_path": "/x/foo.py", "old_string": "a"})], + "2026-06-09T10:00:07Z", + ), + _line("user", [_tool_result("e1", "ok")], "2026-06-09T10:00:08Z"), + # >5min gap before the next message. + _line("user", "back from lunch", "2026-06-09T10:20:00Z"), + ] + proj = tmp_path / "projects" / "-x-demo" + proj.mkdir(parents=True) + (proj / "session1.jsonl").write_text("\n".join(lines)) + return tmp_path / "projects" + + +class TestAuditReads: + def test_metrics(self, transcript_dir): + r = audit_reads(transcript_dir) + + assert r.sessions == 1 + assert r.read_calls == 3 + assert r.dedup_identical_calls == 1 # r2 == r1 + assert r.subset_calls == 1 # r3 ⊂ r1 + # Mechanism rows size each opportunity independently, so a read + # can appear in more than one: r1 and r3 both precede the edit + # (stale), and r3 is also a subset of r1. Only identical-repeat + # excludes from stale (replacing a pointer twice is meaningless). + assert r.stale_calls == 2 + assert r.gaps_over_5m == 1 + assert r.sessions_with_gap == 1 + assert r.linenum_overhead_bytes > 0 + assert r.class_bytes.get("source code", 0) > 0 + assert r.tool_bytes["Read"] == r.read_bytes + assert r.reads_per_file_max == 3 + + def test_render_text_runs(self, transcript_dir): + out = render_text(audit_reads(transcript_dir)) + assert "Read opportunity sizing" in out + assert "identical repeat" in out + assert "cache-death windows" in out + + def test_json_roundtrip(self, transcript_dir): + r = audit_reads(transcript_dir) + data = json.loads(r.to_json()) + assert data["read_calls"] == 3 + + def test_malformed_lines_tolerated(self, tmp_path): + proj = tmp_path / "p" + proj.mkdir() + (proj / "bad.jsonl").write_text("not json\n{\n" + _line("user", "hi")) + r = audit_reads(tmp_path) + assert r.sessions == 1 + assert r.read_calls == 0 + + def test_empty_dir(self, tmp_path): + r = audit_reads(tmp_path) + assert r.sessions == 0 + + +class TestMaturationSim: + def test_metrics(self, transcript_dir): + from headroom.audit.maturation import simulate_maturation + + r = simulate_maturation(transcript_dir) + assert r.read_calls == 3 + # r2 and r3 target the already-read foo.py; r3 is partial. + assert r.rereads_any == 2 + assert r.rereads_partial == 1 + # CONTENT is ~1.2KB — below the 2KB maturation floor — so the + # big-read metrics stay empty on this fixture. + assert r.big_reads == 0 + # The edit follows reads of the same file with touch-gap 1. + assert r.edits_with_prior_read == 1 + assert r.at_risk_edits[1] == 0 + + def test_big_read_metrics(self, tmp_path): + from headroom.audit.maturation import MATURE_FLOOR, simulate_maturation + + big = "x" * (MATURE_FLOOR + 100) + lines = [ + _line("assistant", [_tool_use("r1", "Read", {"file_path": "/x/big.py"})]), + _line("user", [_tool_result("r1", big)]), + ] + proj = tmp_path / "p" + proj.mkdir() + (proj / "s.jsonl").write_text("\n".join(lines)) + r = simulate_maturation(tmp_path) + assert r.big_reads == 1 + assert r.never_touched_again == 1 + + def test_render_runs(self, transcript_dir): + from headroom.audit.maturation import render_sim_text, simulate_maturation + + out = render_sim_text(simulate_maturation(transcript_dir)) + assert "maturation simulation" in out + assert "at-risk edits" in out + + +class TestCli: + def test_cli_text_and_json(self, transcript_dir): + from click.testing import CliRunner + + from headroom.cli.main import main + + runner = CliRunner() + res = runner.invoke(main, ["audit-reads", "--path", str(transcript_dir)]) + assert res.exit_code == 0, res.output + assert "Read opportunity sizing" in res.output + + res = runner.invoke( + main, ["audit-reads", "--path", str(transcript_dir), "--format", "json"] + ) + assert res.exit_code == 0 + assert json.loads(res.output)["sessions"] == 1 + + def test_cli_simulate_maturation(self, transcript_dir): + from click.testing import CliRunner + + from headroom.cli.main import main + + runner = CliRunner() + res = runner.invoke( + main, ["audit-reads", "--path", str(transcript_dir), "--simulate-maturation"] + ) + assert res.exit_code == 0, res.output + assert "maturation simulation" in res.output + + res = runner.invoke( + main, + [ + "audit-reads", + "--path", + str(transcript_dir), + "--simulate-maturation", + "--format", + "json", + ], + ) + assert res.exit_code == 0 + data = json.loads(res.output) + assert data["maturation_simulation"]["read_calls"] == 3 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_ccr.py b/tests/test_ccr.py index c68880725..c0e5f8db4 100644 --- a/tests/test_ccr.py +++ b/tests/test_ccr.py @@ -243,7 +243,7 @@ class TestCCRConfig: config = CCRConfig() assert config.enabled is True assert config.store_max_entries == 1000 - assert config.store_ttl_seconds == 300 + assert config.store_ttl_seconds == 1800 # session-scale (was 300) assert config.inject_retrieval_marker is True assert config.feedback_enabled is True assert config.min_items_to_cache == 20 diff --git a/tests/test_ccr_row_drop_store_bridge.py b/tests/test_ccr_row_drop_store_bridge.py index 9fe78a927..7c61a1494 100644 --- a/tests/test_ccr_row_drop_store_bridge.py +++ b/tests/test_ccr_row_drop_store_bridge.py @@ -368,12 +368,18 @@ def test_v1_compress_then_v1_retrieve_resolves_marker_hash() -> None: # Build a payload similar to the issue's reproducer — 200 items # with enough variation to trigger the lossy path. The Rust # crusher's adaptive_k will keep ~15 and drop the rest. + # + # The blob is unique-per-item and long relative to the key names so + # the lossless Table/CSV path (which wins by stripping repeated keys + # when it saves >= lossless_min_savings_ratio) cannot clear the bar — + # this test exists to exercise the LOSSY row-drop path and its + # Rust -> Python CCR store bridge. items = [ { "id": i, "score": 0.99 if i % 30 == 0 else 0.6, "msg": f"Result {i:03d}{' error' if i % 30 == 0 else ' ok'}", - "blob": "x" * 80, + "blob": f"payload-{i:04d}-" + "".join(chr(97 + (i * 7 + j) % 26) for j in range(240)), } for i in range(200) ] diff --git a/tests/test_ccr_sqlite_backend.py b/tests/test_ccr_sqlite_backend.py new file mode 100644 index 000000000..3c6a88550 --- /dev/null +++ b/tests/test_ccr_sqlite_backend.py @@ -0,0 +1,229 @@ +"""Tests for the SQLite CCR backend and session-scale TTL defaults. + +The SQLite backend is the default for `get_compression_store()` because the +30-minute TTL assumes entries survive proxy restarts and are visible across +worker processes — neither holds for the in-memory dict. +""" + +from __future__ import annotations + +import os +import sqlite3 +import time + +import pytest + +from headroom.cache.backends.sqlite import SQLiteBackend +from headroom.cache.compression_store import CompressionEntry, CompressionStore + + +def make_entry(hash_key: str = "h1", content: str = "x" * 600, ttl: int = 1800) -> CompressionEntry: + return CompressionEntry( + hash=hash_key, + original_content=content, + compressed_content="c", + original_tokens=100, + compressed_tokens=10, + original_item_count=50, + compressed_item_count=5, + tool_name="Read", + tool_call_id="t1", + query_context=None, + created_at=time.time(), + ttl=ttl, + ) + + +@pytest.fixture +def db_path(tmp_path): + return tmp_path / "ccr_test.db" + + +class TestSQLiteBackend: + def test_crud_roundtrip(self, db_path): + b = SQLiteBackend(db_path) + entry = make_entry() + b.set("h1", entry) + + got = b.get("h1") + assert got is not None + assert got.original_content == entry.original_content + assert got.tool_name == "Read" + assert got.ttl == 1800 + + assert b.exists("h1") + assert b.count() == 1 + assert b.keys() == ["h1"] + assert b.delete("h1") + assert not b.exists("h1") + assert not b.delete("h1") + + def test_survives_reopen(self, db_path): + """The restart-survival property the default flip exists for.""" + SQLiteBackend(db_path).set("h1", make_entry()) + + reopened = SQLiteBackend(db_path) + got = reopened.get("h1") + assert got is not None + assert got.original_content == "x" * 600 + + def test_two_connections_share_data(self, db_path): + """Multi-worker property: a second live connection sees writes.""" + writer = SQLiteBackend(db_path) + reader = SQLiteBackend(db_path) + writer.set("h1", make_entry()) + assert reader.get("h1") is not None + + def test_items_and_stats(self, db_path): + b = SQLiteBackend(db_path) + b.set("h1", make_entry("h1")) + b.set("h2", make_entry("h2")) + + items = dict(b.items()) + assert set(items) == {"h1", "h2"} + stats = b.get_stats() + assert stats["backend_type"] == "sqlite" + assert stats["entry_count"] == 2 + assert stats["bytes_used"] > 0 + + def test_clear(self, db_path): + b = SQLiteBackend(db_path) + b.set("h1", make_entry()) + b.clear() + assert b.count() == 0 + + def test_unknown_json_fields_tolerated(self, db_path): + """Forward-compat: entries written by a newer headroom version + (extra fields) must still load.""" + b = SQLiteBackend(db_path) + b.set("h1", make_entry()) + with b._lock: + row = b._conn.execute("SELECT entry_json FROM ccr_entries WHERE hash='h1'").fetchone() + doctored = row[0][:-1] + ', "field_from_the_future": 7}' + b._conn.execute("UPDATE ccr_entries SET entry_json=? WHERE hash='h1'", (doctored,)) + b._conn.commit() + + got = b.get("h1") + assert got is not None + assert got.original_content == "x" * 600 + + def test_store_ttl_enforcement_via_compression_store(self, db_path): + """TTL checks stay in CompressionStore; expired entries miss.""" + store = CompressionStore(backend=SQLiteBackend(db_path)) + expired = make_entry(ttl=1) + expired.created_at = time.time() - 10 + store._backend.set("h1", expired) + + assert store.retrieve("h1") is None + + def test_retrieval_count_persists(self, db_path): + """record_access mutations are re-persisted (store re-sets the + entry after mutating), so feedback counts survive reopen.""" + store = CompressionStore(backend=SQLiteBackend(db_path)) + store._backend.set("h1", make_entry()) + store.retrieve("h1", query="foo") + + reopened = SQLiteBackend(db_path) + got = reopened.get("h1") + assert got is not None + assert got.retrieval_count == 1 + + +class TestMultiWorkerSafety: + def test_busy_error_does_not_delete_database(self, db_path): + """SQLITE_BUSY (OperationalError, a DatabaseError subclass) under + multi-worker write contention must be treated as transient — NOT + as corruption that deletes every stored original.""" + b = SQLiteBackend(db_path) + b.set("h1", make_entry()) + + class BusyOnceConn: + """Delegating wrapper; first SELECT raises 'database is locked'.""" + + def __init__(self, real): + self._real = real + self.raised = False + + def execute(self, *args, **kwargs): + if not self.raised and args and "SELECT" in args[0]: + self.raised = True + raise sqlite3.OperationalError("database is locked") + return self._real.execute(*args, **kwargs) + + def __getattr__(self, name): + return getattr(self._real, name) + + real = b._conn + b._conn = BusyOnceConn(real) # type: ignore[assignment] + assert b.get("h1") is None # transient miss, not a crash + + b._conn = real + # The data and the database file both survived. + assert db_path.exists() + assert b.get("h1") is not None + + def test_corruption_message_triggers_reset(self, db_path): + b = SQLiteBackend(db_path) + b.set("h1", make_entry()) + b._handle_db_error(sqlite3.DatabaseError("database disk image is malformed"), "get") + # Database recreated: empty but functional. + assert b.count() == 0 + b.set("h2", make_entry("h2")) + assert b.exists("h2") + + def test_busy_timeout_configured(self, db_path): + b = SQLiteBackend(db_path) + timeout = b._conn.execute("PRAGMA busy_timeout").fetchone()[0] + assert timeout >= 5000 + + def test_expired_rows_purged_on_open(self, db_path): + b = SQLiteBackend(db_path) + expired = make_entry(ttl=1) + expired.created_at = time.time() - 10 + b.set("old", expired) + b.set("fresh", make_entry("fresh")) + + reopened = SQLiteBackend(db_path) + assert not reopened.exists("old") # swept at open + assert reopened.exists("fresh") + + @pytest.mark.skipif(os.name != "posix", reason="POSIX permissions") + def test_database_file_is_private(self, db_path): + SQLiteBackend(db_path) + mode = db_path.stat().st_mode & 0o777 + assert mode == 0o600 + + +class TestDefaults: + def test_session_scale_ttl_lockstep(self): + """CCRConfig, CompressionEntry, and CompressionStore must agree.""" + from headroom.config import CCRConfig + + assert CCRConfig().store_ttl_seconds == 1800 + assert CompressionEntry.__dataclass_fields__["ttl"].default == 1800 + assert CompressionStore()._default_ttl == 1800 + + def test_default_backend_is_sqlite(self, monkeypatch, tmp_path): + from headroom.cache.compression_store import _create_default_ccr_backend + + monkeypatch.delenv("HEADROOM_CCR_BACKEND", raising=False) + monkeypatch.setenv("HEADROOM_CCR_SQLITE_PATH", str(tmp_path / "d.db")) + backend = _create_default_ccr_backend() + assert backend is not None + assert backend.get_stats()["backend_type"] == "sqlite" + + def test_memory_opt_out(self, monkeypatch): + from headroom.cache.compression_store import _create_default_ccr_backend + + monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory") + assert _create_default_ccr_backend() is None + + def test_miss_message_is_actionable(self): + from headroom.cache.compression_store import CCR_MISS_MESSAGE + + assert "re-read" in CCR_MISS_MESSAGE + assert "re-run" in CCR_MISS_MESSAGE + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_compression_store.py b/tests/test_compression_store.py index 699b052ff..787373d66 100644 --- a/tests/test_compression_store.py +++ b/tests/test_compression_store.py @@ -255,7 +255,7 @@ class TestCompressionEntry: created_at=time.time(), ) assert entry.hash == "abc123" - assert entry.ttl == 300 # Default TTL + assert entry.ttl == 1800 # Default TTL (session-scale) assert entry.retrieval_count == 0 assert entry.search_queries == [] assert entry.last_accessed is None @@ -427,7 +427,7 @@ class TestCompressionStoreInit: store = CompressionStore() assert store._max_entries == 1000 - assert store._default_ttl == 300 + assert store._default_ttl == 1800 assert store._enable_feedback is True assert store._backend is not None @@ -540,7 +540,7 @@ class TestCompressionStoreOperations: entry = store.retrieve(hash_key) assert entry is not None - assert entry.ttl == 300 # Default TTL + assert entry.ttl == 1800 # Default TTL (session-scale) def test_store_accepts_custom_ttl(self, store: CompressionStore): """store() accepts custom TTL override.""" diff --git a/tests/test_compressor_config_exposure.py b/tests/test_compressor_config_exposure.py new file mode 100644 index 000000000..1737b294b --- /dev/null +++ b/tests/test_compressor_config_exposure.py @@ -0,0 +1,143 @@ +"""Tests for the newly exposed Rust compressor knobs. + +Covers: +- lossless_min_savings_ratio plumbing (Python dataclass → PyO3 → Rust) + and the 0.15 lockstep default on both sides. +- CompactConfig heuristics plumbing. +- SearchCompressor group_by_file output mode (`rg --heading` style). +- factor_out_constants config acceptance end-to-end. +- ContentRouter plumbing for both knobs. + +Requires the rebuilt `headroom._core` extension — these tests fail loudly +(not skip) if the installed extension predates the new fields, because a +silent version skew here is exactly the parity drift the lockstep rule +exists to prevent. +""" + +from __future__ import annotations + +import pytest + +from headroom.transforms.search_compressor import SearchCompressor, SearchCompressorConfig +from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig + + +class TestSmartCrusherConfigExposure: + def test_python_defaults(self): + cfg = SmartCrusherConfig() + assert cfg.lossless_min_savings_ratio == 0.15 + assert cfg.compaction_core_field_fraction == 0.8 + assert cfg.compaction_heterogeneous_core_ratio == 0.6 + assert cfg.compaction_max_flatten_inner_keys == 6 + assert cfg.compaction_min_buckets == 2 + assert cfg.compaction_max_buckets == 8 + + def test_rust_default_lockstep(self): + """Rust PyO3 default must equal the Python dataclass default.""" + from headroom._core import SmartCrusherConfig as RustConfig + + assert RustConfig().lossless_min_savings_ratio == 0.15 + assert RustConfig().compaction_core_field_fraction == 0.8 + assert RustConfig().compaction_max_buckets == 8 + + def test_values_reach_rust(self): + from headroom._core import SmartCrusherConfig as RustConfig + + rust_cfg = RustConfig( + lossless_min_savings_ratio=0.42, + compaction_core_field_fraction=0.7, + compaction_heterogeneous_core_ratio=0.5, + compaction_max_flatten_inner_keys=10, + compaction_min_buckets=3, + compaction_max_buckets=12, + ) + assert rust_cfg.lossless_min_savings_ratio == 0.42 + assert rust_cfg.compaction_core_field_fraction == 0.7 + assert rust_cfg.compaction_heterogeneous_core_ratio == 0.5 + assert rust_cfg.compaction_max_flatten_inner_keys == 10 + assert rust_cfg.compaction_min_buckets == 3 + assert rust_cfg.compaction_max_buckets == 12 + + def test_crusher_accepts_new_fields(self): + crusher = SmartCrusher( + config=SmartCrusherConfig( + lossless_min_savings_ratio=0.5, + factor_out_constants=True, + ) + ) + # Construction succeeded and compresses without error. + items = ",".join(f'{{"id": {i}, "status": "ok"}}' for i in range(40)) + result = crusher.crush(f"[{items}]") + assert result.compressed + + def test_foreign_config_object_tolerated(self): + """headroom.config.SmartCrusherConfig (the SDK-surface class) is + structurally similar and flows through getattr fallbacks.""" + from headroom.config import SmartCrusherConfig as SdkConfig + + crusher = SmartCrusher(config=SdkConfig()) # type: ignore[arg-type] + assert crusher is not None + + +class TestSearchGroupedOutput: + INPUT = "\n".join( + [ + "src/very/long/path/to/module.py:10:def alpha():", + "src/very/long/path/to/module.py:20:def beta():", + "src/very/long/path/to/module.py:30:def gamma():", + "src/other.py:5:class Other:", + ] + ) + + def test_standard_format_default(self): + result = SearchCompressor(SearchCompressorConfig()).compress(self.INPUT) + # Classic file:line:content — path on every match line. + assert "src/very/long/path/to/module.py:10:" in result.compressed + + def test_grouped_format(self): + result = SearchCompressor(SearchCompressorConfig(group_by_file=True)).compress(self.INPUT) + lines = result.compressed.splitlines() + # Path appears as a heading... + assert "src/very/long/path/to/module.py" in lines + # ...and match lines carry only line:content. + assert "10:def alpha():" in lines + # No classic-format line remains. + assert not any(line.startswith("src/very/long/path/to/module.py:10:") for line in lines) + + def test_grouped_is_smaller(self): + std = SearchCompressor(SearchCompressorConfig()).compress(self.INPUT) + grp = SearchCompressor(SearchCompressorConfig(group_by_file=True)).compress(self.INPUT) + assert len(grp.compressed) < len(std.compressed) + + def test_grouped_deterministic(self): + c = SearchCompressor(SearchCompressorConfig(group_by_file=True)) + assert c.compress(self.INPUT).compressed == c.compress(self.INPUT).compressed + + +class TestContentRouterPlumbing: + def test_router_passes_smart_crusher_config(self): + from headroom.transforms.content_router import ContentRouter, ContentRouterConfig + + router = ContentRouter( + ContentRouterConfig(smart_crusher=SmartCrusherConfig(lossless_min_savings_ratio=0.33)) + ) + crusher = router._get_smart_crusher() + assert crusher.config.lossless_min_savings_ratio == 0.33 + + def test_router_passes_search_grouping(self): + from headroom.transforms.content_router import ContentRouter, ContentRouterConfig + + router = ContentRouter(ContentRouterConfig(search_group_by_file=True)) + compressor = router._get_search_compressor() + assert compressor.config.group_by_file is True + + def test_router_defaults_off(self): + from headroom.transforms.content_router import ContentRouter, ContentRouterConfig + + router = ContentRouter(ContentRouterConfig()) + compressor = router._get_search_compressor() + assert compressor.config.group_by_file is False + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_live/__init__.py b/tests/test_live/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/test_live/conftest.py b/tests/test_live/conftest.py new file mode 100644 index 000000000..6d4801f64 --- /dev/null +++ b/tests/test_live/conftest.py @@ -0,0 +1,33 @@ +"""Shared fixtures for live API tests. + +These tests hit real provider APIs and cost (small amounts of) money. +They are skipped unless the relevant key is present, and live in their +own directory so broad suite runs can exclude them wholesale: + + python -m pytest tests/ --ignore=tests/test_live + +Keys are loaded from the repo-root .env when present so the suite works +in the same environment the proxy runs in. +""" + +from __future__ import annotations + +import os +from pathlib import Path + + +def _load_dotenv() -> None: + env_path = Path(__file__).resolve().parents[2] / ".env" + if not env_path.exists(): + return + for line in env_path.read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, _, value = line.partition("=") + key, value = key.strip(), value.strip().strip("'\"") + if key and value and key not in os.environ: + os.environ[key] = value + + +_load_dotenv() diff --git a/tests/test_live/test_live_anthropic.py b/tests/test_live/test_live_anthropic.py new file mode 100644 index 000000000..d8a58636c --- /dev/null +++ b/tests/test_live/test_live_anthropic.py @@ -0,0 +1,149 @@ +"""Live Anthropic API tests for the Claude Code (tool_result block) path. + +Validates the two claims unit tests cannot: the real API accepts our +transformed message shapes, and the model can still answer correctly from +them — the no-accuracy-loss contract, end to end. + +Skipped without ANTHROPIC_API_KEY. Costs: a few hundred haiku tokens/run. +""" + +from __future__ import annotations + +import os + +import httpx +import pytest + +from headroom.transforms.search_compressor import SearchCompressor, SearchCompressorConfig + +pytestmark = pytest.mark.skipif( + not os.environ.get("ANTHROPIC_API_KEY"), + reason="ANTHROPIC_API_KEY not set", +) + +MODEL = "claude-haiku-4-5-20251001" +API_URL = "https://api.anthropic.com/v1/messages" + +FILE_CONTENT = ( + ' 1\tdef answer():\n 2\t """Returns the magic number."""\n 3\t return 42\n' +) + "".join(f" {i}\t# padding line {i}\n" for i in range(4, 40)) + +READ_TOOL = { + "name": "Read", + "description": "Read a file", + "input_schema": { + "type": "object", + "properties": {"file_path": {"type": "string"}}, + "required": ["file_path"], + }, +} + + +def call_anthropic(messages: list[dict], system: str | None = None) -> str: + body: dict = { + "model": MODEL, + "max_tokens": 150, + "tools": [READ_TOOL], + "messages": messages, + } + if system: + body["system"] = system + resp = httpx.post( + API_URL, + json=body, + headers={ + "x-api-key": os.environ["ANTHROPIC_API_KEY"], + "anthropic-version": "2023-06-01", + }, + timeout=60, + ) + assert resp.status_code == 200, f"{resp.status_code}: {resp.text[:500]}" + return "".join( + block.get("text", "") for block in resp.json()["content"] if block.get("type") == "text" + ) + + +def read_roundtrip(tc_id: str, content: str) -> list[dict]: + return [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": tc_id, + "name": "Read", + "input": {"file_path": "/src/magic.py"}, + } + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": tc_id, "content": content}], + }, + ] + + +class TestLifecycleMarkerLive: + def test_api_accepts_stale_marker_shape(self): + """A stale-Read marker (file edited after read) must be a valid + message body and not confuse the model into inventing content.""" + messages = [{"role": "user", "content": "Read /src/magic.py"}] + messages += read_roundtrip("toolu_r1", FILE_CONTENT) + messages += [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_e1", + "name": "Read", # registered tool; the marker is what matters + "input": {"file_path": "/src/magic.py"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_e1", + "content": "[Read content stale: /src/magic.py was modified after this " + "read — re-read the file for current content. " + "Retrieve original: hash=abc123def456abc123def456]", + } + ], + }, + { + "role": "user", + "content": "Is the content of /src/magic.py in this conversation current or " + "stale? One word.", + }, + ] + reply = call_anthropic(messages) + assert "stale" in reply.lower(), f"model misread the marker: {reply!r}" + + +class TestGroupedSearchLive: + def test_model_reads_grouped_format(self): + raw = "\n".join( + [ + "src/payments/processor.py:12:def charge(amount):", + "src/payments/processor.py:45:def refund(amount):", + "src/users/auth.py:7:def login(user):", + ] + ) + grouped = SearchCompressor(SearchCompressorConfig(group_by_file=True)).compress(raw) + messages = [ + { + "role": "user", + "content": "Here are grep results for 'def ':\n\n" + + grouped.compressed + + "\n\nWhich file defines refund()? Reply with just the path.", + } + ] + reply = call_anthropic(messages) + assert "src/payments/processor.py" in reply, f"grouped format misread: {reply!r}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_live/test_live_openai.py b/tests/test_live/test_live_openai.py new file mode 100644 index 000000000..e77f22309 --- /dev/null +++ b/tests/test_live/test_live_openai.py @@ -0,0 +1,109 @@ +"""Live OpenAI API tests for the Codex (role="tool" message) path. + +Same contract as the Anthropic live tests: the real API must accept our +transformed message shapes (lifecycle markers inside role="tool" results), +and the model must read them correctly. + +Skipped without OPENAI_API_KEY. Costs: a few hundred gpt-4o-mini tokens/run. +""" + +from __future__ import annotations + +import json +import os + +import httpx +import pytest + +pytestmark = pytest.mark.skipif( + not os.environ.get("OPENAI_API_KEY"), + reason="OPENAI_API_KEY not set", +) + +MODEL = "gpt-4o-mini" +API_URL = "https://api.openai.com/v1/chat/completions" + +FILE_CONTENT = ( + ' 1\tdef answer():\n 2\t """Returns the magic number."""\n 3\t return 42\n' +) + "".join(f" {i}\t# padding line {i}\n" for i in range(4, 40)) + +STALE_MARKER = ( + "[Read content stale: /src/magic.py was modified after this read — " + "re-read the file for current content. " + "Retrieve original: hash=abc123def456abc123def456]" +) + +READ_TOOL = { + "type": "function", + "function": { + "name": "Read", + "description": "Read a file", + "parameters": { + "type": "object", + "properties": {"file_path": {"type": "string"}}, + "required": ["file_path"], + }, + }, +} + + +def call_openai(messages: list[dict]) -> str: + resp = httpx.post( + API_URL, + json={ + "model": MODEL, + "max_tokens": 150, + "tools": [READ_TOOL], + "messages": messages, + }, + headers={"Authorization": f"Bearer {os.environ['OPENAI_API_KEY']}"}, + timeout=60, + ) + assert resp.status_code == 200, f"{resp.status_code}: {resp.text[:500]}" + return resp.json()["choices"][0]["message"]["content"] or "" + + +def read_roundtrip(tc_id: str, content: str) -> list[dict]: + return [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": tc_id, + "type": "function", + "function": { + "name": "Read", + "arguments": json.dumps({"file_path": "/src/magic.py"}), + }, + } + ], + }, + {"role": "tool", "tool_call_id": tc_id, "content": content}, + ] + + +class TestLifecycleMarkerLiveOpenAI: + def test_api_accepts_stale_marker_shape(self): + """A stale-Read marker inside a role='tool' message must be a + valid body and must be read as 'this content is outdated'.""" + messages = [{"role": "user", "content": "Read /src/magic.py"}] + messages += read_roundtrip("call_r1", FILE_CONTENT) + messages += [ + {"role": "assistant", "content": "Read it. Anything else?"}, + {"role": "user", "content": "Check it once more."}, + ] + messages += read_roundtrip("call_r2", STALE_MARKER) + messages.append( + { + "role": "user", + "content": "Is the latest read of /src/magic.py in this conversation " + "current or stale? One word.", + } + ) + reply = call_openai(messages) + assert "stale" in reply.lower(), f"model misread the marker: {reply!r}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_proxy_ccr.py b/tests/test_proxy_ccr.py index dc86fafa9..a553cec52 100644 --- a/tests/test_proxy_ccr.py +++ b/tests/test_proxy_ccr.py @@ -67,7 +67,7 @@ class TestCCRRetrieveEndpoint: response = client.post("/v1/retrieve", json={"hash": "nonexistent123"}) assert response.status_code == 404 assert "Entry not found" in response.json()["detail"] - assert "CCR TTL: 300 seconds" in response.json()["detail"] + assert "CCR TTL: 1800 seconds" in response.json()["detail"] def test_retrieve_expired_hash_reports_expiration_detail(self, client): """Expired entries report expiration separately from missing hashes.""" @@ -263,7 +263,7 @@ class TestCCRStatsEndpoint: data = response.json() assert "store" in data assert data["store"]["entry_count"] == 0 - assert data["store"]["default_ttl_seconds"] == 300 + assert data["store"]["default_ttl_seconds"] == 1800 assert "recent_retrievals" in data def test_stats_exposes_env_configured_ttl(self, client, monkeypatch): diff --git a/tests/test_proxy_scalability.py b/tests/test_proxy_scalability.py index 654c78eec..effa361a0 100644 --- a/tests/test_proxy_scalability.py +++ b/tests/test_proxy_scalability.py @@ -258,15 +258,21 @@ class TestWorkerConfiguration: monkeypatch.delenv(_MULTI_WORKER_CONFIG_ENV, raising=False) - with patch("headroom.proxy.server.uvicorn.run", fake_run): - run_server(config, workers=4, limit_concurrency=250) + try: + with patch("headroom.proxy.server.uvicorn.run", fake_run): + run_server(config, workers=4, limit_concurrency=250) - assert captured["app"] == "headroom.proxy.server:create_app_from_env" - assert captured["kwargs"]["workers"] == 4 - assert captured["kwargs"]["limit_concurrency"] == 250 - assert captured["kwargs"]["factory"] is True - payload = json.loads(os.environ[_MULTI_WORKER_CONFIG_ENV]) - assert payload["host"] == "0.0.0.0" - assert payload["port"] == 8787 - assert payload["max_connections"] == 200 - monkeypatch.delenv(_MULTI_WORKER_CONFIG_ENV, raising=False) + assert captured["app"] == "headroom.proxy.server:create_app_from_env" + assert captured["kwargs"]["workers"] == 4 + assert captured["kwargs"]["limit_concurrency"] == 250 + assert captured["kwargs"]["factory"] is True + payload = json.loads(os.environ[_MULTI_WORKER_CONFIG_ENV]) + assert payload["host"] == "0.0.0.0" + assert payload["port"] == 8787 + assert payload["max_connections"] == 200 + finally: + # run_server sets this via raw os.environ. Pop it directly rather + # than via monkeypatch.delenv: delenv records the current (JSON) + # value and re-restores it on teardown, leaking the config into + # later tests (e.g. _proxy_config_from_env then ignores HEADROOM_*). + os.environ.pop(_MULTI_WORKER_CONFIG_ENV, None)