feat: MCP auto-inject + RDP screenshot tool
Some checks failed
Build & Sign Wraith / Build Windows + Sign (push) Failing after 16s
Some checks failed
Build & Sign Wraith / Build Windows + Sign (push) Failing after 16s
- 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) <noreply@anthropic.com>
This commit is contained in:
parent
8276b0cc59
commit
add0f0628f
1
src-tauri/Cargo.lock
generated
1
src-tauri/Cargo.lock
generated
@ -8894,6 +8894,7 @@ dependencies = [
|
||||
"md5",
|
||||
"pem",
|
||||
"pkcs8 0.10.2",
|
||||
"png",
|
||||
"portable-pty",
|
||||
"rand 0.9.2",
|
||||
"reqwest 0.12.28",
|
||||
|
||||
@ -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"] }
|
||||
|
||||
@ -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)),
|
||||
};
|
||||
|
||||
|
||||
@ -105,9 +105,10 @@ pub fn run() {
|
||||
use tauri::Manager;
|
||||
let state = app.state::<AppState>();
|
||||
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),
|
||||
}
|
||||
|
||||
@ -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<usize>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ScreenshotRequest {
|
||||
session_id: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct TerminalExecuteRequest {
|
||||
session_id: String,
|
||||
@ -49,7 +56,7 @@ fn err_response<T: Serialize>(msg: String) -> Json<McpResponse<T>> {
|
||||
async fn handle_list_sessions(
|
||||
AxumState(state): AxumState<Arc<McpServerState>>,
|
||||
) -> Json<McpResponse<Vec<serde_json::Value>>> {
|
||||
let sessions: Vec<serde_json::Value> = state.ssh.list_sessions()
|
||||
let mut sessions: Vec<serde_json::Value> = 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<Arc<McpServerState>>,
|
||||
Json(req): Json<ScreenshotRequest>,
|
||||
) -> Json<McpResponse<String>> {
|
||||
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<Arc<McpServerState>>,
|
||||
Json(req): Json<TerminalReadRequest>,
|
||||
@ -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<u16, String> {
|
||||
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
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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<String, String> {
|
||||
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))
|
||||
|
||||
Loading…
Reference in New Issue
Block a user