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",
|
"md5",
|
||||||
"pem",
|
"pem",
|
||||||
"pkcs8 0.10.2",
|
"pkcs8 0.10.2",
|
||||||
|
"png",
|
||||||
"portable-pty",
|
"portable-pty",
|
||||||
"rand 0.9.2",
|
"rand 0.9.2",
|
||||||
"reqwest 0.12.28",
|
"reqwest 0.12.28",
|
||||||
|
|||||||
@ -55,6 +55,7 @@ portable-pty = "0.8"
|
|||||||
# MCP HTTP server (for bridge binary communication)
|
# MCP HTTP server (for bridge binary communication)
|
||||||
axum = "0.8"
|
axum = "0.8"
|
||||||
ureq = "3"
|
ureq = "3"
|
||||||
|
png = "0.17"
|
||||||
|
|
||||||
# RDP (IronRDP)
|
# RDP (IronRDP)
|
||||||
ironrdp = { version = "0.14", features = ["connector", "session", "graphics", "input"] }
|
ironrdp = { version = "0.14", features = ["connector", "session", "graphics", "input"] }
|
||||||
|
|||||||
@ -108,6 +108,17 @@ fn handle_tools_list(id: Value) -> JsonRpcResponse {
|
|||||||
"required": ["session_id", "command"]
|
"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",
|
"name": "list_sessions",
|
||||||
"description": "List all active Wraith sessions (SSH, RDP, PTY) with connection details",
|
"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!({})),
|
"list_sessions" => call_wraith(port, "/mcp/sessions", serde_json::json!({})),
|
||||||
"terminal_read" => call_wraith(port, "/mcp/terminal/read", args.clone()),
|
"terminal_read" => call_wraith(port, "/mcp/terminal/read", args.clone()),
|
||||||
"terminal_execute" => call_wraith(port, "/mcp/terminal/execute", 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)),
|
_ => Err(format!("Unknown tool: {}", tool_name)),
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|||||||
@ -105,9 +105,10 @@ pub fn run() {
|
|||||||
use tauri::Manager;
|
use tauri::Manager;
|
||||||
let state = app.state::<AppState>();
|
let state = app.state::<AppState>();
|
||||||
let ssh_clone = state.ssh.clone();
|
let ssh_clone = state.ssh.clone();
|
||||||
|
let rdp_clone = state.rdp.clone();
|
||||||
let scrollback_clone = state.scrollback.clone();
|
let scrollback_clone = state.scrollback.clone();
|
||||||
tauri::async_runtime::spawn(async move {
|
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),
|
Ok(port) => log::info!("MCP server started on localhost:{}", port),
|
||||||
Err(e) => log::error!("Failed to start MCP server: {}", e),
|
Err(e) => log::error!("Failed to start MCP server: {}", e),
|
||||||
}
|
}
|
||||||
|
|||||||
@ -10,11 +10,13 @@ use serde::{Deserialize, Serialize};
|
|||||||
use tokio::net::TcpListener;
|
use tokio::net::TcpListener;
|
||||||
|
|
||||||
use crate::mcp::ScrollbackRegistry;
|
use crate::mcp::ScrollbackRegistry;
|
||||||
|
use crate::rdp::RdpService;
|
||||||
use crate::ssh::session::SshService;
|
use crate::ssh::session::SshService;
|
||||||
|
|
||||||
/// Shared state passed to axum handlers.
|
/// Shared state passed to axum handlers.
|
||||||
pub struct McpServerState {
|
pub struct McpServerState {
|
||||||
pub ssh: SshService,
|
pub ssh: SshService,
|
||||||
|
pub rdp: RdpService,
|
||||||
pub scrollback: ScrollbackRegistry,
|
pub scrollback: ScrollbackRegistry,
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -24,6 +26,11 @@ struct TerminalReadRequest {
|
|||||||
lines: Option<usize>,
|
lines: Option<usize>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Deserialize)]
|
||||||
|
struct ScreenshotRequest {
|
||||||
|
session_id: String,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Deserialize)]
|
#[derive(Deserialize)]
|
||||||
struct TerminalExecuteRequest {
|
struct TerminalExecuteRequest {
|
||||||
session_id: String,
|
session_id: String,
|
||||||
@ -49,7 +56,7 @@ fn err_response<T: Serialize>(msg: String) -> Json<McpResponse<T>> {
|
|||||||
async fn handle_list_sessions(
|
async fn handle_list_sessions(
|
||||||
AxumState(state): AxumState<Arc<McpServerState>>,
|
AxumState(state): AxumState<Arc<McpServerState>>,
|
||||||
) -> Json<McpResponse<Vec<serde_json::Value>>> {
|
) -> 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()
|
.into_iter()
|
||||||
.map(|s| serde_json::json!({
|
.map(|s| serde_json::json!({
|
||||||
"id": s.id,
|
"id": s.id,
|
||||||
@ -59,9 +66,32 @@ async fn handle_list_sessions(
|
|||||||
"username": s.username,
|
"username": s.username,
|
||||||
}))
|
}))
|
||||||
.collect();
|
.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)
|
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(
|
async fn handle_terminal_read(
|
||||||
AxumState(state): AxumState<Arc<McpServerState>>,
|
AxumState(state): AxumState<Arc<McpServerState>>,
|
||||||
Json(req): Json<TerminalReadRequest>,
|
Json(req): Json<TerminalReadRequest>,
|
||||||
@ -134,14 +164,16 @@ async fn handle_terminal_execute(
|
|||||||
/// Start the MCP HTTP server and write the port to disk.
|
/// Start the MCP HTTP server and write the port to disk.
|
||||||
pub async fn start_mcp_server(
|
pub async fn start_mcp_server(
|
||||||
ssh: SshService,
|
ssh: SshService,
|
||||||
|
rdp: RdpService,
|
||||||
scrollback: ScrollbackRegistry,
|
scrollback: ScrollbackRegistry,
|
||||||
) -> Result<u16, String> {
|
) -> Result<u16, String> {
|
||||||
let state = Arc::new(McpServerState { ssh, scrollback });
|
let state = Arc::new(McpServerState { ssh, rdp, scrollback });
|
||||||
|
|
||||||
let app = Router::new()
|
let app = Router::new()
|
||||||
.route("/mcp/sessions", post(handle_list_sessions))
|
.route("/mcp/sessions", post(handle_list_sessions))
|
||||||
.route("/mcp/terminal/read", post(handle_terminal_read))
|
.route("/mcp/terminal/read", post(handle_terminal_read))
|
||||||
.route("/mcp/terminal/execute", post(handle_terminal_execute))
|
.route("/mcp/terminal/execute", post(handle_terminal_execute))
|
||||||
|
.route("/mcp/screenshot", post(handle_screenshot))
|
||||||
.with_state(state);
|
.with_state(state);
|
||||||
|
|
||||||
let listener = TcpListener::bind("127.0.0.1:0").await
|
let listener = TcpListener::bind("127.0.0.1:0").await
|
||||||
|
|||||||
@ -93,8 +93,10 @@ impl PtyService {
|
|||||||
|
|
||||||
let mut cmd = CommandBuilder::new(shell_path);
|
let mut cmd = CommandBuilder::new(shell_path);
|
||||||
|
|
||||||
// Auto-inject MCP server config so AI CLIs discover the bridge
|
// Auto-inject MCP server config so AI CLIs discover the bridge.
|
||||||
cmd.env("WRAITH_MCP_BRIDGE", "wraith-mcp-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
|
let child = pair.slave
|
||||||
.spawn_command(cmd)
|
.spawn_command(cmd)
|
||||||
|
|||||||
@ -199,6 +199,28 @@ impl RdpService {
|
|||||||
Ok(buf.clone())
|
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> {
|
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))?;
|
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))
|
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