consciousness/poc-memory/src/bin/memory-search.rs

// memory-search: combined hook for session context loading + ambient memory retrieval
//
// On first prompt per session: loads full memory context (identity, journal, etc.)
// On subsequent prompts: searches memory for relevant entries
// On post-compaction: reloads full context
//
// Reads JSON from stdin (Claude Code UserPromptSubmit hook format),
// outputs results for injection into the conversation.

use poc_memory::search;
use poc_memory::store;
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;
    }

    let state_dir = PathBuf::from("/tmp/claude-memory-search");
    fs::create_dir_all(&state_dir).ok();

    // Detect post-compaction reload
    let is_compaction = prompt.contains("continued from a previous conversation");

    // 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 {
        // 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()
        };

        // Load full memory context
        if let Ok(output) = Command::new("poc-memory").args(["load-context"]).output() {
            if output.status.success() {
                let ctx = String::from_utf8_lossy(&output.stdout);
                if !ctx.trim().is_empty() {
                    print!("{}", ctx);
                }
            }
        }

        // On first prompt, also bump lookup counter for the cookie
        let _ = cookie; // used for tagging below
    }

    // Always do ambient search (skip on very short or system prompts)
    let word_count = prompt.split_whitespace().count();
    if word_count < 3 {
        return;
    }

    for prefix in &["is AFK", "You're on your own", "IRC mention"] {
        if prompt.starts_with(prefix) {
            return;
        }
    }

    let query = search::extract_query_terms(prompt, 3);
    if query.is_empty() {
        return;
    }

    let store = match store::Store::load() {
        Ok(s) => s,
        Err(_) => return,
    };

    let results = search::search(&query, &store);
    if results.is_empty() {
        return;
    }

    // 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 seen = load_seen(&state_dir, session_id);

    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; }

        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);
            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 { return; }

    println!("Recalled memories [{}]:", cookie);
    print!("{}", result_output);

    // Clean up stale state files (opportunistic)
    cleanup_stale_files(&state_dir, Duration::from_secs(86400));
}


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() || !key.contains('.') {
        None
    } else {
        Some(key.to_string())
    }
}

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();
                }
            }
        }
    }
}