openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 17:38:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

code-mode: linearize cell termination and completion

Browse Source
This commit is contained in:
Channing Conger
2026-06-19 22:56:14 +00:00
parent 0dbe655fec
commit 641452f668
6 changed files with 511 additions and 196 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode/src/cell_actor/mod.rs
+100 -114
View File
@@ -5,12 +5,11 @@ mod types;
use std::collections::HashMap;
use std::future::Future;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use serde_json::Value as JsonValue;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::sync::watch;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
@@ -27,6 +26,7 @@ pub(crate) use self::types::CellError;
pub(crate) use self::types::CellEventFuture;
pub(crate) use self::types::CellHandle;
pub(crate) use self::types::CellHost;
pub(crate) use self::types::CellLifecycle;
pub(crate) use self::types::CellToolCall;
use crate::runtime::PendingRuntimeMode;
use crate::runtime::RuntimeCommand;
@@ -65,25 +65,17 @@ impl CellActor {
event_tx,
PendingRuntimeMode::PauseUntilResumed,
)?;
let cancellation_token = session_shutdown_token.child_token();
let termination_token = CancellationToken::new();
let accepting_requests = Arc::new(AtomicBool::new(true));
let handle = CellHandle::new(
command_tx,
cancellation_token.clone(),
termination_token.clone(),
Arc::clone(&accepting_requests),
);
let lifecycle = Arc::new(CellLifecycle::new(session_shutdown_token));
let (terminal_event_tx, terminal_event_rx) = watch::channel(/*init*/ None);
let handle = CellHandle::new(command_tx, Arc::clone(&lifecycle), terminal_event_rx);
let task = run_cell(
host,
CellContext {
runtime_tx,
runtime_control_tx,
runtime_terminate_handle,
cancellation_token,
termination_token,
accepting_requests,
session_shutdown_token,
lifecycle,
terminal_event_tx,
},
event_rx,
command_rx,
@@ -102,10 +94,8 @@ struct CellContext {
runtime_tx: std::sync::mpsc::Sender<RuntimeCommand>,
runtime_control_tx: std::sync::mpsc::Sender<RuntimeControlCommand>,
runtime_terminate_handle: v8::IsolateHandle,
cancellation_token: CancellationToken,
termination_token: CancellationToken,
accepting_requests: Arc<AtomicBool>,
session_shutdown_token: CancellationToken,
lifecycle: Arc<CellLifecycle>,
terminal_event_tx: watch::Sender<Option<CellEvent>>,
}
struct Observer {
@@ -113,10 +103,6 @@ struct Observer {
response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
}
struct Termination {
response_tx: Option<oneshot::Sender<Result<CellEvent, CellError>>>,
}
async fn run_cell<H: CellHost>(
host: Arc<H>,
context: CellContext,
@@ -128,16 +114,17 @@ async fn run_cell<H: CellHost>(
runtime_tx,
runtime_control_tx,
runtime_terminate_handle,
cancellation_token,
termination_token,
accepting_requests,
session_shutdown_token,
lifecycle,
terminal_event_tx,
} = context;
let cancellation_token = lifecycle.work_cancellation_token();
let termination_token = lifecycle.termination_token();
let session_shutdown_token = lifecycle.session_shutdown_token();
let mut content_items = Vec::new();
let mut pending_tool_call_ids = Vec::new();
let mut completed_event = None;
let mut observer = Some(initial_observer);
let mut termination: Option<Termination> = None;
let mut termination = false;
let mut runtime_closed = false;
let mut runtime_paused = false;
let mut yield_timer: Option<std::pin::Pin<Box<tokio::time::Sleep>>> = None;
@@ -167,14 +154,14 @@ async fn run_cell<H: CellHost>(
} => {
let maybe_command = match command_outcome {
Some(command) => command,
None if termination.is_some() => continue,
None => Some(CellCommand::Terminate { response_tx: None }),
None if termination => continue,
None => Some(CellCommand::Terminate),
};
let Some(command) = maybe_command else {
if completed_event.is_some() {
break;
}
termination = Some(Termination { response_tx: None });
termination = true;
begin_termination(
&runtime_tx,
&runtime_control_tx,
@@ -189,9 +176,17 @@ async fn run_cell<H: CellHost>(
match command {
CellCommand::Observe { mode, response_tx } => {
if let Some(event) = completed_event.take() {
accepting_requests.store(false, Ordering::Release);
let _ = response_tx.send(Ok(event));
break;
observer = Some(Observer { mode, response_tx });
if send_or_buffer_completion(
event,
&mut observer,
&mut completed_event,
&lifecycle,
&terminal_event_tx,
) {
break;
}
continue;
}
if observer
.as_ref()
@@ -200,7 +195,7 @@ async fn run_cell<H: CellHost>(
observer = None;
yield_timer = None;
}
if observer.is_some() || termination.is_some() {
if observer.is_some() || termination {
let _ = response_tx.send(Err(CellError::Busy));
continue;
}
@@ -213,21 +208,15 @@ async fn run_cell<H: CellHost>(
&runtime_control_tx,
);
}
CellCommand::Terminate { response_tx } => {
if let Some(event) = completed_event.take() {
accepting_requests.store(false, Ordering::Release);
if let Some(response_tx) = response_tx {
let _ = response_tx.send(Ok(event));
}
CellCommand::Terminate => {
if completed_event.take().is_some() {
lifecycle.close();
break;
}
if termination.is_some() {
if let Some(response_tx) = response_tx {
let _ = response_tx.send(Err(CellError::AlreadyTerminating));
}
if termination {
continue;
}
termination = Some(Termination { response_tx });
termination = true;
yield_timer = None;
begin_termination(
&runtime_tx,
@@ -243,9 +232,9 @@ async fn run_cell<H: CellHost>(
CallbackCompletion::Cancel,
).await;
send_termination_events(
&accepting_requests,
&lifecycle,
&terminal_event_tx,
observer.take(),
termination.take(),
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
@@ -279,7 +268,7 @@ async fn run_cell<H: CellHost>(
}, if !yield_deadline_elapsed => {
let Some(event) = maybe_event else {
runtime_closed = true;
if termination.is_some() || session_shutdown_token.is_cancelled() {
if termination || session_shutdown_token.is_cancelled() {
finish_callbacks(
&cancellation_token,
&mut notification_tasks,
@@ -287,9 +276,9 @@ async fn run_cell<H: CellHost>(
CallbackCompletion::Cancel,
).await;
send_termination_events(
&accepting_requests,
&lifecycle,
&terminal_event_tx,
observer.take(),
termination.take(),
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
@@ -305,12 +294,10 @@ async fn run_cell<H: CellHost>(
)
.await;
if termination_token.is_cancelled() {
let termination =
receive_queued_termination(&mut command_rx).await;
send_termination_events(
&accepting_requests,
&lifecycle,
&terminal_event_tx,
observer.take(),
termination,
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
@@ -319,9 +306,20 @@ async fn run_cell<H: CellHost>(
}
if session_shutdown_token.is_cancelled() {
send_termination_events(
&accepting_requests,
&lifecycle,
&terminal_event_tx,
observer.take(),
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
);
break;
}
if !lifecycle.commit_completion(|| {}) {
send_termination_events(
&lifecycle,
&terminal_event_tx,
observer.take(),
/*termination*/ None,
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
@@ -336,7 +334,8 @@ async fn run_cell<H: CellHost>(
event,
&mut observer,
&mut completed_event,
&accepting_requests,
&lifecycle,
&terminal_event_tx,
) {
break;
}
@@ -414,7 +413,7 @@ async fn run_cell<H: CellHost>(
RuntimeEvent::Result { stored_value_writes, error_text } => {
runtime_closed = true;
yield_timer = None;
if termination.is_some() || session_shutdown_token.is_cancelled() {
if termination || session_shutdown_token.is_cancelled() {
finish_callbacks(
&cancellation_token,
&mut notification_tasks,
@@ -422,9 +421,9 @@ async fn run_cell<H: CellHost>(
CallbackCompletion::Cancel,
).await;
send_termination_events(
&accepting_requests,
&lifecycle,
&terminal_event_tx,
observer.take(),
termination.take(),
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
@@ -439,12 +438,10 @@ async fn run_cell<H: CellHost>(
)
.await;
if termination_token.is_cancelled() {
let termination =
receive_queued_termination(&mut command_rx).await;
send_termination_events(
&accepting_requests,
&lifecycle,
&terminal_event_tx,
observer.take(),
termination,
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
@@ -453,9 +450,9 @@ async fn run_cell<H: CellHost>(
}
if session_shutdown_token.is_cancelled() {
send_termination_events(
&accepting_requests,
&lifecycle,
&terminal_event_tx,
observer.take(),
/*termination*/ None,
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
@@ -466,18 +463,16 @@ async fn run_cell<H: CellHost>(
biased;
_ = termination_token.cancelled() => false,
_ = session_shutdown_token.cancelled() => false,
_ = host.commit_stored_values(stored_value_writes) => true,
committed = host.commit_stored_values(
stored_value_writes,
Arc::clone(&lifecycle),
) => committed,
};
if !committed || session_shutdown_token.is_cancelled() {
let termination = if termination_token.is_cancelled() {
receive_queued_termination(&mut command_rx).await
} else {
None
};
if !committed {
send_termination_events(
&accepting_requests,
&lifecycle,
&terminal_event_tx,
observer.take(),
termination,
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
@@ -492,7 +487,8 @@ async fn run_cell<H: CellHost>(
event,
&mut observer,
&mut completed_event,
&accepting_requests,
&lifecycle,
&terminal_event_tx,
) {
break;
}
@@ -508,7 +504,7 @@ async fn run_cell<H: CellHost>(
}
}
// Reject requests that arrive while asynchronous terminal cleanup runs.
accepting_requests.store(false, Ordering::Release);
lifecycle.close();
drop(command_rx.take());
begin_termination(
&runtime_tx,
@@ -526,41 +522,33 @@ async fn run_cell<H: CellHost>(
host.closed().await;
}
async fn receive_queued_termination(
command_rx: &mut Option<mpsc::UnboundedReceiver<CellCommand>>,
) -> Option<Termination> {
loop {
let command = command_rx.as_mut()?.recv().await?;
match command {
CellCommand::Observe { response_tx, .. } => {
let _ = response_tx.send(Err(CellError::Busy));
}
CellCommand::Terminate { response_tx } => {
return Some(Termination { response_tx });
}
}
}
}
fn send_or_buffer_completion(
event: CellEvent,
observer: &mut Option<Observer>,
completed_event: &mut Option<CellEvent>,
accepting_requests: &AtomicBool,
lifecycle: &CellLifecycle,
terminal_event_tx: &watch::Sender<Option<CellEvent>>,
) -> bool {
if observer
.as_ref()
.is_some_and(|observer| observer.response_tx.is_closed())
{
*observer = None;
}
if observer.is_some() {
accepting_requests.store(false, Ordering::Release);
send_observer_event(observer.take(), event);
true
} else {
terminal_event_tx.send_replace(Some(event.clone()));
let Some(observer) = observer.take() else {
lifecycle.buffer_completion();
*completed_event = Some(event);
false
return false;
};
match observer.response_tx.send(Ok(event)) {
Ok(()) => {
lifecycle.close();
true
}
Err(Ok(event)) => {
lifecycle.buffer_completion();
*completed_event = Some(event);
false
}
Err(Err(error)) => {
lifecycle.close();
panic!("completion delivery unexpectedly carried an actor error: {error:?}");
}
}
}
@@ -571,16 +559,14 @@ fn send_observer_event(observer: Option<Observer>, event: CellEvent) {
}
fn send_termination_events(
accepting_requests: &AtomicBool,
lifecycle: &CellLifecycle,
terminal_event_tx: &watch::Sender<Option<CellEvent>>,
observer: Option<Observer>,
termination: Option<Termination>,
event: CellEvent,
) {
accepting_requests.store(false, Ordering::Release);
send_observer_event(observer, event.clone());
if let Some(response_tx) = termination.and_then(|termination| termination.response_tx) {
let _ = response_tx.send(Ok(event));
}
lifecycle.close();
terminal_event_tx.send_replace(Some(event.clone()));
send_observer_event(observer, event);
}
fn observer_timer(observer: &Observer) -> Option<std::pin::Pin<Box<tokio::time::Sleep>>> {
codex-rs/code-mode/src/cell_actor/tests.rs
+174 -22
View File
@@ -1,5 +1,6 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Barrier;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::thread;
@@ -36,6 +37,12 @@ struct BlockingCommitHost {
commit_release: Semaphore,
}
struct TerminationCommitRaceHost {
commit_polled: Barrier,
termination_admitted: Barrier,
committed: AtomicBool,
}
impl CellHost for TestHost {
async fn invoke_tool(
&self,
@@ -54,7 +61,13 @@ impl CellHost for TestHost {
Ok(())
}
async fn commit_stored_values(&self, _stored_value_writes: HashMap<String, JsonValue>) {}
async fn commit_stored_values(
&self,
_stored_value_writes: HashMap<String, JsonValue>,
lifecycle: Arc<CellLifecycle>,
) -> bool {
lifecycle.commit_completion(|| {})
}
async fn closed(&self) {}
}
@@ -77,8 +90,14 @@ impl CellHost for RecordingHost {
Ok(())
}
async fn commit_stored_values(&self, _stored_value_writes: HashMap<String, JsonValue>) {
self.committed.store(true, Ordering::Release);
async fn commit_stored_values(
&self,
_stored_value_writes: HashMap<String, JsonValue>,
lifecycle: Arc<CellLifecycle>,
) -> bool {
lifecycle.commit_completion(|| {
self.committed.store(true, Ordering::Release);
})
}
async fn closed(&self) {
@@ -104,8 +123,14 @@ impl CellHost for PendingNotificationHost {
std::future::pending().await
}
async fn commit_stored_values(&self, _stored_value_writes: HashMap<String, JsonValue>) {
self.committed.store(true, Ordering::Release);
async fn commit_stored_values(
&self,
_stored_value_writes: HashMap<String, JsonValue>,
lifecycle: Arc<CellLifecycle>,
) -> bool {
lifecycle.commit_completion(|| {
self.committed.store(true, Ordering::Release);
})
}
async fn closed(&self) {
@@ -144,7 +169,11 @@ impl CellHost for BlockingCommitHost {
Ok(())
}
async fn commit_stored_values(&self, _stored_value_writes: HashMap<String, JsonValue>) {
async fn commit_stored_values(
&self,
_stored_value_writes: HashMap<String, JsonValue>,
lifecycle: Arc<CellLifecycle>,
) -> bool {
self.commit_started_tx
.send(())
.expect("test did not receive commit start");
@@ -153,6 +182,50 @@ impl CellHost for BlockingCommitHost {
.await
.expect("test did not release commit")
.forget();
lifecycle.commit_completion(|| {})
}
async fn closed(&self) {}
}
impl TerminationCommitRaceHost {
fn new() -> Arc<Self> {
Arc::new(Self {
commit_polled: Barrier::new(/*n*/ 2),
termination_admitted: Barrier::new(/*n*/ 2),
committed: AtomicBool::new(false),
})
}
}
impl CellHost for TerminationCommitRaceHost {
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>,
lifecycle: Arc<CellLifecycle>,
) -> bool {
self.commit_polled.wait();
self.termination_admitted.wait();
lifecycle.commit_completion(|| {
self.committed.store(true, Ordering::Release);
})
}
async fn closed(&self) {}
@@ -195,14 +268,12 @@ fn spawn_cell_actor_harness_with_host<H: CellHost>(
)
.unwrap();
let session_shutdown_token = CancellationToken::new();
let cancellation_token = session_shutdown_token.child_token();
let termination_token = CancellationToken::new();
let accepting_requests = Arc::new(AtomicBool::new(true));
let lifecycle = Arc::new(CellLifecycle::new(session_shutdown_token.clone()));
let (terminal_event_tx, terminal_event_rx) = watch::channel(/*init*/ None);
let handle = CellHandle::new(
command_tx.clone(),
cancellation_token.clone(),
termination_token.clone(),
Arc::clone(&accepting_requests),
Arc::clone(&lifecycle),
terminal_event_rx,
);
let task = tokio::spawn(run_cell(
host,
@@ -210,10 +281,8 @@ fn spawn_cell_actor_harness_with_host<H: CellHost>(
runtime_tx,
runtime_control_tx,
runtime_terminate_handle,
cancellation_token,
termination_token,
accepting_requests,
session_shutdown_token: session_shutdown_token.clone(),
lifecycle,
terminal_event_tx,
},
event_rx,
command_rx,
@@ -297,6 +366,34 @@ async fn termination_preempts_result_without_committing_stored_values() {
assert!(host.closed.load(Ordering::Acquire));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn admitted_termination_wins_during_the_commit_poll() {
let host = TerminationCommitRaceHost::new();
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();
host.commit_polled.wait();
let termination = harness.handle.terminate();
host.termination_admitted.wait();
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));
}
#[tokio::test]
async fn session_shutdown_during_result_callback_drain_prevents_stored_value_commit() {
let host = Arc::new(PendingNotificationHost::default());
@@ -319,18 +416,16 @@ async fn session_shutdown_during_result_callback_drain_prevents_stored_value_com
)
.unwrap();
let session_shutdown_token = CancellationToken::new();
let termination_token = CancellationToken::new();
let accepting_requests = Arc::new(AtomicBool::new(true));
let lifecycle = Arc::new(CellLifecycle::new(session_shutdown_token.clone()));
let (terminal_event_tx, _) = watch::channel(/*init*/ None);
let mut task = Box::pin(run_cell(
Arc::clone(&host),
CellContext {
runtime_tx,
runtime_control_tx,
runtime_terminate_handle,
cancellation_token: session_shutdown_token.child_token(),
termination_token,
accepting_requests,
session_shutdown_token: session_shutdown_token.clone(),
lifecycle,
terminal_event_tx,
},
event_rx,
command_rx,
@@ -494,6 +589,63 @@ async fn observer_receives_a_buffered_completion_after_commit_finishes() {
harness.task.await.unwrap();
}
#[test]
fn failed_completion_delivery_rebuffers_the_event() {
let session_shutdown_token = CancellationToken::new();
let lifecycle = CellLifecycle::new(session_shutdown_token);
assert!(lifecycle.commit_completion(|| {}));
let (terminal_event_tx, terminal_event_rx) = watch::channel(/*init*/ None);
let (response_tx, response_rx) = oneshot::channel();
drop(response_rx);
let mut observer = Some(Observer {
mode: ObserveMode::PendingFrontier,
response_tx,
});
let event = CellEvent::Completed {
content_items: Vec::new(),
error_text: None,
};
let mut completed_event = None;
let delivered = send_or_buffer_completion(
event.clone(),
&mut observer,
&mut completed_event,
&lifecycle,
&terminal_event_tx,
);
assert_eq!(
(
delivered,
observer.is_none(),
completed_event,
terminal_event_rx.borrow().clone(),
lifecycle.accepting_requests(),
),
(false, true, Some(event.clone()), Some(event), true),
);
}
#[tokio::test]
async fn session_shutdown_rejects_observation_before_the_actor_closes() {
let session_shutdown_token = CancellationToken::new();
let lifecycle = Arc::new(CellLifecycle::new(session_shutdown_token.clone()));
let (command_tx, mut command_rx) = mpsc::unbounded_channel();
let (_terminal_event_tx, terminal_event_rx) = watch::channel(/*init*/ None);
let handle = CellHandle::new(command_tx, lifecycle, terminal_event_rx);
session_shutdown_token.cancel();
assert_eq!(
handle.observe(ObserveMode::PendingFrontier).await,
Err(CellError::Closed)
);
assert!(matches!(
command_rx.try_recv(),
Err(mpsc::error::TryRecvError::Empty)
));
}
#[tokio::test]
async fn dropped_observer_does_not_block_the_next_observer() {
let (host, mut commit_started_rx) = BlockingCommitHost::new();
codex-rs/code-mode/src/cell_actor/types.rs
+153 -41
View File
@@ -2,12 +2,12 @@ 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::Ordering;
use std::sync::Mutex;
use serde_json::Value as JsonValue;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use crate::session_runtime::CellEvent;
@@ -55,7 +55,8 @@ pub(crate) trait CellHost: Send + Sync + 'static {
fn commit_stored_values(
&self,
stored_value_writes: HashMap<String, JsonValue>,
) -> impl Future<Output = ()> + Send;
lifecycle: Arc<CellLifecycle>,
) -> impl Future<Output = bool> + Send;
fn closed(&self) -> impl Future<Output = ()> + Send;
}
@@ -63,33 +64,28 @@ pub(crate) trait CellHost: Send + Sync + 'static {
#[derive(Clone)]
pub(crate) struct CellHandle {
command_tx: mpsc::UnboundedSender<CellCommand>,
cancellation_token: CancellationToken,
termination_token: CancellationToken,
termination_requested: Arc<AtomicBool>,
accepting_requests: Arc<AtomicBool>,
lifecycle: Arc<CellLifecycle>,
terminal_event_rx: watch::Receiver<Option<CellEvent>>,
}
impl CellHandle {
pub(super) fn new(
command_tx: mpsc::UnboundedSender<CellCommand>,
cancellation_token: CancellationToken,
termination_token: CancellationToken,
accepting_requests: Arc<AtomicBool>,
lifecycle: Arc<CellLifecycle>,
terminal_event_rx: watch::Receiver<Option<CellEvent>>,
) -> Self {
Self {
command_tx,
cancellation_token,
termination_token,
termination_requested: Arc::new(AtomicBool::new(false)),
accepting_requests,
lifecycle,
terminal_event_rx,
}
}
pub(crate) fn observe(&self, mode: ObserveMode) -> CellEventFuture {
if !self.accepting_requests.load(Ordering::Acquire) {
if !self.lifecycle.accepting_requests() {
return closed_event();
}
let (response_tx, response_rx) = oneshot::channel();
let (response_tx, response_rx) = tokio::sync::oneshot::channel();
if self
.command_tx
.send(CellCommand::Observe { mode, response_tx })
@@ -101,31 +97,134 @@ impl CellHandle {
}
pub(crate) fn terminate(&self) -> CellEventFuture {
if !self.accepting_requests.load(Ordering::Acquire) {
match self.lifecycle.request_termination() {
TerminationRequest::Terminate => {}
TerminationRequest::ConsumeCompletion => {}
TerminationRequest::AlreadyTerminating => {
return Box::pin(async { Err(CellError::AlreadyTerminating) });
}
TerminationRequest::Closed => return closed_event(),
}
if self.command_tx.send(CellCommand::Terminate).is_err() {
self.lifecycle.close();
return closed_event();
}
if self
.termination_requested
.compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
.is_err()
{
return Box::pin(async { Err(CellError::AlreadyTerminating) });
terminal_event(self.terminal_event_rx.clone())
}
}
pub(crate) struct CellLifecycle {
state: Mutex<CellLifecycleState>,
session_shutdown_token: CancellationToken,
termination_token: CancellationToken,
work_cancellation_token: CancellationToken,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum CellLifecycleState {
Running,
TerminationRequested,
CompletionCommitted,
CompletionBuffered,
CompletionConsumptionRequested,
Closed,
}
enum TerminationRequest {
Terminate,
ConsumeCompletion,
AlreadyTerminating,
Closed,
}
impl CellLifecycle {
pub(crate) fn new(session_shutdown_token: CancellationToken) -> Self {
let work_cancellation_token = session_shutdown_token.child_token();
Self {
state: Mutex::new(CellLifecycleState::Running),
session_shutdown_token,
termination_token: CancellationToken::new(),
work_cancellation_token,
}
let (response_tx, response_rx) = oneshot::channel();
if self
.command_tx
.send(CellCommand::Terminate {
response_tx: Some(response_tx),
})
.is_err()
{
self.termination_requested.store(false, Ordering::Relaxed);
return closed_event();
}
pub(crate) fn accepting_requests(&self) -> bool {
let accepting_state = matches!(
*self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner),
CellLifecycleState::Running
| CellLifecycleState::CompletionCommitted
| CellLifecycleState::CompletionBuffered
);
accepting_state && !self.session_shutdown_token.is_cancelled()
}
fn request_termination(&self) -> TerminationRequest {
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
match *state {
CellLifecycleState::Running => {
*state = CellLifecycleState::TerminationRequested;
self.termination_token.cancel();
self.work_cancellation_token.cancel();
TerminationRequest::Terminate
}
CellLifecycleState::TerminationRequested => TerminationRequest::AlreadyTerminating,
CellLifecycleState::CompletionCommitted | CellLifecycleState::CompletionBuffered => {
*state = CellLifecycleState::CompletionConsumptionRequested;
TerminationRequest::ConsumeCompletion
}
CellLifecycleState::CompletionConsumptionRequested => {
TerminationRequest::AlreadyTerminating
}
CellLifecycleState::Closed => TerminationRequest::Closed,
}
// The command carries the terminal response; these tokens interrupt in-flight cleanup.
self.termination_token.cancel();
self.cancellation_token.cancel();
response_event(response_rx)
}
pub(crate) fn commit_completion(&self, commit: impl FnOnce()) -> bool {
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if *state != CellLifecycleState::Running || self.session_shutdown_token.is_cancelled() {
return false;
}
commit();
*state = CellLifecycleState::CompletionCommitted;
true
}
pub(crate) fn buffer_completion(&self) {
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if *state == CellLifecycleState::CompletionCommitted {
*state = CellLifecycleState::CompletionBuffered;
}
}
pub(crate) fn close(&self) {
*self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = CellLifecycleState::Closed;
}
pub(crate) fn session_shutdown_token(&self) -> CancellationToken {
self.session_shutdown_token.clone()
}
pub(crate) fn termination_token(&self) -> CancellationToken {
self.termination_token.clone()
}
pub(crate) fn work_cancellation_token(&self) -> CancellationToken {
self.work_cancellation_token.clone()
}
}
@@ -134,15 +233,28 @@ pub(super) enum CellCommand {
mode: ObserveMode,
response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
},
Terminate {
response_tx: Option<oneshot::Sender<Result<CellEvent, CellError>>>,
},
Terminate,
}
fn response_event(response_rx: oneshot::Receiver<Result<CellEvent, CellError>>) -> CellEventFuture {
fn response_event(
response_rx: tokio::sync::oneshot::Receiver<Result<CellEvent, CellError>>,
) -> CellEventFuture {
Box::pin(async move { response_rx.await.unwrap_or(Err(CellError::Closed)) })
}
fn terminal_event(mut terminal_event_rx: watch::Receiver<Option<CellEvent>>) -> CellEventFuture {
Box::pin(async move {
loop {
if let Some(event) = terminal_event_rx.borrow_and_update().clone() {
return Ok(event);
}
if terminal_event_rx.changed().await.is_err() {
return Err(CellError::Closed);
}
}
})
}
fn closed_event() -> CellEventFuture {
Box::pin(async { Err(CellError::Closed) })
}
codex-rs/code-mode/src/session_runtime/mod.rs
+10 -6
View File
@@ -27,6 +27,7 @@ use crate::cell_actor::CellError;
use crate::cell_actor::CellEventFuture;
use crate::cell_actor::CellHandle;
use crate::cell_actor::CellHost;
use crate::cell_actor::CellLifecycle;
use crate::cell_actor::CellToolCall;
type RuntimeEventFuture = Pin<Box<dyn Future<Output = Result<CellEvent, Error>> + Send + 'static>>;
@@ -215,12 +216,15 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
.await
}
async fn commit_stored_values(&self, stored_value_writes: HashMap<String, JsonValue>) {
self.inner
.stored_values
.lock()
.await
.extend(stored_value_writes);
async fn commit_stored_values(
&self,
stored_value_writes: HashMap<String, JsonValue>,
lifecycle: Arc<CellLifecycle>,
) -> bool {
let mut stored_values = self.inner.stored_values.lock().await;
lifecycle.commit_completion(|| {
stored_values.extend(stored_value_writes);
})
}
async fn closed(&self) {
codex-rs/core/src/tools/code_mode/delegate.rs
+35 -13
View File
@@ -1,3 +1,4 @@
use std::future::Future;
use std::sync::Arc;
use codex_code_mode::CellId;
@@ -64,13 +65,11 @@ impl CodeModeDispatchBroker {
cancellation_token,
response_tx,
} => {
let response = tokio::select! {
biased;
_ = cancellation_token.cancelled() => {
Err("code mode notification cancelled".to_string())
}
response = host.notify(call_id, cell_id, text) => response,
};
let response = dispatch_unless_cancelled(&cancellation_token, || {
host.notify(call_id, cell_id, text)
})
.await
.unwrap_or_else(|| Err("code mode notification cancelled".to_string()));
let _ = response_tx.send(response);
}
DispatchMessage::InvokeTool {
@@ -78,14 +77,18 @@ impl CodeModeDispatchBroker {
cancellation_token,
response_tx,
} => {
if cancellation_token.is_cancelled() {
continue;
}
let host = Arc::clone(&host);
tokio::spawn(async move {
let response = tokio::select! {
response = host.invoke_tool(
invocation,
cancellation_token.clone(),
) => response,
_ = cancellation_token.cancelled() => return,
let Some(response) =
dispatch_unless_cancelled(&cancellation_token, || {
host.invoke_tool(invocation, cancellation_token.clone())
})
.await
else {
return;
};
let _ = response_tx.send(response);
});
@@ -99,6 +102,25 @@ impl CodeModeDispatchBroker {
}
}
async fn dispatch_unless_cancelled<T, F, Fut>(
cancellation_token: &CancellationToken,
dispatch: F,
) -> Option<T>
where
F: FnOnce() -> Fut,
Fut: Future<Output = T>,
{
if cancellation_token.is_cancelled() {
return None;
}
let dispatch = dispatch();
tokio::select! {
biased;
_ = cancellation_token.cancelled() => None,
response = dispatch => Some(response),
}
}
impl CodeModeSessionDelegate for CodeModeDispatchBroker {
fn invoke_tool<'a>(
&'a self,
codex-rs/core/src/tools/code_mode/delegate_tests.rs
+39
View File
@@ -1,4 +1,6 @@
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use codex_code_mode::CellId;
@@ -9,12 +11,49 @@ use tokio_util::sync::CancellationToken;
use super::super::ExecContext;
use super::CodeModeDispatchBroker;
use super::dispatch_unless_cancelled;
use crate::session::tests::make_session_and_context;
use crate::state::ActiveTurn;
use crate::tools::registry::ToolRegistry;
use crate::tools::router::ToolRouter;
use crate::turn_diff_tracker::TurnDiffTracker;
#[tokio::test]
async fn cancelled_dispatch_is_rejected_before_construction() {
let cancellation_token = CancellationToken::new();
cancellation_token.cancel();
let constructed = AtomicBool::new(false);
let response = dispatch_unless_cancelled(&cancellation_token, || {
constructed.store(true, Ordering::Release);
async { "unexpected response" }
})
.await;
assert_eq!(response, None);
assert!(!constructed.load(Ordering::Acquire));
}
#[tokio::test]
async fn cancellation_before_dispatch_poll_has_priority() {
let cancellation_token = CancellationToken::new();
let cancellation_for_dispatch = cancellation_token.clone();
let polled = Arc::new(AtomicBool::new(false));
let polled_by_dispatch = Arc::clone(&polled);
let response = dispatch_unless_cancelled(&cancellation_token, || {
cancellation_for_dispatch.cancel();
async move {
polled_by_dispatch.store(true, Ordering::Release);
"unexpected response"
}
})
.await;
assert_eq!(response, None);
assert!(!polled.load(Ordering::Acquire));
}
#[tokio::test]
#[expect(
clippy::await_holding_invalid_type,