mecha_core/tool/

mod.rs

//! Tools: the things an agent can actually do.
//!
//! A tool is a name, a description, a JSON Schema, and an async function. The
//! registry holds them; MCP servers and native Rust functions both land here as
//! the same trait object, so the agent loop never learns the difference.

pub mod ask;
pub mod builtin;
pub mod todo;

use crate::config::{PermissionMode, SecurityConfig, ToolsConfig};
use crate::message::ToolSpec;
use anyhow::Result;
use async_trait::async_trait;
use serde_json::Value;
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;

#[derive(Debug, Clone)]
pub struct ToolOutput {
    pub content: String,
    /// Returned to the model as `is_error: true` so it can recover rather than
    /// treating the failure as a result.
    pub is_error: bool,
    /// True when this content actually came from outside the machine.
    ///
    /// Distinct from the tool's declared `untrusted_input` capability, which
    /// says what the tool *can* return. A refusal generated by mecha's own
    /// guards is not third-party content, and labelling it as such makes the
    /// model invent explanations for its own harness's behaviour.
    pub external: bool,
}

impl ToolOutput {
    pub fn ok(content: impl Into<String>) -> Self {
        ToolOutput {
            content: content.into(),
            is_error: false,
            external: false,
        }
    }

    pub fn err(content: impl Into<String>) -> Self {
        ToolOutput {
            content: content.into(),
            is_error: true,
            external: false,
        }
    }

    /// Mark this content as having come from outside the machine.
    pub fn from_outside(mut self) -> Self {
        self.external = true;
        self
    }
}

/// What a tool can do — the vocabulary MCP standardized (`readOnly`,
/// `destructive`, `openWorld`) plus the two axes that decide whether an agent
/// can be turned into an exfiltration tool.
///
/// The *lethal trifecta* is private data + untrusted content + a way out. Any
/// agent holding all three can be instructed, by text hidden in the content it
/// reads, to take the private data and send it somewhere. Annotating tools on
/// these axes is what lets the loop refuse that combination structurally.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Capabilities {
    /// Returns data the user considers private.
    pub private_data: bool,
    /// Returns content a third party can influence — a web page, an email body,
    /// a calendar invite title. Treat everything it returns as hostile.
    pub untrusted_input: bool,
    /// Can transmit data outside the user's control. Note that a plain HTTP GET
    /// qualifies: the secret goes in the query string.
    pub external_send: bool,
    /// May destroy or overwrite data.
    pub destructive: bool,
}

impl Capabilities {
    pub fn private(mut self) -> Self {
        self.private_data = true;
        self
    }
    pub fn untrusted(mut self) -> Self {
        self.untrusted_input = true;
        self
    }
    pub fn sends(mut self) -> Self {
        self.external_send = true;
        self
    }
    pub fn destructive(mut self) -> Self {
        self.destructive = true;
        self
    }

    /// Everything either side declares.
    ///
    /// Union rather than assignment, because the only safe direction for an
    /// override is *wider*. Letting config narrow a tool's declared
    /// capabilities would disarm the interlock on the strength of a claim
    /// nothing enforces — the same mistake as a sandbox that silently degrades,
    /// and it would make the cheapest configuration the most dangerous one. A
    /// server that genuinely over-declares is what `TrifectaPolicy` is for: one
    /// deliberate, visible decision instead of a quiet per-server exemption.
    pub fn union(self, other: Capabilities) -> Self {
        Capabilities {
            private_data: self.private_data || other.private_data,
            untrusted_input: self.untrusted_input || other.untrusted_input,
            external_send: self.external_send || other.external_send,
            destructive: self.destructive || other.destructive,
        }
    }
}

#[async_trait]
pub trait Tool: Send + Sync {
    fn name(&self) -> &str;
    fn description(&self) -> &str;
    fn input_schema(&self) -> Value;

    /// Read-only tools skip the approval gate and are safe to run in parallel.
    fn read_only(&self) -> bool {
        false
    }

    /// Declared risk surface. The default is the conservative one for a tool
    /// nobody has classified: assume it does nothing special.
    fn capabilities(&self) -> Capabilities {
        Capabilities::default()
    }

