kent/consciousness
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: clean up tool dispatch and extract helpers

Browse source 
- Move working_stack tool to tools/working_stack.rs (was orphaned in agent.rs)
- Create control.rs for pause/switch_model/yield_to_user with Result<ToolOutput>
- Add ToolOutput::error() and ToolOutput::text() helper constructors
- Clean up dispatch() with Option<Result<ToolOutput>> pattern for rich tools
- Refactor memory.rs: extract cmd(), write_node(), supersede(), get_str(), get_f64()
- Merge run_rg() and run_grep() into unified run_search() in grep.rs
- Extract truncate_output() helper shared by bash, grep, glob tools

Net: -77 lines, better structure, less duplication
This commit is contained in:
Kent Overstreet 2026-03-21 15:18:53 -04:00
parent 3fd485a2e9
commit 45b7bba22a
6 changed files with 290 additions and 264 deletions
Download patch file Download diff file Expand all files Collapse all files

8
poc-agent/src/tools/bash.rs
View file

@ -168,13 +168,7 @@ pub async fn run_bash(args: &serde_json::Value, tracker: &ProcessTracker) -> Res
result = "(no output)".to_string();
}
const MAX_OUTPUT: usize = 30000;
if result.len() > MAX_OUTPUT {
result.truncate(MAX_OUTPUT);
result.push_str("\n... (output truncated)");
}
Ok(result)
Ok(super::truncate_output(result, 30000))
}
Ok(Err(e)) => {
Err(anyhow::anyhow!("Command failed: {}", e))

103
poc-agent/src/tools/control.rs Normal file
View file

@ -0,0 +1,103 @@
// tools/control.rs — Agent control tools
//
// Tools that affect agent control flow rather than performing work.
// These return Result<ToolOutput> to maintain consistency with other
// tools that can fail. The dispatch function handles error wrapping.
use anyhow::{Context, Result};
use super::ToolOutput;
use crate::types::ToolDef;
pub fn pause(_args: &serde_json::Value) -> Result<ToolOutput> {
Ok(ToolOutput {
text: "Pausing autonomous behavior. Only user input will wake you.".to_string(),
is_yield: true,
images: Vec::new(),
model_switch: None,
dmn_pause: true,
})
}
pub fn switch_model(args: &serde_json::Value) -> Result<ToolOutput> {
let model = args
.get("model")
.and_then(|v| v.as_str())
.context("'model' parameter is required")?;
if model.is_empty() {
anyhow::bail!("'model' parameter cannot be empty");
}
Ok(ToolOutput {
text: format!("Switching to model '{}' after this turn.", model),
is_yield: false,
images: Vec::new(),
model_switch: Some(model.to_string()),
dmn_pause: false,
})
}
pub fn yield_to_user(args: &serde_json::Value) -> Result<ToolOutput> {
let msg = args
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("Waiting for input.");
Ok(ToolOutput {
text: format!("Yielding. {}", msg),
is_yield: true,
images: Vec::new(),
model_switch: None,
dmn_pause: false,
})
}
pub fn definitions() -> Vec<ToolDef> {
vec![
ToolDef::new(
"switch_model",
"Switch to a different LLM model mid-conversation. The switch \
takes effect after the current turn completes. Use this when \
a task would benefit from a different model's strengths. \
Your memories and conversation history carry over.",
serde_json::json!({
"type": "object",
"properties": {
"model": {
"type": "string",
"description": "Name of the model to switch to (configured in config.json5)"
}
},
"required": ["model"]
}),
),
ToolDef::new(
"pause",
"Pause all autonomous behavior (DMN). You will only run when \
the user types something. Use this as a safety valve when \
you're stuck in a loop, confused, or want to fully stop. \
NOTE: only the user can unpause (Ctrl+P or /wake) — you \
cannot undo this yourself.",
serde_json::json!({
"type": "object",
"properties": {}
}),
),
ToolDef::new(
"yield_to_user",
"Signal that you want to wait for user input before continuing. \
Call this when you have a question for the user, when you've \
completed their request and want feedback, or when you genuinely \
want to pause. This is the ONLY way to enter a waiting state — \
without calling this tool, the agent loop will keep prompting you \
after a brief interval.",
serde_json::json!({
"type": "object",
"properties": {
"message": {
"type": "string",
"description": "Optional status message (e.g., 'Waiting for your thoughts on the design')"
}
}
}),
),
]
}

9
poc-agent/src/tools/glob_tool.rs
View file

@ -73,13 +73,6 @@ pub fn glob_search(args: &serde_json::Value) -> Result<String> {
output.push('\n');
}
// Truncate if too many
const MAX_OUTPUT: usize = 30000;
if output.len() > MAX_OUTPUT {
output.truncate(MAX_OUTPUT);
output.push_str("\n... (output truncated)");
}
output.push_str(&format!("\n({} files matched)", entries.len()));
Ok(output)
Ok(super::truncate_output(output, 30000))
}

