407 lines
13 KiB
Rust
407 lines
13 KiB
Rust
|
//! Built-in tools: shell, list_dir, read_file, write_file.
|
||
|
|
|
||
|
|
use anyhow::{Context, Result, anyhow};
|
||
|
|
use serde_json::{Value, json};
|
||
|
|
use std::path::{Component, Path, PathBuf};
|
||
|
|
use std::process::Stdio;
|
||
|
|
use std::time::Duration;
|
||
|
|
use tokio::process::Command;
|
||
|
|
|
||
|
|
pub const MAX_READ_BYTES: usize = 256 * 1024;
|
||
|
|
pub const SHELL_TIMEOUT_SECS: u64 = 30;
|
||
|
|
|
||
|
|
/// Runtime context for tool execution.
|
||
|
|
#[derive(Debug, Clone)]
|
||
|
|
pub struct ToolContext {
|
||
|
|
/// Default working directory for relative paths / shell.
|
||
|
|
pub cwd: PathBuf,
|
||
|
|
/// Optional workspace root; paths outside it are rejected when set.
|
||
|
|
pub workspace_root: Option<PathBuf>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl ToolContext {
|
||
|
|
pub fn new(cwd: impl Into<PathBuf>) -> Self {
|
||
|
|
let cwd = cwd.into();
|
||
|
|
Self {
|
||
|
|
cwd: cwd.clone(),
|
||
|
|
workspace_root: Some(cwd),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn with_workspace(mut self, root: Option<PathBuf>) -> Self {
|
||
|
|
self.workspace_root = root;
|
||
|
|
self
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// OpenAI-compatible tool definitions for chat completions.
|
||
|
|
pub fn tool_definitions() -> Value {
|
||
|
|
json!([
|
||
|
|
{
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": "shell",
|
||
|
|
"description": "Run a shell command. Captures stdout, stderr, and exit code. Timeout 30s. cwd defaults to the session working directory.",
|
||
|
|
"parameters": {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"command": { "type": "string", "description": "Shell command to run" },
|
||
|
|
"cwd": { "type": "string", "description": "Optional working directory" }
|
||
|
|
},
|
||
|
|
"required": ["command"]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": "list_dir",
|
||
|
|
"description": "List entries in a directory (names only, sorted).",
|
||
|
|
"parameters": {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"path": { "type": "string", "description": "Directory path (relative to cwd or absolute)" }
|
||
|
|
},
|
||
|
|
"required": ["path"]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": "read_file",
|
||
|
|
"description": "Read a text file (max 256KB). Returns contents as UTF-8 (lossy).",
|
||
|
|
"parameters": {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"path": { "type": "string", "description": "File path" }
|
||
|
|
},
|
||
|
|
"required": ["path"]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": "write_file",
|
||
|
|
"description": "Write text to a file, creating parent directories as needed.",
|
||
|
|
"parameters": {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"path": { "type": "string", "description": "File path" },
|
||
|
|
"content": { "type": "string", "description": "File contents" }
|
||
|
|
},
|
||
|
|
"required": ["path", "content"]
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
])
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Resolve a user-supplied path against cwd and optionally enforce workspace_root.
|
||
|
|
pub fn resolve_path(ctx: &ToolContext, path: &str) -> Result<PathBuf> {
|
||
|
|
let raw = Path::new(path);
|
||
|
|
let joined = if raw.is_absolute() {
|
||
|
|
raw.to_path_buf()
|
||
|
|
} else {
|
||
|
|
ctx.cwd.join(raw)
|
||
|
|
};
|
||
|
|
|
||
|
|
// Normalize without requiring the path to exist (for write_file parents).
|
||
|
|
let resolved = normalize_path(&joined);
|
||
|
|
|
||
|
|
if let Some(root) = &ctx.workspace_root {
|
||
|
|
let root_norm = normalize_path(root);
|
||
|
|
if !resolved.starts_with(&root_norm) {
|
||
|
|
return Err(anyhow!(
|
||
|
|
"path {:?} is outside workspace root {:?}",
|
||
|
|
resolved,
|
||
|
|
root_norm
|
||
|
|
));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
Ok(resolved)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn normalize_path(path: &Path) -> PathBuf {
|
||
|
|
let mut out = PathBuf::new();
|
||
|
|
for comp in path.components() {
|
||
|
|
match comp {
|
||
|
|
Component::Prefix(p) => out.push(p.as_os_str()),
|
||
|
|
Component::RootDir => out.push(Component::RootDir.as_os_str()),
|
||
|
|
Component::CurDir => {}
|
||
|
|
Component::ParentDir => {
|
||
|
|
out.pop();
|
||
|
|
}
|
||
|
|
Component::Normal(c) => out.push(c),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
out
|
||
|
|
}
|
||
|
|
|
||
|
|
pub async fn execute_tool(ctx: &ToolContext, name: &str, arguments_json: &str) -> String {
|
||
|
|
match execute_tool_inner(ctx, name, arguments_json).await {
|
||
|
|
Ok(v) => v.to_string(),
|
||
|
|
Err(e) => json!({ "error": format!("{e:#}") }).to_string(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn execute_tool_inner(ctx: &ToolContext, name: &str, arguments_json: &str) -> Result<Value> {
|
||
|
|
let args: Value = serde_json::from_str(arguments_json)
|
||
|
|
.with_context(|| format!("invalid tool arguments JSON for {name}"))?;
|
||
|
|
|
||
|
|
match name {
|
||
|
|
"shell" => tool_shell(ctx, &args).await,
|
||
|
|
"list_dir" => tool_list_dir(ctx, &args).await,
|
||
|
|
"read_file" => tool_read_file(ctx, &args).await,
|
||
|
|
"write_file" => tool_write_file(ctx, &args).await,
|
||
|
|
other => Err(anyhow!("unknown tool: {other}")),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn tool_shell(ctx: &ToolContext, args: &Value) -> Result<Value> {
|
||
|
|
let command = args
|
||
|
|
.get("command")
|
||
|
|
.and_then(|v| v.as_str())
|
||
|
|
.ok_or_else(|| anyhow!("shell: missing 'command'"))?;
|
||
|
|
|
||
|
|
let cwd = if let Some(c) = args.get("cwd").and_then(|v| v.as_str()) {
|
||
|
|
resolve_path(ctx, c)?
|
||
|
|
} else {
|
||
|
|
ctx.cwd.clone()
|
||
|
|
};
|
||
|
|
|
||
|
|
if !cwd.is_dir() {
|
||
|
|
return Err(anyhow!("shell cwd is not a directory: {}", cwd.display()));
|
||
|
|
}
|
||
|
|
|
||
|
|
let child = Command::new("sh")
|
||
|
|
.arg("-c")
|
||
|
|
.arg(command)
|
||
|
|
.current_dir(&cwd)
|
||
|
|
.stdout(Stdio::piped())
|
||
|
|
.stderr(Stdio::piped())
|
||
|
|
.kill_on_drop(true)
|
||
|
|
.spawn()
|
||
|
|
.with_context(|| format!("failed to spawn shell for: {command}"))?;
|
||
|
|
|
||
|
|
let timeout = Duration::from_secs(SHELL_TIMEOUT_SECS);
|
||
|
|
let output = match tokio::time::timeout(timeout, child.wait_with_output()).await {
|
||
|
|
Ok(Ok(out)) => out,
|
||
|
|
Ok(Err(e)) => return Err(anyhow!("shell wait failed: {e}")),
|
||
|
|
Err(_) => {
|
||
|
|
return Ok(json!({
|
||
|
|
"error": format!("command timed out after {SHELL_TIMEOUT_SECS}s"),
|
||
|
|
"timed_out": true,
|
||
|
|
"command": command,
|
||
|
|
}));
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
let stdout = String::from_utf8_lossy(&output.stdout).to_string();
|
||
|
|
let stderr = String::from_utf8_lossy(&output.stderr).to_string();
|
||
|
|
let exit_code = output.status.code().unwrap_or(-1);
|
||
|
|
|
||
|
|
Ok(json!({
|
||
|
|
"command": command,
|
||
|
|
"cwd": cwd.display().to_string(),
|
||
|
|
"exit_code": exit_code,
|
||
|
|
"stdout": truncate_output(&stdout, 64 * 1024),
|
||
|
|
"stderr": truncate_output(&stderr, 32 * 1024),
|
||
|
|
}))
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn tool_list_dir(ctx: &ToolContext, args: &Value) -> Result<Value> {
|
||
|
|
let path = args
|
||
|
|
.get("path")
|
||
|
|
.and_then(|v| v.as_str())
|
||
|
|
.ok_or_else(|| anyhow!("list_dir: missing 'path'"))?;
|
||
|
|
let dir = resolve_path(ctx, path)?;
|
||
|
|
let mut rd = tokio::fs::read_dir(&dir)
|
||
|
|
.await
|
||
|
|
.with_context(|| format!("list_dir: {}", dir.display()))?;
|
||
|
|
|
||
|
|
let mut entries = Vec::new();
|
||
|
|
while let Some(ent) = rd.next_entry().await? {
|
||
|
|
let name = ent.file_name().to_string_lossy().to_string();
|
||
|
|
let file_type = ent.file_type().await?;
|
||
|
|
let kind = if file_type.is_dir() {
|
||
|
|
"dir"
|
||
|
|
} else if file_type.is_symlink() {
|
||
|
|
"symlink"
|
||
|
|
} else {
|
||
|
|
"file"
|
||
|
|
};
|
||
|
|
entries.push(json!({ "name": name, "kind": kind }));
|
||
|
|
}
|
||
|
|
entries.sort_by(|a, b| {
|
||
|
|
a["name"]
|
||
|
|
.as_str()
|
||
|
|
.unwrap_or("")
|
||
|
|
.cmp(b["name"].as_str().unwrap_or(""))
|
||
|
|
});
|
||
|
|
|
||
|
|
Ok(json!({
|
||
|
|
"path": dir.display().to_string(),
|
||
|
|
"entries": entries,
|
||
|
|
}))
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn tool_read_file(ctx: &ToolContext, args: &Value) -> Result<Value> {
|
||
|
|
let path = args
|
||
|
|
.get("path")
|
||
|
|
.and_then(|v| v.as_str())
|
||
|
|
.ok_or_else(|| anyhow!("read_file: missing 'path'"))?;
|
||
|
|
let file = resolve_path(ctx, path)?;
|
||
|
|
let meta = tokio::fs::metadata(&file)
|
||
|
|
.await
|
||
|
|
.with_context(|| format!("read_file: {}", file.display()))?;
|
||
|
|
if !meta.is_file() {
|
||
|
|
return Err(anyhow!("read_file: not a file: {}", file.display()));
|
||
|
|
}
|
||
|
|
if meta.len() as usize > MAX_READ_BYTES {
|
||
|
|
return Err(anyhow!(
|
||
|
|
"read_file: file too large ({} bytes > {} cap)",
|
||
|
|
meta.len(),
|
||
|
|
MAX_READ_BYTES
|
||
|
|
));
|
||
|
|
}
|
||
|
|
let bytes = tokio::fs::read(&file)
|
||
|
|
.await
|
||
|
|
.with_context(|| format!("read_file: {}", file.display()))?;
|
||
|
|
let content = String::from_utf8_lossy(&bytes).to_string();
|
||
|
|
Ok(json!({
|
||
|
|
"path": file.display().to_string(),
|
||
|
|
"bytes": bytes.len(),
|
||
|
|
"content": content,
|
||
|
|
}))
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn tool_write_file(ctx: &ToolContext, args: &Value) -> Result<Value> {
|
||
|
|
let path = args
|
||
|
|
.get("path")
|
||
|
|
.and_then(|v| v.as_str())
|
||
|
|
.ok_or_else(|| anyhow!("write_file: missing 'path'"))?;
|
||
|
|
let content = args
|
||
|
|
.get("content")
|
||
|
|
.and_then(|v| v.as_str())
|
||
|
|
.ok_or_else(|| anyhow!("write_file: missing 'content'"))?;
|
||
|
|
let file = resolve_path(ctx, path)?;
|
||
|
|
if let Some(parent) = file.parent() {
|
||
|
|
tokio::fs::create_dir_all(parent)
|
||
|
|
.await
|
||
|
|
.with_context(|| format!("write_file create_dir_all: {}", parent.display()))?;
|
||
|
|
}
|
||
|
|
tokio::fs::write(&file, content.as_bytes())
|
||
|
|
.await
|
||
|
|
.with_context(|| format!("write_file: {}", file.display()))?;
|
||
|
|
Ok(json!({
|
||
|
|
"path": file.display().to_string(),
|
||
|
|
"bytes_written": content.len(),
|
||
|
|
}))
|
||
|
|
}
|
||
|
|
|
||
|
|
fn truncate_output(s: &str, max: usize) -> String {
|
||
|
|
if s.len() <= max {
|
||
|
|
s.to_string()
|
||
|
|
} else {
|
||
|
|
format!(
|
||
|
|
"{}…\n[truncated {} bytes]",
|
||
|
|
&s[..max],
|
||
|
|
s.len().saturating_sub(max)
|
||
|
|
)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(test)]
|
||
|
|
mod tests {
|
||
|
|
use super::*;
|
||
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||
|
|
|
||
|
|
fn temp_ctx() -> (ToolContext, PathBuf) {
|
||
|
|
let stamp = SystemTime::now()
|
||
|
|
.duration_since(UNIX_EPOCH)
|
||
|
|
.unwrap()
|
||
|
|
.as_nanos();
|
||
|
|
let dir = std::env::temp_dir().join(format!("grokboy-tools-{stamp}"));
|
||
|
|
std::fs::create_dir_all(&dir).unwrap();
|
||
|
|
(ToolContext::new(dir.clone()), dir)
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn write_read_list_shell() {
|
||
|
|
let (ctx, dir) = temp_ctx();
|
||
|
|
|
||
|
|
let w = execute_tool(
|
||
|
|
&ctx,
|
||
|
|
"write_file",
|
||
|
|
&json!({"path": "hello.txt", "content": "你好 GrokBoy"}).to_string(),
|
||
|
|
)
|
||
|
|
.await;
|
||
|
|
let w: Value = serde_json::from_str(&w).unwrap();
|
||
|
|
assert!(w.get("error").is_none(), "{w}");
|
||
|
|
assert_eq!(w["bytes_written"], "你好 GrokBoy".len());
|
||
|
|
|
||
|
|
let r = execute_tool(
|
||
|
|
&ctx,
|
||
|
|
"read_file",
|
||
|
|
&json!({"path": "hello.txt"}).to_string(),
|
||
|
|
)
|
||
|
|
.await;
|
||
|
|
let r: Value = serde_json::from_str(&r).unwrap();
|
||
|
|
assert_eq!(r["content"], "你好 GrokBoy");
|
||
|
|
|
||
|
|
let l = execute_tool(&ctx, "list_dir", &json!({"path": "."}).to_string()).await;
|
||
|
|
let l: Value = serde_json::from_str(&l).unwrap();
|
||
|
|
let names: Vec<&str> = l["entries"]
|
||
|
|
.as_array()
|
||
|
|
.unwrap()
|
||
|
|
.iter()
|
||
|
|
.map(|e| e["name"].as_str().unwrap())
|
||
|
|
.collect();
|
||
|
|
assert!(names.contains(&"hello.txt"));
|
||
|
|
|
||
|
|
let s = execute_tool(
|
||
|
|
&ctx,
|
||
|
|
"shell",
|
||
|
|
&json!({"command": "echo hi && ls hello.txt"}).to_string(),
|
||
|
|
)
|
||
|
|
.await;
|
||
|
|
let s: Value = serde_json::from_str(&s).unwrap();
|
||
|
|
assert_eq!(s["exit_code"], 0);
|
||
|
|
assert!(s["stdout"].as_str().unwrap().contains("hi"));
|
||
|
|
|
||
|
|
let _ = std::fs::remove_dir_all(&dir);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn rejects_path_traversal() {
|
||
|
|
let (ctx, dir) = temp_ctx();
|
||
|
|
let out = execute_tool(
|
||
|
|
&ctx,
|
||
|
|
"read_file",
|
||
|
|
&json!({"path": "../outside.txt"}).to_string(),
|
||
|
|
)
|
||
|
|
.await;
|
||
|
|
let v: Value = serde_json::from_str(&out).unwrap();
|
||
|
|
assert!(v.get("error").is_some(), "{v}");
|
||
|
|
let _ = std::fs::remove_dir_all(&dir);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn tool_defs_include_four() {
|
||
|
|
let defs = tool_definitions();
|
||
|
|
let arr = defs.as_array().unwrap();
|
||
|
|
assert_eq!(arr.len(), 4);
|
||
|
|
let names: Vec<&str> = arr
|
||
|
|
.iter()
|
||
|
|
.map(|t| t["function"]["name"].as_str().unwrap())
|
||
|
|
.collect();
|
||
|
|
assert!(names.contains(&"shell"));
|
||
|
|
assert!(names.contains(&"list_dir"));
|
||
|
|
assert!(names.contains(&"read_file"));
|
||
|
|
assert!(names.contains(&"write_file"));
|
||
|
|
}
|
||
|
|
}
|