mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
code-mode: move cell state into library actor (#28599)
A code-mode cell is a single JavaScript execution that can produce output, call tools, wait for asynchronous work, resume, or be terminated. This PR extracts the existing per-cell run loop into a dedicated actor that owns the cell’s lifecycle state. It is primarily an ownership change rather than a new lifecycle contract: existing behavior now has one clear implementation boundary. ### Architecture The session service remains responsible for session-wide concerns: allocating cell IDs, storing shared values, creating cells, and routing requests to them. Once a cell is created, its execution state belongs to its actor. Callers interact with the actor through a handle. The actor receives two kinds of input: runtime events and control requests. A single event loop serializes these inputs and applies the lifecycle rules. It tracks the current observer—the caller waiting for an update—along with accumulated output, outstanding callbacks, runtime state, yield deadlines, and termination progress. Observation, termination, completion, and cleanup therefore have one consistent owner. When the runtime has no immediately runnable work and is waiting only on timers or tool results, the actor can return accumulated output and information about outstanding tool calls while keeping the cell available to resume. On completion or termination, it performs the appropriate callback cleanup before publishing the final result and removing the cell from the session. A small host interface connects the actor to session-owned facilities such as tool dispatch, notifications, stored values, and final cell removal, keeping those responsibilities outside the actor itself. ### Why Previously, cell lifecycle state and coordination lived alongside session management. The actor boundary makes each cell a self-contained state machine with a single writer, while the service becomes a registry and adapter around it. This makes lifecycle behavior easier to reason about and test in isolation. It also establishes a clean boundary for later changing where cells run or how they communicate without recreating their lifecycle rules.
This commit is contained in:
Channing Conger
committed by
GitHub
Unverified
GitHub
parent
1883dedc0e
commit
e2f074e16c
@@ -0,0 +1,77 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use tokio::task::JoinSet;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::warn;
|
||||
|
||||
use super::CellHost;
|
||||
use super::CellToolCall;
|
||||
use crate::runtime::RuntimeCommand;
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub(super) enum CallbackCompletion {
|
||||
DrainNotifications,
|
||||
Cancel,
|
||||
}
|
||||
|
||||
pub(super) fn spawn_notification<H: CellHost>(
|
||||
tasks: &mut JoinSet<()>,
|
||||
host: Arc<H>,
|
||||
call_id: String,
|
||||
text: String,
|
||||
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}");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub(super) fn spawn_tool<H: CellHost>(
|
||||
tasks: &mut JoinSet<()>,
|
||||
host: Arc<H>,
|
||||
invocation: CellToolCall,
|
||||
runtime_tx: std::sync::mpsc::Sender<RuntimeCommand>,
|
||||
cancellation_token: CancellationToken,
|
||||
) {
|
||||
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 _ = runtime_tx.send(command);
|
||||
});
|
||||
}
|
||||
|
||||
pub(super) async fn finish_callbacks(
|
||||
cancellation_token: &CancellationToken,
|
||||
notification_tasks: &mut JoinSet<()>,
|
||||
tool_tasks: &mut JoinSet<()>,
|
||||
completion: CallbackCompletion,
|
||||
) {
|
||||
if matches!(completion, CallbackCompletion::Cancel) {
|
||||
cancellation_token.cancel();
|
||||
}
|
||||
drain_tasks(notification_tasks, "notification").await;
|
||||
cancellation_token.cancel();
|
||||
drain_tasks(tool_tasks, "tool").await;
|
||||
}
|
||||
|
||||
pub(super) fn log_task_result(
|
||||
task_result: Option<Result<(), tokio::task::JoinError>>,
|
||||
description: &str,
|
||||
) {
|
||||
if let Some(Err(err)) = task_result
|
||||
&& !err.is_cancelled()
|
||||
{
|
||||
warn!("code mode {description} task failed: {err}");
|
||||
}
|
||||
}
|
||||
|
||||
async fn drain_tasks(tasks: &mut JoinSet<()>, description: &str) {
|
||||
while let Some(result) = tasks.join_next().await {
|
||||
log_task_result(Some(result), description);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
use codex_code_mode_protocol::CodeModeToolKind;
|
||||
use codex_code_mode_protocol::ExecuteRequest;
|
||||
use codex_code_mode_protocol::FunctionCallOutputContentItem;
|
||||
use codex_code_mode_protocol::ImageDetail;
|
||||
use codex_code_mode_protocol::ToolDefinition;
|
||||
use codex_protocol::ToolName;
|
||||
|
||||
use super::CellImageDetail;
|
||||
use super::CellOutputItem;
|
||||
use super::CellRequest;
|
||||
use super::CellToolKind;
|
||||
|
||||
pub(super) fn runtime_request(request: CellRequest) -> ExecuteRequest {
|
||||
ExecuteRequest {
|
||||
tool_call_id: request.tool_call_id,
|
||||
enabled_tools: request
|
||||
.enabled_tools
|
||||
.into_iter()
|
||||
.map(|definition| ToolDefinition {
|
||||
name: definition.name,
|
||||
tool_name: ToolName {
|
||||
name: definition.tool_name.name,
|
||||
namespace: definition.tool_name.namespace,
|
||||
},
|
||||
description: definition.description,
|
||||
kind: match definition.kind {
|
||||
CellToolKind::Function => CodeModeToolKind::Function,
|
||||
CellToolKind::Freeform => CodeModeToolKind::Freeform,
|
||||
},
|
||||
input_schema: None,
|
||||
output_schema: None,
|
||||
})
|
||||
.collect(),
|
||||
source: request.source,
|
||||
yield_time_ms: None,
|
||||
max_output_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn cell_tool_kind(kind: CodeModeToolKind) -> CellToolKind {
|
||||
match kind {
|
||||
CodeModeToolKind::Function => CellToolKind::Function,
|
||||
CodeModeToolKind::Freeform => CellToolKind::Freeform,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn output_item(item: FunctionCallOutputContentItem) -> CellOutputItem {
|
||||
match item {
|
||||
FunctionCallOutputContentItem::InputText { text } => CellOutputItem::Text { text },
|
||||
FunctionCallOutputContentItem::InputImage { image_url, detail } => CellOutputItem::Image {
|
||||
image_url,
|
||||
detail: detail.map(|detail| match detail {
|
||||
ImageDetail::Auto => CellImageDetail::Auto,
|
||||
ImageDetail::Low => CellImageDetail::Low,
|
||||
ImageDetail::High => CellImageDetail::High,
|
||||
ImageDetail::Original => CellImageDetail::Original,
|
||||
}),
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,476 @@
|
||||
mod callbacks;
|
||||
mod conversions;
|
||||
mod types;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::Value as JsonValue;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio::task::JoinSet;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use self::callbacks::CallbackCompletion;
|
||||
use self::callbacks::finish_callbacks;
|
||||
use self::callbacks::log_task_result;
|
||||
use self::callbacks::spawn_notification;
|
||||
use self::callbacks::spawn_tool;
|
||||
use self::conversions::cell_tool_kind;
|
||||
use self::conversions::output_item;
|
||||
use self::conversions::runtime_request;
|
||||
use self::types::CellCommand;
|
||||
pub(crate) use self::types::CellError;
|
||||
pub(crate) use self::types::CellEvent;
|
||||
pub(crate) use self::types::CellEventFuture;
|
||||
pub(crate) use self::types::CellHandle;
|
||||
pub(crate) use self::types::CellHost;
|
||||
pub(crate) use self::types::CellImageDetail;
|
||||
pub(crate) use self::types::CellOutputItem;
|
||||
pub(crate) use self::types::CellRequest;
|
||||
pub(crate) use self::types::CellToolCall;
|
||||
pub(crate) use self::types::CellToolDefinition;
|
||||
pub(crate) use self::types::CellToolKind;
|
||||
pub(crate) use self::types::CellToolName;
|
||||
pub(crate) use self::types::ObserveMode;
|
||||
use crate::runtime::PendingRuntimeMode;
|
||||
use crate::runtime::RuntimeCommand;
|
||||
use crate::runtime::RuntimeControlCommand;
|
||||
use crate::runtime::RuntimeEvent;
|
||||
use crate::runtime::spawn_runtime;
|
||||
|
||||
pub(crate) struct CellActor;
|
||||
|
||||
impl CellActor {
|
||||
pub(crate) fn prepare<H: CellHost>(
|
||||
request: CellRequest,
|
||||
stored_values: HashMap<String, JsonValue>,
|
||||
host: Arc<H>,
|
||||
initial_observe_mode: ObserveMode,
|
||||
) -> Result<
|
||||
(
|
||||
CellHandle,
|
||||
CellEventFuture,
|
||||
impl Future<Output = ()> + Send + 'static,
|
||||
),
|
||||
String,
|
||||
> {
|
||||
let (event_tx, event_rx) = mpsc::unbounded_channel();
|
||||
let (command_tx, command_rx) = mpsc::unbounded_channel();
|
||||
let (initial_response_tx, initial_response_rx) = oneshot::channel();
|
||||
let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = spawn_runtime(
|
||||
stored_values,
|
||||
runtime_request(request),
|
||||
event_tx,
|
||||
PendingRuntimeMode::PauseUntilResumed,
|
||||
)?;
|
||||
let cancellation_token = CancellationToken::new();
|
||||
let handle = CellHandle::new(command_tx, cancellation_token.clone());
|
||||
let task = run_cell(
|
||||
host,
|
||||
CellContext {
|
||||
runtime_tx,
|
||||
runtime_control_tx,
|
||||
runtime_terminate_handle,
|
||||
cancellation_token,
|
||||
},
|
||||
event_rx,
|
||||
command_rx,
|
||||
Observer {
|
||||
mode: initial_observe_mode,
|
||||
response_tx: initial_response_tx,
|
||||
},
|
||||
);
|
||||
let initial_response =
|
||||
Box::pin(async move { initial_response_rx.await.unwrap_or(Err(CellError::Closed)) });
|
||||
Ok((handle, initial_response, task))
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
struct Observer {
|
||||
mode: ObserveMode,
|
||||
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,
|
||||
mut event_rx: mpsc::UnboundedReceiver<RuntimeEvent>,
|
||||
mut command_rx: mpsc::UnboundedReceiver<CellCommand>,
|
||||
initial_observer: Observer,
|
||||
) {
|
||||
let CellContext {
|
||||
runtime_tx,
|
||||
runtime_control_tx,
|
||||
runtime_terminate_handle,
|
||||
cancellation_token,
|
||||
} = context;
|
||||
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 runtime_closed = false;
|
||||
let mut runtime_paused = false;
|
||||
let mut yield_timer: Option<std::pin::Pin<Box<tokio::time::Sleep>>> = None;
|
||||
let mut notification_tasks = JoinSet::new();
|
||||
let mut tool_tasks = JoinSet::new();
|
||||
|
||||
loop {
|
||||
let yield_deadline_elapsed = yield_timer
|
||||
.as_ref()
|
||||
.is_some_and(|yield_timer| yield_timer.deadline() <= tokio::time::Instant::now());
|
||||
tokio::select! {
|
||||
biased;
|
||||
maybe_command = command_rx.recv() => {
|
||||
let Some(command) = maybe_command else {
|
||||
if completed_event.is_some() {
|
||||
break;
|
||||
}
|
||||
termination = Some(Termination { response_tx: None });
|
||||
begin_termination(
|
||||
&runtime_tx,
|
||||
&runtime_control_tx,
|
||||
&runtime_terminate_handle,
|
||||
&cancellation_token,
|
||||
);
|
||||
if runtime_closed {
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
};
|
||||
match command {
|
||||
CellCommand::Observe { mode, response_tx } => {
|
||||
if let Some(event) = completed_event.take() {
|
||||
let _ = response_tx.send(Ok(event));
|
||||
break;
|
||||
}
|
||||
if observer.is_some() || termination.is_some() {
|
||||
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 { response_tx } => {
|
||||
if let Some(event) = completed_event.take() {
|
||||
if let Some(response_tx) = response_tx {
|
||||
let _ = response_tx.send(Ok(event));
|
||||
}
|
||||
break;
|
||||
}
|
||||
if termination.is_some() {
|
||||
if let Some(response_tx) = response_tx {
|
||||
let _ = response_tx.send(Err(CellError::AlreadyTerminating));
|
||||
}
|
||||
continue;
|
||||
}
|
||||
termination = Some(Termination { response_tx });
|
||||
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(
|
||||
observer.take(),
|
||||
termination.take(),
|
||||
CellEvent::Terminated {
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
},
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
_ = async {
|
||||
if let Some(yield_timer) = yield_timer.as_mut() {
|
||||
yield_timer.await;
|
||||
} else {
|
||||
std::future::pending::<()>().await;
|
||||
}
|
||||
} => {
|
||||
yield_timer = None;
|
||||
send_observer_event(
|
||||
observer.take(),
|
||||
CellEvent::Yielded {
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
},
|
||||
);
|
||||
}
|
||||
maybe_event = async {
|
||||
if runtime_closed {
|
||||
std::future::pending::<Option<RuntimeEvent>>().await
|
||||
} else {
|
||||
event_rx.recv().await
|
||||
}
|
||||
}, if !yield_deadline_elapsed => {
|
||||
let Some(event) = maybe_event else {
|
||||
runtime_closed = true;
|
||||
if termination.is_some() {
|
||||
finish_callbacks(
|
||||
&cancellation_token,
|
||||
&mut notification_tasks,
|
||||
&mut tool_tasks,
|
||||
CallbackCompletion::Cancel,
|
||||
).await;
|
||||
send_termination_events(
|
||||
observer.take(),
|
||||
termination.take(),
|
||||
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;
|
||||
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,
|
||||
) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
continue;
|
||||
};
|
||||
match event {
|
||||
RuntimeEvent::Started => {
|
||||
yield_timer = observer.as_ref().and_then(observer_timer);
|
||||
}
|
||||
RuntimeEvent::Pending => {
|
||||
runtime_paused = true;
|
||||
if matches!(
|
||||
observer.as_ref().map(|observer| observer.mode),
|
||||
Some(ObserveMode::PendingFrontier)
|
||||
) {
|
||||
yield_timer = None;
|
||||
send_observer_event(
|
||||
observer.take(),
|
||||
CellEvent::Pending {
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
pending_tool_call_ids: std::mem::take(
|
||||
&mut pending_tool_call_ids,
|
||||
),
|
||||
},
|
||||
);
|
||||
} else {
|
||||
pending_tool_call_ids.clear();
|
||||
let _ = runtime_control_tx.send(RuntimeControlCommand::Continue);
|
||||
runtime_paused = false;
|
||||
}
|
||||
}
|
||||
RuntimeEvent::ContentItem(item) => content_items.push(output_item(item)),
|
||||
RuntimeEvent::YieldRequested => {
|
||||
if matches!(
|
||||
observer.as_ref().map(|observer| observer.mode),
|
||||
Some(ObserveMode::YieldAfter(_))
|
||||
) {
|
||||
yield_timer = None;
|
||||
send_observer_event(
|
||||
observer.take(),
|
||||
CellEvent::Yielded {
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
RuntimeEvent::Notify { call_id, text } => {
|
||||
spawn_notification(
|
||||
&mut notification_tasks,
|
||||
Arc::clone(&host),
|
||||
call_id,
|
||||
text,
|
||||
cancellation_token.child_token(),
|
||||
);
|
||||
}
|
||||
RuntimeEvent::ToolCall { id, name, kind, input } => {
|
||||
pending_tool_call_ids.push(id.clone());
|
||||
spawn_tool(
|
||||
&mut tool_tasks,
|
||||
Arc::clone(&host),
|
||||
CellToolCall {
|
||||
id,
|
||||
name: CellToolName {
|
||||
name: name.name,
|
||||
namespace: name.namespace,
|
||||
},
|
||||
kind: cell_tool_kind(kind),
|
||||
input,
|
||||
},
|
||||
runtime_tx.clone(),
|
||||
cancellation_token.child_token(),
|
||||
);
|
||||
}
|
||||
RuntimeEvent::Result { stored_value_writes, error_text } => {
|
||||
runtime_closed = true;
|
||||
yield_timer = None;
|
||||
if termination.is_some() {
|
||||
finish_callbacks(
|
||||
&cancellation_token,
|
||||
&mut notification_tasks,
|
||||
&mut tool_tasks,
|
||||
CallbackCompletion::Cancel,
|
||||
).await;
|
||||
send_termination_events(
|
||||
observer.take(),
|
||||
termination.take(),
|
||||
CellEvent::Terminated {
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
},
|
||||
);
|
||||
break;
|
||||
}
|
||||
finish_callbacks(
|
||||
&cancellation_token,
|
||||
&mut notification_tasks,
|
||||
&mut tool_tasks,
|
||||
CallbackCompletion::DrainNotifications,
|
||||
).await;
|
||||
host.commit_stored_values(stored_value_writes).await;
|
||||
let event = CellEvent::Completed {
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
error_text,
|
||||
};
|
||||
if send_or_buffer_completion(
|
||||
event,
|
||||
&mut observer,
|
||||
&mut completed_event,
|
||||
) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
task_result = notification_tasks.join_next(), if !notification_tasks.is_empty() => {
|
||||
log_task_result(task_result, "notification");
|
||||
}
|
||||
task_result = tool_tasks.join_next(), if !tool_tasks.is_empty() => {
|
||||
log_task_result(task_result, "tool");
|
||||
}
|
||||
}
|
||||
}
|
||||
begin_termination(
|
||||
&runtime_tx,
|
||||
&runtime_control_tx,
|
||||
&runtime_terminate_handle,
|
||||
&cancellation_token,
|
||||
);
|
||||
finish_callbacks(
|
||||
&cancellation_token,
|
||||
&mut notification_tasks,
|
||||
&mut tool_tasks,
|
||||
CallbackCompletion::Cancel,
|
||||
)
|
||||
.await;
|
||||
host.closed().await;
|
||||
}
|
||||
|
||||
fn send_or_buffer_completion(
|
||||
event: CellEvent,
|
||||
observer: &mut Option<Observer>,
|
||||
completed_event: &mut Option<CellEvent>,
|
||||
) -> bool {
|
||||
if observer.is_some() {
|
||||
send_observer_event(observer.take(), event);
|
||||
true
|
||||
} else {
|
||||
*completed_event = Some(event);
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
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(
|
||||
observer: Option<Observer>,
|
||||
termination: Option<Termination>,
|
||||
event: CellEvent,
|
||||
) {
|
||||
send_observer_event(observer, event.clone());
|
||||
if let Some(response_tx) = termination.and_then(|termination| termination.response_tx) {
|
||||
let _ = response_tx.send(Ok(event));
|
||||
}
|
||||
}
|
||||
|
||||
fn observer_timer(observer: &Observer) -> Option<std::pin::Pin<Box<tokio::time::Sleep>>> {
|
||||
match observer.mode {
|
||||
ObserveMode::YieldAfter(duration) => Some(Box::pin(tokio::time::sleep(duration))),
|
||||
ObserveMode::PendingFrontier => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn resume_for_observation(
|
||||
mode: ObserveMode,
|
||||
runtime_paused: &mut bool,
|
||||
runtime_tx: &std::sync::mpsc::Sender<RuntimeCommand>,
|
||||
runtime_control_tx: &std::sync::mpsc::Sender<RuntimeControlCommand>,
|
||||
) {
|
||||
if *runtime_paused {
|
||||
let control = match mode {
|
||||
ObserveMode::YieldAfter(_) => RuntimeControlCommand::Continue,
|
||||
ObserveMode::PendingFrontier => RuntimeControlCommand::Resume,
|
||||
};
|
||||
let _ = runtime_control_tx.send(control);
|
||||
*runtime_paused = false;
|
||||
} else if matches!(mode, ObserveMode::PendingFrontier) {
|
||||
let _ = runtime_tx.send(RuntimeCommand::ObservePendingFrontier);
|
||||
}
|
||||
}
|
||||
|
||||
fn begin_termination(
|
||||
runtime_tx: &std::sync::mpsc::Sender<RuntimeCommand>,
|
||||
runtime_control_tx: &std::sync::mpsc::Sender<RuntimeControlCommand>,
|
||||
runtime_terminate_handle: &v8::IsolateHandle,
|
||||
cancellation_token: &CancellationToken,
|
||||
) {
|
||||
cancellation_token.cancel();
|
||||
let _ = runtime_tx.send(RuntimeCommand::Terminate);
|
||||
let _ = runtime_control_tx.send(RuntimeControlCommand::Terminate);
|
||||
let _ = runtime_terminate_handle.terminate_execution();
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,143 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
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::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use super::*;
|
||||
|
||||
struct TestHost;
|
||||
|
||||
impl CellHost for TestHost {
|
||||
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>) {}
|
||||
|
||||
async fn closed(&self) {}
|
||||
}
|
||||
|
||||
struct CellActorHarness {
|
||||
event_tx: mpsc::UnboundedSender<RuntimeEvent>,
|
||||
handle: CellHandle,
|
||||
initial_event_rx: oneshot::Receiver<Result<CellEvent, CellError>>,
|
||||
task: tokio::task::JoinHandle<()>,
|
||||
_runtime_event_rx: mpsc::UnboundedReceiver<RuntimeEvent>,
|
||||
}
|
||||
|
||||
fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> 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();
|
||||
let (runtime_event_tx, runtime_event_rx) = mpsc::unbounded_channel();
|
||||
let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = spawn_runtime(
|
||||
HashMap::new(),
|
||||
ExecuteRequest {
|
||||
tool_call_id: "call-1".to_string(),
|
||||
enabled_tools: Vec::new(),
|
||||
source: "await new Promise(() => {});".to_string(),
|
||||
yield_time_ms: None,
|
||||
max_output_tokens: None,
|
||||
},
|
||||
runtime_event_tx,
|
||||
PendingRuntimeMode::PauseUntilResumed,
|
||||
)
|
||||
.unwrap();
|
||||
let handle = CellHandle::new(command_tx, CancellationToken::new());
|
||||
let task = tokio::spawn(run_cell(
|
||||
Arc::new(TestHost),
|
||||
CellContext {
|
||||
runtime_tx,
|
||||
runtime_control_tx,
|
||||
runtime_terminate_handle,
|
||||
cancellation_token: CancellationToken::new(),
|
||||
},
|
||||
event_rx,
|
||||
command_rx,
|
||||
Observer {
|
||||
mode: initial_observe_mode,
|
||||
response_tx: initial_event_tx,
|
||||
},
|
||||
));
|
||||
|
||||
CellActorHarness {
|
||||
event_tx,
|
||||
handle,
|
||||
initial_event_rx,
|
||||
task,
|
||||
_runtime_event_rx: runtime_event_rx,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn yield_timer_preempts_buffered_runtime_output() {
|
||||
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::ZERO));
|
||||
harness.event_tx.send(RuntimeEvent::Started).unwrap();
|
||||
harness
|
||||
.event_tx
|
||||
.send(RuntimeEvent::ContentItem(
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "queued output".to_string(),
|
||||
},
|
||||
))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
harness.initial_event_rx.await.unwrap(),
|
||||
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![CellOutputItem::Text {
|
||||
text: "queued output".to_string(),
|
||||
}],
|
||||
})
|
||||
);
|
||||
harness.task.await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn queued_termination_preempts_unobserved_runtime_completion() {
|
||||
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs(60)));
|
||||
harness
|
||||
.event_tx
|
||||
.send(RuntimeEvent::Result {
|
||||
stored_value_writes: HashMap::new(),
|
||||
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();
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
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::time::Duration;
|
||||
|
||||
use serde_json::Value as JsonValue;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
pub(crate) type CellEventFuture =
|
||||
Pin<Box<dyn Future<Output = Result<CellEvent, CellError>> + Send + 'static>>;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub(crate) enum ObserveMode {
|
||||
YieldAfter(Duration),
|
||||
PendingFrontier,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub(crate) enum CellEvent {
|
||||
Yielded {
|
||||
content_items: Vec<CellOutputItem>,
|
||||
},
|
||||
Pending {
|
||||
content_items: Vec<CellOutputItem>,
|
||||
pending_tool_call_ids: Vec<String>,
|
||||
},
|
||||
Completed {
|
||||
content_items: Vec<CellOutputItem>,
|
||||
error_text: Option<String>,
|
||||
},
|
||||
Terminated {
|
||||
content_items: Vec<CellOutputItem>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub(crate) enum CellError {
|
||||
Busy,
|
||||
AlreadyTerminating,
|
||||
Closed,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub(crate) enum CellOutputItem {
|
||||
Text {
|
||||
text: String,
|
||||
},
|
||||
Image {
|
||||
image_url: String,
|
||||
detail: Option<CellImageDetail>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub(crate) enum CellImageDetail {
|
||||
Auto,
|
||||
Low,
|
||||
High,
|
||||
Original,
|
||||
}
|
||||
|
||||
pub(crate) struct CellRequest {
|
||||
pub(crate) tool_call_id: String,
|
||||
pub(crate) enabled_tools: Vec<CellToolDefinition>,
|
||||
pub(crate) source: String,
|
||||
}
|
||||
|
||||
pub(crate) struct CellToolDefinition {
|
||||
pub(crate) name: String,
|
||||
pub(crate) tool_name: CellToolName,
|
||||
pub(crate) description: String,
|
||||
pub(crate) kind: CellToolKind,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub(crate) struct CellToolName {
|
||||
pub(crate) name: String,
|
||||
pub(crate) namespace: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub(crate) enum CellToolKind {
|
||||
Function,
|
||||
Freeform,
|
||||
}
|
||||
|
||||
pub(crate) struct CellToolCall {
|
||||
pub(crate) id: String,
|
||||
pub(crate) name: CellToolName,
|
||||
pub(crate) kind: CellToolKind,
|
||||
pub(crate) input: Option<JsonValue>,
|
||||
}
|
||||
|
||||
/// Connects a cell actor to session-owned callbacks and lifecycle state.
|
||||
///
|
||||
/// Implementations must honor callback cancellation and must not return from
|
||||
/// `closed` until the session can no longer route requests to the cell.
|
||||
pub(crate) trait CellHost: Send + Sync + 'static {
|
||||
fn invoke_tool(
|
||||
&self,
|
||||
invocation: CellToolCall,
|
||||
cancellation_token: CancellationToken,
|
||||
) -> impl Future<Output = Result<JsonValue, String>> + Send;
|
||||
|
||||
fn notify(
|
||||
&self,
|
||||
call_id: String,
|
||||
text: String,
|
||||
cancellation_token: CancellationToken,
|
||||
) -> impl Future<Output = Result<(), String>> + Send;
|
||||
|
||||
fn commit_stored_values(
|
||||
&self,
|
||||
stored_value_writes: HashMap<String, JsonValue>,
|
||||
) -> impl Future<Output = ()> + Send;
|
||||
|
||||
fn closed(&self) -> impl Future<Output = ()> + Send;
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct CellHandle {
|
||||
command_tx: mpsc::UnboundedSender<CellCommand>,
|
||||
cancellation_token: CancellationToken,
|
||||
termination_requested: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl CellHandle {
|
||||
pub(super) fn new(
|
||||
command_tx: mpsc::UnboundedSender<CellCommand>,
|
||||
cancellation_token: CancellationToken,
|
||||
) -> Self {
|
||||
Self {
|
||||
command_tx,
|
||||
cancellation_token,
|
||||
termination_requested: Arc::new(AtomicBool::new(false)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn observe(&self, mode: ObserveMode) -> CellEventFuture {
|
||||
let (response_tx, response_rx) = oneshot::channel();
|
||||
if self
|
||||
.command_tx
|
||||
.send(CellCommand::Observe { mode, response_tx })
|
||||
.is_err()
|
||||
{
|
||||
return closed_event();
|
||||
}
|
||||
response_event(response_rx)
|
||||
}
|
||||
|
||||
pub(crate) fn terminate(&self) -> CellEventFuture {
|
||||
if self
|
||||
.termination_requested
|
||||
.compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
|
||||
.is_err()
|
||||
{
|
||||
return Box::pin(async { Err(CellError::AlreadyTerminating) });
|
||||
}
|
||||
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();
|
||||
}
|
||||
response_event(response_rx)
|
||||
}
|
||||
|
||||
pub(crate) fn shutdown(&self) {
|
||||
self.termination_requested.store(true, Ordering::Relaxed);
|
||||
self.cancellation_token.cancel();
|
||||
let _ = self
|
||||
.command_tx
|
||||
.send(CellCommand::Terminate { response_tx: None });
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) enum CellCommand {
|
||||
Observe {
|
||||
mode: ObserveMode,
|
||||
response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
|
||||
},
|
||||
Terminate {
|
||||
response_tx: Option<oneshot::Sender<Result<CellEvent, CellError>>>,
|
||||
},
|
||||
}
|
||||
|
||||
fn response_event(response_rx: oneshot::Receiver<Result<CellEvent, CellError>>) -> CellEventFuture {
|
||||
Box::pin(async move { response_rx.await.unwrap_or(Err(CellError::Closed)) })
|
||||
}
|
||||
|
||||
fn closed_event() -> CellEventFuture {
|
||||
Box::pin(async { Err(CellError::Closed) })
|
||||
}
|
||||
Reference in New Issue
Block a user