// memory-search: combined hook for session context loading + ambient memory retrieval // // Modes: // --hook Run as Claude Code UserPromptSubmit hook (reads stdin, injects into conversation) // --debug Replay last stashed input, dump every stage to stdout // --seen Show the seen set for current session // (default) No-op (future: manual search modes) use clap::Parser; use poc_memory::search::{self, AlgoStage}; use poc_memory::store; use std::collections::{BTreeMap, HashSet}; use std::fs; use std::io::{self, Read, Write}; use std::path::{Path, PathBuf}; use std::process::Command; use std::time::{Duration, SystemTime}; #[derive(Parser)] #[command(name = "memory-search")] struct Args { /// Run as Claude Code hook (reads stdin, outputs for injection) #[arg(long)] hook: bool, /// Debug mode: replay last stashed input, dump every stage #[arg(short, long)] debug: bool, /// Show session state: seen set, returned memories, compaction info #[arg(long)] seen: bool, /// Max results from search pipeline (filtered by seen set before injection) #[arg(long, default_value = "50")] max_results: usize, /// Search query (bypasses stashed input, uses this as the prompt) #[arg(long, short)] query: Option, /// Algorithm pipeline stages: e.g. spread spectral,k=20 spread,max_hops=4 /// Default: spread. pipeline: Vec, } const STASH_PATH: &str = "/tmp/claude-memory-search/last-input.json"; /// Max bytes per context chunk (hook output limit is ~10K chars) const CHUNK_SIZE: usize = 9000; fn main() { // Daemon agent calls set POC_AGENT=1 — skip memory search. if std::env::var("POC_AGENT").is_ok() { return; } let args = Args::parse(); if args.seen { show_seen(); return; } // --query mode: skip all hook/context machinery, just search if let Some(ref query_str) = args.query { run_query_mode(query_str, &args); return; } let input = if args.hook { // Hook mode: read from stdin, stash for later debug runs let mut buf = String::new(); io::stdin().read_to_string(&mut buf).unwrap_or_default(); fs::create_dir_all("/tmp/claude-memory-search").ok(); fs::write(STASH_PATH, &buf).ok(); buf } else { // All other modes: replay stashed input fs::read_to_string(STASH_PATH).unwrap_or_else(|_| { eprintln!("No stashed input at {}", STASH_PATH); std::process::exit(1); }) }; let debug = args.debug || !args.hook; let json: serde_json::Value = match serde_json::from_str(&input) { Ok(v) => v, Err(_) => return, }; let prompt = json["prompt"].as_str().unwrap_or(""); let session_id = json["session_id"].as_str().unwrap_or(""); if session_id.is_empty() { return; } let state_dir = PathBuf::from("/tmp/claude-memory-search"); fs::create_dir_all(&state_dir).ok(); // Detect post-compaction reload via mmap backward scan let transcript_path = json["transcript_path"].as_str().unwrap_or(""); let is_compaction = poc_memory::transcript::detect_new_compaction( &state_dir, session_id, transcript_path, ); // First prompt or post-compaction: load full context let cookie_path = state_dir.join(format!("cookie-{}", session_id)); let is_first = !cookie_path.exists(); if is_first || is_compaction { // Reset seen set and returned list let seen_path = state_dir.join(format!("seen-{}", session_id)); let returned_path = state_dir.join(format!("returned-{}", session_id)); fs::remove_file(&seen_path).ok(); fs::remove_file(&returned_path).ok(); } if debug { println!("[memory-search] session={} is_first={} is_compaction={}", session_id, is_first, is_compaction); } if is_first || is_compaction { // Create/touch the cookie let cookie = if is_first { let c = generate_cookie(); fs::write(&cookie_path, &c).ok(); c } else { fs::read_to_string(&cookie_path).unwrap_or_default().trim().to_string() }; if debug { println!("[memory-search] loading full context"); } // Load full memory context, chunk it, print first chunk, save rest if let Ok(output) = Command::new("poc-memory").args(["admin", "load-context"]).output() { if output.status.success() { let ctx = String::from_utf8_lossy(&output.stdout).to_string(); if !ctx.trim().is_empty() { // Extract keys from all chunks for seen set for line in ctx.lines() { if line.starts_with("--- ") && line.ends_with(" ---") { let inner = &line[4..line.len() - 4]; if let Some(paren) = inner.rfind(" (") { let key = inner[..paren].trim(); mark_seen(&state_dir, session_id, key); } } } let chunks = chunk_context(&ctx, CHUNK_SIZE); if debug { println!("[memory-search] context: {} bytes, {} chunks", ctx.len(), chunks.len()); } // Print first chunk if let Some(first) = chunks.first() { if args.hook { print!("{}", first); } } // Save remaining chunks for drip-feeding save_pending_chunks(&state_dir, session_id, &chunks[1..]); } } } let _ = cookie; } else { // Not first call: drip-feed next pending chunk if let Some(chunk) = pop_pending_chunk(&state_dir, session_id) { if debug { println!("[memory-search] drip-feeding pending chunk: {} bytes", chunk.len()); } if args.hook { print!("{}", chunk); } } } // Search requires a prompt (PostToolUse events don't have one) if prompt.is_empty() { return; } // Skip system/AFK prompts for prefix in &["is AFK", "You're on your own", "IRC mention"] { if prompt.starts_with(prefix) { return; } } let store = match store::Store::load() { Ok(s) => s, Err(_) => return, }; // Search for node keys in last ~150k tokens of transcript if debug { println!("[memory-search] transcript: {}", transcript_path); } let mut terms = extract_weighted_terms(transcript_path, 150_000, &store); // Also extract terms from the prompt itself (handles fresh sessions // and queries about topics not yet mentioned in the transcript) let prompt_terms = search::extract_query_terms(prompt, 8); if !prompt_terms.is_empty() { if debug { println!("[memory-search] prompt terms: {}", prompt_terms); } for word in prompt_terms.split_whitespace() { let lower = word.to_lowercase(); // Prompt terms get weight 1.0 (same as direct mention) terms.entry(lower).or_insert(1.0); } } // Boost node keys that appear as substrings in the current prompt. // Makes explicit mentions strong seeds for spread — the graph // determines what gets pulled in, this just ensures the seed fires. { let prompt_lower = prompt.to_lowercase(); for (key, node) in &store.nodes { if node.deleted { continue; } let key_lower = key.to_lowercase(); if key_lower.len() < 5 { continue; } if prompt_lower.contains(&key_lower) { *terms.entry(key_lower).or_insert(0.0) += 10.0; if debug { println!("[memory-search] prompt key boost: {} (+10.0)", key); } } } } if debug { println!("[memory-search] {} terms total", terms.len()); let mut by_weight: Vec<_> = terms.iter().collect(); by_weight.sort_by(|a, b| b.1.total_cmp(a.1)); for (term, weight) in by_weight.iter().take(20) { println!(" {:.3} {}", weight, term); } } if terms.is_empty() { if debug { println!("[memory-search] no terms found, done"); } return; } // Parse algorithm pipeline let pipeline: Vec = if args.pipeline.is_empty() { // Default: just spreading activation vec![AlgoStage::parse("spread").unwrap()] } else { let mut stages = Vec::new(); for arg in &args.pipeline { match AlgoStage::parse(arg) { Ok(s) => stages.push(s), Err(e) => { eprintln!("error: {}", e); std::process::exit(1); } } } stages }; if debug { let names: Vec = pipeline.iter().map(|s| format!("{}", s.algo)).collect(); println!("[memory-search] pipeline: {}", names.join(" → ")); } // Extract seeds from terms let graph = poc_memory::graph::build_graph_fast(&store); let (seeds, direct_hits) = search::match_seeds(&terms, &store); if seeds.is_empty() { if debug { println!("[memory-search] no seeds matched, done"); } return; } if debug { println!("[memory-search] {} seeds", seeds.len()); let mut sorted = seeds.clone(); sorted.sort_by(|a, b| b.1.total_cmp(&a.1)); for (key, score) in sorted.iter().take(20) { println!(" {:.4} {}", score, key); } } let raw_results = search::run_pipeline(&pipeline, seeds, &graph, &store, debug, args.max_results); let results: Vec = raw_results.into_iter() .map(|(key, activation)| { let is_direct = direct_hits.contains(&key); search::SearchResult { key, activation, is_direct, snippet: None } }).collect(); if debug { println!("[memory-search] {} search results", results.len()); for r in results.iter().take(10) { let marker = if r.is_direct { "→" } else { " " }; println!(" {} [{:.4}] {}", marker, r.activation, r.key); } } if results.is_empty() { if debug { println!("[memory-search] no results, done"); } return; } let seen = load_seen(&state_dir, session_id); if debug { println!("[memory-search] {} keys in seen set", seen.len()); } // Format results like poc-memory search output let search_output = search::format_results(&results); let cookie = fs::read_to_string(&cookie_path).unwrap_or_default().trim().to_string(); let mut result_output = String::new(); let mut count = 0; let max_entries = 3; for line in search_output.lines() { if count >= max_entries { break; } let trimmed = line.trim(); if trimmed.is_empty() { continue; } if let Some(key) = extract_key_from_line(trimmed) { if seen.contains(&key) { continue; } mark_seen(&state_dir, session_id, &key); mark_returned(&state_dir, session_id, &key); result_output.push_str(line); result_output.push('\n'); count += 1; } else if count > 0 { result_output.push_str(line); result_output.push('\n'); } } if count == 0 { if debug { println!("[memory-search] all results already seen"); } return; } if args.hook { println!("Recalled memories [{}]:", cookie); } print!("{}", result_output); // Record search hits with daemon (fire-and-forget) let hit_keys: Vec<&str> = results.iter().map(|r| r.key.as_str()).collect(); if debug { println!("[memory-search] recording {} search hits", hit_keys.len()); } match poc_memory::agents::daemon::rpc_record_hits(&hit_keys) { Ok(()) => { if debug { println!("[memory-search] hits recorded"); } } Err(e) => { if debug { println!("[memory-search] hit recording failed: {}", e); } } } // Clean up stale state files (opportunistic) cleanup_stale_files(&state_dir, Duration::from_secs(86400)); } /// Direct query mode: search for a term without hook/stash machinery. fn run_query_mode(query: &str, args: &Args) { let store = match poc_memory::store::Store::load() { Ok(s) => s, Err(e) => { eprintln!("failed to load store: {}", e); return; } }; // Build terms from the query string let mut terms: BTreeMap = BTreeMap::new(); let prompt_terms = search::extract_query_terms(query, 8); for word in prompt_terms.split_whitespace() { terms.entry(word.to_lowercase()).or_insert(1.0); } // Also check for exact node key match (the query itself, lowercased) let query_lower = query.to_lowercase(); for (key, node) in &store.nodes { if node.deleted { continue; } if key.to_lowercase() == query_lower { terms.insert(query_lower.clone(), 10.0); break; } } println!("[query] terms: {:?}", terms); if terms.is_empty() { println!("[query] no terms extracted"); return; } let graph = poc_memory::graph::build_graph_fast(&store); let (seeds, direct_hits) = search::match_seeds(&terms, &store); println!("[query] {} seeds", seeds.len()); let mut sorted = seeds.clone(); sorted.sort_by(|a, b| b.1.total_cmp(&a.1)); for (key, score) in sorted.iter().take(20) { let marker = if direct_hits.contains(key) { "→" } else { " " }; println!(" {} {:.4} {}", marker, score, key); } let pipeline: Vec = if args.pipeline.is_empty() { vec![AlgoStage::parse("spread").unwrap()] } else { args.pipeline.iter() .filter_map(|a| AlgoStage::parse(a).ok()) .collect() }; let max_results = args.max_results.max(25); let results = search::run_pipeline(&pipeline, seeds, &graph, &store, true, max_results); println!("\n[query] top {} results:", results.len().min(25)); for (i, (key, score)) in results.iter().take(25).enumerate() { let marker = if direct_hits.contains(key) { "→" } else { " " }; println!(" {:2}. {} [{:.4}] {}", i + 1, marker, score, key); } } /// Split context output into chunks of approximately `max_bytes`, breaking /// at section boundaries ("--- KEY (group) ---" lines). fn chunk_context(ctx: &str, max_bytes: usize) -> Vec { // Split into sections at group boundaries, then merge small adjacent // sections into chunks up to max_bytes. let mut sections: Vec = Vec::new(); let mut current = String::new(); for line in ctx.lines() { // Group headers start new sections if line.starts_with("--- ") && line.ends_with(" ---") && !current.is_empty() { sections.push(std::mem::take(&mut current)); } if !current.is_empty() { current.push('\n'); } current.push_str(line); } if !current.is_empty() { sections.push(current); } // Merge small sections into chunks, respecting max_bytes let mut chunks: Vec = Vec::new(); let mut chunk = String::new(); for section in sections { if !chunk.is_empty() && chunk.len() + section.len() + 1 > max_bytes { chunks.push(std::mem::take(&mut chunk)); } if !chunk.is_empty() { chunk.push('\n'); } chunk.push_str(§ion); } if !chunk.is_empty() { chunks.push(chunk); } chunks } /// Save remaining chunks to disk for drip-feeding on subsequent hook calls. fn save_pending_chunks(dir: &Path, session_id: &str, chunks: &[String]) { let chunks_dir = dir.join(format!("chunks-{}", session_id)); // Clear any old chunks let _ = fs::remove_dir_all(&chunks_dir); if chunks.is_empty() { return; } fs::create_dir_all(&chunks_dir).ok(); for (i, chunk) in chunks.iter().enumerate() { let path = chunks_dir.join(format!("{:04}", i)); fs::write(path, chunk).ok(); } } /// Pop the next pending chunk (lowest numbered file). Returns None if no chunks remain. fn pop_pending_chunk(dir: &Path, session_id: &str) -> Option { let chunks_dir = dir.join(format!("chunks-{}", session_id)); if !chunks_dir.exists() { return None; } let mut entries: Vec<_> = fs::read_dir(&chunks_dir).ok()? .flatten() .filter(|e| e.file_type().map(|t| t.is_file()).unwrap_or(false)) .collect(); entries.sort_by_key(|e| e.file_name()); let first = entries.first()?; let content = fs::read_to_string(first.path()).ok()?; fs::remove_file(first.path()).ok(); // Clean up directory if empty if fs::read_dir(&chunks_dir).ok().map(|mut d| d.next().is_none()).unwrap_or(true) { fs::remove_dir(&chunks_dir).ok(); } Some(content) } /// Reverse-scan the transcript JSONL, extracting text from user/assistant /// messages until we accumulate `max_tokens` tokens of text content. /// Then search for all node keys as substrings, weighted by position. fn extract_weighted_terms( path: &str, max_tokens: usize, store: &poc_memory::store::Store, ) -> BTreeMap { if path.is_empty() { return BTreeMap::new(); } let content = match fs::read_to_string(path) { Ok(c) => c, Err(_) => return BTreeMap::new(), }; // Collect text from messages, scanning backwards, until token budget hit let mut message_texts: Vec = Vec::new(); let mut token_count = 0; for line in content.lines().rev() { if token_count >= max_tokens { break; } let obj: serde_json::Value = match serde_json::from_str(line) { Ok(v) => v, Err(_) => continue, }; let msg_type = obj.get("type").and_then(|v| v.as_str()).unwrap_or(""); if msg_type != "user" && msg_type != "assistant" { continue; } let mut msg_text = String::new(); let msg = obj.get("message").unwrap_or(&obj); match msg.get("content") { Some(serde_json::Value::String(s)) => { msg_text.push_str(s); } Some(serde_json::Value::Array(arr)) => { for block in arr { if block.get("type").and_then(|v| v.as_str()) == Some("text") { if let Some(t) = block.get("text").and_then(|v| v.as_str()) { msg_text.push(' '); msg_text.push_str(t); } } } } _ => {} } token_count += msg_text.len() / 4; message_texts.push(msg_text); } // Reverse so oldest is first (position weighting: later = more recent = higher) message_texts.reverse(); let all_text = message_texts.join(" ").to_lowercase(); let text_len = all_text.len(); if text_len == 0 { return BTreeMap::new(); } // Search for each node key as a substring (casefolded), accumulate position-weighted score let mut terms = BTreeMap::new(); for (key, _node) in &store.nodes { let key_folded = key.to_lowercase(); let mut pos = 0; while let Some(found) = all_text[pos..].find(&key_folded) { let abs_pos = pos + found; let weight = (abs_pos + 1) as f64 / text_len as f64; *terms.entry(key_folded.clone()).or_insert(0.0) += weight; pos = abs_pos + key_folded.len(); } } terms } fn extract_key_from_line(line: &str) -> Option { let after_bracket = line.find("] ")?; let rest = &line[after_bracket + 2..]; let key_end = rest.find(" (c").unwrap_or(rest.len()); let key = rest[..key_end].trim(); if key.is_empty() { None } else { Some(key.to_string()) } } fn generate_cookie() -> String { uuid::Uuid::new_v4().as_simple().to_string()[..12].to_string() } /// Parse a seen-file line: "TIMESTAMP\tKEY" or legacy "KEY" fn parse_seen_line(line: &str) -> &str { line.split_once('\t').map(|(_, key)| key).unwrap_or(line) } fn load_seen(dir: &Path, session_id: &str) -> HashSet { let path = dir.join(format!("seen-{}", session_id)); if path.exists() { fs::read_to_string(path) .unwrap_or_default() .lines() .filter(|s| !s.is_empty()) .map(|s| parse_seen_line(s).to_string()) .collect() } else { HashSet::new() } } fn mark_seen(dir: &Path, session_id: &str, key: &str) { let path = dir.join(format!("seen-{}", session_id)); if let Ok(mut f) = fs::OpenOptions::new().create(true).append(true).open(path) { let ts = chrono::Local::now().format("%Y-%m-%dT%H:%M:%S"); writeln!(f, "{}\t{}", ts, key).ok(); } } fn mark_returned(dir: &Path, session_id: &str, key: &str) { let returned = load_returned(dir, session_id); if returned.contains(&key.to_string()) { return; } let path = dir.join(format!("returned-{}", session_id)); if let Ok(mut f) = fs::OpenOptions::new().create(true).append(true).open(path) { writeln!(f, "{}", key).ok(); } } fn load_returned(dir: &Path, session_id: &str) -> Vec { let path = dir.join(format!("returned-{}", session_id)); if path.exists() { let mut seen = HashSet::new(); fs::read_to_string(path) .unwrap_or_default() .lines() .filter(|s| !s.is_empty()) .filter(|s| seen.insert(s.to_string())) .map(|s| s.to_string()) .collect() } else { Vec::new() } } fn show_seen() { let state_dir = PathBuf::from("/tmp/claude-memory-search"); // Read stashed input for session_id let input = match fs::read_to_string(STASH_PATH) { Ok(s) => s, Err(_) => { eprintln!("No stashed input at {}", STASH_PATH); return; } }; let json: serde_json::Value = match serde_json::from_str(&input) { Ok(v) => v, Err(_) => { eprintln!("Failed to parse stashed input"); return; } }; let session_id = json["session_id"].as_str().unwrap_or(""); if session_id.is_empty() { eprintln!("No session_id in stashed input"); return; } let transcript_path = json["transcript_path"].as_str().unwrap_or(""); println!("Session: {}", session_id); let cookie_path = state_dir.join(format!("cookie-{}", session_id)); if let Ok(cookie) = fs::read_to_string(&cookie_path) { println!("Cookie: {}", cookie.trim()); } // Show last compaction info let compaction_path = state_dir.join(format!("compaction-{}", session_id)); match fs::read_to_string(&compaction_path) { Ok(offset_str) => { let offset: u64 = offset_str.trim().parse().unwrap_or(0); // Try to get a timestamp from the compaction offset in the transcript let ts = if !transcript_path.is_empty() && offset > 0 { poc_memory::transcript::compaction_timestamp(transcript_path, offset) } else { None }; match ts { Some(t) => println!("Last compaction: offset {} ({})", offset, t), None => println!("Last compaction: offset {}", offset), } } Err(_) => println!("Last compaction: none detected"), } // Pending chunks let chunks_dir = state_dir.join(format!("chunks-{}", session_id)); let pending = fs::read_dir(&chunks_dir).ok() .map(|d| d.flatten().count()) .unwrap_or(0); if pending > 0 { println!("Pending chunks: {}", pending); } // Read seen file in insertion order (append-only file) let seen_path = state_dir.join(format!("seen-{}", session_id)); let seen_lines: Vec = fs::read_to_string(&seen_path) .unwrap_or_default() .lines() .filter(|s| !s.is_empty()) .map(|s| s.to_string()) .collect(); let returned = load_returned(&state_dir, session_id); let returned_set: HashSet<_> = returned.iter().cloned().collect(); // Count context-loaded vs search-returned let context_keys: Vec<_> = seen_lines.iter() .map(|l| parse_seen_line(l).to_string()) .filter(|k| !returned_set.contains(k)) .collect(); let search_keys: Vec<_> = seen_lines.iter() .map(|l| parse_seen_line(l).to_string()) .filter(|k| returned_set.contains(k)) .collect(); println!("\nSeen set ({} total):", seen_lines.len()); if !context_keys.is_empty() { println!(" Context-loaded ({}):", context_keys.len()); for key in &context_keys { println!(" {}", key); } } if !search_keys.is_empty() { println!(" Search-returned ({}):", search_keys.len()); for key in &search_keys { println!(" {}", key); } } // Show returned keys that aren't in the seen set (bug indicator) let seen_key_set: HashSet<_> = seen_lines.iter() .map(|l| parse_seen_line(l).to_string()) .collect(); let orphan_returned: Vec<_> = returned.iter() .filter(|k| !seen_key_set.contains(k.as_str())) .collect(); if !orphan_returned.is_empty() { println!("\n WARNING: {} returned keys not in seen set (pre-compaction?):", orphan_returned.len()); for key in &orphan_returned { println!(" {}", key); } } } fn cleanup_stale_files(dir: &Path, max_age: Duration) { let entries = match fs::read_dir(dir) { Ok(e) => e, Err(_) => return, }; let cutoff = SystemTime::now() - max_age; for entry in entries.flatten() { if let Ok(meta) = entry.metadata() { if let Ok(modified) = meta.modified() { if modified < cutoff { fs::remove_file(entry.path()).ok(); } } } } }