codecraft/server/src/tools/browser.rs

//! Interactive headless browser for agents (Chrome DevTools Protocol):
//! open a page, click, type, run JavaScript, then get a screenshot, the
//! visible text and console errors. Public internet and localhost are
//! allowed; private network addresses are blocked for every request the page
//! makes (sub-resources and redirects included).

use super::{Tool, ToolCtx, ToolOutput, arg_str, truncate};
use anyhow::{Context, Result, anyhow, bail};
use async_trait::async_trait;
use base64::Engine;
use futures::{SinkExt, StreamExt};
use serde_json::{Value, json};
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::process::Command;
use tokio::sync::{mpsc, oneshot};
use tokio_tungstenite::tungstenite::Message as WsMessage;

/// Private, link-local and other non-public ranges (loopback is allowed).
pub fn is_blocked_ip(ip: IpAddr) -> bool {
    match ip {
        IpAddr::V4(v4) => {
            let o = v4.octets();
            v4.is_private()
                || v4.is_link_local()
                || v4.is_unspecified()
                || v4.is_broadcast()
                || (o[0] == 100 && (64..128).contains(&o[1])) // CGNAT
                || (o[0] == 192 && o[1] == 0 && o[2] == 0)
                || o[0] >= 240
        }
        IpAddr::V6(v6) => {
            if let Some(v4) = v6.to_ipv4_mapped() {
                return !v4.is_loopback() && is_blocked_ip(IpAddr::V4(v4));
            }
            let s = v6.segments();
            v6.is_unspecified() || (s[0] & 0xfe00) == 0xfc00 || (s[0] & 0xffc0) == 0xfe80
        }
    }
}

/// Is the URL's host allowed? Cached per host.
async fn host_allowed(url: &str, cache: &Mutex<HashMap<String, bool>>) -> bool {
    let Ok(u) = reqwest::Url::parse(url) else { return false };
    match u.scheme() {
        "http" | "https" | "ws" | "wss" => {}
        "data" | "blob" | "about" => return true,
        _ => return false,
    }
    let Some(host) = u.host_str().map(|h| h.trim_matches(['[', ']']).to_string()) else { return false };
    if let Some(&ok) = cache.lock().unwrap().get(&host) {
        return ok;
    }
    let port = u.port_or_known_default().unwrap_or(80);
    let ok = match host.parse::<IpAddr>() {
        Ok(ip) => !is_blocked_ip(ip),
        Err(_) => match tokio::time::timeout(Duration::from_secs(5), tokio::net::lookup_host((host.as_str(), port))).await {
            Ok(Ok(addrs)) => {
                let addrs: Vec<_> = addrs.collect();
                !addrs.is_empty() && addrs.iter().all(|a| !is_blocked_ip(a.ip()))
            }
            _ => false,
        },
    };
    cache.lock().unwrap().insert(host, ok);
    ok
}

type Pending = Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>>;

/// Minimal CDP client on one page target.
struct Cdp {
    out: mpsc::UnboundedSender<String>,
    pending: Pending,
    next: Arc<AtomicU64>,
}

impl Cdp {
    async fn call(&self, method: &str, params: Value) -> Result<Value> {
        call(&self.out, &self.pending, &self.next, method, params).await
    }
}

async fn call(out: &mpsc::UnboundedSender<String>, pending: &Pending, next: &AtomicU64, method: &str, params: Value) -> Result<Value> {
    let id = next.fetch_add(1, Ordering::Relaxed);
    let (tx, rx) = oneshot::channel();
    pending.lock().unwrap().insert(id, tx);
    out.send(json!({"id": id, "method": method, "params": params}).to_string()).map_err(|_| anyhow!("browser connection closed"))?;
    let v = tokio::time::timeout(Duration::from_secs(30), rx).await.map_err(|_| anyhow!("{method}: browser did not answer"))?.map_err(|_| anyhow!("browser connection closed"))?;
    if let Some(e) = v.get("error") {
        bail!("{method}: {}", e["message"].as_str().unwrap_or("error"));
    }
    Ok(v["result"].clone())
}

#[derive(Default)]
struct PageLog {
    console: Vec<String>,
    blocked: Vec<String>,
    loads: u64,
}

pub struct Browser;

const MAX_STEPS: usize = 40;

