#[cfg(unix)]
use std::io;
use std::{
fmt::Display,
io::{Read, Write},
path::Path,
sync::Arc,
thread,
time::{Duration, Instant},
};
use crossbeam_channel::{Receiver, RecvTimeoutError, Sender, TryRecvError, bounded};
use portable_pty::{CommandBuilder, PtySize as PortablePtySize, native_pty_system};
use crate::{
constants::MUSTER_PROJECT_ENV,
domain::{
agent_session::AgentSessionId,
port::{OutputSink, ProcessHandle, ProcessRunner},
process::StopSignal,
pty::{ExitOutcome, ProcessOutput, PtyError, PtySize, SpawnRequest},
value::CommandLine,
},
};
const READ_BUFFER_BYTES: usize = 4096;
#[cfg(unix)]
const SHELL_PROGRAM: &str = "/bin/sh";
#[cfg(windows)]
const SHELL_PROGRAM: &str = "cmd.exe";
#[cfg(unix)]
const SHELL_COMMAND_FLAG: &str = "-c";
#[cfg(windows)]
const SHELL_COMMAND_FLAG: &str = "/C";
const HOOK_SUBCOMMAND: &str = "hook";
const LAUNCH_SUBCOMMAND: &str = "launch";
const SESSION_ARGUMENT: &str = "--session";
const TERM_VAR: &str = "TERM";
const DEFAULT_TERM: &str = "xterm-256color";
const EXIT_DRAIN_GRACE: Duration = Duration::from_millis(200);
const PROCESS_GROUP_POLL_INTERVAL: Duration = Duration::from_millis(10);
#[cfg(not(unix))]
const SUSPEND_UNSUPPORTED: &str = "suspend and resume are only supported on Unix";
#[derive(Clone, Copy, Default)]
pub struct PortablePtyRunner;
impl ProcessRunner for PortablePtyRunner {
fn spawn(
&self,
request: SpawnRequest,
sink: Box<dyn OutputSink>,
) -> Result<Box<dyn ProcessHandle>, PtyError> {
let pair = native_pty_system()
.openpty(to_portable_size(request.size()))
.map_err(system_error)?;
let mut command = match request.command() {
Some(command_line) => {
if let Some(session_id) = request.agent_session_id() {
Self::agent_launcher(session_id, command_line)?
} else {
let mut builder = CommandBuilder::new(SHELL_PROGRAM);
builder.arg(SHELL_COMMAND_FLAG);
builder.arg(command_line.as_ref());
builder
}
},
None => CommandBuilder::new_default_prog(),
};
if let Some(dir) = request.working_dir() {
if !dir.is_dir() {
return Err(PtyError::InvalidWorkingDir(dir.clone()));
}
command.cwd(dir);
} else if let Ok(cwd) = std::env::current_dir() {
command.cwd(cwd);
}
if std::env::var_os(TERM_VAR).is_none() {
command.env(TERM_VAR, DEFAULT_TERM);
}
if let Some(project) = request.project() {
command.env(MUSTER_PROJECT_ENV, project);
}
for (name, value) in request.environment() {
command.env(name, value);
}
let mut child = pair.slave.spawn_command(command).map_err(system_error)?;
let mut killer = child.clone_killer();
let mut waiter_killer = child.clone_killer();
let pid = child.process_id();
drop(pair.slave);
let io = pair
.master
.try_clone_reader()
.and_then(|reader| pair.master.take_writer().map(|writer| (reader, writer)));
let (mut reader, writer) = match io {
Ok(io) => io,
Err(error) => {
let _ = killer.kill();
return Err(system_error(error));
},
};
let sink: Arc<dyn OutputSink> = Arc::from(sink);
let (reader_done_tx, reader_done_rx) = bounded(1);
let (grace_tx, grace_rx) = bounded(1);
let reader_sink = Arc::clone(&sink);
let reader_handle = thread::spawn(move || {
let mut buffer = [0u8; READ_BUFFER_BYTES];
loop {
match reader.read(&mut buffer) {
Ok(0) | Err(_) => break,
Ok(read) => reader_sink.send(ProcessOutput::Chunk(buffer[..read].to_vec())),
}
}
let _ = reader_done_tx.send(());
});
thread::spawn(move || {
let outcome = match child.wait() {
Ok(status) if status.success() => ExitOutcome::Succeeded,
_ => ExitOutcome::Failed,
};
let drain_grace = grace_rx.try_recv().unwrap_or(EXIT_DRAIN_GRACE);
if wait_for_drain(&reader_done_rx, drain_grace, || process_group_is_alive(pid))
== DrainStatus::TimedOut
{
let _ = terminate_group(pid, &mut waiter_killer);
let _ = reader_done_rx.recv();
}
let _ = reader_handle.join();
sink.send(ProcessOutput::Exited(outcome));
});
Ok(Box::new(PtyProcessHandle {
master: pair.master,
writer,
killer,
grace_tx,
pid,
}))
}
}
impl PortablePtyRunner {
fn agent_launcher(
session_id: &AgentSessionId,
command: &CommandLine,
) -> Result<CommandBuilder, PtyError> {
let executable = std::env::current_exe().map_err(system_error)?;
Ok(Self::agent_launcher_for(&executable, session_id, command))
}
fn agent_launcher_for(
executable: &Path,
session_id: &AgentSessionId,
command: &CommandLine,
) -> CommandBuilder {
let mut launcher = CommandBuilder::new(executable);
launcher.args([
HOOK_SUBCOMMAND,
LAUNCH_SUBCOMMAND,
SESSION_ARGUMENT,
session_id.as_ref(),
"--",
command.as_ref(),
]);
launcher
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum DrainStatus {
Complete,
TimedOut,
}
fn wait_for_drain<F>(
reader_done_rx: &Receiver<()>,
grace: Duration,
mut group_is_alive: F,
) -> DrainStatus
where
F: FnMut() -> bool,
{
let started = Instant::now();
let mut reader_done = match reader_done_rx.try_recv() {
Ok(()) | Err(TryRecvError::Disconnected) => true,
Err(TryRecvError::Empty) => false,
};
loop {
if reader_done && !group_is_alive() {
return DrainStatus::Complete;
}
let remaining = grace.saturating_sub(started.elapsed());
if remaining.is_zero() {
return DrainStatus::TimedOut;
}
let poll = remaining.min(PROCESS_GROUP_POLL_INTERVAL);
if reader_done {
thread::sleep(poll);
continue;
}
match reader_done_rx.recv_timeout(poll) {
Ok(()) | Err(RecvTimeoutError::Disconnected) => reader_done = true,
Err(RecvTimeoutError::Timeout) => {},
}
}
}
fn process_group_is_alive(pid: Option<u32>) -> bool {
#[cfg(unix)]
{
let Some(pid) = pid else {
return false;
};
if unsafe { libc::kill(-(pid as libc::pid_t), 0) } == 0 {
return true;
}
io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
}
#[cfg(not(unix))]
{
let _ = pid;
false
}
}
struct PtyProcessHandle {
master: Box<dyn portable_pty::MasterPty + Send>,
writer: Box<dyn Write + Send>,
killer: Box<dyn portable_pty::ChildKiller + Send + Sync>,
grace_tx: Sender<Duration>,
pid: Option<u32>,
}
impl ProcessHandle for PtyProcessHandle {
fn process_id(&self) -> Option<u32> {
self.pid
}
fn write_input(&mut self, bytes: &[u8]) -> Result<(), PtyError> {
self.writer.write_all(bytes)?;
self.writer.flush()?;
Ok(())
}
fn resize(&mut self, size: PtySize) -> Result<(), PtyError> {
self.master
.resize(to_portable_size(size))
.map_err(system_error)
}
fn pause(&mut self) -> Result<(), PtyError> {
#[cfg(unix)]
{
signal_group(self.pid, libc::SIGSTOP)
}
#[cfg(not(unix))]
{
Err(PtyError::Unsupported(SUSPEND_UNSUPPORTED.to_string()))
}
}
fn resume(&mut self) -> Result<(), PtyError> {
#[cfg(unix)]
{
signal_group(self.pid, libc::SIGCONT)
}
#[cfg(not(unix))]
{
Err(PtyError::Unsupported(SUSPEND_UNSUPPORTED.to_string()))
}
}
fn terminate(&mut self, signal: StopSignal, grace: Duration) -> Result<(), PtyError> {
let _ = self.grace_tx.try_send(grace);
#[cfg(unix)]
{
let _ = signal_group(self.pid, libc::SIGCONT);
signal_group(self.pid, unix_signal(signal))
}
#[cfg(not(unix))]
{
let _ = signal;
self.killer.kill().map_err(system_error)
}
}
fn kill(&mut self) -> Result<(), PtyError> {
terminate_group(self.pid, &mut self.killer)
}
}
#[cfg(unix)]
fn unix_signal(signal: StopSignal) -> libc::c_int {
match signal {
StopSignal::Terminate => libc::SIGTERM,
StopSignal::Interrupt => libc::SIGINT,
}
}
fn to_portable_size(size: PtySize) -> PortablePtySize {
PortablePtySize {
rows: size.rows().into_inner(),
cols: size.cols().into_inner(),
pixel_width: 0,
pixel_height: 0,
}
}
fn system_error<E: Display>(error: E) -> PtyError {
PtyError::System(error.to_string())
}
#[cfg(unix)]
fn signal_group(pid: Option<u32>, signal: libc::c_int) -> Result<(), PtyError> {
let pid = pid.ok_or_else(|| PtyError::Unsupported("no pid to signal".to_string()))?;
if unsafe { libc::kill(-(pid as libc::pid_t), signal) } == -1 {
return Err(PtyError::Io(std::io::Error::last_os_error()));
}
Ok(())
}
fn terminate_group(
pid: Option<u32>,
killer: &mut Box<dyn portable_pty::ChildKiller + Send + Sync>,
) -> Result<(), PtyError> {
#[cfg(unix)]
{
let group_result = signal_group(pid, libc::SIGKILL);
let child_result = killer.kill().map_err(system_error);
match (group_result, child_result) {
(Ok(()), _) | (_, Ok(())) => Ok(()),
(Err(error), Err(_)) => Err(error),
}
}
#[cfg(not(unix))]
{
let _ = pid;
killer.kill().map_err(system_error)
}
}
#[cfg(test)]
mod tests {
use std::{
cell::Cell,
path::PathBuf,
time::{Duration, Instant},
};
use super::*;
use crate::domain::{
agent_session::AgentSessionId,
value::{Cols, CommandLine, Rows},
};
const OUTPUT_TIMEOUT: Duration = Duration::from_secs(5);
const TEST_STOP_GRACE: Duration = Duration::from_secs(3);
#[cfg(target_os = "macos")]
const DESCENDANT_CLEANUP_MIN_WAIT: Duration = Duration::from_millis(750);
const DESCENDANT_CLEANUP_COMMAND: &str = r#"sh -c 'trap "" TERM HUP; printf ready; sleep 1; printf descendant-clean; sleep 1' & trap 'exit 0' TERM; wait"#;
struct ChannelSink(crossbeam_channel::Sender<ProcessOutput>);
impl OutputSink for ChannelSink {
fn send(&self, output: ProcessOutput) {
let _ = self.0.send(output);
}
}
fn request(command: &str) -> SpawnRequest {
SpawnRequest::builder()
.command(Some(CommandLine::try_new(command).unwrap()))
.size(
PtySize::builder()
.rows(Rows::new(24))
.cols(Cols::new(80))
.build(),
)
.build()
}
#[test]
fn agent_launch_handoff_preserves_shell_expressions() {
let session = AgentSessionId::try_new("session $(injection)").unwrap();
let command = CommandLine::try_new("FOO=bar codex | tee agent.log").unwrap();
let launch = PortablePtyRunner::agent_launcher_for(
&PathBuf::from("/tmp/muster"),
&session,
&command,
);
assert_eq!(launch.get_argv(), &vec![
PathBuf::from("/tmp/muster").into_os_string(),
HOOK_SUBCOMMAND.into(),
LAUNCH_SUBCOMMAND.into(),
SESSION_ARGUMENT.into(),
"session $(injection)".into(),
"--".into(),
"FOO=bar codex | tee agent.log".into(),
]);
}
#[test]
fn drain_waits_for_the_process_group_after_the_reader_finishes() {
let (tx, rx) = bounded(1);
tx.send(()).unwrap();
drop(tx);
let polls = Cell::new(0);
let status = wait_for_drain(&rx, OUTPUT_TIMEOUT, || {
let next = polls.get() + 1;
polls.set(next);
next < 3
});
assert_eq!(status, DrainStatus::Complete);
assert_eq!(polls.get(), 3);
}
#[test]
fn streams_output_then_reports_success() {
let (tx, rx) = crossbeam_channel::unbounded();
let _handle = PortablePtyRunner
.spawn(request("printf hello"), Box::new(ChannelSink(tx)))
.unwrap();
let mut bytes = Vec::new();
let mut outcome = None;
while let Ok(output) = rx.recv_timeout(OUTPUT_TIMEOUT) {
match output {
ProcessOutput::Chunk(chunk) => bytes.extend_from_slice(&chunk),
ProcessOutput::Exited(exit) => outcome = Some(exit),
}
}
assert!(String::from_utf8_lossy(&bytes).contains("hello"));
assert_eq!(outcome, Some(ExitOutcome::Succeeded));
}
#[test]
fn final_output_is_delivered_before_exit() {
let (tx, rx) = crossbeam_channel::unbounded();
let _handle = PortablePtyRunner
.spawn(request("printf hello; exit 1"), Box::new(ChannelSink(tx)))
.unwrap();
let mut events = Vec::new();
while let Ok(output) = rx.recv_timeout(OUTPUT_TIMEOUT) {
events.push(output);
}
let exit_pos = events
.iter()
.position(|event| matches!(event, ProcessOutput::Exited(_)))
.expect("an exit is reported");
assert_eq!(exit_pos, events.len() - 1, "exit must be the final event");
let output: Vec<u8> = events[..exit_pos]
.iter()
.flat_map(|event| match event {
ProcessOutput::Chunk(chunk) => chunk.clone(),
ProcessOutput::Exited(_) => Vec::new(),
})
.collect();
assert!(String::from_utf8_lossy(&output).contains("hello"));
}
#[cfg(unix)]
#[test]
fn graceful_termination_drains_shutdown_output() {
let (tx, rx) = crossbeam_channel::unbounded();
let mut handle = PortablePtyRunner
.spawn(
request(
"trap 'printf shutdown; exit 0' TERM; printf ready; while :; do sleep 1; done",
),
Box::new(ChannelSink(tx)),
)
.unwrap();
let mut bytes = Vec::new();
while !String::from_utf8_lossy(&bytes).contains("ready") {
match rx.recv_timeout(OUTPUT_TIMEOUT).unwrap() {
ProcessOutput::Chunk(chunk) => bytes.extend(chunk),
ProcessOutput::Exited(_) => panic!("command exited before termination"),
}
}
handle
.terminate(StopSignal::Terminate, TEST_STOP_GRACE)
.unwrap();
while let Ok(output) = rx.recv_timeout(OUTPUT_TIMEOUT) {
match output {
ProcessOutput::Chunk(chunk) => bytes.extend(chunk),
ProcessOutput::Exited(_) => break,
}
}
assert!(String::from_utf8_lossy(&bytes).contains("shutdown"));
}
#[cfg(unix)]
#[test]
fn graceful_termination_preserves_descendant_cleanup_after_shell_exit() {
let (tx, rx) = crossbeam_channel::unbounded();
let mut handle = PortablePtyRunner
.spawn(
request(DESCENDANT_CLEANUP_COMMAND),
Box::new(ChannelSink(tx)),
)
.unwrap();
let mut bytes = Vec::new();
while !String::from_utf8_lossy(&bytes).contains("ready") {
match rx.recv_timeout(OUTPUT_TIMEOUT).unwrap() {
ProcessOutput::Chunk(chunk) => bytes.extend(chunk),
ProcessOutput::Exited(_) => panic!("command exited before termination"),
}
}
#[cfg(target_os = "macos")]
let termination_started = Instant::now();
handle
.terminate(StopSignal::Terminate, TEST_STOP_GRACE)
.unwrap();
let mut exited = false;
while let Ok(output) = rx.recv_timeout(OUTPUT_TIMEOUT) {
match output {
ProcessOutput::Chunk(chunk) => bytes.extend(chunk),
ProcessOutput::Exited(_) => {
exited = true;
break;
},
}
}
assert!(exited, "the runner must report the completed cleanup");
#[cfg(target_os = "macos")]
assert!(
termination_started.elapsed() >= DESCENDANT_CLEANUP_MIN_WAIT,
"an early macOS PTY EOF must not bypass descendant cleanup"
);
#[cfg(not(target_os = "macos"))]
assert!(String::from_utf8_lossy(&bytes).contains("descendant-clean"));
}
#[test]
fn a_slow_drain_delivers_all_output_before_exit() {
const OUTPUT_LEN: usize = 200_000;
let (tx, rx) = crossbeam_channel::bounded(1);
let _handle = PortablePtyRunner
.spawn(
request(r"head -c 200000 /dev/zero | tr '\0' x"),
Box::new(ChannelSink(tx)),
)
.unwrap();
let mut events = Vec::new();
while let Ok(output) = rx.recv_timeout(OUTPUT_TIMEOUT) {
events.push(output);
thread::sleep(Duration::from_millis(5));
}
let exit_pos = events
.iter()
.position(|event| matches!(event, ProcessOutput::Exited(_)))
.expect("an exit is reported");
assert_eq!(exit_pos, events.len() - 1, "exit must be the final event");
let total: usize = events[..exit_pos]
.iter()
.map(|event| match event {
ProcessOutput::Chunk(chunk) => chunk.len(),
ProcessOutput::Exited(_) => 0,
})
.sum();
assert_eq!(
total, OUTPUT_LEN,
"all output delivered, not truncated by the exit"
);
}
#[test]
fn reports_failure_for_nonzero_exit() {
let (tx, rx) = crossbeam_channel::unbounded();
let _handle = PortablePtyRunner
.spawn(request("exit 3"), Box::new(ChannelSink(tx)))
.unwrap();
let mut outcome = None;
while let Ok(output) = rx.recv_timeout(OUTPUT_TIMEOUT) {
if let ProcessOutput::Exited(exit) = output {
outcome = Some(exit);
}
}
assert_eq!(outcome, Some(ExitOutcome::Failed));
}
#[test]
fn invalid_working_directory_is_rejected() {
let (tx, _rx) = crossbeam_channel::unbounded();
let request = SpawnRequest::builder()
.command(Some(CommandLine::try_new("true").unwrap()))
.working_dir(Some(PathBuf::from("/no/such/muster/dir")))
.size(
PtySize::builder()
.rows(Rows::new(24))
.cols(Cols::new(80))
.build(),
)
.build();
let result = PortablePtyRunner.spawn(request, Box::new(ChannelSink(tx)));
assert!(matches!(result, Err(PtyError::InvalidWorkingDir(_))));
}
#[test]
fn inherits_the_current_directory_when_no_working_dir_is_set() {
let (tx, rx) = crossbeam_channel::unbounded();
let _handle = PortablePtyRunner
.spawn(request("pwd -P"), Box::new(ChannelSink(tx)))
.unwrap();
let mut bytes = Vec::new();
while let Ok(output) = rx.recv_timeout(OUTPUT_TIMEOUT) {
if let ProcessOutput::Chunk(chunk) = output {
bytes.extend_from_slice(&chunk);
}
}
let expected = std::env::current_dir().unwrap();
assert!(String::from_utf8_lossy(&bytes).contains(expected.to_str().unwrap()));
}
#[test]
fn exit_is_observed_before_a_backgrounded_descendant_closes_the_pty() {
let (tx, rx) = crossbeam_channel::unbounded();
let start = Instant::now();
let _handle = PortablePtyRunner
.spawn(request("sleep 5 &"), Box::new(ChannelSink(tx)))
.unwrap();
loop {
match rx.recv_timeout(OUTPUT_TIMEOUT) {
Ok(ProcessOutput::Exited(_)) => break,
Ok(_) => {},
Err(_) => break,
}
}
assert!(start.elapsed() < Duration::from_secs(2));
}
}