From d7a5ac6347f319eca3235684ab927c021aa116c2 Mon Sep 17 00:00:00 2001 From: Kent Overstreet Date: Mon, 13 Apr 2026 12:08:46 -0400 Subject: [PATCH] memory tools: simplify provenance handling Move provenance injection to dispatch() entry point - agent provenance is always written to args._provenance before routing. Individual tool functions now just call get_provenance(args) which is sync and simple. Removes agent parameter from: write, link_add, supersede, journal_new, journal_update. Co-Authored-By: Proof of Concept --- src/agent/tools/memory.rs | 59 +++++++++++++++++++-------------------- 1 file changed, 28 insertions(+), 31 deletions(-) diff --git a/src/agent/tools/memory.rs b/src/agent/tools/memory.rs index 65a5ad9..5aa5f81 100644 --- a/src/agent/tools/memory.rs +++ b/src/agent/tools/memory.rs @@ -68,16 +68,12 @@ pub async fn run_with_local_store(tool_name: &str, args: serde_json::Value) -> R result } -/// Get provenance from agent, or from args._provenance, or "manual". -async fn get_provenance(agent: &Option>, args: &serde_json::Value) -> String { - // Check args first (set by RPC path) - if let Some(p) = args.get("_provenance").and_then(|v| v.as_str()) { - return p.to_string(); - } - match agent { - Some(a) => a.state.lock().await.provenance.clone(), - None => "manual".to_string(), - } +/// Get provenance from args._provenance, or "manual". +fn get_provenance(args: &serde_json::Value) -> String { + args.get("_provenance") + .and_then(|v| v.as_str()) + .unwrap_or("manual") + .to_string() } /// Single entry point for all memory/journal tool calls. @@ -87,13 +83,14 @@ async fn dispatch( agent: &Option>, args: serde_json::Value, ) -> Result { + let mut args = args; + if let Some(a) = agent { + let prov = a.state.lock().await.provenance.clone(); + args.as_object_mut().map(|o| o.insert("_provenance".into(), prov.into())); + } + if !is_daemon() { - // Forward to daemon, attaching provenance - let mut args = args; - if let Some(a) = agent { - let prov = a.state.lock().await.provenance.clone(); - args.as_object_mut().map(|o| o.insert("_provenance".into(), prov.into())); - } + // Forward to daemon let name = tool_name.to_string(); return tokio::task::spawn_blocking(move || { crate::mcp_server::memory_rpc(&name, args) @@ -103,16 +100,16 @@ async fn dispatch( // Daemon path - dispatch to implementation match tool_name { "memory_render" => render(&args).await, - "memory_write" => write(agent, &args).await, + "memory_write" => write(&args).await, "memory_search" => search(&args).await, "memory_links" => links(&args).await, "memory_link_set" => link_set(&args).await, - "memory_link_add" => link_add(agent, &args).await, + "memory_link_add" => link_add(&args).await, "memory_delete" => delete(&args).await, "memory_history" => history(&args).await, "memory_weight_set" => weight_set(&args).await, "memory_rename" => rename(&args).await, - "memory_supersede" => supersede(agent, &args).await, + "memory_supersede" => supersede(&args).await, "memory_query" => query(&args).await, "graph_topology" => graph_topology().await, "graph_health" => graph_health().await, @@ -122,8 +119,8 @@ async fn dispatch( "graph_link_impact" => graph_link_impact(&args).await, "graph_hubs" => graph_hubs(&args).await, "journal_tail" => journal_tail(&args).await, - "journal_new" => journal_new(agent, &args).await, - "journal_update" => journal_update(agent, &args).await, + "journal_new" => journal_new(&args).await, + "journal_update" => journal_update(&args).await, _ => anyhow::bail!("unknown tool: {}", tool_name), } } @@ -245,10 +242,10 @@ async fn render(args: &serde_json::Value) -> Result { } } -async fn write(agent: &Option>, args: &serde_json::Value) -> Result { +async fn write(args: &serde_json::Value) -> Result { let key = get_str(args, "key")?; let content = get_str(args, "content")?; - let prov = get_provenance(agent, args).await; + let prov = get_provenance(args); let arc = cached_store().await?; let mut store = arc.lock().await; let result = store.upsert_provenance(key, content, &prov) @@ -322,12 +319,12 @@ async fn link_set(args: &serde_json::Value) -> Result { Ok(format!("{} ↔ {} strength {:.2} → {:.2}", s, t, old, strength)) } -async fn link_add(agent: &Option>, args: &serde_json::Value) -> Result { +async fn link_add(args: &serde_json::Value) -> Result { let arc = cached_store().await?; let mut store = arc.lock().await; 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 prov = get_provenance(agent, args).await; + let prov = get_provenance(args); let strength = store.add_link(&s, &t, &prov).map_err(|e| anyhow::anyhow!("{}", e))?; store.save().map_err(|e| anyhow::anyhow!("{}", e))?; Ok(format!("linked {} → {} (strength={:.2})", s, t, strength)) @@ -419,7 +416,7 @@ async fn rename(args: &serde_json::Value) -> Result { Ok(format!("Renamed '{}' → '{}'", resolved, new_key)) } -async fn supersede(agent: &Option>, args: &serde_json::Value) -> Result { +async fn supersede(args: &serde_json::Value) -> Result { 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"); @@ -430,7 +427,7 @@ async fn supersede(agent: &Option>, args: &s .ok_or_else(|| anyhow::anyhow!("node not found: {}", old_key))?; let notice = format!("**SUPERSEDED** by `{}` — {}\n\n---\n\n{}", new_key, reason, content.trim()); - let prov = get_provenance(agent, args).await; + let prov = get_provenance(args); store.upsert_provenance(old_key, ¬ice, &prov) .map_err(|e| anyhow::anyhow!("{}", e))?; store.set_weight(old_key, 0.01).map_err(|e| anyhow::anyhow!("{}", e))?; @@ -526,7 +523,7 @@ fn level_to_node_type(level: i64) -> crate::store::NodeType { } } -async fn journal_new(agent: &Option>, args: &serde_json::Value) -> Result { +async fn journal_new(args: &serde_json::Value) -> Result { let name = get_str(args, "name")?; let title = get_str(args, "title")?; let body = get_str(args, "body")?; @@ -557,14 +554,14 @@ async fn journal_new(agent: &Option>, args: }; let mut node = crate::store::new_node(&key, &content); node.node_type = level_to_node_type(level); - node.provenance = get_provenance(agent, args).await; + node.provenance = get_provenance(args); 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)) } -async fn journal_update(agent: &Option>, args: &serde_json::Value) -> Result { +async fn journal_update(args: &serde_json::Value) -> Result { let body = get_str(args, "body")?; let level = args.get("level").and_then(|v| v.as_i64()).unwrap_or(0); let node_type = level_to_node_type(level); @@ -579,7 +576,7 @@ async fn journal_update(agent: &Option>, arg }; let existing = store.nodes.get(&key).unwrap().content.clone(); let new_content = format!("{}\n\n{}", existing.trim_end(), body); - let prov = get_provenance(agent, args).await; + let prov = get_provenance(args); store.upsert_provenance(&key, &new_content, &prov) .map_err(|e| anyhow::anyhow!("{}", e))?; store.save().map_err(|e| anyhow::anyhow!("{}", e))?;