mirror of
https://github.com/openai/codex.git
synced 2026-09-14 11:57:03 +00:00
Merge 184d0e7a5b into sapling-pr-archive-bolinfest
This commit is contained in:
@@ -48,12 +48,23 @@ mod imp {
|
||||
use crate::unified_exec::SpawnLifecycle;
|
||||
use codex_shell_escalation::EscalationSession;
|
||||
|
||||
const ESCALATE_SOCKET_ENV_VAR: &str = "CODEX_ESCALATE_SOCKET";
|
||||
|
||||
#[derive(Debug)]
|
||||
struct ZshForkSpawnLifecycle {
|
||||
escalation_session: EscalationSession,
|
||||
}
|
||||
|
||||
impl SpawnLifecycle for ZshForkSpawnLifecycle {
|
||||
fn inherited_fds(&self) -> Vec<std::os::fd::RawFd> {
|
||||
self.escalation_session
|
||||
.env()
|
||||
.get(ESCALATE_SOCKET_ENV_VAR)
|
||||
.and_then(|fd| fd.parse().ok())
|
||||
.into_iter()
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn after_spawn(&mut self) {
|
||||
self.escalation_session.close_client_socket();
|
||||
}
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
#![allow(clippy::module_inception)]
|
||||
|
||||
use std::os::fd::RawFd;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::atomic::Ordering;
|
||||
@@ -25,6 +26,10 @@ use super::UnifiedExecError;
|
||||
use super::head_tail_buffer::HeadTailBuffer;
|
||||
|
||||
pub(crate) trait SpawnLifecycle: std::fmt::Debug + Send + Sync {
|
||||
fn inherited_fds(&self) -> Vec<RawFd> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn after_spawn(&mut self) {}
|
||||
}
|
||||
|
||||
|
||||
@@ -535,14 +535,16 @@ impl UnifiedExecProcessManager {
|
||||
.command
|
||||
.split_first()
|
||||
.ok_or(UnifiedExecError::MissingCommandLine)?;
|
||||
let inherited_fds = spawn_lifecycle.inherited_fds();
|
||||
|
||||
let spawn_result = if tty {
|
||||
codex_utils_pty::pty::spawn_process(
|
||||
codex_utils_pty::pty::spawn_process_with_inherited_fds(
|
||||
program,
|
||||
args,
|
||||
env.cwd.as_path(),
|
||||
&env.env,
|
||||
&env.arg0,
|
||||
&inherited_fds,
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use anyhow::Context;
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
use std::process::Stdio;
|
||||
use std::time::Duration;
|
||||
|
||||
pub async fn wait_for_pid_file(path: &Path) -> anyhow::Result<String> {
|
||||
@@ -24,6 +25,7 @@ pub async fn wait_for_pid_file(path: &Path) -> anyhow::Result<String> {
|
||||
pub fn process_is_alive(pid: &str) -> anyhow::Result<bool> {
|
||||
let status = std::process::Command::new("kill")
|
||||
.args(["-0", pid])
|
||||
.stderr(Stdio::null())
|
||||
.status()
|
||||
.context("failed to probe process liveness with kill -0")?;
|
||||
Ok(status.success())
|
||||
|
||||
@@ -32,6 +32,8 @@ use core_test_support::test_codex::test_codex;
|
||||
use core_test_support::wait_for_event;
|
||||
use core_test_support::wait_for_event_match;
|
||||
use core_test_support::wait_for_event_with_timeout;
|
||||
use core_test_support::zsh_fork::build_unified_exec_zsh_fork_test;
|
||||
use core_test_support::zsh_fork::zsh_fork_runtime;
|
||||
use pretty_assertions::assert_eq;
|
||||
use regex_lite::Regex;
|
||||
use serde_json::Value;
|
||||
@@ -1323,7 +1325,6 @@ async fn exec_command_reports_chunk_and_exit_metadata() -> Result<()> {
|
||||
.into_iter()
|
||||
.map(|request| request.body_json())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let outputs = collect_tool_outputs(&bodies)?;
|
||||
let metadata = outputs
|
||||
.get(call_id)
|
||||
@@ -1445,7 +1446,6 @@ async fn unified_exec_defaults_to_pipe() -> Result<()> {
|
||||
.into_iter()
|
||||
.map(|request| request.body_json())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let outputs = collect_tool_outputs(&bodies)?;
|
||||
let output = outputs
|
||||
.get(call_id)
|
||||
@@ -1539,7 +1539,6 @@ async fn unified_exec_can_enable_tty() -> Result<()> {
|
||||
.into_iter()
|
||||
.map(|request| request.body_json())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let outputs = collect_tool_outputs(&bodies)?;
|
||||
let output = outputs
|
||||
.get(call_id)
|
||||
@@ -1624,7 +1623,6 @@ async fn unified_exec_respects_early_exit_notifications() -> Result<()> {
|
||||
.into_iter()
|
||||
.map(|request| request.body_json())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let outputs = collect_tool_outputs(&bodies)?;
|
||||
let output = outputs
|
||||
.get(call_id)
|
||||
@@ -1823,6 +1821,227 @@ async fn write_stdin_returns_exit_metadata_and_clears_session() -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[cfg(unix)]
|
||||
async fn unified_exec_zsh_fork_keeps_python_repl_attached_to_zsh_session() -> Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
let Some(runtime) = zsh_fork_runtime("unified exec zsh-fork tty session test")? else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let python = match which("python3") {
|
||||
Ok(path) => path,
|
||||
Err(_) => {
|
||||
eprintln!("python3 not found in PATH, skipping zsh-fork python repl test.");
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
|
||||
let server = start_mock_server().await;
|
||||
let test = build_unified_exec_zsh_fork_test(
|
||||
&server,
|
||||
runtime,
|
||||
AskForApproval::Never,
|
||||
SandboxPolicy::new_workspace_write_policy(),
|
||||
|_| {},
|
||||
)
|
||||
.await?;
|
||||
|
||||
let start_call_id = "uexec-zsh-fork-python-start";
|
||||
let send_call_id = "uexec-zsh-fork-python-pid";
|
||||
let exit_call_id = "uexec-zsh-fork-python-exit";
|
||||
|
||||
let start_args = serde_json::json!({
|
||||
"cmd": python.display().to_string(),
|
||||
"yield_time_ms": 500,
|
||||
"tty": true,
|
||||
});
|
||||
let send_args = serde_json::json!({
|
||||
"chars": "import os; print(os.getpid())\r\n",
|
||||
"session_id": 1000,
|
||||
"yield_time_ms": 500,
|
||||
});
|
||||
let exit_args = serde_json::json!({
|
||||
"chars": "import sys; sys.exit(0)\r\n",
|
||||
"session_id": 1000,
|
||||
"yield_time_ms": 500,
|
||||
});
|
||||
|
||||
let responses = vec![
|
||||
sse(vec![
|
||||
ev_response_created("resp-1"),
|
||||
ev_function_call(
|
||||
start_call_id,
|
||||
"exec_command",
|
||||
&serde_json::to_string(&start_args)?,
|
||||
),
|
||||
ev_completed("resp-1"),
|
||||
]),
|
||||
sse(vec![
|
||||
ev_response_created("resp-2"),
|
||||
ev_function_call(
|
||||
send_call_id,
|
||||
"write_stdin",
|
||||
&serde_json::to_string(&send_args)?,
|
||||
),
|
||||
ev_completed("resp-2"),
|
||||
]),
|
||||
sse(vec![
|
||||
ev_assistant_message("msg-1", "python is running"),
|
||||
ev_completed("resp-3"),
|
||||
]),
|
||||
sse(vec![
|
||||
ev_response_created("resp-4"),
|
||||
ev_function_call(
|
||||
exit_call_id,
|
||||
"write_stdin",
|
||||
&serde_json::to_string(&exit_args)?,
|
||||
),
|
||||
ev_completed("resp-4"),
|
||||
]),
|
||||
sse(vec![
|
||||
ev_assistant_message("msg-2", "all done"),
|
||||
ev_completed("resp-5"),
|
||||
]),
|
||||
];
|
||||
let request_log = mount_sse_sequence(&server, responses).await;
|
||||
|
||||
test.codex
|
||||
.submit(Op::UserTurn {
|
||||
items: vec![UserInput::Text {
|
||||
text: "test unified exec zsh-fork tty behavior".into(),
|
||||
text_elements: Vec::new(),
|
||||
}],
|
||||
final_output_json_schema: None,
|
||||
cwd: test.cwd_path().to_path_buf(),
|
||||
approval_policy: AskForApproval::Never,
|
||||
sandbox_policy: SandboxPolicy::new_workspace_write_policy(),
|
||||
model: test.session_configured.model.clone(),
|
||||
effort: None,
|
||||
summary: None,
|
||||
service_tier: None,
|
||||
collaboration_mode: None,
|
||||
personality: None,
|
||||
})
|
||||
.await?;
|
||||
|
||||
wait_for_event(&test.codex, |event| {
|
||||
matches!(event, EventMsg::TurnComplete(_))
|
||||
})
|
||||
.await;
|
||||
|
||||
let requests = request_log.requests();
|
||||
assert!(!requests.is_empty(), "expected at least one POST request");
|
||||
let bodies = requests
|
||||
.into_iter()
|
||||
.map(|request| request.body_json())
|
||||
.collect::<Vec<_>>();
|
||||
let outputs = collect_tool_outputs(&bodies)?;
|
||||
|
||||
let start_output = outputs
|
||||
.get(start_call_id)
|
||||
.expect("missing start output for exec_command");
|
||||
let process_id = start_output
|
||||
.process_id
|
||||
.clone()
|
||||
.expect("expected process id from exec_command");
|
||||
assert!(
|
||||
start_output.exit_code.is_none(),
|
||||
"initial exec_command should leave the PTY session running"
|
||||
);
|
||||
|
||||
let send_output = outputs
|
||||
.get(send_call_id)
|
||||
.expect("missing write_stdin output");
|
||||
let normalized = send_output.output.replace("\r\n", "\n");
|
||||
let python_pid = Regex::new(r"(?m)^(\d+)$")
|
||||
.expect("valid python pid regex")
|
||||
.captures(&normalized)
|
||||
.and_then(|captures| captures.get(1))
|
||||
.map(|value| value.as_str().to_string())
|
||||
.with_context(|| format!("missing python pid in output {normalized:?}"))?;
|
||||
assert!(
|
||||
process_is_alive(&python_pid)?,
|
||||
"python process should still be alive after printing its pid, got output {normalized:?}"
|
||||
);
|
||||
assert_eq!(send_output.process_id.as_deref(), Some(process_id.as_str()));
|
||||
assert!(
|
||||
send_output.exit_code.is_none(),
|
||||
"write_stdin should not report an exit code while the process is still running"
|
||||
);
|
||||
|
||||
let zsh_pid = std::process::Command::new("ps")
|
||||
.args(["-o", "ppid=", "-p", &python_pid])
|
||||
.output()
|
||||
.context("failed to look up python parent pid")?;
|
||||
let zsh_pid = String::from_utf8(zsh_pid.stdout)
|
||||
.context("python parent pid output is not UTF-8")?
|
||||
.trim()
|
||||
.to_string();
|
||||
assert!(
|
||||
!zsh_pid.is_empty(),
|
||||
"expected python parent pid to identify the zsh session"
|
||||
);
|
||||
assert!(
|
||||
process_is_alive(&zsh_pid)?,
|
||||
"expected zsh parent process {zsh_pid} to still be alive"
|
||||
);
|
||||
|
||||
let zsh_command = std::process::Command::new("ps")
|
||||
.args(["-o", "command=", "-p", &zsh_pid])
|
||||
.output()
|
||||
.context("failed to look up zsh parent command")?;
|
||||
let zsh_command =
|
||||
String::from_utf8(zsh_command.stdout).context("zsh parent command output is not UTF-8")?;
|
||||
assert!(
|
||||
zsh_command.contains("zsh"),
|
||||
"expected python parent command to be zsh, got {zsh_command:?}"
|
||||
);
|
||||
|
||||
test.codex
|
||||
.submit(Op::UserTurn {
|
||||
items: vec![UserInput::Text {
|
||||
text: "shut down the python repl".into(),
|
||||
text_elements: Vec::new(),
|
||||
}],
|
||||
final_output_json_schema: None,
|
||||
cwd: test.cwd_path().to_path_buf(),
|
||||
approval_policy: AskForApproval::Never,
|
||||
sandbox_policy: SandboxPolicy::new_workspace_write_policy(),
|
||||
model: test.session_configured.model.clone(),
|
||||
effort: None,
|
||||
summary: None,
|
||||
service_tier: None,
|
||||
collaboration_mode: None,
|
||||
personality: None,
|
||||
})
|
||||
.await?;
|
||||
|
||||
wait_for_event(&test.codex, |event| {
|
||||
matches!(event, EventMsg::TurnComplete(_))
|
||||
})
|
||||
.await;
|
||||
|
||||
let requests = request_log.requests();
|
||||
assert!(!requests.is_empty(), "expected at least one POST request");
|
||||
let bodies = requests
|
||||
.into_iter()
|
||||
.map(|request| request.body_json())
|
||||
.collect::<Vec<_>>();
|
||||
let outputs = collect_tool_outputs(&bodies)?;
|
||||
let exit_output = outputs
|
||||
.get(exit_call_id)
|
||||
.expect("missing exit output after requesting python shutdown");
|
||||
assert!(
|
||||
exit_output.exit_code.is_none() || exit_output.exit_code == Some(0),
|
||||
"exit request should either leave cleanup to the background watcher or report success directly, got {exit_output:?}"
|
||||
);
|
||||
wait_for_process_exit(&python_pid).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn unified_exec_emits_end_event_when_session_dies_via_stdin() -> Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use std::io;
|
||||
use std::os::fd::AsFd;
|
||||
use std::os::fd::AsRawFd;
|
||||
use std::os::fd::FromRawFd as _;
|
||||
use std::os::fd::OwnedFd;
|
||||
|
||||
use anyhow::Context as _;
|
||||
@@ -28,6 +28,12 @@ fn get_escalate_client() -> anyhow::Result<AsyncDatagramSocket> {
|
||||
Ok(unsafe { AsyncDatagramSocket::from_raw_fd(client_fd) }?)
|
||||
}
|
||||
|
||||
fn duplicate_fd_for_transfer(fd: impl AsFd, name: &str) -> anyhow::Result<OwnedFd> {
|
||||
fd.as_fd()
|
||||
.try_clone_to_owned()
|
||||
.with_context(|| format!("failed to duplicate {name} for escalation transfer"))
|
||||
}
|
||||
|
||||
pub async fn run_shell_escalation_execve_wrapper(
|
||||
file: String,
|
||||
argv: Vec<String>,
|
||||
@@ -62,11 +68,18 @@ pub async fn run_shell_escalation_execve_wrapper(
|
||||
.context("failed to receive EscalateResponse")?;
|
||||
match message.action {
|
||||
EscalateAction::Escalate => {
|
||||
// TODO: maybe we should send ALL open FDs (except the escalate client)?
|
||||
// Duplicate stdio before transferring ownership to the server. The
|
||||
// wrapper must keep using its own stdin/stdout/stderr until the
|
||||
// escalated child takes over.
|
||||
let destination_fds = [
|
||||
io::stdin().as_raw_fd(),
|
||||
io::stdout().as_raw_fd(),
|
||||
io::stderr().as_raw_fd(),
|
||||
];
|
||||
let fds_to_send = [
|
||||
unsafe { OwnedFd::from_raw_fd(io::stdin().as_raw_fd()) },
|
||||
unsafe { OwnedFd::from_raw_fd(io::stdout().as_raw_fd()) },
|
||||
unsafe { OwnedFd::from_raw_fd(io::stderr().as_raw_fd()) },
|
||||
duplicate_fd_for_transfer(io::stdin(), "stdin")?,
|
||||
duplicate_fd_for_transfer(io::stdout(), "stdout")?,
|
||||
duplicate_fd_for_transfer(io::stderr(), "stderr")?,
|
||||
];
|
||||
|
||||
// TODO: also forward signals over the super-exec socket
|
||||
@@ -74,7 +87,7 @@ pub async fn run_shell_escalation_execve_wrapper(
|
||||
client
|
||||
.send_with_fds(
|
||||
SuperExecMessage {
|
||||
fds: fds_to_send.iter().map(AsRawFd::as_raw_fd).collect(),
|
||||
fds: destination_fds.into_iter().collect(),
|
||||
},
|
||||
&fds_to_send,
|
||||
)
|
||||
@@ -115,3 +128,23 @@ pub async fn run_shell_escalation_execve_wrapper(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::os::fd::AsRawFd;
|
||||
use std::os::unix::net::UnixStream;
|
||||
|
||||
#[test]
|
||||
fn duplicate_fd_for_transfer_does_not_close_original() {
|
||||
let (left, _right) = UnixStream::pair().expect("socket pair");
|
||||
let original_fd = left.as_raw_fd();
|
||||
|
||||
let duplicate = duplicate_fd_for_transfer(&left, "test fd").expect("duplicate fd");
|
||||
assert_ne!(duplicate.as_raw_fd(), original_fd);
|
||||
|
||||
drop(duplicate);
|
||||
|
||||
assert_ne!(unsafe { libc::fcntl(original_fd, libc::F_GETFD) }, -1);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,11 +1,10 @@
|
||||
use core::fmt;
|
||||
use std::any::Any;
|
||||
use std::io;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex as StdMutex;
|
||||
|
||||
use portable_pty::MasterPty;
|
||||
use portable_pty::SlavePty;
|
||||
use tokio::sync::broadcast;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
@@ -17,8 +16,8 @@ pub(crate) trait ChildTerminator: Send + Sync {
|
||||
}
|
||||
|
||||
pub struct PtyHandles {
|
||||
pub _slave: Option<Box<dyn SlavePty + Send>>,
|
||||
pub _master: Box<dyn MasterPty + Send>,
|
||||
pub _slave: Option<Box<dyn Any + Send>>,
|
||||
pub _master: Box<dyn Any + Send>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for PtyHandles {
|
||||
|
||||
@@ -1,6 +1,18 @@
|
||||
use std::collections::HashMap;
|
||||
#[cfg(unix)]
|
||||
use std::fs::File;
|
||||
use std::io::ErrorKind;
|
||||
#[cfg(unix)]
|
||||
use std::os::fd::FromRawFd;
|
||||
#[cfg(unix)]
|
||||
use std::os::fd::RawFd;
|
||||
#[cfg(unix)]
|
||||
use std::os::unix::process::CommandExt;
|
||||
use std::path::Path;
|
||||
#[cfg(unix)]
|
||||
use std::process::Command as StdCommand;
|
||||
#[cfg(unix)]
|
||||
use std::process::Stdio;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex as StdMutex;
|
||||
@@ -60,6 +72,18 @@ impl ChildTerminator for PtyChildTerminator {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
struct RawPidTerminator {
|
||||
process_group_id: u32,
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
impl ChildTerminator for RawPidTerminator {
|
||||
fn kill(&mut self) -> std::io::Result<()> {
|
||||
crate::process_group::kill_process_group(self.process_group_id)
|
||||
}
|
||||
}
|
||||
|
||||
fn platform_native_pty_system() -> Box<dyn portable_pty::PtySystem + Send> {
|
||||
#[cfg(windows)]
|
||||
{
|
||||
@@ -79,11 +103,39 @@ pub async fn spawn_process(
|
||||
cwd: &Path,
|
||||
env: &HashMap<String, String>,
|
||||
arg0: &Option<String>,
|
||||
) -> Result<SpawnedProcess> {
|
||||
spawn_process_with_inherited_fds(program, args, cwd, env, arg0, &[]).await
|
||||
}
|
||||
|
||||
/// Spawn a process attached to a PTY, preserving any inherited file
|
||||
/// descriptors listed in `inherited_fds` across exec on Unix.
|
||||
pub async fn spawn_process_with_inherited_fds(
|
||||
program: &str,
|
||||
args: &[String],
|
||||
cwd: &Path,
|
||||
env: &HashMap<String, String>,
|
||||
arg0: &Option<String>,
|
||||
inherited_fds: &[i32],
|
||||
) -> Result<SpawnedProcess> {
|
||||
if program.is_empty() {
|
||||
anyhow::bail!("missing program for PTY spawn");
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
if !inherited_fds.is_empty() {
|
||||
return spawn_process_preserving_fds(program, args, cwd, env, arg0, inherited_fds).await;
|
||||
}
|
||||
|
||||
spawn_process_portable(program, args, cwd, env, arg0).await
|
||||
}
|
||||
|
||||
async fn spawn_process_portable(
|
||||
program: &str,
|
||||
args: &[String],
|
||||
cwd: &Path,
|
||||
env: &HashMap<String, String>,
|
||||
arg0: &Option<String>,
|
||||
) -> Result<SpawnedProcess> {
|
||||
let pty_system = platform_native_pty_system();
|
||||
let pair = pty_system.openpty(PtySize {
|
||||
rows: 24,
|
||||
@@ -167,11 +219,11 @@ pub async fn spawn_process(
|
||||
|
||||
let handles = PtyHandles {
|
||||
_slave: if cfg!(windows) {
|
||||
Some(pair.slave)
|
||||
Some(Box::new(pair.slave))
|
||||
} else {
|
||||
None
|
||||
},
|
||||
_master: pair.master,
|
||||
_master: Box::new(pair.master),
|
||||
};
|
||||
|
||||
let (handle, output_rx) = ProcessHandle::new(
|
||||
@@ -198,3 +250,214 @@ pub async fn spawn_process(
|
||||
exit_rx,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
async fn spawn_process_preserving_fds(
|
||||
program: &str,
|
||||
args: &[String],
|
||||
cwd: &Path,
|
||||
env: &HashMap<String, String>,
|
||||
arg0: &Option<String>,
|
||||
inherited_fds: &[RawFd],
|
||||
) -> Result<SpawnedProcess> {
|
||||
let (master, slave) = open_unix_pty()?;
|
||||
let mut command = StdCommand::new(arg0.as_ref().unwrap_or(&program.to_string()));
|
||||
command.current_dir(cwd);
|
||||
command.env_clear();
|
||||
for arg in args {
|
||||
command.arg(arg);
|
||||
}
|
||||
for (key, value) in env {
|
||||
command.env(key, value);
|
||||
}
|
||||
|
||||
let stdin = slave.try_clone()?;
|
||||
let stdout = slave.try_clone()?;
|
||||
let stderr = slave.try_clone()?;
|
||||
let inherited_fds = inherited_fds.to_vec();
|
||||
|
||||
unsafe {
|
||||
command
|
||||
.stdin(Stdio::from(stdin))
|
||||
.stdout(Stdio::from(stdout))
|
||||
.stderr(Stdio::from(stderr))
|
||||
.pre_exec(move || {
|
||||
for signo in &[
|
||||
libc::SIGCHLD,
|
||||
libc::SIGHUP,
|
||||
libc::SIGINT,
|
||||
libc::SIGQUIT,
|
||||
libc::SIGTERM,
|
||||
libc::SIGALRM,
|
||||
] {
|
||||
libc::signal(*signo, libc::SIG_DFL);
|
||||
}
|
||||
|
||||
let empty_set: libc::sigset_t = std::mem::zeroed();
|
||||
libc::sigprocmask(libc::SIG_SETMASK, &empty_set, std::ptr::null_mut());
|
||||
|
||||
if libc::setsid() == -1 {
|
||||
return Err(std::io::Error::last_os_error());
|
||||
}
|
||||
|
||||
#[allow(clippy::cast_lossless)]
|
||||
if libc::ioctl(0, libc::TIOCSCTTY as _, 0) == -1 {
|
||||
return Err(std::io::Error::last_os_error());
|
||||
}
|
||||
|
||||
close_random_fds_except(&inherited_fds);
|
||||
Ok(())
|
||||
});
|
||||
}
|
||||
|
||||
let mut child = command.spawn()?;
|
||||
drop(slave);
|
||||
let process_group_id = child.id();
|
||||
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128);
|
||||
let (output_tx, _) = broadcast::channel::<Vec<u8>>(256);
|
||||
let initial_output_rx = output_tx.subscribe();
|
||||
|
||||
let mut reader = master.try_clone()?;
|
||||
let output_tx_clone = output_tx.clone();
|
||||
let reader_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || {
|
||||
let mut buf = [0u8; 8_192];
|
||||
loop {
|
||||
match std::io::Read::read(&mut reader, &mut buf) {
|
||||
Ok(0) => break,
|
||||
Ok(n) => {
|
||||
let _ = output_tx_clone.send(buf[..n].to_vec());
|
||||
}
|
||||
Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
|
||||
Err(ref e) if e.kind() == ErrorKind::WouldBlock => {
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
continue;
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let writer = Arc::new(tokio::sync::Mutex::new(master.try_clone()?));
|
||||
let writer_handle: JoinHandle<()> = tokio::spawn({
|
||||
let writer = Arc::clone(&writer);
|
||||
async move {
|
||||
while let Some(bytes) = writer_rx.recv().await {
|
||||
let mut guard = writer.lock().await;
|
||||
use std::io::Write;
|
||||
let _ = guard.write_all(&bytes);
|
||||
let _ = guard.flush();
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let (exit_tx, exit_rx) = oneshot::channel::<i32>();
|
||||
let exit_status = Arc::new(AtomicBool::new(false));
|
||||
let wait_exit_status = Arc::clone(&exit_status);
|
||||
let exit_code = Arc::new(StdMutex::new(None));
|
||||
let wait_exit_code = Arc::clone(&exit_code);
|
||||
let wait_handle: JoinHandle<()> = tokio::task::spawn_blocking(move || {
|
||||
let code = match child.wait() {
|
||||
Ok(status) => status.code().unwrap_or(-1),
|
||||
Err(_) => -1,
|
||||
};
|
||||
wait_exit_status.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
if let Ok(mut guard) = wait_exit_code.lock() {
|
||||
*guard = Some(code);
|
||||
}
|
||||
let _ = exit_tx.send(code);
|
||||
});
|
||||
|
||||
let handles = PtyHandles {
|
||||
_slave: None,
|
||||
_master: Box::new(master),
|
||||
};
|
||||
|
||||
let (handle, output_rx) = ProcessHandle::new(
|
||||
writer_tx,
|
||||
output_tx,
|
||||
initial_output_rx,
|
||||
Box::new(RawPidTerminator { process_group_id }),
|
||||
reader_handle,
|
||||
Vec::new(),
|
||||
writer_handle,
|
||||
wait_handle,
|
||||
exit_status,
|
||||
exit_code,
|
||||
Some(handles),
|
||||
);
|
||||
|
||||
Ok(SpawnedProcess {
|
||||
session: handle,
|
||||
output_rx,
|
||||
exit_rx,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn open_unix_pty() -> Result<(File, File)> {
|
||||
let mut master: RawFd = -1;
|
||||
let mut slave: RawFd = -1;
|
||||
let mut size = libc::winsize {
|
||||
ws_row: 24,
|
||||
ws_col: 80,
|
||||
ws_xpixel: 0,
|
||||
ws_ypixel: 0,
|
||||
};
|
||||
|
||||
let result = unsafe {
|
||||
libc::openpty(
|
||||
&mut master,
|
||||
&mut slave,
|
||||
std::ptr::null_mut(),
|
||||
std::ptr::null_mut(),
|
||||
&mut size,
|
||||
)
|
||||
};
|
||||
if result != 0 {
|
||||
anyhow::bail!("failed to openpty: {:?}", std::io::Error::last_os_error());
|
||||
}
|
||||
|
||||
set_cloexec(master)?;
|
||||
set_cloexec(slave)?;
|
||||
|
||||
Ok(unsafe { (File::from_raw_fd(master), File::from_raw_fd(slave)) })
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn set_cloexec(fd: RawFd) -> std::io::Result<()> {
|
||||
let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) };
|
||||
if flags == -1 {
|
||||
return Err(std::io::Error::last_os_error());
|
||||
}
|
||||
let result = unsafe { libc::fcntl(fd, libc::F_SETFD, flags | libc::FD_CLOEXEC) };
|
||||
if result == -1 {
|
||||
return Err(std::io::Error::last_os_error());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn close_random_fds_except(preserved_fds: &[RawFd]) {
|
||||
if let Ok(dir) = std::fs::read_dir("/dev/fd") {
|
||||
let mut fds = Vec::new();
|
||||
for entry in dir {
|
||||
let num = entry
|
||||
.ok()
|
||||
.map(|entry| entry.file_name())
|
||||
.and_then(|name| name.into_string().ok())
|
||||
.and_then(|name| name.parse::<RawFd>().ok());
|
||||
if let Some(num) = num {
|
||||
if num <= 2 || preserved_fds.contains(&num) {
|
||||
continue;
|
||||
}
|
||||
fds.push(num);
|
||||
}
|
||||
}
|
||||
for fd in fds {
|
||||
unsafe {
|
||||
libc::close(fd);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,6 +3,8 @@ use std::path::Path;
|
||||
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
#[cfg(unix)]
|
||||
use crate::pty::spawn_process_with_inherited_fds;
|
||||
use crate::spawn_pipe_process;
|
||||
use crate::spawn_pty_process;
|
||||
|
||||
@@ -183,15 +185,13 @@ async fn wait_for_marker_pid(
|
||||
collected.extend_from_slice(&chunk);
|
||||
|
||||
let text = String::from_utf8_lossy(&collected);
|
||||
if let Some(marker_idx) = text.find(marker) {
|
||||
let suffix = &text[marker_idx + marker.len()..];
|
||||
let digits: String = suffix
|
||||
.chars()
|
||||
.skip_while(|ch| !ch.is_ascii_digit())
|
||||
.take_while(char::is_ascii_digit)
|
||||
.collect();
|
||||
if !digits.is_empty() {
|
||||
return Ok(digits.parse()?);
|
||||
for line in text.lines() {
|
||||
let line = line.trim_end_matches('\r').trim();
|
||||
let Some(suffix) = line.strip_prefix(marker) else {
|
||||
continue;
|
||||
};
|
||||
if !suffix.is_empty() && suffix.chars().all(|ch| ch.is_ascii_digit()) {
|
||||
return Ok(suffix.parse()?);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -444,3 +444,119 @@ async fn pty_terminate_kills_background_children_in_same_process_group() -> anyh
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn pty_spawn_can_preserve_inherited_fds() -> anyhow::Result<()> {
|
||||
use std::io::Read;
|
||||
use std::os::fd::AsRawFd;
|
||||
use std::os::fd::FromRawFd;
|
||||
|
||||
let mut fds = [0; 2];
|
||||
let result = unsafe { libc::pipe(fds.as_mut_ptr()) };
|
||||
if result != 0 {
|
||||
return Err(std::io::Error::last_os_error().into());
|
||||
}
|
||||
|
||||
let mut read_end = unsafe { std::fs::File::from_raw_fd(fds[0]) };
|
||||
let write_end = unsafe { std::fs::File::from_raw_fd(fds[1]) };
|
||||
|
||||
let mut env_map: HashMap<String, String> = std::env::vars().collect();
|
||||
env_map.insert(
|
||||
"PRESERVED_FD".to_string(),
|
||||
write_end.as_raw_fd().to_string(),
|
||||
);
|
||||
|
||||
let script = "printf __preserved__ >&$PRESERVED_FD";
|
||||
let spawned = spawn_process_with_inherited_fds(
|
||||
"/bin/sh",
|
||||
&["-c".to_string(), script.to_string()],
|
||||
Path::new("."),
|
||||
&env_map,
|
||||
&None,
|
||||
&[write_end.as_raw_fd()],
|
||||
)
|
||||
.await?;
|
||||
|
||||
drop(write_end);
|
||||
|
||||
let (_, code) = collect_output_until_exit(spawned.output_rx, spawned.exit_rx, 2_000).await;
|
||||
assert_eq!(code, 0, "expected preserved-fd PTY child to exit cleanly");
|
||||
|
||||
let mut pipe_output = String::new();
|
||||
read_end.read_to_string(&mut pipe_output)?;
|
||||
assert_eq!(pipe_output, "__preserved__");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn pty_preserving_inherited_fds_keeps_python_repl_running() -> anyhow::Result<()> {
|
||||
use std::os::fd::AsRawFd;
|
||||
use std::os::fd::FromRawFd;
|
||||
|
||||
let Some(python) = find_python() else {
|
||||
eprintln!(
|
||||
"python not found; skipping pty_preserving_inherited_fds_keeps_python_repl_running"
|
||||
);
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let mut fds = [0; 2];
|
||||
let result = unsafe { libc::pipe(fds.as_mut_ptr()) };
|
||||
if result != 0 {
|
||||
return Err(std::io::Error::last_os_error().into());
|
||||
}
|
||||
|
||||
let read_end = unsafe { std::fs::File::from_raw_fd(fds[0]) };
|
||||
let preserved_fd = unsafe { std::fs::File::from_raw_fd(fds[1]) };
|
||||
|
||||
let mut env_map: HashMap<String, String> = std::env::vars().collect();
|
||||
env_map.insert(
|
||||
"PRESERVED_FD".to_string(),
|
||||
preserved_fd.as_raw_fd().to_string(),
|
||||
);
|
||||
|
||||
let spawned = spawn_process_with_inherited_fds(
|
||||
&python,
|
||||
&[],
|
||||
Path::new("."),
|
||||
&env_map,
|
||||
&None,
|
||||
&[preserved_fd.as_raw_fd()],
|
||||
)
|
||||
.await?;
|
||||
drop(read_end);
|
||||
drop(preserved_fd);
|
||||
|
||||
let writer = spawned.session.writer_sender();
|
||||
let mut output_rx = spawned.output_rx;
|
||||
let newline = "\n";
|
||||
let mut output = wait_for_python_repl_ready(&writer, &mut output_rx, 5_000, newline).await?;
|
||||
let marker = "__codex_preserved_py_pid:";
|
||||
writer
|
||||
.send(format!("import os; print('{marker}' + str(os.getpid())){newline}").into_bytes())
|
||||
.await?;
|
||||
|
||||
let python_pid = match wait_for_marker_pid(&mut output_rx, marker, 2_000).await {
|
||||
Ok(pid) => pid,
|
||||
Err(err) => {
|
||||
spawned.session.terminate();
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
assert!(
|
||||
process_exists(python_pid)?,
|
||||
"expected python pid {python_pid} to stay alive after prompt output"
|
||||
);
|
||||
|
||||
writer.send(format!("exit(){newline}").into_bytes()).await?;
|
||||
let (remaining_output, code) =
|
||||
collect_output_until_exit(output_rx, spawned.exit_rx, 5_000).await;
|
||||
output.extend_from_slice(&remaining_output);
|
||||
|
||||
assert_eq!(code, 0, "expected python to exit cleanly");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user