wraith/src-tauri/src/pty/mod.rs

//! Local PTY service — spawns shells for the AI copilot panel.

use std::io::{Read, Write};
use std::sync::{Arc, Mutex};

use base64::Engine;
use dashmap::DashMap;
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,
    pub path: String,
}

pub struct PtySession {
    pub id: String,
    pub shell_path: String,
    writer: Mutex<Box<dyn Write + Send>>,
    master: Mutex<Box<dyn MasterPty + Send>>,
    child: Mutex<Box<dyn Child + Send + Sync>>,
}

pub struct PtyService {
    sessions: DashMap<String, Arc<PtySession>>,
}

impl PtyService {
    pub fn new() -> Self {
        Self { sessions: DashMap::new() }
    }

    /// Detect available shells on the system.
    pub fn list_shells(&self) -> Vec<ShellInfo> {
        let mut shells = Vec::new();

        #[cfg(unix)]
        {
            if let Ok(user_shell) = std::env::var("SHELL") {
                if std::path::Path::new(&user_shell).exists() {
                    let name = std::path::Path::new(&user_shell)
                        .file_name()
                        .and_then(|n| n.to_str())
                        .unwrap_or("shell")
                        .to_string();
                    shells.push(ShellInfo { name, path: user_shell });
                }
            }
            for (name, path) in [("bash", "/bin/bash"), ("zsh", "/bin/zsh"), ("sh", "/bin/sh")] {
                if std::path::Path::new(path).exists() && !shells.iter().any(|s| s.path == path) {
                    shells.push(ShellInfo { name: name.to_string(), path: path.to_string() });
                }
            }
        }

        #[cfg(windows)]
        {
            shells.push(ShellInfo { name: "PowerShell".to_string(), path: "powershell.exe".to_string() });
            shells.push(ShellInfo { name: "CMD".to_string(), path: "cmd.exe".to_string() });
            for git_bash in [
                r"C:\Program Files\Git\bin\bash.exe",
                r"C:\Program Files (x86)\Git\bin\bash.exe",
            ] {
                if std::path::Path::new(git_bash).exists() {
                    shells.push(ShellInfo { name: "Git Bash".to_string(), path: git_bash.to_string() });
                    break;
                }
            }
            // WSL (Windows Subsystem for Linux)
            if std::path::Path::new(r"C:\Windows\System32\wsl.exe").exists() {
                shells.push(ShellInfo { name: "WSL".to_string(), path: r"C:\Windows\System32\wsl.exe".to_string() });
            }
        }

        shells
    }

    /// Spawn a local shell and start reading its output.
    pub fn spawn(
        &self,
        shell_path: &str,
        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();

        let pair = pty_system
            .openpty(PtySize { rows, cols, pixel_width: 0, pixel_height: 0 })
            .map_err(|e| format!("Failed to open PTY: {}", e))?;

        let mut cmd = CommandBuilder::new(shell_path);

        // 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)
            .map_err(|e| format!("Failed to spawn shell '{}': {}", shell_path, e))?;

        let reader = pair.master
            .try_clone_reader()
            .map_err(|e| format!("Failed to clone PTY reader: {}", e))?;

        let writer = pair.master
            .take_writer()
            .map_err(|e| format!("Failed to take PTY writer: {}", e))?;

        let session = Arc::new(PtySession {
            id: session_id.clone(),
            shell_path: shell_path.to_string(),
            writer: Mutex::new(writer),
            master: Mutex::new(pair.master),
            child: Mutex::new(child),
        });

        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
        // tokio runtime context (sync Tauri commands may not have one).
        let sid = session_id.clone();
        let app = app_handle;
        std::thread::spawn(move || {
            let mut reader = std::io::BufReader::new(reader);
            let mut buf = [0u8; 4096];
            loop {
                match reader.read(&mut buf) {
                    Ok(0) => {
                        let _ = app.emit(&format!("pty:close:{}", sid), ());
                        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);
                    }
                    Err(_) => {
                        let _ = app.emit(&format!("pty:close:{}", sid), ());
                        break;
                    }
                }
            }
        });

        Ok(session_id)
    }

    /// Write data to a PTY session's stdin.
    pub fn write(&self, session_id: &str, data: &[u8]) -> Result<(), String> {
        let session = self.sessions.get(session_id)
            .ok_or_else(|| format!("PTY session {} not found", session_id))?;
        let mut writer = session.writer.lock()
            .map_err(|e| format!("Failed to lock PTY writer: {}", e))?;
        writer.write_all(data)
            .map_err(|e| format!("Failed to write to PTY {}: {}", session_id, e))
    }

    /// Resize a PTY session.
    pub fn resize(&self, session_id: &str, cols: u16, rows: u16) -> Result<(), String> {
        let session = self.sessions.get(session_id)
            .ok_or_else(|| format!("PTY session {} not found", session_id))?;
        let master = session.master.lock()
            .map_err(|e| format!("Failed to lock PTY master: {}", e))?;
        master.resize(PtySize { rows, cols, pixel_width: 0, pixel_height: 0 })
            .map_err(|e| format!("Failed to resize PTY {}: {}", session_id, e))
    }

    /// Kill and remove a PTY session.
    pub fn disconnect(&self, session_id: &str) -> Result<(), String> {
        let (_, session) = self.sessions.remove(session_id)
            .ok_or_else(|| format!("PTY session {} not found", session_id))?;
        if let Ok(mut child) = session.child.lock() {
            let _ = child.kill();
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn list_shells_returns_at_least_one() {
        let svc = PtyService::new();
        let shells = svc.list_shells();
        assert!(!shells.is_empty(), "should find at least one shell");
        for shell in &shells {
            assert!(!shell.name.is_empty());
            assert!(!shell.path.is_empty());
        }
    }

    #[test]
    fn list_shells_no_duplicates() {
        let svc = PtyService::new();
        let shells = svc.list_shells();
        let mut paths: Vec<&str> = shells.iter().map(|s| s.path.as_str()).collect();
        let original_len = paths.len();
        paths.sort();
        paths.dedup();
        assert_eq!(original_len, paths.len(), "shell list should not contain duplicates");
    }

    #[test]
    fn disconnect_nonexistent_session_errors() {
        let svc = PtyService::new();
        assert!(svc.disconnect("nonexistent").is_err());
    }

    #[test]
    fn write_nonexistent_session_errors() {
        let svc = PtyService::new();
        assert!(svc.write("nonexistent", b"hello").is_err());
    }

    #[test]
    fn resize_nonexistent_session_errors() {
        let svc = PtyService::new();
        assert!(svc.resize("nonexistent", 80, 24).is_err());
    }
}