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

code-mode: harden actor lifecycle races

Browse Source
This commit is contained in:
Channing Conger
2026-06-19 07:32:59 +00:00
parent 2e9a35d530
commit a736ac14ad
3 changed files with 530 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -177,6 +177,13 @@ async fn run_cell<H: CellHost>(
let _ = response_tx.send(Ok(event));
break;
}
if observer
.as_ref()
.is_some_and(|observer| observer.response_tx.is_closed())
{
observer = None;
yield_timer = None;
}
if observer.is_some() || termination.is_some() {
let _ = response_tx.send(Err(CellError::Busy));
continue;

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

@@ -9,6 +9,7 @@ use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::FunctionCallOutputContentItem;
use pretty_assertions::assert_eq;
use serde_json::Value as JsonValue;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
@@ -18,6 +19,17 @@ use crate::session_runtime::OutputItem;
struct TestHost;
#[derive(Default)]
struct RecordingHost {
committed: AtomicBool,
closed: AtomicBool,
}
struct BlockingCommitHost {
commit_started_tx: mpsc::UnboundedSender<()>,
commit_release: Semaphore,
}
impl CellHost for TestHost {
async fn invoke_tool(
&self,
@@ -41,6 +53,78 @@ impl CellHost for TestHost {
async fn closed(&self) {}
}
impl CellHost for RecordingHost {
async fn invoke_tool(
&self,
_invocation: CellToolCall,
_cancellation_token: CancellationToken,
) -> Result<JsonValue, String> {
Err("unexpected tool call".to_string())
}
async fn notify(
&self,
_call_id: String,
_text: String,
_cancellation_token: CancellationToken,
) -> Result<(), String> {
Ok(())
}
async fn commit_stored_values(&self, _stored_value_writes: HashMap<String, JsonValue>) {
self.committed.store(true, Ordering::Release);
}
async fn closed(&self) {
self.closed.store(true, Ordering::Release);
}
}
impl BlockingCommitHost {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<()>) {
let (commit_started_tx, commit_started_rx) = mpsc::unbounded_channel();
(
Arc::new(Self {
commit_started_tx,
commit_release: Semaphore::new(/*permits*/ 0),
}),
commit_started_rx,
)
}
}
impl CellHost for BlockingCommitHost {
async fn invoke_tool(
&self,
_invocation: CellToolCall,
_cancellation_token: CancellationToken,
) -> Result<JsonValue, String> {
Err("unexpected tool call".to_string())
}
async fn notify(
&self,
_call_id: String,
_text: String,
_cancellation_token: CancellationToken,
) -> Result<(), String> {
Ok(())
}
async fn commit_stored_values(&self, _stored_value_writes: HashMap<String, JsonValue>) {
self.commit_started_tx
.send(())
.expect("test did not receive commit start");
self.commit_release
.acquire()
.await
.expect("test did not release commit")
.forget();
}
async fn closed(&self) {}
}
struct CellActorHarness {
event_tx: mpsc::UnboundedSender<RuntimeEvent>,
command_tx: mpsc::UnboundedSender<CellCommand>,
@@ -52,6 +136,13 @@ struct CellActorHarness {
}
fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarness {
spawn_cell_actor_harness_with_host(initial_observe_mode, Arc::new(TestHost))
}
fn spawn_cell_actor_harness_with_host<H: CellHost>(
initial_observe_mode: ObserveMode,
host: Arc<H>,
) -> CellActorHarness {
let (event_tx, event_rx) = mpsc::unbounded_channel();
let (command_tx, command_rx) = mpsc::unbounded_channel();
let (initial_event_tx, initial_event_rx) = oneshot::channel();
@@ -70,10 +161,10 @@ fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarne
)
.unwrap();
let session_shutdown_token = CancellationToken::new();
let cancellation_token = CancellationToken::new();
let cancellation_token = session_shutdown_token.child_token();
let handle = CellHandle::new(command_tx.clone());
let task = tokio::spawn(run_cell(
Arc::new(TestHost),
host,
CellContext {
runtime_tx,
runtime_control_tx,
@@ -153,6 +244,165 @@ async fn queued_termination_preempts_unobserved_runtime_completion() {
harness.task.await.unwrap();
}
#[tokio::test]
async fn termination_preempts_result_without_committing_stored_values() {
let host = Arc::new(RecordingHost::default());
let harness = spawn_cell_actor_harness_with_host(
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
Arc::clone(&host),
);
harness
.event_tx
.send(RuntimeEvent::Result {
stored_value_writes: HashMap::from([("key".to_string(), JsonValue::Bool(true))]),
error_text: None,
})
.unwrap();
let termination = harness.handle.terminate();
let terminated = Ok(CellEvent::Terminated {
content_items: Vec::new(),
});
assert_eq!(termination.await, terminated.clone());
assert_eq!(harness.initial_event_rx.await.unwrap(), terminated);
harness.task.await.unwrap();
assert!(!host.committed.load(Ordering::Acquire));
assert!(host.closed.load(Ordering::Acquire));
}
#[tokio::test]
async fn observer_receives_a_buffered_completion_after_commit_finishes() {
let (host, mut commit_started_rx) = BlockingCommitHost::new();
let harness = spawn_cell_actor_harness_with_host(
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
Arc::clone(&host),
);
harness.event_tx.send(RuntimeEvent::Started).unwrap();
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
assert_eq!(
harness.initial_event_rx.await.unwrap(),
Ok(CellEvent::Yielded {
content_items: Vec::new(),
})
);
harness
.event_tx
.send(RuntimeEvent::Result {
stored_value_writes: HashMap::new(),
error_text: None,
})
.unwrap();
assert_eq!(commit_started_rx.recv().await, Some(()));
let completion = harness
.handle
.observe(ObserveMode::YieldAfter(Duration::ZERO));
host.commit_release.add_permits(/*n*/ 1);
assert_eq!(
completion.await,
Ok(CellEvent::Completed {
content_items: Vec::new(),
error_text: None,
})
);
harness.task.await.unwrap();
}
#[tokio::test]
async fn dropped_observer_does_not_block_the_next_observer() {
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs(
/*secs*/ 60,
)));
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
assert_eq!(
harness.initial_event_rx.await.unwrap(),
Ok(CellEvent::Yielded {
content_items: Vec::new(),
})
);
let abandoned_observer = harness
.handle
.observe(ObserveMode::YieldAfter(Duration::ZERO));
drop(abandoned_observer);
let next_observer = harness
.handle
.observe(ObserveMode::YieldAfter(Duration::ZERO));
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
assert_eq!(
next_observer.await,
Ok(CellEvent::Yielded {
content_items: Vec::new(),
})
);
let termination = harness.handle.terminate();
drop(harness.event_tx);
assert_eq!(
termination.await,
Ok(CellEvent::Terminated {
content_items: Vec::new(),
})
);
harness.task.await.unwrap();
}
#[tokio::test]
async fn session_shutdown_terminates_the_active_observer() {
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs(
/*secs*/ 60,
)));
harness.session_shutdown_token.cancel();
drop(harness.event_tx);
assert_eq!(
harness.initial_event_rx.await.unwrap(),
Ok(CellEvent::Terminated {
content_items: Vec::new(),
})
);
harness.task.await.unwrap();
}
#[tokio::test]
async fn unexpected_runtime_close_reports_a_completed_error_and_closes_the_actor() {
let host = Arc::new(RecordingHost::default());
let harness = spawn_cell_actor_harness_with_host(
ObserveMode::YieldAfter(Duration::from_secs(/*secs*/ 60)),
Arc::clone(&host),
);
drop(harness.event_tx);
assert_eq!(
harness.initial_event_rx.await.unwrap(),
Ok(CellEvent::Completed {
content_items: Vec::new(),
error_text: Some("exec runtime ended unexpectedly".to_string()),
})
);
harness.task.await.unwrap();
assert!(host.closed.load(Ordering::Acquire));
}
#[tokio::test]
async fn termination_wins_when_the_runtime_channel_closes() {
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs(
/*secs*/ 60,
)));
let termination = harness.handle.terminate();
drop(harness.event_tx);
let terminated = Ok(CellEvent::Terminated {
content_items: Vec::new(),
});
assert_eq!(termination.await, terminated.clone());
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(

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

@@ -16,6 +16,7 @@ use crate::ToolDefinition;
#[derive(Debug, PartialEq)]
enum DelegateEvent {
NotificationStarted,
NotificationFinished,
NotificationCancelled,
ToolStarted,
ToolCancelled,
@@ -34,6 +35,11 @@ struct HeldNotificationDelegate {
notification_release: Notify,
}
struct ReleasableNotificationDelegate {
events_tx: mpsc::UnboundedSender<DelegateEvent>,
notification_release: Notify,
}
impl HeldNotificationDelegate {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
let (events_tx, events_rx) = mpsc::unbounded_channel();
@@ -51,6 +57,23 @@ impl HeldNotificationDelegate {
}
}
impl ReleasableNotificationDelegate {
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();
}
}
impl CodeModeSessionDelegate for HeldNotificationDelegate {
fn invoke_tool<'a>(
&'a self,
@@ -58,7 +81,9 @@ impl CodeModeSessionDelegate for HeldNotificationDelegate {
cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async move {
let _ = self.events_tx.send(DelegateEvent::ToolStarted);
cancellation_token.cancelled().await;
let _ = self.events_tx.send(DelegateEvent::ToolCancelled);
Err("cancelled".to_string())
})
}
@@ -86,6 +111,47 @@ impl CodeModeSessionDelegate for HeldNotificationDelegate {
}
}
impl CodeModeSessionDelegate for ReleasableNotificationDelegate {
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())
})
}
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);
tokio::select! {
_ = self.notification_release.notified() => {
let _ = self.events_tx.send(DelegateEvent::NotificationFinished);
Ok(())
}
_ = cancellation_token.cancelled() => {
let _ = self.events_tx.send(DelegateEvent::NotificationCancelled);
Err("cancelled".to_string())
}
}
})
}
fn cell_closed(&self, cell_id: &CellId) {
let _ = self
.events_tx
.send(DelegateEvent::CellClosed(cell_id.clone()));
}
}
impl BlockingDelegate {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
let (events_tx, events_rx) = mpsc::unbounded_channel();
@@ -392,6 +458,159 @@ async fn shutdown_cancels_notifications_while_natural_completion_is_draining() {
);
}
#[tokio::test]
async fn natural_completion_waits_for_notifications_before_responding() {
let (delegate, mut events_rx) = ReleasableNotificationDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
let cell = service
.execute(execute_request(r#"notify("pending");"#))
.await
.unwrap();
let mut initial_response = Box::pin(cell.initial_response());
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationStarted
);
assert!(
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), &mut initial_response)
.await
.is_err()
);
delegate.release_notification();
assert_eq!(
initial_response.await.unwrap(),
RuntimeResponse::Result {
cell_id: cell_id("1"),
content_items: Vec::new(),
error_text: None,
}
);
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationFinished
);
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
);
}
#[tokio::test]
async fn termination_cancels_pending_tools_before_responding() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
let cell = 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);
assert_eq!(
cell.initial_response().await.unwrap(),
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
}
);
assert_eq!(
service.terminate(cell_id("1")).await.unwrap(),
WaitOutcome::LiveCell(RuntimeResponse::Terminated {
cell_id: cell_id("1"),
content_items: Vec::new(),
})
);
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"))
);
}
#[tokio::test]
async fn shutdown_cancels_pending_tools_before_returning() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate.clone()));
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);
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::ToolCancelled
);
assert_eq!(shutdown.await.unwrap(), Ok(()));
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
);
}
#[tokio::test]
async fn shutdown_cancels_mixed_callbacks_while_natural_completion_is_draining() {
let (delegate, mut events_rx) = HeldNotificationDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate.clone()));
service
.execute(ExecuteRequest {
enabled_tools: vec![blocking_tool()],
source: r#"notify("pending"); tools.block({});"#.to_string(),
..execute_request("")
})
.await
.unwrap();
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationStarted
);
assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);
let shutdown_service = Arc::clone(&service);
let shutdown = tokio::spawn(async move { shutdown_service.shutdown().await });
let first_cancelled = next_event(&mut events_rx).await;
let second_cancelled = next_event(&mut events_rx).await;
assert!(matches!(
(&first_cancelled, &second_cancelled),
(
&DelegateEvent::NotificationCancelled,
&DelegateEvent::ToolCancelled,
) | (
&DelegateEvent::ToolCancelled,
&DelegateEvent::NotificationCancelled,
)
));
delegate.release_notification();
assert_eq!(shutdown.await.unwrap(), Ok(()));
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
);
}
#[tokio::test]
async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {
let (delegate, mut events_rx) = HeldNotificationDelegate::new();
@@ -423,6 +642,9 @@ async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {
DelegateEvent::NotificationCancelled
);
let shutdown_service = Arc::clone(&service);
let shutdown = tokio::spawn(async move { shutdown_service.shutdown().await });
let repeated_termination = service.terminate(cell_id("1")).await;
delegate.release_notification();
@@ -437,6 +659,7 @@ async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {
content_items: Vec::new(),
})
);
assert_eq!(shutdown.await.unwrap(), Ok(()));
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
@@ -490,6 +713,54 @@ async fn second_observer_is_rejected_without_displacing_the_first() {
);
}
#[tokio::test]
async fn dropping_a_wait_allows_a_later_wait_to_observe_the_cell() {
let service = CodeModeService::new();
let cell = service
.execute(execute_request(
"yield_control(); await new Promise(() => {});",
))
.await
.unwrap();
assert_eq!(
cell.initial_response().await.unwrap(),
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
}
);
let abandoned_wait = service
.begin_wait(WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 60_000,
})
.await;
drop(abandoned_wait);
assert_eq!(
service
.wait(WaitRequest {
cell_id: cell_id("1"),
yield_time_ms: 0,
})
.await
.unwrap(),
WaitOutcome::LiveCell(RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
})
);
assert_eq!(
service.terminate(cell_id("1")).await.unwrap(),
WaitOutcome::LiveCell(RuntimeResponse::Terminated {
cell_id: cell_id("1"),
content_items: Vec::new(),
})
);
}
#[tokio::test]
async fn natural_completion_cleans_up_callbacks_before_responding() {
let (delegate, mut events_rx) = BlockingDelegate::new();