#[async_trait]
impl Tool for Browser {
    fn name(&self) -> &str {
        "browser"
    }
    fn description(&self) -> String {
        "Use a headless Chromium like a user: open a URL (public site or http://localhost:PORT), run \
         steps (click, type, press keys, wait, run JavaScript, navigate), then get a screenshot, the \
         visible text and console errors. Each call starts a fresh browser – put the whole flow \
         (e.g. log in, open a page, check it) into one call. Private network addresses are blocked."
            .into()
    }
    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "url": {"type": "string", "description": "Start page (http, https or file:// inside your folders)"},
                "steps": {
                    "type": "array",
                    "description": "Executed in order after the page loaded",
                    "items": {
                        "type": "object",
                        "properties": {
                            "action": {"type": "string", "enum": ["click", "type", "press", "wait", "wait_for", "eval", "navigate", "screenshot"]},
                            "selector": {"type": "string", "description": "CSS selector (click, type, wait_for)"},
                            "text": {"type": "string", "description": "Text to type (type); with submit=true Enter is pressed afterwards"},
                            "submit": {"type": "boolean"},
                            "key": {"type": "string", "description": "Enter, Tab, Escape, ArrowDown … (press)"},
                            "ms": {"type": "integer", "description": "Milliseconds (wait), max 15000"},
                            "js": {"type": "string", "description": "JavaScript expression; the (awaited) result is returned (eval)"},
                            "url": {"type": "string", "description": "navigate"}
                        },
                        "required": ["action"]
                    }
                },
                "width": {"type": "integer", "description": "Viewport width, default 1280"},
                "height": {"type": "integer", "description": "Viewport height, default 800"},
                "full_page": {"type": "boolean", "description": "Screenshot of the whole page instead of the viewport"},
                "text": {"type": "boolean", "description": "Include the visible text (default true)"}
            },
            "required": ["url"]
        })
    }
    fn title(&self, input: &Value) -> String {
        let n = input["steps"].as_array().map(|s| s.len()).unwrap_or(0);
        format!("Browser {}{}", input["url"].as_str().unwrap_or("?"), if n > 0 { format!(" ({n} steps)") } else { String::new() })
    }
    async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
        let url = arg_str(&input, "url")?.to_string();
        let steps = input["steps"].as_array().cloned().unwrap_or_default();
        if steps.len() > MAX_STEPS {
            bail!("at most {MAX_STEPS} steps per call");
        }
        let clean_room = ctx.agent == "craft" && ctx.clean_room;
        if clean_room && steps.iter().any(|s| s["action"] == "eval") {
            return Ok(ToolOutput::error("Clean room mode: running JavaScript on the page is not allowed (it could read the source)."));
        }
        if let Some(path) = url.strip_prefix("file://") {
            ctx.resolve(path)?;
        } else if !(url.starts_with("http://") || url.starts_with("https://")) {
            bail!("only http(s) and file:// URLs are supported");
        }
        let cache = Arc::new(Mutex::new(HashMap::new()));
        if !url.starts_with("file://") && !host_allowed(&url, &cache).await {
            bail!("blocked: {url} points to a private network address (only public hosts and localhost are allowed)");
        }
        let fut = session(ctx, &input, &url, &steps, cache);
        tokio::select! {
            r = tokio::time::timeout(Duration::from_secs(180), fut) => r.map_err(|_| anyhow!("the browser session took longer than 3 minutes"))?,
            _ = ctx.cancel.cancelled() => bail!("aborted"),
        }
    }
}

async fn session(ctx: &ToolCtx, input: &Value, url: &str, steps: &[Value], cache: Arc<Mutex<HashMap<String, bool>>>) -> Result<ToolOutput> {
    let profile = std::env::temp_dir().join(format!("cc-browser-{}", crate::types::new_id()));
    let bin = ctx.app.config().craft.chromium.clone();
    let mut child = Command::new(&bin)
        .args(["--headless=new", "--no-sandbox", "--disable-gpu", "--hide-scrollbars", "--mute-audio", "--disable-dev-shm-usage", "--no-first-run", "--no-default-browser-check", "--remote-debugging-port=0"])
        .arg(format!("--user-data-dir={}", profile.display()))
        .arg("about:blank")
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::piped())
        .kill_on_drop(true)
        .spawn()
        .with_context(|| format!("could not start '{bin}' (set craft.chromium in the config)"))?;
    let result = drive(&mut child, input, url, steps, cache).await;
    let _ = child.kill().await;
    let _ = tokio::fs::remove_dir_all(&profile).await;
    result
}

