superintelligence-fork / src/tools.rs

raw file 17365 bytes
//! Tools the model can call: filesystem read/search/write and a shell.
//!
//! Every tool resolves its path relative to a root (the working directory) so
//! the model can't wander off into the rest of the machine. Writes and shell
//! execution additionally require approval, which main.rs supplies as a
//! callback.

use anyhow::{Context as _, Result};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};

use crate::config;

/// A tool call the model asked for.
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct ToolCall {
    #[serde(default)]
    pub id: String,
    #[serde(rename = "type", default)]
    pub kind: String,
    pub function: FunctionCall,
}

#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct FunctionCall {
    pub name: String,
    /// JSON-encoded arguments, as a string (OpenAI's wire format).
    pub arguments: String,
}

impl ToolCall {
    /// One-line rendering of the call, for showing the user what's happening.
    pub fn summary(&self) -> String {
        let args =
            serde_json::from_str::<serde_json::Value>(&self.function.arguments).unwrap_or_default();
        let mut parts = Vec::new();
        if let Some(obj) = args.as_object() {
            for (key, value) in obj {
                let text = match value {
                    serde_json::Value::String(s) => {
                        // Shell commands and file bodies are the interesting
                        // part, so don't truncate those.
                        if key == "command" || key == "content" {
                            s.clone()
                        } else {
                            truncate(s, 80)
                        }
                    }
                    other => other.to_string(),
                };
                parts.push(format!("{key}={text}"));
            }
        }
        format!("{}({})", self.function.name, parts.join(", "))
    }
}

/// JSON schema list sent with the request, declaring the available tools.
pub fn specs() -> serde_json::Value {
    serde_json::json!([
        {
            "type": "function",
            "function": {
                "name": "read_file",
                "description": "Read a UTF-8 text file. Returns the file contents, truncated if long.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "path": {"type": "string", "description": "File path, relative to the working directory."}
                    },
                    "required": ["path"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "list_dir",
                "description": "List the entries of a directory, with sizes and a / suffix on directories.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "path": {"type": "string", "description": "Directory path, relative to the working directory. Defaults to '.'"}
                    }
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "search",
                "description": "Search file contents for a regular expression across the working directory. Returns path:line:match entries.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "pattern": {"type": "string", "description": "Regular expression to look for."},
                        "path": {"type": "string", "description": "Subdirectory to search. Defaults to '.'"},
                        "include": {"type": "string", "description": "Only search files whose name contains this, e.g. '.rs' or 'main'."}
                    },
                    "required": ["pattern"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "write_file",
                "description": "Write text to a file, creating parent directories as needed. Overwrites existing files. Requires user approval.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "path": {"type": "string", "description": "File path, relative to the working directory."},
                        "content": {"type": "string", "description": "Full contents to write."}
                    },
                    "required": ["path", "content"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "run_shell",
                "description": format!("Run a command with {} in the working directory. Requires user approval.", config::SHELL_DESC),
                "parameters": {
                    "type": "object",
                    "properties": {
                        "command": {"type": "string", "description": "The command line to execute."}
                    },
                    "required": ["command"]
                }
            }
        }
    ])
}

/// Whether the user has pre-approved destructive tools (`--yes`).
#[derive(Clone, Copy, Default)]
pub struct Policy {
    pub auto_approve: bool,
}

/// Run one tool call, returning the text to hand back to the model.
///
/// `confirm` is asked before writes and shell commands; it should return
/// whether to proceed.
pub async fn execute<C>(call: &ToolCall, policy: Policy, confirm: &mut C) -> String
where
    C: FnMut(&ToolCall) -> bool,
{
    let args = match serde_json::from_str::<serde_json::Value>(&call.function.arguments) {
        Ok(value) => value,
        Err(err) => return format!("Error: arguments were not valid JSON: {err}"),
    };

    match call.function.name.as_str() {
        "read_file" => match read_file(&arg_str(&args, "path")) {
            Ok(text) => text,
            Err(err) => format!("Error: {err:#}"),
        },
        "list_dir" => match list_dir(&arg_str(&args, "path")) {
            Ok(text) => text,
            Err(err) => format!("Error: {err:#}"),
        },
        "search" => {
            let pattern = arg_str(&args, "pattern");
            if pattern.is_empty() {
                return "Error: `pattern` is required.".into();
            }
            match search(
                &pattern,
                &arg_str(&args, "path"),
                &arg_str(&args, "include"),
            ) {
                Ok(text) => text,
                Err(err) => format!("Error: {err:#}"),
            }
        }
        "write_file" => {
            let path = arg_str(&args, "path");
            let content = arg_str(&args, "content");
            if !approve(call, policy, confirm) {
                return "Skipped: the user declined this write. Do not retry it; ask them what they'd prefer.".into();
            }
            match write_file(&path, &content) {
                Ok(report) => report,
                Err(err) => format!("Error: {err:#}"),
            }
        }
        "run_shell" => {
            let command = arg_str(&args, "command");
            if command.trim().is_empty() {
                return "Error: `command` is required.".into();
            }
            if !approve(call, policy, confirm) {
                return "Skipped: the user declined this command. Do not retry it; describe what you would have run instead.".into();
            }
            match run_shell(&command).await {
                Ok(report) => report,
                Err(err) => format!("Error: {err:#}"),
            }
        }
        other => format!("Error: unknown tool `{other}`."),
    }
}