    async fn call(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput>;

    /// State this tool holds that a compaction must not lose.
    ///
    /// Compaction replaces the middle of a transcript with prose, and the
    /// measured failure mode is that a summariser preserves *what is true* and
    /// drops *how far you got*. Some of "how far you got" does not live in the
    /// messages at all — it lives in a tool — and for that state a summary is
    /// the wrong mechanism twice over: it is lossy, and the tool already has
    /// the exact current answer.
    ///
    /// So a tool may hand its state to the compaction to be carried across
    /// **verbatim**. Three rules make this safe rather than a second source of
    /// truth:
    ///
    /// - It is read at compaction time, so it is current by construction. A
    ///   stale copy is impossible because nothing stores one.
    /// - Exactly one copy survives: the carried block replaces the previous
    ///   one rather than accumulating beside it, or an old list would sit in
    ///   the prompt contradicting the new one.
    /// - It is for state the tool *owns*, not a summary of what happened. A
    ///   tool that returned prose here would be smuggling a second summariser
    ///   into the loop, unvalidated.
    ///
    /// `None` — the default — means "nothing worth carrying", which is the
    /// honest answer for every stateless tool.
    fn carried_state(&self) -> Option<CarriedState> {
        None
    }

    fn spec(&self) -> ToolSpec {
        ToolSpec {
            name: self.name().to_string(),
            description: self.description().to_string(),
            input_schema: self.input_schema(),
        }
    }
}

/// A tool's own state, on its way across a compaction.
///
/// `label` names it in the rebuilt prompt (the tool's name is the obvious
/// choice); `body` is reproduced exactly, because verbatim is the whole point.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CarriedState {
    pub label: String,
    pub body: String,
}

/// What a tool is allowed to touch.
#[derive(Debug, Clone)]
pub struct ToolCtx {
    /// Filesystem tools refuse paths outside this root.
    pub workspace: PathBuf,
    pub shell_timeout: std::time::Duration,
    pub security: SecurityConfig,
    /// The byte budget one *turn's* tool results share, divided equally
    /// across the calls in the batch so one runaway tool cannot starve its
    /// siblings (mecha executes a turn's calls concurrently, so they land
    /// together). The old per-tool cap was 200 KB — ~50k tokens, 1.5× the
    /// whole local context window, which is not a cap so much as a promise
    /// to overflow.
    pub output_budget_bytes: usize,
    /// Where an oversized result is saved in full before its transcript copy
    /// is cut. `None` disables spilling — the cut then names what was lost
    /// instead of where to find it. Per-context on purpose: two eval cases
    /// sharing one spill directory could read each other's output through it.
    pub spill_dir: Option<PathBuf>,
    /// The run's event channel, so a tool that *contains* a run — a subagent —
    /// can surface its progress instead of going dark until it returns.
    ///
    /// Display-only, and treat it that way: any tool (including a third-party
    /// MCP server's) can send fabricated events down this channel, so nothing
    /// that matters may key off it. Conversation state, taint, and run
    /// completion all come from the loop and the caller's join handle, never
    /// from events. Stamped by [`Agent::run_in`] per run; `None` everywhere
    /// nobody is watching (batch, eval).
    ///
    /// [`Agent::run_in`]: crate::agent::Agent::run_in
    pub events: Option<tokio::sync::mpsc::UnboundedSender<crate::agent::AgentEvent>>,
    /// The run's cancellation token. A tool that contains a run passes it on,
    /// so cancelling the parent actually cancels the child instead of politely
    /// waiting out its entire run. Stamped by `Agent::run_in`, like `events`.
    pub cancel: Option<tokio_util::sync::CancellationToken>,
    /// The run's phase. A tool that contains a run passes it on, so delegation
    /// is not the way to get a write executed from a planning run. Stamped by
    /// `Agent::run_in`, like `events`.
    pub phase: crate::agent::Phase,
    /// The `tool_use` id of the call this context was built for. Stamped per
    /// dispatch (only when `events` is watched), so a tool that contains a
    /// run can tag its forwarded events with the call that spawned it — two
    /// subagents running in parallel are otherwise indistinguishable to a
    /// renderer.
    pub call_id: Option<String>,
}

impl Default for ToolCtx {
    fn default() -> Self {
        ToolCtx {
            workspace: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
            shell_timeout: std::time::Duration::from_secs(120),
            security: SecurityConfig::default(),
            output_budget_bytes: 24_000,
            spill_dir: fresh_spill_dir(),
            events: None,
            cancel: None,
            phase: crate::agent::Phase::default(),
            call_id: None,
        }
    }
}

/// A spill directory no other context shares. Not created until first used.
fn fresh_spill_dir() -> Option<PathBuf> {
    Some(std::env::temp_dir().join(format!("mecha-spill-{}", uuid::Uuid::new_v4())))
}

