codex: address PR review feedback (#26206)

codex-rs/Cargo.lock

@@ -3460,6 +3460,7 @@ dependencies = [
 name = "codex-process"
 version = "0.0.0"
 dependencies = [
+ "either",
  "tokio",
  "tracing",
  "tracing-test",

codex-rs/Cargo.toml

@@ -289,6 +289,7 @@ dns-lookup = "3.0.1"
 dotenvy = "0.15.7"
 dunce = "1.0.4"
 ed25519-dalek = { version = "2.2.0", features = ["pkcs8"] }
+either = "1.15.0"
 encoding_rs = "0.8.35"
 env_logger = "0.11.9"
 eventsource-stream = "0.2.3"

codex-rs/process/Cargo.toml

@@ -13,6 +13,7 @@ doctest = false
 workspace = true
 
 [dependencies]
+either = { workspace = true }
 tokio = { workspace = true, features = ["process"] }
 tracing = { workspace = true }
 

codex-rs/process/src/command_ext.rs

@@ -0,0 +1,13 @@
+use std::io;
+
+/// Extends process command types with Codex-specific process spawning.
+///
+/// Callers should use this trait to guarantee that child processes are joined.
+/// Implementations return the corresponding managed child handle.
+pub trait CommandExt {
+    /// The managed child handle returned after spawning.
+    type Child;
+
+    /// Spawns this command and returns a child handle that must be joined.
+    fn spawn_managed(&mut self) -> io::Result<Self::Child>;
+}

codex-rs/process/src/drop_bomb.rs

@@ -0,0 +1,39 @@
+#[derive(Debug)]
+pub(crate) struct DropBomb {
+    armed: bool,
+}
+
+impl DropBomb {
+    pub(crate) fn new() -> Self {
+        Self { armed: true }
+    }
+
+    pub(crate) fn disarm(&mut self) {
+        self.armed = false;
+    }
+
+    #[cfg(test)]
+    pub(crate) fn is_armed(&self) -> bool {
+        self.armed
+    }
+}
+
+impl Drop for DropBomb {
+    fn drop(&mut self) {
+        if !self.armed {
+            return;
+        }
+
+        const UNJOINED_CHILD_MESSAGE: &str = "managed child process dropped without being joined";
+
+        if cfg!(debug_assertions) && !std::thread::panicking() {
+            panic!("{UNJOINED_CHILD_MESSAGE}");
+        }
+
+        tracing::error!("{UNJOINED_CHILD_MESSAGE}");
+    }
+}
+
+#[cfg(test)]
+#[path = "drop_bomb_tests.rs"]
+mod tests;

codex-rs/process/src/drop_bomb_tests.rs

@@ -1,5 +1,5 @@
 use super::DropBomb;
-use super::UNJOINED_CHILD_MESSAGE;
+use crate::test_support::UNJOINED_CHILD_MESSAGE;
 use crate::test_support::panic_message;
 use std::panic::AssertUnwindSafe;
 use std::panic::catch_unwind;
@@ -13,11 +13,9 @@ fn disarmed_drop_bomb_does_not_report_an_error() {
 
 #[cfg(debug_assertions)]
 #[test]
+#[should_panic(expected = "managed child process dropped without being joined")]
 fn armed_drop_bomb_panics() {
-    let panic = catch_unwind(AssertUnwindSafe(|| drop(DropBomb::new())))
-        .expect_err("armed drop bomb should panic");
-
-    assert_eq!(panic_message(panic.as_ref()), UNJOINED_CHILD_MESSAGE);
+    drop(DropBomb::new());
 }
 
 #[cfg(not(debug_assertions))]

codex-rs/process/src/lib.rs

@@ -4,48 +4,11 @@
 //! dropped. Debug builds enforce this with a drop bomb, while release builds
 //! log an error.
 
+mod command_ext;
+mod drop_bomb;
+
 pub mod sync;
 pub mod tokio;
 
-const UNJOINED_CHILD_MESSAGE: &str = "managed child process dropped without being joined";
-
-#[derive(Debug)]
-struct DropBomb {
-    armed: bool,
-}
-
-impl DropBomb {
-    fn new() -> Self {
-        Self { armed: true }
-    }
-
-    fn disarm(&mut self) {
-        self.armed = false;
-    }
-
-    #[cfg(test)]
-    fn is_armed(&self) -> bool {
-        self.armed
-    }
-}
-
-impl Drop for DropBomb {
-    fn drop(&mut self) {
-        if !self.armed {
-            return;
-        }
-
-        if cfg!(debug_assertions) && !std::thread::panicking() {
-            panic!("{UNJOINED_CHILD_MESSAGE}");
-        }
-
-        tracing::error!("{UNJOINED_CHILD_MESSAGE}");
-    }
-}
-
-#[cfg(test)]
-#[path = "lib_tests.rs"]
-mod tests;
-
 #[cfg(test)]
 mod test_support;

codex-rs/process/src/sync.rs

@@ -1,6 +1,8 @@
 //! Managed wrappers for [`std::process`].
 
-use crate::DropBomb;
+pub use crate::command_ext::CommandExt;
+use crate::drop_bomb::DropBomb;
+use either::Either;
 use std::io;
 use std::ops::Deref;
 use std::ops::DerefMut;
@@ -9,17 +11,10 @@ use std::process::Command;
 use std::process::ExitStatus;
 use std::process::Output;
 
-/// Extends [`std::process::Command`] with Codex-specific process spawning.
-///
-/// Callers should use [`CommandExt::spawn_managed`] when the returned child
-/// must be joined before its handle is dropped.
-pub trait CommandExt {
-    /// Spawns this command and returns a child handle that must be joined.
-    fn spawn_managed(&mut self) -> io::Result<Child>;
-}
-
 impl CommandExt for Command {
-    fn spawn_managed(&mut self) -> io::Result<Child> {
+    type Child = Child;
+
+    fn spawn_managed(&mut self) -> io::Result<Self::Child> {
         self.spawn().map(Child::new)
     }
 }
@@ -27,6 +22,8 @@ impl CommandExt for Command {
 /// A synchronous child process handle that must be explicitly joined.
 #[derive(Debug)]
 pub struct Child {
+    // This is an Option only so DropBomb still runs after wait_with_output()
+    // moves the native child into its consuming join.
     child: Option<StdChild>,
     bomb: DropBomb,
 }
@@ -40,7 +37,7 @@ impl Child {
     }
 
     /// Waits for the child to exit and disarms the drop bomb on success.
-    pub fn wait(&mut self) -> io::Result<ExitStatus> {
+    pub fn wait(mut self) -> io::Result<ExitStatus> {
         let result = self.child_mut().wait();
         if result.is_ok() {
             self.bomb.disarm();
@@ -50,13 +47,16 @@ impl Child {
 
     /// Returns the child's exit status without blocking.
     ///
-    /// The drop bomb is disarmed only when an exit status is available.
-    pub fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
-        let result = self.child_mut().try_wait();
-        if matches!(result, Ok(Some(_))) {
-            self.bomb.disarm();
+    /// Returns the still-armed child handle when an exit status is not yet
+    /// available.
+    pub fn try_wait(mut self) -> io::Result<Either<ExitStatus, Self>> {
+        match self.child_mut().try_wait()? {
+            Some(status) => {
+                self.bomb.disarm();
+                Ok(Either::Left(status))
+            }
+            None => Ok(Either::Right(self)),
         }
-        result
     }
 
     /// Waits for the child to exit and collects its output.
@@ -72,21 +72,21 @@ impl Child {
     fn child(&self) -> &StdChild {
         match self.child.as_ref() {
             Some(child) => child,
-            None => panic!("managed child is unavailable while joining"),
+            None => panic!("managed child was made None before its wrapper was dropped"),
         }
     }
 
     fn child_mut(&mut self) -> &mut StdChild {
         match self.child.as_mut() {
             Some(child) => child,
-            None => panic!("managed child is unavailable while joining"),
+            None => panic!("managed child was made None before its wrapper was dropped"),
         }
     }
 
     fn take_child(&mut self) -> StdChild {
         match self.child.take() {
             Some(child) => child,
-            None => panic!("managed child is unavailable while joining"),
+            None => panic!("managed child was made None before its wrapper was dropped"),
         }
     }
 }

codex-rs/process/src/sync_tests.rs

@@ -1,33 +1,26 @@
 use super::Child;
 use super::CommandExt as _;
-use crate::UNJOINED_CHILD_MESSAGE;
 use crate::test_support;
 use crate::test_support::STDERR_TEXT;
 use crate::test_support::STDOUT_TEXT;
-#[cfg(debug_assertions)]
-use crate::test_support::panic_message;
+#[cfg(not(debug_assertions))]
+use crate::test_support::UNJOINED_CHILD_MESSAGE;
+use either::Either;
 use std::io::Read;
 use std::ops::DerefMut;
-#[cfg(debug_assertions)]
-use std::panic::AssertUnwindSafe;
-#[cfg(debug_assertions)]
-use std::panic::catch_unwind;
 use std::process::Stdio;
 use std::time::Duration;
 use std::time::Instant;
 
 #[test]
-fn wait_disarms_bomb_and_can_be_repeated() {
-    let mut child = test_support::command("exit-success")
+fn wait_disarms_bomb() {
+    let child = test_support::command("exit-success")
         .spawn_managed()
         .expect("spawn helper");
     assert!(child.bomb.is_armed());
 
     let status = child.wait().expect("wait for helper");
     assert!(status.success());
-    assert!(!child.bomb.is_armed());
-
-    assert_eq!(child.wait().expect("repeat wait for helper"), status);
 }
 
 #[test]
@@ -46,17 +39,19 @@ fn stdio_is_available_through_deref_mut() {
 
 #[test]
 fn try_wait_keeps_bomb_armed_until_status_is_available() {
-    let mut child = test_support::command("sleep")
+    let child = test_support::command("sleep")
         .spawn_managed()
         .expect("spawn helper");
 
-    assert_eq!(child.try_wait().expect("poll sleeping helper"), None);
+    let mut child = match child.try_wait().expect("poll sleeping helper") {
+        Either::Left(status) => panic!("sleeping helper exited unexpectedly: {status}"),
+        Either::Right(child) => child,
+    };
     assert!(child.bomb.is_armed());
 
     child.kill().expect("kill sleeping helper");
     assert!(child.bomb.is_armed());
     assert!(!child.wait().expect("wait for killed helper").success());
-    assert!(!child.bomb.is_armed());
 }
 
 #[test]
@@ -67,11 +62,14 @@ fn try_wait_disarms_bomb_when_status_is_available() {
     let deadline = Instant::now() + Duration::from_secs(5);
 
     loop {
-        if let Some(status) = child.try_wait().expect("poll helper") {
-            assert!(status.success());
-            assert!(!child.bomb.is_armed());
-            return;
-        }
+        child = match child.try_wait().expect("poll helper") {
+            Either::Left(status) => {
+                assert!(status.success());
+                return;
+            }
+            Either::Right(child) => child,
+        };
+        assert!(child.bomb.is_armed());
         assert!(Instant::now() < deadline, "helper did not exit");
         std::thread::sleep(Duration::from_millis(10));
     }
@@ -94,16 +92,14 @@ fn wait_with_output_collects_output() {
 
 #[cfg(debug_assertions)]
 #[test]
+#[should_panic(expected = "managed child process dropped without being joined")]
 fn dropping_unjoined_child_panics() {
     let mut child = test_support::command("sleep")
         .spawn_managed()
         .expect("spawn helper");
     clean_up_without_disarming(&mut child);
 
-    let panic = catch_unwind(AssertUnwindSafe(|| drop(child)))
-        .expect_err("dropping unjoined child should panic");
-
-    assert_eq!(panic_message(panic.as_ref()), UNJOINED_CHILD_MESSAGE);
+    drop(child);
 }
 
 #[cfg(not(debug_assertions))]

codex-rs/process/src/test_support.rs

@@ -6,6 +6,8 @@ const SUBPROCESS_MODE_ENV: &str = "CODEX_PROCESS_TEST_SUBPROCESS_MODE";
 
 pub(crate) const STDOUT_TEXT: &str = "managed stdout";
 pub(crate) const STDERR_TEXT: &str = "managed stderr";
+pub(crate) const UNJOINED_CHILD_MESSAGE: &str =
+    "managed child process dropped without being joined";
 
 pub(crate) fn command(mode: &str) -> Command {
     let mut command = Command::new(std::env::current_exe().expect("current test binary"));

codex-rs/process/src/tokio.rs

@@ -1,25 +1,20 @@
 //! Managed wrappers for [`tokio::process`].
 
-use crate::DropBomb;
+pub use crate::command_ext::CommandExt;
+use crate::drop_bomb::DropBomb;
 use ::tokio::process::Child as TokioChild;
 use ::tokio::process::Command;
+use either::Either;
 use std::io;
 use std::ops::Deref;
 use std::ops::DerefMut;
 use std::process::ExitStatus;
 use std::process::Output;
 
-/// Extends [`tokio::process::Command`] with Codex-specific process spawning.
-///
-/// Callers should use [`CommandExt::spawn_managed`] when the returned child
-/// must be joined before its handle is dropped.
-pub trait CommandExt {
-    /// Spawns this command and returns a child handle that must be joined.
-    fn spawn_managed(&mut self) -> io::Result<Child>;
-}
-
 impl CommandExt for Command {
-    fn spawn_managed(&mut self) -> io::Result<Child> {
+    type Child = Child;
+
+    fn spawn_managed(&mut self) -> io::Result<Self::Child> {
         self.spawn().map(Child::new)
     }
 }
@@ -27,6 +22,8 @@ impl CommandExt for Command {
 /// An asynchronous child process handle that must be explicitly joined.
 #[derive(Debug)]
 pub struct Child {
+    // This is an Option only so DropBomb still runs after wait_with_output()
+    // moves the native child into its consuming join.
     child: Option<TokioChild>,
     bomb: DropBomb,
 }
@@ -40,7 +37,7 @@ impl Child {
     }
 
     /// Waits for the child to exit and disarms the drop bomb on success.
-    pub async fn wait(&mut self) -> io::Result<ExitStatus> {
+    pub async fn wait(mut self) -> io::Result<ExitStatus> {
         let result = self.child_mut().wait().await;
         if result.is_ok() {
             self.bomb.disarm();
@@ -50,13 +47,16 @@ impl Child {
 
     /// Returns the child's exit status without blocking.
     ///
-    /// The drop bomb is disarmed only when an exit status is available.
-    pub fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
-        let result = self.child_mut().try_wait();
-        if matches!(result, Ok(Some(_))) {
-            self.bomb.disarm();
+    /// Returns the still-armed child handle when an exit status is not yet
+    /// available.
+    pub fn try_wait(mut self) -> io::Result<Either<ExitStatus, Self>> {
+        match self.child_mut().try_wait()? {
+            Some(status) => {
+                self.bomb.disarm();
+                Ok(Either::Left(status))
+            }
+            None => Ok(Either::Right(self)),
         }
-        result
     }
 
     /// Waits for the child to exit and collects its output.
@@ -72,21 +72,21 @@ impl Child {
     fn child(&self) -> &TokioChild {
         match self.child.as_ref() {
             Some(child) => child,
-            None => panic!("managed child is unavailable while joining"),
+            None => panic!("managed child was made None before its wrapper was dropped"),
         }
     }
 
     fn child_mut(&mut self) -> &mut TokioChild {
         match self.child.as_mut() {
             Some(child) => child,
-            None => panic!("managed child is unavailable while joining"),
+            None => panic!("managed child was made None before its wrapper was dropped"),
         }
     }
 
     fn take_child(&mut self) -> TokioChild {
         match self.child.take() {
             Some(child) => child,
-            None => panic!("managed child is unavailable while joining"),
+            None => panic!("managed child was made None before its wrapper was dropped"),
         }
     }
 }

codex-rs/process/src/tokio_tests.rs

@@ -1,17 +1,13 @@
 use super::Child;
 use super::CommandExt as _;
-use crate::UNJOINED_CHILD_MESSAGE;
 use crate::test_support;
 use crate::test_support::STDERR_TEXT;
 use crate::test_support::STDOUT_TEXT;
-#[cfg(debug_assertions)]
-use crate::test_support::panic_message;
+#[cfg(not(debug_assertions))]
+use crate::test_support::UNJOINED_CHILD_MESSAGE;
 use ::tokio as tokio_crate;
+use either::Either;
 use std::ops::DerefMut;
-#[cfg(debug_assertions)]
-use std::panic::AssertUnwindSafe;
-#[cfg(debug_assertions)]
-use std::panic::catch_unwind;
 use std::process::Stdio;
 use std::time::Duration;
 use tokio_crate::io::AsyncReadExt;
@@ -21,17 +17,14 @@ use tokio_crate::time::sleep;
 use tokio_crate::time::timeout;
 
 #[tokio_crate::test]
-async fn wait_disarms_bomb_and_can_be_repeated() {
-    let mut child = command("exit-success")
+async fn wait_disarms_bomb() {
+    let child = command("exit-success")
         .spawn_managed()
         .expect("spawn helper");
     assert!(child.bomb.is_armed());
 
     let status = child.wait().await.expect("wait for helper");
     assert!(status.success());
-    assert!(!child.bomb.is_armed());
-
-    assert_eq!(child.wait().await.expect("repeat wait for helper"), status);
 }
 
 #[tokio_crate::test]
@@ -53,9 +46,12 @@ async fn stdio_is_available_through_deref_mut() {
 
 #[tokio_crate::test]
 async fn try_wait_keeps_bomb_armed_until_status_is_available() {
-    let mut child = command("sleep").spawn_managed().expect("spawn helper");
+    let child = command("sleep").spawn_managed().expect("spawn helper");
 
-    assert_eq!(child.try_wait().expect("poll sleeping helper"), None);
+    let mut child = match child.try_wait().expect("poll sleeping helper") {
+        Either::Left(status) => panic!("sleeping helper exited unexpectedly: {status}"),
+        Either::Right(child) => child,
+    };
     assert!(child.bomb.is_armed());
 
     child.start_kill().expect("kill sleeping helper");
@@ -67,7 +63,6 @@ async fn try_wait_keeps_bomb_armed_until_status_is_available() {
             .expect("wait for killed helper")
             .success()
     );
-    assert!(!child.bomb.is_armed());
 }
 
 #[tokio_crate::test]
@@ -78,11 +73,14 @@ async fn try_wait_disarms_bomb_when_status_is_available() {
     let deadline = Instant::now() + Duration::from_secs(5);
 
     loop {
-        if let Some(status) = child.try_wait().expect("poll helper") {
-            assert!(status.success());
-            assert!(!child.bomb.is_armed());
-            return;
-        }
+        child = match child.try_wait().expect("poll helper") {
+            Either::Left(status) => {
+                assert!(status.success());
+                return;
+            }
+            Either::Right(child) => child,
+        };
+        assert!(child.bomb.is_armed());
         assert!(Instant::now() < deadline, "helper did not exit");
         sleep(Duration::from_millis(10)).await;
     }
@@ -118,33 +116,47 @@ async fn kill_keeps_bomb_armed_until_explicit_wait() {
             .expect("wait for killed helper")
             .success()
     );
-    assert!(!child.bomb.is_armed());
-}
-
-#[tokio_crate::test]
-async fn cancelled_wait_keeps_bomb_armed() {
-    let mut child = command("sleep").spawn_managed().expect("spawn helper");
-
-    assert!(
-        timeout(Duration::from_millis(10), child.wait())
-            .await
-            .is_err()
-    );
-    assert!(child.bomb.is_armed());
-
-    child.start_kill().expect("kill sleeping helper");
-    assert!(
-        !child
-            .wait()
-            .await
-            .expect("wait for killed helper")
-            .success()
-    );
-    assert!(!child.bomb.is_armed());
 }
 
 #[cfg(debug_assertions)]
 #[tokio_crate::test]
+#[should_panic(expected = "managed child process dropped without being joined")]
+async fn cancelled_wait_panics_when_dropped() {
+    let mut command = command("sleep");
+    command.kill_on_drop(true);
+    let child = command.spawn_managed().expect("spawn helper");
+    let mut wait = Box::pin(child.wait());
+
+    assert!(
+        timeout(Duration::from_millis(10), wait.as_mut())
+            .await
+            .is_err()
+    );
+    drop(wait);
+}
+
+#[cfg(not(debug_assertions))]
+#[tokio_crate::test]
+#[tracing_test::traced_test]
+async fn cancelled_wait_logs_an_error_when_dropped() {
+    let mut command = command("sleep");
+    command.kill_on_drop(true);
+    let child = command.spawn_managed().expect("spawn helper");
+    let mut wait = Box::pin(child.wait());
+
+    assert!(
+        timeout(Duration::from_millis(10), wait.as_mut())
+            .await
+            .is_err()
+    );
+    drop(wait);
+
+    assert!(logs_contain(UNJOINED_CHILD_MESSAGE));
+}
+
+#[cfg(debug_assertions)]
+#[tokio_crate::test]
+#[should_panic(expected = "managed child process dropped without being joined")]
 async fn cancelled_wait_with_output_panics_when_dropped() {
     let mut command = command("sleep");
     command.kill_on_drop(true);
@@ -156,10 +168,7 @@ async fn cancelled_wait_with_output_panics_when_dropped() {
             .await
             .is_err()
     );
-    let panic = catch_unwind(AssertUnwindSafe(|| drop(wait)))
-        .expect_err("dropping cancelled wait_with_output should panic");
-
-    assert_eq!(panic_message(panic.as_ref()), UNJOINED_CHILD_MESSAGE);
+    drop(wait);
 }
 
 #[cfg(not(debug_assertions))]
@@ -183,14 +192,12 @@ async fn cancelled_wait_with_output_logs_an_error_when_dropped() {
 
 #[cfg(debug_assertions)]
 #[tokio_crate::test]
+#[should_panic(expected = "managed child process dropped without being joined")]
 async fn dropping_unjoined_child_panics() {
     let mut child = command("sleep").spawn_managed().expect("spawn helper");
     clean_up_without_disarming(&mut child).await;
 
-    let panic = catch_unwind(AssertUnwindSafe(|| drop(child)))
-        .expect_err("dropping unjoined child should panic");
-
-    assert_eq!(panic_message(panic.as_ref()), UNJOINED_CHILD_MESSAGE);
+    drop(child);
 }
 
 #[cfg(not(debug_assertions))]