From 013344f6fd31e968b466833b780eb503840696e8 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Fri, 1 May 2026 14:40:47 -0700 Subject: [PATCH] =?UTF-8?q?fix(rust):=20IntelligentContextManager=20port?= =?UTF-8?q?=20(simplified,=20OSS)=20=E2=80=94=20PR-B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Direct port of headroom/transforms/intelligent_context.py (1077 LOC) with a deliberately simpler architecture: - One built-in strategy: DropByScoreStrategy. Multi-factor scoring + safety rails + CCR-on-drop persistence. Outclasses every gateway competitor's rolling-window behaviour. - Minimal ContextStrategy trait so Enterprise plugs additional strategies (compress-first, summarize, memory tiers) into the same orchestrator without modifying it. - 6 config fields instead of 12+. Cuts: compress_threshold, summarize_threshold, summarization_*, memory_tiers_*, warm/cold_*. Module layout (parallel to scoring/, signals/, transforms/): context/ ├── config.rs — IcmConfig (6 fields) ├── workspace.rs — ContextWorkspace + StrategyOutcome ├── safety.rs — system / last-N-turns / tool-pair │ atomicity / frozen prefix ├── candidate.rs — turn / tool-unit / single candidates ├── ccr_drop.rs — serialize dropped messages, store │ under content-hash, emit marker ├── manager.rs — cascade orchestrator └── strategy/ ├── mod.rs — ContextStrategy trait └── drop_by_score.rs — THE OSS strategy Quirks preserved from Python: - Tool-call/response atomicity (OpenAI + Anthropic + Strands shapes) - frozen_message_count protects prompt-cache prefix - Cascade: each strategy returns tokens_freed + fully_resolved - should_apply strict gate (no tokenization on under-budget requests) - CCR-on-drop default ON (the OSS-defining behaviour) 44 new unit tests (50 total context tests), 811 total in headroom-core. NO PyO3 yet (PR-C). NO Python deletion yet (PR-D). --- crates/headroom-core/src/context/candidate.rs | 415 ++++++++++++++++++ crates/headroom-core/src/context/ccr_drop.rs | 201 +++++++++ crates/headroom-core/src/context/config.rs | 90 ++++ crates/headroom-core/src/context/manager.rs | 387 ++++++++++++++++ crates/headroom-core/src/context/mod.rs | 47 ++ crates/headroom-core/src/context/safety.rs | 357 +++++++++++++++ .../src/context/strategy/drop_by_score.rs | 410 +++++++++++++++++ .../headroom-core/src/context/strategy/mod.rs | 45 ++ crates/headroom-core/src/context/workspace.rs | 85 ++++ crates/headroom-core/src/lib.rs | 1 + 10 files changed, 2038 insertions(+) create mode 100644 crates/headroom-core/src/context/candidate.rs create mode 100644 crates/headroom-core/src/context/ccr_drop.rs create mode 100644 crates/headroom-core/src/context/config.rs create mode 100644 crates/headroom-core/src/context/manager.rs create mode 100644 crates/headroom-core/src/context/mod.rs create mode 100644 crates/headroom-core/src/context/safety.rs create mode 100644 crates/headroom-core/src/context/strategy/drop_by_score.rs create mode 100644 crates/headroom-core/src/context/strategy/mod.rs create mode 100644 crates/headroom-core/src/context/workspace.rs diff --git a/crates/headroom-core/src/context/candidate.rs b/crates/headroom-core/src/context/candidate.rs new file mode 100644 index 000000000..bf7f6fe28 --- /dev/null +++ b/crates/headroom-core/src/context/candidate.rs @@ -0,0 +1,415 @@ +//! 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, + /// 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); + +/// 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 { + use std::collections::HashMap; + + // tool_call_id → response message index + let mut response_map: HashMap = 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 = 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 = 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, + tool_units: &[ToolUnit], +) -> Vec { + // 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 = 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 = 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 = 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); + } +} diff --git a/crates/headroom-core/src/context/ccr_drop.rs b/crates/headroom-core/src/context/ccr_drop.rs new file mode 100644 index 000000000..f5078f082 --- /dev/null +++ b/crates/headroom-core/src/context/ccr_drop.rs @@ -0,0 +1,201 @@ +//! 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()`, 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, +) -> Option { + 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::(); + + // 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 = 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 { + 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 = 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); + } +} diff --git a/crates/headroom-core/src/context/config.rs b/crates/headroom-core/src/context/config.rs new file mode 100644 index 000000000..fed04484c --- /dev/null +++ b/crates/headroom-core/src/context/config.rs @@ -0,0 +1,90 @@ +//! `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); + } +} diff --git a/crates/headroom-core/src/context/manager.rs b/crates/headroom-core/src/context/manager.rs new file mode 100644 index 000000000..b425d142b --- /dev/null +++ b/crates/headroom-core/src/context/manager.rs @@ -0,0 +1,387 @@ +//! `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, + /// 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, + 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, +} + +/// The orchestrator. Constructed once per process and reused across +/// requests; `Send + Sync`. +pub struct IntelligentContextManager { + config: IcmConfig, + tokenizer: Arc, + strategies: Vec>, +} + +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, + ccr: Option>, + ) -> Self { + let scorer = Arc::new(MessageScorer::new( + Some(config.scoring_weights), + None, + None, + 0.1, + )); + let drop_strategy: Box = 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) -> 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, 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 = 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 = Arc::new(EstimatingCounter::default()); + let store: Arc = 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 = 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 = (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 = 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 = 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 = (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 = (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")); + } +} diff --git a/crates/headroom-core/src/context/mod.rs b/crates/headroom-core/src/context/mod.rs new file mode 100644 index 000000000..36a34101f --- /dev/null +++ b/crates/headroom-core/src/context/mod.rs @@ -0,0 +1,47 @@ +//! 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}; diff --git a/crates/headroom-core/src/context/safety.rs b/crates/headroom-core/src/context/safety.rs new file mode 100644 index 000000000..17c34a666 --- /dev/null +++ b/crates/headroom-core/src/context/safety.rs @@ -0,0 +1,357 @@ +//! `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 { + let mut protected: HashSet = 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, + ) -> HashSet { + let mut response_indices: HashSet = 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 { + let mut ids: HashSet = 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) -> 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()); + } +} diff --git a/crates/headroom-core/src/context/strategy/drop_by_score.rs b/crates/headroom-core/src/context/strategy/drop_by_score.rs new file mode 100644 index 000000000..bfe59cafd --- /dev/null +++ b/crates/headroom-core/src/context/strategy/drop_by_score.rs @@ -0,0 +1,410 @@ +//! `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, + tokenizer: Arc, + ccr_store: Option>, + ccr_on_drop: bool, +} + +impl DropByScoreStrategy { + pub fn new( + scorer: Arc, + tokenizer: Arc, + ccr_store: Option>, + 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 = 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 = Vec::new(); + let mut indices_to_remove: std::collections::BTreeSet = + 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 = 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 = 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, 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>, + ccr_on_drop: bool, + tokenizer: Arc, + ) -> 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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); + } +} diff --git a/crates/headroom-core/src/context/strategy/mod.rs b/crates/headroom-core/src/context/strategy/mod.rs new file mode 100644 index 000000000..d2c364abb --- /dev/null +++ b/crates/headroom-core/src/context/strategy/mod.rs @@ -0,0 +1,45 @@ +//! `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; diff --git a/crates/headroom-core/src/context/workspace.rs b/crates/headroom-core/src/context/workspace.rs new file mode 100644 index 000000000..fd2c9c30f --- /dev/null +++ b/crates/headroom-core/src/context/workspace.rs @@ -0,0 +1,85 @@ +//! 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, + + /// 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, + + /// 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, + + /// 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, +} + +impl ContextWorkspace { + pub fn new(messages: Vec, protected: HashSet, 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, + + /// 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, +} diff --git a/crates/headroom-core/src/lib.rs b/crates/headroom-core/src/lib.rs index a80534abe..9b6fb70e4 100644 --- a/crates/headroom-core/src/lib.rs +++ b/crates/headroom-core/src/lib.rs @@ -1,6 +1,7 @@ //! headroom-core: foundation crate for the Rust port of Headroom. pub mod ccr; +pub mod context; pub mod relevance; pub mod scoring; pub mod signals;