diff --git a/codex-rs/code-mode/src/cell_actor/mod.rs b/codex-rs/code-mode/src/cell_actor/mod.rs index 11f0501752..9889d32ce9 100644 --- a/codex-rs/code-mode/src/cell_actor/mod.rs +++ b/codex-rs/code-mode/src/cell_actor/mod.rs @@ -104,11 +104,6 @@ struct Termination { response_tx: Option>>, } -enum CommandEvent { - Received(Option), - SessionShutdown, -} - async fn run_cell( host: Arc, context: CellContext, @@ -140,31 +135,25 @@ async fn run_cell( .is_some_and(|yield_timer| yield_timer.deadline() <= tokio::time::Instant::now()); tokio::select! { biased; - command_event = async { + command_outcome = async { tokio::select! { biased; _ = session_shutdown_token.cancelled(), if command_rx.is_some() => { - CommandEvent::SessionShutdown + drop(command_rx.take()); + None } command = async { match command_rx.as_mut() { Some(command_rx) => command_rx.recv().await, None => std::future::pending::>().await, } - } => { - CommandEvent::Received(command) - } + } => Some(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 maybe_command = match command_outcome { + Some(command) => command, + None if termination.is_some() => continue, + None => Some(CellCommand::Terminate { response_tx: None }), }; let Some(command) = maybe_command else { if completed_event.is_some() { diff --git a/codex-rs/code-mode/src/session_runtime/mod.rs b/codex-rs/code-mode/src/session_runtime/mod.rs index 2f345b9a30..d74e86efb8 100644 --- a/codex-rs/code-mode/src/session_runtime/mod.rs +++ b/codex-rs/code-mode/src/session_runtime/mod.rs @@ -168,7 +168,7 @@ impl SessionRuntime { stored_values, host, initial_observe_mode, - self.inner.shutdown_token.child_token(), + self.inner.shutdown_token.clone(), ) .map_err(Error::Runtime)?; cells.insert(cell_id.clone(), CellState::Live(handle));