Add tool use to the orchestrator (#4)

Add tool use without sandboxing.

Currently available tools are list dir, read file, write file and exec bash.

Reviewed-on: #4
Co-authored-by: Drew Galbraith <drew@tiramisu.one>
Co-committed-by: Drew Galbraith <drew@tiramisu.one>
This commit is contained in:
Drew 2026-03-02 03:00:13 +00:00 committed by Drew
parent 6b85ff3cb8
commit 797d7564b7
20 changed files with 1822 additions and 129 deletions

220
src/tools/mod.rs Normal file
View file

@ -0,0 +1,220 @@
//! Tool system: trait, registry, risk classification, and built-in tools.
//!
//! All tools implement the [`Tool`] trait. The [`ToolRegistry`] collects them
//! and provides lookup by name plus generation of [`ToolDefinition`]s for the
//! model provider.
mod list_directory;
mod read_file;
mod shell_exec;
mod write_file;
use std::path::{Path, PathBuf};
use async_trait::async_trait;
use crate::core::types::ToolDefinition;
/// The output of a tool execution.
#[derive(Debug)]
pub struct ToolOutput {
/// The text content returned to the model.
pub content: String,
/// Whether the tool encountered an error.
pub is_error: bool,
}
/// Risk classification for tool approval gating.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RiskLevel {
/// Safe to execute without user confirmation (e.g. read-only operations).
AutoApprove,
/// Requires explicit user approval before execution (e.g. writes, shell).
RequiresApproval,
}
/// A tool that the model can invoke.
///
/// The `execute` method is async so that tool implementations can use
/// `tokio::fs` and `tokio::process` without blocking a Tokio worker thread.
/// `#[async_trait]` desugars the async fn to a boxed future, which is required
/// for `dyn Tool` to remain object-safe.
#[async_trait]
pub trait Tool: Send + Sync {
/// The name the model uses to invoke this tool.
fn name(&self) -> &str;
/// Human-readable description for the model.
fn description(&self) -> &str;
/// JSON Schema for the tool's input parameters.
fn input_schema(&self) -> serde_json::Value;
/// The risk level of this tool.
fn risk_level(&self) -> RiskLevel;
/// Execute the tool with the given input, confined to `working_dir`.
async fn execute(
&self,
input: &serde_json::Value,
working_dir: &Path,
) -> Result<ToolOutput, ToolError>;
}
/// Errors from tool execution.
#[derive(Debug, thiserror::Error)]
pub enum ToolError {
/// The requested path escapes the working directory.
#[error("path escapes working directory: {0}")]
PathEscape(PathBuf),
/// An I/O error during tool execution.
#[error("I/O error: {0}")]
Io(#[from] std::io::Error),
/// A required input field is missing or has the wrong type.
#[error("invalid input: {0}")]
InvalidInput(String),
}
/// Validate that `requested` resolves to a path inside `working_dir`.
///
/// Joins `working_dir` with `requested`, canonicalizes the result (resolving
/// symlinks and `..` components), and checks that the canonical path starts
/// with the canonical working directory.
///
/// Returns the canonical path on success, or [`ToolError::PathEscape`] if the
/// path would escape the working directory.
pub fn validate_path(working_dir: &Path, requested: &str) -> Result<PathBuf, ToolError> {
let candidate = if Path::new(requested).is_absolute() {
PathBuf::from(requested)
} else {
working_dir.join(requested)
};
// For paths that don't exist yet (e.g. write_file creating a new file),
// canonicalize the parent directory and append the filename.
let canonical = if candidate.exists() {
candidate
.canonicalize()
.map_err(|_| ToolError::PathEscape(candidate.clone()))?
} else {
let parent = candidate
.parent()
.ok_or_else(|| ToolError::PathEscape(candidate.clone()))?;
let file_name = candidate
.file_name()
.ok_or_else(|| ToolError::PathEscape(candidate.clone()))?;
let canonical_parent = parent
.canonicalize()
.map_err(|_| ToolError::PathEscape(candidate.clone()))?;
canonical_parent.join(file_name)
};
let canonical_root = working_dir
.canonicalize()
.map_err(|_| ToolError::PathEscape(candidate.clone()))?;
if canonical.starts_with(&canonical_root) {
Ok(canonical)
} else {
Err(ToolError::PathEscape(candidate))
}
}
/// Collection of available tools with name-based lookup.
pub struct ToolRegistry {
tools: Vec<Box<dyn Tool>>,
}
impl ToolRegistry {
/// Create an empty registry with no tools.
#[allow(dead_code)]
pub fn empty() -> Self {
Self { tools: Vec::new() }
}
/// Create a registry with the default built-in tools.
pub fn default_tools() -> Self {
Self {
tools: vec![
Box::new(read_file::ReadFile),
Box::new(list_directory::ListDirectory),
Box::new(write_file::WriteFile),
Box::new(shell_exec::ShellExec),
],
}
}
/// Look up a tool by name.
pub fn get(&self, name: &str) -> Option<&dyn Tool> {
self.tools.iter().find(|t| t.name() == name).map(|t| &**t)
}
/// Generate [`ToolDefinition`]s for the model provider.
pub fn definitions(&self) -> Vec<ToolDefinition> {
self.tools
.iter()
.map(|t| ToolDefinition {
name: t.name().to_string(),
description: t.description().to_string(),
input_schema: t.input_schema(),
})
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
use tempfile::TempDir;
#[test]
fn validate_path_allows_subpath() {
let dir = TempDir::new().unwrap();
let sub = dir.path().join("sub");
fs::create_dir(&sub).unwrap();
let result = validate_path(dir.path(), "sub");
assert!(result.is_ok());
assert!(
result
.unwrap()
.starts_with(dir.path().canonicalize().unwrap())
);
}
#[test]
fn validate_path_rejects_traversal() {
let dir = TempDir::new().unwrap();
let result = validate_path(dir.path(), "../../../etc/passwd");
assert!(result.is_err());
assert!(matches!(result, Err(ToolError::PathEscape(_))));
}
#[test]
fn validate_path_rejects_absolute_outside() {
let dir = TempDir::new().unwrap();
let result = validate_path(dir.path(), "/etc/passwd");
assert!(result.is_err());
}
#[test]
fn validate_path_allows_new_file_in_working_dir() {
let dir = TempDir::new().unwrap();
let result = validate_path(dir.path(), "new_file.txt");
assert!(result.is_ok());
}
#[test]
fn registry_default_has_all_tools() {
let reg = ToolRegistry::default_tools();
assert!(reg.get("read_file").is_some());
assert!(reg.get("list_directory").is_some());
assert!(reg.get("write_file").is_some());
assert!(reg.get("shell_exec").is_some());
assert!(reg.get("nonexistent").is_none());
}
#[test]
fn registry_definitions_match_tools() {
let reg = ToolRegistry::default_tools();
let defs = reg.definitions();
assert_eq!(defs.len(), 4);
assert!(defs.iter().any(|d| d.name == "read_file"));
}
}