187 lines
5.8 KiB
Rust
187 lines
5.8 KiB
Rust
|
|
// 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::PathBuf;
|
||
|
|
use std::process::Command;
|
||
|
|
|
||
|
|
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();
|
||
|
|
|
||
|
|
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: &PathBuf, 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 {
|
||
|
|
let out = Command::new("head")
|
||
|
|
.args(["-c", "12", "/dev/urandom"])
|
||
|
|
.output()
|
||
|
|
.expect("failed to read urandom");
|
||
|
|
out.stdout.iter()
|
||
|
|
.map(|b| {
|
||
|
|
let idx = (*b as usize) % 62;
|
||
|
|
if idx < 10 { (b'0' + idx as u8) as char }
|
||
|
|
else if idx < 36 { (b'a' + (idx - 10) as u8) as char }
|
||
|
|
else { (b'A' + (idx - 36) as u8) as char }
|
||
|
|
})
|
||
|
|
.take(16)
|
||
|
|
.collect()
|
||
|
|
}
|
||
|
|
|
||
|
|
fn load_seen(dir: &PathBuf, 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: &PathBuf, 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();
|
||
|
|
}
|
||
|
|
}
|