diff --git a/crates/spacesh-pty/src/lib.rs b/crates/spacesh-pty/src/lib.rs index 9736bfa..b2946a1 100644 --- a/crates/spacesh-pty/src/lib.rs +++ b/crates/spacesh-pty/src/lib.rs @@ -1 +1,148 @@ -// PtyHandle implemented in Task 3. +use std::io::{Read, Write}; +use anyhow::Result; +use portable_pty::{CommandBuilder, MasterPty, PtySize, native_pty_system}; +use tokio::sync::mpsc; + +/// A spawned PTY with its child process. Output chunks arrive on `output`. +pub struct PtyHandle { + master: Box, + writer: Box, + child: Box, + /// Raw output chunks read off the PTY master (already on the async side). + pub output: mpsc::Receiver>, +} + +/// Parameters for spawning a surface's process. +pub struct SpawnSpec { + pub command: String, + pub args: Vec, + pub cwd: std::path::PathBuf, + pub cols: u16, + pub rows: u16, + /// Extra environment variables (e.g. SPACESH_SURFACE_ID). + pub env: Vec<(String, String)>, +} + +impl PtyHandle { + pub fn spawn(spec: SpawnSpec) -> Result { + let pty_system = native_pty_system(); + let pair = pty_system.openpty(PtySize { + rows: spec.rows, + cols: spec.cols, + pixel_width: 0, + pixel_height: 0, + })?; + + let mut cmd = CommandBuilder::new(&spec.command); + for a in &spec.args { + cmd.arg(a); + } + cmd.cwd(&spec.cwd); + for (k, v) in &spec.env { + cmd.env(k, v); + } + + let child = pair.slave.spawn_command(cmd)?; + // The slave handle must be dropped so the child is the only holder; otherwise + // EOF is never observed on the master after the child exits. + drop(pair.slave); + + let writer = pair.master.take_writer()?; + let mut reader = pair.master.try_clone_reader()?; + + let (tx, rx) = mpsc::channel::>(256); + std::thread::spawn(move || { + let mut buf = [0u8; 8192]; + loop { + match reader.read(&mut buf) { + Ok(0) => break, // EOF: child closed the pty + Ok(n) => { + if tx.blocking_send(buf[..n].to_vec()).is_err() { + break; // receiver gone + } + } + Err(_) => break, + } + } + }); + + Ok(Self { + master: pair.master, + writer, + child, + output: rx, + }) + } + + pub fn write_input(&mut self, bytes: &[u8]) -> Result<()> { + self.writer.write_all(bytes)?; + self.writer.flush()?; + Ok(()) + } + + pub fn resize(&self, cols: u16, rows: u16) -> Result<()> { + self.master.resize(PtySize { rows, cols, pixel_width: 0, pixel_height: 0 })?; + Ok(()) + } + + /// Best-effort wait for the child's exit code (blocking). + pub fn wait(&mut self) -> i32 { + match self.child.wait() { + Ok(status) => status.exit_code() as i32, + Err(_) => -1, + } + } + + pub fn kill(&mut self) { + let _ = self.child.kill(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn shell_spec(script: &str) -> SpawnSpec { + SpawnSpec { + command: "/bin/sh".into(), + args: vec!["-c".into(), script.into()], + cwd: std::env::temp_dir(), + cols: 80, + rows: 24, + env: vec![("SPACESH_SURFACE_ID".into(), "s_test".into())], + } + } + + #[tokio::test] + async fn spawn_echo_produces_output() { + let mut handle = PtyHandle::spawn(shell_spec("printf SPACESH_OK")).unwrap(); + let mut collected = Vec::new(); + // Drain until EOF (channel closes when the reader thread sees EOF). + while let Some(chunk) = handle.output.recv().await { + collected.extend_from_slice(&chunk); + } + let text = String::from_utf8_lossy(&collected); + assert!(text.contains("SPACESH_OK"), "got: {text:?}"); + } + + #[tokio::test] + async fn resize_does_not_error() { + let handle = PtyHandle::spawn(shell_spec("sleep 0.2")).unwrap(); + handle.resize(120, 40).unwrap(); + } + + #[tokio::test] + async fn input_is_echoed_back() { + // `cat` echoes stdin back to stdout on a pty. + let mut handle = PtyHandle::spawn(shell_spec("cat")).unwrap(); + handle.write_input(b"hello\n").unwrap(); + let mut collected = Vec::new(); + // Read a few chunks then kill cat to end the stream. + if let Some(chunk) = handle.output.recv().await { + collected.extend_from_slice(&chunk); + } + handle.kill(); + let text = String::from_utf8_lossy(&collected); + assert!(text.contains("hello"), "got: {text:?}"); + } +}