/// Ask the callback, unless the user passed `--yes` or there is no terminal
/// able to answer.
fn approve<C>(call: &ToolCall, policy: Policy, confirm: &mut C) -> bool
where
    C: FnMut(&ToolCall) -> bool,
{
    if policy.auto_approve {
        return true;
    }
    confirm(call)
}

fn arg_str(args: &serde_json::Value, key: &str) -> String {
    args.get(key)
        .and_then(|v| v.as_str())
        .unwrap_or_default()
        .to_string()
}

// ─── filesystem ─────────────────────────────────────────────────────────────

fn root() -> Result<PathBuf> {
    std::env::current_dir().context("resolving the working directory")
}

/// Resolve a path against the working directory, refusing to leave it.
fn resolve(path: &str, must_exist: bool) -> Result<PathBuf> {
    let root = root()?;
    let root = root
        .canonicalize()
        .context("canonicalizing the working directory")?;

    let joined = if path.trim().is_empty() || path == "." {
        root.clone()
    } else {
        let candidate = Path::new(path);
        if candidate.is_absolute() {
            candidate.to_path_buf()
        } else {
            root.join(candidate)
        }
    };

    let resolved = if must_exist {
        joined
            .canonicalize()
            .with_context(|| format!("resolving {}", joined.display()))?
    } else if let Some(parent) = joined.parent() {
        // The file may not exist yet, so anchor on the nearest existing ancestor.
        let parent = parent
            .canonicalize()
            .with_context(|| format!("resolving {}", parent.display()))?;
        parent.join(joined.file_name().unwrap_or_default())
    } else {
        joined
    };

    if !resolved.starts_with(&root) {
        anyhow::bail!(
            "{} is outside the working directory ({})",
            resolved.display(),
            root.display()
        );
    }
    Ok(resolved)
}

fn read_file(path: &str) -> Result<String> {
    if path.trim().is_empty() {
        anyhow::bail!("`path` is required");
    }
    let path = resolve(path, true)?;
    let bytes = std::fs::read(&path).with_context(|| format!("reading {}", path.display()))?;
    let text = String::from_utf8(bytes)
        .map_err(|_| anyhow::anyhow!("{} is not a UTF-8 text file", path.display()))?;
    Ok(truncate(&text, config::TOOL_OUTPUT_MAX_CHARS))
}

fn list_dir(path: &str) -> Result<String> {
    let path = resolve(path, true)?;
    if !path.is_dir() {
        anyhow::bail!("{} is not a directory", path.display());
    }

    let mut entries: Vec<(String, bool, u64)> = Vec::new();
    for entry in std::fs::read_dir(&path).with_context(|| format!("listing {}", path.display()))? {
        let entry = entry?;
        let meta = entry.metadata().ok();
        let is_dir = meta.as_ref().is_some_and(|m| m.is_dir());
        let size = meta.as_ref().map(|m| m.len()).unwrap_or(0);
        entries.push((
            entry.file_name().to_string_lossy().into_owned(),
            is_dir,
            size,
        ));
    }

    if entries.is_empty() {
        return Ok("(empty directory)".into());
    }

    entries.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
    let mut out = String::new();
    for (name, is_dir, size) in entries.into_iter().take(500) {
        if is_dir {
            out.push_str(&format!("{name}/\n"));
        } else {
            out.push_str(&format!("{name} ({})\n", human_size(size)));
        }
    }
    Ok(out)
}

fn write_file(path: &str, content: &str) -> Result<String> {
    if path.trim().is_empty() {
        anyhow::bail!("`path` is required");
    }
    let path = resolve(path, false)?;
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)
            .with_context(|| format!("creating {}", parent.display()))?;
    }
    std::fs::write(&path, content).with_context(|| format!("writing {}", path.display()))?;
    Ok(format!(
        "Wrote {} to {}.",
        human_size(content.len() as u64),
        path.display()
    ))
}

// ─── search ──────────────────────────────────────────────────────────────────

