From 4b97bb2f2e39ab96fbb8917d6c47719b5e50cb8d Mon Sep 17 00:00:00 2001 From: ProofOfConcept Date: Wed, 25 Mar 2026 01:48:15 -0400 Subject: [PATCH] runner: context-aware memory tracking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Memory tools now dispatch through a special path in the runner (like working_stack) instead of the generic tools::dispatch. This gives them &mut self access to track loaded nodes: - memory_render/memory_links: loads MemoryNode, registers in context.loaded_nodes (replace if already tracked) - memory_write: refreshes existing tracked node if present - All other memory tools: dispatch directly, no tracking needed The debug screen (context_state_summary) now shows a "Memory nodes" section listing all loaded nodes with version, weight, and link count. This is the agent knowing what it's holding — the foundation for intelligent refresh and eviction. Co-Authored-By: Proof of Concept --- src/agent/runner.rs | 84 ++++++++++++++++++++++++++++++++++++++++++ src/agent/tools/mod.rs | 2 +- src/agent/types.rs | 4 ++ 3 files changed, 89 insertions(+), 1 deletion(-) 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";