boxmaker/crates/loopd/src/tools.rs

//! The tool port, registry, dispatch and the fake tools that stand in for `brokerd` until M3.

use std::sync::Mutex;

use crate::llama::ToolSchema;

/// The seam every tool call goes through. `loopd` answers `find_tool` and `call_tool` itself;
/// anything else, including grants, is passed to the port (`brokerd`).
pub trait ToolPort: Send + Sync {
    /// Returns the final answer, never `PendingApproval`. If the broker says the call is waiting
    /// for the owner, the port calls `on_pending` once and goes on waiting.
    fn call(
        &self,
        request: &proto::ToolRequest,
        on_pending: &mut dyn FnMut(&Pending),
    ) -> proto::ToolResponse;
}

/// One waiting call, as the broker reports it.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Pending {
    pub approval: u64,
    pub expires: proto::Timestamp,
}

#[derive(Debug, Clone, PartialEq)]
pub struct Entry {
    pub schema: ToolSchema,
    pub core: bool,
}

#[derive(Debug, Clone, PartialEq, Default)]
pub struct Registry {
    entries: Vec<Entry>,
}

impl Registry {
    pub fn new(entries: Vec<Entry>) -> Registry {
        Registry { entries }
    }

    /// The tool set for the M2b epoch: `clock` (core) and `echo` (not core). This is the test
    /// registry: `echo` exists only in `FakeTools` and the recorded conversations, and `loopd
    /// serve` never uses it.
    pub fn m2b() -> Registry {
        Registry::new(vec![
            Entry {
                schema: clock_schema(),
                core: true,
            },
            Entry {
                schema: echo_schema(),
                core: false,
            },
        ])
    }

    /// The tool set for the M3a epoch: `clock` (core), then the four tools `brokerd` decides on.
    /// The core entries match `m2b()`, so the tools array is unchanged.
    pub fn m3a() -> Registry {
        Registry::new(vec![
            Entry {
                schema: clock_schema(),
                core: true,
            },
            Entry {
                schema: read_file_schema(),
                core: false,
            },
            Entry {
                schema: write_file_schema(),
                core: false,
            },
            Entry {
                schema: shell_schema(),
                core: false,
            },
            Entry {
                schema: http_fetch_schema(),
                core: false,
            },
        ])
    }

    /// The fixed tools array: core entries, then `find_tool`, then `call_tool`.
    pub fn core_schemas(&self) -> Vec<ToolSchema> {
        let mut schemas: Vec<ToolSchema> = self
            .entries
            .iter()
            .filter(|e| e.core)
            .map(|e| e.schema.clone())
            .collect();
        schemas.push(find_tool_schema());
        schemas.push(call_tool_schema());
        schemas
    }

    pub fn get(&self, name: &str) -> Option<&Entry> {
        self.entries.iter().find(|e| e.schema.name == name)
    }

    /// Case-insensitive substring on name or description; an empty query matches nothing.
    pub fn find(&self, query: &str) -> Vec<&ToolSchema> {
        if query.is_empty() {
            return Vec::new();
        }
        let lower = query.to_lowercase();
        self.entries
            .iter()
            .filter(|e| {
                e.schema.name.to_lowercase().contains(&lower)
                    || e.schema.description.to_lowercase().contains(&lower)
            })
            .map(|e| &e.schema)
            .collect()
    }
}

pub const FIND_TOOL: &str = "find_tool";
pub const CALL_TOOL: &str = "call_tool";
pub const CLOCK: &str = "clock";

pub fn clock_schema() -> ToolSchema {
    ToolSchema {
        name: "clock".to_string(),
        description: "The current date and time, as RFC 3339 in UTC.".to_string(),
        parameters: serde_json::json!({ "type": "object", "properties": {} }),
    }
}

pub fn echo_schema() -> ToolSchema {
    ToolSchema {
        name: "echo".to_string(),
        description: "Returns its text argument unchanged.".to_string(),
        parameters: serde_json::json!({
            "type": "object",
            "properties": {
                "text": {
                    "type": "string",
                    "description": "The text to return unchanged."
                }
            },
            "required": ["text"]
        }),
    }
}

