runner: context-aware memory tracking

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 <poc@bcachefs.org>
This commit is contained in:
ProofOfConcept 2026-03-25 01:48:15 -04:00
parent 2c61a3575d
commit 4b97bb2f2e
3 changed files with 89 additions and 1 deletions

View file

@ -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<ContextSection> = 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)