wraith/src-tauri/src/bin/wraith_mcp_bridge.rs

//! Wraith MCP Bridge — stdio JSON-RPC proxy to Wraith's HTTP API.
//!
//! This binary is spawned by AI CLIs (Claude Code, Gemini CLI) as an MCP
//! server. It reads JSON-RPC requests from stdin, translates them to HTTP
//! calls against the running Wraith instance, and writes responses to stdout.
//!
//! The Wraith instance's MCP HTTP port is read from the data directory's
//! `mcp-port` file.

use std::io::{self, BufRead, Write};

use serde::{Deserialize, Serialize};
use serde_json::Value;

#[derive(Deserialize)]
#[allow(dead_code)]
struct JsonRpcRequest {
    jsonrpc: String,
    id: Value,
    method: String,
    #[serde(default)]
    params: Value,
}

#[derive(Serialize)]
struct JsonRpcResponse {
    jsonrpc: String,
    id: Value,
    #[serde(skip_serializing_if = "Option::is_none")]
    result: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    error: Option<JsonRpcError>,
}

#[derive(Serialize)]
struct JsonRpcError {
    code: i32,
    message: String,
}

fn get_mcp_port() -> Result<u16, String> {
    // Check standard locations for the port file
    let port_file = if let Ok(appdata) = std::env::var("APPDATA") {
        std::path::PathBuf::from(appdata).join("Wraith").join("mcp-port")
    } else if let Ok(home) = std::env::var("HOME") {
        if cfg!(target_os = "macos") {
            std::path::PathBuf::from(home).join("Library").join("Application Support").join("Wraith").join("mcp-port")
        } else {
            std::path::PathBuf::from(home).join(".local").join("share").join("wraith").join("mcp-port")
        }
    } else {
        return Err("Cannot determine data directory".to_string());
    };

    let port_str = std::fs::read_to_string(&port_file)
        .map_err(|e| format!("Cannot read MCP port file at {}: {} — is Wraith running?", port_file.display(), e))?;

    port_str.trim().parse::<u16>()
        .map_err(|e| format!("Invalid port in MCP port file: {}", e))
}

fn handle_initialize(id: Value) -> JsonRpcResponse {
    JsonRpcResponse {
        jsonrpc: "2.0".to_string(),
        id,
        result: Some(serde_json::json!({
            "protocolVersion": "2024-11-05",
            "capabilities": {
                "tools": {}
            },
            "serverInfo": {
                "name": "wraith-terminal",
                "version": "1.0.0"
            }
        })),
        error: None,
    }
}

fn handle_tools_list(id: Value) -> JsonRpcResponse {
    JsonRpcResponse {
        jsonrpc: "2.0".to_string(),
        id,
        result: Some(serde_json::json!({
            "tools": [
                {
                    "name": "terminal_read",
                    "description": "Read recent terminal output from an active SSH or PTY session (ANSI codes stripped)",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "session_id": { "type": "string", "description": "The session ID to read from. Use list_sessions to find IDs." },
                            "lines": { "type": "number", "description": "Number of recent lines to return (default: 50)" }
                        },
                        "required": ["session_id"]
                    }
                },
                {
                    "name": "terminal_execute",
                    "description": "Execute a command in an active SSH session and return the output",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "session_id": { "type": "string", "description": "The SSH session ID to execute in" },
                            "command": { "type": "string", "description": "The command to run" },
                            "timeout_ms": { "type": "number", "description": "Max wait time in ms (default: 5000)" }
                        },
                        "required": ["session_id", "command"]
                    }
                },
                {
                    "name": "terminal_screenshot",
                    "description": "Capture a screenshot of an active RDP session as a base64-encoded PNG image for visual analysis",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "session_id": { "type": "string", "description": "The RDP session ID to screenshot" }
                        },
                        "required": ["session_id"]
                    }
                },
                {
                    "name": "sftp_list",
                    "description": "List files in a directory on a remote host via SFTP",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "session_id": { "type": "string", "description": "The SSH session ID" },
                            "path": { "type": "string", "description": "Remote directory path" }
                        },
                        "required": ["session_id", "path"]
                    }
                },
                {
                    "name": "sftp_read",
                    "description": "Read a file from a remote host via SFTP",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "session_id": { "type": "string", "description": "The SSH session ID" },
                            "path": { "type": "string", "description": "Remote file path" }
                        },
                        "required": ["session_id", "path"]
                    }
                },
                {
                    "name": "sftp_write",
                    "description": "Write content to a file on a remote host via SFTP",
                    "inputSchema": {
                        "type": "object",
                        "properties": {
                            "session_id": { "type": "string", "description": "The SSH session ID" },
                            "path": { "type": "string", "description": "Remote file path" },
                            "content": { "type": "string", "description": "File content to write" }
                        },
                        "required": ["session_id", "path", "content"]
                    }
                },
                {
                    "name": "list_sessions",
                    "description": "List all active Wraith sessions (SSH, RDP, PTY) with connection details",
                    "inputSchema": {
                        "type": "object",
                        "properties": {}
                    }
                }
            ]
        })),
        error: None,
    }
}

