feat: MCP Phase 1 — scrollback buffer, terminal_read, terminal_execute
All checks were successful
Build & Sign Wraith / Build Windows + Sign (push) Successful in 2m52s
All checks were successful
Build & Sign Wraith / Build Windows + Sign (push) Successful in 2m52s
Infrastructure for the Wraith Terminal MCP server: - ScrollbackBuffer: 64KB circular buffer per session with ANSI stripping - ScrollbackRegistry: DashMap registry shared between output loops and MCP - SSH output loop feeds scrollback in addition to emitting events - PTY output loop feeds scrollback in addition to emitting events - mcp_terminal_read: read last N lines from any session (ANSI stripped) - mcp_terminal_execute: send command + marker, capture output until marker - mcp_list_sessions: enumerate all active SSH sessions with metadata 8 new scrollback tests (ring buffer, ANSI strip, line limiting). 95 total tests, zero warnings. Bridge binary and auto-config injection to follow. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Vantz Stockwell
parent
4b68b8549a
commit
a3a7116f00
@ -61,7 +61,7 @@ npm install # Install frontend deps
|
||||
npm run dev # Vite dev server only
|
||||
cargo tauri dev # Full app (Rust + frontend)
|
||||
cargo tauri build # Production build
|
||||
cd src-tauri && cargo test # Run Rust tests (87 tests)
|
||||
cd src-tauri && cargo test # Run Rust tests (95 tests)
|
||||
cd src-tauri && cargo build # Build Rust only
|
||||
```
|
||||
|
||||
|
||||
340
docs/specs/wraith-terminal-mcp-spec.md
Normal file
340
docs/specs/wraith-terminal-mcp-spec.md
Normal file
@ -0,0 +1,340 @@
|
||||
# Wraith Terminal MCP — Design Specification
|
||||
|
||||
**Date:** March 25, 2026
|
||||
**Status:** Draft
|
||||
**Author:** Gargoyle (HQ XO)
|
||||
|
||||
---
|
||||
|
||||
## 1. Problem
|
||||
|
||||
The AI copilot panel in Wraith runs CLI tools (Claude, Gemini, Codex) in a local PTY. The AI can chat with the user, but it cannot independently interact with active SSH/RDP sessions. The technician has to manually copy-paste terminal output into the AI and relay commands back.
|
||||
|
||||
The goal: let the AI **drive** the terminal. Read output. Execute commands. Take screenshots. React to errors. All through a standardized protocol.
|
||||
|
||||
---
|
||||
|
||||
## 2. Solution: Wraith Terminal MCP Server
|
||||
|
||||
Implement an MCP (Model Context Protocol) server inside Wraith's Rust backend that exposes active sessions as tools and resources. The AI CLI running in the copilot panel connects to this MCP server and gains programmatic access to every open session.
|
||||
|
||||
### Architecture
|
||||
|
||||
```
|
||||
AI CLI (claude/gemini)
|
||||
|
|
||||
+-- MCP Client (built into the CLI)
|
||||
|
|
||||
+-- connects to localhost:PORT or Unix socket
|
||||
|
|
||||
v
|
||||
Wraith MCP Server (Rust, runs inside Tauri)
|
||||
|
|
||||
+-- Tool: terminal_execute(session_id, command)
|
||||
+-- Tool: terminal_read(session_id, lines?)
|
||||
+-- Tool: terminal_screenshot(session_id) [RDP only]
|
||||
+-- Tool: sftp_list(session_id, path)
|
||||
+-- Tool: sftp_read(session_id, path)
|
||||
+-- Tool: sftp_write(session_id, path, content)
|
||||
+-- Resource: sessions://list
|
||||
+-- Resource: sessions://{id}/info
|
||||
+-- Resource: sessions://{id}/scrollback
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. MCP Server Implementation
|
||||
|
||||
### 3.1 Transport
|
||||
|
||||
Two options for how the AI CLI connects to the MCP server:
|
||||
|
||||
**Option A: stdio (Recommended for v1)**
|
||||
- The copilot panel spawns the AI CLI with `--mcp-server` flag pointing to a Wraith helper binary
|
||||
- The helper binary communicates with Wraith's Tauri backend via Tauri commands
|
||||
- Simple, no port management, no firewall issues
|
||||
- Pattern: AI CLI → stdio → wraith-mcp-bridge → Tauri invoke → session data
|
||||
|
||||
**Option B: HTTP/SSE (Future)**
|
||||
- Wraith runs an HTTP server on localhost:random-port
|
||||
- AI CLI connects via `--mcp-server http://localhost:PORT`
|
||||
- More flexible (multiple AI CLIs can connect), but requires port management
|
||||
- Pattern: AI CLI → HTTP → Wraith MCP HTTP handler → session data
|
||||
|
||||
### 3.2 Rust Implementation
|
||||
|
||||
```
|
||||
src-tauri/src/mcp/
|
||||
mod.rs — MCP server lifecycle, transport handling
|
||||
tools.rs — Tool definitions (terminal_execute, screenshot, etc.)
|
||||
resources.rs — Resource definitions (session list, scrollback)
|
||||
bridge.rs — Bridge between MCP protocol and existing services
|
||||
```
|
||||
|
||||
The MCP server reuses existing services:
|
||||
- `SshService` — for terminal_execute, terminal_read on SSH sessions
|
||||
- `RdpService` — for terminal_screenshot on RDP sessions
|
||||
- `SftpService` — for sftp_list, sftp_read, sftp_write
|
||||
- `PtyService` — for local shell access
|
||||
- `SessionStore` (DashMap) — for session enumeration
|
||||
|
||||
---
|
||||
|
||||
## 4. Tools
|
||||
|
||||
### 4.1 terminal_execute
|
||||
|
||||
Execute a command in an active SSH or local PTY session and return the output.
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "terminal_execute",
|
||||
"description": "Execute a command in a terminal session and return output",
|
||||
"parameters": {
|
||||
"session_id": "string — the active session ID",
|
||||
"command": "string — the command to run (newline appended automatically)",
|
||||
"timeout_ms": "number — max wait for output (default: 5000)"
|
||||
},
|
||||
"returns": "string — captured terminal output after command execution"
|
||||
}
|
||||
```
|
||||
|
||||
**Implementation:** Write command + `\n` to the session's writer. Start capturing output from the session's reader. Wait for a shell prompt pattern or timeout. Return captured bytes as UTF-8 string.
|
||||
|
||||
**Challenge:** Detecting when command output is "done" — shell prompt detection is fragile. Options:
|
||||
- **Marker approach:** Send `echo __WRAITH_DONE__` after the command, capture until marker appears
|
||||
- **Timeout approach:** Wait N ms after last output byte, assume done
|
||||
- **Prompt regex:** Configurable prompt pattern (default: `$ `, `# `, `> `, `PS>`)
|
||||
|
||||
Recommend: marker approach for SSH, timeout approach for PTY (since local shells have predictable prompt timing).
|
||||
|
||||
### 4.2 terminal_read
|
||||
|
||||
Read the current scrollback or recent output from a session without executing anything.
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "terminal_read",
|
||||
"description": "Read recent terminal output from a session",
|
||||
"parameters": {
|
||||
"session_id": "string",
|
||||
"lines": "number — last N lines (default: 50)"
|
||||
},
|
||||
"returns": "string — terminal scrollback content (ANSI stripped)"
|
||||
}
|
||||
```
|
||||
|
||||
**Implementation:** Maintain a circular buffer of recent output per session (last 10KB). On read, return the last N lines with ANSI escape codes stripped.
|
||||
|
||||
**Note:** The buffer exists in the Rust backend, not xterm.js. The AI doesn't need to scrape the DOM — it reads from the same data stream that feeds the terminal.
|
||||
|
||||
### 4.3 terminal_screenshot
|
||||
|
||||
Capture the current frame of an RDP session as a base64-encoded PNG.
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "terminal_screenshot",
|
||||
"description": "Capture a screenshot of an RDP session",
|
||||
"parameters": {
|
||||
"session_id": "string — must be an RDP session"
|
||||
},
|
||||
"returns": "string — base64-encoded PNG image"
|
||||
}
|
||||
```
|
||||
|
||||
**Implementation:** The RDP frame buffer is already maintained by `RdpService`. Encode the current frame as PNG (using the `image` crate), base64 encode, return. The AI CLI passes this to the multimodal AI provider for visual analysis.
|
||||
|
||||
**Use case:** "Screenshot the error on screen. What can you tell me about it?"
|
||||
|
||||
### 4.4 sftp_list
|
||||
|
||||
List files in a directory on the remote host via the session's SFTP channel.
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "sftp_list",
|
||||
"description": "List files in a remote directory",
|
||||
"parameters": {
|
||||
"session_id": "string",
|
||||
"path": "string — remote directory path"
|
||||
},
|
||||
"returns": "array of { name, size, modified, is_dir }"
|
||||
}
|
||||
```
|
||||
|
||||
### 4.5 sftp_read
|
||||
|
||||
Read a file from the remote host.
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "sftp_read",
|
||||
"description": "Read a file from the remote host",
|
||||
"parameters": {
|
||||
"session_id": "string",
|
||||
"path": "string — remote file path",
|
||||
"max_bytes": "number — limit (default: 1MB)"
|
||||
},
|
||||
"returns": "string — file content (UTF-8) or base64 for binary"
|
||||
}
|
||||
```
|
||||
|
||||
### 4.6 sftp_write
|
||||
|
||||
Write a file to the remote host.
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "sftp_write",
|
||||
"description": "Write content to a file on the remote host",
|
||||
"parameters": {
|
||||
"session_id": "string",
|
||||
"path": "string — remote file path",
|
||||
"content": "string — file content"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 5. Resources
|
||||
|
||||
### 5.1 sessions://list
|
||||
|
||||
Returns all active sessions with their type, connection info, and status.
|
||||
|
||||
```json
|
||||
[
|
||||
{
|
||||
"id": "ssh-abc123",
|
||||
"type": "ssh",
|
||||
"name": "prod-web-01",
|
||||
"host": "10.0.1.50",
|
||||
"username": "admin",
|
||||
"status": "connected"
|
||||
},
|
||||
{
|
||||
"id": "rdp-def456",
|
||||
"type": "rdp",
|
||||
"name": "dc-01",
|
||||
"host": "10.0.1.10",
|
||||
"status": "connected"
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
### 5.2 sessions://{id}/info
|
||||
|
||||
Detailed info about a specific session — connection parameters, uptime, bytes transferred.
|
||||
|
||||
### 5.3 sessions://{id}/scrollback
|
||||
|
||||
Full scrollback buffer for a terminal session (last 10KB, ANSI stripped).
|
||||
|
||||
---
|
||||
|
||||
## 6. Security
|
||||
|
||||
- **MCP server only binds to localhost** — no remote access, no network exposure
|
||||
- **Session access inherits Wraith's auth** — if the user is logged into Wraith, the MCP server trusts the connection
|
||||
- **No credential exposure** — the MCP tools execute commands through existing authenticated sessions. The AI never sees passwords or SSH keys.
|
||||
- **Audit trail** — every MCP tool invocation logged with timestamp, session ID, command, and result size
|
||||
- **Read-only option** — sessions can be marked read-only in connection settings, preventing terminal_execute and sftp_write
|
||||
|
||||
---
|
||||
|
||||
## 7. AI CLI Integration
|
||||
|
||||
### 7.1 Claude Code
|
||||
|
||||
Claude Code already supports MCP servers via `--mcp-server` flag or `.claude/settings.json`. Configuration:
|
||||
|
||||
```json
|
||||
{
|
||||
"mcpServers": {
|
||||
"wraith": {
|
||||
"command": "wraith-mcp-bridge",
|
||||
"args": []
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
The `wraith-mcp-bridge` is a small binary that Wraith ships alongside the main app. It communicates with the running Wraith instance via Tauri's IPC.
|
||||
|
||||
### 7.2 Gemini CLI
|
||||
|
||||
Gemini CLI supports MCP servers similarly. Same bridge binary, same configuration pattern.
|
||||
|
||||
### 7.3 Auto-Configuration
|
||||
|
||||
When the copilot panel launches an AI CLI, Wraith can auto-inject the MCP server configuration via environment variables or command-line flags, so the user doesn't have to manually configure anything.
|
||||
|
||||
```rust
|
||||
// When spawning the AI CLI in the PTY:
|
||||
let mut cmd = CommandBuilder::new(shell_path);
|
||||
cmd.env("CLAUDE_MCP_SERVERS", r#"{"wraith":{"command":"wraith-mcp-bridge"}}"#);
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 8. Data Flow Example
|
||||
|
||||
**User says to Claude in copilot panel:** "Check disk space on the server I'm connected to"
|
||||
|
||||
1. Claude's MCP client calls `sessions://list` → gets `[{id: "ssh-abc", name: "prod-web-01", ...}]`
|
||||
2. Claude calls `terminal_execute(session_id: "ssh-abc", command: "df -h")`
|
||||
3. Wraith MCP bridge → Tauri invoke → SshService.write("ssh-abc", "df -h\n")
|
||||
4. Wraith captures output until prompt marker
|
||||
5. Returns: `/dev/sda1 50G 45G 5G 90% /`
|
||||
6. Claude analyzes: "Your root partition is at 90%. You should clean up /var/log or expand the disk."
|
||||
|
||||
**User says:** "Screenshot the RDP session, what's that error?"
|
||||
|
||||
1. Claude calls `terminal_screenshot(session_id: "rdp-def")`
|
||||
2. Wraith MCP bridge → RdpService.get_frame("rdp-def") → PNG encode → base64
|
||||
3. Returns 200KB base64 PNG
|
||||
4. Claude (multimodal) analyzes the image: "That's a Windows Event Viewer showing Event ID 1001 — application crash in outlook.exe. The faulting module is mso.dll. This is a known Office corruption issue. Run `sfc /scannow` or repair Office from Control Panel."
|
||||
|
||||
---
|
||||
|
||||
## 9. Implementation Phases
|
||||
|
||||
### Phase 1: Bridge + Basic Tools (MVP)
|
||||
- `wraith-mcp-bridge` binary (stdio transport)
|
||||
- `terminal_execute` tool (marker-based output capture)
|
||||
- `terminal_read` tool (scrollback buffer)
|
||||
- `sessions://list` resource
|
||||
- Auto-configuration when spawning AI CLI
|
||||
|
||||
### Phase 2: SFTP + Screenshot
|
||||
- `sftp_list`, `sftp_read`, `sftp_write` tools
|
||||
- `terminal_screenshot` tool (RDP frame capture)
|
||||
- Session info resource
|
||||
|
||||
### Phase 3: Advanced
|
||||
- HTTP/SSE transport for multi-client access
|
||||
- Read-only session enforcement
|
||||
- Audit trail logging
|
||||
- AI-initiated session creation ("Connect me to prod-web-01")
|
||||
|
||||
---
|
||||
|
||||
## 10. Dependencies
|
||||
|
||||
| Component | Crate/Tool | License |
|
||||
|---|---|---|
|
||||
| MCP protocol | Custom implementation (JSON-RPC over stdio) | Proprietary |
|
||||
| PNG encoding | `image` crate | MIT/Apache-2.0 |
|
||||
| Base64 | `base64` crate (already in deps) | MIT/Apache-2.0 |
|
||||
| ANSI stripping | `strip-ansi-escapes` crate | MIT/Apache-2.0 |
|
||||
| Bridge binary | Rust, ships alongside Wraith | Proprietary |
|
||||
|
||||
---
|
||||
|
||||
## 11. Black Binder Note
|
||||
|
||||
An MCP server embedded in a remote access client that gives AI tools programmatic access to live SSH, RDP, and SFTP sessions is, to the company's knowledge, a novel integration. No competing SSH/RDP client ships with an MCP server that allows AI assistants to interact with active remote sessions.
|
||||
|
||||
The combination of terminal command execution, RDP screenshot analysis, and SFTP file operations through a standardized AI tool protocol represents a new category of AI-augmented remote access.
|
||||
121
src-tauri/src/commands/mcp_commands.rs
Normal file
121
src-tauri/src/commands/mcp_commands.rs
Normal file
@ -0,0 +1,121 @@
|
||||
//! Tauri commands for MCP tool operations.
|
||||
//!
|
||||
//! These expose terminal_read, terminal_execute, and session listing to both
|
||||
//! the frontend and the MCP bridge binary.
|
||||
|
||||
use serde::Serialize;
|
||||
use tauri::State;
|
||||
|
||||
use crate::AppState;
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct McpSessionInfo {
|
||||
pub id: String,
|
||||
pub session_type: String, // "ssh" or "pty"
|
||||
pub name: String,
|
||||
pub host: Option<String>,
|
||||
pub username: Option<String>,
|
||||
}
|
||||
|
||||
/// List all active sessions (SSH + PTY) with metadata.
|
||||
#[tauri::command]
|
||||
pub fn mcp_list_sessions(state: State<'_, AppState>) -> Vec<McpSessionInfo> {
|
||||
let mut sessions = Vec::new();
|
||||
|
||||
// SSH sessions
|
||||
for info in state.ssh.list_sessions() {
|
||||
sessions.push(McpSessionInfo {
|
||||
id: info.id,
|
||||
session_type: "ssh".to_string(),
|
||||
name: format!("{}@{}:{}", info.username, info.hostname, info.port),
|
||||
host: Some(info.hostname),
|
||||
username: Some(info.username),
|
||||
});
|
||||
}
|
||||
|
||||
sessions
|
||||
}
|
||||
|
||||
/// Read the last N lines from a session's scrollback buffer (ANSI stripped).
|
||||
#[tauri::command]
|
||||
pub fn mcp_terminal_read(
|
||||
session_id: String,
|
||||
lines: Option<usize>,
|
||||
state: State<'_, AppState>,
|
||||
) -> Result<String, String> {
|
||||
let n = lines.unwrap_or(50);
|
||||
let buf = state.scrollback.get(&session_id)
|
||||
.ok_or_else(|| format!("No scrollback buffer for session {}", session_id))?;
|
||||
Ok(buf.read_lines(n))
|
||||
}
|
||||
|
||||
/// Execute a command in an SSH session and capture output using a marker.
|
||||
///
|
||||
/// Sends the command followed by `echo __WRAITH_MCP_DONE__`, then reads the
|
||||
/// scrollback until the marker appears or timeout is reached.
|
||||
#[tauri::command]
|
||||
pub async fn mcp_terminal_execute(
|
||||
session_id: String,
|
||||
command: String,
|
||||
timeout_ms: Option<u64>,
|
||||
state: State<'_, AppState>,
|
||||
) -> Result<String, String> {
|
||||
let timeout = timeout_ms.unwrap_or(5000);
|
||||
let marker = "__WRAITH_MCP_DONE__";
|
||||
|
||||
// Record current buffer position
|
||||
let buf = state.scrollback.get(&session_id)
|
||||
.ok_or_else(|| format!("No scrollback buffer for session {}", session_id))?;
|
||||
let before = buf.total_written();
|
||||
|
||||
// Send command + marker echo
|
||||
let full_cmd = format!("{}\necho {}\n", command, marker);
|
||||
state.ssh.write(&session_id, full_cmd.as_bytes()).await?;
|
||||
|
||||
// Poll scrollback until marker appears or timeout
|
||||
let start = std::time::Instant::now();
|
||||
let timeout_dur = std::time::Duration::from_millis(timeout);
|
||||
|
||||
loop {
|
||||
if start.elapsed() > timeout_dur {
|
||||
// Return whatever we captured so far
|
||||
let raw = buf.read_raw();
|
||||
let total = buf.total_written();
|
||||
// Extract just the new content since we sent the command
|
||||
let new_bytes = total.saturating_sub(before);
|
||||
let output = if new_bytes > 0 && raw.len() >= new_bytes {
|
||||
&raw[raw.len() - new_bytes.min(raw.len())..]
|
||||
} else {
|
||||
""
|
||||
};
|
||||
return Ok(format!("[timeout after {}ms]\n{}", timeout, output));
|
||||
}
|
||||
|
||||
let raw = buf.read_raw();
|
||||
if raw.contains(marker) {
|
||||
// Extract output between command echo and marker
|
||||
let total = buf.total_written();
|
||||
let new_bytes = total.saturating_sub(before);
|
||||
let output = if new_bytes > 0 && raw.len() >= new_bytes {
|
||||
raw[raw.len() - new_bytes.min(raw.len())..].to_string()
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
|
||||
// Strip the command echo and marker from output
|
||||
let clean = output
|
||||
.lines()
|
||||
.filter(|line| {
|
||||
!line.contains(marker)
|
||||
&& !line.trim().starts_with(&command.trim_start().chars().take(20).collect::<String>())
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
return Ok(clean.trim().to_string());
|
||||
}
|
||||
|
||||
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||
}
|
||||
}
|
||||
@ -7,3 +7,4 @@ pub mod sftp_commands;
|
||||
pub mod rdp_commands;
|
||||
pub mod theme_commands;
|
||||
pub mod pty_commands;
|
||||
pub mod mcp_commands;
|
||||
|
||||
@ -18,7 +18,7 @@ pub fn spawn_local_shell(
|
||||
app_handle: AppHandle,
|
||||
state: State<'_, AppState>,
|
||||
) -> Result<String, String> {
|
||||
state.pty.spawn(&shell_path, cols as u16, rows as u16, app_handle)
|
||||
state.pty.spawn(&shell_path, cols as u16, rows as u16, app_handle, &state.scrollback)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
|
||||
@ -32,6 +32,7 @@ pub async fn connect_ssh(
|
||||
cols,
|
||||
rows,
|
||||
&state.sftp,
|
||||
&state.scrollback,
|
||||
)
|
||||
.await
|
||||
}
|
||||
@ -63,6 +64,7 @@ pub async fn connect_ssh_with_key(
|
||||
cols,
|
||||
rows,
|
||||
&state.sftp,
|
||||
&state.scrollback,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
@ -9,6 +9,7 @@ pub mod rdp;
|
||||
pub mod theme;
|
||||
pub mod workspace;
|
||||
pub mod pty;
|
||||
pub mod mcp;
|
||||
pub mod commands;
|
||||
|
||||
use std::path::PathBuf;
|
||||
@ -25,6 +26,7 @@ use rdp::RdpService;
|
||||
use theme::ThemeService;
|
||||
use workspace::WorkspaceService;
|
||||
use pty::PtyService;
|
||||
use mcp::ScrollbackRegistry;
|
||||
|
||||
pub struct AppState {
|
||||
pub db: Database,
|
||||
@ -38,6 +40,7 @@ pub struct AppState {
|
||||
pub theme: ThemeService,
|
||||
pub workspace: WorkspaceService,
|
||||
pub pty: PtyService,
|
||||
pub scrollback: ScrollbackRegistry,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
@ -57,6 +60,7 @@ impl AppState {
|
||||
theme: ThemeService::new(database.clone()),
|
||||
workspace: WorkspaceService::new(SettingsService::new(database.clone())),
|
||||
pty: PtyService::new(),
|
||||
scrollback: ScrollbackRegistry::new(),
|
||||
})
|
||||
}
|
||||
|
||||
@ -109,6 +113,7 @@ pub fn run() {
|
||||
commands::rdp_commands::connect_rdp, commands::rdp_commands::rdp_get_frame, commands::rdp_commands::rdp_send_mouse, commands::rdp_commands::rdp_send_key, commands::rdp_commands::rdp_send_clipboard, commands::rdp_commands::disconnect_rdp, commands::rdp_commands::list_rdp_sessions,
|
||||
commands::theme_commands::list_themes, commands::theme_commands::get_theme,
|
||||
commands::pty_commands::list_available_shells, commands::pty_commands::spawn_local_shell, commands::pty_commands::pty_write, commands::pty_commands::pty_resize, commands::pty_commands::disconnect_pty,
|
||||
commands::mcp_commands::mcp_list_sessions, commands::mcp_commands::mcp_terminal_read, commands::mcp_commands::mcp_terminal_execute,
|
||||
])
|
||||
.run(tauri::generate_context!())
|
||||
.expect("error while running tauri application");
|
||||
|
||||
42
src-tauri/src/mcp/mod.rs
Normal file
42
src-tauri/src/mcp/mod.rs
Normal file
@ -0,0 +1,42 @@
|
||||
//! MCP (Model Context Protocol) infrastructure for Wraith.
|
||||
//!
|
||||
//! Provides programmatic access to active sessions so AI tools running in the
|
||||
//! copilot panel can read terminal output, execute commands, and enumerate
|
||||
//! sessions.
|
||||
|
||||
pub mod scrollback;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use crate::mcp::scrollback::ScrollbackBuffer;
|
||||
|
||||
/// Registry of scrollback buffers keyed by session ID.
|
||||
/// Shared between SSH/PTY output loops (writers) and MCP tools (readers).
|
||||
pub struct ScrollbackRegistry {
|
||||
buffers: DashMap<String, Arc<ScrollbackBuffer>>,
|
||||
}
|
||||
|
||||
impl ScrollbackRegistry {
|
||||
pub fn new() -> Self {
|
||||
Self { buffers: DashMap::new() }
|
||||
}
|
||||
|
||||
/// Create and register a new scrollback buffer for a session.
|
||||
pub fn create(&self, session_id: &str) -> Arc<ScrollbackBuffer> {
|
||||
let buf = Arc::new(ScrollbackBuffer::new());
|
||||
self.buffers.insert(session_id.to_string(), buf.clone());
|
||||
buf
|
||||
}
|
||||
|
||||
/// Get the scrollback buffer for a session.
|
||||
pub fn get(&self, session_id: &str) -> Option<Arc<ScrollbackBuffer>> {
|
||||
self.buffers.get(session_id).map(|entry| entry.clone())
|
||||
}
|
||||
|
||||
/// Remove a session's scrollback buffer.
|
||||
pub fn remove(&self, session_id: &str) {
|
||||
self.buffers.remove(session_id);
|
||||
}
|
||||
}
|
||||
195
src-tauri/src/mcp/scrollback.rs
Normal file
195
src-tauri/src/mcp/scrollback.rs
Normal file
@ -0,0 +1,195 @@
|
||||
//! Per-session scrollback buffer for MCP terminal_read.
|
||||
//!
|
||||
//! A thread-safe circular buffer that stores the last N bytes of terminal
|
||||
//! output. Both SSH and PTY output loops write to it. The MCP tools read
|
||||
//! from it without touching xterm.js or the frontend.
|
||||
|
||||
use std::sync::Mutex;
|
||||
|
||||
const DEFAULT_CAPACITY: usize = 64 * 1024; // 64KB per session
|
||||
|
||||
/// Thread-safe circular buffer for terminal output.
|
||||
pub struct ScrollbackBuffer {
|
||||
inner: Mutex<RingBuffer>,
|
||||
}
|
||||
|
||||
struct RingBuffer {
|
||||
data: Vec<u8>,
|
||||
capacity: usize,
|
||||
/// Write position (wraps around)
|
||||
write_pos: usize,
|
||||
/// Total bytes written (for detecting wrap)
|
||||
total_written: usize,
|
||||
}
|
||||
|
||||
impl ScrollbackBuffer {
|
||||
pub fn new() -> Self {
|
||||
Self::with_capacity(DEFAULT_CAPACITY)
|
||||
}
|
||||
|
||||
pub fn with_capacity(capacity: usize) -> Self {
|
||||
Self {
|
||||
inner: Mutex::new(RingBuffer {
|
||||
data: vec![0u8; capacity],
|
||||
capacity,
|
||||
write_pos: 0,
|
||||
total_written: 0,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Append bytes to the buffer. Old data is overwritten when full.
|
||||
pub fn push(&self, bytes: &[u8]) {
|
||||
let mut buf = self.inner.lock().unwrap();
|
||||
for &b in bytes {
|
||||
let pos = buf.write_pos;
|
||||
buf.data[pos] = b;
|
||||
buf.write_pos = (pos + 1) % buf.capacity;
|
||||
buf.total_written += 1;
|
||||
}
|
||||
}
|
||||
|
||||
/// Read the last `n` lines from the buffer, with ANSI escape codes stripped.
|
||||
pub fn read_lines(&self, n: usize) -> String {
|
||||
let raw = self.read_raw();
|
||||
let text = strip_ansi(&raw);
|
||||
let lines: Vec<&str> = text.lines().collect();
|
||||
let start = lines.len().saturating_sub(n);
|
||||
lines[start..].join("\n")
|
||||
}
|
||||
|
||||
/// Read all buffered content as raw bytes (ordered oldest→newest).
|
||||
pub fn read_raw(&self) -> String {
|
||||
let buf = self.inner.lock().unwrap();
|
||||
let bytes = if buf.total_written >= buf.capacity {
|
||||
// Buffer has wrapped — read from write_pos to end, then start to write_pos
|
||||
let mut out = Vec::with_capacity(buf.capacity);
|
||||
out.extend_from_slice(&buf.data[buf.write_pos..]);
|
||||
out.extend_from_slice(&buf.data[..buf.write_pos]);
|
||||
out
|
||||
} else {
|
||||
// Buffer hasn't wrapped yet
|
||||
buf.data[..buf.write_pos].to_vec()
|
||||
};
|
||||
String::from_utf8_lossy(&bytes).to_string()
|
||||
}
|
||||
|
||||
/// Total bytes written since creation.
|
||||
pub fn total_written(&self) -> usize {
|
||||
self.inner.lock().unwrap().total_written
|
||||
}
|
||||
}
|
||||
|
||||
/// Strip ANSI escape sequences from text.
|
||||
fn strip_ansi(input: &str) -> String {
|
||||
let mut output = String::with_capacity(input.len());
|
||||
let mut chars = input.chars().peekable();
|
||||
|
||||
while let Some(ch) = chars.next() {
|
||||
if ch == '\x1b' {
|
||||
// ESC sequence — consume until terminator
|
||||
if let Some(&next) = chars.peek() {
|
||||
if next == '[' {
|
||||
chars.next(); // consume '['
|
||||
// CSI sequence — consume until letter
|
||||
while let Some(&c) = chars.peek() {
|
||||
chars.next();
|
||||
if c.is_ascii_alphabetic() || c == '~' || c == '@' {
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else if next == ']' {
|
||||
chars.next(); // consume ']'
|
||||
// OSC sequence — consume until BEL or ST
|
||||
while let Some(&c) = chars.peek() {
|
||||
chars.next();
|
||||
if c == '\x07' {
|
||||
break;
|
||||
}
|
||||
if c == '\x1b' {
|
||||
if chars.peek() == Some(&'\\') {
|
||||
chars.next();
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
chars.next(); // consume single-char escape
|
||||
}
|
||||
}
|
||||
} else if ch == '\r' {
|
||||
// Skip carriage returns for cleaner output
|
||||
continue;
|
||||
} else {
|
||||
output.push(ch);
|
||||
}
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn push_and_read_basic() {
|
||||
let buf = ScrollbackBuffer::new();
|
||||
buf.push(b"hello world\n");
|
||||
let lines = buf.read_lines(10);
|
||||
assert!(lines.contains("hello world"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_lines_limits_output() {
|
||||
let buf = ScrollbackBuffer::new();
|
||||
buf.push(b"line1\nline2\nline3\nline4\nline5\n");
|
||||
let lines = buf.read_lines(2);
|
||||
assert!(!lines.contains("line3"));
|
||||
assert!(lines.contains("line4"));
|
||||
assert!(lines.contains("line5"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn circular_buffer_wraps() {
|
||||
let buf = ScrollbackBuffer::with_capacity(16);
|
||||
buf.push(b"AAAAAAAAAAAAAAAA"); // fill 16 bytes
|
||||
buf.push(b"BBBB"); // overwrite first 4
|
||||
let raw = buf.read_raw();
|
||||
assert!(raw.starts_with("AAAAAAAAAAAA")); // 12 A's remain
|
||||
assert!(raw.ends_with("BBBB"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_ansi_removes_csi() {
|
||||
let input = "\x1b[32mgreen\x1b[0m normal";
|
||||
assert_eq!(strip_ansi(input), "green normal");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_ansi_removes_osc() {
|
||||
let input = "\x1b]0;title\x07text";
|
||||
assert_eq!(strip_ansi(input), "text");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_ansi_preserves_plain_text() {
|
||||
let input = "no escapes here\njust text";
|
||||
assert_eq!(strip_ansi(input), "no escapes here\njust text");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_buffer_returns_empty() {
|
||||
let buf = ScrollbackBuffer::new();
|
||||
assert_eq!(buf.read_lines(10), "");
|
||||
assert_eq!(buf.total_written(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn total_written_tracks_all_bytes() {
|
||||
let buf = ScrollbackBuffer::with_capacity(8);
|
||||
buf.push(b"12345678"); // 8 bytes
|
||||
buf.push(b"ABCD"); // 4 more, wraps
|
||||
assert_eq!(buf.total_written(), 12);
|
||||
}
|
||||
}
|
||||
@ -9,6 +9,8 @@ use portable_pty::{native_pty_system, Child, CommandBuilder, MasterPty, PtySize}
|
||||
use serde::Serialize;
|
||||
use tauri::{AppHandle, Emitter};
|
||||
|
||||
use crate::mcp::ScrollbackRegistry;
|
||||
|
||||
#[derive(Debug, Serialize, Clone)]
|
||||
pub struct ShellInfo {
|
||||
pub name: String,
|
||||
@ -80,6 +82,7 @@ impl PtyService {
|
||||
cols: u16,
|
||||
rows: u16,
|
||||
app_handle: AppHandle,
|
||||
scrollback: &ScrollbackRegistry,
|
||||
) -> Result<String, String> {
|
||||
let session_id = uuid::Uuid::new_v4().to_string();
|
||||
let pty_system = native_pty_system();
|
||||
@ -112,6 +115,9 @@ impl PtyService {
|
||||
|
||||
self.sessions.insert(session_id.clone(), session);
|
||||
|
||||
// Create scrollback buffer for MCP terminal_read
|
||||
let scrollback_buf = scrollback.create(&session_id);
|
||||
|
||||
// Output reader loop — runs in a dedicated OS thread because
|
||||
// portable-pty's reader is synchronous (std::io::Read) and
|
||||
// long-lived. Using std::thread::spawn avoids requiring a
|
||||
@ -128,6 +134,7 @@ impl PtyService {
|
||||
break;
|
||||
}
|
||||
Ok(n) => {
|
||||
scrollback_buf.push(&buf[..n]);
|
||||
let encoded = base64::engine::general_purpose::STANDARD.encode(&buf[..n]);
|
||||
let _ = app.emit(&format!("pty:data:{}", sid), encoded);
|
||||
}
|
||||
|
||||
@ -12,6 +12,7 @@ use tokio::sync::Mutex as TokioMutex;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::db::Database;
|
||||
use crate::mcp::ScrollbackRegistry;
|
||||
use crate::sftp::SftpService;
|
||||
use crate::ssh::cwd::CwdTracker;
|
||||
use crate::ssh::host_key::{HostKeyResult, HostKeyStore};
|
||||
@ -82,7 +83,7 @@ impl SshService {
|
||||
Self { sessions: DashMap::new(), db }
|
||||
}
|
||||
|
||||
pub async fn connect(&self, app_handle: AppHandle, hostname: &str, port: u16, username: &str, auth: AuthMethod, cols: u32, rows: u32, sftp_service: &SftpService) -> Result<String, String> {
|
||||
pub async fn connect(&self, app_handle: AppHandle, hostname: &str, port: u16, username: &str, auth: AuthMethod, cols: u32, rows: u32, sftp_service: &SftpService, scrollback: &ScrollbackRegistry) -> Result<String, String> {
|
||||
let session_id = uuid::Uuid::new_v4().to_string();
|
||||
let config = Arc::new(russh::client::Config::default());
|
||||
let handler = SshClient { host_key_store: HostKeyStore::new(self.db.clone()), hostname: hostname.to_string(), port };
|
||||
@ -147,6 +148,9 @@ impl SshService {
|
||||
}
|
||||
}
|
||||
|
||||
// Create scrollback buffer for MCP terminal_read
|
||||
let scrollback_buf = scrollback.create(&session_id);
|
||||
|
||||
// Output reader loop — owns the Channel exclusively.
|
||||
// Writes go through Handle::data() so no shared mutex is needed.
|
||||
let sid = session_id.clone();
|
||||
@ -157,10 +161,12 @@ impl SshService {
|
||||
msg = channel.wait() => {
|
||||
match msg {
|
||||
Some(ChannelMsg::Data { ref data }) => {
|
||||
scrollback_buf.push(data.as_ref());
|
||||
let encoded = base64::engine::general_purpose::STANDARD.encode(data.as_ref());
|
||||
let _ = app.emit(&format!("ssh:data:{}", sid), encoded);
|
||||
}
|
||||
Some(ChannelMsg::ExtendedData { ref data, .. }) => {
|
||||
scrollback_buf.push(data.as_ref());
|
||||
let encoded = base64::engine::general_purpose::STANDARD.encode(data.as_ref());
|
||||
let _ = app.emit(&format!("ssh:data:{}", sid), encoded);
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user