hyperhive/hive-screen-mcp/src/main.rs

//! `hive-screen-mcp` — stdio MCP bridge for GUI agents.
//!
//! Provides screenshot, keyboard-input, and mouse tools for agents running a
//! Weston Wayland compositor (`services.hyperhive.agent.gui.enable = true`). Each tool:
//!
//! - `screenshot`  → `grim` (always available with gui enabled)
//! - `type_text`   → `wtype` (Wayland virtual-keyboard protocol, no daemon)
//! - `key_press`   → `wtype -k` (same protocol as `type_text`)
//! - `mouse_move`  → RFB `PointerEvent` to the local neatvnc server
//!   (`HIVE_GUI_VNC_PORT`, default 5900)
//! - `mouse_click` → RFB `PointerEvent` sequence: move → button-down →
//!   button-up (all in one TCP connection)
//!
//! All tools operate in userspace. The keyboard tools use the Wayland
//! virtual-keyboard protocol; the mouse tools speak RFB directly to the
//! Weston VNC backend — the compositor's native remote-input path, no
//! `/dev/uinput` or kernel bypass.
//!
//! `WAYLAND_DISPLAY` and `XDG_RUNTIME_DIR` are injected globally by the
//! `weston-vnc` module. `HIVE_GUI_VNC_PORT` is set by the harness service
//! when `services.hyperhive.agent.gui.enable = true` (defaults to 5900 if absent).

use anyhow::Result;
use rmcp::{
    ServerHandler, ServiceExt,
    handler::server::wrapper::Parameters,
    schemars::{self, JsonSchema},
    tool, tool_handler, tool_router,
    transport::stdio,
};
use serde::Deserialize;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::process::Command;

/// Upper bound on one `grim` / `wtype` run. Against a live local compositor
/// these finish in well under a second, a long `type_text` included. The
/// agent's turn waits on the tool call, so a wedged compositor would
/// otherwise hold the turn for as long as it stays wedged.
const CMD_TIMEOUT: Duration = Duration::from_secs(30);

/// Upper bound on one RFB exchange (connect, handshake, pointer events) with
/// the neatvnc server on `127.0.0.1`: a few round trips of a few bytes each.
/// A server that is down refuses the connect at once; this catches one that
/// accepts and then never speaks.
const VNC_TIMEOUT: Duration = Duration::from_secs(10);

/// Run a command and return its result.
/// `Ok(stdout_trimmed)` on success (may be empty string when the tool
/// produces no output). `Err(human_readable_message)` on failure
/// (non-zero exit, failed spawn, or no exit within `limit`, in which case
/// the child is killed).
async fn run_cmd(program: &str, args: &[&str], limit: Duration) -> Result<String, String> {
    let output = Command::new(program).args(args).kill_on_drop(true).output();
    let Ok(result) = tokio::time::timeout(limit, output).await else {
        return Err(format!(
            "`{program}` timed out after {}s and was killed",
            limit.as_secs_f32()
        ));
    };
    match result {
        Ok(out) if out.status.success() => {
            Ok(String::from_utf8_lossy(&out.stdout).trim().to_owned())
        }
        Ok(out) => {
            let msg = String::from_utf8_lossy(&out.stderr).trim().to_owned();
            Err(format!(
                "command failed (exit {}): {}",
                out.status.code().unwrap_or(-1),
                if msg.is_empty() { "(no stderr)" } else { &msg }
            ))
        }
        Err(e) => Err(format!("failed to launch `{program}`: {e}")),
    }
}

/// Format a `run_cmd` result as a tool-response string. `Ok("")` →
/// `"ok"` (most input tools produce no stdout); `Ok(s)` → `s`;
/// `Err(e)` → `e`.
fn cmd_result(r: Result<String, String>) -> String {
    match r {
        Ok(s) if s.is_empty() => "ok".to_owned(),
        Ok(s) => s,
        Err(e) => e,
    }
}

/// Read `HIVE_GUI_VNC_PORT` from the environment, falling back to 5900.
fn vnc_port() -> u16 {
    std::env::var("HIVE_GUI_VNC_PORT")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(5900)
}

