consciousness/poc-memory/src/bin/memory-search.rs
Kent Overstreet 5d6b2021f8 Agent identity, parallel scheduling, memory-search fixes, stemmer optimization
- Agent identity injection: prepend core-personality to all agent prompts
  so agents dream as me, not as generic graph workers. Include instructions
  to walk the graph and connect new nodes to core concepts.

- Parallel agent scheduling: sequential within type, parallel across types.
  Different agent types (linker, organize, replay) run concurrently.

- Linker prompt: graph walking instead of keyword search for connections.
  "Explore the local topology and walk the graph until you find the best
  connections."

- memory-search fixes: format_results no longer truncates to 5 results,
  pipeline default raised to 50, returned file cleared on compaction,
  --seen and --seen-full merged, compaction timestamp in --seen output,
  max_entries=3 per prompt for steady memory drip.

- Stemmer optimization: strip_suffix now works in-place on a single String
  buffer instead of allocating 18 new Strings per word. Note for future:
  reversed-suffix trie for O(suffix_len) instead of O(n_rules).

- Transcript: add compaction_timestamp() for --seen display.

- Agent budget configurable (default 4000 from config).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-15 12:49:10 -04:00

770 lines
26 KiB
Rust

// 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<String>,
/// Algorithm pipeline stages: e.g. spread spectral,k=20 spread,max_hops=4
/// Default: spread.
pipeline: Vec<String>,
}
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<AlgoStage> = 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<String> = 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<search::SearchResult> = 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<String, f64> = 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<AlgoStage> = 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<String> {
// Split into sections at group boundaries, then merge small adjacent
// sections into chunks up to max_bytes.
let mut sections: Vec<String> = 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<String> = 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(&section);
}
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<String> {
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<String, f64> {
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<String> = 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<String> {
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<String> {
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<String> {
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<String> = 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();
}
}
}
}
}