impl ToolCtx {
    /// The same policy pointed at a different root. Used to give one run — an
    /// eval case, a batch item — its own isolated copy of a workspace without
    /// rebuilding the agent around it. The spill directory is re-derived too:
    /// a re-rooted context is a new isolation domain, and inheriting the old
    /// one would let its runs read each other's spilled output.
    pub fn with_workspace(&self, workspace: impl Into<PathBuf>) -> Self {
        ToolCtx {
            workspace: workspace.into(),
            spill_dir: fresh_spill_dir(),
            ..self.clone()
        }
    }

    /// Resolve a model-supplied path against the workspace and prove it stays
    /// inside. The path is untrusted input: `..`, symlinks, and absolute paths
    /// all have to be checked after canonicalization, not before.
    pub fn resolve(&self, raw: &str) -> Result<PathBuf> {
        let candidate = {
            let p = Path::new(raw);
            if p.is_absolute() {
                p.to_path_buf()
            } else {
                self.workspace.join(p)
            }
        };

        // The file may not exist yet (a write), so canonicalize the nearest
        // existing ancestor and re-append the rest.
        let mut existing = candidate.as_path();
        let mut trailing = Vec::new();
        let canonical_root = loop {
            match existing.canonicalize() {
                Ok(c) => break c,
                Err(_) => match existing.parent() {
                    Some(parent) => {
                        if let Some(name) = existing.file_name() {
                            trailing.push(name.to_owned());
                        }
                        existing = parent;
                    }
                    None => anyhow::bail!("cannot resolve path {raw:?}"),
                },
            }
        };
        let mut resolved = canonical_root;
        for part in trailing.iter().rev() {
            resolved.push(part);
        }

        let root = self
            .workspace
            .canonicalize()
            .unwrap_or_else(|_| self.workspace.clone());
        if resolved.starts_with(&root) {
            return Ok(resolved);
        }
        // The spill directory is the one sanctioned exception: oversized tool
        // output is saved there, and the truncation marker tells the model to
        // read the rest from exactly that path. Its contents are this
        // context's own tool results, so nothing new becomes reachable.
        if let Some(spill) = &self.spill_dir {
            let spill_root = spill.canonicalize().unwrap_or_else(|_| spill.clone());
            if resolved.starts_with(&spill_root) {
                return Ok(resolved);
            }
        }
        anyhow::bail!(
            "path {raw:?} resolves outside the workspace ({})",
            root.display()
        )
    }
}

/// Floor under a result's share of the turn budget. A wide batch must not
/// starve every result down to a marker with no content: below this, the
/// division stops and the total budget is allowed to overrun instead.
pub const SPILL_FLOOR_BYTES: usize = 4_096;

/// Cut an oversized tool result down to `cap` bytes, saving the full output
/// where the model can get it back.
///
/// The marker is written for the model, and it names the recovery — a
/// truncation notice that only says "gone" leaves the model to conclude the
/// rest never existed, and the elision line number is what makes the recovery
/// a single call instead of a scan. A failed spill degrades to a plain cut
/// that says the output was *not* saved; losing the tail must never lose the
/// run.
pub fn cap_result(
    content: String,
    cap: usize,
    spill_dir: Option<&Path>,
    tool: &str,
    id: &str,
) -> String {
    if content.len() <= cap {
        return content;
    }
    // Cut on a char boundary, never mid-codepoint.
    let mut cut = cap;
    while cut > 0 && !content.is_char_boundary(cut) {
        cut -= 1;
    }
    let head = &content[..cut];
    // The line the elision starts on. A cut mid-line means that same line —
    // re-reading from it overlaps a little, which is the right direction.
    let line = head.matches('\n').count() + 1;
    let total = content.len();

    let saved = spill_dir.and_then(|dir| {
        // Owner-only: spilled output is tool results in full — the same
        // sensitivity as the transcript, sitting in the shared temp dir.
        crate::create_private_dir(dir).ok()?;
        // A random component, because the call id alone can collide: batch
        // items and non-sandboxed eval cases share one context, and a local
        // server under a pinned seed can hand identical requests identical
        // call ids. A collision would silently overwrite, leaving one
        // conversation's marker pointing at another conversation's content.
        let tag = &uuid::Uuid::new_v4().to_string()[..8];
        let file = dir.join(format!("{}-{}-{tag}.txt", safe_name(tool), safe_name(id)));
        std::fs::write(&file, &content).ok()?;
        Some(file)
    });

    match saved {
        Some(path) => format!(
            "{head}\n\n[truncated by the harness: showing the first {cut} of {total} bytes; \
             the rest begins on line {line}. The full output is saved at {path} — continue \
             with fs_read {{\"path\": \"{path}\", \"offset\": {line}}}, or search it with \
             grep.]",
            path = path.display()
        ),
        None => format!(
            "{head}\n\n[truncated by the harness: {omitted} of {total} bytes were dropped \
             from line {line} on, and the full output could not be saved. Narrow the \
             request and re-run the tool if the rest is needed.]",
            omitted = total - cut
        ),
    }
}