/// Perform the RFB 3.8 handshake on `stream` and return on success.
///
/// Assumes the server advertises security type 1 (None) as set by
/// `auth-method=none` in `weston.ini`. Consumes the full `ServerInit`
/// frame so the caller is ready to send client messages.
async fn rfb_handshake(stream: &mut TcpStream) -> Result<(), String> {
    // Protocol version — server sends 12 bytes ("RFB 003.xxx\n"), we
    // reply with 3.8.
    let mut ver = [0u8; 12];
    stream
        .read_exact(&mut ver)
        .await
        .map_err(|e| format!("VNC: read version: {e}"))?;
    stream
        .write_all(b"RFB 003.008\n")
        .await
        .map_err(|e| format!("VNC: send version: {e}"))?;

    // Security type negotiation (RFB 3.8): server sends count + list.
    let count = stream
        .read_u8()
        .await
        .map_err(|e| format!("VNC: read sec-type count: {e}"))?;
    let mut types = vec![0u8; count as usize];
    stream
        .read_exact(&mut types)
        .await
        .map_err(|e| format!("VNC: read sec-types: {e}"))?;
    if !types.contains(&1) {
        return Err(format!(
            "VNC: server did not offer security type None (got: {types:?})"
        ));
    }
    stream
        .write_all(&[1])
        .await
        .map_err(|e| format!("VNC: send sec-type: {e}"))?;

    // SecurityResult: 4 bytes big-endian; 0 = OK.
    let result = stream
        .read_u32()
        .await
        .map_err(|e| format!("VNC: read sec-result: {e}"))?;
    if result != 0 {
        return Err(format!(
            "VNC: security handshake rejected (result={result})"
        ));
    }

    // ClientInit: 1 byte; 1 = shared session (don't disconnect other clients).
    stream
        .write_all(&[1])
        .await
        .map_err(|e| format!("VNC: send client-init: {e}"))?;

    // ServerInit: width(2) + height(2) + pixel-format(16) +
    //             name-length(4) + name. Read and discard.
    stream
        .read_u16()
        .await
        .map_err(|e| format!("VNC: read width: {e}"))?;
    stream
        .read_u16()
        .await
        .map_err(|e| format!("VNC: read height: {e}"))?;
    let mut pf = [0u8; 16];
    stream
        .read_exact(&mut pf)
        .await
        .map_err(|e| format!("VNC: read pixel-format: {e}"))?;
    let name_len = stream
        .read_u32()
        .await
        .map_err(|e| format!("VNC: read name-len: {e}"))?;
    if name_len > 256 {
        return Err(format!(
            "VNC: server-init name length {name_len} exceeds cap (256)"
        ));
    }
    let mut name = vec![0u8; name_len as usize];
    stream
        .read_exact(&mut name)
        .await
        .map_err(|e| format!("VNC: read name: {e}"))?;

    Ok(())
}

/// Encode an RFB `PointerEvent` message (6 bytes).
///
/// - `button_mask`: bitmask — bit 0 = left button, bit 1 = middle, bit 2 = right
/// - `x`, `y`: absolute pixel position (big-endian in the wire format)
fn rfb_pointer_event(button_mask: u8, x: u16, y: u16) -> [u8; 6] {
    let xb = x.to_be_bytes();
    let yb = y.to_be_bytes();
    [5, button_mask, xb[0], xb[1], yb[0], yb[1]]
}

/// Connect to the local neatvnc server, perform the RFB handshake, and
/// send a sequence of `PointerEvent` messages in one connection.
///
/// `events` is a slice of `(button_mask, x, y)` tuples. Each is encoded
/// as an RFB `PointerEvent` and written in order; the connection is then
/// closed. Errors at any step short-circuit and return a `String`, as does
/// the whole exchange taking longer than `limit`.
async fn rfb_send_pointer_events(
    port: u16,
    events: &[(u8, u16, u16)],
    limit: Duration,
) -> Result<(), String> {
    tokio::time::timeout(limit, rfb_exchange(port, events))
        .await
        .unwrap_or_else(|_| {
            Err(format!(
                "VNC: no reply from localhost:{port} within {}s",
                limit.as_secs_f32()
            ))
        })
}

async fn rfb_exchange(port: u16, events: &[(u8, u16, u16)]) -> Result<(), String> {
    let mut stream = TcpStream::connect(("127.0.0.1", port))
        .await
        .map_err(|e| format!("VNC: connect to localhost:{port}: {e}"))?;
    rfb_handshake(&mut stream).await?;
    for &(mask, x, y) in events {
        stream
            .write_all(&rfb_pointer_event(mask, x, y))
            .await
            .map_err(|e| format!("VNC: send pointer event: {e}"))?;
    }
    stream
        .flush()
        .await
        .map_err(|e| format!("VNC: flush: {e}"))?;
    Ok(())
}

struct ScreenMcp;

