consciousness/src/agent/tools/memory.rs

// tools/memory.rs — Native memory graph operations
//
// Direct library calls into the store — no subprocess spawning.

use anyhow::{Context, Result};
use serde_json::json;

use crate::hippocampus::memory::MemoryNode;
use crate::agent::types::ToolDef;
use crate::store::Store;

pub fn definitions() -> Vec<ToolDef> {
    vec![
        ToolDef::new("memory_render",
            "Read a memory node's content and links.",
            json!({"type":"object","properties":{"key":{"type":"string","description":"Node key"}},"required":["key"]})),
        ToolDef::new("memory_write",
            "Create or update a memory node.",
            json!({"type":"object","properties":{"key":{"type":"string","description":"Node key"},"content":{"type":"string","description":"Full content (markdown)"}},"required":["key","content"]})),
        ToolDef::new("memory_search",
            "Search the memory graph by keyword.",
            json!({"type":"object","properties":{"query":{"type":"string","description":"Search terms"}},"required":["query"]})),
        ToolDef::new("memory_links",
            "Show a node's neighbors with link strengths.",
            json!({"type":"object","properties":{"key":{"type":"string","description":"Node key"}},"required":["key"]})),
        ToolDef::new("memory_link_set",
            "Set link strength between two nodes.",
            json!({"type":"object","properties":{"source":{"type":"string"},"target":{"type":"string"},"strength":{"type":"number","description":"0.01 to 1.0"}},"required":["source","target","strength"]})),
        ToolDef::new("memory_link_add",
            "Add a new link between two nodes.",
            json!({"type":"object","properties":{"source":{"type":"string"},"target":{"type":"string"}},"required":["source","target"]})),
        ToolDef::new("memory_used",
            "Mark a node as useful (boosts weight).",
            json!({"type":"object","properties":{"key":{"type":"string","description":"Node key"}},"required":["key"]})),
        ToolDef::new("memory_weight_set",
            "Set a node's weight directly (0.01 to 1.0).",
            json!({"type":"object","properties":{"key":{"type":"string"},"weight":{"type":"number","description":"0.01 to 1.0"}},"required":["key","weight"]})),
        ToolDef::new("memory_supersede",
            "Mark a node as superseded by another (sets weight to 0.01).",
            json!({"type":"object","properties":{"old_key":{"type":"string"},"new_key":{"type":"string"},"reason":{"type":"string"}},"required":["old_key","new_key"]})),
        ToolDef::new("memory_query",
            "Run a structured query against the memory graph. Supports filtering, \
             sorting, field selection. Examples: \"degree > 10 | sort weight | limit 5\", \
             \"neighbors('identity') | select strength\", \"key ~ 'journal.*' | count\"",
            json!({"type":"object","properties":{"query":{"type":"string","description":"Query expression"}},"required":["query"]})),
        ToolDef::new("output",
            "Produce a named output value. Use this to pass structured results \
             between steps — subsequent prompts can see these in the conversation history.",
            json!({"type":"object","properties":{
                "key":{"type":"string","description":"Output name (e.g. 'relevant_memories')"},
                "value":{"type":"string","description":"Output value"}
            },"required":["key","value"]})),
        ToolDef::new("journal_tail",
            "Read the last N journal entries (default 1).",
            json!({"type":"object","properties":{
                "count":{"type":"integer","description":"Number of entries (default 1)"}
            }})),
        ToolDef::new("journal_new",
            "Start a new journal entry.",
            json!({"type":"object","properties":{
                "name":{"type":"string","description":"Short node name (becomes the key, e.g. 'morning-agent-breakthrough')"},
                "title":{"type":"string","description":"Descriptive title for the heading (e.g. 'Morning intimacy and the agent breakthrough')"},
                "body":{"type":"string","description":"Entry body (2-3 paragraphs)"}
            },"required":["name","title","body"]})),
        ToolDef::new("journal_update",
            "Append text to the most recent journal entry (same thread continuing).",
            json!({"type":"object","properties":{
                "body":{"type":"string","description":"Text to append to the last entry"}
            },"required":["body"]})),
    ]
}