/// Tool names and call ids become file names; anything else becomes `-`.
fn safe_name(s: &str) -> String {
    s.chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
                c
            } else {
                '-'
            }
        })
        .collect()
}

/// The decision an approver hands back for one pending call.
#[derive(Debug, Clone)]
pub enum Decision {
    Allow,
    /// The reason is passed to the model so it can pick another approach.
    Deny(String),
}

/// Gates tool calls that aren't read-only. The CLI implements this with a
/// terminal prompt; a headless caller can auto-allow or auto-deny.
#[async_trait]
pub trait Approver: Send + Sync {
    async fn approve(&self, tool: &dyn Tool, input: &Value) -> Decision;
}

/// Answers from the configured [`PermissionMode`] without asking anyone.
pub struct ModeApprover {
    pub mode: PermissionMode,
}

#[async_trait]
impl Approver for ModeApprover {
    async fn approve(&self, tool: &dyn Tool, _input: &Value) -> Decision {
        match self.mode {
            PermissionMode::Allow => Decision::Allow,
            PermissionMode::ReadOnly if tool.read_only() => Decision::Allow,
            PermissionMode::ReadOnly => Decision::Deny(format!(
                "`{}` modifies state and this run is read-only",
                tool.name()
            )),
            // Nothing is watching to answer, so the safe reading of "ask" is no.
            PermissionMode::Ask => Decision::Deny(format!(
                "`{}` needs approval and this run is non-interactive (use --yes to allow)",
                tool.name()
            )),
        }
    }
}

#[derive(Default)]
pub struct Registry {
    tools: BTreeMap<String, Arc<dyn Tool>>,
}

impl Registry {
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a tool. A later registration with the same name replaces the
    /// earlier one, so MCP servers can shadow built-ins deliberately.
    pub fn insert(&mut self, tool: Arc<dyn Tool>) {
        self.tools.insert(tool.name().to_string(), tool);
    }

    pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
        self.tools.get(name)
    }

    pub fn is_empty(&self) -> bool {
        self.tools.is_empty()
    }

    pub fn len(&self) -> usize {
        self.tools.len()
    }

    pub fn iter(&self) -> impl Iterator<Item = &Arc<dyn Tool>> {
        self.tools.values()
    }

    /// Everything the registered tools want carried across a compaction.
    ///
    /// In the registry's stable order, so a compaction does not reorder the
    /// prompt for a reason nobody can see. Asked of every tool, including an
    /// MCP server's — the loop does not learn which tools have state, only
    /// that some do, which is the same reason it never learns where a tool
    /// came from.
    pub fn carried_state(&self) -> Vec<CarriedState> {
        self.tools
            .values()
            .filter_map(|t| t.carried_state())
            .collect()
    }

    /// Specs in a stable order — the tool list is the very front of the prompt
    /// prefix, so reordering it would invalidate the cache on every request.
    pub fn specs(&self) -> Vec<ToolSpec> {
        self.tools.values().map(|t| t.spec()).collect()
    }

    /// Specs a given phase permits, in the same stable order.
    ///
    /// Note what this does to the prompt cache: planning sends a shorter tool
    /// list, so switching phase changes the front of the prefix and the next
    /// turn re-pays for it. That is the price of the tools being genuinely
    /// absent rather than merely refused, and it is the right trade.
    pub fn specs_for(&self, phase: crate::agent::Phase) -> Vec<ToolSpec> {
        self.tools
            .values()
            .filter(|t| phase.allows(t.read_only()))
            .map(|t| t.spec())
            .collect()
    }

    /// Register the built-ins permitted by config.
    ///
    /// The sandbox is passed in rather than read from config here because it
    /// changes what `shell` *is* — an unconfined shell and a confined one
    /// declare different capabilities, and the loop's interlock reads them.
    pub fn with_builtins(
        mut self,
        cfg: &ToolsConfig,
        sandbox: Arc<crate::sandbox::Sandbox>,
    ) -> Self {
        for tool in builtin::all(sandbox) {
            let name = tool.name();
            let allowed = cfg.enabled.is_empty() || cfg.enabled.iter().any(|e| e == name);
            let blocked = cfg.disabled.iter().any(|d| d == name);
            if allowed && !blocked {
                self.insert(tool);
            }
        }
        self
    }
}

