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

Simplify code mode session runtime lifecycle

Browse Source
This commit is contained in:
Channing Conger
2026-06-19 23:43:56 +00:00
parent 641452f668
commit 2c497327f4
5 changed files with 405 additions and 479 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode/src/cell_actor/mod.rs
+145 -270
View File
@@ -9,7 +9,6 @@ use std::sync::Arc;
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;
@@ -26,8 +25,10 @@ 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::CellState;
pub(crate) use self::types::CellToolCall;
use self::types::CompletionDelivery;
use self::types::ObservationDelivery;
use crate::runtime::PendingRuntimeMode;
use crate::runtime::RuntimeCommand;
use crate::runtime::RuntimeControlCommand;
@@ -47,7 +48,7 @@ impl CellActor {
stored_values: HashMap<String, JsonValue>,
host: Arc<H>,
initial_observe_mode: ObserveMode,
session_shutdown_token: CancellationToken,
cell_state: Arc<CellState>,
) -> Result<
(
CellHandle,
@@ -65,17 +66,14 @@ impl CellActor {
event_tx,
PendingRuntimeMode::PauseUntilResumed,
)?;
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 handle = CellHandle::new(command_tx, Arc::clone(&cell_state));
let task = run_cell(
host,
CellContext {
runtime_tx,
runtime_control_tx,
runtime_terminate_handle,
lifecycle,
terminal_event_tx,
cell_state,
},
event_rx,
command_rx,
@@ -94,8 +92,7 @@ struct CellContext {
runtime_tx: std::sync::mpsc::Sender<RuntimeCommand>,
runtime_control_tx: std::sync::mpsc::Sender<RuntimeControlCommand>,
runtime_terminate_handle: v8::IsolateHandle,
lifecycle: Arc<CellLifecycle>,
terminal_event_tx: watch::Sender<Option<CellEvent>>,
cell_state: Arc<CellState>,
}
struct Observer {
@@ -114,15 +111,12 @@ async fn run_cell<H: CellHost>(
runtime_tx,
runtime_control_tx,
runtime_terminate_handle,
lifecycle,
terminal_event_tx,
cell_state,
} = context;
let cancellation_token = lifecycle.work_cancellation_token();
let termination_token = lifecycle.termination_token();
let session_shutdown_token = lifecycle.session_shutdown_token();
let cancellation_token = cell_state.cancellation_token();
let callback_cancellation_token = cancellation_token.child_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 = false;
let mut runtime_closed = false;
@@ -137,112 +131,67 @@ async fn run_cell<H: CellHost>(
.is_some_and(|yield_timer| yield_timer.deadline() <= tokio::time::Instant::now());
tokio::select! {
biased;
command_outcome = async {
tokio::select! {
biased;
_ = session_shutdown_token.cancelled(), if command_rx.is_some() => {
drop(command_rx.take());
None
}
command = async {
match command_rx.as_mut() {
Some(command_rx) => command_rx.recv().await,
None => std::future::pending::<Option<CellCommand>>().await,
}
} => Some(command),
_ = cancellation_token.cancelled(), if !termination => {
termination = true;
yield_timer = None;
drop(command_rx.take());
begin_termination(
&runtime_tx,
&runtime_control_tx,
&runtime_terminate_handle,
&cancellation_token,
);
if runtime_closed {
finish_callbacks(
&callback_cancellation_token,
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::Cancel,
).await;
finish_termination(
&cell_state,
observer.take().map(|observer| observer.response_tx),
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
);
break;
}
}
maybe_command = async {
match command_rx.as_mut() {
Some(command_rx) => command_rx.recv().await,
None => std::future::pending::<Option<CellCommand>>().await,
}
} => {
let maybe_command = match command_outcome {
Some(command) => command,
None if termination => continue,
None => Some(CellCommand::Terminate),
};
let Some(command) = maybe_command else {
if completed_event.is_some() {
break;
}
termination = true;
begin_termination(
&runtime_tx,
&runtime_control_tx,
&runtime_terminate_handle,
&cancellation_token,
);
if runtime_closed {
break;
}
let Some(CellCommand::Observe { mode, response_tx }) = maybe_command else {
cancellation_token.cancel();
continue;
};
match command {
CellCommand::Observe { mode, response_tx } => {
if let Some(event) = completed_event.take() {
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()
.is_some_and(|observer| observer.response_tx.is_closed())
{
observer = None;
yield_timer = None;
}
if observer.is_some() || termination {
let _ = response_tx.send(Err(CellError::Busy));
continue;
}
observer = Some(Observer { mode, response_tx });
yield_timer = observer.as_ref().and_then(observer_timer);
resume_for_observation(
mode,
&mut runtime_paused,
&runtime_tx,
&runtime_control_tx,
);
}
CellCommand::Terminate => {
if completed_event.take().is_some() {
lifecycle.close();
break;
}
if termination {
continue;
}
termination = true;
yield_timer = None;
begin_termination(
&runtime_tx,
&runtime_control_tx,
&runtime_terminate_handle,
&cancellation_token,
);
if runtime_closed {
finish_callbacks(
&cancellation_token,
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::Cancel,
).await;
send_termination_events(
&lifecycle,
&terminal_event_tx,
observer.take(),
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
);
break;
}
}
let response_tx = match cell_state.route_observation(response_tx) {
ObservationDelivery::Running(response_tx) => response_tx,
ObservationDelivery::Delivered => break,
ObservationDelivery::Buffered | ObservationDelivery::Closed => continue,
};
if observer
.as_ref()
.is_some_and(|observer| observer.response_tx.is_closed())
{
observer = None;
yield_timer = None;
}
if observer.is_some() || termination {
let _ = response_tx.send(Err(CellError::Busy));
continue;
}
observer = Some(Observer { mode, response_tx });
yield_timer = observer.as_ref().and_then(observer_timer);
resume_for_observation(
mode,
&mut runtime_paused,
&runtime_tx,
&runtime_control_tx,
);
}
_ = async {
if let Some(yield_timer) = yield_timer.as_mut() {
@@ -268,77 +217,58 @@ async fn run_cell<H: CellHost>(
}, if !yield_deadline_elapsed => {
let Some(event) = maybe_event else {
runtime_closed = true;
if termination || session_shutdown_token.is_cancelled() {
if termination || cancellation_token.is_cancelled() {
finish_callbacks(
&cancellation_token,
&callback_cancellation_token,
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::Cancel,
).await;
send_termination_events(
&lifecycle,
&terminal_event_tx,
observer.take(),
finish_termination(
&cell_state,
observer.take().map(|observer| observer.response_tx),
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
);
break;
}
if completed_event.is_none() {
finish_callbacks(
&cancellation_token,
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::DrainNotifications,
)
.await;
if termination_token.is_cancelled() {
send_termination_events(
&lifecycle,
&terminal_event_tx,
observer.take(),
finish_callbacks(
&callback_cancellation_token,
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::DrainNotifications,
)
.await;
let completion_content_items = std::mem::take(&mut content_items);
let event = CellEvent::Completed {
content_items: completion_content_items.clone(),
error_text: Some("exec runtime ended unexpectedly".to_string()),
};
let _ = tokio::select! {
biased;
_ = cancellation_token.cancelled() => false,
committed = host.commit_completion(
HashMap::new(),
event.clone(),
Arc::clone(&cell_state),
) => committed,
};
match cell_state.deliver_completion(
observer.take().map(|observer| observer.response_tx),
) {
CompletionDelivery::Delivered => break,
CompletionDelivery::Buffered => {}
CompletionDelivery::Rejected(response_tx) => {
finish_termination(
&cell_state,
response_tx,
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
content_items: completion_content_items,
},
);
break;
}
if session_shutdown_token.is_cancelled() {
send_termination_events(
&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(),
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
);
break;
}
let event = CellEvent::Completed {
content_items: std::mem::take(&mut content_items),
error_text: Some("exec runtime ended unexpectedly".to_string()),
};
if send_or_buffer_completion(
event,
&mut observer,
&mut completed_event,
&lifecycle,
&terminal_event_tx,
) {
break;
}
}
continue;
};
@@ -389,7 +319,7 @@ async fn run_cell<H: CellHost>(
Arc::clone(&host),
call_id,
text,
cancellation_token.child_token(),
callback_cancellation_token.child_token(),
);
}
RuntimeEvent::ToolCall { id, name, kind, input } => {
@@ -407,23 +337,22 @@ async fn run_cell<H: CellHost>(
input,
},
runtime_tx.clone(),
cancellation_token.child_token(),
callback_cancellation_token.child_token(),
);
}
RuntimeEvent::Result { stored_value_writes, error_text } => {
runtime_closed = true;
yield_timer = None;
if termination || session_shutdown_token.is_cancelled() {
if termination || cancellation_token.is_cancelled() {
finish_callbacks(
&cancellation_token,
&callback_cancellation_token,
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::Cancel,
).await;
send_termination_events(
&lifecycle,
&terminal_event_tx,
observer.take(),
finish_termination(
&cell_state,
observer.take().map(|observer| observer.response_tx),
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
@@ -431,66 +360,41 @@ async fn run_cell<H: CellHost>(
break;
}
finish_callbacks(
&cancellation_token,
&callback_cancellation_token,
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::DrainNotifications,
)
.await;
if termination_token.is_cancelled() {
send_termination_events(
&lifecycle,
&terminal_event_tx,
observer.take(),
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
);
break;
}
if session_shutdown_token.is_cancelled() {
send_termination_events(
&lifecycle,
&terminal_event_tx,
observer.take(),
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
);
break;
}
let committed = tokio::select! {
biased;
_ = termination_token.cancelled() => false,
_ = session_shutdown_token.cancelled() => false,
committed = host.commit_stored_values(
stored_value_writes,
Arc::clone(&lifecycle),
) => committed,
};
if !committed {
send_termination_events(
&lifecycle,
&terminal_event_tx,
observer.take(),
CellEvent::Terminated {
content_items: std::mem::take(&mut content_items),
},
);
break;
}
let completion_content_items = std::mem::take(&mut content_items);
let event = CellEvent::Completed {
content_items: std::mem::take(&mut content_items),
content_items: completion_content_items.clone(),
error_text,
};
if send_or_buffer_completion(
event,
&mut observer,
&mut completed_event,
&lifecycle,
&terminal_event_tx,
let _ = tokio::select! {
biased;
_ = cancellation_token.cancelled() => false,
committed = host.commit_completion(
stored_value_writes,
event.clone(),
Arc::clone(&cell_state),
) => committed,
};
match cell_state.deliver_completion(
observer.take().map(|observer| observer.response_tx),
) {
break;
CompletionDelivery::Delivered => break,
CompletionDelivery::Buffered => {}
CompletionDelivery::Rejected(response_tx) => {
finish_termination(
&cell_state,
response_tx,
CellEvent::Terminated {
content_items: completion_content_items,
},
);
break;
}
}
}
}
@@ -504,7 +408,7 @@ async fn run_cell<H: CellHost>(
}
}
// Reject requests that arrive while asynchronous terminal cleanup runs.
lifecycle.close();
cell_state.tombstone();
drop(command_rx.take());
begin_termination(
&runtime_tx,
@@ -513,7 +417,7 @@ async fn run_cell<H: CellHost>(
&cancellation_token,
);
finish_callbacks(
&cancellation_token,
&callback_cancellation_token,
&mut notification_tasks,
&mut tool_tasks,
CallbackCompletion::Cancel,
@@ -522,51 +426,22 @@ async fn run_cell<H: CellHost>(
host.closed().await;
}
fn send_or_buffer_completion(
event: CellEvent,
observer: &mut Option<Observer>,
completed_event: &mut Option<CellEvent>,
lifecycle: &CellLifecycle,
terminal_event_tx: &watch::Sender<Option<CellEvent>>,
) -> bool {
terminal_event_tx.send_replace(Some(event.clone()));
let Some(observer) = observer.take() else {
lifecycle.buffer_completion();
*completed_event = Some(event);
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:?}");
}
}
}
fn send_observer_event(observer: Option<Observer>, event: CellEvent) {
if let Some(observer) = observer {
let _ = observer.response_tx.send(Ok(event));
}
}
fn send_termination_events(
lifecycle: &CellLifecycle,
terminal_event_tx: &watch::Sender<Option<CellEvent>>,
observer: Option<Observer>,
fn finish_termination(
cell_state: &CellState,
observer_tx: Option<oneshot::Sender<Result<CellEvent, CellError>>>,
event: CellEvent,
) {
lifecycle.close();
terminal_event_tx.send_replace(Some(event.clone()));
send_observer_event(observer, event);
if let Some(event) = cell_state.finish_termination(event)
&& let Some(observer_tx) = observer_tx
{
let _ = observer_tx.send(Ok(event));
}
}
fn observer_timer(observer: &Observer) -> Option<std::pin::Pin<Box<tokio::time::Sleep>>> {
codex-rs/code-mode/src/cell_actor/tests.rs
+42 -59
View File
@@ -61,12 +61,13 @@ impl CellHost for TestHost {
Ok(())
}
async fn commit_stored_values(
async fn commit_completion(
&self,
_stored_value_writes: HashMap<String, JsonValue>,
lifecycle: Arc<CellLifecycle>,
event: CellEvent,
cell_state: Arc<CellState>,
) -> bool {
lifecycle.commit_completion(|| {})
cell_state.commit_completion(event, || {})
}
async fn closed(&self) {}
@@ -90,12 +91,13 @@ impl CellHost for RecordingHost {
Ok(())
}
async fn commit_stored_values(
async fn commit_completion(
&self,
_stored_value_writes: HashMap<String, JsonValue>,
lifecycle: Arc<CellLifecycle>,
event: CellEvent,
cell_state: Arc<CellState>,
) -> bool {
lifecycle.commit_completion(|| {
cell_state.commit_completion(event, || {
self.committed.store(true, Ordering::Release);
})
}
@@ -123,12 +125,13 @@ impl CellHost for PendingNotificationHost {
std::future::pending().await
}
async fn commit_stored_values(
async fn commit_completion(
&self,
_stored_value_writes: HashMap<String, JsonValue>,
lifecycle: Arc<CellLifecycle>,
event: CellEvent,
cell_state: Arc<CellState>,
) -> bool {
lifecycle.commit_completion(|| {
cell_state.commit_completion(event, || {
self.committed.store(true, Ordering::Release);
})
}
@@ -169,10 +172,11 @@ impl CellHost for BlockingCommitHost {
Ok(())
}
async fn commit_stored_values(
async fn commit_completion(
&self,
_stored_value_writes: HashMap<String, JsonValue>,
lifecycle: Arc<CellLifecycle>,
event: CellEvent,
cell_state: Arc<CellState>,
) -> bool {
self.commit_started_tx
.send(())
@@ -182,7 +186,7 @@ impl CellHost for BlockingCommitHost {
.await
.expect("test did not release commit")
.forget();
lifecycle.commit_completion(|| {})
cell_state.commit_completion(event, || {})
}
async fn closed(&self) {}
@@ -216,14 +220,15 @@ impl CellHost for TerminationCommitRaceHost {
Ok(())
}
async fn commit_stored_values(
async fn commit_completion(
&self,
_stored_value_writes: HashMap<String, JsonValue>,
lifecycle: Arc<CellLifecycle>,
event: CellEvent,
cell_state: Arc<CellState>,
) -> bool {
self.commit_polled.wait();
self.termination_admitted.wait();
lifecycle.commit_completion(|| {
cell_state.commit_completion(event, || {
self.committed.store(true, Ordering::Release);
})
}
@@ -268,21 +273,15 @@ fn spawn_cell_actor_harness_with_host<H: CellHost>(
)
.unwrap();
let session_shutdown_token = CancellationToken::new();
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(),
Arc::clone(&lifecycle),
terminal_event_rx,
);
let cell_state = Arc::new(CellState::new(session_shutdown_token.child_token()));
let handle = CellHandle::new(command_tx.clone(), Arc::clone(&cell_state));
let task = tokio::spawn(run_cell(
host,
CellContext {
runtime_tx,
runtime_control_tx,
runtime_terminate_handle,
lifecycle,
terminal_event_tx,
cell_state,
},
event_rx,
command_rx,
@@ -416,16 +415,14 @@ async fn session_shutdown_during_result_callback_drain_prevents_stored_value_com
)
.unwrap();
let session_shutdown_token = CancellationToken::new();
let lifecycle = Arc::new(CellLifecycle::new(session_shutdown_token.clone()));
let (terminal_event_tx, _) = watch::channel(/*init*/ None);
let cell_state = Arc::new(CellState::new(session_shutdown_token.child_token()));
let mut task = Box::pin(run_cell(
Arc::clone(&host),
CellContext {
runtime_tx,
runtime_control_tx,
runtime_terminate_handle,
lifecycle,
terminal_event_tx,
cell_state,
},
event_rx,
command_rx,
@@ -592,48 +589,34 @@ async fn observer_receives_a_buffered_completion_after_commit_finishes() {
#[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 cell_state = CellState::new(session_shutdown_token.child_token());
let event = CellEvent::Completed {
content_items: Vec::new(),
error_text: None,
};
let mut completed_event = None;
assert!(cell_state.commit_completion(event.clone(), || {}));
let (response_tx, response_rx) = oneshot::channel();
drop(response_rx);
assert!(matches!(
cell_state.deliver_completion(Some(response_tx)),
CompletionDelivery::Buffered
));
assert!(cell_state.accepting_observations());
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),
);
let (response_tx, mut response_rx) = oneshot::channel();
assert!(matches!(
cell_state.route_observation(response_tx),
ObservationDelivery::Delivered
));
assert_eq!(response_rx.try_recv(), Ok(Ok(event)));
}
#[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 cell_state = Arc::new(CellState::new(session_shutdown_token.child_token()));
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);
let handle = CellHandle::new(command_tx, cell_state);
session_shutdown_token.cancel();
assert_eq!(
codex-rs/code-mode/src/cell_actor/types.rs
+166 -111
View File
@@ -7,7 +7,6 @@ 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;
@@ -32,7 +31,7 @@ pub(crate) struct CellToolCall {
pub(crate) input: Option<JsonValue>,
}
/// Connects a cell actor to session-owned callbacks and lifecycle state.
/// Connects a cell actor to session-owned callbacks and stored values.
///
/// Implementations should forward callback cancellation to downstream work.
/// The actor stops awaiting callbacks once cancellation begins. Implementations
@@ -52,10 +51,11 @@ pub(crate) trait CellHost: Send + Sync + 'static {
cancellation_token: CancellationToken,
) -> impl Future<Output = Result<(), String>> + Send;
fn commit_stored_values(
fn commit_completion(
&self,
stored_value_writes: HashMap<String, JsonValue>,
lifecycle: Arc<CellLifecycle>,
event: CellEvent,
cell_state: Arc<CellState>,
) -> impl Future<Output = bool> + Send;
fn closed(&self) -> impl Future<Output = ()> + Send;
@@ -64,28 +64,22 @@ pub(crate) trait CellHost: Send + Sync + 'static {
#[derive(Clone)]
pub(crate) struct CellHandle {
command_tx: mpsc::UnboundedSender<CellCommand>,
lifecycle: Arc<CellLifecycle>,
terminal_event_rx: watch::Receiver<Option<CellEvent>>,
state: Arc<CellState>,
}
impl CellHandle {
pub(super) fn new(
command_tx: mpsc::UnboundedSender<CellCommand>,
lifecycle: Arc<CellLifecycle>,
terminal_event_rx: watch::Receiver<Option<CellEvent>>,
state: Arc<CellState>,
) -> Self {
Self {
command_tx,
lifecycle,
terminal_event_rx,
}
Self { command_tx, state }
}
pub(crate) fn observe(&self, mode: ObserveMode) -> CellEventFuture {
if !self.lifecycle.accepting_requests() {
if !self.state.accepting_observations() {
return closed_event();
}
let (response_tx, response_rx) = tokio::sync::oneshot::channel();
let (response_tx, response_rx) = oneshot::channel();
if self
.command_tx
.send(CellCommand::Observe { mode, response_tx })
@@ -97,134 +91,207 @@ impl CellHandle {
}
pub(crate) fn terminate(&self) -> CellEventFuture {
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();
}
terminal_event(self.terminal_event_rx.clone())
self.state.request_termination()
}
}
pub(crate) struct CellLifecycle {
state: Mutex<CellLifecycleState>,
session_shutdown_token: CancellationToken,
termination_token: CancellationToken,
work_cancellation_token: CancellationToken,
/// The single linearization point for a cell's terminal outcome.
///
/// 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
/// while it is held.
pub(crate) struct CellState {
phase: Mutex<CellPhase>,
cancellation_token: CancellationToken,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum CellLifecycleState {
enum CellPhase {
Running,
TerminationRequested,
CompletionCommitted,
CompletionBuffered,
CompletionConsumptionRequested,
Terminating {
response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
},
Completed(CellEvent),
Tombstone,
}
pub(crate) enum CompletionDelivery {
Delivered,
Buffered,
Rejected(Option<oneshot::Sender<Result<CellEvent, CellError>>>),
}
pub(crate) enum ObservationDelivery {
Running(oneshot::Sender<Result<CellEvent, CellError>>),
Delivered,
Buffered,
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();
impl CellState {
pub(crate) fn new(cancellation_token: CancellationToken) -> Self {
Self {
state: Mutex::new(CellLifecycleState::Running),
session_shutdown_token,
termination_token: CancellationToken::new(),
work_cancellation_token,
phase: Mutex::new(CellPhase::Running),
cancellation_token,
}
}
pub(crate) fn accepting_requests(&self) -> bool {
let accepting_state = matches!(
pub(crate) fn accepting_observations(&self) -> bool {
let accepting_phase = matches!(
*self
.state
.phase
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner),
CellLifecycleState::Running
| CellLifecycleState::CompletionCommitted
| CellLifecycleState::CompletionBuffered
CellPhase::Running | CellPhase::Completed(_)
);
accepting_state && !self.session_shutdown_token.is_cancelled()
accepting_phase && !self.cancellation_token.is_cancelled()
}
fn request_termination(&self) -> TerminationRequest {
let mut state = self
.state
pub(crate) fn request_termination(&self) -> CellEventFuture {
let mut phase = self
.phase
.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
match std::mem::replace(&mut *phase, CellPhase::Tombstone) {
CellPhase::Running => {
let (response_tx, response_rx) = oneshot::channel();
*phase = CellPhase::Terminating { response_tx };
self.cancellation_token.cancel();
response_event(response_rx)
}
CellLifecycleState::TerminationRequested => TerminationRequest::AlreadyTerminating,
CellLifecycleState::CompletionCommitted | CellLifecycleState::CompletionBuffered => {
*state = CellLifecycleState::CompletionConsumptionRequested;
TerminationRequest::ConsumeCompletion
CellPhase::Terminating { response_tx } => {
*phase = CellPhase::Terminating { response_tx };
Box::pin(async { Err(CellError::AlreadyTerminating) })
}
CellLifecycleState::CompletionConsumptionRequested => {
TerminationRequest::AlreadyTerminating
CellPhase::Completed(event) => {
*phase = CellPhase::Completed(event.clone());
self.cancellation_token.cancel();
ready_event(event)
}
CellLifecycleState::Closed => TerminationRequest::Closed,
CellPhase::Tombstone => closed_event(),
}
}
pub(crate) fn commit_completion(&self, commit: impl FnOnce()) -> bool {
let mut state = self
.state
pub(crate) fn commit_completion(&self, event: CellEvent, commit: impl FnOnce()) -> bool {
let mut phase = self
.phase
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if *state != CellLifecycleState::Running || self.session_shutdown_token.is_cancelled() {
if !matches!(*phase, CellPhase::Running) || self.cancellation_token.is_cancelled() {
return false;
}
commit();
*state = CellLifecycleState::CompletionCommitted;
*phase = CellPhase::Completed(event);
true
}
pub(crate) fn buffer_completion(&self) {
let mut state = self
.state
pub(crate) fn deliver_completion(
&self,
response_tx: Option<oneshot::Sender<Result<CellEvent, CellError>>>,
) -> CompletionDelivery {
let mut phase = self
.phase
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if *state == CellLifecycleState::CompletionCommitted {
*state = CellLifecycleState::CompletionBuffered;
let event = match std::mem::replace(&mut *phase, CellPhase::Tombstone) {
CellPhase::Completed(event) => event,
previous => {
*phase = previous;
return CompletionDelivery::Rejected(response_tx);
}
};
let Some(response_tx) = response_tx else {
*phase = CellPhase::Completed(event);
return CompletionDelivery::Buffered;
};
match response_tx.send(Ok(event)) {
Ok(()) => {
self.cancellation_token.cancel();
CompletionDelivery::Delivered
}
Err(Ok(event)) => {
*phase = CellPhase::Completed(event);
CompletionDelivery::Buffered
}
Err(Err(error)) => {
panic!("completion delivery unexpectedly carried an actor error: {error:?}")
}
}
}
pub(crate) fn close(&self) {
*self
.state
pub(crate) fn route_observation(
&self,
response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
) -> ObservationDelivery {
let mut phase = self
.phase
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = CellLifecycleState::Closed;
.unwrap_or_else(std::sync::PoisonError::into_inner);
match std::mem::replace(&mut *phase, CellPhase::Tombstone) {
CellPhase::Running => {
*phase = CellPhase::Running;
ObservationDelivery::Running(response_tx)
}
CellPhase::Completed(event) => match response_tx.send(Ok(event)) {
Ok(()) => {
self.cancellation_token.cancel();
ObservationDelivery::Delivered
}
Err(Ok(event)) => {
*phase = CellPhase::Completed(event);
ObservationDelivery::Buffered
}
Err(Err(error)) => {
panic!("completion delivery unexpectedly carried an actor error: {error:?}")
}
},
CellPhase::Terminating {
response_tx: termination_tx,
} => {
*phase = CellPhase::Terminating {
response_tx: termination_tx,
};
let _ = response_tx.send(Err(CellError::Closed));
ObservationDelivery::Closed
}
CellPhase::Tombstone => {
let _ = response_tx.send(Err(CellError::Closed));
ObservationDelivery::Closed
}
}
}
pub(crate) fn session_shutdown_token(&self) -> CancellationToken {
self.session_shutdown_token.clone()
pub(crate) fn finish_termination(&self, event: CellEvent) -> Option<CellEvent> {
let mut phase = self
.phase
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let observer_event = match std::mem::replace(&mut *phase, CellPhase::Tombstone) {
CellPhase::Running => Some(event),
CellPhase::Terminating { response_tx } => {
let _ = response_tx.send(Ok(event.clone()));
Some(event)
}
CellPhase::Completed(completed_event) => Some(completed_event),
CellPhase::Tombstone => None,
};
self.cancellation_token.cancel();
observer_event
}
pub(crate) fn termination_token(&self) -> CancellationToken {
self.termination_token.clone()
pub(crate) fn tombstone(&self) {
*self
.phase
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = CellPhase::Tombstone;
self.cancellation_token.cancel();
}
pub(crate) fn work_cancellation_token(&self) -> CancellationToken {
self.work_cancellation_token.clone()
pub(crate) fn cancellation_token(&self) -> CancellationToken {
self.cancellation_token.clone()
}
}
@@ -233,26 +300,14 @@ pub(super) enum CellCommand {
mode: ObserveMode,
response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
},
Terminate,
}
fn response_event(
response_rx: tokio::sync::oneshot::Receiver<Result<CellEvent, CellError>>,
) -> CellEventFuture {
fn response_event(response_rx: 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 ready_event(event: CellEvent) -> CellEventFuture {
Box::pin(async move { Ok(event) })
}
fn closed_event() -> CellEventFuture {
codex-rs/code-mode/src/session_runtime/mod.rs
+28 -24
View File
@@ -7,8 +7,8 @@ use std::sync::Arc;
use serde_json::Value as JsonValue;
use tokio::sync::Mutex;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
pub use self::types::CellEvent;
pub use self::types::CellId;
@@ -27,7 +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::CellState;
use crate::cell_actor::CellToolCall;
type RuntimeEventFuture = Pin<Box<dyn Future<Output = Result<CellEvent, Error>> + Send + 'static>>;
@@ -39,20 +39,25 @@ pub struct SessionRuntime<D: SessionRuntimeDelegate> {
struct Inner<D: SessionRuntimeDelegate> {
stored_values: Mutex<HashMap<String, JsonValue>>,
cells: Mutex<HashMap<CellId, CellHandle>>,
cell_count_tx: watch::Sender<usize>,
cells: Mutex<HashMap<CellId, CellEntry>>,
// Tracks actor cleanup only. CellPhase remains the semantic lifecycle.
cell_tasks: TaskTracker,
shutdown_token: CancellationToken,
delegate: Arc<D>,
}
enum CellEntry {
Active(CellHandle),
Tombstone,
}
impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
pub fn new(delegate: Arc<D>) -> Self {
let (cell_count_tx, _) = watch::channel(/*init*/ 0);
Self {
inner: Arc::new(Inner {
stored_values: Mutex::new(HashMap::new()),
cells: Mutex::new(HashMap::new()),
cell_count_tx,
cell_tasks: TaskTracker::new(),
shutdown_token: CancellationToken::new(),
delegate,
}),
@@ -85,7 +90,7 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
) -> Result<PendingEvent, Error> {
let handle = {
let cells = self.inner.cells.lock().await;
let Some(handle) = cells.get(cell_id) else {
let Some(CellEntry::Active(handle)) = cells.get(cell_id) else {
return Err(Error::MissingCell(cell_id.clone()));
};
handle.clone()
@@ -98,7 +103,7 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
pub async fn terminate(&self, cell_id: &CellId) -> Result<CellEvent, Error> {
let handle = {
let cells = self.inner.cells.lock().await;
let Some(handle) = cells.get(cell_id) else {
let Some(CellEntry::Active(handle)) = cells.get(cell_id) else {
return Err(Error::MissingCell(cell_id.clone()));
};
handle.clone()
@@ -111,15 +116,12 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
pub async fn shutdown(&self) -> Result<(), Error> {
self.begin_shutdown();
// Serialize with cell admission before taking the count snapshot.
// 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;
let mut cell_count = self.inner.cell_count_tx.subscribe();
self.inner.cell_tasks.close();
drop(cells);
while *cell_count.borrow_and_update() != 0 {
if cell_count.changed().await.is_err() {
break;
}
}
self.inner.cell_tasks.wait().await;
Ok(())
}
@@ -141,23 +143,24 @@ impl<D: SessionRuntimeDelegate> SessionRuntime<D> {
if cells.contains_key(&cell_id) {
return Err(Error::DuplicateCell(cell_id));
}
let cell_state = Arc::new(CellState::new(self.inner.shutdown_token.child_token()));
let (handle, initial_event, task) = CellActor::prepare(
request,
stored_values,
host,
initial_observe_mode,
self.inner.shutdown_token.clone(),
cell_state,
)
.map_err(Error::Runtime)?;
cells.insert(cell_id.clone(), handle);
self.inner.cell_count_tx.send_replace(cells.len());
cells.insert(cell_id.clone(), CellEntry::Active(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.shutdown_token.cancel();
self.inner.cell_tasks.close();
}
}
@@ -216,21 +219,22 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
.await
}
async fn commit_stored_values(
async fn commit_completion(
&self,
stored_value_writes: HashMap<String, JsonValue>,
lifecycle: Arc<CellLifecycle>,
event: CellEvent,
cell_state: Arc<CellState>,
) -> bool {
let mut stored_values = self.inner.stored_values.lock().await;
lifecycle.commit_completion(|| {
cell_state.commit_completion(event, || {
stored_values.extend(stored_value_writes);
})
}
async fn closed(&self) {
let mut cells = self.inner.cells.lock().await;
if cells.remove(&self.cell_id).is_some() {
self.inner.cell_count_tx.send_replace(cells.len());
if let Some(entry) = cells.get_mut(&self.cell_id) {
*entry = CellEntry::Tombstone;
}
}
}
codex-rs/code-mode/src/session_runtime/tests.rs
+24 -15
View File
@@ -87,9 +87,9 @@ fn execute_request_with_id(cell_id: &str, source: &str) -> ExecuteRequest {
#[tokio::test]
#[expect(
clippy::await_holding_invalid_type,
reason = "test holds the registry lock to observe terminal routing before removal"
reason = "test holds the registry lock to observe terminal routing before tombstoning"
)]
async fn terminal_cells_are_unrouted_before_they_are_unregistered() {
async fn terminal_cells_are_unrouted_before_the_registry_is_tombstoned() {
let runtime = Arc::new(SessionRuntime::new(Arc::new(RecordingDelegate)));
let cell_id = CellId::new("1");
assert_eq!(
@@ -104,9 +104,11 @@ async fn terminal_cells_are_unrouted_before_they_are_unregistered() {
})
);
let mut cell_count = runtime.inner.cell_count_tx.subscribe();
let cells = runtime.inner.cells.lock().await;
let handle = cells.get(&cell_id).unwrap().clone();
let CellEntry::Active(handle) = cells.get(&cell_id).unwrap() else {
panic!("running cell was unexpectedly tombstoned");
};
let handle = handle.clone();
assert_eq!(
handle.terminate().await,
Ok(CellEvent::Terminated {
@@ -121,9 +123,11 @@ async fn terminal_cells_are_unrouted_before_they_are_unregistered() {
assert_eq!(handle.terminate().await, Err(CellError::Closed));
drop(cells);
cell_count.changed().await.unwrap();
assert_eq!(*cell_count.borrow_and_update(), 0);
assert_eq!(runtime.shutdown().await, Ok(()));
assert!(matches!(
runtime.inner.cells.lock().await.get(&cell_id),
Some(CellEntry::Tombstone)
));
}
#[tokio::test]
@@ -254,16 +258,14 @@ async fn drop_terminates_cells_when_the_registry_is_locked() {
);
let inner = Arc::clone(&runtime.inner);
let mut cell_count = inner.cell_count_tx.subscribe();
let cells = inner.cells.lock().await;
drop(runtime);
drop(cells);
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), cell_count.changed())
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), inner.cell_tasks.wait())
.await
.unwrap()
.unwrap();
assert_eq!(*cell_count.borrow_and_update(), 0);
assert!(inner.cell_tasks.is_empty());
}
#[tokio::test]
@@ -291,20 +293,18 @@ async fn drop_cancels_notifications_during_natural_completion_when_registry_is_l
);
let inner = Arc::clone(&runtime.inner);
let mut cell_count = inner.cell_count_tx.subscribe();
let cells = inner.cells.lock().await;
drop(runtime);
drop(cells);
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), cell_count.changed())
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), inner.cell_tasks.wait())
.await
.unwrap()
.unwrap();
assert_eq!(*cell_count.borrow_and_update(), 0);
assert!(inner.cell_tasks.is_empty());
}
#[tokio::test]
async fn client_owned_cell_id_is_preserved_and_active_duplicates_are_rejected() {
async fn client_owned_cell_id_is_preserved_and_ids_are_never_reused() {
let runtime = SessionRuntime::new(Arc::new(RecordingDelegate));
let cell_id = CellId::new("client-generated");
@@ -335,4 +335,13 @@ async fn client_owned_cell_id_is_preserved_and_active_duplicates_are_rejected()
content_items: Vec::new(),
})
);
assert_eq!(
runtime
.execute(
execute_request_with_id(cell_id.as_str(), "text('resurrected');"),
ObserveMode::YieldAfter(Duration::from_millis(/*millis*/ 1)),
)
.await,
Err(Error::DuplicateCell(cell_id))
);
}