openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-06 15:29:32 +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-20 18:07:21 +00:00
parent 82722d45e4
commit 442d204af4
5 changed files with 235 additions and 125 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
codex-rs/code-mode/src/cell_actor/callbacks.rs
View File

@@ -22,8 +22,14 @@ pub(super) fn spawn_notification<H: CellHost>(
cancellation_token: CancellationToken,
) {
tasks.spawn(async move {
if let Err(err) = host.notify(call_id, text, cancellation_token).await {
warn!("failed to deliver code mode notification: {err}");
tokio::select! {
biased;
_ = cancellation_token.cancelled() => {}
result = host.notify(call_id, text, cancellation_token.clone()) => {
if let Err(err) = result {
warn!("failed to deliver code mode notification: {err}");
}
}
}
});
}
@@ -37,9 +43,15 @@ pub(super) fn spawn_tool<H: CellHost>(
) {
tasks.spawn(async move {
let id = invocation.id.clone();
let command = match host.invoke_tool(invocation, cancellation_token).await {
Ok(result) => RuntimeCommand::ToolResponse { id, result },
Err(error_text) => RuntimeCommand::ToolError { id, error_text },
let command = tokio::select! {
biased;
_ = cancellation_token.cancelled() => {
return;
}
result = host.invoke_tool(invocation, cancellation_token.clone()) => match result {
Ok(result) => RuntimeCommand::ToolResponse { id, result },
Err(error_text) => RuntimeCommand::ToolError { id, error_text },
}
};
let _ = runtime_tx.send(command);
});

9
codex-rs/code-mode/src/cell_actor/types.rs
View File

@@ -93,16 +93,13 @@ impl CellHandle {
pub(crate) fn terminate(&self) -> CellEventFuture {
self.state.request_termination()
}
pub(crate) fn shutdown(&self) {
self.state.cancellation_token().cancel();
}
}
/// The single linearization point for a cell's terminal outcome.
///
/// Callback cancellation tokens are children of the cell token, so a terminal
/// decision cancels runtime work and its callbacks together.
/// The cancellation token is a child of the owning session token. Callback
/// tokens are children of this token, so cancellation flows strictly from the
/// session to the cell and then to its callbacks.
///
/// The mutex is held only for synchronous phase transitions and terminal
/// delivery. Runtime execution, observation waits, and callbacks never run

203
codex-rs/code-mode/src/service_contract_tests.rs
View File

@@ -16,50 +16,84 @@ use crate::ToolDefinition;
#[derive(Debug, PartialEq)]
enum DelegateEvent {
NotificationStarted,
NotificationCancelled,
ToolStarted,
ToolCancelled,
CellClosed(CellId),
}
struct BlockingDelegate {
events_tx: mpsc::UnboundedSender<DelegateEvent>,
notification_finished: AtomicBool,
tool_finished: AtomicBool,
tool_future_dropped: AtomicBool,
tool_release: Notify,
}
struct HeldNotificationDelegate {
events_tx: mpsc::UnboundedSender<DelegateEvent>,
notification_release: Notify,
struct DropFlag<'a>(&'a AtomicBool);
impl Drop for DropFlag<'_> {
fn drop(&mut self) {
self.0.store(true, Ordering::Release);
}
}
impl HeldNotificationDelegate {
struct NeverResolvingNotificationDelegate {
events_tx: mpsc::UnboundedSender<DelegateEvent>,
}
struct NeverResolvingToolDelegate {
events_tx: mpsc::UnboundedSender<DelegateEvent>,
}
impl NeverResolvingNotificationDelegate {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
let (events_tx, events_rx) = mpsc::unbounded_channel();
(
Arc::new(Self {
events_tx,
notification_release: Notify::new(),
}),
events_rx,
)
}
fn release_notification(&self) {
self.notification_release.notify_one();
(Arc::new(Self { events_tx }), events_rx)
}
}
impl CodeModeSessionDelegate for HeldNotificationDelegate {
impl NeverResolvingToolDelegate {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
let (events_tx, events_rx) = mpsc::unbounded_channel();
(Arc::new(Self { events_tx }), events_rx)
}
}
impl CodeModeSessionDelegate for NeverResolvingNotificationDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
_cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async { Err("unexpected tool call".to_string()) })
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
_cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async move {
let _ = self.events_tx.send(DelegateEvent::NotificationStarted);
std::future::pending().await
})
}
fn cell_closed(&self, cell_id: &CellId) {
let _ = self
.events_tx
.send(DelegateEvent::CellClosed(cell_id.clone()));
}
}
impl CodeModeSessionDelegate for NeverResolvingToolDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
_cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async move {
cancellation_token.cancelled().await;
Err("cancelled".to_string())
let _ = self.events_tx.send(DelegateEvent::ToolStarted);
std::future::pending().await
})
}
@@ -68,15 +102,9 @@ impl CodeModeSessionDelegate for HeldNotificationDelegate {
_call_id: String,
_cell_id: CellId,
_text: String,
cancellation_token: CancellationToken,
_cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async move {
let _ = self.events_tx.send(DelegateEvent::NotificationStarted);
cancellation_token.cancelled().await;
let _ = self.events_tx.send(DelegateEvent::NotificationCancelled);
self.notification_release.notified().await;
Ok(())
})
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, cell_id: &CellId) {
@@ -92,8 +120,7 @@ impl BlockingDelegate {
(
Arc::new(Self {
events_tx,
notification_finished: AtomicBool::new(false),
tool_finished: AtomicBool::new(false),
tool_future_dropped: AtomicBool::new(false),
tool_release: Notify::new(),
}),
events_rx,
@@ -111,16 +138,15 @@ impl CodeModeSessionDelegate for BlockingDelegate {
_invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
let drop_flag = DropFlag(&self.tool_future_dropped);
Box::pin(async move {
let _drop_flag = drop_flag;
let _ = self.events_tx.send(DelegateEvent::ToolStarted);
tokio::select! {
_ = self.tool_release.notified() => {
self.tool_finished.store(true, Ordering::Release);
Ok(serde_json::Value::Null)
}
_ = cancellation_token.cancelled() => {
self.tool_finished.store(true, Ordering::Release);
let _ = self.events_tx.send(DelegateEvent::ToolCancelled);
Err("cancelled".to_string())
}
}
@@ -137,8 +163,6 @@ impl CodeModeSessionDelegate for BlockingDelegate {
Box::pin(async move {
let _ = self.events_tx.send(DelegateEvent::NotificationStarted);
cancellation_token.cancelled().await;
self.notification_finished.store(true, Ordering::Release);
let _ = self.events_tx.send(DelegateEvent::NotificationCancelled);
Err("cancelled".to_string())
})
}
@@ -289,7 +313,10 @@ async fn observed_natural_completion_wins_over_termination() {
tokio::time::timeout(Duration::from_secs(1), async {
loop {
let response = service
.execute(execute_request(r#"text(String(load("finished")));"#))
.execute(ExecuteRequest {
yield_time_ms: Some(60_000),
..execute_request(r#"text(String(load("finished")));"#)
})
.await
.unwrap()
.initial_response()
@@ -323,7 +350,7 @@ async fn observed_natural_completion_wins_over_termination() {
}
#[tokio::test]
async fn termination_cancels_pending_callbacks_before_responding() {
async fn termination_discards_pending_callbacks_before_responding() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
let cell = service
@@ -351,11 +378,6 @@ async fn termination_cancels_pending_callbacks_before_responding() {
content_items: Vec::new(),
})
);
assert!(delegate.notification_finished.load(Ordering::Acquire));
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationCancelled
);
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
@@ -363,9 +385,9 @@ async fn termination_cancels_pending_callbacks_before_responding() {
}
#[tokio::test]
async fn shutdown_cancels_notifications_while_natural_completion_is_draining() {
let (delegate, mut events_rx) = HeldNotificationDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate.clone()));
async fn shutdown_does_not_await_notifications_during_natural_completion() {
let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate));
service
.execute(execute_request(r#"notify("pending");"#))
.await
@@ -377,15 +399,12 @@ async fn shutdown_cancels_notifications_while_natural_completion_is_draining() {
);
let shutdown_service = Arc::clone(&service);
let shutdown = tokio::spawn(async move { shutdown_service.shutdown().await });
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationCancelled
);
delegate.release_notification();
assert_eq!(shutdown.await.unwrap(), Ok(()));
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), async move {
shutdown_service.shutdown().await
})
.await
.expect("shutdown should not await a non-cooperative notification")
.unwrap();
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
@@ -393,9 +412,33 @@ async fn shutdown_cancels_notifications_while_natural_completion_is_draining() {
}
#[tokio::test]
async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {
let (delegate, mut events_rx) = HeldNotificationDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate.clone()));
async fn shutdown_does_not_await_a_non_cooperative_nested_tool() {
let (delegate, mut events_rx) = NeverResolvingToolDelegate::new();
let service = CodeModeService::with_delegate(delegate);
let _started = service
.execute(ExecuteRequest {
enabled_tools: vec![blocking_tool()],
source: r#"await tools.block({});"#.to_string(),
..execute_request("")
})
.await
.unwrap();
assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), service.shutdown())
.await
.expect("shutdown should not await a non-cooperative nested tool")
.unwrap();
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
);
}
#[tokio::test]
async fn termination_does_not_await_a_non_cooperative_notification() {
let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new();
let service = CodeModeService::with_delegate(delegate);
let cell = service
.execute(execute_request(
r#"notify("pending"); await new Promise(() => {});"#,
@@ -415,23 +458,14 @@ async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {
}
);
let terminating_service = Arc::clone(&service);
let first_termination =
tokio::spawn(async move { terminating_service.terminate(cell_id("1")).await });
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationCancelled
);
let repeated_termination = service.terminate(cell_id("1")).await;
delegate.release_notification();
assert_eq!(
repeated_termination.unwrap_err(),
"exec cell 1 is already terminating"
);
assert_eq!(
first_termination.await.unwrap().unwrap(),
tokio::time::timeout(
Duration::from_millis(/*millis*/ 100),
service.terminate(cell_id("1")),
)
.await
.expect("termination should not await a non-cooperative notification")
.unwrap(),
WaitOutcome::LiveCell(RuntimeResponse::Terminated {
cell_id: cell_id("1"),
content_items: Vec::new(),
@@ -497,7 +531,12 @@ async fn natural_completion_cleans_up_callbacks_before_responding() {
let cell = service
.execute(ExecuteRequest {
enabled_tools: vec![blocking_tool()],
source: r#"tools.block({}); text("done");"#.to_string(),
source: concat!(
"tools.block({});",
"await new Promise(resolve => setTimeout(resolve, 100));",
"text('done');",
)
.to_string(),
yield_time_ms: Some(60_000),
..execute_request("")
})
@@ -515,13 +554,5 @@ async fn natural_completion_cleans_up_callbacks_before_responding() {
error_text: None,
}
);
assert!(delegate.tool_finished.load(Ordering::Acquire));
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::ToolCancelled
);
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
);
assert!(delegate.tool_future_dropped.load(Ordering::Acquire));
}

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

@@ -4,13 +4,13 @@ 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;
use serde_json::Value as JsonValue;
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
pub(crate) use self::types::CellEvent;
pub(crate) use self::types::CellId;
@@ -43,8 +43,9 @@ pub(crate) struct SessionRuntime<D: SessionRuntimeDelegate> {
struct Inner<D: SessionRuntimeDelegate> {
stored_values: Mutex<HashMap<String, JsonValue>>,
cells: Mutex<HashMap<CellId, CellHandle>>,
cell_tasks: TaskTracker,
shutdown_token: CancellationToken,
delegate: Arc<D>,
shutting_down: AtomicBool,
next_cell_id: AtomicU64,
}
@@ -54,15 +55,16 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
inner: Arc::new(Inner {
stored_values: Mutex::new(HashMap::new()),
cells: Mutex::new(HashMap::new()),
cell_tasks: TaskTracker::new(),
shutdown_token: CancellationToken::new(),
delegate,
shutting_down: AtomicBool::new(false),
next_cell_id: AtomicU64::new(1),
}),
}
}
pub(crate) fn is_alive(&self) -> bool {
!self.inner.shutting_down.load(Ordering::Acquire)
!self.inner.shutdown_token.is_cancelled()
}
pub(crate) async fn execute(
@@ -70,6 +72,9 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
request: CreateCellRequest,
initial_observe_mode: ObserveMode,
) -> Result<StartedCell, Error> {
if self.inner.shutdown_token.is_cancelled() {
return Err(Error::ShuttingDown);
}
let cell_id = self.allocate_cell_id();
let initial_event = self
.start_cell(cell_id.clone(), request, initial_observe_mode)
@@ -122,21 +127,13 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
}
pub(crate) async fn shutdown(&self) -> Result<(), Error> {
self.inner.shutting_down.store(true, Ordering::Release);
let handles = self
.inner
.cells
.lock()
.await
.values()
.cloned()
.collect::<Vec<_>>();
for handle in handles {
handle.shutdown();
}
while !self.inner.cells.lock().await.is_empty() {
tokio::task::yield_now().await;
}
self.begin_shutdown();
// Taking the registry lock ensures every cell that passed the shutdown
// check has registered its actor with the tracker before we wait.
let cells = self.inner.cells.lock().await;
self.inner.cell_tasks.close();
drop(cells);
self.inner.cell_tasks.wait().await;
Ok(())
}
@@ -161,13 +158,13 @@ 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) {
return Err(Error::DuplicateCell(cell_id));
}
let cell_state = Arc::new(CellState::new(CancellationToken::new()));
let cell_state = Arc::new(CellState::new(self.inner.shutdown_token.child_token()));
let (handle, initial_event, task) = CellActor::prepare(
request,
stored_values,
@@ -177,18 +174,14 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
)
.map_err(Error::Runtime)?;
cells.insert(cell_id.clone(), handle);
self.inner.cell_tasks.spawn(task);
drop(cells);
tokio::spawn(task);
Ok(map_actor_event(cell_id, initial_event))
}
fn begin_shutdown(&self) {
self.inner.shutting_down.store(true, Ordering::Release);
if let Ok(cells) = self.inner.cells.try_lock() {
for handle in cells.values() {
handle.shutdown();
}
}
self.inner.shutdown_token.cancel();
self.inner.cell_tasks.close();
}
}

77
codex-rs/code-mode/src/session_runtime/tests.rs
View File

@@ -108,3 +108,80 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa
);
runtime.shutdown().await.unwrap();
}
fn execute_request(source: &str) -> CreateCellRequest {
CreateCellRequest {
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: source.to_string(),
}
}
#[tokio::test]
#[expect(
clippy::await_holding_invalid_type,
reason = "test holds the registry lock to force admission ahead of shutdown"
)]
async fn shutdown_rejects_cell_admission_queued_before_the_registry_lock() {
let runtime = Arc::new(SessionRuntime::new(Arc::new(RecordingDelegate)));
let cells = runtime.inner.cells.lock().await;
let execution = runtime.execute(
execute_request("while (true) {}"),
ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)),
);
tokio::pin!(execution);
std::future::poll_fn(|context| match execution.as_mut().poll(context) {
Poll::Pending => Poll::Ready(()),
Poll::Ready(Ok(_)) => panic!("execution completed before the registry lock was released"),
Poll::Ready(Err(error)) => {
panic!("execution failed before the registry lock was released: {error}")
}
})
.await;
let shutdown = runtime.shutdown();
tokio::pin!(shutdown);
std::future::poll_fn(|context| match shutdown.as_mut().poll(context) {
Poll::Pending => Poll::Ready(()),
Poll::Ready(Ok(())) => panic!("shutdown completed before acquiring the registry lock"),
Poll::Ready(Err(error)) => {
panic!("shutdown failed before acquiring the registry lock: {error}")
}
})
.await;
assert!(!runtime.is_alive());
drop(cells);
assert!(matches!(execution.await, Err(Error::ShuttingDown)));
assert_eq!(shutdown.await, Ok(()));
}
#[tokio::test]
async fn drop_terminates_cells_when_the_registry_is_locked() {
let runtime = SessionRuntime::new(Arc::new(RecordingDelegate));
let started = runtime
.execute(
execute_request("while (true) {}"),
ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)),
)
.await
.unwrap();
assert_eq!(started.cell_id, CellId::new("1"));
assert_eq!(
started.initial_event().await,
Ok(CellEvent::Yielded {
content_items: Vec::new(),
})
);
let inner = Arc::clone(&runtime.inner);
let cells = inner.cells.lock().await;
drop(runtime);
drop(cells);
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), inner.cell_tasks.wait())
.await
.unwrap();
assert!(inner.cell_tasks.is_empty());
}