infinite/src-tauri/src/pty.rs

use base64::{engine::general_purpose::STANDARD as B64, Engine as _};
use portable_pty::{native_pty_system, CommandBuilder, MasterPty, PtySize};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::io::{Read, Write};
use std::sync::Mutex;
use tauri::{AppHandle, Emitter, State};
use uuid::Uuid;

pub struct PtyHandle {
    master: Box<dyn MasterPty + Send>,
    writer: Box<dyn Write + Send>,
}

#[derive(Default)]
pub struct PtyState {
    handles: Mutex<HashMap<String, PtyHandle>>,
}

#[derive(Deserialize)]
pub struct SpawnOpts {
    pub cols: u16,
    pub rows: u16,
    pub shell: Option<String>,
    pub cwd: Option<String>,
}

#[derive(Serialize, Clone)]
struct DataPayload {
    id: String,
    data: String,
}

#[derive(Serialize, Clone)]
struct ExitPayload {
    id: String,
    status: Option<i32>,
}

#[tauri::command]
pub async fn pty_spawn(
    app: AppHandle,
    state: State<'_, PtyState>,
    opts: SpawnOpts,
) -> Result<String, String> {
    let id = Uuid::new_v4().to_string();

    let pty_system = native_pty_system();
    let pair = pty_system
        .openpty(PtySize {
            rows: opts.rows.max(1),
            cols: opts.cols.max(1),
            pixel_width: 0,
            pixel_height: 0,
        })
        .map_err(|e| e.to_string())?;

    let shell = opts
        .shell
        .or_else(|| std::env::var("SHELL").ok())
        .unwrap_or_else(|| "/bin/bash".into());

    let mut cmd = CommandBuilder::new(&shell);
    if let Some(cwd) = opts.cwd {
        cmd.cwd(cwd);
    } else if let Ok(home) = std::env::var("HOME") {
        cmd.cwd(home);
    }
    cmd.env("TERM", "xterm-256color");
    cmd.env("COLORTERM", "truecolor");

    let mut child = pair
        .slave
        .spawn_command(cmd)
        .map_err(|e| e.to_string())?;
    drop(pair.slave);

    let reader = pair
        .master
        .try_clone_reader()
        .map_err(|e| e.to_string())?;
    let writer = pair
        .master
        .take_writer()
        .map_err(|e| e.to_string())?;

    let app_for_thread = app.clone();
    let id_for_thread = id.clone();
    std::thread::spawn(move || {
        let mut reader = reader;
        let mut buf = vec![0u8; 4096];
        loop {
            match reader.read(&mut buf) {
                Ok(0) => break,
                Ok(n) => {
                    let payload = DataPayload {
                        id: id_for_thread.clone(),
                        data: B64.encode(&buf[..n]),
                    };
                    let _ = app_for_thread.emit("pty:data", payload);
                }
                Err(_) => break,
            }
        }
        let status = child.wait().ok().map(|s| s.exit_code() as i32);
        let _ = app_for_thread.emit(
            "pty:exit",
            ExitPayload {
                id: id_for_thread,
                status,
            },
        );
    });

    state
        .handles
        .lock()
        .unwrap()
        .insert(id.clone(), PtyHandle { master: pair.master, writer });

    Ok(id)
}

#[tauri::command]
pub fn pty_write(
    state: State<'_, PtyState>,
    id: String,
    data: String,
) -> Result<(), String> {
    let bytes = B64.decode(&data).map_err(|e| e.to_string())?;
    let mut handles = state.handles.lock().unwrap();
    let handle = handles.get_mut(&id).ok_or("pty not found")?;
    handle.writer.write_all(&bytes).map_err(|e| e.to_string())?;
    handle.writer.flush().map_err(|e| e.to_string())?;
    Ok(())
}

#[tauri::command]
pub fn pty_resize(
    state: State<'_, PtyState>,
    id: String,
    cols: u16,
    rows: u16,
) -> Result<(), String> {
    let handles = state.handles.lock().unwrap();
    let handle = handles.get(&id).ok_or("pty not found")?;
    handle
        .master
        .resize(PtySize {
            cols: cols.max(1),
            rows: rows.max(1),
            pixel_width: 0,
            pixel_height: 0,
        })
        .map_err(|e| e.to_string())?;
    Ok(())
}

#[tauri::command]
pub fn pty_kill(state: State<'_, PtyState>, id: String) -> Result<(), String> {
    // Dropping the handle closes the master fd; child receives SIGHUP and exits.
    state.handles.lock().unwrap().remove(&id);
    Ok(())
}