Improve Windows process management edge cases (#19211)

## Summary

Some improvements to Windows process-management issues from
#15578

- bound the elevated runner pipe-connect handshake instead of waiting
forever on blocking pipe connects
- terminate the spawned runner if that handshake fails, so timeout/error
paths do not leave a stray `codex-command-runner.exe`
- loop on partial `WriteFile` results when forwarding stdin in the
elevated runner, so input is not silently truncated
- fix the concrete HANDLE/SID cleanup paths in the runner setup code
- keep draining driver-backed stdout/stderr after exit until the backend
closes, instead of dropping the tail after a fixed 200ms grace period
- reuse `LocalSid` for SID ownership and add more explanatory comments
around the ownership/concurrency-sensitive code paths

## Why

The original PR fixed a lot of Windows session plumbing, but there were
still a few sharp process-lifecycle edges:

- some elevated runner handshakes could block forever
- the new timeout path could still orphan the spawned runner process
- stdin forwarding still assumed a single `WriteFile` consumed the whole
buffer
- a few raw HANDLE/SID error paths still leaked
- driver-backed output could still lose the last chunk of stdout/stderr
on slower backends

## Validation

- `cargo fmt -p codex-windows-sandbox -p codex-utils-pty`
- `cargo test -p codex-utils-pty`
- `cargo test -p codex-windows-sandbox finish_driver_spawn`
- `cargo test -p codex-windows-sandbox runner_`

Ran a local test matrix of unified-exec and shell_tool tests, all
passing

Author: iceweasel-oai

codex-rs/utils/pty/src/process.rs

@@ -330,22 +330,20 @@ pub fn spawn_from_driver(driver: ProcessDriver) -> SpawnedProcess {
          output_tx: mpsc::Sender<Vec<u8>>,
          mut exit_seen_rx: watch::Receiver<bool>| {
             tokio::spawn(async move {
-                let mut process_exited = false;
                 loop {
-                    let recv_result = if process_exited {
-                        match tokio::time::timeout(
-                            std::time::Duration::from_millis(200),
-                            output_rx.recv(),
-                        )
-                        .await
-                        {
-                            Ok(result) => result,
-                            Err(_) => break,
-                        }
+                    let recv_result = if *exit_seen_rx.borrow() {
+                        // Once exit has been observed, we no longer want a timer here. Some
+                        // backends publish the exit code before their final stdout/stderr bytes
+                        // have been forwarded through the broadcast channel, so a fixed grace
+                        // period can still drop the tail of the stream under load.
+                        //
+                        // Instead, keep waiting until the driver closes the broadcast sender.
+                        // That makes the shutdown contract explicit: the backend is responsible
+                        // for dropping its sender when it has truly finished forwarding output.
+                        output_rx.recv().await
                     } else {
                         tokio::select! {
                             _ = exit_seen_rx.changed() => {
-                                process_exited = *exit_seen_rx.borrow();
                                 continue;
                             }
                             result = output_rx.recv() => result,

codex-rs/utils/pty/src/tests.rs

@@ -688,6 +688,51 @@ async fn driver_backed_process_can_resize_via_resizer_hook() -> anyhow::Result<(
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn driver_backed_process_drains_output_that_arrives_after_exit_signal() -> anyhow::Result<()>
+{
+    let (writer_tx, _writer_rx) = tokio::sync::mpsc::channel::<Vec<u8>>(1);
+    let (stdout_tx, stdout_driver_rx) = tokio::sync::broadcast::channel::<Vec<u8>>(8);
+    let (exit_tx, exit_rx) = tokio::sync::oneshot::channel::<i32>();
+
+    let spawned = spawn_from_driver(ProcessDriver {
+        writer_tx,
+        stdout_rx: stdout_driver_rx,
+        stderr_rx: None,
+        exit_rx,
+        terminator: None,
+        writer_handle: None,
+        resizer: None,
+    });
+
+    let SpawnedProcess {
+        session: _session,
+        stdout_rx,
+        stderr_rx: _stderr_rx,
+        exit_rx,
+    } = spawned;
+    let stdout_task = tokio::spawn(async move { collect_split_output(stdout_rx).await });
+
+    exit_tx.send(0).expect("send exit code");
+    tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
+    stdout_tx.send(b"tail".to_vec())?;
+    drop(stdout_tx);
+
+    let timeout = tokio::time::Duration::from_secs(2);
+    let code = tokio::time::timeout(timeout, exit_rx)
+        .await
+        .map_err(|_| anyhow::anyhow!("timed out waiting for driver exit"))?
+        .unwrap_or(-1);
+    let stdout = tokio::time::timeout(timeout, stdout_task)
+        .await
+        .map_err(|_| anyhow::anyhow!("timed out waiting to drain driver stdout"))??;
+
+    assert_eq!(stdout, b"tail".to_vec());
+    assert_eq!(code, 0);
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn pipe_terminate_aborts_detached_readers() -> anyhow::Result<()> {
     if !setsid_available() {

codex-rs/windows-sandbox-rs/src/elevated/command_runner_win.rs

@@ -16,6 +16,7 @@ use codex_windows_sandbox::ErrorPayload;
 use codex_windows_sandbox::ExitPayload;
 use codex_windows_sandbox::FramedMessage;
 use codex_windows_sandbox::LaunchDesktop;
+use codex_windows_sandbox::LocalSid;
 use codex_windows_sandbox::Message;
 use codex_windows_sandbox::OutputPayload;
 use codex_windows_sandbox::OutputStream;
@@ -27,7 +28,6 @@ use codex_windows_sandbox::SpawnRequest;
 use codex_windows_sandbox::StderrMode;
 use codex_windows_sandbox::StdinMode;
 use codex_windows_sandbox::allow_null_device;
-use codex_windows_sandbox::convert_string_sid_to_sid;
 use codex_windows_sandbox::create_readonly_token_with_caps_from;
 use codex_windows_sandbox::create_workspace_write_token_with_caps_from;
 use codex_windows_sandbox::decode_bytes;
@@ -41,7 +41,6 @@ use codex_windows_sandbox::read_handle_loop;
 use codex_windows_sandbox::spawn_process_with_pipes;
 use codex_windows_sandbox::to_wide;
 use codex_windows_sandbox::write_frame;
-use std::ffi::c_void;
 use std::fs::File;
 use std::os::windows::io::FromRawHandle;
 use std::path::Path;
@@ -52,8 +51,7 @@ use std::sync::Mutex as StdMutex;
 use windows_sys::Win32::Foundation::CloseHandle;
 use windows_sys::Win32::Foundation::GetLastError;
 use windows_sys::Win32::Foundation::HANDLE;
-use windows_sys::Win32::Foundation::HLOCAL;
-use windows_sys::Win32::Foundation::LocalFree;
+use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE;
 use windows_sys::Win32::Storage::FileSystem::CreateFileW;
 use windows_sys::Win32::Storage::FileSystem::FILE_GENERIC_READ;
 use windows_sys::Win32::Storage::FileSystem::FILE_GENERIC_WRITE;
@@ -94,23 +92,58 @@ struct IpcSpawnedProcess {
     _pipe_handles: Option<PipeSpawnHandles>,
 }
 
+/// Small RAII wrapper for raw Win32 handles.
+///
+/// The elevated runner has a few early-return paths where we acquire a token, job, or pipe
+/// handle and then may fail while preparing the child. Keeping those handles in a guard makes
+/// the error paths read more directly and closes the gaps that were previously leaking them.
+struct OwnedWinHandle(HANDLE);
+
+impl OwnedWinHandle {
+    fn new(handle: HANDLE) -> Self {
+        Self(handle)
+    }
+
+    fn raw(&self) -> HANDLE {
+        self.0
+    }
+
+    fn into_raw(mut self) -> HANDLE {
+        // Transfer ownership to the caller. After this point the caller is responsible for
+        // eventually closing the returned HANDLE.
+        let handle = self.0;
+        self.0 = 0;
+        handle
+    }
+}
+
+impl Drop for OwnedWinHandle {
+    fn drop(&mut self) {
+        if self.0 != 0 && self.0 != INVALID_HANDLE_VALUE {
+            unsafe {
+                CloseHandle(self.0);
+            }
+        }
+    }
+}
+
 unsafe fn create_job_kill_on_close() -> Result<HANDLE> {
-    let h = CreateJobObjectW(std::ptr::null_mut(), std::ptr::null());
-    if h == 0 {
+    let h_job = OwnedWinHandle::new(CreateJobObjectW(std::ptr::null_mut(), std::ptr::null()));
+    if h_job.raw() == 0 {
         return Err(anyhow::anyhow!("CreateJobObjectW failed"));
     }
     let mut limits: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = std::mem::zeroed();
     limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
     let ok = SetInformationJobObject(
-        h,
+        h_job.raw(),
         JobObjectExtendedLimitInformation,
         &mut limits as *mut _ as *mut _,
         std::mem::size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
     );
     if ok == 0 {
         return Err(anyhow::anyhow!("SetInformationJobObject failed"));
     }
-    Ok(h)
+    Ok(h_job.into_raw())
 }
 
 /// Open a named pipe created by the parent process.
@@ -190,45 +223,42 @@ fn spawn_ipc_process(req: &SpawnRequest) -> Result<IpcSpawnedProcess> {
     let log_dir = req.codex_home.clone();
     hide_current_user_profile_dir(req.codex_home.as_path());
     let policy = parse_policy(&req.policy_json_or_preset).context("parse policy_json_or_preset")?;
-    let mut cap_psids: Vec<*mut c_void> = Vec::new();
+    let mut cap_psids: Vec<LocalSid> = Vec::new();
     for sid in &req.cap_sids {
-        let Some(psid) = (unsafe { convert_string_sid_to_sid(sid) }) else {
-            anyhow::bail!("ConvertStringSidToSidW failed for capability SID");
-        };
-        cap_psids.push(psid);
+        cap_psids.push(
+            LocalSid::from_string(sid)
+                .context("ConvertStringSidToSidW failed for capability SID")?,
+        );
     }
     if cap_psids.is_empty() {
         anyhow::bail!("runner: empty capability SID list");
     }
 
-    let base = unsafe { get_current_token_for_restriction()? };
-    let token_res: Result<(HANDLE, *mut c_void)> = unsafe {
+    // The token helpers still take raw SID pointers, but we keep ownership in `LocalSid`
+    // wrappers for as long as possible. That way any failure after SID parsing but before the
+    // child is fully spawned still releases the backing LocalAlloc memory automatically.
+    let cap_psid_ptrs: Vec<*mut _> = cap_psids.iter().map(LocalSid::as_ptr).collect();
+    let base = OwnedWinHandle::new(unsafe { get_current_token_for_restriction()? });
+    let h_token = OwnedWinHandle::new(unsafe {
         match &policy {
             SandboxPolicy::ReadOnly { .. } => {
-                create_readonly_token_with_caps_from(base, &cap_psids)
-                    .map(|h_token| (h_token, cap_psids[0]))
+                create_readonly_token_with_caps_from(base.raw(), &cap_psid_ptrs)
             }
             SandboxPolicy::WorkspaceWrite { .. } => {
-                create_workspace_write_token_with_caps_from(base, &cap_psids)
-                    .map(|h_token| (h_token, cap_psids[0]))
+                create_workspace_write_token_with_caps_from(base.raw(), &cap_psid_ptrs)
             }
             SandboxPolicy::DangerFullAccess | SandboxPolicy::ExternalSandbox { .. } => {
                 unreachable!()
             }
         }
-    };
-    let (h_token, psid_to_use) = token_res?;
+    }?);
     unsafe {
-        CloseHandle(base);
-        allow_null_device(psid_to_use);
-        for psid in &cap_psids {
+        // These ACL adjustments need the raw SID values, but ownership stays with `cap_psids`.
+        // We do not manually `LocalFree` anything here; the wrappers handle every return path.
+        allow_null_device(cap_psid_ptrs[0]);
+        for psid in &cap_psid_ptrs {
             allow_null_device(*psid);
         }
-        for psid in cap_psids {
-            if !psid.is_null() {
-                LocalFree(psid as HLOCAL);
-            }
-        }
     }
 
     let effective_cwd = effective_cwd(&req.cwd, Some(log_dir.as_path()));
@@ -238,7 +268,7 @@ fn spawn_ipc_process(req: &SpawnRequest) -> Result<IpcSpawnedProcess> {
     let mut pipe_handles = None;
     let (pi, stdout_handle, stderr_handle, stdin_handle) = if req.tty {
         let (pi, conpty) = codex_windows_sandbox::spawn_conpty_process_as_user(
-            h_token,
+            h_token.raw(),
             &req.command,
             &effective_cwd,
             &req.env,
@@ -269,7 +299,7 @@ fn spawn_ipc_process(req: &SpawnRequest) -> Result<IpcSpawnedProcess> {
             StdinMode::Closed
         };
         let spawned_pipes: PipeSpawnHandles = spawn_process_with_pipes(
-            h_token,
+            h_token.raw(),
             &req.command,
             &effective_cwd,
             &req.env,
@@ -287,10 +317,6 @@ fn spawn_ipc_process(req: &SpawnRequest) -> Result<IpcSpawnedProcess> {
         pipe_handles = Some(spawned_pipes);
         (pi, stdout_handle, stderr_handle, stdin_handle)
     };
-
-    unsafe {
-        CloseHandle(h_token);
-    }
     Ok(IpcSpawnedProcess {
         log_dir,
         pi,
@@ -337,7 +363,7 @@ fn spawn_input_loop(
     stdin_handle: Option<HANDLE>,
     hpc_handle: Arc<StdMutex<Option<HANDLE>>>,
     process_handle: Arc<StdMutex<Option<HANDLE>>>,
-    _log_dir: Option<PathBuf>,
+    log_dir: Option<PathBuf>,
 ) -> std::thread::JoinHandle<()> {
     std::thread::spawn(move || {
         let mut stdin_handle = stdin_handle;
@@ -353,15 +379,54 @@ fn spawn_input_loop(
                         continue;
                     };
                     if let Some(handle) = stdin_handle {
-                        let mut written: u32 = 0;
-                        unsafe {
-                            let _ = windows_sys::Win32::Storage::FileSystem::WriteFile(
-                                handle,
-                                bytes.as_ptr(),
-                                bytes.len() as u32,
-                                &mut written,
-                                ptr::null_mut(),
-                            );
+                        let mut offset = 0usize;
+                        // `WriteFile` can report success after consuming only part of the buffer
+                        // when the target is a pipe. Treat this like a normal partial write and
+                        // keep advancing until every decoded stdin byte has been forwarded.
+                        //
+                        // If the child closes stdin or the pipe enters an error state, we log
+                        // that fact, close our local HANDLE, and stop trying to forward later
+                        // `Stdin` frames. That prevents silent truncation while also avoiding an
+                        // endless stream of failing writes after the child is already gone.
+                        while offset < bytes.len() {
+                            let chunk = &bytes[offset..];
+                            let chunk_len = chunk.len().min(u32::MAX as usize);
+                            let mut written = 0u32;
+                            let ok = unsafe {
+                                windows_sys::Win32::Storage::FileSystem::WriteFile(
+                                    handle,
+                                    chunk.as_ptr(),
+                                    chunk_len as u32,
+                                    &mut written,
+                                    ptr::null_mut(),
+                                )
+                            };
+                            if ok == 0 {
+                                log_note(
+                                    &format!(
+                                        "runner stdin write failed after {offset} bytes: {}",
+                                        unsafe { GetLastError() }
+                                    ),
+                                    log_dir.as_deref(),
+                                );
+                                unsafe {
+                                    CloseHandle(handle);
+                                }
+                                stdin_handle = None;
+                                break;
+                            }
+                            if written == 0 {
+                                log_note(
+                                    "runner stdin write made no progress; closing child stdin",
+                                    log_dir.as_deref(),
+                                );
+                                unsafe {
+                                    CloseHandle(handle);
+                                }
+                                stdin_handle = None;
+                                break;
+                            }
+                            offset += written as usize;
                         }
                     }
                 }
@@ -432,11 +497,14 @@ pub fn main() -> Result<()> {
         anyhow::bail!("runner: no pipe-out provided");
     };
 
-    let h_pipe_in = open_pipe(&pipe_in, FILE_GENERIC_READ)?;
-    let h_pipe_out = open_pipe(&pipe_out, FILE_GENERIC_WRITE)?;
-    let mut pipe_read = unsafe { File::from_raw_handle(h_pipe_in as _) };
+    // Open both pipe ends under guards first so a failure on the second open cannot leak the
+    // first HANDLE. Only after both opens succeed do we transfer ownership into `File`, which
+    // then becomes responsible for closing them.
+    let h_pipe_in = OwnedWinHandle::new(open_pipe(&pipe_in, FILE_GENERIC_READ)?);
+    let h_pipe_out = OwnedWinHandle::new(open_pipe(&pipe_out, FILE_GENERIC_WRITE)?);
+    let mut pipe_read = unsafe { File::from_raw_handle(h_pipe_in.into_raw() as _) };
     let pipe_write = Arc::new(StdMutex::new(unsafe {
-        File::from_raw_handle(h_pipe_out as _)
+        File::from_raw_handle(h_pipe_out.into_raw() as _)
     }));
 
     let req = match read_spawn_request(&mut pipe_read) {