diff --git a/Cargo.lock b/Cargo.lock index 5d0ddc5..f6c8d8e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2063,7 +2063,9 @@ dependencies = [ "anyhow", "futures", "hex", + "indexmap", "magika", + "once_cell", "ort", "regex", "serde", diff --git a/crates/mem-core/Cargo.toml b/crates/mem-core/Cargo.toml index 40ca688..e6ce638 100644 --- a/crates/mem-core/Cargo.toml +++ b/crates/mem-core/Cargo.toml @@ -19,3 +19,5 @@ time = { version = "0.3", features = ["serde", "formatting", "parsing", "macros" magika = "1.1.0" ort = { version = "2.0.0-rc.12", default-features = true } regex = "1.10" +once_cell = "1.19" +indexmap = "2.0" diff --git a/crates/mem-core/src/lib.rs b/crates/mem-core/src/lib.rs index 2346387..fe9b455 100644 --- a/crates/mem-core/src/lib.rs +++ b/crates/mem-core/src/lib.rs @@ -20,4 +20,4 @@ pub use lesson::{ pub use query::{Query, QuerySet, SynthesisQuery}; pub use prompt::{PromptBuilder, PromptMessages}; pub use symptom_projection::{project_symptom, SymptomVector}; -pub use optimizer::{ContextOptimizer, ContextOptimizerConfig, ContentType, OptimizedChunk}; +pub use optimizer::{ContextOptimizer, ContextOptimizerConfig, ContentType, OptimizedChunk, CacheAligner, AlignedContent, CcrStore}; diff --git a/crates/mem-core/src/optimizer/cache_align.rs b/crates/mem-core/src/optimizer/cache_align.rs new file mode 100644 index 0000000..0787322 --- /dev/null +++ b/crates/mem-core/src/optimizer/cache_align.rs @@ -0,0 +1,191 @@ +//! CacheAligner — Stabilize prompt prefix for LLM provider KV cache hits +//! +//! LLM providers (Anthropic, OpenAI) use prefix-based KV caching. A single +//! changing timestamp early in the prompt invalidates the entire cache. +//! +//! CacheAligner detects dynamic patterns and moves them to the context tail, +//! preserving the stable prefix for cache hits. + +use regex::Regex; +use once_cell::sync::Lazy; + +static ISO_TIMESTAMP: Lazy = Lazy::new(|| Regex::new(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}").unwrap()); +static UUID: Lazy = Lazy::new(|| Regex::new(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}").unwrap()); +static SESSION_ID: Lazy = Lazy::new(|| Regex::new(r"session[_-]?id[=:]\s*([a-zA-Z0-9]+)").unwrap()); +static RUN_ID: Lazy = Lazy::new(|| Regex::new(r"run[_-]?id[=:]\s*([a-zA-Z0-9_-]+)").unwrap()); +static TEMP_PATH: Lazy = Lazy::new(|| Regex::new(r"(/tmp|/var/tmp|C:\\Users\\[^\\]+\\AppData|~)/[^\s]+").unwrap()); +static SHA256: Lazy = Lazy::new(|| Regex::new(r"[a-f0-9]{64}").unwrap()); +static LINE_COL: Lazy = Lazy::new(|| Regex::new(r":\d{1,5}:\d{1,5}").unwrap()); + +pub struct CacheAligner; + +#[derive(Debug, Clone)] +pub struct AlignedContent { + /// Stable prefix (should be cached by LLM provider) + pub stable_prefix: String, + /// Dynamic tail (timestamps, UUIDs, session IDs, etc.) + pub dynamic_tail: String, + /// How much of the prefix changed (0.0 = identical, 1.0 = completely different) + pub drift_metric: f32, +} + +impl CacheAligner { + /// Stabilize prompt by moving dynamic content to tail + pub fn align(content: &str) -> AlignedContent { + let lines: Vec<&str> = content.lines().collect(); + let total_lines = lines.len(); + let mut static_lines = Vec::new(); + let mut dynamic_lines = Vec::new(); + + for line in &lines { + if Self::is_dynamic_line(line) { + dynamic_lines.push(*line); + } else { + static_lines.push(*line); + } + } + + let stable_prefix = static_lines.join("\n"); + let dynamic_tail = if dynamic_lines.is_empty() { + String::new() + } else { + format!("\n\n{}", dynamic_lines.join("\n")) + }; + + // Drift metric: ratio of dynamic lines + let drift_metric = if total_lines > 0 { + dynamic_lines.len() as f32 / total_lines as f32 + } else { + 0.0 + }; + + AlignedContent { + stable_prefix, + dynamic_tail, + drift_metric, + } + } + + /// Check if a line contains dynamic content + fn is_dynamic_line(line: &str) -> bool { + // Timestamps + if ISO_TIMESTAMP.is_match(line) { + return true; + } + + // UUIDs + if UUID.is_match(line) { + return true; + } + + // Session/run IDs + if SESSION_ID.is_match(line) || RUN_ID.is_match(line) { + return true; + } + + // Temp paths + if TEMP_PATH.is_match(line) { + return true; + } + + // SHA256 hashes (but not in common markers like "ccr:" prefix) + if SHA256.is_match(line) && !line.starts_with("", hash); + + assert!(hint.len() > 10); + assert!(hint.contains("CCR:")); + assert!(hint.contains(&hash)); + } + + #[test] + fn test_ccr_large_content() { + let store = CcrStore::new(); + let large = "x".repeat(100_000); + + let hash = store.store(&large).unwrap(); + let retrieved = store.retrieve(&hash).unwrap(); + + assert_eq!(retrieved.unwrap().len(), 100_000); + } + + #[test] + fn test_ccr_multiple_stores_same_content() { + let store = CcrStore::new(); + + let h1 = store.store("shared").unwrap(); + let h2 = store.store("shared").unwrap(); + + // Same content should produce same hash + assert_eq!(h1, h2); + + // But they should have separate cache entries (LRU) + // Last one should be retrievable + assert!(store.retrieve(&h2).unwrap().is_some()); + } +} diff --git a/crates/mem-core/src/optimizer/mod.rs b/crates/mem-core/src/optimizer/mod.rs index b7ce175..26cad17 100644 --- a/crates/mem-core/src/optimizer/mod.rs +++ b/crates/mem-core/src/optimizer/mod.rs @@ -8,6 +8,8 @@ pub mod router; pub mod log; pub mod json; pub mod diff; +pub mod cache_align; +pub mod ccr; use anyhow::Result; use serde::{Deserialize, Serialize}; @@ -16,6 +18,8 @@ pub use router::ContentRouter; pub use log::LogCompressor; pub use json::JsonCrusher; pub use diff::DiffCompressor; +pub use cache_align::{CacheAligner, AlignedContent}; +pub use ccr::CcrStore; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct OptimizedChunk { @@ -84,6 +88,7 @@ pub struct ContextOptimizer { log_compressor: LogCompressor, json_crusher: JsonCrusher, diff_compressor: DiffCompressor, + ccr_store: CcrStore, config: ContextOptimizerConfig, } @@ -99,16 +104,23 @@ impl ContextOptimizer { let log_compressor = LogCompressor::new(); let json_crusher = JsonCrusher::new(); let diff_compressor = DiffCompressor::new(); + let ccr_store = CcrStore::new(); Ok(Self { router, log_compressor, json_crusher, diff_compressor, + ccr_store, config, }) } + /// Get reference to CCR store for retrieval + pub fn ccr_store(&self) -> &CcrStore { + &self.ccr_store + } + /// Optimize a chunk of content pub fn optimize(&self, content: &str) -> Result { if !self.config.enabled { @@ -143,12 +155,19 @@ impl ContextOptimizer { let original_tokens = estimate_tokens(content); let compressed_tokens = estimate_tokens(&compressed); + // Store original in CCR if compression happened and CCR is enabled + let ccr_hash = if self.config.ccr_enabled && compressed != content { + self.ccr_store.store(content).ok() + } else { + None + }; + Ok(OptimizedChunk { compressed, original_tokens, compressed_tokens, content_type, - ccr_hash: None, // TODO: Implement CCR + ccr_hash, }) } } @@ -185,4 +204,26 @@ mod tests { let chunk = optimizer.optimize("test content").unwrap(); assert_eq!(chunk.compressed, "test content"); } + + #[test] + fn test_optimizer_with_ccr() { + let optimizer = ContextOptimizer::new().unwrap(); + let log_content = "ERROR: failed\nWARN: ignored"; + let chunk = optimizer.optimize(log_content).unwrap(); + + // If compression happened and CCR enabled, should have hash + if chunk.compressed != log_content { + assert!(chunk.ccr_hash.is_some()); + } + } + + #[test] + fn test_optimizer_cache_align() { + let content = "System prompt\nAt 2026-08-28T09:15:00Z query was:\nContext"; + let aligned = CacheAligner::align(content); + + assert!(aligned.stable_prefix.contains("System")); + assert!(aligned.stable_prefix.contains("Context")); + assert!(aligned.dynamic_tail.contains("2026-08-28")); + } }