diff --git a/codex-rs/core/src/unified_exec/async_watcher.rs b/codex-rs/core/src/unified_exec/async_watcher.rs index d20a85843b..5ffbfaaafd 100644 --- a/codex-rs/core/src/unified_exec/async_watcher.rs +++ b/codex-rs/core/src/unified_exec/async_watcher.rs @@ -53,7 +53,7 @@ pub(crate) fn start_streaming_output( output_closed, output_closed_notify, .. - } = process.output_handles(); + } = process.output_handles().clone(); let session_ref = Arc::clone(&context.session); let turn_ref = Arc::clone(&context.turn); diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index 1aac39f8f8..fa6ae4e5a0 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -10,7 +10,6 @@ use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnContext; use crate::tools::context::ExecCommandToolOutput; use crate::unified_exec::WriteStdinRequest; -use crate::unified_exec::process::OutputHandles; use codex_exec_server::ExecProcess; use codex_exec_server::ExecProcessEventReceiver; use codex_exec_server::ExecProcessFuture; @@ -150,20 +149,9 @@ async fn exec_command_with_tty( .insert(process_id, entry); } - let OutputHandles { - output_buffer, - output_notify, - output_closed, - output_closed_notify, - cancellation_token, - } = process.output_handles(); let deadline = started_at + Duration::from_millis(yield_time_ms); let collected_output = UnifiedExecProcessManager::collect_output_until_deadline( - &output_buffer, - &output_notify, - &output_closed, - &output_closed_notify, - &cancellation_token, + process.output_handles(), Some(session.subscribe_elicitation_pause_state()), deadline, ) @@ -809,19 +797,8 @@ async fn unified_exec_uses_remote_exec_server_when_configured() -> anyhow::Resul process.write(b"printf 'remote-unified-exec\\n'\n").await?; tokio::time::sleep(Duration::from_millis(100)).await; - let crate::unified_exec::process::OutputHandles { - output_buffer, - output_notify, - output_closed, - output_closed_notify, - cancellation_token, - } = process.output_handles(); let collected = UnifiedExecProcessManager::collect_output_until_deadline( - &output_buffer, - &output_notify, - &output_closed, - &output_closed_notify, - &cancellation_token, + process.output_handles(), /*pause_state*/ None, Instant::now() + Duration::from_millis(2_500), ) diff --git a/codex-rs/core/src/unified_exec/process.rs b/codex-rs/core/src/unified_exec/process.rs index dd10930547..e8ec6b9882 100644 --- a/codex-rs/core/src/unified_exec/process.rs +++ b/codex-rs/core/src/unified_exec/process.rs @@ -58,6 +58,7 @@ impl SpawnLifecycle for NoopSpawnLifecycle {} pub(crate) type OutputBuffer = Arc>; /// Shared output state exposed to polling and streaming consumers. +#[derive(Clone)] pub(crate) struct OutputHandles { pub(crate) output_buffer: OutputBuffer, pub(crate) output_notify: Arc, @@ -89,11 +90,7 @@ enum ProcessHandle { pub(crate) struct UnifiedExecProcess { process_handle: ProcessHandle, output_tx: broadcast::Sender>, - output_buffer: OutputBuffer, - output_notify: Arc, - output_closed: Arc, - output_closed_notify: Arc, - cancellation_token: CancellationToken, + output: OutputHandles, output_drained: Arc, interaction_lock: Arc>, state_tx: watch::Sender, @@ -119,11 +116,13 @@ impl UnifiedExecProcess { sandbox_type: SandboxType, spawn_lifecycle: Option, ) -> Self { - let output_buffer = Arc::new(Mutex::new(HeadTailBuffer::default())); - let output_notify = Arc::new(Notify::new()); - let output_closed = Arc::new(AtomicBool::new(false)); - let output_closed_notify = Arc::new(Notify::new()); - let cancellation_token = CancellationToken::new(); + let output = OutputHandles { + output_buffer: Arc::new(Mutex::new(HeadTailBuffer::default())), + output_notify: Arc::new(Notify::new()), + output_closed: Arc::new(AtomicBool::new(false)), + output_closed_notify: Arc::new(Notify::new()), + cancellation_token: CancellationToken::new(), + }; let output_drained = Arc::new(Notify::new()); let (output_tx, _) = broadcast::channel(64); let (state_tx, state_rx) = watch::channel(ProcessState::default()); @@ -131,11 +130,7 @@ impl UnifiedExecProcess { Self { process_handle, output_tx, - output_buffer, - output_notify, - output_closed, - output_closed_notify, - cancellation_token, + output, output_drained, interaction_lock: Arc::new(Mutex::new(())), state_tx, @@ -160,7 +155,7 @@ impl UnifiedExecProcess { WriteStatus::UnknownProcess | WriteStatus::StdinClosed => { let state = self.state_rx.borrow().clone(); let _ = self.state_tx.send_replace(state.exited(state.exit_code)); - self.cancellation_token.cancel(); + self.output.cancellation_token.cancel(); Err(UnifiedExecError::WriteToStdin) } WriteStatus::Starting => Err(UnifiedExecError::WriteToStdin), @@ -171,14 +166,8 @@ impl UnifiedExecProcess { } } - pub(super) fn output_handles(&self) -> OutputHandles { - OutputHandles { - output_buffer: Arc::clone(&self.output_buffer), - output_notify: Arc::clone(&self.output_notify), - output_closed: Arc::clone(&self.output_closed), - output_closed_notify: Arc::clone(&self.output_closed_notify), - cancellation_token: self.cancellation_token.clone(), - } + pub(super) fn output_handles(&self) -> &OutputHandles { + &self.output } pub(super) fn output_receiver(&self) -> tokio::sync::broadcast::Receiver> { @@ -186,7 +175,7 @@ impl UnifiedExecProcess { } pub(super) fn cancellation_token(&self) -> CancellationToken { - self.cancellation_token.clone() + self.output.cancellation_token.clone() } pub(super) fn output_drained_notify(&self) -> Arc { @@ -216,7 +205,7 @@ impl UnifiedExecProcess { } fn finish_termination(&self) { - self.cancellation_token.cancel(); + self.output.cancellation_token.cancel(); if let Some(output_task) = &self.output_task { output_task.abort(); } @@ -271,7 +260,7 @@ impl UnifiedExecProcess { } async fn snapshot_output(&self) -> Vec> { - let guard = self.output_buffer.lock().await; + let guard = self.output.output_buffer.lock().await; guard.snapshot_chunks() } @@ -284,8 +273,11 @@ impl UnifiedExecProcess { } pub(super) async fn check_for_sandbox_denial(&self) -> Result<(), UnifiedExecError> { - let _ = - tokio::time::timeout(Duration::from_millis(20), self.output_notify.notified()).await; + let _ = tokio::time::timeout( + Duration::from_millis(20), + self.output.output_notify.notified(), + ) + .await; let collected_chunks = self.snapshot_output().await; let mut aggregated: Vec = Vec::new(); @@ -354,10 +346,7 @@ impl UnifiedExecProcess { ); managed.output_task = Some(Self::spawn_local_output_task( output_rx, - Arc::clone(&managed.output_buffer), - Arc::clone(&managed.output_notify), - Arc::clone(&managed.output_closed), - Arc::clone(&managed.output_closed_notify), + managed.output_handles().clone(), managed.output_tx.clone(), )); @@ -383,7 +372,7 @@ impl UnifiedExecProcess { tokio::spawn({ let state_tx = managed.state_tx.clone(); - let cancellation_token = managed.cancellation_token.clone(); + let cancellation_token = managed.output.cancellation_token.clone(); async move { let exit_code = exit_rx.await.ok(); let state = state_tx.borrow().clone(); @@ -403,7 +392,7 @@ impl UnifiedExecProcess { // classification rather than attributing a violation to a guessed backend. let sandbox_type = started.sandbox_type.unwrap_or(SandboxType::None); let mut managed = Self::new(process_handle, sandbox_type, /*spawn_lifecycle*/ None); - let output_handles = managed.output_handles(); + let output_handles = managed.output_handles().clone(); managed.output_task = Some(Self::spawn_exec_server_output_task( started, output_handles, @@ -601,12 +590,16 @@ impl UnifiedExecProcess { fn spawn_local_output_task( mut receiver: tokio::sync::broadcast::Receiver>, - buffer: OutputBuffer, - output_notify: Arc, - output_closed: Arc, - output_closed_notify: Arc, + output_handles: OutputHandles, output_tx: broadcast::Sender>, ) -> JoinHandle<()> { + let OutputHandles { + output_buffer, + output_notify, + output_closed, + output_closed_notify, + .. + } = output_handles; tokio::spawn(async move { let _output_task_guard = OutputTaskGuard { output_closed: Arc::clone(&output_closed), @@ -615,7 +608,7 @@ impl UnifiedExecProcess { loop { match receiver.recv().await { Ok(chunk) => { - let mut guard = buffer.lock().await; + let mut guard = output_buffer.lock().await; guard.push_chunk(chunk.clone()); drop(guard); let _ = output_tx.send(chunk); @@ -635,7 +628,7 @@ impl UnifiedExecProcess { fn signal_exit(&self, exit_code: Option) { let state = self.state_rx.borrow().clone(); let _ = self.state_tx.send_replace(state.exited(exit_code)); - self.cancellation_token.cancel(); + self.output.cancellation_token.cancel(); } } diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index c46c4d1926..f40b89d195 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -6,7 +6,6 @@ use std::num::NonZeroUsize; use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; -use tokio::sync::Notify; use tokio::sync::watch; use tokio::time::Duration; use tokio::time::Instant; @@ -54,7 +53,6 @@ use crate::unified_exec::async_watcher::start_streaming_output; use crate::unified_exec::clamp_yield_time; use crate::unified_exec::generate_chunk_id; use crate::unified_exec::head_tail_buffer::HeadTailBuffer; -use crate::unified_exec::process::OutputBuffer; use crate::unified_exec::process::OutputHandles; use crate::unified_exec::process::SpawnLifecycleHandle; use crate::unified_exec::process::UnifiedExecProcess; @@ -211,11 +209,7 @@ fn exec_server_params_for_request( /// Borrowed process state prepared for a `write_stdin` or poll operation. struct PreparedProcessHandles { process: Arc, - output_buffer: OutputBuffer, - output_notify: Arc, - output_closed: Arc, - output_closed_notify: Arc, - cancellation_token: CancellationToken, + output: OutputHandles, pause_state: Option>, session: Option>, network_approval: Option, @@ -502,20 +496,9 @@ impl UnifiedExecProcessManager { // For the initial exec_command call, we both stream output to events // (via start_streaming_output above) and collect a snapshot here for // the tool response body. - let OutputHandles { - output_buffer, - output_notify, - output_closed, - output_closed_notify, - cancellation_token, - } = process.output_handles(); let deadline = start + Duration::from_millis(yield_time_ms); let collected_output = Self::collect_output_until_deadline( - &output_buffer, - &output_notify, - &output_closed, - &output_closed_notify, - &cancellation_token, + process.output_handles(), Some(context.session.subscribe_elicitation_pause_state()), deadline, ) @@ -700,11 +683,7 @@ impl UnifiedExecProcessManager { let PreparedProcessHandles { process, - output_buffer, - output_notify, - output_closed, - output_closed_notify, - cancellation_token, + output, pause_state, session, network_approval, @@ -760,16 +739,8 @@ impl UnifiedExecProcessManager { }; let start = Instant::now(); let deadline = start + Duration::from_millis(yield_time_ms); - let collected_output = Self::collect_output_until_deadline( - &output_buffer, - &output_notify, - &output_closed, - &output_closed_notify, - &cancellation_token, - pause_state, - deadline, - ) - .await; + let collected_output = + Self::collect_output_until_deadline(&output, pause_state, deadline).await; let wall_time = Instant::now().saturating_duration_since(start); let original_token_count = usize::try_from(approx_tokens_from_byte_count( @@ -911,13 +882,7 @@ impl UnifiedExecProcessManager { return Err(UnifiedExecError::UnknownProcessId { process_id }); } entry.last_used = Instant::now(); - let OutputHandles { - output_buffer, - output_notify, - output_closed, - output_closed_notify, - cancellation_token, - } = entry.process.output_handles(); + let output = entry.process.output_handles().clone(); let pause_state = entry .session .upgrade() @@ -926,11 +891,7 @@ impl UnifiedExecProcessManager { Ok(PreparedProcessHandles { process: Arc::clone(&entry.process), - output_buffer, - output_notify, - output_closed, - output_closed_notify, - cancellation_token, + output, pause_state, session, network_approval: entry.network_approval.clone(), @@ -1264,16 +1225,19 @@ impl UnifiedExecProcessManager { } pub(super) async fn collect_output_until_deadline( - output_buffer: &OutputBuffer, - output_notify: &Arc, - output_closed: &Arc, - output_closed_notify: &Arc, - cancellation_token: &CancellationToken, + output: &OutputHandles, mut pause_state: Option>, mut deadline: Instant, ) -> HeadTailBuffer { const POST_EXIT_CLOSE_WAIT_CAP: Duration = Duration::from_millis(50); + let OutputHandles { + output_buffer, + output_notify, + output_closed, + output_closed_notify, + cancellation_token, + } = output; let mut collected = HeadTailBuffer::default(); let mut exit_signal_received = cancellation_token.is_cancelled(); let mut post_exit_deadline: Option = None; diff --git a/codex-rs/core/src/unified_exec/process_manager_tests.rs b/codex-rs/core/src/unified_exec/process_manager_tests.rs index 7c6f484de3..ead8d6257b 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -2,6 +2,7 @@ use super::*; use crate::unified_exec::clamp_yield_time; use codex_network_proxy::ManagedNetworkSandboxContext; use pretty_assertions::assert_eq; +use tokio::sync::Notify; use tokio::time::Duration; use tokio::time::Instant; @@ -258,13 +259,16 @@ async fn output_collection_stays_bounded_across_repeated_drains() { let output_closed = Arc::new(AtomicBool::new(false)); let output_closed_notify = Arc::new(Notify::new()); let cancellation_token = CancellationToken::new(); + let output = OutputHandles { + output_buffer: Arc::clone(&output_buffer), + output_notify: Arc::clone(&output_notify), + output_closed: Arc::clone(&output_closed), + output_closed_notify: Arc::clone(&output_closed_notify), + cancellation_token: cancellation_token.clone(), + }; let collect = UnifiedExecProcessManager::collect_output_until_deadline( - &output_buffer, - &output_notify, - &output_closed, - &output_closed_notify, - &cancellation_token, + &output, /*pause_state*/ None, Instant::now() + Duration::from_secs(5), ); @@ -323,13 +327,16 @@ async fn output_collection_preserves_omissions_from_drained_buffer() { let output_closed_notify = Arc::new(Notify::new()); let cancellation_token = CancellationToken::new(); cancellation_token.cancel(); + let output = OutputHandles { + output_buffer, + output_notify, + output_closed, + output_closed_notify, + cancellation_token, + }; let collected = UnifiedExecProcessManager::collect_output_until_deadline( - &output_buffer, - &output_notify, - &output_closed, - &output_closed_notify, - &cancellation_token, + &output, /*pause_state*/ None, Instant::now() + Duration::from_secs(1), )