#[tool_router]
impl ScreenMcp {
    #[tool(
        description = "Take a screenshot of the agent's Wayland display and save it as a PNG \
            file under `/tmp/`. Returns the path to the saved file — pass that path to the \
            `Read` tool to view the image visually. Requires `grim` and a live \
            `WAYLAND_DISPLAY` (provided automatically when `services.hyperhive.agent.gui.enable = true`)."
    )]
    async fn screenshot(&self) -> String {
        let ms = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map_or(0, |d| d.as_millis());
        let path = format!("/tmp/hive-screenshot-{ms}.png");
        match run_cmd("grim", &["-t", "png", &path], CMD_TIMEOUT).await {
            Ok(_) => format!("screenshot saved to `{path}` — use the Read tool to view it"),
            Err(e) => e,
        }
    }

    #[tool(
        description = "Type text into the currently focused Wayland window. Uses `wtype` \
            via the Wayland virtual-keyboard protocol (compositor-mediated, no kernel \
            injection). Supports arbitrary Unicode. For special keys (Enter, Tab, Escape, \
            function keys, arrow keys) or modifier combos, use `key_press` instead. \
            Requires a live `WAYLAND_DISPLAY` (provided automatically when \
            `services.hyperhive.agent.gui.enable = true`)."
    )]
    async fn type_text(&self, Parameters(args): Parameters<TypeTextArgs>) -> String {
        cmd_result(run_cmd("wtype", &[&args.text], CMD_TIMEOUT).await)
    }

    #[tool(
        description = "Press a key or key combination via the Wayland virtual-keyboard \
            protocol (compositor-mediated, no kernel injection). `keys` uses XKB keysym \
            syntax: a single key name (`Return`, `Tab`, `Escape`, `BackSpace`, `Delete`, \
            `space`, `F1`–`F12`, `Left`, `Right`, `Up`, `Down`) or a modifier+key combo \
            with `+`-separated parts where all but the last are modifiers (`ctrl+c`, \
            `ctrl+shift+t`, `super+l`, `alt+F4`). Modifier names: `ctrl`, `shift`, `alt`, \
            `super`. Requires a live `WAYLAND_DISPLAY` (provided automatically when \
            `services.hyperhive.agent.gui.enable = true`)."
    )]
    async fn key_press(&self, Parameters(args): Parameters<KeyPressArgs>) -> String {
        // Parse "mod1+mod2+key" into: -M mod1 -M mod2 -k key -m mod2 -m mod1
        // split('+') always yields at least one element, so split_last() is safe.
        let parts: Vec<&str> = args.keys.split('+').collect();
        let (key, mods) = parts.split_last().expect("split yields ≥1 element");
        let mut wtype_args: Vec<String> = Vec::new();
        for &m in mods {
            wtype_args.push("-M".to_owned());
            wtype_args.push(m.to_owned());
        }
        wtype_args.push("-k".to_owned());
        wtype_args.push((*key).to_owned());
        for &m in mods.iter().rev() {
            wtype_args.push("-m".to_owned());
            wtype_args.push(m.to_owned());
        }
        let arg_refs: Vec<&str> = wtype_args.iter().map(String::as_str).collect();
        cmd_result(run_cmd("wtype", &arg_refs, CMD_TIMEOUT).await)
    }

    #[tool(
        description = "Move the mouse cursor to an absolute pixel position on the display. \
            `x` is measured from the left edge, `y` from the top edge. Coordinates are in \
            display pixels — use `screenshot` first to see the current resolution (the \
            display size can vary). Injects a RFB `PointerEvent` into the Weston VNC \
            backend — compositor-mediated, no kernel bypass."
    )]
    async fn mouse_move(&self, Parameters(args): Parameters<MouseMoveArgs>) -> String {
        let (Ok(x), Ok(y)) = (u16::try_from(args.x), u16::try_from(args.y)) else {
            return "mouse_move: x and y must be in range 0–65535".to_owned();
        };
        match rfb_send_pointer_events(vnc_port(), &[(0, x, y)], VNC_TIMEOUT).await {
            Ok(()) => "ok".to_owned(),
            Err(e) => e,
        }
    }

    #[tool(
        description = "Click a mouse button at an absolute pixel position. `x` and `y` are \
            the target coordinates (0-based from the top-left) in display pixels — use \
            `screenshot` first to determine the current resolution (display size can vary). \
            `button` is `\"left\"` (default), `\"right\"`, or `\"middle\"`. Sends: move to \
            position → button-down → button-up as RFB `PointerEvent` messages to the Weston \
            VNC backend — compositor-mediated, no kernel bypass."
    )]
    async fn mouse_click(&self, Parameters(args): Parameters<MouseClickArgs>) -> String {
        let (Ok(x), Ok(y)) = (u16::try_from(args.x), u16::try_from(args.y)) else {
            return "mouse_click: x and y must be in range 0–65535".to_owned();
        };
        let btn_mask: u8 = match args.button.as_deref().unwrap_or("left") {
            "left" => 0b0000_0001,
            "right" => 0b0000_0100,
            "middle" => 0b0000_0010,
            other => {
                return format!(
                    "mouse_click: unknown button {other:?} — use \"left\", \"right\", or \"middle\""
                );
            }
        };
        // Move to position, press, release — all in one VNC connection.
        let events = [(0, x, y), (btn_mask, x, y), (0, x, y)];
        match rfb_send_pointer_events(vnc_port(), &events, VNC_TIMEOUT).await {
            Ok(()) => "ok".to_owned(),
            Err(e) => e,
        }
    }
}