#[cfg(test)]
mod cap_tests {
    use super::*;

    fn scratch(name: &str) -> PathBuf {
        let dir = std::env::temp_dir().join(format!("mecha-cap-{name}-{}", uuid::Uuid::new_v4()));
        std::fs::create_dir_all(&dir).unwrap();
        dir
    }

    #[test]
    fn a_result_under_the_cap_is_untouched() {
        let out = cap_result("short".into(), 100, None, "shell", "t1");
        assert_eq!(out, "short");
    }

    #[test]
    fn an_oversized_result_is_spilled_whole_and_the_marker_names_the_recovery() {
        let dir = scratch("spill");
        let body: String = (1..=100).map(|i| format!("line {i}\n")).collect();

        let out = cap_result(body.clone(), 200, Some(&dir), "shell", "t1");

        // The transcript copy is bounded...
        assert!(out.len() < body.len());
        assert!(out.starts_with("line 1\n"));
        // ...the disk copy is not: byte-identical, so nothing was lost. The
        // name carries a random tag, so it is discovered rather than assumed.
        let file = std::fs::read_dir(&dir)
            .unwrap()
            .next()
            .unwrap()
            .unwrap()
            .path();
        assert!(file
            .file_name()
            .unwrap()
            .to_str()
            .unwrap()
            .starts_with("shell-t1-"));
        assert_eq!(std::fs::read_to_string(&file).unwrap(), body);

        // The marker gives the model a single call back to the rest: the
        // path, and the line the elision starts on.
        let line = body[..200].matches('\n').count() + 1;
        assert!(out.contains(&file.display().to_string()), "{out}");
        assert!(out.contains(&format!("\"offset\": {line}")), "{out}");
        assert!(out.contains("fs_read"), "the recovery must be named: {out}");

        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn a_failed_spill_degrades_to_a_cut_that_admits_the_loss() {
        // A directory that cannot exist: spilling fails, the run must not.
        let impossible = PathBuf::from("/dev/null/not-a-dir");
        let body = "x".repeat(1000);
        let out = cap_result(body, 100, Some(&impossible), "shell", "t1");

        assert!(out.contains("could not be saved"), "{out}");
        assert!(
            out.contains("re-run the tool"),
            "the fallback still names a recovery: {out}"
        );
        assert!(
            !out.contains("/dev/null"),
            "no path is promised that does not exist"
        );
    }

    #[test]
    fn the_cut_lands_on_a_char_boundary() {
        // A cap that falls mid-codepoint must back up, not panic.
        let body = "é".repeat(100); // 2 bytes per char
        let out = cap_result(body, 33, None, "shell", "t1");
        assert!(out.starts_with(&"é".repeat(16)));
    }

    #[test]
    fn the_jail_admits_the_spill_directory_and_nothing_else_new() {
        let workspace = scratch("ws");
        let spill = scratch("spilldir");
        let ctx = ToolCtx {
            workspace: workspace.clone(),
            spill_dir: Some(spill.clone()),
            ..ToolCtx::default()
        };

        // The marker names an absolute spill path; fs_read must be able to
        // follow it, or the recovery the model was promised is a lie.
        std::fs::write(spill.join("shell-t1.txt"), "spilled").unwrap();
        let resolved = ctx
            .resolve(&spill.join("shell-t1.txt").display().to_string())
            .unwrap();
        assert!(resolved.ends_with("shell-t1.txt"));

        // The exception is the spill directory, not the temp dir around it.
        let elsewhere = std::env::temp_dir().join("mecha-cap-elsewhere.txt");
        std::fs::write(&elsewhere, "no").unwrap();
        assert!(ctx.resolve(&elsewhere.display().to_string()).is_err());

        // And with spilling disabled there is no exception at all.
        let no_spill = ToolCtx {
            workspace,
            spill_dir: None,
            ..ToolCtx::default()
        };
        assert!(no_spill
            .resolve(&spill.join("shell-t1.txt").display().to_string())
            .is_err());

        std::fs::remove_dir_all(&spill).ok();
        std::fs::remove_file(&elsewhere).ok();
    }

    #[test]
    fn a_rerooted_context_gets_its_own_spill_directory() {
        // Two eval cases sharing one spill directory could read each other's
        // output through it — the same isolation rule as the workspace copy.
        let ctx = ToolCtx::default();
        let rerooted = ctx.with_workspace(std::env::temp_dir());
        assert_ne!(ctx.spill_dir, rerooted.spill_dir);
    }
}