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 <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-13 12:08:46 -04:00
parent dc1049f62d
commit d7a5ac6347

View file

@ -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<std::sync::Arc<crate::agent::Agent>>, 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<std::sync::Arc<crate::agent::Agent>>,
args: serde_json::Value,
) -> Result<String> {
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<String> {
}
}
async fn write(agent: &Option<std::sync::Arc<crate::agent::Agent>>, args: &serde_json::Value) -> Result<String> {
async fn write(args: &serde_json::Value) -> Result<String> {
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<String> {
Ok(format!("{}{} strength {:.2}{:.2}", s, t, old, strength))
}
async fn link_add(agent: &Option<std::sync::Arc<crate::agent::Agent>>, args: &serde_json::Value) -> Result<String> {
async fn link_add(args: &serde_json::Value) -> Result<String> {
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<String> {
Ok(format!("Renamed '{}' → '{}'", resolved, new_key))
}
async fn supersede(agent: &Option<std::sync::Arc<crate::agent::Agent>>, args: &serde_json::Value) -> Result<String> {
async fn supersede(args: &serde_json::Value) -> Result<String> {
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<std::sync::Arc<crate::agent::Agent>>, 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, &notice, &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<std::sync::Arc<crate::agent::Agent>>, args: &serde_json::Value) -> Result<String> {
async fn journal_new(args: &serde_json::Value) -> Result<String> {
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<std::sync::Arc<crate::agent::Agent>>, 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<std::sync::Arc<crate::agent::Agent>>, args: &serde_json::Value) -> Result<String> {
async fn journal_update(args: &serde_json::Value) -> Result<String> {
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<std::sync::Arc<crate::agent::Agent>>, 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))?;