consciousness/src/bin/memory-search.rs

196 lines
6.1 KiB
Rust
Raw Normal View History

// memory-search: hook binary for ambient memory retrieval
//
// Reads JSON from stdin (Claude Code UserPromptSubmit hook format),
// searches memory for relevant entries, outputs results tagged with
// an anti-injection cookie.
//
// This is a thin wrapper that delegates to the poc-memory search
// engine but formats output for the hook protocol.
use std::collections::HashSet;
use std::fs;
use std::io::{self, Read, Write};
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::{Duration, SystemTime};
fn main() {
let mut input = String::new();
io::stdin().read_to_string(&mut input).unwrap_or_default();
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 prompt.is_empty() || session_id.is_empty() {
return;
}
// Skip short prompts
let word_count = prompt.split_whitespace().count();
if word_count < 3 {
return;
}
// Skip system/idle prompts
for prefix in &["Kent is AFK", "You're on your own", "IRC mention"] {
if prompt.starts_with(prefix) {
return;
}
}
// Extract search terms (strip stop words)
let query = extract_query_terms(prompt, 3);
if query.is_empty() {
return;
}
// Run poc-memory search
let output = Command::new("poc-memory")
.args(["search", &query])
.output();
let search_output = match output {
Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout).to_string(),
_ => return,
};
if search_output.trim().is_empty() {
return;
}
// Session state for dedup
let state_dir = PathBuf::from("/tmp/claude-memory-search");
fs::create_dir_all(&state_dir).ok();
// Clean up state files older than 24h (opportunistic, best-effort)
cleanup_stale_files(&state_dir, Duration::from_secs(86400));
let cookie = load_or_create_cookie(&state_dir, session_id);
let seen = load_seen(&state_dir, session_id);
// Parse search output and filter
let mut result_output = String::new();
let mut count = 0;
let max_entries = 5;
for line in search_output.lines() {
if count >= max_entries { break; }
// Lines starting with → or space+number are results
let trimmed = line.trim();
if trimmed.is_empty() { continue; }
// Extract key from result line like "→ 1. [0.83/0.83] identity.md (c4)"
if let Some(key) = extract_key_from_line(trimmed) {
if seen.contains(&key) { continue; }
mark_seen(&state_dir, session_id, &key);
result_output.push_str(line);
result_output.push('\n');
count += 1;
} else if count > 0 {
// Snippet line following a result
result_output.push_str(line);
result_output.push('\n');
}
}
if count == 0 { return; }
println!("Recalled memories [{}]:", cookie);
print!("{}", result_output);
}
fn extract_query_terms(text: &str, max_terms: usize) -> String {
const STOP_WORDS: &[&str] = &[
"the", "a", "an", "is", "are", "was", "were", "do", "does", "did",
"have", "has", "had", "will", "would", "could", "should", "can",
"may", "might", "shall", "been", "being", "to", "of", "in", "for",
"on", "with", "at", "by", "from", "as", "but", "or", "and", "not",
"no", "if", "then", "than", "that", "this", "it", "its", "my",
"your", "our", "we", "you", "i", "me", "he", "she", "they", "them",
"what", "how", "why", "when", "where", "about", "just", "let",
"want", "tell", "show", "think", "know", "see", "look", "make",
"get", "go", "some", "any", "all", "very", "really", "also", "too",
"so", "up", "out", "here", "there",
];
text.to_lowercase()
.split(|c: char| !c.is_alphanumeric())
.filter(|w| !w.is_empty() && w.len() > 2 && !STOP_WORDS.contains(w))
.take(max_terms)
.collect::<Vec<_>>()
.join(" ")
}
fn extract_key_from_line(line: &str) -> Option<String> {
// Match lines like "→ 1. [0.83/0.83] identity.md (c4)"
// or " 1. [0.83/0.83] identity.md (c4)"
let after_bracket = line.find("] ")?;
let rest = &line[after_bracket + 2..];
// Key is from here until optional " (c" or end of line
let key_end = rest.find(" (c").unwrap_or(rest.len());
let key = rest[..key_end].trim();
if key.is_empty() || !key.contains('.') {
None
} else {
Some(key.to_string())
}
}
fn load_or_create_cookie(dir: &Path, session_id: &str) -> String {
let path = dir.join(format!("cookie-{}", session_id));
if path.exists() {
fs::read_to_string(&path).unwrap_or_default().trim().to_string()
} else {
let cookie = generate_cookie();
fs::write(&path, &cookie).ok();
cookie
}
}
fn generate_cookie() -> String {
uuid::Uuid::new_v4().as_simple().to_string()[..12].to_string()
}
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()
.map(|s| 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) {
writeln!(f, "{}", key).ok();
}
}
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();
}
}
}
}
}