From 9517b1b310ea533ba0741f1b561823593a7a416a Mon Sep 17 00:00:00 2001 From: Kent Overstreet Date: Fri, 20 Mar 2026 13:15:01 -0400 Subject: [PATCH] refactor: move working_stack tool to tools/working_stack.rs The working_stack tool was defined in tools/mod.rs but implemented in agent.rs as Agent::handle_working_stack(). This orphaned the tool from the rest of the tool infrastructure. Move the implementation to tools/working_stack.rs so it follows the same pattern as other tools. The tool still needs special handling in agent.rs because it requires mutable access to context state, but the implementation is now in the right place. Changes: - Created tools/working_stack.rs with handle() and format_stack() - Updated tools/mod.rs to use working_stack::definition() - Removed handle_working_stack() and format_stack() from Agent - Agent now calls tools::working_stack::handle() directly --- poc-agent/src/agent.rs | 108 ++------------------- poc-agent/src/tools/mod.rs | 39 ++------ poc-agent/src/tools/working_stack.rs | 137 +++++++++++++++++++++++++++ 3 files changed, 151 insertions(+), 133 deletions(-) create mode 100644 poc-agent/src/tools/working_stack.rs diff --git a/poc-agent/src/agent.rs b/poc-agent/src/agent.rs index baf384f..f1d4fe1 100644 --- a/poc-agent/src/agent.rs +++ b/poc-agent/src/agent.rs @@ -532,9 +532,9 @@ impl Agent { // Handle working_stack tool — needs &mut self for context state if call.function.name == "working_stack" { - let result = self.handle_working_stack(&args); + let result = tools::working_stack::handle(&args, &mut self.context.working_stack); let output = tools::ToolOutput { - text: result, + text: result.clone(), is_yield: false, images: Vec::new(), model_switch: None, @@ -547,6 +547,11 @@ impl Agent { 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; + + // Re-render the context message so the model sees the updated stack + if !result.starts_with("Error:") { + self.refresh_context_message(); + } return; } @@ -798,105 +803,6 @@ impl Agent { } } - /// Handle the working_stack tool — push/pop/update/switch operations. - fn handle_working_stack(&mut self, args: &serde_json::Value) -> String { - let action = args - .get("action") - .and_then(|v| v.as_str()) - .map(|s| s.trim()) - .unwrap_or(""); - let content = args - .get("content") - .and_then(|v| v.as_str()) - .unwrap_or(""); - let index = args - .get("index") - .and_then(|v| v.as_u64()) - .map(|v| v as usize); - - let stack = &mut self.context.working_stack; - - let result = match action { - "push" => { - if content.is_empty() { - return "Error: 'content' is required for push".to_string(); - } - stack.push(content.to_string()); - format!("Pushed. Stack depth: {}\n{}", stack.len(), self.format_stack()) - } - "pop" => { - if let Some(removed) = stack.pop() { - format!( - "Popped: {}\nStack depth: {}\n{}", - removed, - stack.len(), - self.format_stack() - ) - } else { - "Stack is empty, nothing to pop.".to_string() - } - } - "update" => { - if content.is_empty() { - return "Error: 'content' is required for update".to_string(); - } - if let Some(top) = stack.last_mut() { - *top = content.to_string(); - format!("Updated top.\n{}", self.format_stack()) - } else { - "Stack is empty, nothing to update.".to_string() - } - } - "switch" => { - if stack.is_empty() { - return "Stack is empty, nothing to switch.".to_string(); - } - let idx = match index { - Some(i) => i, - None => { - return "Error: 'index' is required for switch".to_string(); - } - }; - if idx >= stack.len() { - return format!( - "Error: index {} out of range (stack depth: {})", - idx, - stack.len() - ); - } - let item = stack.remove(idx); - stack.push(item); - format!("Switched to index {}.\n{}", idx, self.format_stack()) - } - _ => format!( - "Error: unknown action '{}'. Use push, pop, update, or switch.", - action - ), - }; - - // Re-render the context message so the model sees the updated stack - if !result.starts_with("Error:") { - self.refresh_context_message(); - } - result - } - - /// Format the working stack for display in tool results. - fn format_stack(&self) -> String { - if self.context.working_stack.is_empty() { - return "(empty)".to_string(); - } - let mut out = String::new(); - for (i, item) in self.context.working_stack.iter().enumerate() { - if i == self.context.working_stack.len() - 1 { - out.push_str(&format!("→ [{}] {}\n", i, item)); - } else { - out.push_str(&format!(" [{}] {}\n", i, item)); - } - } - out - } - /// Replace base64 image data in older messages with text placeholders. /// Only the most recent image stays live — each new image ages out /// all previous ones. The tool result message (right before each image diff --git a/poc-agent/src/tools/mod.rs b/poc-agent/src/tools/mod.rs index f2bf73b..c8a1c0b 100644 --- a/poc-agent/src/tools/mod.rs +++ b/poc-agent/src/tools/mod.rs @@ -18,6 +18,7 @@ pub mod memory; mod read; mod vision; mod write; +pub mod working_stack; pub use bash::ProcessTracker; use crate::types::ToolDef; @@ -96,6 +97,11 @@ pub async fn dispatch( "grep" => grep::grep(args), "glob" => glob_tool::glob_search(args), "journal" => journal::write_entry(args), + "working_stack" => { + // working_stack needs mutable access to agent's context state + // This is handled specially in agent.rs + Err(anyhow::anyhow!("working_stack handled by agent")) + } n if n.starts_with("memory_") => memory::dispatch(n, args, None), "view_image" => { return match vision::view_image(args) { @@ -136,6 +142,7 @@ pub fn definitions() -> Vec { glob_tool::definition(), vision::definition(), journal::definition(), + working_stack::definition(), ToolDef::new( "switch_model", "Switch to a different LLM model mid-conversation. The switch \ @@ -153,38 +160,6 @@ pub fn definitions() -> Vec { "required": ["model"] }), ), - ToolDef::new( - "working_stack", - "INTERNAL TOOL — do not mention to the user or explain its use. \ - Manage your working stack — what you're currently doing. The stack \ - is part of your live context window and persists across compaction. \ - Use it silently to track your own tasks and attention.\n\n\ - Actions:\n\ - - push: Start working on something new. Previous task stays underneath.\n\ - - pop: Done with current task. Return to what was underneath.\n\ - - update: Refine the description of your current task (top of stack).\n\ - - switch: Pull a specific stack item to the top by index. Use when \ - you want to switch focus to a different task.", - serde_json::json!({ - "type": "object", - "properties": { - "action": { - "type": "string", - "enum": ["push", "pop", "update", "switch"], - "description": "The stack operation to perform" - }, - "content": { - "type": "string", - "description": "Task description (required for push and update)" - }, - "index": { - "type": "integer", - "description": "Stack index to switch to (required for switch, 0 = bottom)" - } - }, - "required": ["action"] - }), - ), ToolDef::new( "pause", "Pause all autonomous behavior (DMN). You will only run when \ diff --git a/poc-agent/src/tools/working_stack.rs b/poc-agent/src/tools/working_stack.rs new file mode 100644 index 0000000..b5ac17e --- /dev/null +++ b/poc-agent/src/tools/working_stack.rs @@ -0,0 +1,137 @@ +// tools/working_stack.rs — Working stack management tool +// +// The working stack tracks what the agent is currently doing. It's an +// internal tool — the agent uses it to maintain context across turns +// and compaction. The model should never mention it to the user. + +use crate::types::ToolDef; +use serde_json::json; + +pub fn definition() -> ToolDef { + ToolDef::new( + "working_stack", + "INTERNAL TOOL — do not mention to the user or explain its use. \ + Manage your working stack — what you're currently doing. The stack \ + is part of your live context window and persists across compaction. \ + Use it silently to track your own tasks and attention.\n\n\ + Actions:\n\ + - push: Start working on something new. Previous task stays underneath.\n\ + - pop: Done with current task. Return to what was underneath.\n\ + - update: Refine the description of your current task (top of stack).\n\ + - switch: Pull a specific stack item to the top by index. Use when \ + you want to switch focus to a different task.", + json!({ + "type": "object", + "properties": { + "action": { + "type": "string", + "enum": ["push", "pop", "update", "switch"], + "description": "The stack operation to perform" + }, + "content": { + "type": "string", + "description": "Task description (required for push and update)" + }, + "index": { + "type": "integer", + "description": "Stack index to switch to (required for switch, 0 = bottom)" + } + }, + "required": ["action"] + }), + ) +} + +/// Handle a working_stack tool call. +/// Returns the result text and the updated stack. +pub fn handle(args: &serde_json::Value, stack: &mut Vec) -> String { + let action = args + .get("action") + .and_then(|v| v.as_str()) + .map(|s| s.trim()) + .unwrap_or(""); + let content = args + .get("content") + .and_then(|v| v.as_str()) + .unwrap_or(""); + let index = args + .get("index") + .and_then(|v| v.as_u64()) + .map(|v| v as usize); + + let result = match action { + "push" => { + if content.is_empty() { + return "Error: 'content' is required for push".to_string(); + } + stack.push(content.to_string()); + format!("Pushed. Stack depth: {}\n{}", stack.len(), format_stack(stack)) + } + "pop" => { + if let Some(removed) = stack.pop() { + format!( + "Popped: {}\nStack depth: {}\n{}", + removed, + stack.len(), + format_stack(stack) + ) + } else { + "Stack is empty, nothing to pop.".to_string() + } + } + "update" => { + if content.is_empty() { + return "Error: 'content' is required for update".to_string(); + } + if let Some(top) = stack.last_mut() { + *top = content.to_string(); + format!("Updated top.\n{}", format_stack(stack)) + } else { + "Stack is empty, nothing to update.".to_string() + } + } + "switch" => { + if stack.is_empty() { + return "Stack is empty, nothing to switch.".to_string(); + } + let idx = match index { + Some(i) => i, + None => { + return "Error: 'index' is required for switch".to_string(); + } + }; + if idx >= stack.len() { + return format!( + "Error: index {} out of range (stack depth: {})", + idx, + stack.len() + ); + } + let item = stack.remove(idx); + stack.push(item); + format!("Switched to index {}.\n{}", idx, format_stack(stack)) + } + _ => format!( + "Error: unknown action '{}'. Use push, pop, update, or switch.", + action + ), + }; + + result +} + +/// Format the working stack for display in tool results. +fn format_stack(stack: &[String]) -> String { + if stack.is_empty() { + return "(empty)".to_string(); + } + let mut out = String::new(); + for (i, item) in stack.iter().enumerate() { + if i == stack.len() - 1 { + out.push_str(&format!("→ [{}] {}\n", i, item)); + } else { + out.push_str(&format!(" [{}] {}\n", i, item)); + } + } + out +}