55
poc-agent/src/tools/grep.rs
View file

@ -51,41 +51,39 @@ fn has_rg() -> bool {
}
pub fn grep(args: &serde_json::Value) -> Result<String> {
let pattern = args["pattern"].as_str().context("pattern is required")?;
let pattern = get_str(args, "pattern")?;
let path = args["path"].as_str().unwrap_or(".");
let file_glob = args["glob"].as_str();
let show_content = args["show_content"].as_bool().unwrap_or(false);
let context = args["context_lines"].as_u64();
let output = if has_rg() {
run_rg(pattern, path, file_glob, show_content, context)?
run_search("rg", pattern, path, file_glob, show_content, context, true)?
} else {
run_grep(pattern, path, file_glob, show_content, context)?
run_search("grep", pattern, path, file_glob, show_content, context, false)?
};
if output.is_empty() {
return Ok("No matches found.".to_string());
}
let mut result = output;
const MAX_OUTPUT: usize = 30000;
if result.len() > MAX_OUTPUT {
result.truncate(MAX_OUTPUT);
result.push_str("\n... (output truncated)");
Ok(super::truncate_output(output, 30000))
}
Ok(result)
}
fn run_rg(
/// Run a grep/rg search. Unified implementation for both tools.
fn run_search(
tool: &str,
pattern: &str,
path: &str,
file_glob: Option<&str>,
show_content: bool,
context: Option<u64>,
use_rg: bool,
) -> Result<String> {
let mut cmd = Command::new("rg");
let mut cmd = Command::new(tool);
if use_rg {
// ripgrep args
if show_content {
cmd.arg("-n");
if let Some(c) = context {
@ -94,26 +92,12 @@ fn run_rg(
} else {
cmd.arg("--files-with-matches");
}
if let Some(g) = file_glob {
cmd.arg("--glob").arg(g);
}
cmd.arg(pattern).arg(path);
let output = cmd.output().context("Failed to run rg")?;
Ok(String::from_utf8_lossy(&output.stdout).to_string())
}
fn run_grep(
pattern: &str,
path: &str,
file_glob: Option<&str>,
show_content: bool,
context: Option<u64>,
) -> Result<String> {
let mut cmd = Command::new("grep");
} else {
// grep args
cmd.arg("-r"); // recursive
if show_content {
cmd.arg("-n"); // line numbers
if let Some(c) = context {
@ -122,13 +106,20 @@ fn run_grep(
} else {
cmd.arg("-l"); // files-with-matches
}
if let Some(g) = file_glob {
cmd.arg("--include").arg(g);
}
cmd.arg("-E"); // extended regex
}
cmd.arg(pattern).arg(path);
let output = cmd.output().context("Failed to run grep")?;
let output = cmd.output().with_context(|| format!("Failed to run {}", tool))?;
Ok(String::from_utf8_lossy(&output.stdout).to_string())
}
/// Helper: get required string argument.
fn get_str<'a>(args: &'a serde_json::Value, name: &'a str) -> Result<&'a str> {
args.get(name)
.and_then(|v| v.as_str())
.context(format!("{} is required", name))
}

166
poc-agent/src/tools/memory.rs
View file

@ -6,7 +6,8 @@
use anyhow::{Context, Result};
use serde_json::json;
use std::process::Command;
use std::io::Write;
use std::process::{Command, Stdio};
use crate::types::ToolDef;
@ -177,106 +178,58 @@ pub fn definitions() -> Vec<ToolDef> {
/// Dispatch a memory tool call. Shells out to poc-memory CLI.
pub fn dispatch(name: &str, args: &serde_json::Value, provenance: Option<&str>) -> Result<String> {
match name {
let result = match name {
"memory_render" => {
let key = args["key"].as_str().context("key is required")?;
run_poc_memory(&["render", key], provenance)
let key = get_str(args, "key")?;
cmd(&["render", key], provenance)?
}
"memory_write" => {
let key = args["key"].as_str().context("key is required")?;
let content = args["content"].as_str().context("content is required")?;
let mut cmd = Command::new("poc-memory");
cmd.args(["write", key])
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped());
if let Some(prov) = provenance {
cmd.env("POC_PROVENANCE", prov);
}
let mut child = cmd.spawn()
.context("spawn poc-memory write")?;
use std::io::Write;
child.stdin.take().unwrap().write_all(content.as_bytes())
.context("write content to stdin")?;
let output = child.wait_with_output().context("wait poc-memory write")?;
Ok(String::from_utf8_lossy(&output.stdout).to_string()
+ &String::from_utf8_lossy(&output.stderr))
let key = get_str(args, "key")?;
let content = get_str(args, "content")?;
write_node(key, content, provenance)?
}
"memory_search" => {
let query = args["query"].as_str().context("query is required")?;
run_poc_memory(&["search", query], provenance)
let query = get_str(args, "query")?;
cmd(&["search", query], provenance)?
}
"memory_links" => {
let key = args["key"].as_str().context("key is required")?;
run_poc_memory(&["graph", "link", key], provenance)
let key = get_str(args, "key")?;
cmd(&["graph", "link", key], provenance)?
}
"memory_link_set" => {
let source = args["source"].as_str().context("source is required")?;
let target = args["target"].as_str().context("target is required")?;
let strength = args["strength"].as_f64().context("strength is required")?;
run_poc_memory(&["graph", "link-set", source, target, &format!("{:.2}", strength)], provenance)
let source = get_str(args, "source")?;
let target = get_str(args, "target")?;
let strength = get_f64(args, "strength")?;
cmd(&["graph", "link-set", source, target, &format!("{:.2}", strength)], provenance)?
}
"memory_link_add" => {
let source = args["source"].as_str().context("source is required")?;
let target = args["target"].as_str().context("target is required")?;
run_poc_memory(&["graph", "link-add", source, target], provenance)
let source = get_str(args, "source")?;
let target = get_str(args, "target")?;
cmd(&["graph", "link-add", source, target], provenance)?
}
"memory_used" => {
let key = args["key"].as_str().context("key is required")?;
run_poc_memory(&["used", key], provenance)
let key = get_str(args, "key")?;
cmd(&["used", key], provenance)?
}
"memory_weight_set" => {
let key = args["key"].as_str().context("key is required")?;
let weight = args["weight"].as_f64().context("weight is required")?;
run_poc_memory(&["weight-set", key, &format!("{:.2}", weight)], provenance)
}
"memory_supersede" => {
let old_key = args["old_key"].as_str().context("old_key is required")?;
let new_key = args["new_key"].as_str().context("new_key is required")?;
let reason = args.get("reason").and_then(|v| v.as_str()).unwrap_or("superseded");
// Read old node, prepend superseded notice, write back, set weight to 0.01
let old_content = run_poc_memory(&["render", old_key], provenance).unwrap_or_default();
// Strip the links section from render output
let content_only = old_content.split("\n\n---\nLinks:").next().unwrap_or(&old_content);
let notice = format!(
"**SUPERSEDED** by `{}` — {}\n\nOriginal content preserved below for reference.\n\n---\n\n{}",
new_key, reason, content_only.trim()
);
let mut cmd = Command::new("poc-memory");
cmd.args(["write", old_key])
.stdin(std::process::Stdio::piped())
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped());
if let Some(prov) = provenance {
cmd.env("POC_PROVENANCE", prov);
}
let mut child = cmd.spawn()
.context("spawn poc-memory write")?;
use std::io::Write;
child.stdin.take().unwrap().write_all(notice.as_bytes())
.context("write supersede notice")?;
let output = child.wait_with_output().context("wait poc-memory write")?;
let write_result = String::from_utf8_lossy(&output.stdout).to_string();
// Set weight to 0.01
let weight_result = run_poc_memory(&["weight-set", old_key, "0.01"], provenance)
.unwrap_or_else(|e| format!("weight-set failed: {}", e));
Ok(format!("{}\n{}", write_result.trim(), weight_result.trim()))
}
_ => Err(anyhow::anyhow!("Unknown memory tool: {}", name)),
let key = get_str(args, "key")?;
let weight = get_f64(args, "weight")?;
cmd(&["weight-set", key, &format!("{:.2}", weight)], provenance)?
}
"memory_supersede" => supersede(args, provenance)?,
_ => anyhow::bail!("Unknown memory tool: {}", name),
};
Ok(result)
}
fn run_poc_memory(args: &[&str], provenance: Option<&str>) -> Result<String> {
/// Run poc-memory command and return stdout.
fn cmd(args: &[&str], provenance: Option<&str>) -> Result<String> {
let mut cmd = Command::new("poc-memory");
cmd.args(args);
if let Some(prov) = provenance {
cmd.env("POC_PROVENANCE", prov);
}
let output = cmd.output()
.context("run poc-memory")?;
let output = cmd.output().context("run poc-memory")?;
let stdout = String::from_utf8_lossy(&output.stdout);
let stderr = String::from_utf8_lossy(&output.stderr);
if output.status.success() {
@ -285,3 +238,60 @@ fn run_poc_memory(args: &[&str], provenance: Option<&str>) -> Result<String> {
Ok(format!("{}{}", stdout, stderr))
}
}
/// Write content to a node via stdin.
fn write_node(key: &str, content: &str, provenance: Option<&str>) -> Result<String> {
let mut cmd = Command::new("poc-memory");
cmd.args(["write", key])
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
if let Some(prov) = provenance {
cmd.env("POC_PROVENANCE", prov);
}
let mut child = cmd.spawn().context("spawn poc-memory write")?;
child.stdin.take().unwrap().write_all(content.as_bytes())
.context("write content to stdin")?;
let output = child.wait_with_output().context("wait poc-memory write")?;
Ok(String::from_utf8_lossy(&output.stdout).to_string()
+ &String::from_utf8_lossy(&output.stderr))
}
/// Handle memory_supersede - reads old node, prepends notice, writes back, sets weight.
fn supersede(args: &serde_json::Value, provenance: Option<&str>) -> Result<String> {
let old_key = get_str(args, "old_key")?;
let new_key = get_str(args, "new_key")?;
let reason = args.get("reason").and_then(|v| v.as_str()).unwrap_or("superseded");
// Read old node
let old_content = cmd(&["render", old_key], provenance)?;
let content_only = old_content.split("\n\n---\nLinks:").next().unwrap_or(&old_content);
// Prepend superseded notice
let notice = format!(
"**SUPERSEDED** by `{}` — {}\n\nOriginal content preserved below for reference.\n\n---\n\n{}",
new_key, reason, content_only.trim()
);
// Write back
let write_result = write_node(old_key, &notice, provenance)?;
// Set weight to 0.01
let weight_result = cmd(&["weight-set", old_key, "0.01"], provenance)?;
Ok(format!("{}\n{}", write_result.trim(), weight_result.trim()))
}
/// Helper: get required string argument.
fn get_str<'a>(args: &'a serde_json::Value, name: &'a str) -> Result<&'a str> {
args.get(name)
.and_then(|v| v.as_str())
.context(format!("{} is required", name))
}
/// Helper: get required f64 argument.
fn get_f64(args: &serde_json::Value, name: &str) -> Result<f64> {
args.get(name)
.and_then(|v| v.as_f64())
.context(format!("{} is required", name))
}

175
poc-agent/src/tools/mod.rs
View file

@ -10,6 +10,7 @@
// immediately from an async fn.
mod bash;
mod control;
mod edit;
mod glob_tool;
mod grep;
@ -35,60 +36,64 @@ pub struct ToolOutput {
pub dmn_pause: bool,
}
/// Dispatch a tool call by name, returning the result as a string.
/// Returns (output, is_yield) — is_yield is true only for yield_to_user.
impl ToolOutput {
fn error(e: impl std::fmt::Display) -> Self {
Self {
text: format!("Error: {}", e),
is_yield: false,
images: Vec::new(),
model_switch: None,
dmn_pause: false,
}
}
fn text(s: String) -> Self {
Self {
text: s,
is_yield: false,
images: Vec::new(),
model_switch: None,
dmn_pause: false,
}
}
}
/// Truncate output if it exceeds max length, appending a truncation notice.
/// Used by tools that can produce large amounts of output (bash, grep, glob, etc).
pub fn truncate_output(mut s: String, max: usize) -> String {
if s.len() > max {
s.truncate(max);
s.push_str("\n... (output truncated)");
}
s
}
/// Dispatch a tool call by name.
///
/// Control tools (pause, switch_model, yield_to_user) and view_image
/// return Result<ToolOutput>. Regular tools return Result<String> and
/// get wrapped in a text-only ToolOutput.
///
/// Note: working_stack is handled in agent.rs before reaching this
/// function (it needs mutable context access).
pub async fn dispatch(
name: &str,
args: &serde_json::Value,
tracker: &ProcessTracker,
) -> ToolOutput {
if name == "pause" {
return ToolOutput {
text: "Pausing autonomous behavior. Only user input will wake you.".to_string(),
is_yield: true,
images: Vec::new(),
model_switch: None,
dmn_pause: true,
};
}
if name == "switch_model" {
let model = args
.get("model")
.and_then(|v| v.as_str())
.unwrap_or("");
if model.is_empty() {
return ToolOutput {
text: "Error: 'model' parameter is required".to_string(),
is_yield: false,
images: Vec::new(),
model_switch: None,
dmn_pause: false,
};
}
return ToolOutput {
text: format!("Switching to model '{}' after this turn.", model),
is_yield: false,
images: Vec::new(),
model_switch: Some(model.to_string()),
dmn_pause: false,
};
}
if name == "yield_to_user" {
let msg = args
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("Waiting for input.");
return ToolOutput {
text: format!("Yielding. {}", msg),
is_yield: true,
images: Vec::new(),
model_switch: None,
dmn_pause: false,
// Tools that return Result<ToolOutput> directly
let rich_result = match name {
"pause" => Some(control::pause(args)),
"switch_model" => Some(control::switch_model(args)),
"yield_to_user" => Some(control::yield_to_user(args)),
"view_image" => Some(vision::view_image(args)),
_ => None,
};
if let Some(result) = rich_result {
return result.unwrap_or_else(ToolOutput::error);
}
// Regular tools — return Result<String>
let result = match name {
"read_file" => read::read_file(args),
"write_file" => write::write_file(args),
@ -97,37 +102,13 @@ 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) {
Ok(output) => output,
Err(e) => ToolOutput {
text: format!("Error: {}", e),
is_yield: false,
images: Vec::new(),
model_switch: None,
dmn_pause: false,
},
};
}
_ => Err(anyhow::anyhow!("Unknown tool: {}", name)),
};
let text = match result {
Ok(output) => output,
Err(e) => format!("Error: {}", e),
};
ToolOutput {
text,
is_yield: false,
images: Vec::new(),
model_switch: None,
dmn_pause: false,
match result {
Ok(s) => ToolOutput::text(s),
Err(e) => ToolOutput::error(e),
}
}
@ -143,54 +124,8 @@ pub fn definitions() -> Vec<ToolDef> {
vision::definition(),
journal::definition(),
working_stack::definition(),
ToolDef::new(
"switch_model",
"Switch to a different LLM model mid-conversation. The switch \
takes effect after the current turn completes. Use this when \
a task would benefit from a different model's strengths. \
Your memories and conversation history carry over.",
serde_json::json!({
"type": "object",
"properties": {
"model": {
"type": "string",
"description": "Name of the model to switch to (configured in config.json5)"
}
},
"required": ["model"]
}),
),
ToolDef::new(
"pause",
"Pause all autonomous behavior (DMN). You will only run when \
the user types something. Use this as a safety valve when \
you're stuck in a loop, confused, or want to fully stop. \
NOTE: only the user can unpause (Ctrl+P or /wake) — you \
cannot undo this yourself.",
serde_json::json!({
"type": "object",
"properties": {}
}),
),
ToolDef::new(
"yield_to_user",
"Signal that you want to wait for user input before continuing. \
Call this when you have a question for the user, when you've \
completed their request and want feedback, or when you genuinely \
want to pause. This is the ONLY way to enter a waiting state — \
without calling this tool, the agent loop will keep prompting you \
after a brief interval.",
serde_json::json!({
"type": "object",
"properties": {
"message": {
"type": "string",
"description": "Optional status message (e.g., 'Waiting for your thoughts on the design')"
}
}
}),
),
].into_iter()
.chain(control::definitions())
.chain(memory::definitions())
.collect()
}