pub fn read_file_schema() -> ToolSchema {
    ToolSchema {
        name: "read_file".to_string(),
        description: "Read a text file. Needs a grant from the owner for the path.".to_string(),
        parameters: serde_json::json!({
            "type": "object",
            "properties": {
                "path": { "type": "string", "description": "The absolute path of the file." }
            },
            "required": ["path"]
        }),
    }
}

pub fn write_file_schema() -> ToolSchema {
    ToolSchema {
        name: "write_file".to_string(),
        description: "Write a text file, replacing it. Needs a grant from the owner for the path."
            .to_string(),
        parameters: serde_json::json!({
            "type": "object",
            "properties": {
                "path": { "type": "string", "description": "The absolute path of the file." },
                "content": { "type": "string", "description": "The whole new content of the file." }
            },
            "required": ["path", "content"]
        }),
    }
}

pub fn shell_schema() -> ToolSchema {
    ToolSchema {
        name: "shell".to_string(),
        description: "Run a shell command in a sandbox. Needs a grant from the owner.".to_string(),
        parameters: serde_json::json!({
            "type": "object",
            "properties": {
                "command": { "type": "string", "description": "The command line to run." },
                "cwd": { "type": "string", "description": "The absolute path of the directory to run it in." }
            },
            "required": ["command"]
        }),
    }
}

pub fn http_fetch_schema() -> ToolSchema {
    ToolSchema {
        name: "http_fetch".to_string(),
        description: "Fetch an https URL. Needs a grant from the owner for the host.".to_string(),
        parameters: serde_json::json!({
            "type": "object",
            "properties": {
                "url": { "type": "string", "description": "The URL. It must start with https://" }
            },
            "required": ["url"]
        }),
    }
}

fn find_tool_schema() -> ToolSchema {
    ToolSchema {
        name: "find_tool".to_string(),
        description:
            "Search for more tools by keyword. Returns the schemas of the tools that match."
                .to_string(),
        parameters: serde_json::json!({
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The keyword to search for."
                }
            },
            "required": ["query"]
        }),
    }
}