fn call_wraith(port: u16, endpoint: &str, body: Value) -> Result<Value, String> {
    let url = format!("http://127.0.0.1:{}{}", port, endpoint);
    let body_str = serde_json::to_string(&body).unwrap_or_default();

    let mut resp = ureq::post(url)
        .header("Content-Type", "application/json")
        .send(body_str.as_bytes())
        .map_err(|e| format!("HTTP request to Wraith failed: {}", e))?;

    let resp_str = resp.body_mut().read_to_string()
        .map_err(|e| format!("Failed to read Wraith response: {}", e))?;

    let json: Value = serde_json::from_str(&resp_str)
        .map_err(|e| format!("Failed to parse Wraith response: {}", e))?;

    if json.get("ok").and_then(|v| v.as_bool()) == Some(true) {
        Ok(json.get("data").cloned().unwrap_or(Value::Null))
    } else {
        let err_msg = json.get("error").and_then(|e| e.as_str()).unwrap_or("Unknown error");
        Err(err_msg.to_string())
    }
}

fn handle_tool_call(id: Value, port: u16, tool_name: &str, args: &Value) -> JsonRpcResponse {
    let result = match tool_name {
        "list_sessions" => call_wraith(port, "/mcp/sessions", serde_json::json!({})),
        "terminal_read" => call_wraith(port, "/mcp/terminal/read", args.clone()),
        "terminal_execute" => call_wraith(port, "/mcp/terminal/execute", args.clone()),
        "sftp_list" => call_wraith(port, "/mcp/sftp/list", args.clone()),
        "sftp_read" => call_wraith(port, "/mcp/sftp/read", args.clone()),
        "sftp_write" => call_wraith(port, "/mcp/sftp/write", args.clone()),
        "terminal_screenshot" => {
            let result = call_wraith(port, "/mcp/screenshot", args.clone());
            // Screenshot returns base64 PNG — wrap as image content for multimodal AI
            return match result {
                Ok(b64) => JsonRpcResponse {
                    jsonrpc: "2.0".to_string(),
                    id,
                    result: Some(serde_json::json!({
                        "content": [{
                            "type": "image",
                            "data": b64,
                            "mimeType": "image/png"
                        }]
                    })),
                    error: None,
                },
                Err(e) => JsonRpcResponse {
                    jsonrpc: "2.0".to_string(),
                    id,
                    result: None,
                    error: Some(JsonRpcError { code: -32000, message: e }),
                },
            };
        }
        _ => Err(format!("Unknown tool: {}", tool_name)),
    };

    match result {
        Ok(data) => JsonRpcResponse {
            jsonrpc: "2.0".to_string(),
            id,
            result: Some(serde_json::json!({
                "content": [{
                    "type": "text",
                    "text": if data.is_string() {
                        data.as_str().unwrap().to_string()
                    } else {
                        serde_json::to_string_pretty(&data).unwrap_or_default()
                    }
                }]
            })),
            error: None,
        },
        Err(e) => JsonRpcResponse {
            jsonrpc: "2.0".to_string(),
            id,
            result: None,
            error: Some(JsonRpcError { code: -32000, message: e }),
        },
    }
}

fn main() {
    let port = match get_mcp_port() {
        Ok(p) => p,
        Err(e) => {
            eprintln!("wraith-mcp-bridge: {}", e);
            std::process::exit(1);
        }
    };

    let stdin = io::stdin();
    let mut stdout = io::stdout();

    for line in stdin.lock().lines() {
        let line = match line {
            Ok(l) => l,
            Err(_) => break,
        };

        if line.trim().is_empty() {
            continue;
        }

        let request: JsonRpcRequest = match serde_json::from_str(&line) {
            Ok(r) => r,
            Err(e) => {
                let err_resp = JsonRpcResponse {
                    jsonrpc: "2.0".to_string(),
                    id: Value::Null,
                    result: None,
                    error: Some(JsonRpcError { code: -32700, message: format!("Parse error: {}", e) }),
                };
                let _ = writeln!(stdout, "{}", serde_json::to_string(&err_resp).unwrap());
                let _ = stdout.flush();
                continue;
            }
        };

        let response = match request.method.as_str() {
            "initialize" => handle_initialize(request.id),
            "tools/list" => handle_tools_list(request.id),
            "tools/call" => {
                let tool_name = request.params.get("name")
                    .and_then(|v| v.as_str())
                    .unwrap_or("");
                let args = request.params.get("arguments")
                    .cloned()
                    .unwrap_or(Value::Object(serde_json::Map::new()));
                handle_tool_call(request.id, port, tool_name, &args)
            }
            "notifications/initialized" | "notifications/cancelled" => {
                // Notifications don't get responses
                continue;
            }
            _ => JsonRpcResponse {
                jsonrpc: "2.0".to_string(),
                id: request.id,
                result: None,
                error: Some(JsonRpcError { code: -32601, message: format!("Method not found: {}", request.method) }),
            },
        };

        let _ = writeln!(stdout, "{}", serde_json::to_string(&response).unwrap());
        let _ = stdout.flush();
    }
}