From add0f0628f2157094ab20dd2f335faf997c37726 Mon Sep 17 00:00:00 2001 From: Vantz Stockwell Date: Tue, 24 Mar 2026 23:17:36 -0400 Subject: [PATCH] feat: MCP auto-inject + RDP screenshot tool - Auto-inject CLAUDE_MCP_SERVERS env var when copilot PTY spawns, so Claude Code auto-discovers wraith-mcp-bridge without manual config - RDP screenshot_png_base64() encodes frame buffer as PNG via png crate - Bridge binary exposes terminal_screenshot tool returning MCP image content (base64 PNG with mimeType) for multimodal AI analysis - MCP session list now includes RDP sessions with dimensions - /mcp/screenshot HTTP endpoint on the internal server "Screenshot that RDP session, what's the error?" now works. Co-Authored-By: Claude Opus 4.6 (1M context) --- src-tauri/Cargo.lock | 1 + src-tauri/Cargo.toml | 1 + src-tauri/src/bin/wraith_mcp_bridge.rs | 35 +++++++++++++++++++++++++ src-tauri/src/lib.rs | 3 ++- src-tauri/src/mcp/server.rs | 36 ++++++++++++++++++++++++-- src-tauri/src/pty/mod.rs | 6 +++-- src-tauri/src/rdp/mod.rs | 22 ++++++++++++++++ 7 files changed, 99 insertions(+), 5 deletions(-) diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index f42e7de..86237cc 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -8894,6 +8894,7 @@ dependencies = [ "md5", "pem", "pkcs8 0.10.2", + "png", "portable-pty", "rand 0.9.2", "reqwest 0.12.28", diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index 1f03188..7f8746d 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -55,6 +55,7 @@ portable-pty = "0.8" # MCP HTTP server (for bridge binary communication) axum = "0.8" ureq = "3" +png = "0.17" # RDP (IronRDP) ironrdp = { version = "0.14", features = ["connector", "session", "graphics", "input"] } diff --git a/src-tauri/src/bin/wraith_mcp_bridge.rs b/src-tauri/src/bin/wraith_mcp_bridge.rs index 7c2f041..51e8306 100644 --- a/src-tauri/src/bin/wraith_mcp_bridge.rs +++ b/src-tauri/src/bin/wraith_mcp_bridge.rs @@ -108,6 +108,17 @@ fn handle_tools_list(id: Value) -> JsonRpcResponse { "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": "list_sessions", "description": "List all active Wraith sessions (SSH, RDP, PTY) with connection details", @@ -150,6 +161,30 @@ fn handle_tool_call(id: Value, port: u16, tool_name: &str, args: &Value) -> Json "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()), + "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)), }; diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index c3e4b71..98d3748 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -105,9 +105,10 @@ pub fn run() { use tauri::Manager; let state = app.state::(); let ssh_clone = state.ssh.clone(); + let rdp_clone = state.rdp.clone(); let scrollback_clone = state.scrollback.clone(); tauri::async_runtime::spawn(async move { - match mcp::server::start_mcp_server(ssh_clone, scrollback_clone).await { + match mcp::server::start_mcp_server(ssh_clone, rdp_clone, scrollback_clone).await { Ok(port) => log::info!("MCP server started on localhost:{}", port), Err(e) => log::error!("Failed to start MCP server: {}", e), } diff --git a/src-tauri/src/mcp/server.rs b/src-tauri/src/mcp/server.rs index 9427c2c..87c7aaa 100644 --- a/src-tauri/src/mcp/server.rs +++ b/src-tauri/src/mcp/server.rs @@ -10,11 +10,13 @@ use serde::{Deserialize, Serialize}; use tokio::net::TcpListener; use crate::mcp::ScrollbackRegistry; +use crate::rdp::RdpService; use crate::ssh::session::SshService; /// Shared state passed to axum handlers. pub struct McpServerState { pub ssh: SshService, + pub rdp: RdpService, pub scrollback: ScrollbackRegistry, } @@ -24,6 +26,11 @@ struct TerminalReadRequest { lines: Option, } +#[derive(Deserialize)] +struct ScreenshotRequest { + session_id: String, +} + #[derive(Deserialize)] struct TerminalExecuteRequest { session_id: String, @@ -49,7 +56,7 @@ fn err_response(msg: String) -> Json> { async fn handle_list_sessions( AxumState(state): AxumState>, ) -> Json>> { - let sessions: Vec = state.ssh.list_sessions() + let mut sessions: Vec = state.ssh.list_sessions() .into_iter() .map(|s| serde_json::json!({ "id": s.id, @@ -59,9 +66,32 @@ async fn handle_list_sessions( "username": s.username, })) .collect(); + + // Include RDP sessions + for s in state.rdp.list_sessions() { + sessions.push(serde_json::json!({ + "id": s.id, + "type": "rdp", + "name": s.hostname.clone(), + "host": s.hostname, + "width": s.width, + "height": s.height, + })); + } + ok_response(sessions) } +async fn handle_screenshot( + AxumState(state): AxumState>, + Json(req): Json, +) -> Json> { + match state.rdp.screenshot_png_base64(&req.session_id).await { + Ok(b64) => ok_response(b64), + Err(e) => err_response(e), + } +} + async fn handle_terminal_read( AxumState(state): AxumState>, Json(req): Json, @@ -134,14 +164,16 @@ async fn handle_terminal_execute( /// Start the MCP HTTP server and write the port to disk. pub async fn start_mcp_server( ssh: SshService, + rdp: RdpService, scrollback: ScrollbackRegistry, ) -> Result { - let state = Arc::new(McpServerState { ssh, scrollback }); + let state = Arc::new(McpServerState { ssh, rdp, scrollback }); let app = Router::new() .route("/mcp/sessions", post(handle_list_sessions)) .route("/mcp/terminal/read", post(handle_terminal_read)) .route("/mcp/terminal/execute", post(handle_terminal_execute)) + .route("/mcp/screenshot", post(handle_screenshot)) .with_state(state); let listener = TcpListener::bind("127.0.0.1:0").await diff --git a/src-tauri/src/pty/mod.rs b/src-tauri/src/pty/mod.rs index 83769f1..7c5509c 100644 --- a/src-tauri/src/pty/mod.rs +++ b/src-tauri/src/pty/mod.rs @@ -93,8 +93,10 @@ impl PtyService { let mut cmd = CommandBuilder::new(shell_path); - // Auto-inject MCP server config so AI CLIs discover the bridge - cmd.env("WRAITH_MCP_BRIDGE", "wraith-mcp-bridge"); + // Auto-inject MCP server config so AI CLIs discover the bridge. + // Claude Code reads CLAUDE_MCP_SERVERS env var for server config. + let mcp_config = r#"{"wraith":{"command":"wraith-mcp-bridge","args":[]}}"#; + cmd.env("CLAUDE_MCP_SERVERS", mcp_config); let child = pair.slave .spawn_command(cmd) diff --git a/src-tauri/src/rdp/mod.rs b/src-tauri/src/rdp/mod.rs index 6c9bb12..a2dbb1a 100644 --- a/src-tauri/src/rdp/mod.rs +++ b/src-tauri/src/rdp/mod.rs @@ -199,6 +199,28 @@ impl RdpService { Ok(buf.clone()) } + /// Capture the current RDP frame as a base64-encoded PNG. + pub async fn screenshot_png_base64(&self, session_id: &str) -> Result { + let handle = self.sessions.get(session_id).ok_or_else(|| format!("RDP session {} not found", session_id))?; + let width = handle.width as u32; + let height = handle.height as u32; + let buf = handle.frame_buffer.lock().await; + + // Encode RGBA raw bytes to PNG + let mut png_data = Vec::new(); + { + let mut encoder = png::Encoder::new(&mut png_data, width, height); + encoder.set_color(png::ColorType::Rgba); + encoder.set_depth(png::BitDepth::Eight); + let mut writer = encoder.write_header() + .map_err(|e| format!("PNG header error: {}", e))?; + writer.write_image_data(&buf) + .map_err(|e| format!("PNG encode error: {}", e))?; + } + + Ok(base64::engine::general_purpose::STANDARD.encode(&png_data)) + } + pub fn send_clipboard(&self, session_id: &str, text: &str) -> Result<(), String> { let handle = self.sessions.get(session_id).ok_or_else(|| format!("RDP session {} not found", session_id))?; handle.input_tx.send(InputEvent::Clipboard(text.to_string())).map_err(|_| format!("RDP session {} input channel closed", session_id))