async fn drive(child: &mut tokio::process::Child, input: &Value, url: &str, steps: &[Value], cache: Arc<Mutex<HashMap<String, bool>>>) -> Result<ToolOutput> {
    // "DevTools listening on ws://127.0.0.1:PORT/devtools/browser/…"
    let stderr = child.stderr.take().ok_or_else(|| anyhow!("no stderr"))?;
    let mut lines = BufReader::new(stderr).lines();
    let port = tokio::time::timeout(Duration::from_secs(20), async {
        while let Some(l) = lines.next_line().await? {
            if let Some(rest) = l.split("ws://127.0.0.1:").nth(1) {
                return Ok::<_, anyhow::Error>(rest.split('/').next().unwrap_or("").to_string());
            }
        }
        bail!("chromium exited before DevTools was ready")
    })
    .await
    .map_err(|_| anyhow!("chromium did not start"))??;
    tokio::spawn(async move { while let Ok(Some(_)) = lines.next_line().await {} });

    let list: Value = reqwest::get(format!("http://127.0.0.1:{port}/json/list")).await?.json().await?;
    let ws_url = list.as_array().and_then(|a| a.iter().find(|t| t["type"] == "page")).and_then(|t| t["webSocketDebuggerUrl"].as_str()).ok_or_else(|| anyhow!("no browser page"))?.to_string();
    let (ws, _) = tokio_tungstenite::connect_async(&ws_url).await.context("DevTools connection")?;
    let (mut sink, mut stream) = ws.split();
    let (out_tx, mut out_rx) = mpsc::unbounded_channel::<String>();
    tokio::spawn(async move {
        while let Some(m) = out_rx.recv().await {
            if sink.send(WsMessage::Text(m.into())).await.is_err() {
                break;
            }
        }
    });
    let pending: Pending = Arc::default();
    let next = Arc::new(AtomicU64::new(1));
    let log = Arc::new(Mutex::new(PageLog::default()));
    {
        let pending = pending.clone();
        let out = out_tx.clone();
        let next = next.clone();
        let log = log.clone();
        let cache = cache.clone();
        tokio::spawn(async move {
            while let Some(Ok(msg)) = stream.next().await {
                let WsMessage::Text(text) = msg else { continue };
                let Ok(v) = serde_json::from_str::<Value>(&text) else { continue };
                if let Some(id) = v["id"].as_u64() {
                    if let Some(tx) = pending.lock().unwrap().remove(&id) {
                        let _ = tx.send(v);
                    }
                    continue;
                }
                let p = &v["params"];
                match v["method"].as_str().unwrap_or("") {
                    // Every request passes here: block private addresses.
                    "Fetch.requestPaused" => {
                        let (out, pending, next, cache, log) = (out.clone(), pending.clone(), next.clone(), cache.clone(), log.clone());
                        let req_id = p["requestId"].as_str().unwrap_or("").to_string();
                        let req_url = p["request"]["url"].as_str().unwrap_or("").to_string();
                        tokio::spawn(async move {
                            if host_allowed(&req_url, &cache).await || req_url.starts_with("file://") {
                                let _ = call(&out, &pending, &next, "Fetch.continueRequest", json!({"requestId": req_id})).await;
                            } else {
                                log.lock().unwrap().blocked.push(req_url);
                                let _ = call(&out, &pending, &next, "Fetch.failRequest", json!({"requestId": req_id, "errorReason": "BlockedByClient"})).await;
                            }
                        });
                    }
                    "Runtime.consoleAPICalled" => {
                        let kind = p["type"].as_str().unwrap_or("");
                        if matches!(kind, "error" | "warning" | "assert") {
                            let args: Vec<String> = p["args"].as_array().into_iter().flatten().map(|a| a["value"].as_str().map(String::from).unwrap_or_else(|| a["description"].as_str().unwrap_or("").to_string())).collect();
                            log.lock().unwrap().console.push(format!("console.{kind}: {}", args.join(" ")));
                        }
                    }
                    "Runtime.exceptionThrown" => {
                        let d = &p["exceptionDetails"];
                        let msg = d["exception"]["description"].as_str().or(d["text"].as_str()).unwrap_or("exception");
                        log.lock().unwrap().console.push(format!("uncaught: {}", msg.lines().next().unwrap_or(msg)));
                    }
                    "Log.entryAdded" => {
                        let e = &p["entry"];
                        if e["level"] == "error" {
                            log.lock().unwrap().console.push(format!("{}: {} {}", e["source"].as_str().unwrap_or("log"), e["text"].as_str().unwrap_or(""), e["url"].as_str().unwrap_or("")));
                        }
                    }
                    "Page.loadEventFired" => log.lock().unwrap().loads += 1,
                    _ => {}
                }
            }
        });
    }
    let cdp = Cdp { out: out_tx, pending, next };
    let width = input["width"].as_u64().unwrap_or(1280).clamp(320, 2560);
    let height = input["height"].as_u64().unwrap_or(800).clamp(320, 2400);
    for (m, p) in [
        ("Page.enable", json!({})),
        ("Runtime.enable", json!({})),
        ("Log.enable", json!({})),
        ("Fetch.enable", json!({"patterns": [{"urlPattern": "*"}]})),
        ("Emulation.setDeviceMetricsOverride", json!({"width": width, "height": height, "deviceScaleFactor": 1, "mobile": width < 700})),
    ] {
        cdp.call(m, p).await?;
    }

    let navigate = |target: String| {
        let cdp = &cdp;
        let log = log.clone();
        async move {
            let before = log.lock().unwrap().loads;
            let r = cdp.call("Page.navigate", json!({"url": target})).await?;
            if let Some(e) = r["errorText"].as_str() {
                bail!("could not open {target}: {e}");
            }
            wait_load(&log, before, 20_000).await;
            Ok::<_, anyhow::Error>(())
        }
    };
    navigate(url.to_string()).await?;

    let eval = |js: String| {
        let cdp = &cdp;
        async move {
            let r = cdp.call("Runtime.evaluate", json!({"expression": js, "awaitPromise": true, "returnByValue": true, "userGesture": true})).await?;
            if let Some(d) = r["exceptionDetails"].as_object() {
                bail!("{}", d.get("exception").and_then(|e| e["description"].as_str()).or(d.get("text").and_then(|t| t.as_str())).unwrap_or("JavaScript error"));
            }
            Ok::<_, anyhow::Error>(r["result"]["value"].clone())
        }
    };

    let mut report = Vec::new();
    let mut images = Vec::new();
    for (i, s) in steps.iter().enumerate() {
        let n = i + 1;
        let sel = s["selector"].as_str().unwrap_or("");
        let sel_js = serde_json::to_string(sel)?;
        let before = log.lock().unwrap().loads;
        let res: Result<String> = async {
            match s["action"].as_str().unwrap_or("") {
                "click" => {
                    let r = eval(format!("(()=>{{const e=document.querySelector({sel_js});if(!e)return 'not found';e.scrollIntoView({{block:'center'}});e.click();return 'ok'}})()")).await?;
                    if r != "ok" {
                        bail!("no element matches {sel}");
                    }
                    settle(&log, before).await;
                    Ok(format!("clicked {sel}"))
                }
                "type" => {
                    let r = eval(format!("(()=>{{const e=document.querySelector({sel_js});if(!e)return 'not found';e.focus();if('select' in e)e.select();return 'ok'}})()")).await?;
                    if r != "ok" {
                        bail!("no element matches {sel}");
                    }
                    let text = s["text"].as_str().unwrap_or("");
                    cdp.call("Input.insertText", json!({"text": text})).await?;
                    if s["submit"].as_bool().unwrap_or(false) {
                        press(&cdp, "Enter").await?;
                        settle(&log, before).await;
                    }
                    Ok(format!("typed {} characters into {sel}", text.chars().count()))
                }
                "press" => {
                    let key = s["key"].as_str().unwrap_or("Enter");
                    press(&cdp, key).await?;
                    settle(&log, before).await;
                    Ok(format!("pressed {key}"))
                }
                "wait" => {
                    let ms = s["ms"].as_u64().unwrap_or(1000).min(15_000);
                    tokio::time::sleep(Duration::from_millis(ms)).await;
                    Ok(format!("waited {ms} ms"))
                }
                "wait_for" => {
                    for _ in 0..50 {
                        if eval(format!("!!document.querySelector({sel_js})")).await? == true {
                            return Ok(format!("{sel} is there"));
                        }
                        tokio::time::sleep(Duration::from_millis(200)).await;
                    }
                    bail!("{sel} did not appear within 10 s")
                }
                "eval" => {
                    let v = eval(s["js"].as_str().unwrap_or("undefined").to_string()).await?;
                    Ok(format!("eval → {}", truncate(&serde_json::to_string(&v).unwrap_or_default(), 4000)))
                }
                "navigate" => {
                    let target = s["url"].as_str().ok_or_else(|| anyhow!("navigate needs url"))?.to_string();
                    if !host_allowed(&target, &cache).await {
                        bail!("blocked: {target} points to a private network address");
                    }
                    navigate(target.clone()).await?;
                    Ok(format!("opened {target}"))
                }
                "screenshot" => {
                    images.push(screenshot(&cdp, false).await?);
                    Ok("screenshot taken".into())
                }
                other => bail!("unknown action '{other}'"),
            }
        }
        .await;
        match res {
            Ok(r) => report.push(format!("{n}. {r}")),
            Err(e) => {
                report.push(format!("{n}. FAILED: {e:#} – remaining steps skipped"));
                break;
            }
        }
    }

    let page = eval("({t:document.title,u:location.href,x:document.body?document.body.innerText:''})".into()).await.unwrap_or(Value::Null);
    if images.len() < 4 {
        images.push(screenshot(&cdp, input["full_page"].as_bool().unwrap_or(false)).await?);
    }
    let mut text = format!("Page: {} – {}\n", page["t"].as_str().unwrap_or(""), page["u"].as_str().unwrap_or(url));
    if !report.is_empty() {
        text.push_str(&format!("\nSteps:\n{}\n", report.join("\n")));
    }
    let log = std::mem::take(&mut *log.lock().unwrap());
    if !log.console.is_empty() {
        text.push_str(&format!("\nConsole errors ({}):\n{}\n", log.console.len(), truncate(&log.console.join("\n"), 6000)));
    }
    if !log.blocked.is_empty() {
        text.push_str(&format!("\nBlocked requests to private addresses: {}\n", log.blocked.iter().take(10).cloned().collect::<Vec<_>>().join(", ")));
    }
    if input["text"].as_bool().unwrap_or(true) {
        text.push_str(&format!("\nVisible text:\n{}", truncate(page["x"].as_str().unwrap_or(""), 15_000)));
    }
    Ok(ToolOutput { content: text, images, ..Default::default() })
}