fn call_tool_schema() -> ToolSchema {
    ToolSchema {
        name: "call_tool".to_string(),
        description: "Call a tool that find_tool returned. Pass its name and an arguments object that fits its schema."
            .to_string(),
        parameters: serde_json::json!({
            "type": "object",
            "properties": {
                "name": {
                    "type": "string",
                    "description": "The name of the tool to call."
                },
                "arguments": {
                    "type": "object",
                    "description": "The arguments object that fits the tool's schema."
                }
            },
            "required": ["name", "arguments"]
        }),
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Dispatch {
    Local(String),
    Port { tool: String, arguments: String },
}

/// Decides whether `loopd` handles a tool request itself or sends it to the port.
pub fn dispatch(registry: &Registry, name: &str, arguments: &str) -> Dispatch {
    if name == FIND_TOOL {
        return dispatch_find_tool(registry, arguments);
    }
    if name == CALL_TOOL {
        return dispatch_call_tool(registry, arguments);
    }
    if name == CLOCK {
        return Dispatch::Local(proto::Timestamp::now().to_rfc3339());
    }
    Dispatch::Port {
        tool: name.to_string(),
        arguments: arguments.to_string(),
    }
}

fn dispatch_find_tool(registry: &Registry, arguments: &str) -> Dispatch {
    let query = match serde_json::from_str::<serde_json::Value>(arguments) {
        Ok(serde_json::Value::Object(obj)) => obj
            .get("query")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string(),
        _ => return Dispatch::Local("No tool matches.".to_string()),
    };

    let matches = registry.find(&query);
    if matches.is_empty() {
        return Dispatch::Local("No tool matches.".to_string());
    }

    let mut text = format!("{} tool(s) match:\n", matches.len());
    for schema in matches {
        let json = serde_json::to_string(schema).unwrap_or_else(|_| "{}".to_string());
        text.push_str(&json);
        text.push('\n');
    }
    text.push_str("Call it with call_tool.");
    Dispatch::Local(text)
}

fn dispatch_call_tool(registry: &Registry, arguments: &str) -> Dispatch {
    let obj = match serde_json::from_str::<serde_json::Value>(arguments) {
        Ok(serde_json::Value::Object(obj)) => obj,
        _ => return Dispatch::Local("call_tool needs a JSON object.".to_string()),
    };

    let name = obj.get("name").and_then(|v| v.as_str()).unwrap_or("");
    let args_value = obj.get("arguments").unwrap_or(&serde_json::Value::Null);

    match registry.get(name) {
        Some(entry) if entry.core => {
            Dispatch::Local(format!("{name} is a core tool; call it directly."))
        }
        Some(_) => {
            let args_text =
                serde_json::to_string(args_value).unwrap_or_else(|_| "null".to_string());
            Dispatch::Port {
                tool: name.to_string(),
                arguments: args_text,
            }
        }
        None => Dispatch::Local(format!("No tool named \"{name}\"")),
    }
}

/// Cuts to at most `cap` bytes on a character boundary and appends "\n[truncated]" if it cut.
pub fn cap_result(text: &str, cap: usize) -> (String, bool) {
    if text.len() <= cap {
        return (text.to_string(), false);
    }
    let boundary = text.floor_char_boundary(cap);
    let mut result = text[..boundary].to_string();
    result.push_str("\n[truncated]");
    (result, true)
}

/// The one fixed sentence a denial reaches the model as, named by its reason.
pub fn denial_text(reason: proto::DenyReason) -> &'static str {
    match reason {
        proto::DenyReason::NoGrant => "Denied: no grant allows this call.",
        proto::DenyReason::GrantExpired => "Denied: the grant for this call has expired.",
        proto::DenyReason::TaintTooHigh => {
            "Denied: this session has seen data too sensitive for this call."
        }
        proto::DenyReason::DeniedByGrant => "Denied: a grant forbids this call.",
        proto::DenyReason::ApprovalRefused => "Denied: the owner refused this call.",
        proto::DenyReason::ApprovalExpired => {
            "Denied: the approval request expired without an answer."
        }
        proto::DenyReason::InvalidArguments => "Denied: the arguments are not valid for this tool.",
        proto::DenyReason::GrantsInvalid => {
            "Denied: the grant files have an error; the owner has been told."
        }
        proto::DenyReason::AuditUnavailable => {
            "Denied: the audit log cannot be written; the owner has been told."
        }
        proto::DenyReason::StateUnreadable => {
            "Denied: this session's broker state is damaged; the owner has been told."
        }
    }
}

pub struct FakeTools {
    calls: Mutex<Vec<proto::ToolRequest>>,
}

impl Default for FakeTools {
    fn default() -> Self {
        FakeTools::new()
    }
}

impl FakeTools {
    pub fn new() -> Self {
        FakeTools {
            calls: Mutex::new(Vec::new()),
        }
    }

    pub fn calls(&self) -> Vec<proto::ToolRequest> {
        self.calls.lock().unwrap_or_else(|p| p.into_inner()).clone()
    }
}

impl ToolPort for FakeTools {
    fn call(
        &self,
        request: &proto::ToolRequest,
        _on_pending: &mut dyn FnMut(&Pending),
    ) -> proto::ToolResponse {
        self.calls
            .lock()
            .unwrap_or_else(|p| p.into_inner())
            .push(request.clone());

        match request.tool.as_str() {
            "echo" => match serde_json::from_str::<serde_json::Value>(&request.arguments) {
                Ok(serde_json::Value::Object(obj)) => {
                    match obj.get("text").and_then(|v| v.as_str()) {
                        Some(text) => proto::ToolResponse::Result {
                            content: text.to_string(),
                            class: proto::DataClass::Public,
                            untrusted: false,
                            truncated: false,
                        },
                        None => proto::ToolResponse::Failed {
                            message: "missing text argument".to_string(),
                        },
                    }
                }
                _ => proto::ToolResponse::Failed {
                    message: "arguments must be a JSON object".to_string(),
                },
            },
            _ => proto::ToolResponse::Denied {
                reason: proto::DenyReason::NoGrant,
            },
        }
    }
}