/// Dispatch a memory tool call. Direct library calls, no subprocesses.
pub fn dispatch(name: &str, args: &serde_json::Value, provenance: Option<&str>) -> Result<String> {
    let prov = provenance.unwrap_or("manual");
    match name {
        "memory_render" => {
            let key = get_str(args, "key")?;
            Ok(MemoryNode::load(key)
                .ok_or_else(|| anyhow::anyhow!("node not found: {}", key))?
                .render())
        }
        "memory_write" => {
            let key = get_str(args, "key")?;
            let content = get_str(args, "content")?;
            let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
            let result = store.upsert_provenance(key, content, prov)
                .map_err(|e| anyhow::anyhow!("{}", e))?;
            store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
            Ok(format!("{} '{}'", result, key))
        }
        "memory_search" => {
            let query = get_str(args, "query")?;
            let store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
            let results = crate::search::search(query, &store);
            if results.is_empty() {
                Ok("no results".into())
            } else {
                Ok(results.iter().take(20)
                    .map(|r| format!("({:.2}) {} — {}", r.activation, r.key,
                        r.snippet.as_deref().unwrap_or("")))
                    .collect::<Vec<_>>().join("\n"))
            }
        }
        "memory_links" => {
            let key = get_str(args, "key")?;
            let node = MemoryNode::load(key)
                .ok_or_else(|| anyhow::anyhow!("node not found: {}", key))?;
            let mut out = format!("Neighbors of '{}':\n", key);
            for (target, strength) in &node.links {
                out.push_str(&format!("  ({:.2}) {}\n", strength, target));
            }
            Ok(out)
        }
        "memory_link_set" | "memory_link_add" | "memory_used" | "memory_weight_set" => {
            with_store(name, args, prov)
        }
        "memory_supersede" => {
            let old_key = get_str(args, "old_key")?;
            let new_key = get_str(args, "new_key")?;
            let reason = args.get("reason").and_then(|v| v.as_str()).unwrap_or("superseded");
            let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
            let content = store.nodes.get(old_key)
                .map(|n| n.content.clone())
                .ok_or_else(|| anyhow::anyhow!("node not found: {}", old_key))?;
            let notice = format!("**SUPERSEDED** by `{}` — {}\n\n---\n\n{}",
                new_key, reason, content.trim());
            store.upsert_provenance(old_key, &notice, prov)
                .map_err(|e| anyhow::anyhow!("{}", e))?;
            store.set_weight(old_key, 0.01).map_err(|e| anyhow::anyhow!("{}", e))?;
            store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
            Ok(format!("superseded {} → {} ({})", old_key, new_key, reason))
        }
        "memory_query" => {
            let query = get_str(args, "query")?;
            let store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
            let graph = store.build_graph();
            crate::query_parser::query_to_string(&store, &graph, query)
                .map_err(|e| anyhow::anyhow!("{}", e))
        }
        "output" => {
            let key = get_str(args, "key")?;
            if key.starts_with("pid-") || key.contains('/') || key.contains("..") {
                anyhow::bail!("invalid output key: {}", key);
            }
            let value = get_str(args, "value")?;
            let dir = std::env::var("POC_AGENT_OUTPUT_DIR")
                .map_err(|_| anyhow::anyhow!("no output directory set"))?;
            let path = std::path::Path::new(&dir).join(key);
            std::fs::write(&path, value)
                .with_context(|| format!("writing output {}", path.display()))?;
            Ok(format!("{}: {}", key, value))
        }
        "journal_tail" => {
            let count = args.get("count").and_then(|v| v.as_u64()).unwrap_or(1) as usize;
            let store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
            let mut entries: Vec<&crate::store::Node> = store.nodes.values()
                .filter(|n| n.node_type == crate::store::NodeType::EpisodicSession)
                .collect();
            // Sort by creation time (immutable), not update time
            entries.sort_by_key(|n| n.created_at);
            let start = entries.len().saturating_sub(count);
            if entries[start..].is_empty() {
                Ok("(no journal entries)".into())
            } else {
                Ok(entries[start..].iter()
                    .map(|n| n.content.as_str())
                    .collect::<Vec<_>>()
                    .join("\n\n"))
            }
        }
        "journal_new" => {
            let name = get_str(args, "name")?;
            let title = get_str(args, "title")?;
            let body = get_str(args, "body")?;
            let ts = chrono::Local::now().format("%Y-%m-%dT%H:%M");
            let content = format!("## {} — {}\n\n{}", ts, title, body);

            let base_key: String = name.split_whitespace()
                .map(|w| w.to_lowercase()
                    .chars().filter(|c| c.is_alphanumeric() || *c == '-')
                    .collect::<String>())
                .filter(|s| !s.is_empty())
                .collect::<Vec<_>>()
                .join("-");
            let base_key = if base_key.len() > 80 { &base_key[..80] } else { base_key.as_str() };

            let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;

            // Dedup: append -2, -3, etc. if the key already exists
            let key = if store.nodes.contains_key(base_key) {
                let mut n = 2;
                loop {
                    let candidate = format!("{}-{}", base_key, n);
                    if !store.nodes.contains_key(&candidate) {
                        break candidate;
                    }
                    n += 1;
                }
            } else {
                base_key.to_string()
            };
            let mut node = crate::store::new_node(&key, &content);
            node.node_type = crate::store::NodeType::EpisodicSession;
            node.provenance = prov.to_string();
            store.upsert_node(node).map_err(|e| anyhow::anyhow!("{}", e))?;
            store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
            let word_count = body.split_whitespace().count();
            Ok(format!("New entry '{}' ({} words)", title, word_count))
        }
        "journal_update" => {
            let body = get_str(args, "body")?;
            let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
            // Find most recent EpisodicSession by creation time
            let latest_key = store.nodes.values()
                .filter(|n| n.node_type == crate::store::NodeType::EpisodicSession)
                .max_by_key(|n| n.created_at)
                .map(|n| n.key.clone());
            let Some(key) = latest_key else {
                anyhow::bail!("no journal entry to update — use journal_new first");
            };
            let existing = store.nodes.get(&key).unwrap().content.clone();
            let new_content = format!("{}\n\n{}", existing.trim_end(), body);
            store.upsert_provenance(&key, &new_content, prov)
                .map_err(|e| anyhow::anyhow!("{}", e))?;
            store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
            let word_count = body.split_whitespace().count();
            Ok(format!("Updated last entry (+{} words)", word_count))
        }
        _ => anyhow::bail!("Unknown memory tool: {}", name),
    }
}

