wraith/src-tauri/src/commands/mcp_commands.rs

//! 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;
    }
}