LazyBoy2/crates/grokboy-core/src/tools.rs

//! Built-in tools: shell, files, completion, and optional Playwright browser (P3).

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;

use crate::browser::{self, LastUrlSlot};

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>,
    /// Last navigated browser URL (shared across clones).
    pub last_browser_url: LastUrlSlot,
}

impl ToolContext {
    pub fn new(cwd: impl Into<PathBuf>) -> Self {
        let cwd = cwd.into();
        Self {
            cwd: cwd.clone(),
            workspace_root: Some(cwd),
            last_browser_url: std::sync::Arc::new(std::sync::Mutex::new(None)),
        }
    }

    pub fn with_workspace(mut self, root: Option<PathBuf>) -> Self {
        self.workspace_root = root;
        self
    }

    pub fn last_browser_url_value(&self) -> Option<String> {
        self.last_browser_url.lock().ok().and_then(|g| g.clone())
    }
}

/// OpenAI-compatible tool definitions for chat completions.
pub fn tool_definitions() -> Value {
    let mut defs = 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"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "report_done",
                "description": "Signal that the task is complete. Stops the agent loop with a done verdict. Call once when finished; include a short summary in message.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "message": { "type": "string", "description": "Final summary for the user" }
                    },
                    "required": ["message"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "report_blocked",
                "description": "Signal that the task cannot proceed. Stops the agent loop with a blocked verdict and reason. Prefer this over spinning or inventing results.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "reason": { "type": "string", "description": "Why the agent is blocked" }
                    },
                    "required": ["reason"]
                }
            }
        }
    ]);
    // Optional Playwright DOM tools (always registered; fail closed with install hint if missing).
    if let Some(arr) = defs.as_array_mut() {
        arr.extend(browser::browser_tool_definitions());
    }
    defs
}

/// True if the tool name ends the ReAct loop (completion contract).
pub fn is_completion_tool(name: &str) -> bool {
    matches!(name, "report_done" | "report_blocked")
}

/// 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,
        "report_done" => tool_report_done(&args),
        "report_blocked" => tool_report_blocked(&args),
        name if browser::is_browser_tool(name) => {
            browser::execute_browser_tool(&ctx.cwd, &ctx.last_browser_url, name, &args).await
        }
        other => Err(anyhow!("unknown tool: {other}")),
    }
}

fn tool_report_done(args: &Value) -> Result<Value> {
    let message = args
        .get("message")
        .and_then(|v| v.as_str())
        .ok_or_else(|| anyhow!("report_done: missing 'message'"))?;
    Ok(json!({
        "status": "done",
        "message": message,
    }))
}

fn tool_report_blocked(args: &Value) -> Result<Value> {
    let reason = args
        .get("reason")
        .and_then(|v| v.as_str())
        .ok_or_else(|| anyhow!("report_blocked: missing 'reason'"))?;
    Ok(json!({
        "status": "blocked",
        "reason": reason,
    }))
}

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);
    }

    #[tokio::test]
    async fn report_done_and_blocked() {
        let (ctx, dir) = temp_ctx();
        let done = execute_tool(
            &ctx,
            "report_done",
            &json!({"message": "all good"}).to_string(),
        )
        .await;
        let done: Value = serde_json::from_str(&done).unwrap();
        assert_eq!(done["status"], "done");
        assert_eq!(done["message"], "all good");

        let blocked = execute_tool(
            &ctx,
            "report_blocked",
            &json!({"reason": "no access"}).to_string(),
        )
        .await;
        let blocked: Value = serde_json::from_str(&blocked).unwrap();
        assert_eq!(blocked["status"], "blocked");
        assert_eq!(blocked["reason"], "no access");

        assert!(is_completion_tool("report_done"));
        assert!(is_completion_tool("report_blocked"));
        assert!(!is_completion_tool("shell"));
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn tool_defs_include_core_and_browser() {
        let defs = tool_definitions();
        let arr = defs.as_array().unwrap();
        assert_eq!(arr.len(), 11); // 6 core + 5 browser
        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"));
        assert!(names.contains(&"report_done"));
        assert!(names.contains(&"report_blocked"));
        assert!(names.contains(&"browser_navigate"));
        assert!(names.contains(&"browser_snapshot"));
        assert!(names.contains(&"browser_click"));
        assert!(names.contains(&"browser_type"));
        assert!(names.contains(&"browser_eval"));
    }

    #[tokio::test]
    async fn browser_tool_fails_closed_without_chromium_ok_with_helper() {
        // Missing required args should fail closed via helper protocol (no Chromium needed).
        let (ctx, dir) = temp_ctx();
        // Point cwd at repo root so helper is found when running under cargo test.
        let root = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../..");
        let root = root.canonicalize().unwrap_or(root);
        let ctx = ToolContext {
            cwd: root.clone(),
            workspace_root: ctx.workspace_root,
            last_browser_url: ctx.last_browser_url,
        };
        let out = execute_tool(
            &ctx,
            "browser_type",
            &json!({"selector": "#x"}).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);
    }
}