From 269ac1ca0485f419bedbd59f8d449bbc3020df77 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Thu, 18 Jun 2026 23:42:02 +0000 Subject: [PATCH] code-mode: make session shutdown authoritative --- codex-rs/code-mode/src/cell_actor/mod.rs | 35 +++++++++-- codex-rs/code-mode/src/cell_actor/tests.rs | 60 ++++++++++++++++++- codex-rs/code-mode/src/session_runtime/mod.rs | 12 ++-- 3 files changed, 91 insertions(+), 16 deletions(-) diff --git a/codex-rs/code-mode/src/cell_actor/mod.rs b/codex-rs/code-mode/src/cell_actor/mod.rs index 6fcac36d62..c1f6c03c70 100644 --- a/codex-rs/code-mode/src/cell_actor/mod.rs +++ b/codex-rs/code-mode/src/cell_actor/mod.rs @@ -104,11 +104,16 @@ struct Termination { response_tx: Option>>, } +enum CommandEvent { + Received(Option), + SessionShutdown, +} + async fn run_cell( host: Arc, context: CellContext, mut event_rx: mpsc::UnboundedReceiver, - mut command_rx: mpsc::UnboundedReceiver, + command_rx: mpsc::UnboundedReceiver, initial_observer: Observer, ) { let CellContext { @@ -128,21 +133,39 @@ async fn run_cell( let mut yield_timer: Option>> = None; let mut notification_tasks = JoinSet::new(); let mut tool_tasks = JoinSet::new(); + let mut command_rx = Some(command_rx); loop { let yield_deadline_elapsed = yield_timer .as_ref() .is_some_and(|yield_timer| yield_timer.deadline() <= tokio::time::Instant::now()); tokio::select! { biased; - maybe_command = async { + command_event = async { tokio::select! { biased; - command = command_rx.recv() => command, - _ = session_shutdown_token.cancelled(), if termination.is_none() => { - Some(CellCommand::Terminate { response_tx: None }) + _ = session_shutdown_token.cancelled(), if command_rx.is_some() => { + CommandEvent::SessionShutdown + } + command = async { + match command_rx.as_mut() { + Some(command_rx) => command_rx.recv().await, + None => std::future::pending::>().await, + } + } => { + CommandEvent::Received(command) } } } => { + let maybe_command = match command_event { + CommandEvent::Received(command) => command, + CommandEvent::SessionShutdown => { + drop(command_rx.take()); + if termination.is_some() { + continue; + } + Some(CellCommand::Terminate { response_tx: None }) + } + }; let Some(command) = maybe_command else { if completed_event.is_some() { break; @@ -396,7 +419,7 @@ async fn run_cell( } } // Reject requests that arrive while asynchronous terminal cleanup runs. - drop(command_rx); + drop(command_rx.take()); begin_termination( &runtime_tx, &runtime_control_tx, diff --git a/codex-rs/code-mode/src/cell_actor/tests.rs b/codex-rs/code-mode/src/cell_actor/tests.rs index 2d37d678ba..d3c10b249a 100644 --- a/codex-rs/code-mode/src/cell_actor/tests.rs +++ b/codex-rs/code-mode/src/cell_actor/tests.rs @@ -1,5 +1,8 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::thread; use std::time::Duration; use codex_code_mode_protocol::ExecuteRequest; @@ -40,8 +43,10 @@ impl CellHost for TestHost { struct CellActorHarness { event_tx: mpsc::UnboundedSender, + command_tx: mpsc::UnboundedSender, handle: CellHandle, initial_event_rx: oneshot::Receiver>, + session_shutdown_token: CancellationToken, task: tokio::task::JoinHandle<()>, _runtime_event_rx: mpsc::UnboundedReceiver, } @@ -65,7 +70,8 @@ fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarne ) .unwrap(); let cancellation_token = CancellationToken::new(); - let handle = CellHandle::new(command_tx, cancellation_token.clone()); + let session_shutdown_token = CancellationToken::new(); + let handle = CellHandle::new(command_tx.clone(), cancellation_token.clone()); let task = tokio::spawn(run_cell( Arc::new(TestHost), CellContext { @@ -73,7 +79,7 @@ fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarne runtime_control_tx, runtime_terminate_handle, cancellation_token, - session_shutdown_token: CancellationToken::new(), + session_shutdown_token: session_shutdown_token.clone(), }, event_rx, command_rx, @@ -85,8 +91,10 @@ fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarne CellActorHarness { event_tx, + command_tx, handle, initial_event_rx, + session_shutdown_token, task, _runtime_event_rx: runtime_event_rx, } @@ -144,3 +152,51 @@ async fn queued_termination_preempts_unobserved_runtime_completion() { assert_eq!(harness.initial_event_rx.await.unwrap(), terminated); harness.task.await.unwrap(); } + +#[tokio::test] +async fn session_shutdown_preempts_continuous_command_traffic() { + let mut harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs( + /*secs*/ 60, + ))); + let keep_sending = Arc::new(AtomicBool::new(true)); + let producer_keep_sending = Arc::clone(&keep_sending); + let command_tx = harness.command_tx.clone(); + let (producer_started_tx, producer_started_rx) = oneshot::channel(); + let producer = thread::spawn(move || { + let mut producer_started_tx = Some(producer_started_tx); + while producer_keep_sending.load(Ordering::Relaxed) { + let (response_tx, _response_rx) = oneshot::channel(); + if command_tx + .send(CellCommand::Observe { + mode: ObserveMode::PendingFrontier, + response_tx, + }) + .is_err() + { + break; + } + if let Some(producer_started_tx) = producer_started_tx.take() { + let _ = producer_started_tx.send(()); + } + } + }); + producer_started_rx.await.unwrap(); + + harness.session_shutdown_token.cancel(); + drop(harness.event_tx); + + let shutdown_result = + tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut harness.task).await; + + keep_sending.store(false, Ordering::Relaxed); + producer.join().unwrap(); + if shutdown_result.is_err() { + harness.task.abort(); + let _ = harness.task.await; + } + + match shutdown_result { + Ok(task_result) => task_result.unwrap(), + Err(_) => panic!("session shutdown did not finish while commands were queued"), + } +} diff --git a/codex-rs/code-mode/src/session_runtime/mod.rs b/codex-rs/code-mode/src/session_runtime/mod.rs index 893e54ef31..315b54bbf0 100644 --- a/codex-rs/code-mode/src/session_runtime/mod.rs +++ b/codex-rs/code-mode/src/session_runtime/mod.rs @@ -4,7 +4,6 @@ use std::collections::HashMap; use std::future::Future; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; @@ -46,7 +45,6 @@ struct Inner { cell_count_tx: watch::Sender, shutdown_token: CancellationToken, delegate: Arc, - shutting_down: AtomicBool, next_cell_id: AtomicU64, } @@ -65,14 +63,13 @@ impl SessionRuntime { cell_count_tx, shutdown_token: CancellationToken::new(), delegate, - shutting_down: AtomicBool::new(false), next_cell_id: AtomicU64::new(1), }), } } pub fn is_alive(&self) -> bool { - !self.inner.shutting_down.load(Ordering::Acquire) + !self.inner.shutdown_token.is_cancelled() } pub async fn execute( @@ -80,7 +77,7 @@ impl SessionRuntime { request: ExecuteRequest, initial_observe_mode: ObserveMode, ) -> Result { - if self.inner.shutting_down.load(Ordering::Acquire) { + if self.inner.shutdown_token.is_cancelled() { return Err(Error::ShuttingDown); } let cell_id = self.allocate_cell_id(); @@ -174,7 +171,7 @@ impl SessionRuntime { inner: Arc::clone(&self.inner), }); let mut cells = self.inner.cells.lock().await; - if self.inner.shutting_down.load(Ordering::Acquire) { + if self.inner.shutdown_token.is_cancelled() { return Err(Error::ShuttingDown); } if cells.contains_key(&cell_id) { @@ -185,7 +182,7 @@ impl SessionRuntime { stored_values, host, initial_observe_mode, - self.inner.shutdown_token.clone(), + self.inner.shutdown_token.child_token(), ) .map_err(Error::Runtime)?; cells.insert(cell_id.clone(), CellState::Live(handle)); @@ -196,7 +193,6 @@ impl SessionRuntime { } fn begin_shutdown(&self) { - self.inner.shutting_down.store(true, Ordering::Release); self.inner.shutdown_token.cancel(); // The token reaches every cell if the registry is temporarily locked. if let Ok(cells) = self.inner.cells.try_lock() {