/// Directories we never descend into.
const SKIP_DIRS: &[&str] = &[
    "target",
    ".git",
    "node_modules",
    "__pycache__",
    ".venv",
    "dist",
    "build",
];

fn search(pattern: &str, path: &str, include: &str) -> Result<String> {
    let re = regex::Regex::new(pattern).with_context(|| format!("compiling regex `{pattern}`"))?;
    let start = resolve(path, true)?;
    if !start.is_dir() {
        anyhow::bail!("{} is not a directory", start.display());
    }

    let mut results: Vec<String> = Vec::new();
    let mut files_scanned = 0usize;
    let mut truncated = false;

    let mut stack = vec![start];
    while let Some(dir) = stack.pop() {
        let Ok(read) = std::fs::read_dir(&dir) else {
            continue;
        };
        for entry in read.flatten() {
            let name = entry.file_name().to_string_lossy().into_owned();
            let Ok(meta) = entry.metadata() else { continue };

            if meta.is_dir() {
                if !SKIP_DIRS.contains(&name.as_str()) && !name.starts_with('.') {
                    stack.push(entry.path());
                }
                continue;
            }
            if !include.is_empty() && !name.contains(include) {
                continue;
            }
            if meta.len() > 1_000_000 {
                continue; // skip binaries and huge files
            }

            let Ok(bytes) = std::fs::read(entry.path()) else {
                continue;
            };
            let Ok(text) = String::from_utf8(bytes) else {
                continue;
            };
            files_scanned += 1;

            for (lineno, line) in text.lines().enumerate() {
                if re.is_match(line) {
                    if results.len() >= 200 {
                        truncated = true;
                        break;
                    }
                    results.push(format!(
                        "{}:{}:{}",
                        entry.path().display(),
                        lineno + 1,
                        truncate(line.trim(), 300)
                    ));
                }
            }
            if truncated {
                break;
            }
        }
        if truncated {
            break;
        }
    }

    if results.is_empty() {
        return Ok(format!(
            "No matches for /{pattern}/ in {files_scanned} file(s)."
        ));
    }

    let mut out = results.join("\n");
    if truncated {
        out.push_str("\n… (truncated; narrow the pattern or path)");
    }
    Ok(truncate(&out, config::TOOL_OUTPUT_MAX_CHARS))
}

// ─── shell ──────────────────────────────────────────────────────────────────

async fn run_shell(command: &str) -> Result<String> {
    use tokio::process::Command;

    let root = root()?;
    let child = Command::new(config::SHELL_PROGRAM)
        .arg(config::SHELL_FLAG)
        .arg(command)
        .current_dir(&root)
        .stdin(std::process::Stdio::null())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()
        .with_context(|| {
            format!(
                "spawning `{} {} {command}`",
                config::SHELL_PROGRAM,
                config::SHELL_FLAG
            )
        })?;

    let output = match tokio::time::timeout(
        std::time::Duration::from_secs(config::SHELL_TIMEOUT_SECS),
        child.wait_with_output(),
    )
    .await
    {
        Ok(result) => result.context("collecting command output")?,
        Err(_) => anyhow::bail!("timed out after {}s", config::SHELL_TIMEOUT_SECS),
    };

    let mut out = String::new();
    let stdout = String::from_utf8_lossy(&output.stdout);
    let stderr = String::from_utf8_lossy(&output.stderr);
    if !stdout.trim().is_empty() {
        out.push_str(stdout.trim_end());
    }
    if !stderr.trim().is_empty() {
        if !out.is_empty() {
            out.push_str("\n\n[stderr]\n");
        }
        out.push_str(stderr.trim_end());
    }
    if out.is_empty() {
        out.push_str("(no output)");
    }
    if !output.status.success() {
        out.push_str(&format!("\n\n[exit status: {}]", output.status));
    }
    Ok(truncate(&out, config::TOOL_OUTPUT_MAX_CHARS))
}

// ─── helpers ────────────────────────────────────────────────────────────────

/// Cut at a char boundary so the result fits, appending a marker.
fn truncate(text: &str, max: usize) -> String {
    if text.chars().count() <= max {
        return text.to_string();
    }
    let mut out: String = text.chars().take(max).collect();
    out.push_str("\n… (truncated)");
    out
}

fn human_size(bytes: u64) -> String {
    const UNITS: [&str; 5] = ["B", "KB", "MB", "GB", "TB"];
    let mut value = bytes as f64;
    let mut unit = 0;
    while value >= 1024.0 && unit < UNITS.len() - 1 {
        value /= 1024.0;
        unit += 1;
    }
    if unit == 0 {
        format!("{bytes} B")
    } else {
        format!("{value:.1} {}", UNITS[unit])
    }
}