Add stdio exec-server client transport (#20664)

## Why Configured environments need to connect to exec-server instances that are not necessarily already listening on a websocket URL. A command-backed stdio transport lets Codex start an exec-server process, speak JSON-RPC over its stdio streams, and clean up that child process with the client lifetime. **Stack position:** this is PR 2 of 5. It builds on the server-side stdio listener from PR 1 and provides the client transport used by later environment/config PRs. ## What Changed - Add `ExecServerTransport` variants for websocket URLs and stdio shell commands. - Add stdio command connection support for `ExecServerClient`. - Move websocket/stdio transport setup into `client_transport.rs` so `client.rs` stays focused on shared JSON-RPC client, session, HTTP, and notification behavior. - Tie stdio child process cleanup to the JSON-RPC connection lifetime with a RAII lifetime guard. - Keep existing websocket environment behavior by adapting URL-backed remotes to `ExecServerTransport::WebSocketUrl`. ## Stack - 1. https://github.com/openai/codex/pull/20663 - Add stdio exec-server listener - **2. This PR:** https://github.com/openai/codex/pull/20664 - Add stdio exec-server client transport - 3. https://github.com/openai/codex/pull/20665 - Make environment providers own default selection - 4. https://github.com/openai/codex/pull/20666 - Add CODEX_HOME environments TOML provider - 5. https://github.com/openai/codex/pull/20667 - Load configured environments from CODEX_HOME Split from original draft: https://github.com/openai/codex/pull/20508 ## Validation Not run locally; this was split out of the original draft stack and then refactored to separate transport setup from the base client. --------- Co-authored-by: Codex <noreply@openai.com>
2026-07-01 00:31:56 +08:00 · 2026-05-07 23:48:50 +00:00
parent 79154e6952
commit a3de5bde6e
8 changed files with 692 additions and 79 deletions
@@ -1,12 +1,21 @@
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+use std::time::Duration;
+
 use codex_app_server_protocol::JSONRPCMessage;
 use futures::SinkExt;
 use futures::StreamExt;
 use tokio::io::AsyncRead;
 use tokio::io::AsyncWrite;
+use tokio::process::Child;
 use tokio::sync::mpsc;
 use tokio::sync::watch;
+use tokio::time::timeout;
 use tokio_tungstenite::WebSocketStream;
 use tokio_tungstenite::tungstenite::Message;
+use tracing::debug;
+use tracing::warn;

 use tokio::io::AsyncBufReadExt;
 use tokio::io::AsyncWriteExt;
@@ -14,6 +23,7 @@ use tokio::io::BufReader;
 use tokio::io::BufWriter;

 pub(crate) const CHANNEL_CAPACITY: usize = 128;
+const STDIO_TERMINATION_GRACE_PERIOD: Duration = Duration::from_secs(2);

 #[derive(Debug)]
 pub(crate) enum JsonRpcConnectionEvent {
@@ -22,11 +32,186 @@ pub(crate) enum JsonRpcConnectionEvent {
    Disconnected { reason: Option<String> },
 }

+#[derive(Clone)]
+pub(crate) enum JsonRpcTransport {
+    Plain,
+    Stdio { transport: StdioTransport },
+}
+
+impl JsonRpcTransport {
+    fn from_child_process(child_process: Child) -> Self {
+        Self::Stdio {
+            transport: StdioTransport::spawn(child_process),
+        }
+    }
+
+    pub(crate) fn terminate(&self) {
+        match self {
+            Self::Plain => {}
+            Self::Stdio { transport } => transport.terminate(),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub(crate) struct StdioTransport {
+    handle: Arc<StdioTransportHandle>,
+}
+
+struct StdioTransportHandle {
+    terminate_tx: watch::Sender<bool>,
+    terminate_requested: AtomicBool,
+}
+
+impl StdioTransport {
+    fn spawn(child_process: Child) -> Self {
+        let (terminate_tx, terminate_rx) = watch::channel(false);
+        let handle = Arc::new(StdioTransportHandle {
+            terminate_tx,
+            terminate_requested: AtomicBool::new(false),
+        });
+        spawn_stdio_child_supervisor(child_process, terminate_rx);
+        Self { handle }
+    }
+
+    fn terminate(&self) {
+        self.handle.terminate();
+    }
+}
+
+impl StdioTransportHandle {
+    fn terminate(&self) {
+        if !self.terminate_requested.swap(true, Ordering::AcqRel) {
+            let _ = self.terminate_tx.send(true);
+        }
+    }
+}
+
+impl Drop for StdioTransportHandle {
+    fn drop(&mut self) {
+        self.terminate();
+    }
+}
+
+fn spawn_stdio_child_supervisor(mut child_process: Child, mut terminate_rx: watch::Receiver<bool>) {
+    let process_group_id = child_process.id();
+    tokio::spawn(async move {
+        tokio::select! {
+            result = child_process.wait() => {
+                log_stdio_child_wait_result(result);
+                kill_process_tree(&mut child_process, process_group_id);
+            }
+            () = wait_for_stdio_termination(&mut terminate_rx) => {
+                terminate_stdio_child(&mut child_process, process_group_id).await;
+            }
+        }
+    });
+}
+
+async fn wait_for_stdio_termination(terminate_rx: &mut watch::Receiver<bool>) {
+    loop {
+        if *terminate_rx.borrow() {
+            return;
+        }
+        if terminate_rx.changed().await.is_err() {
+            return;
+        }
+    }
+}
+
+async fn terminate_stdio_child(child_process: &mut Child, process_group_id: Option<u32>) {
+    terminate_process_tree(child_process, process_group_id);
+    match timeout(STDIO_TERMINATION_GRACE_PERIOD, child_process.wait()).await {
+        Ok(result) => {
+            log_stdio_child_wait_result(result);
+        }
+        Err(_) => {
+            kill_process_tree(child_process, process_group_id);
+            log_stdio_child_wait_result(child_process.wait().await);
+        }
+    }
+}
+
+fn terminate_process_tree(child_process: &mut Child, process_group_id: Option<u32>) {
+    let Some(process_group_id) = process_group_id else {
+        kill_direct_child(child_process, "terminate");
+        return;
+    };
+
+    #[cfg(unix)]
+    if let Err(err) = codex_utils_pty::process_group::terminate_process_group(process_group_id) {
+        warn!("failed to terminate exec-server stdio process group {process_group_id}: {err}");
+        kill_direct_child(child_process, "terminate");
+    }
+
+    #[cfg(windows)]
+    if !kill_windows_process_tree(process_group_id) {
+        kill_direct_child(child_process, "terminate");
+    }
+
+    #[cfg(not(any(unix, windows)))]
+    {
+        let _ = process_group_id;
+        kill_direct_child(child_process, "terminate");
+    }
+}
+
+fn kill_process_tree(child_process: &mut Child, process_group_id: Option<u32>) {
+    let Some(process_group_id) = process_group_id else {
+        kill_direct_child(child_process, "kill");
+        return;
+    };
+
+    #[cfg(unix)]
+    if let Err(err) = codex_utils_pty::process_group::kill_process_group(process_group_id) {
+        warn!("failed to kill exec-server stdio process group {process_group_id}: {err}");
+    }
+
+    #[cfg(windows)]
+    if !kill_windows_process_tree(process_group_id) {
+        kill_direct_child(child_process, "kill");
+    }
+
+    #[cfg(not(any(unix, windows)))]
+    {
+        let _ = process_group_id;
+        kill_direct_child(child_process, "kill");
+    }
+}
+
+fn kill_direct_child(child_process: &mut Child, action: &str) {
+    if let Err(err) = child_process.start_kill() {
+        debug!("failed to {action} exec-server stdio child: {err}");
+    }
+}
+
+#[cfg(windows)]
+fn kill_windows_process_tree(pid: u32) -> bool {
+    let pid = pid.to_string();
+    match std::process::Command::new("taskkill")
+        .args(["/PID", pid.as_str(), "/T", "/F"])
+        .status()
+    {
+        Ok(status) => status.success(),
+        Err(err) => {
+            warn!("failed to run taskkill for exec-server stdio process tree {pid}: {err}");
+            false
+        }
+    }
+}
+
+fn log_stdio_child_wait_result(result: std::io::Result<std::process::ExitStatus>) {
+    if let Err(err) = result {
+        debug!("failed to wait for exec-server stdio child: {err}");
+    }
+}
+
 pub(crate) struct JsonRpcConnection {
-    outgoing_tx: mpsc::Sender<JSONRPCMessage>,
-    incoming_rx: mpsc::Receiver<JsonRpcConnectionEvent>,
-    disconnected_rx: watch::Receiver<bool>,
-    task_handles: Vec<tokio::task::JoinHandle<()>>,
+    pub(crate) outgoing_tx: mpsc::Sender<JSONRPCMessage>,
+    pub(crate) incoming_rx: mpsc::Receiver<JsonRpcConnectionEvent>,
+    pub(crate) disconnected_rx: watch::Receiver<bool>,
+    pub(crate) task_handles: Vec<tokio::task::JoinHandle<()>>,
+    pub(crate) transport: JsonRpcTransport,
 }

 impl JsonRpcConnection {
@@ -117,6 +302,7 @@ impl JsonRpcConnection {
            incoming_rx,
            disconnected_rx,
            task_handles: vec![reader_task, writer_task],
+            transport: JsonRpcTransport::Plain,
        }
    }

@@ -251,23 +437,13 @@ impl JsonRpcConnection {
            incoming_rx,
            disconnected_rx,
            task_handles: vec![reader_task, writer_task],
+            transport: JsonRpcTransport::Plain,
        }
    }

-    pub(crate) fn into_parts(
-        self,
-    ) -> (
-        mpsc::Sender<JSONRPCMessage>,
-        mpsc::Receiver<JsonRpcConnectionEvent>,
-        watch::Receiver<bool>,
-        Vec<tokio::task::JoinHandle<()>>,
-    ) {
-        (
-            self.outgoing_tx,
-            self.incoming_rx,
-            self.disconnected_rx,
-            self.task_handles,
-        )
+    pub(crate) fn with_child_process(mut self, child_process: Child) -> Self {
+        self.transport = JsonRpcTransport::from_child_process(child_process);
+        self
    }
 }