From 933221f4827d8476364b410c9e999fbf147b48cf Mon Sep 17 00:00:00 2001 From: Kent Overstreet Date: Mon, 13 Apr 2026 13:12:11 -0400 Subject: [PATCH] memory tools: generate public typed API via macro The memory_tool! macro now generates two functions: - jsonargs_*() - internal, takes JSON args for dispatch table - pub fn name() - typed args, handles RPC-vs-local automatically Callers can now use typed Rust API: memory::write(Some(&agent), "key", "content").await?; memory::query(None, "all | type:semantic", Some("full")).await?; No more manual JSON construction for memory tool calls. Co-Authored-By: Proof of Concept --- src/agent/tools/memory.rs | 179 ++++++++++++++++++++++++++++---------- 1 file changed, 131 insertions(+), 48 deletions(-) diff --git a/src/agent/tools/memory.rs b/src/agent/tools/memory.rs index b8f73ff..4596e38 100644 --- a/src/agent/tools/memory.rs +++ b/src/agent/tools/memory.rs @@ -75,33 +75,14 @@ fn get_provenance(args: &serde_json::Value) -> String { // ── Macro for generating tool wrappers ───────────────────────── // -// memory_tool!(name, mut, arg1: str, arg2: f32, arg3: ?str) +// memory_tool!(name, mut, arg1: [str], arg2: [Option]) // - mut/ref for store mutability -// - type suffixes: str, f32, f64, u64, i64, bool -// - ?type for optional args with default +// - generates jsonargs_* (internal, JSON args) and public typed API macro_rules! memory_tool { - // Mutable store variant - ($name:ident, mut $(, $($arg:ident : [$($typ:tt)+]),* $(,)?)?) => { - async fn $name(args: &serde_json::Value) -> Result { - $($(let $arg = memory_tool!(@extract args, $arg, $($typ)+);)*)? - let prov = get_provenance(args); - let arc = cached_store().await?; - let mut store = arc.lock().await; - crate::hippocampus::$name(&mut store, &prov $($(, $arg)*)?) - } - }; - // Immutable store variant - ($name:ident, ref $(, $($arg:ident : [$($typ:tt)+]),* $(,)?)?) => { - async fn $name(args: &serde_json::Value) -> Result { - $($(let $arg = memory_tool!(@extract args, $arg, $($typ)+);)*)? - let prov = get_provenance(args); - let arc = cached_store().await?; - let store = arc.lock().await; - crate::hippocampus::$name(&store, &prov $($(, $arg)*)?) - } - }; - // Required extractors - fail if missing + // ── Helper rules (must come first) ───────────────────────────── + + // Extract from JSON (@extract $args:ident, $name:ident, str) => { get_str($args, stringify!($name))? }; @@ -114,8 +95,6 @@ macro_rules! memory_tool { .map(|arr| arr.iter().filter_map(|v| v.as_str().map(String::from)).collect::>()) .unwrap_or_default() }; - - // Optional extractors - return Option (@extract $args:ident, $name:ident, Option<&str>) => { $args.get(stringify!($name)).and_then(|v| v.as_str()) }; @@ -137,6 +116,110 @@ macro_rules! memory_tool { (@extract $args:ident, $name:ident, Option) => { $args.get(stringify!($name)).and_then(|v| v.as_f64()) }; + + // Parameter types for function signatures + (@param_type str) => { &str }; + (@param_type f32) => { f32 }; + (@param_type Vec) => { Vec }; + (@param_type Option<&str>) => { Option<&str> }; + (@param_type Option) => { Option }; + (@param_type Option) => { Option }; + (@param_type Option) => { Option }; + (@param_type Option) => { Option }; + (@param_type Option) => { Option }; + (@param_type Option) => { Option }; + + // Serialize to JSON for RPC + (@insert_json $map:ident, $name:ident, str) => { + $map.insert(stringify!($name).into(), serde_json::json!($name)); + }; + (@insert_json $map:ident, $name:ident, f32) => { + $map.insert(stringify!($name).into(), serde_json::json!($name)); + }; + (@insert_json $map:ident, $name:ident, Vec) => { + $map.insert(stringify!($name).into(), serde_json::json!($name)); + }; + (@insert_json $map:ident, $name:ident, Option<&str>) => { + if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); } + }; + (@insert_json $map:ident, $name:ident, Option) => { + if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); } + }; + (@insert_json $map:ident, $name:ident, Option) => { + if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); } + }; + (@insert_json $map:ident, $name:ident, Option) => { + if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); } + }; + (@insert_json $map:ident, $name:ident, Option) => { + if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); } + }; + (@insert_json $map:ident, $name:ident, Option) => { + if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); } + }; + (@insert_json $map:ident, $name:ident, Option) => { + if let Some(v) = $name { $map.insert(stringify!($name).into(), serde_json::json!(v)); } + }; + + // ── Main rules ───────────────────────────────────────────────── + + // Mutable store variant + ($name:ident, mut $(, $($arg:ident : [$($typ:tt)+]),* $(,)?)?) => { + paste::paste! { + async fn [](args: &serde_json::Value) -> Result { + $($(let $arg = memory_tool!(@extract args, $arg, $($typ)+);)*)? + let prov = get_provenance(args); + let arc = cached_store().await?; + let mut store = arc.lock().await; + crate::hippocampus::$name(&mut store, &prov $($(, $arg)*)?) + } + + pub async fn $name(agent: Option<&crate::agent::Agent> $($(, $arg: memory_tool!(@param_type $($typ)+))*)?) -> Result { + if !is_daemon() { + #[allow(unused_mut)] + let mut map = serde_json::Map::new(); + $($(memory_tool!(@insert_json map, $arg, $($typ)+);)*)? + return crate::mcp_server::memory_rpc(concat!("memory_", stringify!($name)), serde_json::Value::Object(map)); + } + let prov = match agent { + Some(a) => a.state.lock().await.provenance.clone(), + None => "manual".to_string(), + }; + let arc = cached_store().await?; + let mut store = arc.lock().await; + crate::hippocampus::$name(&mut store, &prov $($(, $arg)*)?) + } + } + }; + + // Immutable store variant + ($name:ident, ref $(, $($arg:ident : [$($typ:tt)+]),* $(,)?)?) => { + paste::paste! { + async fn [](args: &serde_json::Value) -> Result { + $($(let $arg = memory_tool!(@extract args, $arg, $($typ)+);)*)? + let prov = get_provenance(args); + let arc = cached_store().await?; + let store = arc.lock().await; + crate::hippocampus::$name(&store, &prov $($(, $arg)*)?) + } + + pub async fn $name(agent: Option<&crate::agent::Agent> $($(, $arg: memory_tool!(@param_type $($typ)+))*)?) -> Result { + if !is_daemon() { + #[allow(unused_mut)] + let mut map = serde_json::Map::new(); + $($(memory_tool!(@insert_json map, $arg, $($typ)+);)*)? + return crate::mcp_server::memory_rpc(concat!("memory_", stringify!($name)), serde_json::Value::Object(map)); + } + let prov = match agent { + Some(a) => a.state.lock().await.provenance.clone(), + None => "manual".to_string(), + }; + let arc = cached_store().await?; + let store = arc.lock().await; + crate::hippocampus::$name(&store, &prov $($(, $arg)*)?) + } + } + }; } // ── Memory tools ─────────────────────────────────────────────── @@ -193,28 +276,28 @@ async fn dispatch( // Daemon path - dispatch to implementation match tool_name { - "memory_render" => render(&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(&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(&args).await, - "memory_query" => query(&args).await, - "graph_topology" => graph_topology(&args).await, - "graph_health" => graph_health(&args).await, - "graph_communities" => graph_communities(&args).await, - "graph_normalize_strengths" => graph_normalize_strengths(&args).await, - "graph_trace" => graph_trace(&args).await, - "graph_link_impact" => graph_link_impact(&args).await, - "graph_hubs" => graph_hubs(&args).await, - "journal_tail" => journal_tail(&args).await, - "journal_new" => journal_new(&args).await, - "journal_update" => journal_update(&args).await, + "memory_render" => jsonargs_render(&args).await, + "memory_write" => jsonargs_write(&args).await, + "memory_search" => jsonargs_search(&args).await, + "memory_links" => jsonargs_links(&args).await, + "memory_link_set" => jsonargs_link_set(&args).await, + "memory_link_add" => jsonargs_link_add(&args).await, + "memory_delete" => jsonargs_delete(&args).await, + "memory_history" => jsonargs_history(&args).await, + "memory_weight_set" => jsonargs_weight_set(&args).await, + "memory_rename" => jsonargs_rename(&args).await, + "memory_supersede" => jsonargs_supersede(&args).await, + "memory_query" => jsonargs_query(&args).await, + "graph_topology" => jsonargs_graph_topology(&args).await, + "graph_health" => jsonargs_graph_health(&args).await, + "graph_communities" => jsonargs_graph_communities(&args).await, + "graph_normalize_strengths" => jsonargs_graph_normalize_strengths(&args).await, + "graph_trace" => jsonargs_graph_trace(&args).await, + "graph_link_impact" => jsonargs_graph_link_impact(&args).await, + "graph_hubs" => jsonargs_graph_hubs(&args).await, + "journal_tail" => jsonargs_journal_tail(&args).await, + "journal_new" => jsonargs_journal_new(&args).await, + "journal_update" => jsonargs_journal_update(&args).await, _ => anyhow::bail!("unknown tool: {}", tool_name), } }