diff --git a/src/agent/runner.rs b/src/agent/runner.rs index 122dcb6..d3964d6 100644 --- a/src/agent/runner.rs +++ b/src/agent/runner.rs @@ -94,6 +94,7 @@ impl Agent { personality, journal: String::new(), working_stack: Vec::new(), + loaded_nodes: Vec::new(), }; let session_id = format!("poc-agent-{}", chrono::Utc::now().format("%Y%m%d-%H%M%S")); let mut agent = Self { @@ -431,6 +432,66 @@ impl Agent { return; } + // Handle memory tools — needs &mut self for node tracking + if call.function.name.starts_with("memory_") { + let result = tools::memory::dispatch(&call.function.name, &args, None); + let text = match &result { + Ok(s) => s.clone(), + Err(e) => format!("Error: {:#}", e), + }; + + // Track loaded/updated nodes + if result.is_ok() { + match call.function.name.as_str() { + "memory_render" | "memory_links" => { + if let Some(key) = args.get("key").and_then(|v| v.as_str()) { + if let Some(node) = crate::agent::memory::MemoryNode::load(key) { + // Replace if already tracked, otherwise add + if let Some(existing) = self.context.loaded_nodes.iter_mut() + .find(|n| n.key == node.key) { + *existing = node; + } else { + self.context.loaded_nodes.push(node); + } + } + } + } + "memory_write" => { + if let Some(key) = args.get("key").and_then(|v| v.as_str()) { + if let Some(node) = crate::agent::memory::MemoryNode::load(key) { + // Refresh if already tracked + if let Some(existing) = self.context.loaded_nodes.iter_mut() + .find(|n| n.key == node.key) { + *existing = node; + } + // Don't auto-add writes — only renders register nodes + } + } + } + _ => {} + } + } + + let output = tools::ToolOutput { + text, + is_yield: false, + images: Vec::new(), + model_switch: None, + dmn_pause: false, + }; + let _ = ui_tx.send(UiMessage::ToolResult { + name: call.function.name.clone(), + result: output.text.clone(), + }); + let _ = ui_tx.send(UiMessage::ToolFinished { id: call.id.clone() }); + self.push_message(Message::tool_result(&call.id, &output.text)); + ds.had_tool_calls = true; + if output.text.starts_with("Error:") { + ds.tool_errors += 1; + } + return; + } + let output = tools::dispatch(&call.function.name, &args, &self.process_tracker).await; @@ -569,6 +630,29 @@ impl Agent { children: stack_children, }); + // Loaded memory nodes — tracked by memory tools + if !self.context.loaded_nodes.is_empty() { + let node_children: Vec = self.context.loaded_nodes.iter() + .map(|node| { + let rendered = node.render(); + ContextSection { + name: format!("{} (v{}, w={:.2}, {} links)", + node.key, node.version, node.weight, node.links.len()), + tokens: count(&rendered), + content: String::new(), // don't duplicate in debug view + children: Vec::new(), + } + }) + .collect(); + let node_tokens: usize = node_children.iter().map(|c| c.tokens).sum(); + sections.push(ContextSection { + name: format!("Memory nodes ({} loaded)", self.context.loaded_nodes.len()), + tokens: node_tokens, + content: String::new(), + children: node_children, + }); + } + // Conversation — each message as a child let conv_start = self.messages.iter() .position(|m| m.role == Role::Assistant || m.role == Role::Tool) diff --git a/src/agent/tools/mod.rs b/src/agent/tools/mod.rs index 94c76d4..5813526 100644 --- a/src/agent/tools/mod.rs +++ b/src/agent/tools/mod.rs @@ -102,7 +102,7 @@ pub async fn dispatch( "grep" => grep::grep(args), "glob" => glob_tool::glob_search(args), "journal" => journal::write_entry(args), - n if n.starts_with("memory_") => memory::dispatch(n, args, None), + // memory_* tools are dispatched in runner.rs for context tracking _ => Err(anyhow::anyhow!("Unknown tool: {}", name)), }; diff --git a/src/agent/types.rs b/src/agent/types.rs index 8995f0f..ed83a98 100644 --- a/src/agent/types.rs +++ b/src/agent/types.rs @@ -324,6 +324,10 @@ pub struct ContextState { pub personality: Vec<(String, String)>, pub journal: String, pub working_stack: Vec, + /// Memory nodes currently loaded in the context window. + /// Tracked so the agent knows what it's "seeing" and can + /// refresh nodes after writes. + pub loaded_nodes: Vec, } pub const WORKING_STACK_INSTRUCTIONS: &str = "/home/kent/.config/poc-agent/working-stack.md";