#[derive(Debug, Deserialize, JsonSchema)]
struct TypeTextArgs {
    /// Text to type into the focused window. Supports arbitrary Unicode.
    text: String,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct KeyPressArgs {
    /// Key or key combination in XKB keysym syntax — e.g. `"Return"`, `"Tab"`,
    /// `"ctrl+c"`, `"ctrl+shift+t"`, `"super+l"`, `"alt+F4"`.
    keys: String,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct MouseMoveArgs {
    /// Horizontal pixel coordinate from the left edge of the display (0-based).
    x: i32,
    /// Vertical pixel coordinate from the top edge of the display (0-based).
    y: i32,
}

#[derive(Debug, Deserialize, JsonSchema)]
struct MouseClickArgs {
    /// Horizontal pixel coordinate from the left edge of the display (0-based).
    x: i32,
    /// Vertical pixel coordinate from the top edge of the display (0-based).
    y: i32,
    /// Button to click: `"left"` (default), `"right"`, or `"middle"`.
    button: Option<String>,
}

#[tool_handler]
impl ServerHandler for ScreenMcp {}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_env("RUST_LOG")
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn")),
        )
        .with_writer(std::io::stderr)
        // This is a systemd-managed daemon — stderr always goes to journald,
        // never a human terminal, and journald doesn't strip ANSI escapes:
        // they land in victorialogs as raw byte-array spam otherwise.
        .with_ansi(false)
        .init();

    let service = ScreenMcp.serve(stdio()).await?;
    service.waiting().await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::{Duration, rfb_send_pointer_events, run_cmd};
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::net::TcpListener;

    /// A child that would create `marker` half a second after it starts.
    fn marker_cmd(marker: &std::path::Path) -> String {
        format!("sleep 0.5; touch '{}'", marker.display())
    }

    fn marker_path(test: &str) -> std::path::PathBuf {
        let path =
            std::env::temp_dir().join(format!("hive-screen-mcp-{test}-{}", std::process::id()));
        let _ = std::fs::remove_file(&path);
        path
    }

    #[tokio::test]
    async fn run_cmd_times_out_and_kills_the_child() {
        let marker = marker_path("killed");
        let err = run_cmd(
            "sh",
            &["-c", &marker_cmd(&marker)],
            Duration::from_millis(50),
        )
        .await
        .expect_err("a child running past the limit is an error");
        assert!(err.contains("timed out after"), "{err}");
        tokio::time::sleep(Duration::from_secs(1)).await;
        assert!(!marker.exists(), "the timed-out child kept running");
    }

    #[tokio::test]
    async fn run_cmd_waits_for_a_child_within_the_limit() {
        let marker = marker_path("finished");
        run_cmd("sh", &["-c", &marker_cmd(&marker)], Duration::from_secs(10))
            .await
            .expect("a child inside the limit succeeds");
        assert!(marker.exists());
        let _ = std::fs::remove_file(&marker);
    }

    #[tokio::test]
    async fn rfb_times_out_on_a_server_that_never_speaks() {
        let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
        let port = listener.local_addr().unwrap().port();
        let server = tokio::spawn(async move {
            let (conn, _) = listener.accept().await.unwrap();
            tokio::time::sleep(Duration::from_secs(5)).await;
            drop(conn);
        });
        let err = rfb_send_pointer_events(port, &[(0, 1, 1)], Duration::from_millis(200))
            .await
            .expect_err("a silent server is an error");
        assert!(err.contains("no reply"), "{err}");
        server.abort();
    }

    #[tokio::test]
    async fn rfb_completes_against_a_server_that_answers() {
        let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
        let port = listener.local_addr().unwrap().port();
        let server = tokio::spawn(async move {
            let (mut conn, _) = listener.accept().await.unwrap();
            let mut buf = [0u8; 12];
            conn.write_all(b"RFB 003.008\n").await.unwrap();
            conn.read_exact(&mut buf).await.unwrap();
            conn.write_all(&[1, 1]).await.unwrap();
            conn.read_exact(&mut buf[..1]).await.unwrap();
            conn.write_all(&0u32.to_be_bytes()).await.unwrap();
            conn.read_exact(&mut buf[..1]).await.unwrap();
            // ServerInit: width, height, pixel format, empty name.
            conn.write_all(&[0u8; 2 + 2 + 16 + 4]).await.unwrap();
            let mut event = [0u8; 6];
            conn.read_exact(&mut event).await.unwrap();
            event
        });
        rfb_send_pointer_events(port, &[(0, 1, 2)], Duration::from_secs(10))
            .await
            .expect("a well-behaved server completes");
        assert_eq!(server.await.unwrap(), [5, 0, 0, 1, 0, 2]);
    }
}