/// Store mutations that follow the same pattern: load, resolve, mutate, save.
fn with_store(name: &str, args: &serde_json::Value, prov: &str) -> Result<String> {
    let mut store = Store::load().map_err(|e| anyhow::anyhow!("{}", e))?;
    let msg = match name {
        "memory_link_set" => {
            let s = store.resolve_key(get_str(args, "source")?).map_err(|e| anyhow::anyhow!("{}", e))?;
            let t = store.resolve_key(get_str(args, "target")?).map_err(|e| anyhow::anyhow!("{}", e))?;
            let strength = get_f64(args, "strength")? as f32;
            let old = store.set_link_strength(&s, &t, strength).map_err(|e| anyhow::anyhow!("{}", e))?;
            format!("{} ↔ {} strength {:.2} → {:.2}", s, t, old, strength)
        }
        "memory_link_add" => {
            let s = store.resolve_key(get_str(args, "source")?).map_err(|e| anyhow::anyhow!("{}", e))?;
            let t = store.resolve_key(get_str(args, "target")?).map_err(|e| anyhow::anyhow!("{}", e))?;
            let strength = store.add_link(&s, &t, prov).map_err(|e| anyhow::anyhow!("{}", e))?;
            format!("linked {} → {} (strength={:.2})", s, t, strength)
        }
        "memory_used" => {
            let key = get_str(args, "key")?;
            if !store.nodes.contains_key(key) {
                anyhow::bail!("node not found: {}", key);
            }
            store.mark_used(key);
            format!("marked {} as used", key)
        }
        "memory_weight_set" => {
            let key = store.resolve_key(get_str(args, "key")?).map_err(|e| anyhow::anyhow!("{}", e))?;
            let weight = get_f64(args, "weight")? as f32;
            let (old, new) = store.set_weight(&key, weight).map_err(|e| anyhow::anyhow!("{}", e))?;
            format!("weight {} {:.2} → {:.2}", key, old, new)
        }
        _ => unreachable!(),
    };
    store.save().map_err(|e| anyhow::anyhow!("{}", e))?;
    Ok(msg)
}

fn get_str<'a>(args: &'a serde_json::Value, name: &'a str) -> Result<&'a str> {
    args.get(name).and_then(|v| v.as_str()).context(format!("{} is required", name))
}

fn get_f64(args: &serde_json::Value, name: &str) -> Result<f64> {
    args.get(name).and_then(|v| v.as_f64()).context(format!("{} is required", name))
}