/// Wait until a new load event fired (or the timeout passed).
async fn wait_load(log: &Mutex<PageLog>, before: u64, ms: u64) {
    let deadline = tokio::time::Instant::now() + Duration::from_millis(ms);
    while log.lock().unwrap().loads <= before && tokio::time::Instant::now() < deadline {
        tokio::time::sleep(Duration::from_millis(100)).await;
    }
    tokio::time::sleep(Duration::from_millis(400)).await;
}

/// After an interaction: give a navigation the chance to start and finish.
async fn settle(log: &Mutex<PageLog>, before: u64) {
    tokio::time::sleep(Duration::from_millis(400)).await;
    if log.lock().unwrap().loads > before {
        return wait_load(log, before, 15_000).await;
    }
    wait_load(log, before, 800).await;
}

async fn press(cdp: &Cdp, key: &str) -> Result<()> {
    let (code, vk, text) = match key {
        "Enter" => ("Enter", 13, "\r"),
        "Tab" => ("Tab", 9, ""),
        "Escape" => ("Escape", 27, ""),
        "Backspace" => ("Backspace", 8, ""),
        "ArrowDown" => ("ArrowDown", 40, ""),
        "ArrowUp" => ("ArrowUp", 38, ""),
        "ArrowLeft" => ("ArrowLeft", 37, ""),
        "ArrowRight" => ("ArrowRight", 39, ""),
        "Space" | " " => ("Space", 32, " "),
        other => bail!("unsupported key '{other}'"),
    };
    let key_name = if code == "Space" { " " } else { code };
    let mut down = json!({"type": "keyDown", "key": key_name, "code": code, "windowsVirtualKeyCode": vk});
    if !text.is_empty() {
        down["text"] = json!(text);
    }
    cdp.call("Input.dispatchKeyEvent", down).await?;
    cdp.call("Input.dispatchKeyEvent", json!({"type": "keyUp", "key": key_name, "code": code, "windowsVirtualKeyCode": vk})).await?;
    Ok(())
}

async fn screenshot(cdp: &Cdp, full: bool) -> Result<crate::types::ImageData> {
    let r = cdp.call("Page.captureScreenshot", json!({"format": "png", "captureBeyondViewport": full})).await?;
    let bytes = base64::engine::general_purpose::STANDARD.decode(r["data"].as_str().unwrap_or(""))?;
    let img = image::load_from_memory(&bytes)?;
    super::craft::png_for_model(&img)
}

#[cfg(test)]
mod tests {
    use super::is_blocked_ip;

    #[test]
    fn private_ranges() {
        for ip in ["10.1.2.3", "192.168.1.10", "172.17.0.1", "169.254.169.254", "100.64.0.1", "fd00::1", "fe80::1", "::ffff:192.168.0.1"] {
            assert!(is_blocked_ip(ip.parse().unwrap()), "{ip}");
        }
        for ip in ["127.0.0.1", "::1", "8.8.8.8", "1.1.1.1", "2606:4700::1111"] {
            assert!(!is_blocked_ip(ip.parse().unwrap()), "{ip}");
        }
    }
}