openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-05 15:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

code-mode: make session shutdown authoritative

Browse Source
This commit is contained in:
Channing Conger
2026-06-18 23:42:02 +00:00
parent f8ba23d96a
commit 269ac1ca04
3 changed files with 91 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
codex-rs/code-mode/src/cell_actor/mod.rs
View File

@@ -104,11 +104,16 @@ struct Termination {
response_tx: Option<oneshot::Sender<Result<CellEvent, CellError>>>,
}
enum CommandEvent {
Received(Option<CellCommand>),
SessionShutdown,
}
async fn run_cell<H: CellHost>(
host: Arc<H>,
context: CellContext,
mut event_rx: mpsc::UnboundedReceiver<RuntimeEvent>,
mut command_rx: mpsc::UnboundedReceiver<CellCommand>,
command_rx: mpsc::UnboundedReceiver<CellCommand>,
initial_observer: Observer,
) {
let CellContext {
@@ -128,21 +133,39 @@ async fn run_cell<H: CellHost>(
let mut yield_timer: Option<std::pin::Pin<Box<tokio::time::Sleep>>> = 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::<Option<CellCommand>>().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<H: CellHost>(
}
}
// Reject requests that arrive while asynchronous terminal cleanup runs.
drop(command_rx);
drop(command_rx.take());
begin_termination(
&runtime_tx,
&runtime_control_tx,

60
codex-rs/code-mode/src/cell_actor/tests.rs
View File

@@ -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<RuntimeEvent>,
command_tx: mpsc::UnboundedSender<CellCommand>,
handle: CellHandle,
initial_event_rx: oneshot::Receiver<Result<CellEvent, CellError>>,
session_shutdown_token: CancellationToken,
task: tokio::task::JoinHandle<()>,
_runtime_event_rx: mpsc::UnboundedReceiver<RuntimeEvent>,
}
@@ -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"),
}
}

12
codex-rs/code-mode/src/session_runtime/mod.rs
View File

@@ -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<D: SessionRuntimeDelegate> {
cell_count_tx: watch::Sender<usize>,
shutdown_token: CancellationToken,
delegate: Arc<D>,
shutting_down: AtomicBool,
next_cell_id: AtomicU64,
}
@@ -65,14 +63,13 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
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<D: SessionRuntimeDelegate> SessionRuntime<D> {
request: ExecuteRequest,
initial_observe_mode: ObserveMode,
) -> Result<StartedCell, Error> {
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<D: SessionRuntimeDelegate> SessionRuntime<D> {
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<D: SessionRuntimeDelegate> SessionRuntime<D> {
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<D: SessionRuntimeDelegate> SessionRuntime<D> {
}
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() {