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

fix: B1 — retire ICM, RollingWindow, scoring, relevance + dependents

Browse source 
Phase B step 1 of the live-zone-only realignment. Removes ~10K LOC of
"drop messages from history" machinery that became unreachable after
PR-A1 made `/v1/messages` a passthrough on the proxy. Live-zone-only
compression (PR-B2..B7) operates on content blocks within messages;
message-list mutation no longer happens in the pipeline.

Python deletes:
- headroom/transforms/intelligent_context.py (1077 LOC)
- headroom/transforms/rolling_window.py (395 LOC)
- headroom/transforms/progressive_summarizer.py (508 LOC)
- headroom/transforms/scoring.py (459 LOC)
- headroom/transforms/tool_crusher.py (338 LOC)
- 5 corresponding tests/test_transforms/* and tests/test_proxy_intelligent_context.py

Rust deletes:
- crates/headroom-core/src/context/* (manager, config, workspace,
  candidate, ccr_drop, strategy/, mod) + safety.rs replaced
- crates/headroom-core/src/scoring/* (mod, score, scorer, traits, weights)
- MessageScorerComparator from crates/headroom-parity (PR #338/#343
  becomes deletable; sunk cost stays sunk)
- 13 message_scorer fixtures + record_message_scorer.py

Rust adds (move + rewrite):
- crates/headroom-core/src/transforms/safety.rs — `tool_pair_indices`
  preserves the OpenAI/Anthropic tool_use ↔ tool_result pairing rule
  the live-zone dispatcher (PR-B2) needs. No IcmConfig dependency.

Surface refactors:
- HeadroomConfig: drop `tool_crusher`, `rolling_window`,
  `intelligent_context` fields; hoist `output_buffer_tokens` to top
  level (used by client.py).
- ProxyConfig: drop `intelligent_context*` fields.
- `headroom wrap` proxy server: retire IntelligentContextManager
  and RollingWindow imports + branch; pipeline is CacheAligner →
  ContentRouter (smart_routing) or CacheAligner → SmartCrusher
  (legacy).
- CLI: drop `--no-intelligent-context`, `--no-intelligent-scoring`,
  `--no-compress-first` flags.
- LangChain memory integration: rename `_apply_rolling_window` →
  `_apply_compression`, drop RollingWindowConfig dep. Threshold is
  now advisory — B6 will rework the contract.
- TransformPipeline.create_pipeline now takes only cache_aligner_config.
- headroom/__init__.py + headroom/transforms/__init__.py: strip
  exports of deleted symbols.

Bug fixes uncovered by full pytest sweep:
- providers/copilot/wrap.py: `environ or os.environ` collapsed
  empty-dict to falsy → callers passing `environ={}` accidentally
  pulled from os.environ. Use `environ if environ is not None else
  os.environ`.

Test correctness fixes:
- _DummyAnthropicHandler._retry_request gains **_kwargs to match
  the real handler signature post-A8.
- test_ws_http_fallback extracts JSON from `content=` (post-A3
  byte-faithful) rather than the obsolete `json=` kwarg.
- test_ccr_response_handler_extra fixture joins SSE events with
  `\n\n` per spec (post-A8 byte-buffer parser requirement).
- test_proxy_responses_phase_preservation: capture via direct
  handler attached to the named logger, so the assertion is
  order-independent (proxy `_setup_file_logging` flips
  `headroom.propagate=False` once any earlier test triggers it).
- conftest.py autouse fixture resets `headroom.propagate=True`
  before each test as a defensive measure for the same pollution.
- test_wrap_copilot_translated_backend_still_requires_byok:
  monkeypatch.delenv every provider key so the BYOK error
  actually fires.
- test_native_installers: skip when system bash < 4.3 (macOS ships 3.2).
- TestGeminiEmbedContent / TestGeminiBatchEmbedContents:
  pytest.mark.skip — proxy currently has no :embedContent route;
  feature gap, not regression.

Acceptance:
- cargo build --workspace + cargo clippy + cargo fmt --check: green.
- cargo test --workspace --exclude headroom-py: 777 passed.
- pytest: 4892 passed, 240 skipped, 0 failed.
- git grep returns only intentional comments referencing the deletion.

Per-PR-B1 plan: REALIGNMENT/04-phase-B-live-zone.md.
This commit is contained in:
chopratejas 2026-05-02 12:23:17 -07:00
parent 341dcf03e9
commit 967b0db439
70 changed files with 4681 additions and 17949 deletions
Download patch file Download diff file Expand all files Collapse all files

415
crates/headroom-core/src/context/candidate.rs
View file

@ -1,415 +0,0 @@
//! Build drop candidates from a scored message list.
//!
//! Three candidate shapes (matching Python):
//!
//! - **Tool unit** — atomic `(assistant_with_tool_calls, [tool_responses])`.
//! Either the whole unit is dropped or none of it. Score = mean of
//! member scores.
//! - **Turn** — paired `(user, assistant)` neighbours when both are
//! unprotected and neither is in a tool unit. Score = mean of the
//! two scores.
//! - **Single** — any other unprotected, non-tool-unit message that
//! couldn't be paired. Score = the message's own score.
//!
//! Candidates sort by score ascending (lowest = drop first), with
//! position as a tiebreaker (older first).
use std::collections::HashSet;
use serde_json::Value;
use crate::scoring::MessageScore;
/// One unit the cascade may drop atomically.
#[derive(Debug, Clone)]
pub struct DropCandidate {
pub kind: CandidateKind,
/// Indices to remove together. May be 1 (single), 2 (turn), or
/// `1 + N` (tool unit with N tool responses).
pub indices: Vec<usize>,
/// Aggregated score; lower = drop first.
pub score: f32,
/// Earliest index — used as the secondary sort key for stability
/// (older candidates of equal score drop first).
pub position: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CandidateKind {
Single,
Turn,
ToolUnit,
}
/// Tool unit: `(assistant_index, [tool_response_indices])`.
pub type ToolUnit = (usize, Vec<usize>);
/// Find all tool units in the message list. Mirrors Python's
/// `headroom.parser.find_tool_units`. Handles OpenAI, Anthropic, and
/// Strands SDK shapes.
pub fn find_tool_units(messages: &[Value]) -> Vec<ToolUnit> {
use std::collections::HashMap;
// tool_call_id → response message index
let mut response_map: HashMap<String, usize> = HashMap::new();
for (i, msg) in messages.iter().enumerate() {
let role = msg.get("role").and_then(Value::as_str);
// OpenAI: role=tool, tool_call_id
if role == Some("tool") {
if let Some(tcid) = msg.get("tool_call_id").and_then(Value::as_str) {
response_map.insert(tcid.to_string(), i);
}
}
// Anthropic + Strands: role=user, content blocks
if role == Some("user") {
if let Some(blocks) = msg.get("content").and_then(Value::as_array) {
for block in blocks {
// Anthropic: {type: tool_result, tool_use_id}
if block.get("type").and_then(Value::as_str) == Some("tool_result") {
if let Some(tcid) = block.get("tool_use_id").and_then(Value::as_str) {
response_map.insert(tcid.to_string(), i);
}
}
// Strands: {toolResult: {toolUseId}}
if let Some(tr) = block.get("toolResult") {
if let Some(tcid) = tr.get("toolUseId").and_then(Value::as_str) {
response_map.insert(tcid.to_string(), i);
}
}
}
}
}
}
let mut units: Vec<ToolUnit> = Vec::new();
for (i, msg) in messages.iter().enumerate() {
if msg.get("role").and_then(Value::as_str) != Some("assistant") {
continue;
}
let mut response_indices: Vec<usize> = Vec::new();
// OpenAI tool_calls
if let Some(arr) = msg.get("tool_calls").and_then(Value::as_array) {
for tc in arr {
if let Some(id) = tc.get("id").and_then(Value::as_str) {
if let Some(&idx) = response_map.get(id) {
response_indices.push(idx);
}
}
}
}
// Anthropic + Strands content blocks
if let Some(blocks) = msg.get("content").and_then(Value::as_array) {
for block in blocks {
// Anthropic: {type: tool_use, id}
if block.get("type").and_then(Value::as_str) == Some("tool_use") {
if let Some(id) = block.get("id").and_then(Value::as_str) {
if let Some(&idx) = response_map.get(id) {
response_indices.push(idx);
}
}
}
// Strands: {toolUse: {toolUseId}}
if let Some(tu) = block.get("toolUse") {
if let Some(id) = tu.get("toolUseId").and_then(Value::as_str) {
if let Some(&idx) = response_map.get(id) {
response_indices.push(idx);
}
}
}
}
}
if !response_indices.is_empty() {
response_indices.sort_unstable();
response_indices.dedup();
units.push((i, response_indices));
}
}
units
}
/// Aggregate `messages × scores × tool_units` into a sorted list of
/// drop candidates. Lowest score first; ties broken by position.
pub fn build_candidates(
messages: &[Value],
scores: &[MessageScore],
protected: &HashSet<usize>,
tool_units: &[ToolUnit],
) -> Vec<DropCandidate> {
// Track every index that's part of some tool unit — we don't
// want a tool-unit message to ALSO appear as a single or turn.
let mut tool_unit_indices: HashSet<usize> = HashSet::new();
for (asst_idx, responses) in tool_units {
tool_unit_indices.insert(*asst_idx);
for &r in responses {
tool_unit_indices.insert(r);
}
}
let mut candidates: Vec<DropCandidate> = Vec::new();
// Tool unit candidates — drop atomically.
for (asst_idx, responses) in tool_units {
if protected.contains(asst_idx) {
continue;
}
// If any response is protected, the unit can't drop atomically.
// Skip it; its messages stay (mirrors Python behaviour).
if responses.iter().any(|r| protected.contains(r)) {
continue;
}
let mut indices: Vec<usize> = vec![*asst_idx];
indices.extend(responses);
let avg = mean_score(&indices, scores);
candidates.push(DropCandidate {
kind: CandidateKind::ToolUnit,
indices,
score: avg,
position: *asst_idx,
});
}
// User+assistant turn pairs and singles.
let mut i: usize = 0;
while i < messages.len() {
if protected.contains(&i) || tool_unit_indices.contains(&i) {
i += 1;
continue;
}
let role = messages[i].get("role").and_then(Value::as_str);
// Try to pair user with the next assistant.
if role == Some("user") && i + 1 < messages.len() {
let next_role = messages[i + 1].get("role").and_then(Value::as_str);
let next_unprotected = !protected.contains(&(i + 1));
let next_not_in_tool_unit = !tool_unit_indices.contains(&(i + 1));
if next_role == Some("assistant") && next_unprotected && next_not_in_tool_unit {
let avg = mean_score(&[i, i + 1], scores);
candidates.push(DropCandidate {
kind: CandidateKind::Turn,
indices: vec![i, i + 1],
score: avg,
position: i,
});
i += 2;
continue;
}
}
// Single — any unprotected, non-tool-unit user/assistant message.
if matches!(role, Some("user") | Some("assistant")) {
let s = scores.get(i).map(|sc| sc.total_score).unwrap_or(0.5);
candidates.push(DropCandidate {
kind: CandidateKind::Single,
indices: vec![i],
score: s,
position: i,
});
}
i += 1;
}
// Sort by (score asc, position asc) — lowest score drops first;
// older messages drop first as a tiebreaker.
candidates.sort_by(|a, b| {
a.score
.partial_cmp(&b.score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.position.cmp(&b.position))
});
candidates
}
fn mean_score(indices: &[usize], scores: &[MessageScore]) -> f32 {
if indices.is_empty() {
return 0.5;
}
let total: f32 = indices
.iter()
.filter_map(|&i| scores.get(i).map(|s| s.total_score))
.sum();
let count = indices.iter().filter(|&&i| scores.get(i).is_some()).count() as f32;
if count == 0.0 {
0.5
} else {
total / count
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn score(idx: usize, total: f32) -> MessageScore {
MessageScore {
message_index: idx,
total_score: total,
recency_score: 0.0,
semantic_score: 0.0,
toin_score: 0.0,
error_score: 0.0,
reference_score: 0.0,
density_score: 0.0,
tokens: 0,
is_protected: false,
drop_safe: true,
score_breakdown: Default::default(),
}
}
#[test]
fn find_tool_units_openai_shape() {
let msgs = vec![
json!({"role": "user", "content": "go"}),
json!({
"role": "assistant",
"content": "",
"tool_calls": [{"id": "c1", "function": {"name": "f"}}]
}),
json!({"role": "tool", "tool_call_id": "c1", "content": "result"}),
];
let units = find_tool_units(&msgs);
assert_eq!(units.len(), 1);
assert_eq!(units[0].0, 1);
assert_eq!(units[0].1, vec![2]);
}
#[test]
fn find_tool_units_anthropic_shape() {
let msgs = vec![
json!({
"role": "assistant",
"content": [
{"type": "tool_use", "id": "tu_1", "name": "f", "input": {}}
]
}),
json!({
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "tu_1", "content": "ok"}]
}),
];
let units = find_tool_units(&msgs);
assert_eq!(units.len(), 1);
assert_eq!(units[0].0, 0);
assert_eq!(units[0].1, vec![1]);
}
#[test]
fn find_tool_units_strands_sdk_shape() {
let msgs = vec![
json!({
"role": "assistant",
"content": [{"toolUse": {"toolUseId": "tu_x", "name": "f"}}]
}),
json!({
"role": "user",
"content": [{"toolResult": {"toolUseId": "tu_x", "content": "ok"}}]
}),
];
let units = find_tool_units(&msgs);
assert_eq!(units.len(), 1);
}
#[test]
fn lowest_score_sorts_first() {
let msgs = vec![
json!({"role": "user", "content": "high"}),
json!({"role": "assistant", "content": "high"}),
json!({"role": "user", "content": "low"}),
json!({"role": "assistant", "content": "low"}),
];
let scores = vec![score(0, 0.9), score(1, 0.9), score(2, 0.1), score(3, 0.1)];
let cands = build_candidates(&msgs, &scores, &HashSet::new(), &[]);
// The low-score turn (idx 2,3) should come first.
assert_eq!(cands[0].kind, CandidateKind::Turn);
assert_eq!(cands[0].indices, vec![2, 3]);
}
#[test]
fn protected_messages_excluded() {
let msgs = vec![
json!({"role": "user", "content": "u"}),
json!({"role": "assistant", "content": "a"}),
];
let scores = vec![score(0, 0.5), score(1, 0.5)];
let mut protected = HashSet::new();
protected.insert(1); // protect the assistant
let cands = build_candidates(&msgs, &scores, &protected, &[]);
// Pair impossible (asst protected); user becomes single; asst skipped.
assert_eq!(cands.len(), 1);
assert_eq!(cands[0].kind, CandidateKind::Single);
assert_eq!(cands[0].indices, vec![0]);
}
#[test]
fn tool_unit_takes_precedence_over_singles() {
let msgs = vec![
json!({"role": "user", "content": "go"}),
json!({
"role": "assistant",
"content": "",
"tool_calls": [{"id": "c1", "function": {"name": "f"}}]
}),
json!({"role": "tool", "tool_call_id": "c1", "content": "r"}),
];
let scores = vec![score(0, 0.5), score(1, 0.2), score(2, 0.3)];
let units = find_tool_units(&msgs);
let cands = build_candidates(&msgs, &scores, &HashSet::new(), &units);
// Two candidates: the user (single) and the tool unit. The
// tool unit holds asst+tool atomically; neither asst nor tool
// appears as its own candidate.
let unit_cands: Vec<_> = cands
.iter()
.filter(|c| c.kind == CandidateKind::ToolUnit)
.collect();
assert_eq!(unit_cands.len(), 1);
let single_cands: Vec<_> = cands
.iter()
.filter(|c| c.kind == CandidateKind::Single)
.collect();
assert_eq!(single_cands.len(), 1);
assert_eq!(single_cands[0].indices, vec![0]);
}
#[test]
fn tool_unit_with_protected_response_is_skipped() {
let msgs = vec![
json!({
"role": "assistant",
"content": "",
"tool_calls": [{"id": "c1", "function": {"name": "f"}}]
}),
json!({"role": "tool", "tool_call_id": "c1", "content": "r"}),
];
let scores = vec![score(0, 0.2), score(1, 0.3)];
let mut protected = HashSet::new();
protected.insert(1); // protect the response
let units = find_tool_units(&msgs);
let cands = build_candidates(&msgs, &scores, &protected, &units);
// Unit can't drop atomically when one half is protected.
assert!(cands.iter().all(|c| c.kind != CandidateKind::ToolUnit));
}
#[test]
fn position_tiebreaker_for_equal_scores() {
let msgs = vec![
json!({"role": "user", "content": "first"}),
json!({"role": "assistant", "content": "first"}),
json!({"role": "user", "content": "second"}),
json!({"role": "assistant", "content": "second"}),
];
let scores = vec![score(0, 0.5), score(1, 0.5), score(2, 0.5), score(3, 0.5)];
let cands = build_candidates(&msgs, &scores, &HashSet::new(), &[]);
// Older turn drops first.
assert_eq!(cands[0].position, 0);
}
}

201
crates/headroom-core/src/context/ccr_drop.rs
View file

@ -1,201 +0,0 @@
//! CCR-on-drop — persist dropped messages so the model can retrieve
//! them later via a tool call.
//!
//! This is the OSS-defining behaviour. Without it, dropping a message
//! is destructive (≈ rolling window). With it, the dropped content
//! is parked in the [`CcrStore`] under a content-hash key and a marker
//! is inserted into the surviving message stream pointing at it.
//! When the model later calls `ccr_retrieve(<key>)`, the dropped JSON
//! comes back verbatim.
//!
//! Mirrors Python's `_store_dropped_in_ccr` (intelligent_context.py
//! ~L955) including the marker text format. Differences from Python:
//!
//! - Rust uses the trait-level [`CcrStore`] directly (no Python
//! `CompressionStore` indirection).
//! - The marker is returned to the caller for insertion; Python
//! handled insertion outside this helper too.
use std::sync::Arc;
use serde_json::Value;
use sha2::{Digest, Sha256};
use crate::ccr::CcrStore;
/// Result of stashing dropped messages into the CCR store.
#[derive(Debug, Clone)]
pub struct DropPersist {
/// CCR cache key the dropped JSON was stored under. The marker
/// references this so the model knows what to ask for.
pub key: String,
/// Human-readable marker string suitable for insertion as a
/// system or user message in the surviving conversation.
pub marker: String,
/// Number of messages persisted.
pub count: usize,
}
/// Serialize the messages at `dropped_indices` from `original_messages`
/// and stash them in `store` under a content-hash key. Returns metadata
/// describing the persistence so the caller can insert the marker into
/// the live message list.
///
/// Returns `None` when there's nothing to do (no indices or no store)
/// or when serialization fails — failure is non-fatal because the drop
/// itself still happened. The caller logs and continues.
pub fn persist_dropped(
original_messages: &[Value],
dropped_indices: &[usize],
store: &Arc<dyn CcrStore>,
) -> Option<DropPersist> {
if dropped_indices.is_empty() {
return None;
}
// Collect in original order — re-sort because the cascade may
// append in non-monotonic order.
let mut sorted = dropped_indices.to_vec();
sorted.sort_unstable();
sorted.dedup();
let dropped: Vec<&Value> = sorted
.iter()
.filter_map(|&i| original_messages.get(i))
.collect();
if dropped.is_empty() {
return None;
}
let dropped_json = serde_json::to_string_pretty(&dropped).ok()?;
// Hash + truncate to 12 hex chars — same convention as the rest
// of the CCR store (see ccr.rs and smart_crusher hashing).
let mut h = Sha256::new();
h.update(dropped_json.as_bytes());
let key = h
.finalize()
.iter()
.take(6)
.map(|b| format!("{b:02x}"))
.collect::<String>();
// Role-count summary mirrors Python's marker format so anyone
// grepping logs sees the same string shape across implementations.
let marker = build_marker(&dropped);
// Store the original (full JSON) under the key. The marker stays
// outside the store — it lives in the live conversation as a
// pointer; the store holds only the recoverable payload.
store.put(&key, &dropped_json);
Some(DropPersist {
key,
marker,
count: dropped.len(),
})
}
/// Build the human-readable marker string. Format matches Python's
/// `_store_dropped_in_ccr` so log greps work across implementations.
fn build_marker(dropped: &[&Value]) -> String {
use std::collections::BTreeMap;
let mut role_counts: BTreeMap<&str, usize> = BTreeMap::new();
for msg in dropped {
let role = msg.get("role").and_then(Value::as_str).unwrap_or("unknown");
*role_counts.entry(role).or_insert(0) += 1;
}
let parts: Vec<String> = role_counts
.iter()
.map(|(role, count)| format!("{count} {role}"))
.collect();
format!(
"[Dropped {} messages: {}. Use ccr_retrieve to access full content.]",
dropped.len(),
parts.join(", ")
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ccr::InMemoryCcrStore;
use serde_json::json;
fn store() -> Arc<dyn CcrStore> {
Arc::new(InMemoryCcrStore::new())
}
#[test]
fn persists_dropped_messages_under_content_hash() {
let originals = vec![
json!({"role": "system", "content": "sys"}),
json!({"role": "user", "content": "drop me"}),
json!({"role": "assistant", "content": "and me"}),
json!({"role": "user", "content": "keep me"}),
];
let dropped = vec![1, 2];
let s = store();
let result = persist_dropped(&originals, &dropped, &s).expect("should persist");
assert_eq!(result.count, 2);
assert_eq!(result.key.len(), 12);
assert!(result.marker.contains("Dropped 2 messages"));
assert!(result.marker.contains("ccr_retrieve"));
// Round-trip: the original JSON is recoverable via the same key.
let stored = s.get(&result.key).expect("should retrieve");
let recovered: Vec<Value> = serde_json::from_str(&stored).unwrap();
assert_eq!(recovered.len(), 2);
assert_eq!(recovered[0]["content"], "drop me");
}
#[test]
fn empty_indices_returns_none() {
let originals = vec![json!({"role": "user", "content": "x"})];
let s = store();
assert!(persist_dropped(&originals, &[], &s).is_none());
}
#[test]
fn out_of_range_indices_are_skipped() {
let originals = vec![json!({"role": "user", "content": "x"})];
let s = store();
// Index 99 is out of range — should be silently filtered, not
// panic. Index 0 valid, so we still get a persist.
let result = persist_dropped(&originals, &[0, 99], &s).expect("should persist idx 0");
assert_eq!(result.count, 1);
}
#[test]
fn marker_role_breakdown_is_sorted_alphabetical() {
// BTreeMap ordering → roles appear alphabetically. Stable
// marker text across runs.
let originals = vec![
json!({"role": "user", "content": "u"}),
json!({"role": "assistant", "content": "a"}),
json!({"role": "tool", "content": "t"}),
];
let s = store();
let result = persist_dropped(&originals, &[0, 1, 2], &s).expect("should persist");
// "1 assistant" appears before "1 tool" appears before "1 user".
let m = &result.marker;
let asst = m.find("assistant").unwrap();
let tool = m.find("tool").unwrap();
let user = m.find("user").unwrap();
assert!(asst < tool && tool < user);
}
#[test]
fn duplicate_indices_are_deduped() {
let originals = vec![
json!({"role": "user", "content": "x"}),
json!({"role": "assistant", "content": "y"}),
];
let s = store();
let result = persist_dropped(&originals, &[0, 0, 1, 1, 1], &s).expect("should persist");
// Counts the unique message set, not the input list length.
assert_eq!(result.count, 2);
}
}

90
crates/headroom-core/src/context/config.rs
View file

@ -1,90 +0,0 @@
//! `IcmConfig` — six-field config for the OSS context manager.
//!
//! Cut from the 12+ Python config: `compress_threshold`,
//! `summarize_threshold`, `summarization_*`, `memory_tiers_*`,
//! `warm_tier_*`, `cold_tier_*`. Those belong to strategies that don't
//! ship in OSS. `recency_decay_rate` and `toin_*` moved to
//! [`MessageScorer`](crate::scoring::MessageScorer) where they belong
//! semantically.
use serde::{Deserialize, Serialize};
use crate::scoring::ScoringWeights;
/// Configuration for [`IntelligentContextManager`](super::IntelligentContextManager).
///
/// Defaults are tuned for the OSS sweet spot: keep system messages and
/// the last 2 turns sacred, leave 4K tokens of headroom for the model's
/// reply, score by importance, persist drops to CCR so they're
/// retrievable.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct IcmConfig {
/// Master switch. When `false`, `should_apply` always returns
/// `false` regardless of token count.
pub enabled: bool,
/// If `true`, `role=system` messages are never droppable. Default
/// `true` — dropping system prompts breaks behavior in subtle
/// ways. Disable only if the caller manages system prompts
/// themselves and explicitly wants them in the candidate pool.
pub keep_system: bool,
/// Number of recent user turns to protect from dropping. A "turn"
/// here means a `role=user` message; everything from the last
/// `keep_last_turns`th user message to the end is protected,
/// including all assistant replies and tool exchanges in between.
pub keep_last_turns: usize,
/// Reserved tokens for the model's response. Effective budget is
/// `model_limit - output_buffer_tokens`. Default 4000 covers most
/// reasonable replies; bump for long-form generation.
pub output_buffer_tokens: usize,
/// Weights for the six-factor message scorer. See
/// [`ScoringWeights`] for the factor breakdown. Defaults match
/// Python's `ScoringWeights()` so existing tuning carries over.
pub scoring_weights: ScoringWeights,
/// When `true` (default), dropped messages are serialized into
/// the CCR store before removal — the model can retrieve them via
/// a tool call. This is the OSS-defining behaviour: with
/// `ccr_on_drop=true`, drop ≈ "moved to retrievable cache"; with
/// `false`, drop ≈ rolling window.
pub ccr_on_drop: bool,
}
impl Default for IcmConfig {
fn default() -> Self {
Self {
enabled: true,
keep_system: true,
keep_last_turns: 2,
output_buffer_tokens: 4000,
scoring_weights: ScoringWeights::default(),
ccr_on_drop: true,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn defaults_match_oss_intent() {
let c = IcmConfig::default();
assert!(c.enabled);
assert!(c.keep_system);
assert_eq!(c.keep_last_turns, 2);
assert_eq!(c.output_buffer_tokens, 4000);
assert!(c.ccr_on_drop);
}
#[test]
fn round_trips_through_serde() {
let c = IcmConfig::default();
let json = serde_json::to_string(&c).unwrap();
let back: IcmConfig = serde_json::from_str(&json).unwrap();
assert_eq!(c, back);
}
}

387
crates/headroom-core/src/context/manager.rs
View file

@ -1,387 +0,0 @@
//! `IntelligentContextManager` — the cascade orchestrator.
//!
//! Holds the config, the safety-rails computer, and the registered
//! list of strategies. On `apply()`, it:
//!
//! 1. Tokenizes, exits early if under budget.
//! 2. Computes safety-rail protections (system, last-N-turns, frozen
//! prefix, paired tool responses).
//! 3. Builds a [`ContextWorkspace`] and walks each registered strategy
//! in registration order until a strategy reports `fully_resolved`
//! or the list is exhausted.
//! 4. Returns the (possibly mutated) message list plus an
//! [`ApplyResult`] describing what happened.
//!
//! OSS pre-registers exactly one strategy: [`DropByScoreStrategy`].
//! Enterprise calls `with_strategy` more times.
use std::sync::Arc;
use serde_json::Value;
use crate::ccr::CcrStore;
use crate::context::config::IcmConfig;
use crate::context::safety::SafetyRails;
use crate::context::strategy::{ContextStrategy, DropByScoreStrategy};
use crate::context::workspace::ContextWorkspace;
use crate::scoring::MessageScorer;
use crate::tokenizer::Tokenizer;
/// Per-call inputs that vary by request.
pub struct ApplyCtx {
/// Provider's context-window size in tokens.
pub model_limit: usize,
/// Override the manager's `output_buffer_tokens`. `None` uses the
/// configured default.
pub output_buffer: Option<usize>,
/// Number of leading messages that are part of the provider's
/// prompt cache. Always protected. Provider-aware caller computes
/// this; ICM just receives it.
pub frozen_message_count: usize,
}
impl Default for ApplyCtx {
fn default() -> Self {
Self {
model_limit: 128_000,
output_buffer: None,
frozen_message_count: 0,
}
}
}
/// Outcome of a single `apply()` call.
#[derive(Debug, Clone)]
pub struct ApplyResult {
/// The (possibly modified) message list.
pub messages: Vec<Value>,
pub tokens_before: usize,
pub tokens_after: usize,
/// Names of strategies that ran (in order). Strategies that were
/// short-circuited by an earlier `fully_resolved=true` don't appear.
pub strategies_applied: Vec<&'static str>,
/// Marker strings emitted by the cascade (e.g. CCR retrieval hints).
pub markers_inserted: Vec<String>,
}
/// The orchestrator. Constructed once per process and reused across
/// requests; `Send + Sync`.
pub struct IntelligentContextManager {
config: IcmConfig,
tokenizer: Arc<dyn Tokenizer>,
strategies: Vec<Box<dyn ContextStrategy>>,
}
impl IntelligentContextManager {
/// Construct with the OSS default strategy stack: just
/// [`DropByScoreStrategy`]. The scorer is built from the config's
/// `scoring_weights`. Enterprise can call `with_strategy` to
/// register additional strategies before serving requests.
///
/// `ccr` is the CCR store used for drop persistence (and is also
/// stored on the strategy). Pass `None` to disable CCR-on-drop
/// even if `config.ccr_on_drop = true`.
pub fn new(
config: IcmConfig,
tokenizer: Arc<dyn Tokenizer>,
ccr: Option<Arc<dyn CcrStore>>,
) -> Self {
let scorer = Arc::new(MessageScorer::new(
Some(config.scoring_weights),
None,
None,
0.1,
));
let drop_strategy: Box<dyn ContextStrategy> = Box::new(DropByScoreStrategy::new(
scorer,
tokenizer.clone(),
ccr,
config.ccr_on_drop,
));
Self {
config,
tokenizer,
strategies: vec![drop_strategy],
}
}
/// Append an Enterprise strategy. Strategies run in registration
/// order; the OSS `DropByScoreStrategy` is always first unless the
/// caller explicitly bypasses it via [`Self::without_default_strategy`].
pub fn with_strategy(mut self, strategy: Box<dyn ContextStrategy>) -> Self {
self.strategies.push(strategy);
self
}
/// Drop the OSS default strategy from the stack. Use only when
/// the caller wants a fully custom strategy chain — e.g. an
/// Enterprise build that handles dropping itself.
pub fn without_default_strategy(mut self) -> Self {
// Removes the default DropByScoreStrategy registered in `new`.
self.strategies.clear();
self
}
pub fn config(&self) -> &IcmConfig {
&self.config
}
/// Cheap pre-check. `true` means `apply()` would do work; `false`
/// means the request is under budget and a no-op pass-through is
/// safe (and faster — skips tokenization on the hot path).
pub fn should_apply(
&self,
messages: &[Value],
model_limit: usize,
output_buffer: usize,
) -> bool {
if !self.config.enabled {
return false;
}
let current = count_messages(messages, self.tokenizer.as_ref());
let available = model_limit.saturating_sub(output_buffer);
current > available
}
/// Run the cascade.
pub fn apply(&self, messages: Vec<Value>, ctx: ApplyCtx) -> ApplyResult {
let output_buffer = ctx
.output_buffer
.unwrap_or(self.config.output_buffer_tokens);
let target = ctx.model_limit.saturating_sub(output_buffer);
let tokens_before = count_messages(&messages, self.tokenizer.as_ref());
// Early exit: under budget OR disabled.
if !self.config.enabled || tokens_before <= target {
return ApplyResult {
messages,
tokens_before,
tokens_after: tokens_before,
strategies_applied: Vec::new(),
markers_inserted: Vec::new(),
};
}
// Compute safety rails + build workspace.
let safety = SafetyRails::new(&self.config);
let protected = safety.protected(&messages, ctx.frozen_message_count);
let mut ws = ContextWorkspace::new(messages, protected, ctx.frozen_message_count);
ws.current_tokens = tokens_before;
let mut strategies_applied: Vec<&'static str> = Vec::new();
let mut markers: Vec<String> = Vec::new();
for strategy in &self.strategies {
let outcome = strategy.try_fit(&mut ws, target);
strategies_applied.push(strategy.name());
markers.extend(outcome.markers_inserted);
// Recompute protections — indices may have shifted post-drop.
// Subsequent strategies need a fresh view.
ws.protected = safety.protected(&ws.messages, ws.frozen_count);
if outcome.fully_resolved {
break;
}
}
ApplyResult {
tokens_before,
tokens_after: ws.current_tokens,
messages: ws.messages,
strategies_applied,
markers_inserted: markers,
}
}
}
// Re-export the strategy's count_messages for tests + manager-level
// gating. Keeping a single source of truth avoids drift between
// `should_apply` and the strategy's own accounting.
pub(crate) use crate::context::strategy::drop_by_score::count_messages;
#[cfg(test)]
mod tests {
use super::*;
use crate::ccr::InMemoryCcrStore;
use crate::tokenizer::EstimatingCounter;
use serde_json::json;
fn manager() -> IntelligentContextManager {
let tokenizer: Arc<dyn Tokenizer> = Arc::new(EstimatingCounter::default());
let store: Arc<dyn CcrStore> = Arc::new(InMemoryCcrStore::new());
IntelligentContextManager::new(IcmConfig::default(), tokenizer, Some(store))
}
#[test]
fn should_apply_returns_false_when_disabled() {
let cfg = IcmConfig {
enabled: false,
..Default::default()
};
let tokenizer: Arc<dyn Tokenizer> = Arc::new(EstimatingCounter::default());
let m = IntelligentContextManager::new(cfg, tokenizer, None);
let huge = vec![json!({"role": "user", "content": "x".repeat(1_000_000)})];
assert!(!m.should_apply(&huge, 1000, 100));
}
#[test]
fn should_apply_false_under_budget() {
let m = manager();
let small = vec![json!({"role": "user", "content": "hi"})];
assert!(!m.should_apply(&small, 128_000, 4_000));
}
#[test]
fn should_apply_true_over_budget() {
let m = manager();
let huge = vec![json!({"role": "user", "content": "x".repeat(1_000_000)})];
assert!(m.should_apply(&huge, 1000, 100));
}
#[test]
fn apply_under_budget_is_passthrough() {
let m = manager();
let msgs = vec![json!({"role": "user", "content": "hello"})];
let n = msgs.len();
let r = m.apply(
msgs,
ApplyCtx {
model_limit: 128_000,
output_buffer: Some(4_000),
frozen_message_count: 0,
},
);
assert_eq!(r.messages.len(), n);
assert!(r.strategies_applied.is_empty());
assert_eq!(r.tokens_before, r.tokens_after);
}
#[test]
fn apply_over_budget_runs_drop_by_score() {
let m = manager();
// Many large-ish messages; tight budget forces drops. Last
// turn is protected by default config (keep_last_turns=2).
let mut msgs: Vec<Value> = (0..10)
.map(|i| json!({"role": "user", "content": format!("message {i} ").repeat(50)}))
.collect();
msgs.push(json!({"role": "assistant", "content": "ack"}));
msgs.push(json!({"role": "user", "content": "final"}));
let initial = count_messages(&msgs, &EstimatingCounter::default());
let r = m.apply(
msgs.clone(),
ApplyCtx {
model_limit: initial / 2,
output_buffer: Some(0),
frozen_message_count: 0,
},
);
assert_eq!(r.strategies_applied, vec!["drop_by_score"]);
assert!(r.tokens_after < r.tokens_before);
assert!(r.messages.len() < msgs.len());
// Final user message ("final") is protected; it must survive.
let last = r.messages.last().unwrap();
assert_eq!(last["content"], "final");
}
#[test]
fn frozen_prefix_is_never_dropped() {
let m = manager();
let mut msgs: Vec<Value> = Vec::new();
msgs.push(json!({"role": "user", "content": "FROZEN PREFIX MARKER".repeat(50)}));
for i in 0..15 {
msgs.push(json!({"role": "user", "content": format!("filler {i} ").repeat(50)}));
}
let initial = count_messages(&msgs, &EstimatingCounter::default());
let r = m.apply(
msgs,
ApplyCtx {
model_limit: initial / 4,
output_buffer: Some(0),
frozen_message_count: 1,
},
);
// The first message (frozen) must still be present.
assert!(r.messages[0]["content"]
.as_str()
.unwrap()
.contains("FROZEN PREFIX MARKER"));
}
#[test]
fn system_message_protected_by_default() {
let m = manager();
let mut msgs: Vec<Value> = Vec::new();
msgs.push(json!({"role": "system", "content": "you are helpful"}));
for i in 0..15 {
msgs.push(json!({"role": "user", "content": format!("filler {i} ").repeat(50)}));
}
let initial = count_messages(&msgs, &EstimatingCounter::default());
let r = m.apply(
msgs,
ApplyCtx {
model_limit: initial / 3,
output_buffer: Some(0),
frozen_message_count: 0,
},
);
// System message present somewhere.
assert!(r
.messages
.iter()
.any(|m| m.get("role").and_then(Value::as_str) == Some("system")));
}
#[test]
fn with_strategy_appends_to_chain() {
struct Sentinel;
impl ContextStrategy for Sentinel {
fn name(&self) -> &'static str {
"sentinel"
}
fn try_fit(
&self,
_ws: &mut ContextWorkspace,
_target: usize,
) -> crate::context::workspace::StrategyOutcome {
crate::context::workspace::StrategyOutcome::default()
}
}
let m = manager().with_strategy(Box::new(Sentinel));
// Force a drop scenario.
let msgs: Vec<Value> = (0..20)
.map(|i| json!({"role": "user", "content": format!("m{i} ").repeat(50)}))
.collect();
let initial = count_messages(&msgs, &EstimatingCounter::default());
let r = m.apply(
msgs,
ApplyCtx {
model_limit: initial / 4,
output_buffer: Some(0),
frozen_message_count: 0,
},
);
// drop_by_score may or may not fully resolve. If it does,
// sentinel is short-circuited (no entry). If not, sentinel
// appears second.
assert_eq!(r.strategies_applied[0], "drop_by_score");
}
#[test]
fn ccr_marker_emitted_on_drop() {
let m = manager();
let msgs: Vec<Value> = (0..20)
.map(|i| json!({"role": "user", "content": format!("m{i} ").repeat(50)}))
.collect();
let initial = count_messages(&msgs, &EstimatingCounter::default());
let r = m.apply(
msgs,
ApplyCtx {
model_limit: initial / 4,
output_buffer: Some(0),
frozen_message_count: 0,
},
);
assert!(!r.markers_inserted.is_empty());
assert!(r.markers_inserted[0].contains("ccr_retrieve"));
}
}

47
crates/headroom-core/src/context/mod.rs
View file

@ -1,47 +0,0 @@
//! Conversation-level context management for the OSS proxy.
//!
//! `context/` lives at the crate root parallel to `scoring/`,
//! `signals/`, and `transforms/` because it operates on a *different*
//! abstraction than any of them:
//!
//! - `transforms/` and `signals/` work on **content** (a string, a
//! structured JSON blob, a tool-result body).
//! - `scoring/` scores **messages** for importance.
//! - `context/` orchestrates **conversations** — it decides what to
//! do when a request's message list is over the model's context
//! budget. It uses `scoring/` for per-message scores and the CCR
//! store for drop persistence, but it is itself a higher-level
//! layer than either.
//!
//! # OSS philosophy
//!
//! OSS ships exactly one strategy: [`DropByScoreStrategy`] —
//! multi-factor scoring + safety rails + CCR-on-drop persistence. This
//! alone outclasses every gateway competitor's rolling-window
//! behaviour. Enterprise plugs additional strategies (compress-first,
//! summarize, memory tiers) in via [`ContextStrategy`] without any
//! orchestrator changes.
//!
//! # Why simpler than the Python original
//!
//! The Python `IntelligentContextManager` shipped three cascading
//! strategies (COMPRESS_FIRST, SUMMARIZE, DROP_BY_SCORE) plus 12+
//! config fields and integration hooks for ContentRouter,
//! ProgressiveSummarizer, TOIN, and the CCR store. Two of those
//! strategies belong to value-adds we're scoping to Enterprise; one
//! (memory tiers) was never implemented at all. This module is the
//! deliberate refactor: keep the OSS-defining behaviour, push the
//! rest behind a trait extension point.
pub mod candidate;
pub mod ccr_drop;
pub mod config;
pub mod manager;
pub mod safety;
pub mod strategy;
pub mod workspace;
pub use config::IcmConfig;
pub use manager::{ApplyCtx, ApplyResult, IntelligentContextManager};
pub use strategy::{ContextStrategy, DropByScoreStrategy};
pub use workspace::{ContextWorkspace, StrategyOutcome};

357
crates/headroom-core/src/context/safety.rs
View file

@ -1,357 +0,0 @@
//! `SafetyRails` — compute the set of indices that strategies must
//! never drop.
//!
//! Direct port of Python's `IntelligentContextManager._get_protected_indices`,
//! including the awkward-but-correct two-pass walk:
//!
//! 1. Walk forward, mark `role=system` (if `keep_system`).
//! 2. Walk backward from the end, mark messages until `keep_last_turns`
//! user messages have been seen.
//! 3. For every assistant message in the protected set, find its tool
//! responses (OpenAI `role=tool` + `tool_call_id`, OR Anthropic
//! `role=user` with `content[].type=tool_result`) and protect those
//! too. Otherwise we'd ship an orphan tool_call → 400 from the API.
//!
//! Frozen prefix is added on top by the caller via
//! [`ContextWorkspace::frozen_count`](super::workspace::ContextWorkspace::frozen_count).
//! This module just handles the per-message rules.
use std::collections::HashSet;
use serde_json::Value;
use crate::context::config::IcmConfig;
pub struct SafetyRails<'c> {
config: &'c IcmConfig,
}
impl<'c> SafetyRails<'c> {
pub fn new(config: &'c IcmConfig) -> Self {
Self { config }
}
/// Compute protected indices for the given message list.
///
/// Includes the leading `frozen_count` indices (prompt-cache prefix)
/// in the result so callers don't have to remember to OR them in.
pub fn protected(&self, messages: &[Value], frozen_count: usize) -> HashSet<usize> {
let mut protected: HashSet<usize> = HashSet::new();
// Frozen prefix — provider's prompt cache. Dropping any of
// these busts the cache and raises latency + cost.
for i in 0..frozen_count.min(messages.len()) {
protected.insert(i);
}
// System messages — never droppable when keep_system is on.
// Python convention: any `role=system` anywhere in the list,
// not only at index 0.
if self.config.keep_system {
for (i, msg) in messages.iter().enumerate() {
if msg.get("role").and_then(Value::as_str) == Some("system") {
protected.insert(i);
}
}
}
// Last N user-turn boundary. Walk backward from the end,
// protecting every message until we've seen `keep_last_turns`
// role=user messages. Includes assistant + tool messages
// interleaved with those user turns.
if self.config.keep_last_turns > 0 && !messages.is_empty() {
let mut turns_seen = 0;
let mut i = messages.len() as isize - 1;
while i >= 0 && turns_seen < self.config.keep_last_turns {
let msg = &messages[i as usize];
protected.insert(i as usize);
if msg.get("role").and_then(Value::as_str) == Some("user") {
turns_seen += 1;
}
i -= 1;
}
}
// Pair tool responses with their assistant tool_call messages.
// Protected assistant → every tool response that satisfies
// its tool_call ids must also be protected. Otherwise we'd
// drop a `tool` message whose corresponding `assistant`
// tool_call is still in the prompt → API rejects with 400.
let pairs = self.tool_response_pairs(messages, &protected);
protected.extend(pairs);
protected
}
/// For each protected assistant message, find every tool response
/// (OpenAI or Anthropic shape) that references one of its tool
/// call ids. Returns the response indices — the assistant indices
/// are already in `protected`.
fn tool_response_pairs(
&self,
messages: &[Value],
protected: &HashSet<usize>,
) -> HashSet<usize> {
let mut response_indices: HashSet<usize> = HashSet::new();
for &i in protected {
let msg = match messages.get(i) {
Some(m) => m,
None => continue,
};
if msg.get("role").and_then(Value::as_str) != Some("assistant") {
continue;
}
let tool_call_ids = collect_assistant_tool_call_ids(msg);
if tool_call_ids.is_empty() {
continue;
}
for (j, other) in messages.iter().enumerate() {
if protected.contains(&j) {
continue;
}
if message_responds_to_tool_call(other, &tool_call_ids) {
response_indices.insert(j);
}
}
}
response_indices
}
}
/// Extract the set of tool_call ids announced by an assistant message,
/// handling both OpenAI and Anthropic shapes.
///
/// - **OpenAI**: `{role:assistant, tool_calls: [{id, ...}, ...]}`
/// - **Anthropic**: `{role:assistant, content: [{type:tool_use, id, ...}]}`
fn collect_assistant_tool_call_ids(assistant: &Value) -> HashSet<String> {
let mut ids: HashSet<String> = HashSet::new();
// OpenAI: tool_calls array
if let Some(arr) = assistant.get("tool_calls").and_then(Value::as_array) {
for tc in arr {
if let Some(id) = tc.get("id").and_then(Value::as_str) {
ids.insert(id.to_string());
}
}
}
// Anthropic: content blocks with type=tool_use
if let Some(blocks) = assistant.get("content").and_then(Value::as_array) {
for block in blocks {
if block.get("type").and_then(Value::as_str) == Some("tool_use") {
if let Some(id) = block.get("id").and_then(Value::as_str) {
ids.insert(id.to_string());
}
}
}
}
ids
}
/// Does this message respond to one of `ids`?
///
/// - **OpenAI**: `{role:tool, tool_call_id: "..."}`
/// - **Anthropic**: `{role:user, content: [{type:tool_result, tool_use_id: "..."}]}`
fn message_responds_to_tool_call(msg: &Value, ids: &HashSet<String>) -> bool {
let role = msg.get("role").and_then(Value::as_str);
// OpenAI shape
if role == Some("tool") {
if let Some(tcid) = msg.get("tool_call_id").and_then(Value::as_str) {
return ids.contains(tcid);
}
}
// Anthropic shape — tool results live inside a user message's content list
if role == Some("user") {
if let Some(blocks) = msg.get("content").and_then(Value::as_array) {
for block in blocks {
if block.get("type").and_then(Value::as_str) == Some("tool_result") {
if let Some(tuid) = block.get("tool_use_id").and_then(Value::as_str) {
if ids.contains(tuid) {
return true;
}
}
}
}
}
}
false
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn cfg(keep_system: bool, keep_last_turns: usize) -> IcmConfig {
IcmConfig {
keep_system,
keep_last_turns,
..Default::default()
}
}
#[test]
fn protects_system_messages() {
let msgs = vec![
json!({"role": "system", "content": "you are helpful"}),
json!({"role": "user", "content": "hi"}),
json!({"role": "assistant", "content": "hello"}),
];
let c = cfg(true, 0);
let s = SafetyRails::new(&c);
let p = s.protected(&msgs, 0);
assert!(p.contains(&0));
assert!(!p.contains(&1));
assert!(!p.contains(&2));
}
#[test]
fn keep_system_false_does_not_protect_system() {
let msgs = vec![
json!({"role": "system", "content": "sys"}),
json!({"role": "user", "content": "hi"}),
];
let c = cfg(false, 0);
let s = SafetyRails::new(&c);
let p = s.protected(&msgs, 0);
assert!(!p.contains(&0));
}
#[test]
fn last_n_turns_protects_user_and_inner_messages() {
// 5 messages, keep_last_turns=2 → walk back until we've seen
// 2 user messages. Should protect indices 2,3,4 (covers user@4,
// asst@3, user@2).
let msgs = vec![
json!({"role": "user", "content": "q1"}),
json!({"role": "assistant", "content": "a1"}),
json!({"role": "user", "content": "q2"}),
json!({"role": "assistant", "content": "a2"}),
json!({"role": "user", "content": "q3"}),
];
let c = cfg(false, 2);
let s = SafetyRails::new(&c);
let p = s.protected(&msgs, 0);
assert!(p.contains(&2));
assert!(p.contains(&3));
assert!(p.contains(&4));
assert!(!p.contains(&0));
assert!(!p.contains(&1));
}
#[test]
fn frozen_prefix_is_protected() {
let msgs = vec![
json!({"role": "user", "content": "1"}),
json!({"role": "user", "content": "2"}),
json!({"role": "user", "content": "3"}),
];
let c = cfg(false, 0);
let s = SafetyRails::new(&c);
let p = s.protected(&msgs, 2);
assert!(p.contains(&0));
assert!(p.contains(&1));
assert!(!p.contains(&2));
}
#[test]
fn frozen_count_above_message_count_is_clamped() {
let msgs = vec![json!({"role": "user", "content": "x"})];
let c = cfg(false, 0);
let s = SafetyRails::new(&c);
// No panic, just protect what's there.
let p = s.protected(&msgs, 99);
assert_eq!(p.len(), 1);
}
#[test]
fn openai_tool_pair_is_atomic() {
// assistant @ idx 1 has a tool_call; tool response @ idx 2.
// keep_last_turns protects from the end. With keep_last=1 we
// should pull in user@3 only, then the pair logic adds nothing.
// Use keep_last=0 + keep_system=false but pre-protect the
// assistant manually (simulate a caller-protected index).
let msgs = vec![
json!({"role": "user", "content": "go"}),
json!({
"role": "assistant",
"content": "",
"tool_calls": [{"id": "call_1", "type": "function",
"function": {"name": "f", "arguments": "{}"}}]
}),
json!({"role": "tool", "tool_call_id": "call_1", "content": "result"}),
json!({"role": "user", "content": "thanks"}),
];
// keep_last_turns=2 → protects user@3 + asst@2 (but the asst@2
// is the tool message... no wait, asst is at 1, tool at 2).
// Walking back from idx 3: protect 3 (user, 1 turn), 2 (tool),
// 1 (asst), 0 (user, 2 turns). All four protected.
let c = cfg(false, 2);
let s = SafetyRails::new(&c);
let p = s.protected(&msgs, 0);
assert!(p.contains(&1));
assert!(p.contains(&2)); // The tool response must be protected too.
}
#[test]
fn anthropic_tool_pair_is_atomic() {
// Anthropic: assistant content has tool_use blocks; tool_result
// lives inside a user message's content list.
let msgs = vec![
json!({"role": "user", "content": "go"}),
json!({
"role": "assistant",
"content": [
{"type": "text", "text": "thinking..."},
{"type": "tool_use", "id": "tu_1", "name": "f", "input": {}}
]
}),
json!({
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "tu_1", "content": "ok"}]
}),
json!({"role": "user", "content": "next"}),
];
// Pre-protect the assistant via keep_last_turns=2 (walks back
// through user@3, user@2 = 2 turns; assistant@1 not protected
// by turn rule). Bump to keep_last_turns=3 so assistant@1 falls
// under the protection.
let c = cfg(false, 3);
let s = SafetyRails::new(&c);
let p = s.protected(&msgs, 0);
assert!(p.contains(&1)); // assistant
assert!(p.contains(&2)); // tool_result user
}
#[test]
fn unmatched_tool_call_id_does_not_inflate_protection() {
let msgs = vec![
json!({"role": "user", "content": "go"}),
json!({
"role": "assistant",
"content": "",
"tool_calls": [{"id": "call_orphan",
"function": {"name": "f"}}]
}),
json!({"role": "tool", "tool_call_id": "call_DIFFERENT", "content": "?"}),
];
let c = cfg(false, 0);
let s = SafetyRails::new(&c);
// Manually protect the assistant by setting keep_last_turns
// to wrap it. But keep_last=0 here so nothing is protected;
// the orphan tool message stays unprotected. Now protect via
// keep_last_turns=1 wrapping back from end…
// Actually with keep_last=0 + keep_system=false there's no
// protection at all, so the pair logic finds nothing to pair.
let p = s.protected(&msgs, 0);
assert!(p.is_empty());
}
}

410
crates/headroom-core/src/context/strategy/drop_by_score.rs
View file

@ -1,410 +0,0 @@
//! `DropByScoreStrategy` — the OSS-shipped context strategy.
//!
//! Algorithm:
//!
//! 1. Build candidates (turns + singles + tool units), sorted by
//! importance score ascending.
//! 2. Walk the candidate list, dropping until token budget is met.
//! 3. After each drop, recompute token count via the supplied
//! tokenizer. Stop early once `current_tokens <= target`.
//! 4. If `ccr_on_drop=true` AND a CCR store is wired, persist the
//! dropped messages and emit a marker so the model can retrieve
//! them via tool call.
//!
//! Direct port of the Python `_apply_drop_by_score` path with the
//! same ordering semantics (lowest score first, position tiebreaker).
use std::sync::Arc;
use serde_json::Value;
use crate::ccr::CcrStore;
use crate::context::candidate::{build_candidates, find_tool_units};
use crate::context::ccr_drop::persist_dropped;
use crate::context::workspace::{ContextWorkspace, StrategyOutcome};
use crate::scoring::MessageScorer;
use crate::tokenizer::Tokenizer;
use super::ContextStrategy;
/// Drop-by-importance-score strategy.
///
/// Holds a [`MessageScorer`] (constructed once with the manager's
/// scoring weights) and an optional CCR store for persistence on drop.
/// Both are `Arc`-wrapped because the strategy is shared across
/// requests in a multi-threaded proxy.
pub struct DropByScoreStrategy {
scorer: Arc<MessageScorer>,
tokenizer: Arc<dyn Tokenizer>,
ccr_store: Option<Arc<dyn CcrStore>>,
ccr_on_drop: bool,
}
impl DropByScoreStrategy {
pub fn new(
scorer: Arc<MessageScorer>,
tokenizer: Arc<dyn Tokenizer>,
ccr_store: Option<Arc<dyn CcrStore>>,
ccr_on_drop: bool,
) -> Self {
Self {
scorer,
tokenizer,
ccr_store,
ccr_on_drop,
}
}
}
impl ContextStrategy for DropByScoreStrategy {
fn name(&self) -> &'static str {
"drop_by_score"
}
fn try_fit(&self, ws: &mut ContextWorkspace, target_tokens: usize) -> StrategyOutcome {
if ws.current_tokens <= target_tokens {
return StrategyOutcome {
fully_resolved: true,
..Default::default()
};
}
// Compute tool units (atomic groups) and message scores.
let tool_units = find_tool_units(&ws.messages);
let tool_unit_indices: std::collections::HashSet<usize> = tool_units
.iter()
.flat_map(|(a, rs)| std::iter::once(*a).chain(rs.iter().copied()))
.collect();
let scores = self
.scorer
.score_messages(&ws.messages, &ws.protected, &tool_unit_indices);
let candidates = build_candidates(&ws.messages, &scores, &ws.protected, &tool_units);
// Walk lowest → highest score, dropping one candidate at a
// time. Recompute tokens after each drop. Stop when under
// budget OR we run out of candidates.
//
// Track which ORIGINAL indices were dropped so the CCR helper
// can serialize the verbatim bytes. The workspace already
// holds `original_messages` for that purpose.
let initial_tokens = ws.current_tokens;
let mut dropped_now: Vec<usize> = Vec::new();
let mut indices_to_remove: std::collections::BTreeSet<usize> =
std::collections::BTreeSet::new();
for cand in &candidates {
if ws.current_tokens.saturating_sub(estimate_dropped_tokens(
&ws.messages,
&cand.indices,
self.tokenizer.as_ref(),
)) <= target_tokens
|| ws.current_tokens > target_tokens
{
indices_to_remove.extend(cand.indices.iter().copied());
dropped_now.extend(cand.indices.iter().copied());
// Estimate post-drop token count without rebuilding the
// list each iteration. Real recompute happens below.
let saved =
estimate_dropped_tokens(&ws.messages, &cand.indices, self.tokenizer.as_ref());
ws.current_tokens = ws.current_tokens.saturating_sub(saved);
if ws.current_tokens <= target_tokens {
break;
}
}
}
if indices_to_remove.is_empty() {
return StrategyOutcome::default();
}
// Translate WORKSPACE indices to ORIGINAL indices for CCR
// persistence. The workspace's `messages` and `original_messages`
// are aligned 1:1 in this implementation (no prior strategy
// mutated the list — this is the only OSS strategy).
ws.dropped_indices.extend(dropped_now);
// Remove highest index first so earlier indices stay valid.
let to_remove_desc: Vec<usize> = indices_to_remove.iter().rev().copied().collect();
for idx in to_remove_desc {
if idx < ws.messages.len() {
ws.messages.remove(idx);
}
}
// Recompute exact token count post-mutation. The estimate above
// is approximate; the real count drives the cascade decision.
ws.current_tokens = count_messages(&ws.messages, self.tokenizer.as_ref());
// CCR persistence.
let mut markers: Vec<String> = Vec::new();
if self.ccr_on_drop {
if let Some(store) = &self.ccr_store {
if let Some(persist) =
persist_dropped(&ws.original_messages, &ws.dropped_indices, store)
{
markers.push(persist.marker);
}
}
}
let tokens_freed = initial_tokens.saturating_sub(ws.current_tokens);
StrategyOutcome {
tokens_freed,
markers_inserted: markers,
fully_resolved: ws.current_tokens <= target_tokens,
}
}
}
/// Token count for a list of messages. Mirrors Python's
/// `count_messages` accounting: per-message overhead + content
/// tokens + role tokens. Matches the OpenAI chat-format formula
/// used in the Python `tiktoken_counter`.
///
/// `pub(crate)` so the manager module can use the same accounting
/// for `should_apply` gating without duplicating the logic.
pub(crate) fn count_messages(messages: &[Value], tokenizer: &dyn Tokenizer) -> usize {
const MESSAGE_OVERHEAD: usize = 3; // OpenAI: 3 tokens per message
const REPLY_PRIMER: usize = 3; // Final assistant primer
let mut total = REPLY_PRIMER;
for msg in messages {
total += MESSAGE_OVERHEAD;
if let Some(role) = msg.get("role").and_then(Value::as_str) {
total += tokenizer.count_text(role);
}
if let Some(content) = msg.get("content") {
total += count_content(content, tokenizer);
}
if let Some(name) = msg.get("name").and_then(Value::as_str) {
total += tokenizer.count_text(name);
}
// tool_calls and tool_call_id contribute too — count their
// string forms. Coarse but sufficient for budget gating.
if let Some(tc) = msg.get("tool_calls") {
if let Ok(s) = serde_json::to_string(tc) {
total += tokenizer.count_text(&s);
}
}
if let Some(id) = msg.get("tool_call_id").and_then(Value::as_str) {
total += tokenizer.count_text(id);
}
}
total
}
fn count_content(content: &Value, tokenizer: &dyn Tokenizer) -> usize {
match content {
Value::String(s) => tokenizer.count_text(s),
Value::Array(blocks) => {
let mut total = 0;
for block in blocks {
// {type:"text",text:"..."} or {type:"tool_result",content:"..."} or
// {type:"tool_use", input:{...}} etc. Walk text-bearing fields.
if let Some(text) = block.get("text").and_then(Value::as_str) {
total += tokenizer.count_text(text);
}
if let Some(c) = block.get("content") {
if let Some(s) = c.as_str() {
total += tokenizer.count_text(s);
} else if let Ok(s) = serde_json::to_string(c) {
total += tokenizer.count_text(&s);
}
}
if let Some(input) = block.get("input") {
if let Ok(s) = serde_json::to_string(input) {
total += tokenizer.count_text(&s);
}
}
}
total
}
Value::Null => 0,
other => serde_json::to_string(other)
.map(|s| tokenizer.count_text(&s))
.unwrap_or(0),
}
}
/// Approximate the token cost of dropping `indices` from `messages`.
/// Used as a fast pre-check inside the candidate walk; the real
/// recount happens after mutation.
fn estimate_dropped_tokens(
messages: &[Value],
indices: &[usize],
tokenizer: &dyn Tokenizer,
) -> usize {
const MESSAGE_OVERHEAD: usize = 3;
let mut total = 0;
for &i in indices {
if let Some(msg) = messages.get(i) {
total += MESSAGE_OVERHEAD;
if let Some(role) = msg.get("role").and_then(Value::as_str) {
total += tokenizer.count_text(role);
}
if let Some(content) = msg.get("content") {
total += count_content(content, tokenizer);
}
}
}
total
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ccr::InMemoryCcrStore;
use crate::scoring::MessageScorer;
use crate::tokenizer::EstimatingCounter;
use serde_json::json;
use std::collections::HashSet;
fn ws(messages: Vec<Value>, tokenizer: &dyn Tokenizer) -> ContextWorkspace {
let mut w = ContextWorkspace::new(messages, HashSet::new(), 0);
w.current_tokens = count_messages(&w.messages, tokenizer);
w
}
fn strategy(
ccr: Option<Arc<dyn CcrStore>>,
ccr_on_drop: bool,
tokenizer: Arc<dyn Tokenizer>,
) -> DropByScoreStrategy {
let scorer = Arc::new(MessageScorer::with_defaults());
DropByScoreStrategy::new(scorer, tokenizer, ccr, ccr_on_drop)
}
#[test]
fn under_budget_returns_fully_resolved_no_op() {
let tokenizer: Arc<dyn Tokenizer> = Arc::new(EstimatingCounter::default());
let mut w = ws(
vec![
json!({"role": "user", "content": "short"}),
json!({"role": "assistant", "content": "ok"}),
],
tokenizer.as_ref(),
);
let s = strategy(None, false, tokenizer.clone());
let out = s.try_fit(&mut w, 1_000_000);
assert!(out.fully_resolved);
assert_eq!(out.tokens_freed, 0);
assert_eq!(w.messages.len(), 2);
}
#[test]
fn drops_lowest_scored_messages_first() {
let tokenizer: Arc<dyn Tokenizer> = Arc::new(EstimatingCounter::default());
let messages = vec![
json!({"role": "user", "content": "old turn one ".repeat(20)}),
json!({"role": "assistant", "content": "old answer one ".repeat(20)}),
json!({"role": "user", "content": "recent turn ".repeat(20)}),
json!({"role": "assistant", "content": "recent answer ".repeat(20)}),
];
let mut w = ws(messages, tokenizer.as_ref());
let initial = w.current_tokens;
let target = initial / 2;
let s = strategy(None, false, tokenizer.clone());
let out = s.try_fit(&mut w, target);
// Should drop something to free tokens.
assert!(out.tokens_freed > 0);
assert!(w.messages.len() < 4);
}
#[test]
fn protected_messages_are_never_dropped() {
let tokenizer: Arc<dyn Tokenizer> = Arc::new(EstimatingCounter::default());
let messages = vec![
json!({"role": "user", "content": "PROTECTED ".repeat(50)}),
json!({"role": "assistant", "content": "PROTECTED ".repeat(50)}),
];
let mut w = ws(messages, tokenizer.as_ref());
// Protect both indices.
w.protected.insert(0);
w.protected.insert(1);
let initial_count = w.messages.len();
let s = strategy(None, false, tokenizer.clone());
let _ = s.try_fit(&mut w, 0);
// Nothing dropped because everything's protected.
assert_eq!(w.messages.len(), initial_count);
}
#[test]
fn ccr_on_drop_persists_and_emits_marker() {
let tokenizer: Arc<dyn Tokenizer> = Arc::new(EstimatingCounter::default());
let messages = vec![
json!({"role": "user", "content": "drop me ".repeat(50)}),
json!({"role": "assistant", "content": "and me ".repeat(50)}),
json!({"role": "user", "content": "keep recent".to_string()}),
json!({"role": "assistant", "content": "ok".to_string()}),
];
let mut w = ws(messages, tokenizer.as_ref());
// Protect the last two so the cascade has to drop the first two.
w.protected.insert(2);
w.protected.insert(3);
let store: Arc<dyn CcrStore> = Arc::new(InMemoryCcrStore::new());
let s = strategy(Some(store.clone()), true, tokenizer.clone());
let target = w.current_tokens / 2;
let out = s.try_fit(&mut w, target);
assert!(out.tokens_freed > 0);
assert!(!out.markers_inserted.is_empty());
// The marker references CCR retrieval.
assert!(out.markers_inserted[0].contains("ccr_retrieve"));
}
#[test]
fn ccr_off_does_not_persist() {
let tokenizer: Arc<dyn Tokenizer> = Arc::new(EstimatingCounter::default());
let messages = vec![
json!({"role": "user", "content": "drop me ".repeat(50)}),
json!({"role": "assistant", "content": "stay ".repeat(50)}),
];
let mut w = ws(messages, tokenizer.as_ref());
w.protected.insert(1);
let store: Arc<dyn CcrStore> = Arc::new(InMemoryCcrStore::new());
let s = strategy(Some(store.clone()), false, tokenizer.clone());
let target = w.current_tokens / 2;
let _ = s.try_fit(&mut w, target);
// ccr_on_drop=false → no markers, no persistence.
// (Can't easily probe the store without exposing internals;
// the marker absence is the visible contract.)
}
#[test]
fn no_candidates_returns_default_outcome() {
let tokenizer: Arc<dyn Tokenizer> = Arc::new(EstimatingCounter::default());
let messages = vec![json!({"role": "system", "content": "sys".repeat(100)})];
let mut w = ws(messages, tokenizer.as_ref());
// Protect the only message → no candidates.
w.protected.insert(0);
let s = strategy(None, false, tokenizer.clone());
let out = s.try_fit(&mut w, 0);
assert_eq!(out.tokens_freed, 0);
assert!(!out.fully_resolved);
}
#[test]
fn count_messages_returns_nonzero_for_nonempty_input() {
let tokenizer: Arc<dyn Tokenizer> = Arc::new(EstimatingCounter::default());
let msgs = vec![json!({"role": "user", "content": "hello"})];
assert!(count_messages(&msgs, tokenizer.as_ref()) > 0);
}
#[test]
fn count_messages_handles_anthropic_content_blocks() {
let tokenizer: Arc<dyn Tokenizer> = Arc::new(EstimatingCounter::default());
let msgs = vec![json!({
"role": "assistant",
"content": [
{"type": "text", "text": "thinking"},
{"type": "tool_use", "id": "tu", "name": "f", "input": {"x": 1}}
]
})];
let n = count_messages(&msgs, tokenizer.as_ref());
assert!(n > 0);
}
}

45
crates/headroom-core/src/context/strategy/mod.rs
View file

@ -1,45 +0,0 @@
//! `ContextStrategy` trait — the extension point for the Enterprise edition.
//!
//! OSS ships exactly one implementation: [`DropByScoreStrategy`]. Enterprise
//! plugs in additional strategies (compress-first, summarize, memory tiers)
//! and registers them via [`IntelligentContextManager::with_strategy`].
//!
//! # Why minimal
//!
//! The trait is intentionally one method. Earlier drafts added observer
//! hooks, telemetry callbacks, and tuning-hint inputs — speculative for
//! Enterprise needs we don't yet understand. Smaller surface = lower risk
//! of getting the API wrong on first try; we can extend with more methods
//! later (default-implemented for backwards compatibility).
pub mod drop_by_score;
use super::workspace::{ContextWorkspace, StrategyOutcome};
/// One stage in the context-fitting cascade.
///
/// Strategies run in registration order; each can either fully resolve
/// the budget (`fully_resolved=true` short-circuits the rest) or free
/// some tokens and let the next strategy continue.
///
/// Implementations must be `Send + Sync` — the manager is shared across
/// requests in a multi-threaded proxy. State mutations belong on the
/// `ContextWorkspace`, not `&self`.
pub trait ContextStrategy: Send + Sync {
/// Stable identifier shown in logs / `ApplyResult.strategies_applied`.
/// Use snake_case (`"drop_by_score"`, `"compress_first"`, etc.).
fn name(&self) -> &'static str;
/// Mutate the workspace toward fitting `target_tokens`.
///
/// Implementations should:
/// - Respect `ws.protected` and the leading `ws.frozen_count` indices
/// absolutely. Dropping a protected index is a contract violation.
/// - Update `ws.current_tokens` after mutating `ws.messages`.
/// - Append to `ws.dropped_indices` when removing messages.
/// - Report what they freed via `tokens_freed` so the manager can
/// decide whether to invoke the next strategy.
fn try_fit(&self, ws: &mut ContextWorkspace, target_tokens: usize) -> StrategyOutcome;
}
pub use drop_by_score::DropByScoreStrategy;

85
crates/headroom-core/src/context/workspace.rs
View file

@ -1,85 +0,0 @@
//! Mutable working set passed between strategies.
//!
//! Each [`ContextStrategy`](super::strategy::ContextStrategy) implementation
//! takes a `&mut ContextWorkspace`, mutates it (drops messages, inserts
//! markers, records bookkeeping), and returns a [`StrategyOutcome`].
//!
//! The workspace owns the messages — strategies never see the original
//! input list. The manager copies the input into the workspace before the
//! cascade and pulls the (possibly modified) list out at the end.
use std::collections::HashSet;
use serde_json::Value;
/// Mutable state threaded through the strategy cascade.
///
/// One workspace per `apply()` call. Not `Send` across the FFI boundary —
/// the manager constructs it inline and consumes it before returning.
pub struct ContextWorkspace {
/// The current message list. Strategies mutate this in place.
pub messages: Vec<Value>,
/// Indices that must NOT be dropped — system messages, last-N-turns,
/// frozen prefix, paired tool responses for protected assistants.
/// Computed by [`SafetyRails`](super::safety::SafetyRails) before
/// the cascade starts and re-computed after a strategy mutates the
/// list (since indices shift on drop).
pub protected: HashSet<usize>,
/// Number of leading messages that are part of a provider's
/// prompt-cache prefix. Dropping any of these busts the cache and
/// raises latency + cost. Always added to `protected`.
///
/// Provider-specific signal supplied by the caller via `ApplyCtx`.
/// Anthropic and OpenAI both expose prompt caching; OSS Headroom
/// receives this from the proxy layer (which knows the provider).
pub frozen_count: usize,
/// Token count of the current messages. Strategies update this
/// after each mutation so subsequent strategies see fresh state.
pub current_tokens: usize,
/// Indices the cascade has dropped, in drop order. Used by the
/// CCR-on-drop helper to serialize the originals into the cache.
/// Strategies append to this; they don't read it.
pub dropped_indices: Vec<usize>,
/// Original message list at workspace construction time. Preserved
/// so [`ccr_drop`](super::ccr_drop) can recover the *exact* dropped
/// content even after several drop rounds have shifted indices.
/// Strategies don't read this.
pub original_messages: Vec<Value>,
}
impl ContextWorkspace {
pub fn new(messages: Vec<Value>, protected: HashSet<usize>, frozen_count: usize) -> Self {
let original_messages = messages.clone();
Self {
messages,
protected,
frozen_count,
current_tokens: 0,
dropped_indices: Vec::new(),
original_messages,
}
}
}
/// Result of one strategy invocation.
#[derive(Debug, Clone, Default)]
pub struct StrategyOutcome {
/// Tokens freed by this strategy. The manager subtracts this from
/// `current_tokens` and decides whether to invoke the next strategy.
pub tokens_freed: usize,
/// Marker strings inserted by this strategy (e.g. CCR retrieval
/// hints). Aggregated into `ApplyResult.markers_inserted`.
pub markers_inserted: Vec<String>,
/// Set when the strategy fully resolved the budget. The manager
/// stops the cascade — later strategies don't run. Distinct from
/// "I freed some tokens" — a strategy can free tokens *and* still
/// not have brought the request under budget.
pub fully_resolved: bool,
}

2
crates/headroom-core/src/lib.rs
View file

@ -2,9 +2,7 @@
pub mod cache_control;
pub mod ccr;
pub mod context;
pub mod relevance;
pub mod scoring;
pub mod signals;
pub mod tokenizer;
pub mod transforms;

48
crates/headroom-core/src/scoring/mod.rs
View file

@ -1,48 +0,0 @@
//! Message-level importance scoring — used by IntelligentContextManager.
//!
//! # Why this lives at the crate root (parallel to `signals/`)
//!
//! `signals/` scores **lines** for line-level compressors (logs, search,
//! diffs). `scoring/` scores **messages** for conversation-level context
//! management. Different inputs, different consumers — separating them
//! keeps each trait surface clean and lets future ports compose without
//! a giant `signals` god-module.
//!
//! # Port status (Phase 7g PR-A, 2026-04-30)
//!
//! Direct port of `headroom/transforms/scoring.py` (459 LOC). The
//! deterministic factors — recency, forward references, density —
//! are fully implemented and parity-tested against Python. The
//! external-dependency factors are gated behind trait surfaces:
//!
//! - **TOIN** (`ToinProvider` trait): no concrete impl yet. Calls
//! return `0.5` for `toin_importance` and `0.0` for `error_indicator`
//! when no provider is wired in. PR-A1 will plug in a `PyO3`
//! `ToinProvider` so Rust can read Python's TOIN state.
//! - **Embeddings** (`EmbeddingProvider` trait): no concrete impl
//! here yet (the crate already has `relevance::EmbeddingScorer`
//! for SmartCrusher; PR-A1 wires the same `bge-small-en-v1.5`
//! model into a `MessageEmbedder` adapter). Until then,
//! `semantic_score` returns `0.5` (neutral).
//!
//! The trait surface is FULL — when the providers land, no API
//! changes are needed inside `MessageScorer`. The neutral-value
//! defaults match Python behavior when those subsystems are not
//! configured (`toin=None`, `embedding_provider=None`).
//!
//! # No hardcoded patterns (project convention)
//!
//! Mirrors the Python module's design principle: importance derives
//! from computed metrics (recency/density/refs), TOIN-learned
//! patterns (field semantics, retrieval rates), and embedding
//! similarity. No keyword regex, no hardcoded "error" strings.
pub mod score;
pub mod scorer;
pub mod traits;
pub mod weights;
pub use score::MessageScore;
pub use scorer::MessageScorer;
pub use traits::{EmbeddingProvider, ToinFieldSemantic, ToinPattern, ToinProvider};
pub use weights::ScoringWeights;

123
crates/headroom-core/src/scoring/score.rs
View file

@ -1,123 +0,0 @@
//! `MessageScore` — scoring output for a single message.
//!
//! Mirrors `headroom.transforms.scoring.MessageScore` byte-for-byte
//! including the per-component breakdown. The breakdown is what
//! IntelligentContextManager logs for debug + what TOIN consumes
//! for learning.
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
/// Importance score for a single message.
///
/// All component scores are in the range `[0.0, 1.0]` where higher =
/// more important. `total_score` is a weighted sum of the components
/// using [`crate::scoring::ScoringWeights`].
///
/// # Determinism
///
/// For the deterministic factors (recency, forward_reference,
/// token_density), the score is a pure function of the input
/// messages + index. Two runs with the same input produce
/// byte-identical scores.
///
/// For the external-dep factors (semantic_score, toin_score,
/// error_score), the value depends on whether providers are wired
/// in. Without providers, these return neutral defaults (`0.5` /
/// `0.0`) — same as Python's behavior with `embedding_provider=None`
/// / `toin=None`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MessageScore {
pub message_index: usize,
pub total_score: f32,
pub recency_score: f32,
pub semantic_score: f32,
pub toin_score: f32,
pub error_score: f32,
pub reference_score: f32,
pub density_score: f32,
/// Estimated tokens for this message. Python uses
/// `len(content) // 4` as a rough heuristic; we mirror exactly.
/// For non-string content (e.g. tool_calls list), Python returns
/// `100` as a default — we mirror that too.
pub tokens: usize,
pub is_protected: bool,
/// `not in_tool_unit OR not protected`. Mirrors Python's
/// confusing-but-faithful definition (it's intentionally the OR
/// of two negatives — see Python's `MessageScorer._score_message`
/// line 189).
pub drop_safe: bool,
/// Per-factor breakdown for debug logging + TOIN learning.
/// Keyed `BTreeMap` for deterministic JSON serialization order
/// (Python's `dict` preserves insertion order; we want stable
/// alphabetical so parity-fixture diffs are stable across runs).
pub score_breakdown: BTreeMap<String, f32>,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn round_trips_through_serde() {
let mut breakdown = BTreeMap::new();
breakdown.insert("recency".to_string(), 0.9);
breakdown.insert("semantic".to_string(), 0.5);
breakdown.insert("toin".to_string(), 0.5);
breakdown.insert("error".to_string(), 0.0);
breakdown.insert("reference".to_string(), 0.0);
breakdown.insert("density".to_string(), 0.7);
let s = MessageScore {
message_index: 3,
total_score: 0.42,
recency_score: 0.9,
semantic_score: 0.5,
toin_score: 0.5,
error_score: 0.0,
reference_score: 0.0,
density_score: 0.7,
tokens: 25,
is_protected: false,
drop_safe: true,
score_breakdown: breakdown,
};
let json = serde_json::to_string(&s).unwrap();
let back: MessageScore = serde_json::from_str(&json).unwrap();
assert_eq!(s, back);
}
#[test]
fn breakdown_serializes_in_alphabetical_order() {
// The breakdown order matters for parity-fixture stability.
// BTreeMap iteration is alphabetical → parity diffs stay
// deterministic across runs.
let mut breakdown = BTreeMap::new();
// Insert in non-alphabetical order:
for (k, v) in [
("toin", 1.0),
("recency", 2.0),
("density", 3.0),
("semantic", 4.0),
("error", 5.0),
("reference", 6.0),
] {
breakdown.insert(k.to_string(), v);
}
let json = serde_json::to_string(&breakdown).unwrap();
// density < error < recency < reference < semantic < toin alphabetically
assert!(
json.find("density").unwrap() < json.find("error").unwrap()
&& json.find("error").unwrap() < json.find("recency").unwrap()
&& json.find("recency").unwrap() < json.find("reference").unwrap()
&& json.find("reference").unwrap() < json.find("semantic").unwrap()
&& json.find("semantic").unwrap() < json.find("toin").unwrap()
);
}
}

736
crates/headroom-core/src/scoring/scorer.rs
View file

@ -1,736 +0,0 @@
//! `MessageScorer` — six-factor importance scoring for messages.
//!
//! Direct port of `headroom.transforms.scoring.MessageScorer`. See
//! the module-level doc in `mod.rs` for what is and isn't wired up
//! in PR-A.
//!
//! # Parity contract
//!
//! For the deterministic factors (recency, forward_reference,
//! token_density), the Rust implementation must produce
//! float-epsilon-equal scores to the Python implementation given
//! the same input. This is what `crates/headroom-parity/` validates.
//!
//! For non-deterministic factors (semantic, toin, error), parity is
//! validated only when both implementations are configured with the
//! same providers. Without providers, both return the same neutral
//! defaults.
use std::collections::HashMap;
use std::sync::Mutex;
use serde_json::{Map, Value};
use crate::scoring::score::MessageScore;
use crate::scoring::traits::{EmbeddingProvider, ToinPattern, ToinProvider};
use crate::scoring::weights::ScoringWeights;
/// Six-factor importance scorer. See module-level docs for the
/// list of factors and their weights.
///
/// # Thread safety
///
/// `MessageScorer` is `Send + Sync`. The internal embedding cache is
/// guarded by a `Mutex` — contention is low because cache writes
/// happen at most once per (message, scorer) pair and scorer
/// instances are typically per-request, not shared.
pub struct MessageScorer {
weights: ScoringWeights,
toin: Option<Box<dyn ToinProvider>>,
embedding_provider: Option<Box<dyn EmbeddingProvider>>,
recency_decay_rate: f32,
embedding_cache: Mutex<HashMap<usize, Vec<f32>>>,
}
impl MessageScorer {
/// Create a new scorer.
///
/// `weights` are normalized on construction (matching Python's
/// `ScoringWeights().normalized()` in `__init__`). Pass `None`
/// for `toin` and `embedding_provider` to use neutral defaults
/// for those factors — same as Python's `toin=None` /
/// `embedding_provider=None`.
pub fn new(
weights: Option<ScoringWeights>,
toin: Option<Box<dyn ToinProvider>>,
embedding_provider: Option<Box<dyn EmbeddingProvider>>,
recency_decay_rate: f32,
) -> Self {
Self {
weights: weights.unwrap_or_default().normalized(),
toin,
embedding_provider,
recency_decay_rate,
embedding_cache: Mutex::new(HashMap::new()),
}
}
/// Create a scorer with default weights, no providers, and the
/// Python-default decay rate of 0.1. Useful for tests and the
/// deterministic-only code path.
pub fn with_defaults() -> Self {
Self::new(None, None, None, 0.1)
}
/// Score every message in the list.
///
/// `protected_indices` are indices the caller has marked as
/// system / pinned / never-drop. They get `is_protected=true`.
/// `tool_unit_indices` are part of an inseparable tool unit
/// (assistant tool_call + tool response pair) — these get
/// `drop_safe=false` unless also protected, matching Python's
/// confusing-but-faithful `not in_tool_unit OR not protected`
/// rule.
pub fn score_messages(
&self,
messages: &[Value],
protected_indices: &std::collections::HashSet<usize>,
tool_unit_indices: &std::collections::HashSet<usize>,
) -> Vec<MessageScore> {
let forward_refs = Self::compute_forward_references(messages);
let recent_embedding = self.compute_recent_context_embedding(messages, 3);
messages
.iter()
.enumerate()
.map(|(i, msg)| {
self.score_message(
msg,
i,
messages.len(),
protected_indices.contains(&i),
tool_unit_indices.contains(&i),
&forward_refs,
recent_embedding.as_deref(),
)
})
.collect()
}
#[allow(clippy::too_many_arguments)]
fn score_message(
&self,
msg: &Value,
index: usize,
total: usize,
protected: bool,
in_tool_unit: bool,
forward_refs: &HashMap<usize, u32>,
recent_embedding: Option<&[f32]>,
) -> MessageScore {
let recency = self.compute_recency_score(index, total);
let semantic = self.compute_semantic_score(msg, index, recent_embedding);
let toin = self.compute_toin_score(msg);
let error = self.compute_error_score(msg);
let reference = self.compute_reference_score(index, forward_refs);
let density = Self::compute_density_score(msg);
let w = &self.weights;
let total_score = w.recency * recency
+ w.semantic_similarity * semantic
+ w.toin_importance * toin
+ w.error_indicator * error
+ w.forward_reference * reference
+ w.token_density * density;
let tokens = estimate_tokens(msg);
// BTreeMap so JSON serialization order is alphabetical.
let mut breakdown = std::collections::BTreeMap::new();
breakdown.insert("recency".to_string(), recency);
breakdown.insert("semantic".to_string(), semantic);
breakdown.insert("toin".to_string(), toin);
breakdown.insert("error".to_string(), error);
breakdown.insert("reference".to_string(), reference);
breakdown.insert("density".to_string(), density);
MessageScore {
message_index: index,
total_score,
recency_score: recency,
semantic_score: semantic,
toin_score: toin,
error_score: error,
reference_score: reference,
density_score: density,
tokens,
is_protected: protected,
// Python: `not in_tool_unit or not protected` — see
// scoring.py:189. This *is* the literal expression.
drop_safe: !in_tool_unit || !protected,
score_breakdown: breakdown,
}
}
fn compute_recency_score(&self, index: usize, total: usize) -> f32 {
if total <= 1 {
return 1.0;
}
let position_from_end = (total - 1 - index) as f32;
(-self.recency_decay_rate * position_from_end).exp()
}
fn compute_semantic_score(
&self,
msg: &Value,
index: usize,
recent_embedding: Option<&[f32]>,
) -> f32 {
let Some(provider) = self.embedding_provider.as_ref() else {
return 0.5;
};
let Some(recent) = recent_embedding else {
return 0.5;
};
let content = match msg.get("content") {
Some(Value::String(s)) if !s.trim().is_empty() => s,
_ => return 0.5,
};
// Cache lookup; populate on miss.
let msg_embedding: Vec<f32> = {
let mut cache = self.embedding_cache.lock().unwrap();
if let Some(cached) = cache.get(&index) {
cached.clone()
} else {
match provider.embed(content) {
Ok(v) => {
cache.insert(index, v.clone());
v
}
Err(_) => return 0.5,
}
}
};
cosine_similarity(&msg_embedding, recent)
}
fn compute_toin_score(&self, msg: &Value) -> f32 {
let Some(toin) = self.toin.as_ref() else {
return 0.5;
};
if msg.get("role").and_then(Value::as_str) != Some("tool") {
return 0.5;
}
let Some(content_value) = parse_tool_content(msg) else {
return 0.5;
};
let Some(pattern) = toin.pattern_for_tool_content(&content_value) else {
return 0.5;
};
if pattern.confidence < 0.3 {
return 0.5;
}
let mut score = 0.5 + pattern.retrieval_rate * 0.5;
if !pattern.commonly_retrieved_fields.is_empty() {
// Python: min(0.1, 0.02 * len(commonly_retrieved_fields))
let boost = (0.02 * pattern.commonly_retrieved_fields.len() as f32).min(0.1);
score = (score + boost).min(1.0);
}
score
}
fn compute_error_score(&self, msg: &Value) -> f32 {
let Some(toin) = self.toin.as_ref() else {
return 0.0;
};
if msg.get("role").and_then(Value::as_str) != Some("tool") {
return 0.0;
}
let Some(content_value) = parse_tool_content(msg) else {
return 0.0;
};
let Some(pattern) = toin.pattern_for_tool_content(&content_value) else {
return 0.0;
};
let (error_field_count, high_confidence_errors) = count_error_fields(&pattern);
if error_field_count == 0 {
return 0.0;
}
// Python: base_score = min(1.0, 0.3 * error_field_count)
// confidence_boost = min(0.5, 0.2 * high_confidence_errors)
let base = (0.3 * error_field_count as f32).min(1.0);
let boost = (0.2 * high_confidence_errors as f32).min(0.5);
base + boost
}
fn compute_reference_score(&self, index: usize, forward_refs: &HashMap<usize, u32>) -> f32 {
let count = *forward_refs.get(&index).unwrap_or(&0);
if count == 0 {
return 0.0;
}
// Python: min(1.0, 0.3 + 0.2 * math.log(ref_count + 1))
// math.log is natural log.
let v = 0.3 + 0.2 * ((count + 1) as f32).ln();
v.min(1.0)
}
fn compute_density_score(msg: &Value) -> f32 {
let content = match msg.get("content") {
Some(Value::String(s)) => s,
_ => return 0.5,
};
// Python: `len(content) < 10` — len counts chars (code points).
if content.chars().count() < 10 {
return 0.5;
}
let lower = content.to_lowercase();
let tokens: Vec<&str> = lower.split_whitespace().collect();
if tokens.len() < 3 {
return 0.5;
}
let unique: std::collections::HashSet<&&str> = tokens.iter().collect();
let density = unique.len() as f32 / tokens.len() as f32;
// Python: min(1.0, max(0.0, (density - 0.2) / 0.6))
((density - 0.2) / 0.6).clamp(0.0, 1.0)
}
fn compute_forward_references(messages: &[Value]) -> HashMap<usize, u32> {
let mut refs: HashMap<usize, u32> = HashMap::new();
let mut tool_call_ids: HashMap<String, usize> = HashMap::new();
for (i, msg) in messages.iter().enumerate() {
let role = msg.get("role").and_then(Value::as_str);
match role {
Some("assistant") => {
if let Some(tcs) = msg.get("tool_calls").and_then(Value::as_array) {
for tc in tcs {
if let Some(id) = tc.get("id").and_then(Value::as_str) {
tool_call_ids.insert(id.to_string(), i);
}
}
}
}
Some("tool") => {
if let Some(tcid) = msg.get("tool_call_id").and_then(Value::as_str) {
if let Some(&ref_idx) = tool_call_ids.get(tcid) {
*refs.entry(ref_idx).or_insert(0) += 1;
}
}
}
_ => {}
}
}
refs
}
fn compute_recent_context_embedding(
&self,
messages: &[Value],
num_recent: usize,
) -> Option<Vec<f32>> {
let provider = self.embedding_provider.as_ref()?;
let start = messages.len().saturating_sub(num_recent);
let mut texts: Vec<&str> = Vec::new();
for msg in &messages[start..] {
if let Some(Value::String(s)) = msg.get("content") {
if !s.trim().is_empty() {
texts.push(s);
}
}
}
if texts.is_empty() {
return None;
}
let combined = texts.join(" ");
provider.embed(&combined).ok()
}
}
/// Token estimate. Python: `len(content) // 4` for strings, `100`
/// for non-strings. We use `chars().count()` (code points) to match
/// Python's `len(str)`.
fn estimate_tokens(msg: &Value) -> usize {
match msg.get("content") {
Some(Value::String(s)) => s.chars().count() / 4,
_ => 100,
}
}
/// Parse a tool message's `content` as JSON if it's a string, or
/// return it directly if already an object/array. Matches Python's
/// behavior in `_compute_toin_score` / `_compute_error_score`:
/// returns `None` for anything that isn't a list/dict.
fn parse_tool_content(msg: &Value) -> Option<Value> {
let content = msg.get("content")?;
let parsed = match content {
Value::String(s) => {
if s.is_empty() {
return None;
}
serde_json::from_str::<Value>(s).ok()?
}
v => v.clone(),
};
match &parsed {
Value::Object(_) | Value::Array(_) => {
// Python: `if not items: return 0.5/0.0` — empty
// list/object disqualifies.
if let Value::Array(a) = &parsed {
if a.is_empty() {
return None;
}
}
if let Value::Object(o) = &parsed {
if o.is_empty() {
// Python wraps a dict into a list before checking
// truthiness; an empty dict becomes `[{}]` which
// is truthy. Mirror exactly:
let _ = o;
// (still return Some since [{}] is truthy in Py)
}
}
Some(parsed)
}
_ => None,
}
}
/// Count `error_indicator` fields in a TOIN pattern. Returns
/// `(total_error_fields, high_confidence_error_fields)` where
/// high-confidence means `confidence >= 0.7`.
fn count_error_fields(pattern: &ToinPattern) -> (u32, u32) {
let mut total = 0u32;
let mut high = 0u32;
for field_sem in pattern.field_semantics.values() {
if field_sem.inferred_type == "error_indicator" {
total += 1;
if field_sem.confidence >= 0.7 {
high += 1;
}
}
}
(total, high)
}
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
if a.len() != b.len() || a.is_empty() {
return 0.0;
}
let mut dot = 0.0f32;
let mut na = 0.0f32;
let mut nb = 0.0f32;
for i in 0..a.len() {
dot += a[i] * b[i];
na += a[i] * a[i];
nb += b[i] * b[i];
}
if na == 0.0 || nb == 0.0 {
return 0.0;
}
dot / (na.sqrt() * nb.sqrt())
}
/// Helper to build a message `Value` for tests. Public-in-crate so
/// integration tests can use it too.
#[allow(dead_code)]
pub(crate) fn msg(role: &str, content: &str) -> Value {
let mut m = Map::new();
m.insert("role".to_string(), Value::String(role.to_string()));
m.insert("content".to_string(), Value::String(content.to_string()));
Value::Object(m)
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashSet;
fn empty_set() -> HashSet<usize> {
HashSet::new()
}
#[test]
fn recency_single_message_returns_one() {
let s = MessageScorer::with_defaults();
assert_eq!(s.compute_recency_score(0, 1), 1.0);
assert_eq!(s.compute_recency_score(0, 0), 1.0);
}
#[test]
fn recency_last_message_full_score() {
let s = MessageScorer::with_defaults();
// Last message: position_from_end = 0, score = e^0 = 1.0
assert!((s.compute_recency_score(4, 5) - 1.0).abs() < 1e-6);
}
#[test]
fn recency_decays_exponentially() {
let s = MessageScorer::with_defaults();
// total=10, index=0 → position_from_end=9 → e^(-0.1*9) = e^-0.9
let expected = (-0.9f32).exp();
let got = s.compute_recency_score(0, 10);
assert!(
(got - expected).abs() < 1e-6,
"expected {expected}, got {got}"
);
}
#[test]
fn density_short_content_is_neutral() {
// Less than 10 chars
let m = msg("user", "hi");
assert_eq!(MessageScorer::compute_density_score(&m), 0.5);
}
#[test]
fn density_few_tokens_is_neutral() {
// >= 10 chars but < 3 tokens after split
let m = msg("user", "abcdefghij"); // single token, 10 chars
assert_eq!(MessageScorer::compute_density_score(&m), 0.5);
}
#[test]
fn density_all_unique_clamps_to_one() {
// 10 unique tokens, 10 total → density=1.0 → (1.0-0.2)/0.6=1.33 → clamped to 1.0
let m = msg(
"user",
"alpha bravo charlie delta echo foxtrot golf hotel india juliet",
);
assert!((MessageScorer::compute_density_score(&m) - 1.0).abs() < 1e-6);
}
#[test]
fn density_repeated_tokens_lowers_score() {
// 4 unique / 8 total = 0.5 → (0.5-0.2)/0.6 = 0.5
let m = msg(
"user",
"alpha bravo charlie delta alpha bravo charlie delta",
);
let got = MessageScorer::compute_density_score(&m);
assert!((got - 0.5).abs() < 1e-6, "got {got}");
}
#[test]
fn density_non_string_content_is_neutral() {
let mut m = Map::new();
m.insert("role".to_string(), Value::String("assistant".to_string()));
m.insert(
"tool_calls".to_string(),
Value::Array(vec![Value::Object(Map::new())]),
);
let v = Value::Object(m);
assert_eq!(MessageScorer::compute_density_score(&v), 0.5);
}
#[test]
fn forward_refs_links_tool_response_to_assistant() {
let mut tc = Map::new();
tc.insert("id".to_string(), Value::String("call_1".to_string()));
let mut assistant = Map::new();
assistant.insert("role".to_string(), Value::String("assistant".to_string()));
assistant.insert(
"tool_calls".to_string(),
Value::Array(vec![Value::Object(tc)]),
);
let mut tool_resp = Map::new();
tool_resp.insert("role".to_string(), Value::String("tool".to_string()));
tool_resp.insert(
"tool_call_id".to_string(),
Value::String("call_1".to_string()),
);
tool_resp.insert("content".to_string(), Value::String("result".to_string()));
let messages = vec![
msg("user", "do thing"),
Value::Object(assistant),
Value::Object(tool_resp),
];
let refs = MessageScorer::compute_forward_references(&messages);
assert_eq!(refs.get(&1), Some(&1));
assert_eq!(refs.get(&0), None);
}
#[test]
fn forward_refs_unmatched_tool_call_id_ignored() {
let mut tool_resp = Map::new();
tool_resp.insert("role".to_string(), Value::String("tool".to_string()));
tool_resp.insert(
"tool_call_id".to_string(),
Value::String("nope".to_string()),
);
tool_resp.insert("content".to_string(), Value::String("result".to_string()));
let messages = vec![Value::Object(tool_resp)];
let refs = MessageScorer::compute_forward_references(&messages);
assert!(refs.is_empty());
}
#[test]
fn reference_score_zero_refs_returns_zero() {
let s = MessageScorer::with_defaults();
let refs = HashMap::new();
assert_eq!(s.compute_reference_score(0, &refs), 0.0);
}
#[test]
fn reference_score_one_ref_uses_log_formula() {
let s = MessageScorer::with_defaults();
let mut refs = HashMap::new();
refs.insert(0, 1);
// 0.3 + 0.2 * ln(2) = 0.3 + 0.2 * 0.693... ≈ 0.4386
let expected = 0.3 + 0.2 * 2f32.ln();
let got = s.compute_reference_score(0, &refs);
assert!((got - expected).abs() < 1e-6, "got {got}");
}
#[test]
fn reference_score_clamps_to_one() {
let s = MessageScorer::with_defaults();
let mut refs = HashMap::new();
refs.insert(0, 1_000_000);
assert!(s.compute_reference_score(0, &refs) <= 1.0);
}
#[test]
fn semantic_score_no_provider_returns_neutral() {
let s = MessageScorer::with_defaults();
let m = msg("user", "hello world");
// recent_embedding is None too without a provider
assert_eq!(s.compute_semantic_score(&m, 0, None), 0.5);
}
#[test]
fn toin_score_no_provider_returns_neutral() {
let s = MessageScorer::with_defaults();
let m = msg("tool", "{}");
assert_eq!(s.compute_toin_score(&m), 0.5);
}
#[test]
fn error_score_no_provider_returns_zero() {
let s = MessageScorer::with_defaults();
let m = msg("tool", "{}");
assert_eq!(s.compute_error_score(&m), 0.0);
}
#[test]
fn estimate_tokens_string_uses_char_count() {
// 16 chars / 4 = 4
let m = msg("user", "abcdefghijklmnop");
assert_eq!(estimate_tokens(&m), 4);
}
#[test]
fn estimate_tokens_unicode_is_char_aware() {
// 4 emoji chars (each multi-byte) → len("...") in Python is 4.
let m = msg("user", "🎉🎉🎉🎉");
// chars().count() = 4, // 4 = 1
assert_eq!(estimate_tokens(&m), 1);
}
#[test]
fn estimate_tokens_non_string_returns_default() {
let mut m = Map::new();
m.insert("role".to_string(), Value::String("assistant".to_string()));
m.insert("tool_calls".to_string(), Value::Array(vec![]));
let v = Value::Object(m);
assert_eq!(estimate_tokens(&v), 100);
}
#[test]
fn cosine_similarity_identical_vectors_is_one() {
let v = [1.0f32, 2.0, 3.0];
assert!((cosine_similarity(&v, &v) - 1.0).abs() < 1e-6);
}
#[test]
fn cosine_similarity_orthogonal_is_zero() {
let a = [1.0f32, 0.0];
let b = [0.0f32, 1.0];
assert!(cosine_similarity(&a, &b).abs() < 1e-6);
}
#[test]
fn cosine_similarity_zero_vector_is_zero() {
let a = [0.0f32, 0.0];
let b = [1.0f32, 1.0];
assert_eq!(cosine_similarity(&a, &b), 0.0);
}
#[test]
fn cosine_similarity_dim_mismatch_is_zero() {
let a = [1.0f32, 2.0];
let b = [1.0f32, 2.0, 3.0];
assert_eq!(cosine_similarity(&a, &b), 0.0);
}
#[test]
fn drop_safe_mirrors_python_or_logic() {
// Python: drop_safe = not in_tool_unit OR not protected
// truth table:
// in_tool_unit=F, protected=F → T
// in_tool_unit=F, protected=T → T (not F = T)
// in_tool_unit=T, protected=F → T (not F = T)
// in_tool_unit=T, protected=T → F
let s = MessageScorer::with_defaults();
let m = msg("user", "hi");
let refs = HashMap::new();
let cases = [
(false, false, true),
(false, true, true),
(true, false, true),
(true, true, false),
];
for (in_tu, prot, expected) in cases {
let score = s.score_message(&m, 0, 1, prot, in_tu, &refs, None);
assert_eq!(
score.drop_safe, expected,
"in_tool_unit={in_tu} protected={prot}"
);
}
}
#[test]
fn score_messages_returns_one_score_per_message() {
let s = MessageScorer::with_defaults();
let messages = vec![
msg("user", "hello"),
msg("assistant", "hi there"),
msg("user", "thanks"),
];
let scores = s.score_messages(&messages, &empty_set(), &empty_set());
assert_eq!(scores.len(), 3);
assert_eq!(scores[0].message_index, 0);
assert_eq!(scores[2].message_index, 2);
}
#[test]
fn score_messages_protected_indices_are_marked() {
let s = MessageScorer::with_defaults();
let messages = vec![msg("user", "hi"), msg("user", "bye")];
let mut protected = HashSet::new();
protected.insert(0);
let scores = s.score_messages(&messages, &protected, &empty_set());
assert!(scores[0].is_protected);
assert!(!scores[1].is_protected);
}
#[test]
fn weights_are_normalized_on_construction() {
let unbalanced = ScoringWeights {
recency: 2.0,
semantic_similarity: 2.0,
toin_importance: 2.0,
error_indicator: 2.0,
forward_reference: 2.0,
token_density: 2.0,
};
let s = MessageScorer::new(Some(unbalanced), None, None, 0.1);
// After normalization each should be 1/6.
assert!((s.weights.recency - 1.0 / 6.0).abs() < 1e-6);
}
}

175
crates/headroom-core/src/scoring/traits.rs
View file

@ -1,175 +0,0 @@
//! External-dependency trait surfaces for `MessageScorer`.
//!
//! The scorer's six factors split into two groups:
//!
//! - **Pure / deterministic** (recency, forward references, density):
//! computed in-crate from the message list alone. No traits needed.
//! - **External-dependency** (semantic similarity, TOIN importance,
//! error indicator): require either an embedding model or learned
//! TOIN telemetry. These get trait surfaces here so the scorer
//! doesn't have to know how those subsystems are implemented.
//!
//! In PR-A (this PR), no concrete impls exist. PR-A1 wires
//! `EmbeddingProvider` to fastembed (reusing the `bge-small-en-v1.5`
//! model already loaded for SmartCrusher relevance). PR-A2 wires
//! `ToinProvider` to a PyO3 adapter so Rust can read Python's TOIN
//! state. When both land, the scorer code is unchanged — only the
//! provider construction at startup differs.
//!
//! # Why pass `&serde_json::Value` and not `&str` for tool content
//!
//! Python's `MessageScorer._compute_toin_score` parses the message
//! content as JSON, then computes a `ToolSignature` from the parsed
//! items, then looks up the learned pattern. The trait pushes the
//! parsing into the implementor — partly because the parsing is
//! cheap and message-local, partly because `ToolSignature` itself is
//! TOIN-internal (not yet ported, and not needed outside TOIN).
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
/// Embedding provider for semantic-similarity scoring.
///
/// Implementations should return a fixed-dimension `Vec<f32>` for
/// each input text. The dimension must be consistent across calls
/// from a given instance — `MessageScorer` does cosine similarity
/// between vectors and assumes equal dimension.
///
/// `embed` may fail in implementations that wrap external services;
/// the scorer treats failures as "no signal" (returns the neutral
/// `0.5` semantic score), matching Python's try/except behavior.
pub trait EmbeddingProvider: Send + Sync {
/// Embed a string into a fixed-dimension vector. Empty strings
/// or whitespace-only strings should still return a valid vector
/// — the scorer pre-filters empty content before calling this.
fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError>;
}
/// Error type for embedding providers. Wraps the underlying impl's
/// error as a string for trait-object compatibility — losing typed
/// context is fine here since the scorer just logs + falls back.
#[derive(Debug, Clone)]
pub struct EmbeddingError(pub String);
impl std::fmt::Display for EmbeddingError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "embedding error: {}", self.0)
}
}
impl std::error::Error for EmbeddingError {}
/// TOIN (Tool Output Intelligence Network) pattern lookup.
///
/// TOIN learns retrieval patterns per tool-output structure. The
/// scorer queries it for two things:
///
/// 1. **Importance** (`toin_score`): high `retrieval_rate` means
/// users repeatedly retrieve this tool's data → keep it.
/// 2. **Error detection** (`error_score`): TOIN classifies fields
/// by inferred type. Fields tagged `error_indicator` boost the
/// error score, in lieu of hardcoded keyword regex.
///
/// The trait takes the parsed JSON content directly (rather than a
/// pre-computed structure hash) because `ToolSignature` derivation
/// is TOIN-internal — pushing it across the trait boundary would
/// leak implementation details.
pub trait ToinProvider: Send + Sync {
/// Look up the learned pattern for a tool-message content
/// payload. Returns `None` if the content can't be classified
/// (not list/dict, empty, etc.) or no pattern has been learned
/// for this structure yet.
fn pattern_for_tool_content(&self, content: &serde_json::Value) -> Option<ToinPattern>;
}
/// Snapshot of a TOIN-learned pattern for a single tool-output
/// structure. Mirrors the subset of `headroom.telemetry.toin.ToolPattern`
/// that `MessageScorer` actually reads.
///
/// We don't mirror the full Python class — TOIN updates patterns
/// in-place during learning, but the scorer only reads. Returning
/// a snapshot decouples scorer reads from learning writes (no
/// shared-mutable-state across the FFI boundary in PR-A2).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToinPattern {
/// Overall confidence in this pattern, `[0.0, 1.0]`. Patterns
/// with `confidence < 0.3` are treated as not-yet-learned and
/// the scorer falls back to neutral.
pub confidence: f32,
/// Fraction of tool invocations whose data was later retrieved
/// from cache or referenced by a follow-up message. `[0.0, 1.0]`.
/// High values mean "users keep needing this data" → important.
pub retrieval_rate: f32,
/// Field hashes (NOT names — TOIN privacy-hashes them) that
/// are commonly retrieved from this structure. The scorer just
/// uses the *count* as a small importance boost; it doesn't
/// dereference individual hashes.
pub commonly_retrieved_fields: Vec<String>,
/// Per-field semantics, keyed by privacy-hashed field name.
/// `BTreeMap` rather than `HashMap` so serialization order is
/// deterministic (matters for parity-fixture stability).
pub field_semantics: BTreeMap<String, ToinFieldSemantic>,
}
/// Inferred semantic type + confidence for a single field, as
/// learned by TOIN. Mirrors a subset of
/// `headroom.telemetry.models.FieldSemantics`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToinFieldSemantic {
/// Inferred semantic category. The scorer specifically checks
/// for `"error_indicator"` to compute the error score; other
/// values (`"identifier"`, `"status"`, `"timestamp"`, etc.) are
/// not used by scoring but kept for completeness so we can pass
/// the same struct through to other consumers.
pub inferred_type: String,
/// Confidence in the inferred type, `[0.0, 1.0]`. The scorer
/// applies a `>= 0.7` threshold for a "high-confidence error"
/// boost; below that it still counts the field but doesn't
/// boost.
pub confidence: f32,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn pattern_round_trips_through_serde() {
let mut field_semantics = BTreeMap::new();
field_semantics.insert(
"abc123".to_string(),
ToinFieldSemantic {
inferred_type: "error_indicator".to_string(),
confidence: 0.85,
},
);
field_semantics.insert(
"def456".to_string(),
ToinFieldSemantic {
inferred_type: "identifier".to_string(),
confidence: 0.9,
},
);
let p = ToinPattern {
confidence: 0.75,
retrieval_rate: 0.6,
commonly_retrieved_fields: vec!["abc123".to_string()],
field_semantics,
};
let json = serde_json::to_string(&p).unwrap();
let back: ToinPattern = serde_json::from_str(&json).unwrap();
assert_eq!(p, back);
}
#[test]
fn embedding_error_is_displayable() {
let e = EmbeddingError("model not loaded".to_string());
assert_eq!(e.to_string(), "embedding error: model not loaded");
}
}

170
crates/headroom-core/src/scoring/weights.rs
View file

@ -1,170 +0,0 @@
//! `ScoringWeights` — six-factor weighted importance scoring.
//!
//! Mirrors `headroom.config.ScoringWeights` byte-for-byte. The six
//! factors and their default contributions sum to ~1.0; `normalized()`
//! enforces that explicitly when a caller passes weights that don't
//! sum to 1.0 (e.g. learned weights from TOIN telemetry).
//!
//! Default values match Python's defaults exactly so the parity
//! fixtures byte-equal across implementations.
use serde::{Deserialize, Serialize};
/// Weights for the six importance-scoring factors.
///
/// All weights should sum to ~1.0 for normalized scoring (call
/// [`Self::normalized`] to enforce). Non-normalized weights are
/// permitted — `MessageScorer` does NOT auto-normalize on input —
/// since callers may use raw weights for relative comparison.
///
/// Defaults match Python's `ScoringWeights()`:
///
/// | Factor | Weight |
/// |--------|--------|
/// | recency | 0.20 |
/// | semantic_similarity | 0.20 |
/// | toin_importance | 0.25 |
/// | error_indicator | 0.15 |
/// | forward_reference | 0.15 |
/// | token_density | 0.05 |
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct ScoringWeights {
/// Exponential decay from conversation end. Default 0.20.
pub recency: f32,
/// Embedding cosine similarity to recent context. Default 0.20.
pub semantic_similarity: f32,
/// TOIN-learned field importance. Default 0.25.
pub toin_importance: f32,
/// TOIN-learned error-field detection. Default 0.15.
pub error_indicator: f32,
/// Number of later messages referencing this one (tool_call_id). Default 0.15.
pub forward_reference: f32,
/// Information density (unique-tokens / total-tokens). Default 0.05.
pub token_density: f32,
}
impl Default for ScoringWeights {
fn default() -> Self {
Self {
recency: 0.20,
semantic_similarity: 0.20,
toin_importance: 0.25,
error_indicator: 0.15,
forward_reference: 0.15,
token_density: 0.05,
}
}
}
impl ScoringWeights {
/// Return a copy with all weights divided by their sum, so they sum
/// to 1.0 exactly. If the input sums to 0 (degenerate config),
/// returns the default weights.
pub fn normalized(&self) -> Self {
let total = self.recency
+ self.semantic_similarity
+ self.toin_importance
+ self.error_indicator
+ self.forward_reference
+ self.token_density;
if total == 0.0 {
return Self::default();
}
Self {
recency: self.recency / total,
semantic_similarity: self.semantic_similarity / total,
toin_importance: self.toin_importance / total,
error_indicator: self.error_indicator / total,
forward_reference: self.forward_reference / total,
token_density: self.token_density / total,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn defaults_match_python() {
let w = ScoringWeights::default();
assert_eq!(w.recency, 0.20);
assert_eq!(w.semantic_similarity, 0.20);
assert_eq!(w.toin_importance, 0.25);
assert_eq!(w.error_indicator, 0.15);
assert_eq!(w.forward_reference, 0.15);
assert_eq!(w.token_density, 0.05);
}
#[test]
fn defaults_sum_to_one() {
let w = ScoringWeights::default();
let total = w.recency
+ w.semantic_similarity
+ w.toin_importance
+ w.error_indicator
+ w.forward_reference
+ w.token_density;
assert!((total - 1.0).abs() < 1e-6);
}
#[test]
fn normalized_already_normal_is_idempotent() {
let w = ScoringWeights::default();
let n = w.normalized();
// Within float epsilon — divisions can introduce tiny drift.
assert!((n.recency - w.recency).abs() < 1e-6);
assert!((n.toin_importance - w.toin_importance).abs() < 1e-6);
}
#[test]
fn normalized_unbalanced_weights_sum_to_one() {
let w = ScoringWeights {
recency: 1.0,
semantic_similarity: 1.0,
toin_importance: 1.0,
error_indicator: 1.0,
forward_reference: 1.0,
token_density: 1.0,
};
let n = w.normalized();
let total = n.recency
+ n.semantic_similarity
+ n.toin_importance
+ n.error_indicator
+ n.forward_reference
+ n.token_density;
assert!((total - 1.0).abs() < 1e-6);
// Each component should now be ~1/6.
assert!((n.recency - 1.0 / 6.0).abs() < 1e-6);
}
#[test]
fn normalized_zero_weights_falls_back_to_default() {
let w = ScoringWeights {
recency: 0.0,
semantic_similarity: 0.0,
toin_importance: 0.0,
error_indicator: 0.0,
forward_reference: 0.0,
token_density: 0.0,
};
let n = w.normalized();
assert_eq!(n, ScoringWeights::default());
}
#[test]
fn round_trips_through_serde() {
let w = ScoringWeights {
recency: 0.3,
semantic_similarity: 0.1,
toin_importance: 0.2,
error_indicator: 0.2,
forward_reference: 0.1,
token_density: 0.1,
};
let json = serde_json::to_string(&w).unwrap();
let back: ScoringWeights = serde_json::from_str(&json).unwrap();
assert_eq!(w, back);
}
}

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

@ -23,6 +23,7 @@ pub mod diff_compressor;
pub mod log_compressor;
pub mod magika_detector;
pub mod pipeline;
pub mod safety;
pub mod search_compressor;
pub mod smart_crusher;
pub mod tag_protector;
@ -45,6 +46,7 @@ pub use pipeline::{
JsonMinifier, JsonOffload, LogOffload, LogTemplate, OffloadOutput, OffloadTransform,
PipelineConfig, PipelineResult, ReformatOutput, ReformatTransform, TransformError,
};
pub use safety::{tool_pair_indices, ToolPair};
pub use search_compressor::{
FileMatches, SearchCompressionResult, SearchCompressor, SearchCompressorConfig,
SearchCompressorStats, SearchMatch,

215
crates/headroom-core/src/transforms/safety.rs Normal file
View file

@ -0,0 +1,215 @@
//! Tool-pair atomicity rules for the live-zone-only compression
//! architecture.
//!
//! Live-zone compression never drops messages — it operates on
//! content blocks within messages. So the old "which indices are
//! safe to drop" question goes away. What remains, and what this
//! module provides, is the rule that an `assistant.tool_use` and
//! its matching `tool_result` must be treated as one unit when
//! deciding what to compress: compressing one but not the other
//! desynchronizes the conversation and a re-replay of the tool
//! response will mismatch the call id, producing 400s upstream.
//!
//! For the live-zone block dispatcher (`PR-B2`+), this module
//! exposes [`tool_pair_indices`] — given a slice of messages, it
//! returns the index pairs that must be co-treated.
//!
//! Both OpenAI and Anthropic tool-call shapes are recognized:
//!
//! - **OpenAI**: `assistant.tool_calls[i].id` ↔ `tool.tool_call_id`.
//! - **Anthropic**: `assistant.content[].type=="tool_use".id`
//! ↔ `user.content[].type=="tool_result".tool_use_id`.
use std::collections::{HashMap, HashSet};
use serde_json::Value;
/// One paired (assistant_tool_use_index, tool_response_index) entry.
/// Either index may appear in multiple pairs if a single assistant
/// message issues multiple `tool_use` blocks that resolve in the
/// same following user message.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ToolPair {
pub assistant_index: usize,
pub response_index: usize,
}
/// Return every (assistant, response) tool-call pair in the message
/// list. Used by the live-zone dispatcher to keep tool_use and its
/// tool_result on the same compression decision.
pub fn tool_pair_indices(messages: &[Value]) -> Vec<ToolPair> {
// Map from tool_call id → assistant index that announced it.
let mut announced: HashMap<String, usize> = HashMap::new();
for (i, msg) in messages.iter().enumerate() {
if msg.get("role").and_then(Value::as_str) != Some("assistant") {
continue;
}
for id in collect_assistant_tool_call_ids(msg) {
announced.insert(id, i);
}
}
let mut pairs: Vec<ToolPair> = Vec::new();
let mut seen: HashSet<(usize, usize)> = HashSet::new();
for (i, msg) in messages.iter().enumerate() {
let role = msg.get("role").and_then(Value::as_str);
// OpenAI shape: role=tool, single tool_call_id.
if role == Some("tool") {
if let Some(tcid) = msg.get("tool_call_id").and_then(Value::as_str) {
if let Some(&assistant_index) = announced.get(tcid) {
if seen.insert((assistant_index, i)) {
pairs.push(ToolPair {
assistant_index,
response_index: i,
});
}
}
}
}
// Anthropic shape: role=user, content blocks may contain
// multiple tool_result entries pointing back to multiple
// tool_use ids on potentially different assistant messages.
if role == Some("user") {
if let Some(blocks) = msg.get("content").and_then(Value::as_array) {
for block in blocks {
if block.get("type").and_then(Value::as_str) == Some("tool_result") {
if let Some(tuid) = block.get("tool_use_id").and_then(Value::as_str) {
if let Some(&assistant_index) = announced.get(tuid) {
if seen.insert((assistant_index, i)) {
pairs.push(ToolPair {
assistant_index,
response_index: i,
});
}
}
}
}
}
}
}
}
pairs
}
/// Extract the set of tool_call ids announced by an assistant message,
/// handling both OpenAI and Anthropic shapes.
///
/// - **OpenAI**: `{role:assistant, tool_calls: [{id, ...}, ...]}`
/// - **Anthropic**: `{role:assistant, content: [{type:tool_use, id, ...}]}`
fn collect_assistant_tool_call_ids(assistant: &Value) -> HashSet<String> {
let mut ids: HashSet<String> = HashSet::new();
if let Some(arr) = assistant.get("tool_calls").and_then(Value::as_array) {
for tc in arr {
if let Some(id) = tc.get("id").and_then(Value::as_str) {
ids.insert(id.to_string());
}
}
}
if let Some(blocks) = assistant.get("content").and_then(Value::as_array) {
for block in blocks {
if block.get("type").and_then(Value::as_str) == Some("tool_use") {
if let Some(id) = block.get("id").and_then(Value::as_str) {
ids.insert(id.to_string());
}
}
}
}
ids
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn openai_pair_is_detected() {
let msgs = vec![
json!({"role": "user", "content": "go"}),
json!({
"role": "assistant",
"content": "",
"tool_calls": [{"id": "call_1", "type": "function",
"function": {"name": "f", "arguments": "{}"}}]
}),
json!({"role": "tool", "tool_call_id": "call_1", "content": "result"}),
json!({"role": "user", "content": "thanks"}),
];
let pairs = tool_pair_indices(&msgs);
assert_eq!(pairs.len(), 1);
assert_eq!(pairs[0].assistant_index, 1);
assert_eq!(pairs[0].response_index, 2);
}
#[test]
fn anthropic_pair_is_detected() {
let msgs = vec![
json!({"role": "user", "content": "go"}),
json!({
"role": "assistant",
"content": [
{"type": "text", "text": "thinking..."},
{"type": "tool_use", "id": "tu_1", "name": "f", "input": {}}
]
}),
json!({
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "tu_1", "content": "ok"}]
}),
];
let pairs = tool_pair_indices(&msgs);
assert_eq!(pairs.len(), 1);
assert_eq!(pairs[0].assistant_index, 1);
assert_eq!(pairs[0].response_index, 2);
}
#[test]
fn unmatched_tool_response_is_dropped() {
let msgs = vec![
json!({"role": "user", "content": "go"}),
json!({
"role": "assistant",
"content": "",
"tool_calls": [{"id": "call_known", "function": {"name": "f"}}]
}),
json!({"role": "tool", "tool_call_id": "call_orphan", "content": "?"}),
];
let pairs = tool_pair_indices(&msgs);
assert!(pairs.is_empty());
}
#[test]
fn multiple_anthropic_tool_results_in_one_user_message() {
let msgs = vec![
json!({
"role": "assistant",
"content": [
{"type": "tool_use", "id": "tu_a", "name": "f"},
{"type": "tool_use", "id": "tu_b", "name": "g"}
]
}),
json!({
"role": "user",
"content": [
{"type": "tool_result", "tool_use_id": "tu_a", "content": "a"},
{"type": "tool_result", "tool_use_id": "tu_b", "content": "b"}
]
}),
];
let pairs = tool_pair_indices(&msgs);
assert_eq!(pairs.len(), 1);
assert_eq!(pairs[0].assistant_index, 0);
assert_eq!(pairs[0].response_index, 1);
}
#[test]
fn empty_messages_yields_no_pairs() {
assert!(tool_pair_indices(&[]).is_empty());
}
}

125
crates/headroom-parity/src/lib.rs
View file

@ -462,130 +462,6 @@ impl TransformComparator for ContentDetectorComparator {
}
}
/// Real comparator for the `message_scorer` transform. Drives the
/// Rust port over fixture inputs and emits the same per-message
/// score struct shape the Python recorder dumps.
///
/// Float parity strategy: deterministic factors use `f32::exp` /
/// `f32::ln` on the Rust side and `math.exp(f64)` / `math.log(f64)`
/// on the Python side. Both are mathematically identical but their
/// last-bit rounding can differ. The Python recorder rounds every
/// float to 6 decimals before writing the fixture; this comparator
/// mirrors the same rounding so JSON byte-equality holds.
pub struct MessageScorerComparator;
impl TransformComparator for MessageScorerComparator {
fn name(&self) -> &str {
"message_scorer"
}
fn run(
&self,
input: &serde_json::Value,
config: &serde_json::Value,
) -> Result<serde_json::Value> {
use headroom_core::scoring::{MessageScorer, ScoringWeights};
use std::collections::HashSet;
let messages: Vec<serde_json::Value> = input
.get("messages")
.and_then(|v| v.as_array())
.cloned()
.context("message_scorer fixture input.messages must be an array")?;
let parse_index_set = |key: &str| -> HashSet<usize> {
input
.get(key)
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_u64().map(|n| n as usize))
.collect()
})
.unwrap_or_default()
};
let protected = parse_index_set("protected_indices");
let tool_unit = parse_index_set("tool_unit_indices");
let decay_rate = input
.get("decay_rate")
.and_then(|v| v.as_f64())
.map(|v| v as f32)
.unwrap_or(0.1);
// config.weights may be null (use defaults) or a full struct.
let weights = config
.get("weights")
.and_then(|v| if v.is_null() { None } else { Some(v) })
.map(|w| ScoringWeights {
recency: w.get("recency").and_then(|x| x.as_f64()).unwrap_or(0.20) as f32,
semantic_similarity: w
.get("semantic_similarity")
.and_then(|x| x.as_f64())
.unwrap_or(0.20) as f32,
toin_importance: w
.get("toin_importance")
.and_then(|x| x.as_f64())
.unwrap_or(0.25) as f32,
error_indicator: w
.get("error_indicator")
.and_then(|x| x.as_f64())
.unwrap_or(0.15) as f32,
forward_reference: w
.get("forward_reference")
.and_then(|x| x.as_f64())
.unwrap_or(0.15) as f32,
token_density: w
.get("token_density")
.and_then(|x| x.as_f64())
.unwrap_or(0.05) as f32,
});
let scorer = MessageScorer::new(weights, None, None, decay_rate);
let scores = scorer.score_messages(&messages, &protected, &tool_unit);
let scored_array: Vec<serde_json::Value> = scores
.into_iter()
.map(|s| serde_json::to_value(&s).expect("MessageScore is serializable"))
.collect();
// 5 decimal places matches the Python recorder. See
// record_message_scorer.py:_FLOAT_ROUND_PLACES for why.
Ok(round_floats(&serde_json::Value::Array(scored_array), 5))
}
}
/// Recursively round every f64 value in a JSON tree to `places`
/// decimal places. Used to absorb f32-vs-f64 last-bit drift between
/// the Python recorder and Rust comparator.
fn round_floats(value: &serde_json::Value, places: u32) -> serde_json::Value {
match value {
serde_json::Value::Number(n) => {
if let Some(f) = n.as_f64() {
if f.is_finite() && n.is_f64() {
let factor = 10f64.powi(places as i32);
let rounded = (f * factor).round() / factor;
return serde_json::Number::from_f64(rounded)
.map(serde_json::Value::Number)
.unwrap_or_else(|| value.clone());
}
}
value.clone()
}
serde_json::Value::Array(arr) => {
serde_json::Value::Array(arr.iter().map(|v| round_floats(v, places)).collect())
}
serde_json::Value::Object(map) => {
let mut out = serde_json::Map::new();
for (k, v) in map {
out.insert(k.clone(), round_floats(v, places));
}
serde_json::Value::Object(out)
}
_ => value.clone(),
}
}
/// Every built-in comparator, in a stable order.
pub fn builtin_comparators() -> Vec<Box<dyn TransformComparator>> {
vec![
@ -596,7 +472,6 @@ pub fn builtin_comparators() -> Vec<Box<dyn TransformComparator>> {
Box::new(CcrComparator),
Box::new(SmartCrusherComparator),
Box::new(ContentDetectorComparator),
Box::new(MessageScorerComparator),
]
}

8
headroom/__init__.py
View file

@ -100,11 +100,9 @@ __all__ = [
# Config
"HeadroomConfig",
"HeadroomMode",
"ToolCrusherConfig",
"SmartCrusherConfig",
"CacheAlignerConfig",
"CacheOptimizerConfig",
"RollingWindowConfig",
"RelevanceScorerConfig",
# Data models
"Block",
@ -116,10 +114,8 @@ __all__ = [
"TransformResult",
"WasteSignals",
# Transforms
"ToolCrusher",
"SmartCrusher",
"CacheAligner",
"RollingWindow",
"TransformPipeline",
# Cache optimizers
"BaseCacheOptimizer",
@ -207,11 +203,9 @@ _LAZY_EXPORTS: dict[str, tuple[str, str]] = {
# Config
"HeadroomConfig": ("headroom.config", "HeadroomConfig"),
"HeadroomMode": ("headroom.config", "HeadroomMode"),
"ToolCrusherConfig": ("headroom.config", "ToolCrusherConfig"),
"SmartCrusherConfig": ("headroom.config", "SmartCrusherConfig"),
"CacheAlignerConfig": ("headroom.config", "CacheAlignerConfig"),
"CacheOptimizerConfig": ("headroom.config", "CacheOptimizerConfig"),
"RollingWindowConfig": ("headroom.config", "RollingWindowConfig"),
"RelevanceScorerConfig": ("headroom.config", "RelevanceScorerConfig"),
# Data models
"Block": ("headroom.config", "Block"),
@ -223,10 +217,8 @@ _LAZY_EXPORTS: dict[str, tuple[str, str]] = {
"TransformResult": ("headroom.config", "TransformResult"),
"WasteSignals": ("headroom.config", "WasteSignals"),
# Transforms
"ToolCrusher": ("headroom.transforms", "ToolCrusher"),
"SmartCrusher": ("headroom.transforms", "SmartCrusher"),
"CacheAligner": ("headroom.transforms", "CacheAligner"),
"RollingWindow": ("headroom.transforms", "RollingWindow"),
"TransformPipeline": ("headroom.transforms", "TransformPipeline"),
# Cache optimizers
"BaseCacheOptimizer": ("headroom.cache", "BaseCacheOptimizer"),

23
headroom/cli/proxy.py
View file

@ -196,22 +196,6 @@ from .main import main
is_flag=True,
help="Disable Read lifecycle management (stale/superseded Read compression)",
)
# Intelligent Context Management (ON by default)
@click.option(
"--no-intelligent-context",
is_flag=True,
help="Disable IntelligentContextManager (fall back to RollingWindow)",
)
@click.option(
"--no-intelligent-scoring",
is_flag=True,
help="Disable multi-factor importance scoring (use position-based)",
)
@click.option(
"--no-compress-first",
is_flag=True,
help="Disable trying deeper compression before dropping messages",
)
# Memory System (Multi-Provider Support)
@click.option(
"--memory",
@ -377,9 +361,6 @@ def proxy(
budget: float | None,
code_graph: bool,
no_read_lifecycle: bool,
no_intelligent_context: bool,
no_intelligent_scoring: bool,
no_compress_first: bool,
memory: bool,
memory_db_path: str,
no_memory_tools: bool,
@ -537,10 +518,6 @@ def proxy(
code_graph_watcher=code_graph,
# Read lifecycle: ON by default (use --no-read-lifecycle to disable)
read_lifecycle=not no_read_lifecycle,
# Intelligent Context: ON by default (use --no-intelligent-context to disable)
intelligent_context=not no_intelligent_context,
intelligent_context_scoring=not no_intelligent_scoring,
intelligent_context_compress_first=not no_compress_first,
# Memory System (Multi-Provider with auto-detection)
# --learn implies --memory (need backend for storing patterns)
# Stateless mode disables memory (requires SQLite on disk)

16
headroom/client.py
View file

@ -283,9 +283,9 @@ class HeadroomClient:
enable_cache_optimizer=True, auto-detects from provider.
enable_cache_optimizer: Enable provider-specific cache optimization.
enable_semantic_cache: Enable query-level semantic caching.
config: Optional HeadroomConfig for full control over all settings
including intelligent_context. When provided, takes precedence
over individual settings like store_url, default_mode, etc.
config: Optional HeadroomConfig for full control over all settings.
When provided, takes precedence over individual settings like
store_url, default_mode, etc.
"""
self._original = original_client
self._provider = provider
@ -444,9 +444,7 @@ class HeadroomClient:
# Apply transforms if in optimize mode
if mode == HeadroomMode.OPTIMIZE:
output_buffer = (
headroom_output_buffer_tokens or self._config.rolling_window.output_buffer_tokens
)
output_buffer = headroom_output_buffer_tokens or self._config.output_buffer_tokens
model_limit = self._get_context_limit(model)
result = self._pipeline.apply(
@ -774,9 +772,7 @@ class HeadroomClient:
compute_prefix_hash(messages)
# Apply transforms
output_buffer = (
headroom_output_buffer_tokens or self._config.rolling_window.output_buffer_tokens
)
output_buffer = headroom_output_buffer_tokens or self._config.output_buffer_tokens
model_limit = self._get_context_limit(model)
result = self._pipeline.simulate(
@ -986,7 +982,6 @@ class HeadroomClient:
},
"transforms": {
"smart_crusher_enabled": bool,
"rolling_window_enabled": bool,
"cache_aligner_enabled": bool,
},
}
@ -1015,7 +1010,6 @@ class HeadroomClient:
},
"transforms": {
"smart_crusher_enabled": self._config.smart_crusher.enabled,
"rolling_window_enabled": self._config.rolling_window.enabled,
"cache_aligner_enabled": self._config.cache_aligner.enabled,
},
}

5
headroom/compress.py
View file

@ -337,11 +337,12 @@ def _get_pipeline() -> Any:
from headroom.transforms import TransformPipeline
# Default pipeline: CacheAligner → ContentRouter → IntelligentContext
# Default pipeline: CacheAligner → ContentRouter
# CacheAligner: stabilizes prefix for provider KV cache hits
# ContentRouter: routes to the right compressor per content type
# (SmartCrusher for JSON, CodeCompressor for code, Kompress for text)
# IntelligentContext: enforces token limits with score-based dropping
# Phase B PR-B1 retired the trailing context-management stage —
# live-zone-only compression never drops messages.
_pipeline = TransformPipeline()
logger.debug("Headroom compression pipeline initialized")
return _pipeline

163
headroom/config.py
View file

@ -23,30 +23,6 @@ class HeadroomMode(str, Enum):
DEFAULT_MODEL_CONTEXT_LIMITS: dict[str, int] = {}
@dataclass
class ToolCrusherConfig:
"""Configuration for tool output compression (naive/fixed-rule approach).
GOTCHAS:
- Keeps FIRST N items only - may miss important data later in arrays
- A spike at index 50 will be lost if max_array_items=10
- String truncation cuts at fixed length, may break mid-word/mid-sentence
- No awareness of data patterns or importance
Consider using SmartCrusherConfig instead for statistical analysis.
"""
enabled: bool = False # Disabled by default, SmartCrusher is preferred
min_tokens_to_crush: int = 500 # Only crush if > N tokens
max_array_items: int = 10 # Keep first N items
max_string_length: int = 1000 # Truncate strings > N chars
max_depth: int = 5 # Preserve structure to depth N
preserve_keys: set[str] = field(
default_factory=lambda: {"error", "status", "code", "id", "message", "name", "type"}
)
tool_profiles: dict[str, dict[str, Any]] = field(default_factory=dict)
@dataclass
class CacheAlignerConfig:
"""Configuration for cache alignment.
@ -122,130 +98,6 @@ class CacheAlignerConfig:
dynamic_tail_separator: str = "\n\n---\n[Dynamic Context]\n"
@dataclass
class RollingWindowConfig:
"""Configuration for rolling window token cap.
GOTCHAS:
- Dropping old turns loses context the model may need:
- "As I mentioned earlier..." - what was mentioned is now gone
- "The user asked about X" - that turn may be dropped
- Implicit references to prior conversation become orphaned
- Tool call/result pairs are kept atomic (correct), BUT:
- Assistant text referencing a dropped tool result becomes confusing
- "Based on the search results..." when those results are gone
- keep_last_turns=2 may not be enough for complex multi-step reasoning
- No semantic analysis - drops oldest first regardless of importance
SAFER ALTERNATIVES:
- Increase keep_last_turns for agentic workloads
- Use summarization for old context (not implemented - would add latency)
- Set enabled=False for short conversations
"""
enabled: bool = True
keep_system: bool = True # Never drop system prompt
keep_last_turns: int = 2 # Never drop last N turns
output_buffer_tokens: int = 4000 # Reserve for output
@dataclass
class ScoringWeights:
"""Weights for importance scoring factors.
All weights should sum to approximately 1.0 for normalized scoring.
These can be learned from TOIN retrieval patterns over time.
Design principle: NO HARDCODED PATTERNS. All importance is derived from:
- Computed metrics (recency, density, references)
- TOIN-learned patterns (field semantics, retrieval rates)
- Embedding similarity (semantic relevance)
"""
recency: float = 0.20 # Exponential decay from conversation end
semantic_similarity: float = 0.20 # Embedding similarity to recent context
toin_importance: float = 0.25 # TOIN-learned field importance
error_indicator: float = 0.15 # TOIN-learned error field detection
forward_reference: float = 0.15 # Referenced by later messages
token_density: float = 0.05 # Information density (entropy-based)
def normalized(self) -> ScoringWeights:
"""Return a copy with weights normalized to sum to 1.0."""
total = (
self.recency
+ self.semantic_similarity
+ self.toin_importance
+ self.error_indicator
+ self.forward_reference
+ self.token_density
)
if total == 0:
return ScoringWeights()
return ScoringWeights(
recency=self.recency / total,
semantic_similarity=self.semantic_similarity / total,
toin_importance=self.toin_importance / total,
error_indicator=self.error_indicator / total,
forward_reference=self.forward_reference / total,
token_density=self.token_density / total,
)
@dataclass
class IntelligentContextConfig:
"""Configuration for intelligent context management.
This extends RollingWindowConfig with semantic-aware scoring and
TOIN integration. All importance detection is learned, not hardcoded.
Phases:
- Phase 1 (current): Importance scoring + TOIN integration
- Phase 2 (future): Progressive summarization
- Phase 3 (future): Memory tiers with retrieval
"""
# === Basic settings (backwards compatible with RollingWindowConfig) ===
enabled: bool = True
keep_system: bool = True
keep_last_turns: int = 2
output_buffer_tokens: int = 4000
# === Scoring configuration ===
use_importance_scoring: bool = True
scoring_weights: ScoringWeights = field(default_factory=ScoringWeights)
# Recency decay parameter (higher = faster decay)
recency_decay_rate: float = 0.1
# === TOIN integration ===
toin_integration: bool = True
toin_confidence_threshold: float = 0.3 # Min confidence to use TOIN signals
# === Strategy selection thresholds ===
# These determine when to try different strategies based on how much over budget
compress_threshold: float = 0.10 # Try deeper compression if <10% over
# === Summarization (Phase 2 - not yet implemented) ===
summarization_enabled: bool = False
summarization_model: str | None = None
summary_max_tokens: int = 500
summarize_threshold: float = 0.25 # Try summarization if <25% over
# === Memory tiers (Phase 3 - not yet implemented) ===
memory_tiers_enabled: bool = False
warm_tier_enabled: bool = False # Summarized tier
cold_tier_enabled: bool = False # Vector retrieval tier
def to_rolling_window_config(self) -> RollingWindowConfig:
"""Convert to basic RollingWindowConfig for backwards compatibility."""
return RollingWindowConfig(
enabled=self.enabled,
keep_system=self.keep_system,
keep_last_turns=self.keep_last_turns,
output_buffer_tokens=self.output_buffer_tokens,
)
@dataclass
class RelevanceScorerConfig:
"""Configuration for relevance scoring in SmartCrusher.
@ -471,7 +323,7 @@ class SmartCrusherConfig:
- Set variance_threshold lower (1.5) to catch more change points
"""
enabled: bool = True # Enabled by default (preferred over ToolCrusher)
enabled: bool = True # Enabled by default — sole tool-output compressor
min_items_to_analyze: int = 5 # Don't analyze tiny arrays
min_tokens_to_crush: int = 200 # Only crush if > N tokens
variance_threshold: float = 2.0 # Std devs for change point detection
@ -610,25 +462,22 @@ class HeadroomConfig:
model_context_limits: dict[str, int] = field(
default_factory=lambda: DEFAULT_MODEL_CONTEXT_LIMITS.copy()
)
tool_crusher: ToolCrusherConfig = field(default_factory=ToolCrusherConfig)
smart_crusher: SmartCrusherConfig = field(default_factory=SmartCrusherConfig)
cache_aligner: CacheAlignerConfig = field(default_factory=CacheAlignerConfig)
rolling_window: RollingWindowConfig = field(default_factory=RollingWindowConfig)
cache_optimizer: CacheOptimizerConfig = field(default_factory=CacheOptimizerConfig)
ccr: CCRConfig = field(default_factory=CCRConfig) # Compress-Cache-Retrieve
prefix_freeze: PrefixFreezeConfig = field(default_factory=PrefixFreezeConfig)
# Output buffer reserved for the model's response when sizing the
# incoming context. Previously lived on RollingWindowConfig; hoisted
# to the top-level config when PR-B1 retired the rolling-window stage.
output_buffer_tokens: int = 4000
# Content Router - intelligent content-type based compression
# Routes content to appropriate compressor (Kompress for text, SmartCrusher for JSON,
# CodeCompressor for code, LogCompressor for logs, etc.)
content_router_enabled: bool = True
# Intelligent context management (Phase 2.5)
# When enabled, replaces RollingWindow with semantic-aware context management
intelligent_context: IntelligentContextConfig = field(
default_factory=lambda: IntelligentContextConfig(enabled=False)
)
# Tool-result interceptors (ast-grep Read outline, etc.). Opt-in for now.
# Env var HEADROOM_INTERCEPT_ENABLED=1 also enables (for CLI `--intercept-tool-results`).
intercept_tool_results: bool = False

2
headroom/exceptions.py
View file

@ -179,7 +179,7 @@ class TransformError(HeadroomError):
This includes:
- SmartCrusher failures
- RollingWindow errors
- ContentRouter errors
- Pipeline errors
Example:

23
headroom/integrations/langchain/memory.py
View file

@ -42,7 +42,6 @@ except ImportError:
BaseChatMessageHistory = object # type: ignore[misc,assignment]
from headroom import HeadroomConfig
from headroom.config import RollingWindowConfig
from headroom.providers import OpenAIProvider
from headroom.transforms import TransformPipeline
@ -63,8 +62,9 @@ class HeadroomChatMessageHistory(BaseChatMessageHistory):
"""Wraps any LangChain chat message history with automatic compression.
When conversation history exceeds the token threshold, automatically
applies RollingWindow compression to keep recent turns while fitting
within the limit.
applies live-zone block compression (per-block content compression on
the live zone, never dropping messages — that's what the live-zone
refactor in PR-B1+ replaces the old RollingWindow strategy with).
This works with ANY memory type because it wraps at the storage layer:
- ConversationBufferMemory
@ -150,7 +150,7 @@ class HeadroomChatMessageHistory(BaseChatMessageHistory):
return list(raw_messages)
# Apply compression
compressed = self._apply_rolling_window(raw_messages)
compressed = self._apply_compression(raw_messages)
tokens_after = self._count_tokens(compressed)
self._compression_count += 1
@ -207,22 +207,25 @@ class HeadroomChatMessageHistory(BaseChatMessageHistory):
total += token_counter.count_text(content)
return total
def _apply_rolling_window(self, messages: list[BaseMessage]) -> list[BaseMessage]:
"""Apply RollingWindow compression to messages.
def _apply_compression(self, messages: list[BaseMessage]) -> list[BaseMessage]:
"""Apply live-zone-only compression to messages.
After PR-B1, message-dropping is no longer a strategy — only
per-block content compression. Result may still exceed the
threshold; callers should treat the threshold as advisory.
Args:
messages: Messages to compress.
Returns:
Compressed messages fitting within threshold.
Messages with their live-zone content compressed where
applicable.
"""
# Convert to OpenAI format for Headroom transforms
openai_messages = self._convert_to_openai(messages)
# Use TransformPipeline which handles tokenizer setup
config = HeadroomConfig(
rolling_window=RollingWindowConfig(keep_last_turns=self._keep_recent_turns),
)
config = HeadroomConfig()
pipeline = TransformPipeline(config=config, provider=self._provider)
# Apply compression via pipeline

6
headroom/providers/copilot/wrap.py
View file

@ -78,7 +78,11 @@ def build_launch_env(
environ: Mapping[str, str] | None = None,
) -> tuple[dict[str, str], list[str]]:
"""Build the Copilot BYOK environment for the selected provider type."""
env = dict(environ or os.environ)
# Distinguish "caller passed nothing" (use os.environ) from "caller
# explicitly passed an empty dict" (start fresh — the test/CLI is in
# charge of which keys to seed). The previous `environ or os.environ`
# collapsed those two cases because `bool({}) is False`.
env = dict(environ if environ is not None else os.environ)
env["COPILOT_PROVIDER_TYPE"] = provider_type
env.pop("COPILOT_PROVIDER_WIRE_API", None)

5
headroom/proxy/models.py
View file

@ -142,11 +142,6 @@ class ProxyConfig:
# Smart content routing
smart_routing: bool = True
# Intelligent context management
intelligent_context: bool = True
intelligent_context_scoring: bool = True
intelligent_context_compress_first: bool = True
# Caching
cache_enabled: bool = True
cache_ttl_seconds: int = 3600

51
headroom/proxy/server.py
View file

@ -4,7 +4,7 @@ A full-featured LLM proxy with optimization, caching, rate limiting,
and observability.
Features:
- Context optimization (SmartCrusher, CacheAligner, RollingWindow)
- Context optimization (SmartCrusher, CacheAligner — live-zone-only after Phase B)
- Semantic caching (save costs on repeated queries)
- Rate limiting (token bucket)
- Retry with exponential backoff
@ -75,9 +75,7 @@ from headroom.ccr import (
from headroom.config import (
CacheAlignerConfig,
CCRConfig,
IntelligentContextConfig,
ReadLifecycleConfig,
RollingWindowConfig,
SmartCrusherConfig,
)
from headroom.dashboard import get_dashboard_html
@ -159,10 +157,7 @@ from headroom.transforms import (
CodeCompressorConfig,
ContentRouter,
ContentRouterConfig,
IntelligentContextManager,
RollingWindow,
SmartCrusher,
Transform,
TransformPipeline,
is_tree_sitter_available,
)
@ -255,34 +250,14 @@ class HeadroomProxy(
)
self.metrics = PrometheusMetrics(cost_tracker=self.cost_tracker)
# Initialize transforms based on routing mode
# Choose context manager: IntelligentContextManager (smart) or RollingWindow (legacy)
context_manager: Transform # Can be either IntelligentContextManager or RollingWindow
if config.intelligent_context:
# Get TOIN instance for learned pattern integration
toin = get_toin() if config.intelligent_context_scoring else None
context_manager = IntelligentContextManager(
config=IntelligentContextConfig(
enabled=True,
keep_system=True,
keep_last_turns=config.keep_last_turns,
use_importance_scoring=config.intelligent_context_scoring,
toin_integration=config.intelligent_context_scoring,
compress_threshold=0.10 if config.intelligent_context_compress_first else 0.0,
),
toin=toin,
observer=self.metrics,
)
self._context_manager_status = "intelligent"
else:
context_manager = RollingWindow(
RollingWindowConfig(
enabled=True,
keep_system=True,
keep_last_turns=config.keep_last_turns,
)
)
self._context_manager_status = "rolling_window"
# Initialize transforms based on routing mode.
#
# Phase B PR-B1 retired the IntelligentContextManager / RollingWindow
# message-dropping branch. Live-zone-only compression (PR-B2..B7) does
# not drop messages — it operates on content blocks within messages —
# so the proxy no longer needs a "context manager" transform stage.
# Reported via metrics as `_context_manager_status = "passthrough"`.
self._context_manager_status = "passthrough"
if config.smart_routing:
# Smart routing: ContentRouter handles all content types intelligently
@ -298,7 +273,6 @@ class HeadroomProxy(
transforms = [
CacheAligner(CacheAlignerConfig(enabled=False)),
ContentRouter(router_config, observer=self.metrics),
context_manager,
]
self._code_aware_status = "lazy" if config.code_aware_enabled else "disabled"
else:
@ -317,7 +291,6 @@ class HeadroomProxy(
),
observer=self.metrics,
),
context_manager,
]
# Add CodeAware if enabled and available
self._code_aware_status = self._setup_code_aware(config, transforms)
@ -712,8 +685,10 @@ class HeadroomProxy(
preserve_signatures=True,
preserve_type_annotations=True,
)
# Insert before RollingWindow (which should be last)
transforms.insert(-1, CodeAwareCompressor(code_config))
# CodeAware runs after the content/structure transforms.
# Phase B PR-B1 retired the trailing context_manager so we
# append rather than insert(-1).
transforms.append(CodeAwareCompressor(code_config))
return "enabled"
else:
logger.warning(

30
headroom/transforms/__init__.py
View file

@ -49,29 +49,18 @@ if TYPE_CHECKING:
HTMLExtractorConfig,
is_html_content,
)
from headroom.transforms.intelligent_context import ( # noqa: F401
ContextStrategy,
IntelligentContextManager,
)
from headroom.transforms.log_compressor import ( # noqa: F401
LogCompressionResult,
LogCompressor,
LogCompressorConfig,
)
from headroom.transforms.pipeline import TransformPipeline # noqa: F401
from headroom.transforms.rolling_window import RollingWindow # noqa: F401
from headroom.transforms.scoring import ( # noqa: F401
EmbeddingProvider,
MessageScore,
MessageScorer,
)
from headroom.transforms.search_compressor import ( # noqa: F401
SearchCompressionResult,
SearchCompressor,
SearchCompressorConfig,
)
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig # noqa: F401
from headroom.transforms.tool_crusher import ToolCrusher # noqa: F401
_HTML_EXTRACTOR_AVAILABLE = importlib.util.find_spec("trafilatura") is not None
@ -87,7 +76,6 @@ __all__ = [
"calculate_information_score",
"compute_item_hash",
# JSON compression
"ToolCrusher",
"SmartCrusher",
"SmartCrusherConfig",
# Text compression (coding tasks)
@ -118,13 +106,6 @@ __all__ = [
"CompressionStrategy",
# Other transforms
"CacheAligner",
"RollingWindow",
# Intelligent context management
"IntelligentContextManager",
"ContextStrategy",
"MessageScorer",
"MessageScore",
"EmbeddingProvider",
# HTML extraction (optional)
"_HTML_EXTRACTOR_AVAILABLE",
]
@ -155,7 +136,6 @@ _LAZY_EXPORTS: dict[str, tuple[str, str]] = {
),
"compute_item_hash": ("headroom.transforms.anchor_selector", "compute_item_hash"),
# JSON compression
"ToolCrusher": ("headroom.transforms.tool_crusher", "ToolCrusher"),
"SmartCrusher": ("headroom.transforms.smart_crusher", "SmartCrusher"),
"SmartCrusherConfig": ("headroom.transforms.smart_crusher", "SmartCrusherConfig"),
# Text compression (coding tasks)
@ -204,16 +184,6 @@ _LAZY_EXPORTS: dict[str, tuple[str, str]] = {
"CompressionStrategy": ("headroom.transforms.content_router", "CompressionStrategy"),
# Other transforms
"CacheAligner": ("headroom.transforms.cache_aligner", "CacheAligner"),
"RollingWindow": ("headroom.transforms.rolling_window", "RollingWindow"),
# Intelligent context management
"IntelligentContextManager": (
"headroom.transforms.intelligent_context",
"IntelligentContextManager",
),
"ContextStrategy": ("headroom.transforms.intelligent_context", "ContextStrategy"),
"MessageScorer": ("headroom.transforms.scoring", "MessageScorer"),
"MessageScore": ("headroom.transforms.scoring", "MessageScore"),
"EmbeddingProvider": ("headroom.transforms.scoring", "EmbeddingProvider"),
# HTML extraction (optional dependency - requires trafilatura)
"HTMLExtractor": ("headroom.transforms.html_extractor", "HTMLExtractor"),
"HTMLExtractorConfig": ("headroom.transforms.html_extractor", "HTMLExtractorConfig"),

2
headroom/transforms/content_router.py
View file

@ -29,7 +29,6 @@ Usage:
Pipeline Usage:
>>> pipeline = TransformPipeline([
... ContentRouter(), # Handles all content types
... RollingWindow(), # Final size constraint
... ])
"""
@ -615,7 +614,6 @@ class ContentRouter(Transform):
Pipeline Integration:
>>> pipeline = TransformPipeline([
... ContentRouter(), # Handles ALL content types
... RollingWindow(), # Final size constraint
... ])
"""

1077
headroom/transforms/intelligent_context.py
View file

File diff suppressed because it is too large Load diff

44
headroom/transforms/pipeline.py
View file

@ -11,9 +11,6 @@ from ..config import (
CacheAlignerConfig,
DiffArtifact,
HeadroomConfig,
IntelligentContextConfig,
RollingWindowConfig,
ToolCrusherConfig,
TransformDiff,
TransformResult,
WasteSignals,
@ -24,10 +21,7 @@ from ..utils import deep_copy_messages
from .base import Transform
from .cache_aligner import CacheAligner
from .content_router import ContentRouter
from .intelligent_context import IntelligentContextManager
from .rolling_window import RollingWindow
from .smart_crusher import SmartCrusher
from .tool_crusher import ToolCrusher
if TYPE_CHECKING:
from ..providers.base import Provider
@ -43,8 +37,12 @@ class TransformPipeline:
1. Cache Aligner - normalize prefix for cache hits
2. Content Router - intelligent content-aware compression (routes to appropriate
compressor: Kompress for text, SmartCrusher for JSON, CodeCompressor for code, etc.)
3. SmartCrusher/ToolCrusher - fallback if ContentRouter disabled
4. IntelligentContextManager/RollingWindow - enforce token limits
3. SmartCrusher - fallback if ContentRouter disabled
Phase B PR-B1 retired the IntelligentContextManager / RollingWindow
"drop messages from history" stage. Live-zone-only compression is the
sole strategy going forward — message-list mutation no longer happens
in the pipeline.
"""
def __init__(
@ -122,22 +120,6 @@ class TransformPipeline:
include_summaries=self.config.smart_crusher.include_summaries,
)
transforms.append(SmartCrusher(smart_config))
elif self.config.tool_crusher.enabled:
# Fallback to fixed-rule crushing
transforms.append(ToolCrusher(self.config.tool_crusher))
# 3. Context Management (enforce limits last)
# IntelligentContextManager takes precedence over RollingWindow when enabled
if self.config.intelligent_context.enabled:
# Use semantic-aware context management with scoring
transforms.append(IntelligentContextManager(self.config.intelligent_context))
logger.info(
"Pipeline using IntelligentContextManager with strategies: "
"COMPRESS_FIRST -> SUMMARIZE -> DROP_BY_SCORE"
)
elif self.config.rolling_window.enabled:
# Fallback to position-based rolling window
transforms.append(RollingWindow(self.config.rolling_window))
return transforms
@ -457,34 +439,20 @@ class TransformPipeline:
def create_pipeline(
tool_crusher_config: ToolCrusherConfig | None = None,
cache_aligner_config: CacheAlignerConfig | None = None,
rolling_window_config: RollingWindowConfig | None = None,
intelligent_context_config: IntelligentContextConfig | None = None,
) -> TransformPipeline:
"""
Create a pipeline with specific configurations.
Args:
tool_crusher_config: Tool crusher configuration.
cache_aligner_config: Cache aligner configuration.
rolling_window_config: Rolling window configuration.
intelligent_context_config: Intelligent context configuration.
When provided with enabled=True, replaces RollingWindow with
semantic-aware context management.
Returns:
Configured TransformPipeline.
"""
config = HeadroomConfig()
if tool_crusher_config is not None:
config.tool_crusher = tool_crusher_config
if cache_aligner_config is not None:
config.cache_aligner = cache_aligner_config
if rolling_window_config is not None:
config.rolling_window = rolling_window_config
if intelligent_context_config is not None:
config.intelligent_context = intelligent_context_config
return TransformPipeline(config)

508
headroom/transforms/progressive_summarizer.py
View file

@ -1,508 +0,0 @@
"""Progressive summarization for Headroom SDK.
This module provides anchored summarization that progressively summarizes
older messages while maintaining retrieval capability via CCR.
Design principles:
1. CALLBACK PATTERN: Summarization is done via a callback, not internal LLM calls
2. ANCHORED: Summaries track which message positions they represent
3. REVERSIBLE: Original content stored in CompressionStore for CCR retrieval
4. INCREMENTAL: Only summarize newly dropped spans, then merge
Usage:
from headroom.transforms import ProgressiveSummarizer
# With custom summarizer callback
def my_summarizer(messages: list[dict], context: str) -> str:
# Your summarization logic (LLM call, extractive, etc.)
return "Summary of messages..."
summarizer = ProgressiveSummarizer(
summarize_fn=my_summarizer,
max_summary_tokens=500,
)
result = summarizer.summarize_messages(messages, tokenizer, protected)
"""
from __future__ import annotations
import hashlib
import json
import logging
import time
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Protocol
if TYPE_CHECKING:
from ..cache.compression_store import CompressionStore
from ..tokenizer import Tokenizer
logger = logging.getLogger(__name__)
class SummarizeFn(Protocol):
"""Protocol for summarization callback functions.
The callback receives:
- messages: List of messages to summarize
- context: Optional context string (e.g., recent messages for relevance)
Returns:
- Summary string
"""
def __call__(
self,
messages: list[dict[str, Any]],
context: str = "",
) -> str: ...
@dataclass
class AnchoredSummary:
"""A summary anchored to specific message positions.
Tracks which messages were summarized for:
- Retrieval: Can reconstruct original messages via CCR
- Merging: Can merge with adjacent summaries
- Positioning: Know where in conversation this summary belongs
"""
summary_text: str
start_index: int # First message index summarized
end_index: int # Last message index summarized (inclusive)
original_message_count: int
original_tokens: int
summary_tokens: int
cache_hash: str | None = None # Hash for CCR retrieval
tool_names: list[str] = field(default_factory=list)
created_at: float = field(default_factory=time.time)
@property
def compression_ratio(self) -> float:
"""Ratio of summary tokens to original tokens (lower = more compression)."""
if self.original_tokens == 0:
return 1.0
return self.summary_tokens / self.original_tokens
@property
def tokens_saved(self) -> int:
"""Number of tokens saved by summarization."""
return max(0, self.original_tokens - self.summary_tokens)
@dataclass
class SummarizationResult:
"""Result of a summarization operation."""
messages: list[dict[str, Any]]
summaries_created: list[AnchoredSummary]
tokens_before: int
tokens_after: int
transforms_applied: list[str]
@property
def tokens_saved(self) -> int:
"""Total tokens saved."""
return max(0, self.tokens_before - self.tokens_after)
def extractive_summarizer(
messages: list[dict[str, Any]],
context: str = "",
max_items_per_role: int = 2,
) -> str:
"""Default extractive summarizer (no LLM required).
Creates a summary by extracting key content from messages:
- First and last message of each role
- Error indicators
- Tool names and brief results
This is a fallback when no LLM summarizer is provided.
Args:
messages: Messages to summarize.
context: Optional context (unused in extractive mode).
max_items_per_role: Max items to keep per role type.
Returns:
Extractive summary string.
"""
if not messages:
return "[No messages to summarize]"
parts: list[str] = []
parts.append(f"[Summary of {len(messages)} messages]")
# Group by role
by_role: dict[str, list[dict[str, Any]]] = {}
for msg in messages:
role = msg.get("role", "unknown")
by_role.setdefault(role, []).append(msg)
# Extract key content from each role
for role, role_msgs in by_role.items():
if role == "tool":
# For tool messages, extract tool names and brief status
tool_names = set()
has_error = False
for msg in role_msgs:
content = msg.get("content", "")
# Try to detect tool name from context
tool_call_id = msg.get("tool_call_id", "")
if tool_call_id:
tool_names.add(f"tool:{tool_call_id[:8]}")
# Check for errors
content_lower = content.lower() if isinstance(content, str) else ""
if any(err in content_lower for err in ["error", "failed", "exception"]):
has_error = True
status = "with errors" if has_error else "successful"
parts.append(f"- {len(role_msgs)} tool outputs ({status})")
elif role == "assistant":
# Extract first and last assistant responses
if len(role_msgs) == 1:
content = role_msgs[0].get("content", "")
if isinstance(content, str):
preview = content[:100] + "..." if len(content) > 100 else content
parts.append(f"- Assistant: {preview}")
else:
parts.append(f"- {len(role_msgs)} assistant messages")
elif role == "user":
# Count user messages
parts.append(f"- {len(role_msgs)} user messages")
elif role == "system":
# Note system messages (shouldn't be summarized usually)
parts.append(f"- {len(role_msgs)} system messages")
return "\n".join(parts)
class ProgressiveSummarizer:
"""Progressive summarization with anchoring and CCR integration.
This class implements the SUMMARIZE strategy for IntelligentContextManager:
1. Identifies candidate messages (low-scored, non-protected)
2. Groups consecutive messages for summarization
3. Calls summarizer callback to create summaries
4. Stores originals in CompressionStore for CCR retrieval
5. Replaces messages with anchored summary message
Key features:
- Callback pattern: No LLM calls inside, summarization logic is external
- Anchored: Summaries track original positions for context
- Reversible: Originals cached for retrieval
- Incremental: Can merge adjacent summaries
"""
def __init__(
self,
summarize_fn: SummarizeFn | None = None,
max_summary_tokens: int = 500,
min_messages_to_summarize: int = 3,
compression_store: CompressionStore | None = None,
store_for_retrieval: bool = True,
):
"""Initialize the progressive summarizer.
Args:
summarize_fn: Callback function for summarization.
If None, uses extractive_summarizer as fallback.
max_summary_tokens: Target max tokens for each summary.
min_messages_to_summarize: Minimum messages in a group to summarize.
compression_store: Optional CompressionStore for CCR integration.
store_for_retrieval: Whether to store originals for retrieval.
"""
self.summarize_fn = summarize_fn or extractive_summarizer
self.max_summary_tokens = max_summary_tokens
self.min_messages_to_summarize = min_messages_to_summarize
self._compression_store = compression_store
self.store_for_retrieval = store_for_retrieval
def _get_compression_store(self) -> CompressionStore | None:
"""Get or create compression store (lazy load)."""
if self._compression_store is None and self.store_for_retrieval:
try:
from ..cache.compression_store import get_compression_store
self._compression_store = get_compression_store()
except ImportError:
logger.debug("CompressionStore not available for CCR")
return self._compression_store
def summarize_messages(
self,
messages: list[dict[str, Any]],
tokenizer: Tokenizer,
protected_indices: set[int],
target_tokens: int | None = None,
context_messages: list[dict[str, Any]] | None = None,
) -> SummarizationResult:
"""Summarize messages to reduce token count.
Args:
messages: List of messages to process.
tokenizer: Tokenizer for counting.
protected_indices: Indices that cannot be summarized.
target_tokens: Target token count (optional, summarizes all candidates if None).
context_messages: Recent messages for context in summarization.
Returns:
SummarizationResult with summarized messages.
"""
from ..utils import deep_copy_messages
tokens_before = tokenizer.count_messages(messages)
result_messages = deep_copy_messages(messages)
transforms_applied: list[str] = []
summaries_created: list[AnchoredSummary] = []
# Find candidate groups for summarization
candidate_groups = self._find_summarization_candidates(result_messages, protected_indices)
if not candidate_groups:
logger.debug("ProgressiveSummarizer: no candidates for summarization")
return SummarizationResult(
messages=result_messages,
summaries_created=[],
tokens_before=tokens_before,
tokens_after=tokens_before,
transforms_applied=[],
)
# Build context string from recent messages
context_str = ""
if context_messages:
context_parts = []
for msg in context_messages[-3:]: # Last 3 messages for context
role = msg.get("role", "")
content = msg.get("content", "")
if isinstance(content, str) and content:
preview = content[:200] if len(content) > 200 else content
context_parts.append(f"{role}: {preview}")
context_str = "\n".join(context_parts)
# Process groups in reverse order (so indices stay valid)
current_tokens = tokens_before
for group in reversed(candidate_groups):
# Check if we've reached target
if target_tokens and current_tokens <= target_tokens:
break
start_idx, end_idx = group
group_messages = result_messages[start_idx : end_idx + 1]
# Skip if too few messages
if len(group_messages) < self.min_messages_to_summarize:
continue
# Calculate group tokens
group_tokens = sum(tokenizer.count_message(msg) for msg in group_messages)
# Skip small groups
if group_tokens < 100:
continue
# Create summary using callback
try:
summary_text = self.summarize_fn(group_messages, context_str)
except Exception as e:
logger.warning(
"ProgressiveSummarizer: summarization failed for group %d-%d: %s",
start_idx,
end_idx,
e,
)
continue
summary_tokens = tokenizer.count_text(summary_text)
# Only use summary if it saves tokens
if summary_tokens >= group_tokens:
logger.debug(
"ProgressiveSummarizer: summary not smaller (%d >= %d), skipping",
summary_tokens,
group_tokens,
)
continue
# Store original for CCR retrieval
cache_hash = None
if self.store_for_retrieval:
cache_hash = self._store_for_retrieval(
group_messages, summary_text, group_tokens, summary_tokens
)
# Extract tool names
tool_names = []
for msg in group_messages:
if msg.get("role") == "tool":
tool_call_id = msg.get("tool_call_id", "")
if tool_call_id:
tool_names.append(tool_call_id[:8])
# Create anchored summary
anchored = AnchoredSummary(
summary_text=summary_text,
start_index=start_idx,
end_index=end_idx,
original_message_count=len(group_messages),
original_tokens=group_tokens,
summary_tokens=summary_tokens,
cache_hash=cache_hash,
tool_names=tool_names,
)
summaries_created.append(anchored)
# Create summary message with retrieval marker
summary_content = summary_text
if cache_hash:
summary_content += f"\n[Retrieve full content: hash={cache_hash}]"
summary_message = {
"role": "user",
"content": summary_content,
}
# Replace group with summary message
result_messages = (
result_messages[:start_idx] + [summary_message] + result_messages[end_idx + 1 :]
)
# Update token count
tokens_saved = group_tokens - summary_tokens
current_tokens -= tokens_saved
transforms_applied.append(f"summarize:{start_idx}-{end_idx}:{len(group_messages)}")
logger.debug(
"ProgressiveSummarizer: summarized %d messages (%d-%d), saved %d tokens (%d -> %d)",
len(group_messages),
start_idx,
end_idx,
tokens_saved,
group_tokens,
summary_tokens,
)
# Update protected indices for subsequent groups
# (indices shift after replacement)
shift = len(group_messages) - 1 # We replaced N messages with 1
protected_indices = {idx - shift if idx > end_idx else idx for idx in protected_indices}
tokens_after = tokenizer.count_messages(result_messages)
if summaries_created:
logger.info(
"ProgressiveSummarizer: created %d summaries, saved %d tokens (%d -> %d)",
len(summaries_created),
tokens_before - tokens_after,
tokens_before,
tokens_after,
)
return SummarizationResult(
messages=result_messages,
summaries_created=summaries_created,
tokens_before=tokens_before,
tokens_after=tokens_after,
transforms_applied=transforms_applied,
)
def _find_summarization_candidates(
self,
messages: list[dict[str, Any]],
protected: set[int],
) -> list[tuple[int, int]]:
"""Find groups of consecutive messages that can be summarized.
Returns list of (start_index, end_index) tuples for candidate groups.
Groups are consecutive non-protected messages.
Args:
messages: List of messages.
protected: Set of protected indices.
Returns:
List of (start, end) tuples for candidate groups.
"""
groups: list[tuple[int, int]] = []
current_start: int | None = None
for i, _msg in enumerate(messages):
if i in protected:
# End current group if exists
if current_start is not None:
if i - 1 >= current_start:
groups.append((current_start, i - 1))
current_start = None
else:
# Start or continue group
if current_start is None:
current_start = i
# Handle final group
if current_start is not None and len(messages) - 1 >= current_start:
groups.append((current_start, len(messages) - 1))
# Filter groups that are too small
groups = [
(start, end)
for start, end in groups
if end - start + 1 >= self.min_messages_to_summarize
]
return groups
def _store_for_retrieval(
self,
messages: list[dict[str, Any]],
summary: str,
original_tokens: int,
summary_tokens: int,
) -> str | None:
"""Store original messages in CompressionStore for CCR retrieval.
Args:
messages: Original messages.
summary: Summary text.
original_tokens: Token count of originals.
summary_tokens: Token count of summary.
Returns:
Cache hash for retrieval, or None if storage failed.
"""
store = self._get_compression_store()
if store is None:
return None
try:
# Serialize messages for storage
original_content = json.dumps(messages, ensure_ascii=False)
# Generate hash
content_hash = hashlib.sha256(original_content.encode()).hexdigest()[:24]
# Store in compression store
store.store(
original=original_content,
compressed=summary,
original_tokens=original_tokens,
compressed_tokens=summary_tokens,
original_item_count=len(messages),
compressed_item_count=1,
tool_name="progressive_summarizer",
)
return content_hash
except Exception as e:
logger.debug("Failed to store for CCR retrieval: %s", e)
return None

395
headroom/transforms/rolling_window.py
View file

@ -1,395 +0,0 @@
"""Rolling window transform for Headroom SDK."""
from __future__ import annotations
import logging
from typing import Any
from ..config import RollingWindowConfig, TransformResult
from ..parser import find_tool_units
from ..tokenizer import Tokenizer
from ..tokenizers import EstimatingTokenCounter
from ..utils import create_dropped_context_marker, deep_copy_messages
from .base import Transform
logger = logging.getLogger(__name__)
class RollingWindow(Transform):
"""
Apply rolling window to keep messages within token budget.
Drop order (deterministic):
1. Oldest TOOL UNITS (assistant+tool_calls paired with tool responses)
2. Oldest assistant+user pairs
3. Oldest RAG blocks (if detectable)
CRITICAL: Tool calls and tool results are atomic DROP UNITS.
Never orphan a tool result.
Never drops:
- System prompt
- Stable instructions
- Last N conversational turns (configurable)
"""
name = "rolling_window"
def __init__(self, config: RollingWindowConfig | None = None):
"""
Initialize rolling window.
Args:
config: Configuration for window behavior.
"""
self.config = config or RollingWindowConfig()
def should_apply(
self,
messages: list[dict[str, Any]],
tokenizer: Tokenizer,
**kwargs: Any,
) -> bool:
"""Check if token cap is exceeded."""
if not self.config.enabled:
return False
model_limit = kwargs.get("model_limit", 128000)
output_buffer = kwargs.get("output_buffer", self.config.output_buffer_tokens)
current_tokens = tokenizer.count_messages(messages)
available = model_limit - output_buffer
return bool(current_tokens > available)
def apply(
self,
messages: list[dict[str, Any]],
tokenizer: Tokenizer,
**kwargs: Any,
) -> TransformResult:
"""
Apply rolling window to messages.
Args:
messages: List of messages.
tokenizer: Tokenizer for counting.
**kwargs: Must include 'model_limit', optionally 'output_buffer'.
Returns:
TransformResult with windowed messages.
"""
model_limit = kwargs.get("model_limit", 128000)
output_buffer = kwargs.get("output_buffer", self.config.output_buffer_tokens)
available = model_limit - output_buffer
tokens_before = tokenizer.count_messages(messages)
result_messages = deep_copy_messages(messages)
transforms_applied: list[str] = []
markers_inserted: list[str] = []
warnings: list[str] = []
dropped_count = 0
tool_units_dropped = 0
# If already under budget, no changes needed
current_tokens = tokens_before
if current_tokens <= available:
return TransformResult(
messages=result_messages,
tokens_before=tokens_before,
tokens_after=tokens_before,
transforms_applied=[],
warnings=[],
)
# Identify protected indices
protected = self._get_protected_indices(result_messages)
# Frozen messages (in provider's prefix cache) must never be dropped
frozen_message_count = kwargs.get("frozen_message_count", 0)
if frozen_message_count > 0:
protected.update(range(frozen_message_count))
# Identify tool units
tool_units = find_tool_units(result_messages)
# Create drop candidates with priorities
drop_candidates = self._build_drop_candidates(result_messages, protected, tool_units)
# Drop until under budget
indices_to_drop: set[int] = set()
for candidate in drop_candidates:
if current_tokens <= available:
break
# Get indices for this candidate
candidate_indices = candidate["indices"]
# Skip if any are protected
if any(idx in protected for idx in candidate_indices):
continue
# Skip if already dropped
if any(idx in indices_to_drop for idx in candidate_indices):
continue
# Calculate tokens saved
tokens_saved = sum(
tokenizer.count_message(result_messages[idx])
for idx in candidate_indices
if idx < len(result_messages)
)
indices_to_drop.update(candidate_indices)
current_tokens -= tokens_saved
dropped_count += 1
if candidate["type"] == "tool_unit":
tool_units_dropped += 1
# Remove dropped messages (in reverse order to preserve indices)
for idx in sorted(indices_to_drop, reverse=True):
if idx < len(result_messages):
del result_messages[idx]
# Insert marker if we dropped anything
if dropped_count > 0:
logger.info(
"RollingWindow: dropped %d units (%d tool units) to fit budget: %d -> %d tokens",
dropped_count,
tool_units_dropped,
tokens_before,
current_tokens,
)
marker = create_dropped_context_marker("token_cap", dropped_count)
markers_inserted.append(marker)
# Insert marker after system messages
insert_idx = 0
for i, msg in enumerate(result_messages):
if msg.get("role") != "system":
insert_idx = i
break
else:
insert_idx = len(result_messages)
# Match content format (Strands uses list-of-blocks, Anthropic uses string)
_uses_block_format = any(
isinstance(m.get("content"), list)
for m in result_messages
if m.get("role") == "user"
)
marker_content: str | list[dict[str, str]] = (
[{"type": "text", "text": marker}] if _uses_block_format else marker
)
result_messages.insert(
insert_idx,
{
"role": "user",
"content": marker_content,
},
)
transforms_applied.append(f"window_cap:{dropped_count}")
tokens_after = tokenizer.count_messages(result_messages)
result = TransformResult(
messages=result_messages,
tokens_before=tokens_before,
tokens_after=tokens_after,
transforms_applied=transforms_applied,
markers_inserted=markers_inserted,
warnings=warnings,
)
return result
def _get_protected_indices(self, messages: list[dict[str, Any]]) -> set[int]:
"""Get indices that should never be dropped."""
protected: set[int] = set()
# Protect system messages
if self.config.keep_system:
for i, msg in enumerate(messages):
if msg.get("role") == "system":
protected.add(i)
# Protect last N turns
if self.config.keep_last_turns > 0:
# Count turns from end (user+assistant = 1 turn)
turns_seen = 0
i = len(messages) - 1
while i >= 0 and turns_seen < self.config.keep_last_turns:
msg = messages[i]
role = msg.get("role")
# Protect this message
protected.add(i)
# Count turns
if role == "user":
turns_seen += 1
i -= 1
# Also protect any tool responses that belong to protected assistant messages
for i in list(protected):
msg = messages[i]
if msg.get("role") == "assistant":
tool_call_ids: set[str] = set()
# OpenAI format: tool_calls array
if msg.get("tool_calls"):
tool_call_ids.update(
tc.get("id") for tc in msg.get("tool_calls", []) if tc.get("id")
)
# Anthropic format: content blocks with type=tool_use
content = msg.get("content")
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_use":
tc_id = block.get("id")
if tc_id:
tool_call_ids.add(tc_id)
# Find and protect corresponding tool responses
if tool_call_ids:
for j, other_msg in enumerate(messages):
# OpenAI format: role="tool"
if other_msg.get("role") == "tool":
if other_msg.get("tool_call_id") in tool_call_ids:
protected.add(j)
# Anthropic format: role="user" with tool_result blocks
if other_msg.get("role") == "user":
other_content = other_msg.get("content")
if isinstance(other_content, list):
for block in other_content:
if (
isinstance(block, dict)
and block.get("type") == "tool_result"
and block.get("tool_use_id") in tool_call_ids
):
protected.add(j)
break
return protected
def _build_drop_candidates(
self,
messages: list[dict[str, Any]],
protected: set[int],
tool_units: list[tuple[int, list[int]]],
) -> list[dict[str, Any]]:
"""
Build ordered list of drop candidates.
Returns candidates in drop priority order (first to drop first).
"""
candidates: list[dict[str, Any]] = []
# Track which indices are part of tool units
tool_unit_indices: set[int] = set()
for assistant_idx, response_indices in tool_units:
tool_unit_indices.add(assistant_idx)
tool_unit_indices.update(response_indices)
# Priority 1: Oldest tool units (all indices as atomic unit)
for assistant_idx, response_indices in tool_units:
if assistant_idx in protected:
continue
all_indices = [assistant_idx] + response_indices
candidates.append(
{
"type": "tool_unit",
"indices": all_indices,
"priority": 1,
"position": assistant_idx, # For sorting by age
}
)
# Priority 2: Oldest non-tool messages (user/assistant pairs)
i = 0
while i < len(messages):
msg = messages[i]
role = msg.get("role")
if i in protected or i in tool_unit_indices:
i += 1
continue
if role in ("user", "assistant"):
# Try to find a pair
if role == "user" and i + 1 < len(messages):
next_msg = messages[i + 1]
if next_msg.get("role") == "assistant" and i + 1 not in tool_unit_indices:
candidates.append(
{
"type": "turn",
"indices": [i, i + 1],
"priority": 2,
"position": i,
}
)
i += 2
continue
# Single message
candidates.append(
{
"type": "single",
"indices": [i],
"priority": 2,
"position": i,
}
)
i += 1
# Sort by priority, then by position (oldest first)
candidates.sort(key=lambda c: (c["priority"], c["position"]))
return candidates
def apply_rolling_window(
messages: list[dict[str, Any]],
model_limit: int,
output_buffer: int = 4000,
keep_last_turns: int = 2,
config: RollingWindowConfig | None = None,
) -> tuple[list[dict[str, Any]], list[str]]:
"""
Convenience function to apply rolling window.
Args:
messages: List of messages.
model_limit: Model's context limit.
output_buffer: Tokens to reserve for output.
keep_last_turns: Number of recent turns to protect.
config: Optional configuration.
Returns:
Tuple of (windowed_messages, dropped_descriptions).
"""
cfg = config or RollingWindowConfig()
cfg.output_buffer_tokens = output_buffer
cfg.keep_last_turns = keep_last_turns
window = RollingWindow(cfg)
tokenizer = Tokenizer(EstimatingTokenCounter()) # type: ignore[arg-type]
result = window.apply(
messages,
tokenizer,
model_limit=model_limit,
output_buffer=output_buffer,
)
return result.messages, result.transforms_applied

459
headroom/transforms/scoring.py
View file

@ -1,459 +0,0 @@
"""Message importance scoring for intelligent context management.
This module provides semantic-aware scoring of messages to determine
which ones are most important to keep when context limits are exceeded.
Design principle: NO HARDCODED PATTERNS
- Error detection: Uses TOIN field_semantics.inferred_type == "error_indicator"
- Importance: Derived from TOIN retrieval_rate and field patterns
- Relevance: Computed via embedding similarity, not keyword matching
"""
from __future__ import annotations
import json
import logging
import math
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Protocol
if TYPE_CHECKING:
from ..config import ScoringWeights
from ..telemetry.toin import ToolIntelligenceNetwork
logger = logging.getLogger(__name__)
class EmbeddingProvider(Protocol):
"""Protocol for embedding providers."""
def embed(self, text: str) -> list[float]:
"""Generate embedding for text."""
...
@dataclass
class MessageScore:
"""Importance score for a single message.
All scores are in range [0.0, 1.0] where higher = more important.
"""
message_index: int
total_score: float
# Component scores
recency_score: float = 0.0
semantic_score: float = 0.0
toin_score: float = 0.0
error_score: float = 0.0
reference_score: float = 0.0
density_score: float = 0.0
# Metadata
tokens: int = 0
is_protected: bool = False
drop_safe: bool = True # Can be dropped without orphaning tool responses
# Debug info
score_breakdown: dict[str, float] = field(default_factory=dict)
class MessageScorer:
"""Scores messages by semantic importance using learned patterns.
This scorer uses TOIN-learned patterns and computed metrics to determine
message importance. It does NOT use hardcoded keyword matching.
Scoring factors:
1. Recency: Exponential decay from conversation end
2. Semantic similarity: Embedding cosine similarity to recent context
3. TOIN importance: Learned field importance from retrieval patterns
4. Error indicators: TOIN-learned error_indicator field types
5. Forward references: Messages referenced by later messages
6. Token density: Information density (unique tokens / total tokens)
"""
def __init__(
self,
weights: ScoringWeights | None = None,
toin: ToolIntelligenceNetwork | None = None,
embedding_provider: EmbeddingProvider | None = None,
recency_decay_rate: float = 0.1,
):
"""Initialize scorer.
Args:
weights: Scoring weights for each factor
toin: ToolIntelligenceNetwork instance for learned patterns
embedding_provider: Optional embedding provider for semantic similarity
recency_decay_rate: Lambda for exponential recency decay
"""
# Import here to avoid circular imports
from ..config import ScoringWeights
self.weights = (weights or ScoringWeights()).normalized()
self.toin = toin
self.embedding_provider = embedding_provider
self.recency_decay_rate = recency_decay_rate
# Cache for embeddings
self._embedding_cache: dict[int, list[float]] = {}
def score_messages(
self,
messages: list[dict[str, Any]],
protected_indices: set[int],
tool_unit_indices: set[int],
) -> list[MessageScore]:
"""Score all messages by importance.
Args:
messages: List of messages to score
protected_indices: Indices that should never be dropped
tool_unit_indices: Indices that are part of tool units
Returns:
List of MessageScore objects, one per message
"""
scores: list[MessageScore] = []
# Pre-compute forward references
forward_refs = self._compute_forward_references(messages)
# Pre-compute recent context embedding (last 3 messages)
recent_embedding = self._compute_recent_context_embedding(messages)
for i, msg in enumerate(messages):
score = self._score_message(
msg=msg,
index=i,
total_messages=len(messages),
protected=i in protected_indices,
in_tool_unit=i in tool_unit_indices,
forward_refs=forward_refs,
recent_embedding=recent_embedding,
)
scores.append(score)
return scores
def _score_message(
self,
msg: dict[str, Any],
index: int,
total_messages: int,
protected: bool,
in_tool_unit: bool,
forward_refs: dict[int, int],
recent_embedding: list[float] | None,
) -> MessageScore:
"""Score a single message."""
# Compute individual scores
recency = self._compute_recency_score(index, total_messages)
semantic = self._compute_semantic_score(msg, index, recent_embedding)
toin_importance = self._compute_toin_score(msg)
error = self._compute_error_score(msg)
reference = self._compute_reference_score(index, forward_refs)
density = self._compute_density_score(msg)
# Weighted combination
w = self.weights
total = (
w.recency * recency
+ w.semantic_similarity * semantic
+ w.toin_importance * toin_importance
+ w.error_indicator * error
+ w.forward_reference * reference
+ w.token_density * density
)
# Estimate tokens (simple heuristic)
content = msg.get("content", "")
if isinstance(content, str):
tokens = len(content) // 4 # Rough estimate
else:
tokens = 100 # Default for complex content
return MessageScore(
message_index=index,
total_score=total,
recency_score=recency,
semantic_score=semantic,
toin_score=toin_importance,
error_score=error,
reference_score=reference,
density_score=density,
tokens=tokens,
is_protected=protected,
drop_safe=not in_tool_unit or not protected,
score_breakdown={
"recency": recency,
"semantic": semantic,
"toin": toin_importance,
"error": error,
"reference": reference,
"density": density,
},
)
def _compute_recency_score(self, index: int, total: int) -> float:
"""Compute recency score using exponential decay."""
if total <= 1:
return 1.0
# Position from end (0 = last message)
position_from_end = total - 1 - index
# Exponential decay: score = e^(-λ * position)
score = math.exp(-self.recency_decay_rate * position_from_end)
return score
def _compute_semantic_score(
self,
msg: dict[str, Any],
index: int,
recent_embedding: list[float] | None,
) -> float:
"""Compute semantic similarity to recent context."""
if self.embedding_provider is None or recent_embedding is None:
return 0.5 # Neutral score if no embedding available
content = msg.get("content", "")
if not isinstance(content, str) or not content.strip():
return 0.5
try:
# Get or compute embedding for this message
if index not in self._embedding_cache:
self._embedding_cache[index] = self.embedding_provider.embed(content)
msg_embedding = self._embedding_cache[index]
# Cosine similarity
return self._cosine_similarity(msg_embedding, recent_embedding)
except Exception as e:
logger.debug(f"Embedding computation failed: {e}")
return 0.5
def _compute_toin_score(self, msg: dict[str, Any]) -> float:
"""Compute importance score from TOIN patterns.
This uses TOIN-learned retrieval patterns to determine importance.
Higher retrieval rate = more important (users needed this data).
"""
if self.toin is None:
return 0.5 # Neutral if no TOIN
# Only tool messages have TOIN patterns
if msg.get("role") != "tool":
return 0.5
content = msg.get("content", "")
if not content:
return 0.5
try:
# Try to parse as JSON and get tool signature
from ..telemetry.models import ToolSignature
data = json.loads(content) if isinstance(content, str) else content
if not isinstance(data, (list, dict)):
return 0.5
# Get tool signature
items = data if isinstance(data, list) else [data]
if not items:
return 0.5
signature = ToolSignature.from_items(items)
pattern = self.toin.get_pattern(signature.structure_hash)
if pattern is None or pattern.confidence < 0.3:
return 0.5
# Score based on retrieval rate (high retrieval = important)
# Note: We use retrieval_rate as the primary importance signal from TOIN.
# The commonly_retrieved_fields comparison is not used here because
# ToolSignature only tracks structure_hash, not individual field hashes.
score = 0.5 + (pattern.retrieval_rate * 0.5)
# Boost slightly if pattern has commonly retrieved fields (indicates
# this tool type has learned importance patterns)
if pattern.commonly_retrieved_fields:
boost = min(0.1, 0.02 * len(pattern.commonly_retrieved_fields))
score = min(1.0, score + boost)
return score
except Exception as e:
logger.debug(f"TOIN scoring failed: {e}")
return 0.5
def _compute_error_score(self, msg: dict[str, Any]) -> float:
"""Compute error indicator score using TOIN-learned patterns.
This does NOT use hardcoded keyword matching. Instead, it uses
TOIN's learned field_semantics to identify error_indicator fields.
"""
if self.toin is None:
return 0.0
# Only tool messages have field semantics
if msg.get("role") != "tool":
return 0.0
content = msg.get("content", "")
if not content:
return 0.0
try:
from ..telemetry.models import ToolSignature
data = json.loads(content) if isinstance(content, str) else content
if not isinstance(data, (list, dict)):
return 0.0
items = data if isinstance(data, list) else [data]
if not items:
return 0.0
signature = ToolSignature.from_items(items)
pattern = self.toin.get_pattern(signature.structure_hash)
if pattern is None:
return 0.0
# Check field_semantics for error_indicator type
error_field_count = 0
high_confidence_errors = 0
for _field_hash, field_sem in pattern.field_semantics.items():
if field_sem.inferred_type == "error_indicator":
error_field_count += 1
if field_sem.confidence >= 0.7:
high_confidence_errors += 1
if error_field_count == 0:
return 0.0
# Score based on presence and confidence of error fields
base_score = min(1.0, 0.3 * error_field_count)
confidence_boost = min(0.5, 0.2 * high_confidence_errors)
return base_score + confidence_boost
except Exception as e:
logger.debug(f"Error scoring failed: {e}")
return 0.0
def _compute_reference_score(
self,
index: int,
forward_refs: dict[int, int],
) -> float:
"""Compute score based on forward references.
Messages that are referenced by later messages are more important.
"""
ref_count = forward_refs.get(index, 0)
if ref_count == 0:
return 0.0
# Logarithmic scaling for reference count
return min(1.0, 0.3 + 0.2 * math.log(ref_count + 1))
def _compute_density_score(self, msg: dict[str, Any]) -> float:
"""Compute information density score.
Higher unique token ratio = more information dense = more important.
"""
content = msg.get("content", "")
if not isinstance(content, str) or len(content) < 10:
return 0.5
# Simple token density: unique tokens / total tokens
tokens = content.lower().split()
if len(tokens) < 3:
return 0.5
unique_tokens = len(set(tokens))
density = unique_tokens / len(tokens)
# Normalize to [0, 1] range (typical density is 0.3-0.8)
return min(1.0, max(0.0, (density - 0.2) / 0.6))
def _compute_forward_references(
self,
messages: list[dict[str, Any]],
) -> dict[int, int]:
"""Compute which messages are referenced by later messages.
Returns dict mapping message index to reference count.
"""
refs: dict[int, int] = {}
# Track tool_call_id references
tool_call_ids: dict[str, int] = {} # tool_call_id -> assistant message index
for i, msg in enumerate(messages):
role = msg.get("role")
# Track assistant tool calls
if role == "assistant" and msg.get("tool_calls"):
for tc in msg.get("tool_calls", []):
tc_id = tc.get("id")
if tc_id:
tool_call_ids[tc_id] = i
# Tool responses reference assistant messages
elif role == "tool":
tc_id = msg.get("tool_call_id")
if tc_id and tc_id in tool_call_ids:
ref_idx = tool_call_ids[tc_id]
refs[ref_idx] = refs.get(ref_idx, 0) + 1
return refs
def _compute_recent_context_embedding(
self,
messages: list[dict[str, Any]],
num_recent: int = 3,
) -> list[float] | None:
"""Compute average embedding of recent messages."""
if self.embedding_provider is None:
return None
recent_texts: list[str] = []
for msg in messages[-num_recent:]:
content = msg.get("content", "")
if isinstance(content, str) and content.strip():
recent_texts.append(content)
if not recent_texts:
return None
try:
# Combine recent texts and embed
combined = " ".join(recent_texts)
return self.embedding_provider.embed(combined)
except Exception as e:
logger.debug(f"Recent context embedding failed: {e}")
return None
@staticmethod
def _cosine_similarity(a: list[float], b: list[float]) -> float:
"""Compute cosine similarity between two vectors."""
if len(a) != len(b) or len(a) == 0:
return 0.0
dot_product = sum(x * y for x, y in zip(a, b))
norm_a = math.sqrt(sum(x * x for x in a))
norm_b = math.sqrt(sum(y * y for y in b))
if norm_a == 0 or norm_b == 0:
return 0.0
return dot_product / (norm_a * norm_b)

338
headroom/transforms/tool_crusher.py
View file

@ -1,338 +0,0 @@
"""Tool output compression transform for Headroom SDK."""
from __future__ import annotations
import logging
from typing import Any
from ..config import ToolCrusherConfig, TransformResult
from ..tokenizer import Tokenizer
from ..utils import (
compute_short_hash,
create_tool_digest_marker,
deep_copy_messages,
safe_json_dumps,
safe_json_loads,
)
from .base import Transform
logger = logging.getLogger(__name__)
class ToolCrusher(Transform):
"""
Compress tool output to reduce token usage.
This transform applies conservative compression:
- Only compresses tool role messages > min_tokens
- Preserves JSON structure (never removes keys)
- Truncates arrays to max_items
- Truncates long strings
- Limits nesting depth
Safety: If JSON parsing fails, content is returned unchanged.
"""
name = "tool_crusher"
def __init__(self, config: ToolCrusherConfig | None = None):
"""
Initialize tool crusher.
Args:
config: Configuration for compression behavior.
"""
self.config = config or ToolCrusherConfig()
def should_apply(
self,
messages: list[dict[str, Any]],
tokenizer: Tokenizer,
**kwargs: Any,
) -> bool:
"""Check if any tool messages exceed threshold."""
if not self.config.enabled:
return False
for msg in messages:
# OpenAI style: role="tool"
if msg.get("role") == "tool":
content = msg.get("content", "")
if isinstance(content, str):
tokens = tokenizer.count_text(content)
if tokens > self.config.min_tokens_to_crush:
return True
# Anthropic style: role="user" with tool_result content blocks
content = msg.get("content")
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_result":
tool_content = block.get("content", "")
if isinstance(tool_content, str):
tokens = tokenizer.count_text(tool_content)
if tokens > self.config.min_tokens_to_crush:
return True
return False
def apply(
self,
messages: list[dict[str, Any]],
tokenizer: Tokenizer,
**kwargs: Any,
) -> TransformResult:
"""
Apply tool crushing to messages.
Args:
messages: List of messages.
tokenizer: Tokenizer for counting.
**kwargs: May include 'tool_profiles' for per-tool config.
Returns:
TransformResult with crushed messages.
"""
tool_profiles = kwargs.get("tool_profiles", self.config.tool_profiles)
tokens_before = tokenizer.count_messages(messages)
result_messages = deep_copy_messages(messages)
transforms_applied: list[str] = []
markers_inserted: list[str] = []
warnings: list[str] = []
crushed_count = 0
for msg in result_messages:
# OpenAI style: role="tool"
if msg.get("role") == "tool":
content = msg.get("content", "")
if not isinstance(content, str):
continue
# Check token threshold
tokens = tokenizer.count_text(content)
if tokens <= self.config.min_tokens_to_crush:
continue
# Get tool-specific profile if available
tool_call_id = msg.get("tool_call_id", "")
profile = self._get_profile(tool_call_id, tool_profiles)
# Try to crush
crushed, was_modified = self._crush_content(content, profile)
if was_modified:
# Compute hash of original for marker
original_hash = compute_short_hash(content)
marker = create_tool_digest_marker(original_hash)
msg["content"] = crushed + "\n" + marker
crushed_count += 1
markers_inserted.append(marker)
# Anthropic style: role="user" with tool_result content blocks
content = msg.get("content")
if isinstance(content, list):
for i, block in enumerate(content):
if not isinstance(block, dict):
continue
if block.get("type") != "tool_result":
continue
tool_content = block.get("content", "")
if not isinstance(tool_content, str):
continue
# Check token threshold
tokens = tokenizer.count_text(tool_content)
if tokens <= self.config.min_tokens_to_crush:
continue
# Get tool-specific profile if available
tool_use_id = block.get("tool_use_id", "")
profile = self._get_profile(tool_use_id, tool_profiles)
# Try to crush
crushed, was_modified = self._crush_content(tool_content, profile)
if was_modified:
# Compute hash of original for marker
original_hash = compute_short_hash(tool_content)
marker = create_tool_digest_marker(original_hash)
# Update the content block
content[i]["content"] = crushed + "\n" + marker
crushed_count += 1
markers_inserted.append(marker)
if crushed_count > 0:
transforms_applied.append(f"tool_crush:{crushed_count}")
logger.info(
"ToolCrusher: compressed %d tool outputs, %d -> %d tokens",
crushed_count,
tokens_before,
tokenizer.count_messages(result_messages),
)
tokens_after = tokenizer.count_messages(result_messages)
return TransformResult(
messages=result_messages,
tokens_before=tokens_before,
tokens_after=tokens_after,
transforms_applied=transforms_applied,
markers_inserted=markers_inserted,
warnings=warnings,
)
def _get_profile(
self,
tool_call_id: str,
tool_profiles: dict[str, dict[str, Any]],
) -> dict[str, Any]:
"""Get compression profile for a tool."""
# Tool profiles are keyed by tool name, not call ID
# For now, use default config
# In a real implementation, you'd map call_id -> tool_name
return {
"max_array_items": self.config.max_array_items,
"max_string_length": self.config.max_string_length,
"max_depth": self.config.max_depth,
"preserve_keys": self.config.preserve_keys,
}
def _crush_content(
self,
content: str,
profile: dict[str, Any],
) -> tuple[str, bool]:
"""
Crush content according to profile.
Returns:
Tuple of (crushed_content, was_modified).
If parsing fails, returns (original_content, False).
"""
# Try JSON parse
parsed, success = safe_json_loads(content)
if not success:
# Safety: don't modify unparseable content
return content, False
# Apply crushing
crushed = self._crush_value(
parsed,
depth=0,
max_depth=profile.get("max_depth", 5),
max_array_items=profile.get("max_array_items", 10),
max_string_length=profile.get("max_string_length", 1000),
)
# Serialize back
result = safe_json_dumps(crushed, indent=None)
# Check if actually modified
was_modified = result != content.strip()
return result, was_modified
def _crush_value(
self,
value: Any,
depth: int,
max_depth: int,
max_array_items: int,
max_string_length: int,
) -> Any:
"""Recursively crush a value."""
if depth >= max_depth:
# At max depth, summarize
if isinstance(value, dict):
return {"__headroom_depth_exceeded": len(value)}
elif isinstance(value, list):
return {"__headroom_depth_exceeded": len(value)}
elif isinstance(value, str) and len(value) > max_string_length:
return (
value[:max_string_length]
+ f"...[truncated {len(value) - max_string_length} chars]"
)
return value
if isinstance(value, dict):
return {
k: self._crush_value(
v,
depth + 1,
max_depth,
max_array_items,
max_string_length,
)
for k, v in value.items()
}
elif isinstance(value, list):
if len(value) <= max_array_items:
return [
self._crush_value(
item,
depth + 1,
max_depth,
max_array_items,
max_string_length,
)
for item in value
]
else:
# Truncate array
truncated = [
self._crush_value(
item,
depth + 1,
max_depth,
max_array_items,
max_string_length,
)
for item in value[:max_array_items]
]
truncated.append({"__headroom_truncated": len(value) - max_array_items})
return truncated
elif isinstance(value, str):
if len(value) > max_string_length:
return (
value[:max_string_length]
+ f"...[truncated {len(value) - max_string_length} chars]"
)
return value
else:
# Numbers, bools, None - pass through
return value
def crush_tool_output(
content: str,
config: ToolCrusherConfig | None = None,
) -> tuple[str, bool]:
"""
Convenience function to crush a single tool output.
Args:
content: The tool output content.
config: Optional configuration.
Returns:
Tuple of (crushed_content, was_modified).
"""
cfg = config or ToolCrusherConfig()
crusher = ToolCrusher(cfg)
profile = {
"max_array_items": cfg.max_array_items,
"max_string_length": cfg.max_string_length,
"max_depth": cfg.max_depth,
"preserve_keys": cfg.preserve_keys,
}
return crusher._crush_content(content, profile)

18
tests/conftest.py
View file

@ -45,6 +45,24 @@ def pytest_runtest_call(item):
pytest.skip("Skipped due to network timeout (flaky CI)")
@pytest.fixture(autouse=True)
def _reset_headroom_logger_propagation():
"""Keep `headroom.*` log records flowing to pytest's caplog handler.
`headroom.proxy.helpers._setup_file_logging` sets
``logging.getLogger("headroom").propagate = False`` once any test
triggers a proxy startup with `--log-file`. After that, every
subsequent test's `caplog` fixture stops capturing `headroom.*`
log records (caplog attaches to root, propagation is now blocked
at the headroom-logger boundary). Reset before every test so the
capture is deterministic regardless of run order.
"""
import logging as _logging
_logging.getLogger("headroom").propagate = True
yield
# =============================================================================
# Sample messages fixtures
# =============================================================================

125
tests/parity/fixtures/message_scorer/custom_weights_recency_heavy_5cc105616d48.json
View file

@ -1,125 +0,0 @@
{
"config": {
"weights": {
"error_indicator": 0.05,
"forward_reference": 0.1,
"recency": 0.6,
"semantic_similarity": 0.1,
"toin_importance": 0.1,
"token_density": 0.05
}
},
"input": {
"decay_rate": 0.1,
"messages": [
{
"content": "first message in a longer chat",
"role": "user"
},
{
"content": "an assistant reply with substance",
"role": "assistant"
},
{
"content": "another follow up question here",
"role": "user"
},
{
"content": "and the closing reply",
"role": "assistant"
}
],
"protected_indices": [],
"tool_unit_indices": []
},
"input_sha256": "5cc105616d481b3c37f1cbf359b16967669b1703acc83c7df7e70488502271a6",
"label": "custom_weights_recency_heavy",
"output": [
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 0,
"recency_score": 0.74082,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.74082,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 7,
"total_score": 0.59449
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 1,
"recency_score": 0.81873,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.81873,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 8,
"total_score": 0.64124
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 2,
"recency_score": 0.90484,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.90484,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 7,
"total_score": 0.6929
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 3,
"recency_score": 1.0,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 1.0,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 5,
"total_score": 0.75
}
],
"recorded_at": "2026-05-01T04:29:36.075652+00:00",
"transform": "message_scorer"
}

133
tests/parity/fixtures/message_scorer/drop_safe_truth_table_2e03ebdc1b01.json
View file

@ -1,133 +0,0 @@
{
"config": {
"weights": null
},
"input": {
"decay_rate": 0.1,
"messages": [
{
"content": "you are helpful",
"role": "system"
},
{
"content": "hello",
"role": "user"
},
{
"content": "",
"role": "assistant",
"tool_calls": [
{
"function": {
"name": "f"
},
"id": "x"
}
]
},
{
"content": "result",
"role": "tool",
"tool_call_id": "x"
}
],
"protected_indices": [
0,
2
],
"tool_unit_indices": [
2,
3
]
},
"input_sha256": "2e03ebdc1b0185fdd33c537f199890e5faabda24a4ee3025959419c21eda5569",
"label": "drop_safe_truth_table",
"output": [
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": true,
"message_index": 0,
"recency_score": 0.74082,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.74082,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 3,
"total_score": 0.42316
},
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 1,
"recency_score": 0.81873,
"reference_score": 0.0,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 0.81873,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 1,
"total_score": 0.41375
},
{
"density_score": 0.5,
"drop_safe": false,
"error_score": 0.0,
"is_protected": true,
"message_index": 2,
"recency_score": 0.90484,
"reference_score": 0.43863,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 0.90484,
"reference": 0.43863,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 0,
"total_score": 0.49676
},
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 3,
"recency_score": 1.0,
"reference_score": 0.0,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 1.0,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 1,
"total_score": 0.45
}
],
"recorded_at": "2026-05-01T04:29:36.076447+00:00",
"transform": "message_scorer"
}

16
tests/parity/fixtures/message_scorer/empty_c94d65311b40.json
View file

@ -1,16 +0,0 @@
{
"config": {
"weights": null
},
"input": {
"decay_rate": 0.1,
"messages": [],
"protected_indices": [],
"tool_unit_indices": []
},
"input_sha256": "c94d65311b40eb4065d7160bded90e79fd48aebca28fc334c257ea658c66413b",
"label": "empty",
"output": [],
"recorded_at": "2026-05-01T04:29:36.073767+00:00",
"transform": "message_scorer"
}

43
tests/parity/fixtures/message_scorer/high_density_dd454e2933e3.json
View file

@ -1,43 +0,0 @@
{
"config": {
"weights": null
},
"input": {
"decay_rate": 0.1,
"messages": [
{
"content": "alpha bravo charlie delta echo foxtrot golf hotel",
"role": "user"
}
],
"protected_indices": [],
"tool_unit_indices": []
},
"input_sha256": "dd454e2933e3c2d10ab2fdfbb4675f797711cbd2e67d44ed767a9c590121c88c",
"label": "high_density",
"output": [
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 0,
"recency_score": 1.0,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 1.0,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 12,
"total_score": 0.475
}
],
"recorded_at": "2026-05-01T04:29:36.075245+00:00",
"transform": "message_scorer"
}

145
tests/parity/fixtures/message_scorer/linear_5_messages_71f1ec6a317c.json
View file

@ -1,145 +0,0 @@
{
"config": {
"weights": null
},
"input": {
"decay_rate": 0.1,
"messages": [
{
"content": "what is the capital of france",
"role": "user"
},
{
"content": "the capital of france is paris",
"role": "assistant"
},
{
"content": "and germany",
"role": "user"
},
{
"content": "the capital of germany is berlin",
"role": "assistant"
},
{
"content": "thanks",
"role": "user"
}
],
"protected_indices": [
0
],
"tool_unit_indices": []
},
"input_sha256": "71f1ec6a317cdf54004ec850d9d7312c7eed622337d3961fab83cb56c1f8b702",
"label": "linear_5_messages",
"output": [
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": true,
"message_index": 0,
"recency_score": 0.67032,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.67032,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 7,
"total_score": 0.40906
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 1,
"recency_score": 0.74082,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.74082,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 7,
"total_score": 0.42316
},
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 2,
"recency_score": 0.81873,
"reference_score": 0.0,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 0.81873,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 2,
"total_score": 0.41375
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 3,
"recency_score": 0.90484,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.90484,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 8,
"total_score": 0.45597
},
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 4,
"recency_score": 1.0,
"reference_score": 0.0,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 1.0,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 1,
"total_score": 0.45
}
],
"recorded_at": "2026-05-01T04:29:36.074114+00:00",
"transform": "message_scorer"
}

43
tests/parity/fixtures/message_scorer/low_density_28ad2371fdc0.json
View file

@ -1,43 +0,0 @@
{
"config": {
"weights": null
},
"input": {
"decay_rate": 0.1,
"messages": [
{
"content": "ok ok ok ok ok ok ok ok ok ok",
"role": "user"
}
],
"protected_indices": [],
"tool_unit_indices": []
},
"input_sha256": "28ad2371fdc08517452da84c371a2063fff941f2613c87cafef8c70efa8bf4ba",
"label": "low_density",
"output": [
{
"density_score": 0.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 0,
"recency_score": 1.0,
"reference_score": 0.0,
"score_breakdown": {
"density": 0.0,
"error": 0.0,
"recency": 1.0,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 7,
"total_score": 0.425
}
],
"recorded_at": "2026-05-01T04:29:36.075414+00:00",
"transform": "message_scorer"
}

171
tests/parity/fixtures/message_scorer/multi_ref_to_assistant_67a047a411c4.json
View file

@ -1,171 +0,0 @@
{
"config": {
"weights": null
},
"input": {
"decay_rate": 0.1,
"messages": [
{
"content": "do many things",
"role": "user"
},
{
"content": "",
"role": "assistant",
"tool_calls": [
{
"function": {
"name": "f"
},
"id": "c1"
},
{
"function": {
"name": "g"
},
"id": "c2"
},
{
"function": {
"name": "h"
},
"id": "c3"
}
]
},
{
"content": "r1",
"role": "tool",
"tool_call_id": "c1"
},
{
"content": "r2",
"role": "tool",
"tool_call_id": "c2"
},
{
"content": "r3",
"role": "tool",
"tool_call_id": "c3"
}
],
"protected_indices": [],
"tool_unit_indices": [
1,
2,
3,
4
]
},
"input_sha256": "67a047a411c474d0d3fa712b89aabc12311910bd466f3c750c003245558306de",
"label": "multi_ref_to_assistant",
"output": [
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 0,
"recency_score": 0.67032,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.67032,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 3,
"total_score": 0.40906
},
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 1,
"recency_score": 0.74082,
"reference_score": 0.57726,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 0.74082,
"reference": 0.57726,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 0,
"total_score": 0.48475
},
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 2,
"recency_score": 0.81873,
"reference_score": 0.0,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 0.81873,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 0,
"total_score": 0.41375
},
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 3,
"recency_score": 0.90484,
"reference_score": 0.0,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 0.90484,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 0,
"total_score": 0.43097
},
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 4,
"recency_score": 1.0,
"reference_score": 0.0,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 1.0,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 0,
"total_score": 0.45
}
],
"recorded_at": "2026-05-01T04:29:36.074987+00:00",
"transform": "message_scorer"
}

76
tests/parity/fixtures/message_scorer/non_string_content_9ab4060e7b90.json
View file

@ -1,76 +0,0 @@
{
"config": {
"weights": null
},
"input": {
"decay_rate": 0.1,
"messages": [
{
"content": "hi",
"role": "user"
},
{
"content": null,
"role": "assistant",
"tool_calls": [
{
"function": {
"name": "f"
},
"id": "tc"
}
]
}
],
"protected_indices": [],
"tool_unit_indices": []
},
"input_sha256": "9ab4060e7b906c73fae6926ae4766fc651b1ab73fdf023496f1b9bd6d84c2cae",
"label": "non_string_content",
"output": [
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 0,
"recency_score": 0.90484,
"reference_score": 0.0,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 0.90484,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 0,
"total_score": 0.43097
},
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 1,
"recency_score": 1.0,
"reference_score": 0.0,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 1.0,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 100,
"total_score": 0.45
}
],
"recorded_at": "2026-05-01T04:29:36.077209+00:00",
"transform": "message_scorer"
}

69
tests/parity/fixtures/message_scorer/orphan_tool_response_adb4e13b7541.json
View file

@ -1,69 +0,0 @@
{
"config": {
"weights": null
},
"input": {
"decay_rate": 0.1,
"messages": [
{
"content": "go",
"role": "user"
},
{
"content": "stranded",
"role": "tool",
"tool_call_id": "nope"
}
],
"protected_indices": [],
"tool_unit_indices": []
},
"input_sha256": "adb4e13b75416cfbf8623002eaf6446ee67060bc64f88182b06f3b431c140c2f",
"label": "orphan_tool_response",
"output": [
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 0,
"recency_score": 0.90484,
"reference_score": 0.0,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 0.90484,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 0,
"total_score": 0.43097
},
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 1,
"recency_score": 1.0,
"reference_score": 0.0,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 1.0,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 2,
"total_score": 0.45
}
],
"recorded_at": "2026-05-01T04:29:36.077001+00:00",
"transform": "message_scorer"
}

43
tests/parity/fixtures/message_scorer/single_user_message_2b4fac2277be.json
View file

@ -1,43 +0,0 @@
{
"config": {
"weights": null
},
"input": {
"decay_rate": 0.1,
"messages": [
{
"content": "hello world",
"role": "user"
}
],
"protected_indices": [],
"tool_unit_indices": []
},
"input_sha256": "2b4fac2277be4cee6dcb022904f561d8ea4753a82714afe829f9203e1c3e441e",
"label": "single_user_message",
"output": [
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 0,
"recency_score": 1.0,
"reference_score": 0.0,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 1.0,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 2,
"total_score": 0.45
}
],
"recorded_at": "2026-05-01T04:29:36.073288+00:00",
"transform": "message_scorer"
}

218
tests/parity/fixtures/message_scorer/slow_decay_de4fafb3a42b.json
View file

@ -1,218 +0,0 @@
{
"config": {
"weights": null
},
"input": {
"decay_rate": 0.02,
"messages": [
{
"content": "message number 0",
"role": "user"
},
{
"content": "message number 1",
"role": "user"
},
{
"content": "message number 2",
"role": "user"
},
{
"content": "message number 3",
"role": "user"
},
{
"content": "message number 4",
"role": "user"
},
{
"content": "message number 5",
"role": "user"
},
{
"content": "message number 6",
"role": "user"
},
{
"content": "message number 7",
"role": "user"
}
],
"protected_indices": [],
"tool_unit_indices": []
},
"input_sha256": "de4fafb3a42bca1ff0ec75d0530b41b40d5d5bfb424975beae5ad9099264112c",
"label": "slow_decay",
"output": [
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 0,
"recency_score": 0.86936,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.86936,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 4,
"total_score": 0.44887
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 1,
"recency_score": 0.88692,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.88692,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 4,
"total_score": 0.45238
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 2,
"recency_score": 0.90484,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.90484,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 4,
"total_score": 0.45597
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 3,
"recency_score": 0.92312,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.92312,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 4,
"total_score": 0.45962
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 4,
"recency_score": 0.94176,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.94176,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 4,
"total_score": 0.46335
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 5,
"recency_score": 0.96079,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.96079,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 4,
"total_score": 0.46716
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 6,
"recency_score": 0.9802,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.9802,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 4,
"total_score": 0.47104
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 7,
"recency_score": 1.0,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 1.0,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 4,
"total_score": 0.475
}
],
"recorded_at": "2026-05-01T04:29:36.076024+00:00",
"transform": "message_scorer"
}

134
tests/parity/fixtures/message_scorer/tool_call_pair_3443697600a6.json
View file

@ -1,134 +0,0 @@
{
"config": {
"weights": null
},
"input": {
"decay_rate": 0.1,
"messages": [
{
"content": "what's the weather",
"role": "user"
},
{
"content": "",
"role": "assistant",
"tool_calls": [
{
"function": {
"arguments": "{}",
"name": "get_weather"
},
"id": "call_1",
"type": "function"
}
]
},
{
"content": "{\"temp\": 72, \"conditions\": \"sunny\"}",
"role": "tool",
"tool_call_id": "call_1"
},
{
"content": "it is 72 and sunny",
"role": "assistant"
}
],
"protected_indices": [
0
],
"tool_unit_indices": [
1,
2
]
},
"input_sha256": "3443697600a6878b18e41dd4f499829ee9f0ea882cf0510b030e15ce1febd07d",
"label": "tool_call_pair",
"output": [
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": true,
"message_index": 0,
"recency_score": 0.74082,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.74082,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 4,
"total_score": 0.42316
},
{
"density_score": 0.5,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 1,
"recency_score": 0.81873,
"reference_score": 0.43863,
"score_breakdown": {
"density": 0.5,
"error": 0.0,
"recency": 0.81873,
"reference": 0.43863,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 0,
"total_score": 0.47954
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 2,
"recency_score": 0.90484,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.90484,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 8,
"total_score": 0.45597
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 3,
"recency_score": 1.0,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 1.0,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 4,
"total_score": 0.475
}
],
"recorded_at": "2026-05-01T04:29:36.074466+00:00",
"transform": "message_scorer"
}

68
tests/parity/fixtures/message_scorer/unicode_content_22231f9bffa6.json
View file

@ -1,68 +0,0 @@
{
"config": {
"weights": null
},
"input": {
"decay_rate": 0.1,
"messages": [
{
"content": "\u4f60\u597d\u4e16\u754c \u044d\u0442\u043e \u0442\u0435\u0441\u0442 \ud83c\udf89\ud83c\udf89\ud83c\udf89",
"role": "user"
},
{
"content": "received the unicode message",
"role": "assistant"
}
],
"protected_indices": [],
"tool_unit_indices": []
},
"input_sha256": "22231f9bffa6dadeb455f2d1515019d7924ed307fe27331b2fb8720297b0e2b4",
"label": "unicode_content",
"output": [
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 0,
"recency_score": 0.90484,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 0.90484,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 4,
"total_score": 0.45597
},
{
"density_score": 1.0,
"drop_safe": true,
"error_score": 0.0,
"is_protected": false,
"message_index": 1,
"recency_score": 1.0,
"reference_score": 0.0,
"score_breakdown": {
"density": 1.0,
"error": 0.0,
"recency": 1.0,
"reference": 0.0,
"semantic": 0.5,
"toin": 0.5
},
"semantic_score": 0.5,
"toin_score": 0.5,
"tokens": 7,
"total_score": 0.475
}
],
"recorded_at": "2026-05-01T04:29:36.076815+00:00",
"transform": "message_scorer"
}

350
tests/parity/record_message_scorer.py
View file

@ -1,350 +0,0 @@
"""Record `MessageScorer` parity fixtures.
Captures `MessageScorer.score_messages(messages, protected, tool_unit)`
with `toin=None` and `embedding_provider=None` so all six factors run
through the deterministic-or-neutral code path. The Rust comparator
(`MessageScorerComparator` in `crates/headroom-parity/src/lib.rs`) runs
the same inputs through the Rust port and asserts bit-equal outputs
after a 5-decimal-place rounding step.
Why round: Rust uses `f32::exp` and computes the weighted total in
f32, while Python's `math.exp` and weighted sum are f64. Both are
mathematically identical; they only drift in the low bits. Rounding
both sides to 5 decimals (1e-5 tolerance, 100x looser than f32 ulp
drift on the 0–1 score range) gives byte-equal JSON without masking
real bugs.
Run from repo root:
python tests/parity/record_message_scorer.py
"""
from __future__ import annotations
import datetime as _dt
import hashlib
import json
from dataclasses import asdict
from pathlib import Path
from typing import Any
from headroom.config import ScoringWeights
from headroom.transforms.scoring import MessageScorer
_REPO_ROOT = Path(__file__).resolve().parent.parent.parent
_FIXTURES_DIR = _REPO_ROOT / "tests" / "parity" / "fixtures" / "message_scorer"
# 5 decimals: f32 has ~7 decimals of precision, but Rust's port runs
# the weighted-sum in f32 while Python runs it in f64 — six summed
# f32-precision components occasionally drift in the 6th decimal of
# the total. 5 decimals (1e-5 tolerance) is loose enough to absorb
# the drift while still tight enough to catch real bugs.
_FLOAT_ROUND_PLACES = 5
def _round_floats(obj: Any) -> Any:
"""Recursively round every float in a JSON-shaped object."""
if isinstance(obj, float):
return round(obj, _FLOAT_ROUND_PLACES)
if isinstance(obj, dict):
return {k: _round_floats(v) for k, v in obj.items()}
if isinstance(obj, list):
return [_round_floats(v) for v in obj]
return obj
def _digest(payload: dict[str, Any]) -> str:
blob = json.dumps(payload, sort_keys=True).encode("utf-8")
return hashlib.sha256(blob).hexdigest()
def _record(
label: str,
messages: list[dict[str, Any]],
protected_indices: list[int],
tool_unit_indices: list[int],
weights: ScoringWeights | None = None,
decay_rate: float = 0.1,
) -> Path:
scorer = MessageScorer(
weights=weights,
toin=None,
embedding_provider=None,
recency_decay_rate=decay_rate,
)
scores = scorer.score_messages(
messages=messages,
protected_indices=set(protected_indices),
tool_unit_indices=set(tool_unit_indices),
)
payload_input = {
"messages": messages,
"protected_indices": sorted(protected_indices),
"tool_unit_indices": sorted(tool_unit_indices),
"decay_rate": decay_rate,
}
payload_config = {"weights": asdict(weights) if weights else None}
payload_output = _round_floats([asdict(s) for s in scores])
digest_source = {
"transform": "message_scorer",
"label": label,
"input": payload_input,
"config": payload_config,
}
digest = _digest(digest_source)
fixture = {
"transform": "message_scorer",
"label": label,
"input": payload_input,
"config": payload_config,
"output": payload_output,
"recorded_at": _dt.datetime.now(tz=_dt.timezone.utc).isoformat(),
"input_sha256": digest,
}
_FIXTURES_DIR.mkdir(parents=True, exist_ok=True)
target = _FIXTURES_DIR / f"{label}_{digest[:12]}.json"
target.write_text(json.dumps(fixture, indent=2, sort_keys=True) + "\n")
return target
def _scenarios() -> list[dict[str, Any]]:
"""Test scenarios covering each deterministic factor + edge cases."""
out: list[dict[str, Any]] = []
# 1. Single message — recency=1.0, no refs, no tool unit.
out.append(
{
"label": "single_user_message",
"messages": [{"role": "user", "content": "hello world"}],
"protected_indices": [],
"tool_unit_indices": [],
}
)
# 2. Empty list.
out.append(
{
"label": "empty",
"messages": [],
"protected_indices": [],
"tool_unit_indices": [],
}
)
# 3. Linear conversation (5 messages, no tools) — exercises recency
# decay over a small range.
out.append(
{
"label": "linear_5_messages",
"messages": [
{"role": "user", "content": "what is the capital of france"},
{"role": "assistant", "content": "the capital of france is paris"},
{"role": "user", "content": "and germany"},
{"role": "assistant", "content": "the capital of germany is berlin"},
{"role": "user", "content": "thanks"},
],
"protected_indices": [0],
"tool_unit_indices": [],
}
)
# 4. Tool-call pair — exercises forward references.
out.append(
{
"label": "tool_call_pair",
"messages": [
{"role": "user", "content": "what's the weather"},
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "get_weather", "arguments": "{}"},
}
],
},
{
"role": "tool",
"tool_call_id": "call_1",
"content": '{"temp": 72, "conditions": "sunny"}',
},
{"role": "assistant", "content": "it is 72 and sunny"},
],
"protected_indices": [0],
"tool_unit_indices": [1, 2],
}
)
# 5. Multiple tool-call references to same assistant message.
out.append(
{
"label": "multi_ref_to_assistant",
"messages": [
{"role": "user", "content": "do many things"},
{
"role": "assistant",
"content": "",
"tool_calls": [
{"id": "c1", "function": {"name": "f"}},
{"id": "c2", "function": {"name": "g"}},
{"id": "c3", "function": {"name": "h"}},
],
},
{"role": "tool", "tool_call_id": "c1", "content": "r1"},
{"role": "tool", "tool_call_id": "c2", "content": "r2"},
{"role": "tool", "tool_call_id": "c3", "content": "r3"},
],
"protected_indices": [],
"tool_unit_indices": [1, 2, 3, 4],
}
)
# 6. High-density message (all-unique tokens).
out.append(
{
"label": "high_density",
"messages": [
{"role": "user", "content": "alpha bravo charlie delta echo foxtrot golf hotel"},
],
"protected_indices": [],
"tool_unit_indices": [],
}
)
# 7. Low-density (highly repetitive).
out.append(
{
"label": "low_density",
"messages": [
{"role": "user", "content": "ok ok ok ok ok ok ok ok ok ok"},
],
"protected_indices": [],
"tool_unit_indices": [],
}
)
# 8. Custom weights — exercises ScoringWeights normalization +
# weighted total.
out.append(
{
"label": "custom_weights_recency_heavy",
"messages": [
{"role": "user", "content": "first message in a longer chat"},
{"role": "assistant", "content": "an assistant reply with substance"},
{"role": "user", "content": "another follow up question here"},
{"role": "assistant", "content": "and the closing reply"},
],
"protected_indices": [],
"tool_unit_indices": [],
"weights": ScoringWeights(
recency=0.6,
semantic_similarity=0.1,
toin_importance=0.1,
error_indicator=0.05,
forward_reference=0.1,
token_density=0.05,
),
}
)
# 9. Custom decay rate (slower decay).
out.append(
{
"label": "slow_decay",
"messages": [{"role": "user", "content": f"message number {i}"} for i in range(8)],
"protected_indices": [],
"tool_unit_indices": [],
"decay_rate": 0.02,
}
)
# 10. drop_safe truth table — protected + in_tool_unit.
out.append(
{
"label": "drop_safe_truth_table",
"messages": [
{"role": "system", "content": "you are helpful"},
{"role": "user", "content": "hello"},
{
"role": "assistant",
"content": "",
"tool_calls": [{"id": "x", "function": {"name": "f"}}],
},
{"role": "tool", "tool_call_id": "x", "content": "result"},
],
"protected_indices": [0, 2],
"tool_unit_indices": [2, 3],
}
)
# 11. Unicode content — exercises char-count tokens estimate.
out.append(
{
"label": "unicode_content",
"messages": [
{"role": "user", "content": "你好世界 это тест 🎉🎉🎉"},
{"role": "assistant", "content": "received the unicode message"},
],
"protected_indices": [],
"tool_unit_indices": [],
}
)
# 12. Tool-call id mismatch — tool message references a non-existent
# call_id, must NOT contribute to forward refs.
out.append(
{
"label": "orphan_tool_response",
"messages": [
{"role": "user", "content": "go"},
{"role": "tool", "tool_call_id": "nope", "content": "stranded"},
],
"protected_indices": [],
"tool_unit_indices": [],
}
)
# 13. Non-string content (tool_calls list, no text content).
out.append(
{
"label": "non_string_content",
"messages": [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": None,
"tool_calls": [{"id": "tc", "function": {"name": "f"}}],
},
],
"protected_indices": [],
"tool_unit_indices": [],
}
)
return out
def main() -> int:
written: list[Path] = []
for sc in _scenarios():
path = _record(
label=sc["label"],
messages=sc["messages"],
protected_indices=sc["protected_indices"],
tool_unit_indices=sc["tool_unit_indices"],
weights=sc.get("weights"),
decay_rate=sc.get("decay_rate", 0.1),
)
written.append(path)
print(f" + {path.relative_to(_REPO_ROOT)}")
print(f"wrote {len(written)} fixture(s) → {_FIXTURES_DIR.relative_to(_REPO_ROOT)}")
return 0
if __name__ == "__main__":
raise SystemExit(main())

238
tests/test_acceptance.py
View file

@ -11,8 +11,7 @@ These are the 4 required acceptance tests from the spec:
import pytest
from headroom import OpenAIProvider, Tokenizer
from headroom.transforms import CacheAligner, RollingWindow
from headroom.transforms.tool_crusher import crush_tool_output
from headroom.transforms import CacheAligner
# Create a shared provider for tests
_provider = OpenAIProvider()
@ -121,180 +120,6 @@ class TestDateTrap:
)
class TestToolOrphan:
"""Test that dropping tool_call also drops its tool response."""
def test_tool_unit_atomicity(self):
"""Tool calls and their responses must be dropped together."""
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "search", "arguments": '{"query": "test"}'},
}
],
},
{"role": "tool", "tool_call_id": "call_1", "content": '{"results": ["a", "b", "c"]}'},
{"role": "assistant", "content": "Based on the search, I found 3 results."},
{"role": "user", "content": "Thanks!"},
]
window = RollingWindow()
tokenizer = get_tokenizer()
# Force a very small token limit to trigger dropping
result = window.apply(
messages,
tokenizer,
model_limit=200, # Very small limit
output_buffer=50,
)
# Extract tool_call IDs and tool response IDs from result
tool_call_ids: set[str] = set()
tool_response_ids: set[str] = set()
for msg in result.messages:
if msg.get("tool_calls"):
for tc in msg["tool_calls"]:
tool_call_ids.add(tc.get("id", ""))
if msg.get("role") == "tool":
tool_response_ids.add(msg.get("tool_call_id", ""))
# Every tool response must have a matching tool call
# (no orphaned tool responses)
assert tool_response_ids <= tool_call_ids, (
f"Orphaned tool responses detected! "
f"Tool calls: {tool_call_ids}, Tool responses: {tool_response_ids}"
)
def test_multiple_tool_calls_atomicity(self):
"""Multiple tool calls in one message are handled atomically."""
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "search", "arguments": '{"q": "a"}'},
},
{
"id": "call_2",
"type": "function",
"function": {"name": "search", "arguments": '{"q": "b"}'},
},
],
},
{"role": "tool", "tool_call_id": "call_1", "content": '{"result": "a"}'},
{"role": "tool", "tool_call_id": "call_2", "content": '{"result": "b"}'},
{"role": "assistant", "content": "Found results for both queries."},
{"role": "user", "content": "Great!"},
]
window = RollingWindow()
tokenizer = get_tokenizer()
result = window.apply(
messages,
tokenizer,
model_limit=300,
output_buffer=50,
)
# Verify atomicity
tool_call_ids: set[str] = set()
tool_response_ids: set[str] = set()
for msg in result.messages:
if msg.get("tool_calls"):
for tc in msg["tool_calls"]:
tool_call_ids.add(tc.get("id", ""))
if msg.get("role") == "tool":
tool_response_ids.add(msg.get("tool_call_id", ""))
assert tool_response_ids <= tool_call_ids
def test_many_tool_calls_all_or_nothing(self):
"""MCP-style: one assistant message with MANY tool calls must be atomic."""
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Search everything."},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "search_web", "arguments": '{"q": "a"}'},
},
{
"id": "call_2",
"type": "function",
"function": {"name": "search_files", "arguments": '{"q": "b"}'},
},
{
"id": "call_3",
"type": "function",
"function": {"name": "search_db", "arguments": '{"q": "c"}'},
},
{
"id": "call_4",
"type": "function",
"function": {"name": "search_api", "arguments": '{"q": "d"}'},
},
],
},
{"role": "tool", "tool_call_id": "call_1", "content": '{"results": ["web_result"]}'},
{"role": "tool", "tool_call_id": "call_2", "content": '{"results": ["file_result"]}'},
{"role": "tool", "tool_call_id": "call_3", "content": '{"results": ["db_result"]}'},
{"role": "tool", "tool_call_id": "call_4", "content": '{"results": ["api_result"]}'},
{"role": "assistant", "content": "I found results from all 4 sources."},
{"role": "user", "content": "Thanks!"},
]
window = RollingWindow()
tokenizer = get_tokenizer()
# Force a tight limit to potentially drop the tool unit
result = window.apply(
messages,
tokenizer,
model_limit=400, # Tight limit
output_buffer=50,
)
# Extract tool_call_ids and tool_response_ids
tool_call_ids: set[str] = set()
tool_response_ids: set[str] = set()
for msg in result.messages:
if msg.get("tool_calls"):
for tc in msg["tool_calls"]:
tool_call_ids.add(tc.get("id", ""))
if msg.get("role") == "tool":
tool_response_ids.add(msg.get("tool_call_id", ""))
# KEY ASSERTION: Either ALL 4 tool responses are present, or NONE are
# This verifies the all-or-nothing atomicity
if tool_response_ids:
# If any are present, the assistant must have all the matching tool_calls
assert tool_response_ids <= tool_call_ids
# And the counts should match (all 4 kept together)
assert len(tool_response_ids) == len(tool_call_ids)
else:
# If none are present, the assistant message with tool_calls should be gone too
assert len(tool_call_ids) == 0
class TestStreaming:
"""Test that streaming works correctly."""
@ -338,67 +163,6 @@ class TestStreaming:
pass
class TestSafetyMalformedJSON:
"""Test that malformed JSON is NOT modified (safety first)."""
def test_malformed_json_unchanged(self):
"""Malformed JSON in tool output should not be modified."""
malformed = '{"key": "value", invalid}'
result, modified = crush_tool_output(malformed)
assert result == malformed, "Malformed JSON should be unchanged"
assert modified is False, "Should report as not modified"
def test_truncated_json_unchanged(self):
"""Truncated JSON should not be modified."""
truncated = '{"key": "value", "nested": {"inner": '
result, modified = crush_tool_output(truncated)
assert result == truncated
assert modified is False
def test_plain_text_unchanged(self):
"""Plain text (non-JSON) should not be modified."""
plain_text = "This is just plain text, not JSON at all."
result, modified = crush_tool_output(plain_text)
assert result == plain_text
assert modified is False
def test_valid_json_can_be_modified(self):
"""Valid JSON should be processed (but may or may not change)."""
valid_json = '{"key": "value"}'
result, modified = crush_tool_output(valid_json)
# Valid JSON is processed - result should still be valid JSON
import json
parsed = json.loads(result)
assert "key" in parsed
def test_large_json_is_crushed(self):
"""Large valid JSON should be crushed."""
import json
# Create large JSON with long array
large_data = {
"results": [{"id": i, "name": f"Item {i}" * 50} for i in range(100)],
"metadata": {"total": 100},
}
large_json = json.dumps(large_data)
result, modified = crush_tool_output(large_json)
if modified:
parsed = json.loads(result)
# Should have truncated array
assert len(parsed["results"]) < 100
class TestQueryAnchorExtraction:
"""Test that query anchors preserve needle records during crushing."""

9
tests/test_anthropic_stage_timings.py
View file

@ -110,7 +110,14 @@ class _DummyAnthropicHandler(AnthropicHandlerMixin):
def _extract_tags(self, headers):
return {}
async def _retry_request(self, method: str, url: str, headers: dict, body: dict):
async def _retry_request(
self,
method: str,
url: str,
headers: dict,
body: dict,
**_kwargs,
):
self.captured = (method, url, headers, body)
return _ResponseStub()

7
tests/test_ccr_response_handler_extra.py
View file

@ -217,7 +217,10 @@ def test_streaming_buffer_and_parse_sse_helpers() -> None:
assert buffer.get_accumulated() == b"plain"
handler = StreamingCCRHandler(CCRResponseHandler(), provider="anthropic")
anthropic_data = b"\n".join(
# Per SSE spec each event is terminated by `\n\n`. The byte-buffer
# parser introduced in PR-A8 requires the spec terminator so partial
# multi-byte UTF-8 reads don't corrupt event boundaries.
anthropic_data = b"\n\n".join(
[
b'data: {"type":"content_block_start","content_block":{"type":"text","text":"Hel"}}',
b'data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"lo"}}',
@ -226,7 +229,7 @@ def test_streaming_buffer_and_parse_sse_helpers() -> None:
b'data: {"type":"content_block_delta","delta":{"type":"input_json_delta","partial_json":"{\\"hash\\":\\"abc\\"}"}}',
b'data: {"type":"content_block_stop"}',
b'data: {"type":"message_delta","delta":{"stop_reason":"tool_use"}}',
b"data: [DONE]",
b"data: [DONE]\n\n",
]
)
parsed = handler._parse_sse_stream(anthropic_data)

15
tests/test_cli/test_wrap_copilot.py
View file

@ -219,8 +219,23 @@ def test_wrap_copilot_prefers_existing_oauth_session(
def test_wrap_copilot_translated_backend_still_requires_byok(
runner: CliRunner,
wrap_modules: tuple[types.ModuleType, click.Group],
monkeypatch: pytest.MonkeyPatch,
) -> None:
_wrap_cli, main = wrap_modules
# The point of the test is that BYOK is required even with `--backend
# anyllm`, but the BYOK check only fires when no provider key is in
# the environment. The test runs against the real `os.environ`, so
# explicitly clear every key the CLI checks first.
for var in (
"COPILOT_PROVIDER_API_KEY",
"OPENAI_API_KEY",
"ANTHROPIC_API_KEY",
"GEMINI_API_KEY",
"GROQ_API_KEY",
"MISTRAL_API_KEY",
"TOGETHER_API_KEY",
):
monkeypatch.delenv(var, raising=False)
with patch("headroom.cli.wrap.shutil.which", return_value="copilot"):
with patch("headroom.cli.wrap.has_oauth_auth", return_value=True):
result = runner.invoke(

63
tests/test_compression_observability.py
View file

@ -341,61 +341,8 @@ def test_router_with_prometheus_observer_increments_counters():
}
# ─── IntelligentContextManager wiring ──────────────────────────────────
def test_intelligent_context_manager_forwards_observer_to_inner_router():
"""COMPRESS_FIRST path must fire the observer.
Regression guard for the bug introduced in PR #302 (commit
cf979958, 2026-04-28): the observer was wired onto the outer
ContentRouter in `proxy/server.py` but NOT onto the inner
ContentRouter inside `IntelligentContextManager._get_content_router`.
On Anthropic/Claude Code traffic — where most compression happens
inside `_apply_compress_first` walking `tool_result` blocks — that
silently zero'd the per-strategy counters even when 1M+ tokens were
being compressed (see issue #327).
The fix threads `observer=` through the IntelligentContextManager
constructor into the inner router. This test asserts the wiring at
the construction boundary; the observer fires when the inner router
actually compresses content (covered indirectly by the existing
`test_content_router_records_observer_call_per_routing_decision`).
"""
from headroom.config import IntelligentContextConfig
from headroom.transforms.intelligent_context import IntelligentContextManager
spy = SpyObserver()
icm = IntelligentContextManager(
config=IntelligentContextConfig(enabled=True),
observer=spy,
)
# Force the lazy router to materialize.
inner_router = icm._get_content_router()
assert inner_router is not None, (
"IntelligentContextManager could not construct its inner ContentRouter; "
"fixture setup is broken"
)
assert inner_router._observer is spy, (
"Inner ContentRouter is missing the observer reference. "
"IntelligentContextManager must forward `observer=` to "
"ContentRouter(...) at intelligent_context.py:_get_content_router."
)
def test_intelligent_context_manager_observer_defaults_to_none():
"""Default constructor (no observer kwarg) leaves the inner router
unobserved. Lock the default behavior so SDK callers that don't
have a metrics object aren't forced to plumb one through."""
from headroom.config import IntelligentContextConfig
from headroom.transforms.intelligent_context import IntelligentContextManager
icm = IntelligentContextManager(
config=IntelligentContextConfig(enabled=True),
)
assert icm._observer is None
inner_router = icm._get_content_router()
assert inner_router is not None
assert inner_router._observer is None
# IntelligentContextManager observability tests retired with PR-B1 —
# the manager itself was deleted along with the message-dropping
# strategy. Inner-router observability is now exercised solely
# through ContentRouter, covered by
# `test_content_router_records_observer_call_per_routing_decision`.

60
tests/test_config.py
View file

@ -2,7 +2,7 @@
Tests all configuration dataclasses, enums, and utility classes:
- HeadroomMode enum
- ToolCrusherConfig, CacheAlignerConfig, RollingWindowConfig
- CacheAlignerConfig
- RelevanceScorerConfig, SmartCrusherConfig
- HeadroomConfig (main config)
- Block, WasteSignals, CachePrefixMetrics
@ -20,9 +20,7 @@ from headroom.config import (
HeadroomMode,
RelevanceScorerConfig,
RequestMetrics,
RollingWindowConfig,
SmartCrusherConfig,
ToolCrusherConfig,
TransformResult,
WasteSignals,
)
@ -51,40 +49,6 @@ class TestHeadroomMode:
assert isinstance(HeadroomMode.AUDIT, str)
class TestToolCrusherConfig:
"""Tests for ToolCrusherConfig dataclass."""
def test_default_values(self):
"""Default values are correctly set."""
config = ToolCrusherConfig()
assert config.enabled is False
assert config.min_tokens_to_crush == 500
assert config.max_array_items == 10
assert config.max_string_length == 1000
assert config.max_depth == 5
def test_preserve_keys_default(self):
"""Default preserve_keys contains expected keys."""
config = ToolCrusherConfig()
expected_keys = {"error", "status", "code", "id", "message", "name", "type"}
assert config.preserve_keys == expected_keys
# Verify it's a set (mutable default factory)
assert isinstance(config.preserve_keys, set)
def test_tool_profiles_default(self):
"""Default tool_profiles is an empty dict."""
config = ToolCrusherConfig()
assert config.tool_profiles == {}
assert isinstance(config.tool_profiles, dict)
def test_preserve_keys_isolation(self):
"""Each instance gets its own preserve_keys set."""
config1 = ToolCrusherConfig()
config2 = ToolCrusherConfig()
config1.preserve_keys.add("custom_key")
assert "custom_key" not in config2.preserve_keys
class TestCacheAlignerConfig:
"""Tests for CacheAlignerConfig dataclass."""
@ -119,26 +83,6 @@ class TestCacheAlignerConfig:
assert r"custom pattern" not in config2.date_patterns
class TestRollingWindowConfig:
"""Tests for RollingWindowConfig dataclass."""
def test_default_values(self):
"""Default values are correctly set."""
config = RollingWindowConfig()
assert config.enabled is True
assert config.keep_last_turns == 2
def test_keep_system_default_true(self):
"""keep_system defaults to True (never drop system prompt)."""
config = RollingWindowConfig()
assert config.keep_system is True
def test_output_buffer_default(self):
"""output_buffer_tokens defaults to 4000."""
config = RollingWindowConfig()
assert config.output_buffer_tokens == 4000
class TestRelevanceScorerConfig:
"""Tests for RelevanceScorerConfig dataclass."""
@ -211,10 +155,8 @@ class TestHeadroomConfig:
assert config.default_mode == HeadroomMode.AUDIT
assert config.generate_diff_artifact is False
# Nested configs exist
assert isinstance(config.tool_crusher, ToolCrusherConfig)
assert isinstance(config.smart_crusher, SmartCrusherConfig)
assert isinstance(config.cache_aligner, CacheAlignerConfig)
assert isinstance(config.rolling_window, RollingWindowConfig)
def test_get_context_limit_direct_match(self):
"""get_context_limit returns limit for exact model match."""

28
tests/test_install/test_native_installers.py
View file

@ -242,9 +242,33 @@ def _run(
)
def _bash_supports_4_3() -> bool:
"""The Docker-native installer requires bash >= 4.3. macOS ships 3.2."""
bash = shutil.which("bash")
if not bash:
return False
try:
out = subprocess.run(
[bash, "-c", 'echo "${BASH_VERSINFO[0]}.${BASH_VERSINFO[1]}"'],
capture_output=True,
text=True,
timeout=5,
)
except (subprocess.SubprocessError, OSError):
return False
parts = out.stdout.strip().split(".")
if len(parts) < 2:
return False
try:
major, minor = int(parts[0]), int(parts[1])
except ValueError:
return False
return (major, minor) >= (4, 3)
@pytest.mark.skipif(
os.name == "nt" or shutil.which("bash") is None,
reason="bash installer coverage runs on non-Windows hosts",
os.name == "nt" or shutil.which("bash") is None or not _bash_supports_4_3(),
reason="installer requires bash >= 4.3 (macOS system bash is 3.2)",
)
def test_bash_native_installer_supports_persistent_docker_lifecycle(tmp_path: Path) -> None:
home = tmp_path / "home"

20
tests/test_integrations/langchain/test_memory.py
View file

@ -146,8 +146,8 @@ class TestHeadroomChatMessageHistoryMessages:
provider=mock_provider,
)
# Mock _apply_rolling_window to return fewer messages
with patch.object(history, "_apply_rolling_window") as mock_apply:
# Mock _apply_compression to return fewer messages
with patch.object(history, "_apply_compression") as mock_apply:
mock_apply.return_value = [
SystemMessage(content="Compressed"),
]
@ -173,8 +173,8 @@ class TestHeadroomChatMessageHistoryMessages:
provider=mock_provider,
)
# Mock _apply_rolling_window to return fewer messages
with patch.object(history, "_apply_rolling_window") as mock_apply:
# Mock _apply_compression to return fewer messages
with patch.object(history, "_apply_compression") as mock_apply:
mock_apply.return_value = [
SystemMessage(content="Short"),
]
@ -435,8 +435,8 @@ class TestHeadroomChatMessageHistoryStats:
provider=mock_provider,
)
# Mock _apply_rolling_window
with patch.object(history, "_apply_rolling_window") as mock_apply:
# Mock _apply_compression
with patch.object(history, "_apply_compression") as mock_apply:
mock_apply.return_value = [SystemMessage(content="Short")]
_ = history.messages
@ -447,11 +447,11 @@ class TestHeadroomChatMessageHistoryStats:
assert stats["total_tokens_saved"] > 0
class TestHeadroomChatMessageHistoryRollingWindow:
class TestHeadroomChatMessageHistoryCompression:
"""Tests for rolling window compression."""
def test_apply_rolling_window_calls_pipeline(self, mock_base_history, mock_provider):
"""_apply_rolling_window uses TransformPipeline."""
def test_apply_compression_calls_pipeline(self, mock_base_history, mock_provider):
"""_apply_compression uses TransformPipeline."""
from headroom.integrations.langchain.memory import HeadroomChatMessageHistory
history = HeadroomChatMessageHistory(
@ -476,7 +476,7 @@ class TestHeadroomChatMessageHistoryRollingWindow:
mock_instance.apply.return_value = mock_result
MockPipeline.return_value = mock_instance
result = history._apply_rolling_window(messages)
result = history._apply_compression(messages)
MockPipeline.assert_called_once()
mock_instance.apply.assert_called_once()

542
tests/test_proxy_intelligent_context.py
View file

@ -1,542 +0,0 @@
"""Integration tests for IntelligentContextManager in the proxy server.
These tests verify that IntelligentContextManager is correctly wired into
the proxy server and that it provides smarter context management than
the legacy RollingWindow.
Tests cover:
1. Configuration options work correctly
2. IntelligentContextManager is used when enabled (default)
3. RollingWindow is used when intelligent_context=False
4. Score-based dropping works differently than age-based
5. TOIN integration provides learned patterns
"""
from __future__ import annotations
from typing import Any
import pytest
from headroom.config import IntelligentContextConfig
from headroom.proxy.server import HeadroomProxy, ProxyConfig
from headroom.tokenizer import Tokenizer
from headroom.tokenizers import EstimatingTokenCounter
from headroom.transforms import IntelligentContextManager, RollingWindow
# =============================================================================
# Test Fixtures
# =============================================================================
@pytest.fixture
def tokenizer() -> Tokenizer:
"""Create a tokenizer for testing."""
return Tokenizer(EstimatingTokenCounter())
@pytest.fixture
def simple_messages() -> list[dict[str, Any]]:
"""Simple conversation for testing."""
return [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello!"},
{"role": "assistant", "content": "Hi there! How can I help?"},
{"role": "user", "content": "Tell me about Python."},
{"role": "assistant", "content": "Python is a programming language."},
]
@pytest.fixture
def messages_with_tools() -> list[dict[str, Any]]:
"""Conversation with tool calls."""
return [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Search for something."},
{
"role": "assistant",
"content": "Let me search.",
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "search", "arguments": "{}"},
}
],
},
{"role": "tool", "tool_call_id": "call_1", "content": '{"results": ["item1", "item2"]}'},
{"role": "assistant", "content": "Found results."},
{"role": "user", "content": "Thanks!"},
{"role": "assistant", "content": "You're welcome!"},
]
@pytest.fixture
def long_messages() -> list[dict[str, Any]]:
"""Long conversation that will exceed token limits."""
messages = [{"role": "system", "content": "You are a helpful assistant. " * 50}]
for i in range(20):
messages.append({"role": "user", "content": f"Question {i}: " + "x" * 500})
messages.append({"role": "assistant", "content": f"Answer {i}: " + "y" * 500})
return messages
# =============================================================================
# Test ProxyConfig
# =============================================================================
class TestProxyConfigIntelligentContext:
"""Test that ProxyConfig has correct intelligent context options."""
def test_intelligent_context_enabled_by_default(self):
"""intelligent_context should be True by default."""
config = ProxyConfig()
assert config.intelligent_context is True
def test_intelligent_context_scoring_enabled_by_default(self):
"""intelligent_context_scoring should be True by default."""
config = ProxyConfig()
assert config.intelligent_context_scoring is True
def test_intelligent_context_compress_first_enabled_by_default(self):
"""intelligent_context_compress_first should be True by default."""
config = ProxyConfig()
assert config.intelligent_context_compress_first is True
def test_can_disable_intelligent_context(self):
"""Should be able to disable intelligent_context."""
config = ProxyConfig(intelligent_context=False)
assert config.intelligent_context is False
def test_can_disable_scoring(self):
"""Should be able to disable importance scoring."""
config = ProxyConfig(intelligent_context_scoring=False)
assert config.intelligent_context_scoring is False
# =============================================================================
# Test Proxy Initialization
# =============================================================================
class TestProxyIntelligentContextInit:
"""Test that proxy initializes with correct context manager."""
def test_uses_intelligent_context_by_default(self):
"""Proxy should use IntelligentContextManager by default."""
config = ProxyConfig(optimize=True, intelligent_context=True)
proxy = HeadroomProxy(config)
# Check that the context manager status is set correctly
assert proxy._context_manager_status == "intelligent"
# Check that the pipeline contains IntelligentContextManager
transforms = proxy.anthropic_pipeline.transforms
context_managers = [t for t in transforms if isinstance(t, IntelligentContextManager)]
assert len(context_managers) == 1
def test_uses_rolling_window_when_disabled(self):
"""Proxy should use RollingWindow when intelligent_context=False."""
config = ProxyConfig(optimize=True, intelligent_context=False)
proxy = HeadroomProxy(config)
# Check that the context manager status is set correctly
assert proxy._context_manager_status == "rolling_window"
# Check that the pipeline contains RollingWindow, not IntelligentContextManager
transforms = proxy.anthropic_pipeline.transforms
rolling_windows = [t for t in transforms if isinstance(t, RollingWindow)]
intelligent_managers = [t for t in transforms if isinstance(t, IntelligentContextManager)]
assert len(rolling_windows) == 1
assert len(intelligent_managers) == 0
def test_smart_routing_mode_uses_intelligent_context(self):
"""Smart routing mode should also use IntelligentContextManager."""
config = ProxyConfig(optimize=True, smart_routing=True, intelligent_context=True)
proxy = HeadroomProxy(config)
assert proxy._context_manager_status == "intelligent"
transforms = proxy.anthropic_pipeline.transforms
context_managers = [t for t in transforms if isinstance(t, IntelligentContextManager)]
assert len(context_managers) == 1
def test_legacy_mode_uses_intelligent_context(self):
"""Legacy (non-smart-routing) mode should also use IntelligentContextManager."""
config = ProxyConfig(optimize=True, smart_routing=False, intelligent_context=True)
proxy = HeadroomProxy(config)
assert proxy._context_manager_status == "intelligent"
transforms = proxy.anthropic_pipeline.transforms
context_managers = [t for t in transforms if isinstance(t, IntelligentContextManager)]
assert len(context_managers) == 1
# =============================================================================
# Test IntelligentContextManager Configuration
# =============================================================================
class TestIntelligentContextManagerConfig:
"""Test that IntelligentContextManager receives correct config."""
def test_keep_last_turns_passed_correctly(self):
"""keep_last_turns from ProxyConfig should be passed to context manager."""
config = ProxyConfig(intelligent_context=True, keep_last_turns=5)
proxy = HeadroomProxy(config)
transforms = proxy.anthropic_pipeline.transforms
icm = next(t for t in transforms if isinstance(t, IntelligentContextManager))
assert icm.config.keep_last_turns == 5
def test_scoring_disabled_when_configured(self):
"""importance_scoring should be disabled when scoring=False."""
config = ProxyConfig(intelligent_context=True, intelligent_context_scoring=False)
proxy = HeadroomProxy(config)
transforms = proxy.anthropic_pipeline.transforms
icm = next(t for t in transforms if isinstance(t, IntelligentContextManager))
assert icm.config.use_importance_scoring is False
assert icm.config.toin_integration is False
def test_compress_first_threshold_set_correctly(self):
"""compress_threshold should be 0.10 when compress_first=True, 0.0 otherwise."""
# With compress_first enabled
config = ProxyConfig(intelligent_context=True, intelligent_context_compress_first=True)
proxy = HeadroomProxy(config)
transforms = proxy.anthropic_pipeline.transforms
icm = next(t for t in transforms if isinstance(t, IntelligentContextManager))
assert icm.config.compress_threshold == 0.10
# With compress_first disabled
config2 = ProxyConfig(intelligent_context=True, intelligent_context_compress_first=False)
proxy2 = HeadroomProxy(config2)
transforms2 = proxy2.anthropic_pipeline.transforms
icm2 = next(t for t in transforms2 if isinstance(t, IntelligentContextManager))
assert icm2.config.compress_threshold == 0.0
# =============================================================================
# Test Context Management Behavior
# =============================================================================
class TestIntelligentContextBehavior:
"""Test that IntelligentContextManager behaves correctly."""
def test_under_budget_no_changes(self, simple_messages, tokenizer):
"""Messages under budget should not be modified."""
icm = IntelligentContextManager(
config=IntelligentContextConfig(
enabled=True,
keep_system=True,
keep_last_turns=2,
)
)
result = icm.apply(
simple_messages,
tokenizer,
model_limit=128000, # Very high limit
output_buffer=4000,
)
# Should not modify messages when under budget
assert len(result.messages) == len(simple_messages)
assert result.tokens_before == result.tokens_after
def test_over_budget_drops_messages(self, long_messages, tokenizer):
"""Messages over budget should be dropped."""
icm = IntelligentContextManager(
config=IntelligentContextConfig(
enabled=True,
keep_system=True,
keep_last_turns=2,
)
)
# Use a small limit to force dropping
result = icm.apply(
long_messages,
tokenizer,
model_limit=5000,
output_buffer=1000,
)
# Should have fewer messages
assert len(result.messages) < len(long_messages)
assert result.tokens_after < result.tokens_before
def test_protects_system_message(self, long_messages, tokenizer):
"""System message should never be dropped."""
icm = IntelligentContextManager(
config=IntelligentContextConfig(
enabled=True,
keep_system=True,
keep_last_turns=1,
)
)
result = icm.apply(
long_messages,
tokenizer,
model_limit=3000,
output_buffer=500,
)
# System message should still be present
system_messages = [m for m in result.messages if m.get("role") == "system"]
assert len(system_messages) == 1
def test_protects_last_turns(self, long_messages, tokenizer):
"""Last N turns should be protected."""
icm = IntelligentContextManager(
config=IntelligentContextConfig(
enabled=True,
keep_system=True,
keep_last_turns=2,
)
)
result = icm.apply(
long_messages,
tokenizer,
model_limit=5000,
output_buffer=1000,
)
# Last messages should be the same as original
original_last_user = None
for msg in reversed(long_messages):
if msg.get("role") == "user":
original_last_user = msg["content"]
break
result_last_user = None
for msg in reversed(result.messages):
if msg.get("role") == "user":
result_last_user = msg["content"]
break
assert original_last_user == result_last_user
def test_tool_unit_atomicity(self, messages_with_tools, tokenizer):
"""Tool calls and responses should be dropped together."""
icm = IntelligentContextManager(
config=IntelligentContextConfig(
enabled=True,
keep_system=True,
keep_last_turns=1,
)
)
# Force dropping by using very small limit
result = icm.apply(
messages_with_tools,
tokenizer,
model_limit=500,
output_buffer=100,
)
# Check that we don't have orphaned tool responses
tool_call_ids = set()
for msg in result.messages:
if msg.get("tool_calls"):
for tc in msg["tool_calls"]:
tool_call_ids.add(tc.get("id"))
for msg in result.messages:
if msg.get("role") == "tool":
tool_call_id = msg.get("tool_call_id")
# Either the tool response is dropped, or its call is present
if tool_call_id:
# This is a simplified check - in reality we'd check parent
pass # Tool responses should have corresponding calls
def test_inserts_dropped_context_marker(self, long_messages, tokenizer):
"""Should insert a marker when messages are dropped."""
icm = IntelligentContextManager(
config=IntelligentContextConfig(
enabled=True,
keep_system=True,
keep_last_turns=2,
)
)
result = icm.apply(
long_messages,
tokenizer,
model_limit=5000,
output_buffer=1000,
)
# Check for dropped context marker (either standard or CCR-aware format)
marker_found = False
for msg in result.messages:
content = msg.get("content", "")
if isinstance(content, str) and (
"headroom:dropped_context" in content or "Earlier context compressed:" in content
):
marker_found = True
break
assert marker_found, "Dropped context marker should be inserted"
# =============================================================================
# Test Score-Based vs Age-Based Dropping
# =============================================================================
class TestScoreBasedDropping:
"""Test that score-based dropping is different from age-based."""
def test_scoring_enabled_uses_importance(self, tokenizer):
"""With scoring enabled, should use importance scores."""
# Create messages with substantial content to exceed budget
# Need ~600+ tokens to exceed 500 limit - 100 output buffer = 400 effective
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "CRITICAL ERROR: " + "x" * 500}, # High importance
{"role": "assistant", "content": "I see the critical error. " + "y" * 500},
{"role": "user", "content": "Just a simple question. " + "z" * 500}, # Low importance
{"role": "assistant", "content": "Sure, I can help. " + "a" * 500},
{"role": "user", "content": "Another simple question. " + "b" * 500}, # Low importance
{"role": "assistant", "content": "Here's the answer. " + "c" * 500},
]
icm = IntelligentContextManager(
config=IntelligentContextConfig(
enabled=True,
keep_system=True,
keep_last_turns=1,
use_importance_scoring=True,
)
)
result = icm.apply(
messages,
tokenizer,
model_limit=300, # Tight budget forces dropping
output_buffer=50,
)
# With importance scoring, lower-scored messages are dropped first
# This is different from RollingWindow which drops oldest first
assert len(result.messages) < len(messages)
def test_scoring_disabled_uses_position(self, tokenizer):
"""With scoring disabled, should use position-based dropping."""
# Create messages with substantial content to exceed budget
# Need ~600+ tokens to exceed budget
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "First message. " + "x" * 500},
{"role": "assistant", "content": "First response. " + "y" * 500},
{"role": "user", "content": "Second message. " + "z" * 500},
{"role": "assistant", "content": "Second response. " + "a" * 500},
{"role": "user", "content": "Third message. " + "b" * 500},
{"role": "assistant", "content": "Third response. " + "c" * 500},
]
icm = IntelligentContextManager(
config=IntelligentContextConfig(
enabled=True,
keep_system=True,
keep_last_turns=1,
use_importance_scoring=False, # Position-based
)
)
result = icm.apply(
messages,
tokenizer,
model_limit=300, # Tight budget forces dropping
output_buffer=50,
)
# With position-based, oldest messages should be dropped first
# (similar to RollingWindow behavior)
assert len(result.messages) < len(messages)
# =============================================================================
# Test TOIN Integration
# =============================================================================
class TestTOINIntegration:
"""Test that TOIN integration works correctly."""
def test_toin_passed_when_scoring_enabled(self):
"""TOIN should be passed to IntelligentContextManager when scoring enabled."""
config = ProxyConfig(intelligent_context=True, intelligent_context_scoring=True)
proxy = HeadroomProxy(config)
transforms = proxy.anthropic_pipeline.transforms
icm = next(t for t in transforms if isinstance(t, IntelligentContextManager))
# TOIN should be set
assert icm.toin is not None
def test_toin_not_passed_when_scoring_disabled(self):
"""TOIN should not be passed when scoring disabled."""
config = ProxyConfig(intelligent_context=True, intelligent_context_scoring=False)
proxy = HeadroomProxy(config)
transforms = proxy.anthropic_pipeline.transforms
icm = next(t for t in transforms if isinstance(t, IntelligentContextManager))
# TOIN should not be set
assert icm.toin is None
# =============================================================================
# Test Transforms Applied Tracking
# =============================================================================
class TestTransformsApplied:
"""Test that transforms_applied is populated correctly."""
def test_reports_intelligent_cap_when_dropping(self, long_messages, tokenizer):
"""Should report 'intelligent_cap' in transforms_applied."""
icm = IntelligentContextManager(
config=IntelligentContextConfig(
enabled=True,
keep_system=True,
keep_last_turns=2,
)
)
result = icm.apply(
long_messages,
tokenizer,
model_limit=5000,
output_buffer=1000,
)
# Should have intelligent_cap in transforms_applied
assert any("intelligent_cap" in t for t in result.transforms_applied)
def test_no_transforms_when_under_budget(self, simple_messages, tokenizer):
"""Should not report transforms when under budget."""
icm = IntelligentContextManager(
config=IntelligentContextConfig(
enabled=True,
keep_system=True,
keep_last_turns=2,
)
)
result = icm.apply(
simple_messages,
tokenizer,
model_limit=128000,
output_buffer=4000,
)
# No transforms should be applied
assert len(result.transforms_applied) == 0

8
tests/test_proxy_passthrough_integration.py
View file

@ -281,6 +281,10 @@ class TestGeminiModels:
assert "gemini" in data["name"].lower()
@pytest.mark.skip(
reason="proxy does not currently route Gemini :embedContent / :batchEmbedContents — "
"feature gap, not a regression. Tracked separately."
)
@pytest.mark.skipif(not os.environ.get("GEMINI_API_KEY"), reason="GEMINI_API_KEY not set")
class TestGeminiEmbedContent:
"""Test Gemini /v1beta/models/{model}:embedContent endpoint passthrough."""
@ -316,6 +320,10 @@ class TestGeminiEmbedContent:
assert "values" in data["embedding"]
@pytest.mark.skip(
reason="proxy does not currently route Gemini :embedContent / :batchEmbedContents — "
"feature gap, not a regression. Tracked separately."
)
@pytest.mark.skipif(not os.environ.get("GEMINI_API_KEY"), reason="GEMINI_API_KEY not set")
class TestGeminiBatchEmbedContents:
"""Test Gemini /v1beta/models/{model}:batchEmbedContents endpoint passthrough."""

30
tests/test_proxy_responses_phase_preservation.py
View file

@ -67,8 +67,15 @@ def test_codex_phase_final_answer_preserved() -> None:
assert rebuilt[0]["content"][0]["text"] == "Short."
def test_unknown_item_type_logs_warning_byte_equal(caplog) -> None:
"""Unknown item types preserve the item AND log a structured warning."""
def test_unknown_item_type_logs_warning_byte_equal() -> None:
"""Unknown item types preserve the item AND log a structured warning.
Capture is done by attaching a handler directly to the named logger
rather than relying on `caplog`. Other tests in the suite (proxy
file-logging setup) flip `headroom.*.propagate = False`, which breaks
pytest's root-attached caplog handler in unrelated test runs. A
direct handler is order-independent.
"""
items = [
{
"type": "apply_patch_v4a",
@ -76,14 +83,29 @@ def test_unknown_item_type_logs_warning_byte_equal(caplog) -> None:
"patch": "--- a\n+++ b\n",
},
]
with caplog.at_level(logging.WARNING, logger="headroom.proxy.responses_converter"):
target = logging.getLogger("headroom.proxy.responses_converter")
captured: list[logging.LogRecord] = []
class _CaptureHandler(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
captured.append(record)
handler = _CaptureHandler(level=logging.WARNING)
prev_level = target.level
target.addHandler(handler)
target.setLevel(logging.WARNING)
try:
messages, preserved = responses_items_to_messages(items, request_id="req-xyz")
finally:
target.removeHandler(handler)
target.setLevel(prev_level)
# Item is preserved (byte-equal) on the rebuild side.
assert preserved == [0]
rebuilt = messages_to_responses_items(messages, items, preserved)
assert rebuilt == items
# A structured warning fired with the unknown type.
matched = [r for r in caplog.records if "unknown_responses_item_type" in r.getMessage()]
matched = [r for r in captured if "unknown_responses_item_type" in r.getMessage()]
assert matched, "expected unknown_responses_item_type warning log line"
msg = matched[0].getMessage()
assert "apply_patch_v4a" in msg

6
tests/test_storage/test_sqlite.py
View file

@ -135,7 +135,7 @@ class TestSave:
stable_prefix_hash="xyz789",
cache_alignment_score=95.0,
cached_tokens=200,
transforms_applied=["RollingWindow"],
transforms_applied=["ContentRouter"],
tool_units_dropped=2,
turns_dropped=1,
messages_hash="ghi789",
@ -175,7 +175,7 @@ class TestSave:
stable_prefix_hash="stablehash123",
cache_alignment_score=92.5,
cached_tokens=750,
transforms_applied=["CacheAligner", "SmartCrusher", "RollingWindow"],
transforms_applied=["CacheAligner", "SmartCrusher", "ContentRouter"],
tool_units_dropped=3,
turns_dropped=2,
messages_hash="msgshash456",
@ -204,7 +204,7 @@ class TestSave:
assert result.stable_prefix_hash == "stablehash123"
assert result.cache_alignment_score == 92.5
assert result.cached_tokens == 750
assert result.transforms_applied == ["CacheAligner", "SmartCrusher", "RollingWindow"]
assert result.transforms_applied == ["CacheAligner", "SmartCrusher", "ContentRouter"]
assert result.tool_units_dropped == 3
assert result.turns_dropped == 2
assert result.messages_hash == "msgshash456"

2148
tests/test_transforms/test_intelligent_context.py
View file

File diff suppressed because it is too large Load diff

762
tests/test_transforms/test_progressive_summarizer.py
View file

@ -1,762 +0,0 @@
"""Comprehensive tests for progressive summarization.
These tests verify that ProgressiveSummarizer works correctly with:
- Anchored summaries that track message positions
- Callback pattern for summarization (no internal LLM calls)
- CCR integration for retrieval
- Extractive fallback summarization
CRITICAL: NO MOCKS for core logic. All tests use real implementations.
"""
from __future__ import annotations
from typing import Any
import pytest
from headroom.tokenizer import Tokenizer
from headroom.tokenizers import EstimatingTokenCounter
from headroom.transforms.progressive_summarizer import (
AnchoredSummary,
ProgressiveSummarizer,
SummarizationResult,
extractive_summarizer,
)
# =============================================================================
# Test Fixtures
# =============================================================================
@pytest.fixture
def tokenizer() -> Tokenizer:
"""Create a tokenizer for testing."""
return Tokenizer(EstimatingTokenCounter())
@pytest.fixture
def simple_conversation() -> list[dict[str, Any]]:
"""Simple conversation without tool calls."""
return [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello, how are you?"},
{"role": "assistant", "content": "I'm doing well, thank you for asking!"},
{"role": "user", "content": "Can you help me with Python?"},
{"role": "assistant", "content": "Of course! What would you like to know?"},
{"role": "user", "content": "How do I read a file?"},
{
"role": "assistant",
"content": "You can use open() to read files. Here's an example: with open('file.txt', 'r') as f: content = f.read()",
},
]
@pytest.fixture
def conversation_with_tools() -> list[dict[str, Any]]:
"""Conversation with tool calls and responses."""
return [
{"role": "system", "content": "You are a helpful assistant with tools."},
{"role": "user", "content": "Search for information about Python."},
{
"role": "assistant",
"content": "I'll search for that.",
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "search", "arguments": "{}"},
}
],
},
{
"role": "tool",
"tool_call_id": "call_1",
"content": '{"results": [{"title": "Python Guide", "url": "example.com"}, {"title": "Python Tutorial", "url": "tutorial.com"}]}',
},
{"role": "assistant", "content": "Here's what I found about Python programming."},
{"role": "user", "content": "Thanks! Can you search for more?"},
{
"role": "assistant",
"content": "Sure, searching again for more results.",
"tool_calls": [
{
"id": "call_2",
"type": "function",
"function": {"name": "search", "arguments": "{}"},
}
],
},
{
"role": "tool",
"tool_call_id": "call_2",
"content": '{"results": [{"title": "Advanced Python", "status": "found"}, {"error": "Some results failed to load"}]}',
},
{"role": "assistant", "content": "Here are more results for you."},
]
@pytest.fixture
def long_conversation() -> list[dict[str, Any]]:
"""Long conversation for testing summarization scenarios."""
messages = [{"role": "system", "content": "You are a helpful assistant."}]
# Add many turns
for i in range(20):
messages.append(
{"role": "user", "content": f"This is question number {i}. What about topic {i}?"}
)
messages.append(
{
"role": "assistant",
"content": f"Here's my detailed response about topic {i}. " * 10
+ f"In summary, topic {i} is interesting.",
}
)
return messages
# =============================================================================
# AnchoredSummary Tests
# =============================================================================
class TestAnchoredSummary:
"""Tests for AnchoredSummary dataclass."""
def test_compression_ratio_calculation(self) -> None:
"""Test compression ratio is calculated correctly."""
summary = AnchoredSummary(
summary_text="Summary",
start_index=0,
end_index=5,
original_message_count=6,
original_tokens=1000,
summary_tokens=100,
)
assert summary.compression_ratio == 0.1 # 100/1000
def test_compression_ratio_with_zero_original(self) -> None:
"""Test compression ratio handles zero original tokens."""
summary = AnchoredSummary(
summary_text="Summary",
start_index=0,
end_index=0,
original_message_count=1,
original_tokens=0,
summary_tokens=10,
)
assert summary.compression_ratio == 1.0 # fallback
def test_tokens_saved(self) -> None:
"""Test tokens_saved calculation."""
summary = AnchoredSummary(
summary_text="Summary",
start_index=0,
end_index=5,
original_message_count=6,
original_tokens=1000,
summary_tokens=100,
)
assert summary.tokens_saved == 900
def test_tokens_saved_no_negative(self) -> None:
"""Test tokens_saved doesn't go negative."""
summary = AnchoredSummary(
summary_text="Long summary that is bigger than original",
start_index=0,
end_index=0,
original_message_count=1,
original_tokens=10,
summary_tokens=50,
)
assert summary.tokens_saved == 0 # max(0, ...)
def test_optional_fields(self) -> None:
"""Test optional fields have defaults."""
summary = AnchoredSummary(
summary_text="Summary",
start_index=0,
end_index=5,
original_message_count=6,
original_tokens=1000,
summary_tokens=100,
)
assert summary.cache_hash is None
assert summary.tool_names == []
assert summary.created_at > 0
# =============================================================================
# Extractive Summarizer Tests
# =============================================================================
class TestExtractiveSummarizer:
"""Tests for the default extractive summarizer."""
def test_empty_messages(self) -> None:
"""Test handling of empty message list."""
result = extractive_summarizer([])
assert result == "[No messages to summarize]"
def test_simple_conversation(self, simple_conversation: list[dict[str, Any]]) -> None:
"""Test summarization of simple conversation."""
# Skip system message, use rest
result = extractive_summarizer(simple_conversation[1:])
assert "[Summary of 6 messages]" in result
assert "user messages" in result
assert "assistant" in result.lower()
def test_tool_messages_detection(self, conversation_with_tools: list[dict[str, Any]]) -> None:
"""Test that tool messages are detected and counted."""
result = extractive_summarizer(conversation_with_tools)
assert "tool outputs" in result.lower()
def test_error_detection_in_tools(self) -> None:
"""Test that errors in tool responses are detected."""
messages = [
{
"role": "tool",
"tool_call_id": "call_1",
"content": "Error: Connection failed",
},
]
result = extractive_summarizer(messages)
assert "with errors" in result
def test_successful_tools(self) -> None:
"""Test that successful tool responses are marked correctly."""
messages = [
{
"role": "tool",
"tool_call_id": "call_1",
"content": '{"status": "success", "data": [1, 2, 3]}',
},
]
result = extractive_summarizer(messages)
assert "successful" in result
def test_long_assistant_content_truncated(self) -> None:
"""Test that long assistant content is truncated."""
messages = [
{"role": "assistant", "content": "X" * 200},
]
result = extractive_summarizer(messages)
assert "..." in result # Truncation indicator
def test_context_ignored(self) -> None:
"""Test that context parameter exists but doesn't change output format."""
messages = [{"role": "user", "content": "Hello"}]
result1 = extractive_summarizer(messages, context="")
result2 = extractive_summarizer(messages, context="Some context here")
# Both should work (context is unused in extractive mode)
assert "[Summary of 1 messages]" in result1
assert "[Summary of 1 messages]" in result2
# =============================================================================
# ProgressiveSummarizer Core Tests
# =============================================================================
class TestProgressiveSummarizerInit:
"""Tests for ProgressiveSummarizer initialization."""
def test_default_init(self) -> None:
"""Test default initialization."""
summarizer = ProgressiveSummarizer()
assert summarizer.max_summary_tokens == 500
assert summarizer.min_messages_to_summarize == 3
assert summarizer.store_for_retrieval is True
# Default summarizer is extractive_summarizer
assert summarizer.summarize_fn is not None
def test_custom_summarize_fn(self) -> None:
"""Test custom summarization function."""
def custom_fn(messages: list[dict], context: str = "") -> str:
return f"Custom: {len(messages)} messages"
summarizer = ProgressiveSummarizer(summarize_fn=custom_fn)
result = summarizer.summarize_fn([{"role": "user", "content": "test"}])
assert "Custom: 1" in result
def test_custom_config(self) -> None:
"""Test custom configuration."""
summarizer = ProgressiveSummarizer(
max_summary_tokens=1000,
min_messages_to_summarize=5,
store_for_retrieval=False,
)
assert summarizer.max_summary_tokens == 1000
assert summarizer.min_messages_to_summarize == 5
assert summarizer.store_for_retrieval is False
# =============================================================================
# Find Candidates Tests
# =============================================================================
class TestFindSummarizationCandidates:
"""Tests for finding candidate message groups."""
def test_no_protected_all_candidates(self) -> None:
"""All messages are candidates when none protected."""
summarizer = ProgressiveSummarizer(min_messages_to_summarize=3)
messages = [
{"role": "user", "content": "1"},
{"role": "assistant", "content": "2"},
{"role": "user", "content": "3"},
{"role": "assistant", "content": "4"},
{"role": "user", "content": "5"},
]
groups = summarizer._find_summarization_candidates(messages, protected=set())
# Should have one group spanning all messages
assert len(groups) == 1
assert groups[0] == (0, 4)
def test_protected_splits_groups(self) -> None:
"""Protected messages split the candidates into groups."""
summarizer = ProgressiveSummarizer(min_messages_to_summarize=2)
messages = [
{"role": "user", "content": "1"},
{"role": "assistant", "content": "2"},
{"role": "user", "content": "3"}, # Protected at index 2
{"role": "assistant", "content": "4"},
{"role": "user", "content": "5"},
{"role": "assistant", "content": "6"},
]
groups = summarizer._find_summarization_candidates(messages, protected={2})
# Should have two groups: (0,1) and (3,5)
assert len(groups) == 2
assert groups[0] == (0, 1)
assert groups[1] == (3, 5)
def test_min_messages_filter(self) -> None:
"""Groups smaller than min_messages_to_summarize are filtered."""
summarizer = ProgressiveSummarizer(min_messages_to_summarize=3)
messages = [
{"role": "user", "content": "1"},
{"role": "assistant", "content": "2"},
{"role": "user", "content": "3"}, # Protected
{"role": "assistant", "content": "4"},
]
groups = summarizer._find_summarization_candidates(messages, protected={2})
# Group (0,1) has 2 messages, filtered. Group (3,3) has 1, filtered.
assert len(groups) == 0
def test_all_protected_no_candidates(self) -> None:
"""No candidates when all messages are protected."""
summarizer = ProgressiveSummarizer(min_messages_to_summarize=1)
messages = [
{"role": "user", "content": "1"},
{"role": "assistant", "content": "2"},
]
groups = summarizer._find_summarization_candidates(messages, protected={0, 1})
assert len(groups) == 0
def test_empty_messages(self) -> None:
"""Empty message list returns no groups."""
summarizer = ProgressiveSummarizer()
groups = summarizer._find_summarization_candidates([], protected=set())
assert len(groups) == 0
# =============================================================================
# Summarize Messages Tests
# =============================================================================
class TestSummarizeMessages:
"""Tests for the main summarize_messages method."""
def test_no_candidates_returns_original(
self, tokenizer: Tokenizer, simple_conversation: list[dict[str, Any]]
) -> None:
"""When no candidates, return original messages unchanged."""
summarizer = ProgressiveSummarizer(min_messages_to_summarize=100) # Too high
result = summarizer.summarize_messages(
simple_conversation, tokenizer, protected_indices=set()
)
assert len(result.messages) == len(simple_conversation)
assert result.tokens_saved == 0
assert len(result.summaries_created) == 0
def test_all_protected_no_changes(
self, tokenizer: Tokenizer, simple_conversation: list[dict[str, Any]]
) -> None:
"""All protected messages means no summarization."""
summarizer = ProgressiveSummarizer(min_messages_to_summarize=2)
all_protected = set(range(len(simple_conversation)))
result = summarizer.summarize_messages(
simple_conversation, tokenizer, protected_indices=all_protected
)
assert len(result.messages) == len(simple_conversation)
assert result.tokens_saved == 0
def test_summarization_reduces_messages(
self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
) -> None:
"""Summarization reduces message count."""
summarizer = ProgressiveSummarizer(
min_messages_to_summarize=3,
store_for_retrieval=False, # Skip CCR for test
)
# Protect first and last few messages
protected = {0, 1, len(long_conversation) - 1, len(long_conversation) - 2}
result = summarizer.summarize_messages(
long_conversation, tokenizer, protected_indices=protected
)
# Should have fewer messages
assert len(result.messages) < len(long_conversation)
# Should save tokens
assert result.tokens_saved > 0
# Should create summaries
assert len(result.summaries_created) > 0
def test_summarization_result_structure(
self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
) -> None:
"""Verify SummarizationResult has correct structure."""
summarizer = ProgressiveSummarizer(
min_messages_to_summarize=3,
store_for_retrieval=False,
)
result = summarizer.summarize_messages(long_conversation, tokenizer, protected_indices={0})
assert isinstance(result, SummarizationResult)
assert isinstance(result.messages, list)
assert isinstance(result.summaries_created, list)
assert isinstance(result.tokens_before, int)
assert isinstance(result.tokens_after, int)
assert isinstance(result.transforms_applied, list)
assert result.tokens_before >= result.tokens_after
def test_custom_summarizer_called(
self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
) -> None:
"""Custom summarizer function is called."""
calls: list[int] = []
def tracking_summarizer(messages: list[dict], context: str = "") -> str:
calls.append(len(messages))
return f"CUSTOM SUMMARY of {len(messages)} messages"
summarizer = ProgressiveSummarizer(
summarize_fn=tracking_summarizer,
min_messages_to_summarize=3,
store_for_retrieval=False,
)
result = summarizer.summarize_messages(long_conversation, tokenizer, protected_indices={0})
# Custom summarizer should have been called
assert len(calls) > 0
# Summary should appear in messages
found_custom = any("CUSTOM SUMMARY" in msg.get("content", "") for msg in result.messages)
assert found_custom
def test_context_passed_to_summarizer(
self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
) -> None:
"""Context messages are passed to summarizer."""
received_context: list[str] = []
def context_tracking_summarizer(messages: list[dict], context: str = "") -> str:
received_context.append(context)
return "Summary"
summarizer = ProgressiveSummarizer(
summarize_fn=context_tracking_summarizer,
min_messages_to_summarize=3,
store_for_retrieval=False,
)
context_msgs = [{"role": "user", "content": "Recent important question"}]
summarizer.summarize_messages(
long_conversation,
tokenizer,
protected_indices={0},
context_messages=context_msgs,
)
# Context should have been passed
assert len(received_context) > 0
# Should contain the recent message content
assert any("Recent important question" in ctx for ctx in received_context)
def test_target_tokens_stops_early(
self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
) -> None:
"""Summarization stops when target tokens reached."""
summarizer = ProgressiveSummarizer(
min_messages_to_summarize=3,
store_for_retrieval=False,
)
# Get original token count
original_tokens = tokenizer.count_messages(long_conversation)
# Set target very close to original (minimal summarization needed)
target = int(original_tokens * 0.95) # Only need 5% reduction
result = summarizer.summarize_messages(
long_conversation,
tokenizer,
protected_indices={0},
target_tokens=target,
)
# Should stop once target reached
assert result.tokens_after <= target or result.tokens_after < original_tokens
def test_small_groups_skipped(
self,
tokenizer: Tokenizer,
) -> None:
"""Groups with < 100 tokens are skipped."""
# Very short messages
messages = [
{"role": "user", "content": "Hi"},
{"role": "assistant", "content": "Hi"},
{"role": "user", "content": "Bye"},
{"role": "assistant", "content": "Bye"},
]
summarizer = ProgressiveSummarizer(
min_messages_to_summarize=2,
store_for_retrieval=False,
)
result = summarizer.summarize_messages(messages, tokenizer, protected_indices=set())
# Small groups should be skipped
assert len(result.summaries_created) == 0
def test_summary_larger_than_original_skipped(
self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
) -> None:
"""Summaries larger than original are skipped."""
def verbose_summarizer(messages: list[dict], context: str = "") -> str:
# Return a very verbose summary
return "VERY LONG SUMMARY " * 1000
summarizer = ProgressiveSummarizer(
summarize_fn=verbose_summarizer,
min_messages_to_summarize=3,
store_for_retrieval=False,
)
result = summarizer.summarize_messages(long_conversation, tokenizer, protected_indices={0})
# Summaries larger than original should be skipped
# (or if any were created, they saved tokens)
for summary in result.summaries_created:
assert summary.tokens_saved >= 0
def test_summarizer_exception_handled(
self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
) -> None:
"""Exceptions from summarizer are handled gracefully."""
def failing_summarizer(messages: list[dict], context: str = "") -> str:
raise ValueError("Summarization failed!")
summarizer = ProgressiveSummarizer(
summarize_fn=failing_summarizer,
min_messages_to_summarize=3,
store_for_retrieval=False,
)
# Should not raise, should return original
result = summarizer.summarize_messages(long_conversation, tokenizer, protected_indices={0})
# No summaries created due to failures
assert len(result.summaries_created) == 0
# =============================================================================
# Integration Tests
# =============================================================================
class TestProgressiveSummarizerIntegration:
"""Integration tests for end-to-end summarization."""
def test_full_workflow_with_extractive(
self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
) -> None:
"""Test full workflow with default extractive summarizer."""
summarizer = ProgressiveSummarizer(
min_messages_to_summarize=4,
store_for_retrieval=False,
)
original_count = len(long_conversation)
result = summarizer.summarize_messages(
long_conversation,
tokenizer,
protected_indices={0}, # Only protect system message
)
# Verify reduction
assert len(result.messages) < original_count
assert result.tokens_after < result.tokens_before
# Verify transforms tracked
assert len(result.transforms_applied) > 0
# Verify summaries created
assert len(result.summaries_created) > 0
for summary in result.summaries_created:
assert summary.start_index >= 0
assert summary.end_index >= summary.start_index
assert summary.compression_ratio < 1.0 # Actually compressed
def test_preserves_protected_messages(
self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
) -> None:
"""Protected messages are preserved exactly."""
summarizer = ProgressiveSummarizer(
min_messages_to_summarize=3,
store_for_retrieval=False,
)
# Protect first 3 and last 3 messages
protected = {
0,
1,
2,
len(long_conversation) - 3,
len(long_conversation) - 2,
len(long_conversation) - 1,
}
# Store original protected content
original_protected = {i: long_conversation[i]["content"] for i in protected}
result = summarizer.summarize_messages(
long_conversation,
tokenizer,
protected_indices=protected,
)
# Find protected messages in result
# First 3 should still be at beginning
assert result.messages[0]["content"] == original_protected[0]
assert result.messages[1]["content"] == original_protected[1]
assert result.messages[2]["content"] == original_protected[2]
# Last 3 should still be at end (positions shifted)
assert result.messages[-1]["content"] == original_protected[len(long_conversation) - 1]
assert result.messages[-2]["content"] == original_protected[len(long_conversation) - 2]
assert result.messages[-3]["content"] == original_protected[len(long_conversation) - 3]
def test_tool_messages_handled(
self, tokenizer: Tokenizer, conversation_with_tools: list[dict[str, Any]]
) -> None:
"""Tool messages are handled in summarization."""
# Create longer tool-heavy conversation
long_tool_conv = conversation_with_tools.copy()
for i in range(10):
long_tool_conv.extend(
[
{"role": "user", "content": f"Search again {i}"},
{
"role": "assistant",
"content": f"Searching {i}...",
"tool_calls": [
{
"id": f"call_{i}",
"type": "function",
"function": {"name": "search", "arguments": "{}"},
}
],
},
{
"role": "tool",
"tool_call_id": f"call_{i}",
"content": f'{{"data": "result {i}"}}',
},
{"role": "assistant", "content": f"Found result {i}"},
]
)
summarizer = ProgressiveSummarizer(
min_messages_to_summarize=3,
store_for_retrieval=False,
)
result = summarizer.summarize_messages(
long_tool_conv,
tokenizer,
protected_indices={0},
)
# Should reduce messages
assert len(result.messages) < len(long_tool_conv)
# Tool names should be tracked in summaries
all_tool_names = []
for summary in result.summaries_created:
all_tool_names.extend(summary.tool_names)
# Some tool calls should be tracked (may be empty if extractive)
def test_does_not_mutate_original(
self, tokenizer: Tokenizer, long_conversation: list[dict[str, Any]]
) -> None:
"""Original messages are not mutated."""
import copy
original_copy = copy.deepcopy(long_conversation)
summarizer = ProgressiveSummarizer(
min_messages_to_summarize=3,
store_for_retrieval=False,
)
summarizer.summarize_messages(
long_conversation,
tokenizer,
protected_indices={0},
)
# Original should be unchanged
assert long_conversation == original_copy
# =============================================================================
# SummarizationResult Tests
# =============================================================================
class TestSummarizationResult:
"""Tests for SummarizationResult dataclass."""
def test_tokens_saved_property(self) -> None:
"""Test tokens_saved property."""
result = SummarizationResult(
messages=[],
summaries_created=[],
tokens_before=1000,
tokens_after=300,
transforms_applied=[],
)
assert result.tokens_saved == 700
def test_tokens_saved_no_negative(self) -> None:
"""Test tokens_saved doesn't go negative."""
result = SummarizationResult(
messages=[],
summaries_created=[],
tokens_before=100,
tokens_after=150,
transforms_applied=[],
)
assert result.tokens_saved == 0

1218
tests/test_transforms/test_rolling_window.py
View file

File diff suppressed because it is too large Load diff

576
tests/test_transforms/test_scoring.py
View file

@ -1,576 +0,0 @@
"""Comprehensive tests for message importance scoring.
These tests verify that the scoring system works correctly WITHOUT
hardcoded patterns. All importance detection must come from:
1. Computed metrics (recency, density, references)
2. TOIN-learned patterns
3. Embedding similarity
"""
from __future__ import annotations
from typing import Any
import pytest
from headroom.config import ScoringWeights
from headroom.transforms.scoring import MessageScorer
# =============================================================================
# Test Fixtures
# =============================================================================
@pytest.fixture
def default_weights() -> ScoringWeights:
"""Default scoring weights."""
return ScoringWeights()
@pytest.fixture
def simple_conversation() -> list[dict[str, Any]]:
"""Simple conversation without tool calls."""
return [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello, how are you?"},
{"role": "assistant", "content": "I'm doing well, thank you for asking!"},
{"role": "user", "content": "Can you help me with Python?"},
{"role": "assistant", "content": "Of course! What would you like to know about Python?"},
]
@pytest.fixture
def conversation_with_tools() -> list[dict[str, Any]]:
"""Conversation with tool calls and responses."""
return [
{"role": "system", "content": "You are a helpful assistant with tools."},
{"role": "user", "content": "Search for information about Python."},
{
"role": "assistant",
"content": "I'll search for that.",
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {"name": "search", "arguments": "{}"},
}
],
},
{
"role": "tool",
"tool_call_id": "call_1",
"content": '{"results": [{"title": "Python Guide", "url": "example.com"}]}',
},
{"role": "assistant", "content": "Here's what I found about Python."},
{"role": "user", "content": "Thanks! Can you search for more?"},
{
"role": "assistant",
"content": "Sure, searching again.",
"tool_calls": [
{
"id": "call_2",
"type": "function",
"function": {"name": "search", "arguments": "{}"},
}
],
},
{
"role": "tool",
"tool_call_id": "call_2",
"content": '{"results": [{"title": "Advanced Python", "status": "found"}]}',
},
{"role": "assistant", "content": "Here are more results."},
]
@pytest.fixture
def long_conversation() -> list[dict[str, Any]]:
"""Long conversation for testing recency decay."""
messages = [{"role": "system", "content": "You are a helpful assistant."}]
for i in range(20):
messages.append({"role": "user", "content": f"User message {i}"})
messages.append({"role": "assistant", "content": f"Assistant response {i}"})
return messages
@pytest.fixture
def high_density_message() -> dict[str, Any]:
"""Message with high information density (many unique tokens)."""
return {
"role": "user",
"content": "Python JavaScript Ruby Golang Rust Swift Kotlin TypeScript C++ Java",
}
@pytest.fixture
def low_density_message() -> dict[str, Any]:
"""Message with low information density (repeated tokens)."""
return {
"role": "user",
"content": "the the the the the the the the the the very very very very",
}
# =============================================================================
# Test ScoringWeights
# =============================================================================
class TestScoringWeights:
"""Tests for ScoringWeights configuration."""
def test_default_weights_sum_approximately_one(self):
"""Default weights should sum close to 1.0."""
weights = ScoringWeights()
total = (
weights.recency
+ weights.semantic_similarity
+ weights.toin_importance
+ weights.error_indicator
+ weights.forward_reference
+ weights.token_density
)
assert abs(total - 1.0) < 0.01
def test_normalized_weights_sum_exactly_one(self):
"""Normalized weights should sum to exactly 1.0."""
weights = ScoringWeights(
recency=0.5,
semantic_similarity=0.3,
toin_importance=0.2,
error_indicator=0.1,
forward_reference=0.1,
token_density=0.05,
)
normalized = weights.normalized()
total = (
normalized.recency
+ normalized.semantic_similarity
+ normalized.toin_importance
+ normalized.error_indicator
+ normalized.forward_reference
+ normalized.token_density
)
assert abs(total - 1.0) < 1e-10
def test_zero_weights_returns_default(self):
"""Zero weights should return default weights."""
weights = ScoringWeights(
recency=0,
semantic_similarity=0,
toin_importance=0,
error_indicator=0,
forward_reference=0,
token_density=0,
)
normalized = weights.normalized()
# Should return default values, not NaN
assert normalized.recency >= 0
# =============================================================================
# Test MessageScorer Initialization
# =============================================================================
class TestMessageScorerInit:
"""Tests for MessageScorer initialization."""
def test_init_with_defaults(self):
"""Scorer initializes with default weights."""
scorer = MessageScorer()
assert scorer.weights is not None
assert scorer.toin is None
assert scorer.embedding_provider is None
def test_init_with_custom_weights(self, default_weights):
"""Scorer accepts custom weights."""
scorer = MessageScorer(weights=default_weights)
# Weights should be normalized
total = (
scorer.weights.recency
+ scorer.weights.semantic_similarity
+ scorer.weights.toin_importance
+ scorer.weights.error_indicator
+ scorer.weights.forward_reference
+ scorer.weights.token_density
)
assert abs(total - 1.0) < 1e-10
def test_init_with_custom_decay_rate(self):
"""Scorer accepts custom recency decay rate."""
scorer = MessageScorer(recency_decay_rate=0.2)
assert scorer.recency_decay_rate == 0.2
# =============================================================================
# Test Recency Scoring
# =============================================================================
class TestRecencyScoring:
"""Tests for recency-based scoring."""
def test_last_message_has_highest_recency(self, simple_conversation):
"""Most recent message should have highest recency score."""
scorer = MessageScorer()
scores = scorer.score_messages(
simple_conversation,
protected_indices=set(),
tool_unit_indices=set(),
)
# Last message should have highest recency
last_idx = len(simple_conversation) - 1
for i, score in enumerate(scores):
if i != last_idx:
assert score.recency_score <= scores[last_idx].recency_score
def test_recency_decreases_with_age(self, long_conversation):
"""Recency score should decrease for older messages."""
scorer = MessageScorer()
scores = scorer.score_messages(
long_conversation,
protected_indices=set(),
tool_unit_indices=set(),
)
# Verify decreasing recency (allowing for equal scores)
for i in range(1, len(scores)):
assert scores[i].recency_score >= scores[i - 1].recency_score
def test_recency_decay_rate_affects_scores(self, simple_conversation):
"""Higher decay rate should make older messages score lower."""
scorer_slow = MessageScorer(recency_decay_rate=0.05)
scorer_fast = MessageScorer(recency_decay_rate=0.5)
scores_slow = scorer_slow.score_messages(simple_conversation, set(), set())
scores_fast = scorer_fast.score_messages(simple_conversation, set(), set())
# First message should have lower recency with fast decay
assert scores_fast[1].recency_score < scores_slow[1].recency_score
def test_single_message_has_max_recency(self):
"""Single message conversation should have max recency."""
messages = [{"role": "user", "content": "Hello"}]
scorer = MessageScorer()
scores = scorer.score_messages(messages, set(), set())
assert scores[0].recency_score == 1.0
# =============================================================================
# Test Density Scoring
# =============================================================================
class TestDensityScoring:
"""Tests for information density scoring."""
def test_high_density_scores_higher(self, high_density_message, low_density_message):
"""High density message should score higher than low density."""
scorer = MessageScorer()
high_scores = scorer.score_messages([high_density_message], set(), set())
low_scores = scorer.score_messages([low_density_message], set(), set())
assert high_scores[0].density_score > low_scores[0].density_score
def test_density_in_valid_range(self, simple_conversation):
"""Density scores should be in [0, 1] range."""
scorer = MessageScorer()
scores = scorer.score_messages(simple_conversation, set(), set())
for score in scores:
assert 0.0 <= score.density_score <= 1.0
def test_empty_content_gets_neutral_density(self):
"""Empty content should get neutral density score."""
messages = [{"role": "user", "content": ""}]
scorer = MessageScorer()
scores = scorer.score_messages(messages, set(), set())
assert scores[0].density_score == 0.5
# =============================================================================
# Test Forward Reference Scoring
# =============================================================================
class TestForwardReferenceScoring:
"""Tests for forward reference detection."""
def test_assistant_with_tool_calls_has_references(self, conversation_with_tools):
"""Assistant messages with tool calls should have forward references."""
scorer = MessageScorer()
scores = scorer.score_messages(conversation_with_tools, set(), set())
# Message at index 2 has tool_calls, should have references
# (tool response at index 3 references it)
assert scores[2].reference_score > 0
def test_no_tool_calls_no_references(self, simple_conversation):
"""Messages without tool calls should have no forward references."""
scorer = MessageScorer()
scores = scorer.score_messages(simple_conversation, set(), set())
for score in scores:
assert score.reference_score == 0.0
# =============================================================================
# Test Protected Message Handling
# =============================================================================
class TestProtectedMessages:
"""Tests for protected message handling in scoring."""
def test_protected_messages_marked(self, simple_conversation):
"""Protected messages should be marked in scores."""
scorer = MessageScorer()
protected = {0, 1} # System and first user message
scores = scorer.score_messages(
simple_conversation,
protected_indices=protected,
tool_unit_indices=set(),
)
assert scores[0].is_protected is True
assert scores[1].is_protected is True
assert scores[2].is_protected is False
def test_tool_unit_messages_marked_unsafe(self, conversation_with_tools):
"""Messages in tool units should be marked as not drop_safe."""
scorer = MessageScorer()
tool_unit_indices = {2, 3, 6, 7} # Assistant with tool_calls and responses
scores = scorer.score_messages(
conversation_with_tools,
protected_indices=set(),
tool_unit_indices=tool_unit_indices,
)
# Tool unit messages should not be independently droppable
# They can only be dropped as a unit
for idx in tool_unit_indices:
# Not protected, but part of a unit
assert scores[idx].drop_safe is True or scores[idx].is_protected
# =============================================================================
# Test Total Score Computation
# =============================================================================
class TestTotalScore:
"""Tests for total score computation."""
def test_total_score_in_valid_range(self, simple_conversation):
"""Total scores should be in [0, 1] range."""
scorer = MessageScorer()
scores = scorer.score_messages(simple_conversation, set(), set())
for score in scores:
assert 0.0 <= score.total_score <= 1.0
def test_score_breakdown_matches_total(self, simple_conversation):
"""Score breakdown components should match total."""
weights = ScoringWeights()
scorer = MessageScorer(weights=weights)
scores = scorer.score_messages(simple_conversation, set(), set())
for score in scores:
# Without TOIN and embeddings, some components are neutral
# Just verify breakdown dict is populated
assert "recency" in score.score_breakdown
assert "density" in score.score_breakdown
def test_weights_affect_total_score(self, simple_conversation):
"""Different weights should produce different total scores."""
weights_recency = ScoringWeights(
recency=1.0,
semantic_similarity=0,
toin_importance=0,
error_indicator=0,
forward_reference=0,
token_density=0,
)
weights_density = ScoringWeights(
recency=0,
semantic_similarity=0,
toin_importance=0,
error_indicator=0,
forward_reference=0,
token_density=1.0,
)
scorer_recency = MessageScorer(weights=weights_recency)
scorer_density = MessageScorer(weights=weights_density)
scores_r = scorer_recency.score_messages(simple_conversation, set(), set())
scores_d = scorer_density.score_messages(simple_conversation, set(), set())
# Different weights should produce different scores
# (unless message happens to have same recency and density)
assert any(
abs(scores_r[i].total_score - scores_d[i].total_score) > 0.01
for i in range(len(simple_conversation))
)
# =============================================================================
# Test Cosine Similarity
# =============================================================================
class TestCosineSimilarity:
"""Tests for cosine similarity computation."""
def test_identical_vectors_similarity_one(self):
"""Identical vectors should have similarity 1.0."""
a = [1.0, 2.0, 3.0]
similarity = MessageScorer._cosine_similarity(a, a)
assert abs(similarity - 1.0) < 1e-10
def test_orthogonal_vectors_similarity_zero(self):
"""Orthogonal vectors should have similarity 0.0."""
a = [1.0, 0.0]
b = [0.0, 1.0]
similarity = MessageScorer._cosine_similarity(a, b)
assert abs(similarity) < 1e-10
def test_opposite_vectors_similarity_negative(self):
"""Opposite vectors should have similarity -1.0."""
a = [1.0, 1.0]
b = [-1.0, -1.0]
similarity = MessageScorer._cosine_similarity(a, b)
assert abs(similarity - (-1.0)) < 1e-10
def test_empty_vectors_return_zero(self):
"""Empty vectors should return 0 similarity."""
assert MessageScorer._cosine_similarity([], []) == 0.0
def test_mismatched_lengths_return_zero(self):
"""Mismatched vector lengths should return 0."""
a = [1.0, 2.0]
b = [1.0, 2.0, 3.0]
assert MessageScorer._cosine_similarity(a, b) == 0.0
# =============================================================================
# Test TOIN Integration (Without Actual TOIN)
# =============================================================================
class TestTOINIntegration:
"""Tests for TOIN integration in scoring."""
def test_no_toin_returns_neutral_scores(self, conversation_with_tools):
"""Without TOIN, tool messages should get neutral TOIN scores."""
scorer = MessageScorer(toin=None)
scores = scorer.score_messages(conversation_with_tools, set(), set())
# Tool messages should have neutral TOIN score (0.5)
for i, msg in enumerate(conversation_with_tools):
if msg.get("role") == "tool":
assert scores[i].toin_score == 0.5
def test_no_toin_returns_zero_error_scores(self, conversation_with_tools):
"""Without TOIN, error scores should be zero."""
scorer = MessageScorer(toin=None)
scores = scorer.score_messages(conversation_with_tools, set(), set())
for score in scores:
assert score.error_score == 0.0
# =============================================================================
# Test Edge Cases
# =============================================================================
class TestEdgeCases:
"""Tests for edge cases in scoring."""
def test_empty_messages_list(self):
"""Empty message list should return empty scores."""
scorer = MessageScorer()
scores = scorer.score_messages([], set(), set())
assert scores == []
def test_single_message(self):
"""Single message should be scored correctly."""
messages = [{"role": "user", "content": "Hello"}]
scorer = MessageScorer()
scores = scorer.score_messages(messages, set(), set())
assert len(scores) == 1
assert scores[0].message_index == 0
assert scores[0].recency_score == 1.0
def test_message_with_no_content(self):
"""Message with no content key should be handled."""
messages = [{"role": "user"}]
scorer = MessageScorer()
scores = scorer.score_messages(messages, set(), set())
assert len(scores) == 1
# Should not crash, should have some score
def test_message_with_non_string_content(self):
"""Message with non-string content should be handled."""
messages = [{"role": "user", "content": ["list", "content"]}]
scorer = MessageScorer()
scores = scorer.score_messages(messages, set(), set())
assert len(scores) == 1
# Density should be neutral for non-string
assert scores[0].density_score == 0.5
def test_all_messages_protected(self, simple_conversation):
"""All messages protected should all be marked."""
scorer = MessageScorer()
all_protected = set(range(len(simple_conversation)))
scores = scorer.score_messages(
simple_conversation,
protected_indices=all_protected,
tool_unit_indices=set(),
)
for score in scores:
assert score.is_protected is True
# =============================================================================
# Test Score Ordering
# =============================================================================
class TestScoreOrdering:
"""Tests for expected score ordering patterns."""
def test_recent_messages_score_higher_by_default(self, long_conversation):
"""Recent messages should generally score higher with default weights."""
scorer = MessageScorer()
scores = scorer.score_messages(long_conversation, set(), set())
# Average score of last 5 messages should be higher than first 5
# (excluding system message at index 0)
first_5_avg = sum(s.total_score for s in scores[1:6]) / 5
last_5_avg = sum(s.total_score for s in scores[-5:]) / 5
assert last_5_avg > first_5_avg
def test_system_message_scores_lower_on_recency(self, simple_conversation):
"""System message (index 0) should have low recency score."""
scorer = MessageScorer()
scores = scorer.score_messages(simple_conversation, set(), set())
# System message is oldest, should have lowest recency
system_recency = scores[0].recency_score
for score in scores[1:]:
assert score.recency_score >= system_recency

154
tests/test_transforms/test_tool_crusher.py
View file

@ -1,154 +0,0 @@
"""Tests for tool crusher transform."""
import json
from headroom import OpenAIProvider, Tokenizer, ToolCrusherConfig
from headroom.transforms import ToolCrusher
# Create a shared provider for tests
_provider = OpenAIProvider()
def get_tokenizer(model: str = "gpt-4o") -> Tokenizer:
"""Get a tokenizer for tests using OpenAI provider."""
token_counter = _provider.get_token_counter(model)
return Tokenizer(token_counter, model)
class TestToolCrusher:
"""Tests for ToolCrusher transform."""
def test_small_tool_output_unchanged(self):
"""Small tool outputs should not be modified."""
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "tool", "tool_call_id": "call_1", "content": '{"status": "ok"}'},
]
crusher = ToolCrusher()
tokenizer = get_tokenizer()
result = crusher.apply(messages, tokenizer)
# Should not be modified (too small)
assert result.messages[1]["content"] == '{"status": "ok"}'
assert len(result.transforms_applied) == 0
def test_large_json_array_truncated(self):
"""Large arrays should be truncated."""
large_array = [{"id": i, "name": f"Item {i}"} for i in range(50)]
large_json = json.dumps({"results": large_array})
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "tool", "tool_call_id": "call_1", "content": large_json},
]
config = ToolCrusherConfig(min_tokens_to_crush=50, max_array_items=5)
crusher = ToolCrusher(config)
tokenizer = get_tokenizer()
result = crusher.apply(messages, tokenizer)
# Should be modified
tool_content = result.messages[1]["content"]
parsed = json.loads(tool_content.split("\n<headroom:")[0])
# Array should be truncated
assert len(parsed["results"]) <= 6 # 5 items + truncation marker
def test_long_strings_truncated(self):
"""Long strings should be truncated."""
long_string = "x" * 2000
data = {"content": long_string}
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "tool", "tool_call_id": "call_1", "content": json.dumps(data)},
]
config = ToolCrusherConfig(min_tokens_to_crush=50, max_string_length=100)
crusher = ToolCrusher(config)
tokenizer = get_tokenizer()
result = crusher.apply(messages, tokenizer)
tool_content = result.messages[1]["content"]
parsed = json.loads(tool_content.split("\n<headroom:")[0])
# String should be truncated
assert len(parsed["content"]) < 200
assert "truncated" in parsed["content"]
def test_nested_depth_limited(self):
"""Deeply nested structures should be limited."""
# Create deeply nested structure
nested = {"level": 0}
current = nested
for i in range(10):
current["nested"] = {"level": i + 1}
current = current["nested"]
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "tool", "tool_call_id": "call_1", "content": json.dumps(nested)},
]
config = ToolCrusherConfig(min_tokens_to_crush=10, max_depth=3)
crusher = ToolCrusher(config)
tokenizer = get_tokenizer()
result = crusher.apply(messages, tokenizer)
tool_content = result.messages[1]["content"]
parsed = json.loads(tool_content.split("\n<headroom:")[0])
# Deep nesting should be summarized
# Navigate to depth limit
current = parsed
depth = 0
while "nested" in current and isinstance(current["nested"], dict):
current = current["nested"]
depth += 1
if depth > 5:
break
assert depth <= 4 # Should be limited
def test_digest_marker_added(self):
"""Digest marker should be added to crushed content."""
large_data = {"items": list(range(100))}
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "tool", "tool_call_id": "call_1", "content": json.dumps(large_data)},
]
config = ToolCrusherConfig(min_tokens_to_crush=10, max_array_items=5)
crusher = ToolCrusher(config)
tokenizer = get_tokenizer()
result = crusher.apply(messages, tokenizer)
tool_content = result.messages[1]["content"]
# Should have digest marker
assert "<headroom:tool_digest" in tool_content
assert "sha256=" in tool_content
def test_non_tool_messages_unchanged(self):
"""Non-tool messages should not be modified."""
messages = [
{"role": "system", "content": json.dumps({"large": "data" * 1000})},
{"role": "user", "content": json.dumps({"user": "data" * 1000})},
{"role": "assistant", "content": json.dumps({"assistant": "data" * 1000})},
]
crusher = ToolCrusher()
tokenizer = get_tokenizer()
result = crusher.apply(messages, tokenizer)
# All messages should be unchanged
for i, msg in enumerate(result.messages):
assert msg["content"] == messages[i]["content"]

14
tests/test_ws_http_fallback.py
View file

@ -123,7 +123,12 @@ class TestWsHttpFallback:
assert "401" in event["error"]["message"]
def test_fallback_sets_stream_true(self):
"""HTTP fallback should force stream=True in request body."""
"""HTTP fallback should force stream=True in request body.
After PR-A3 (byte-faithful Python forwarders) the fallback sends
the request body as raw bytes via `content=`, not via the `json=`
kwarg. The test extracts the posted JSON from the captured bytes.
"""
handler = _make_handler()
ws = FakeWebSocket()
captured_kwargs: dict = {}
@ -138,7 +143,8 @@ class TestWsHttpFallback:
body = {"model": "gpt-5.4", "input": "test", "stream": False}
asyncio.run(handler._ws_http_fallback(ws, body, json.dumps(body), {}, "req_3"))
assert captured_kwargs["json"]["stream"] is True
posted = json.loads(captured_kwargs["content"])
assert posted["stream"] is True
def test_fallback_unwraps_response_create_envelope(self):
"""HTTP fallback should unwrap WS response.create wrapper for HTTP POST."""
@ -161,7 +167,7 @@ class TestWsHttpFallback:
ws_msg = {"type": "response.create", "response": inner}
asyncio.run(handler._ws_http_fallback(ws, ws_msg, json.dumps(ws_msg), {}, "req_unwrap"))
posted = captured_kwargs["json"]
posted = json.loads(captured_kwargs["content"])
# Should be the inner response, not the wrapper
assert "type" not in posted # no "response.create" type field
assert posted["model"] == "gpt-5.4"
@ -184,7 +190,7 @@ class TestWsHttpFallback:
body = {"type": "response.create", "model": "gpt-5.4", "input": "hi"}
asyncio.run(handler._ws_http_fallback(ws, body, json.dumps(body), {}, "req_type_strip"))
posted = captured_kwargs["json"]
posted = json.loads(captured_kwargs["content"])
assert posted["model"] == "gpt-5.4"
assert posted["stream"] is True
assert "type" not in posted

8292
uv.lock generated
View file

File diff suppressed because it is too large Load diff