mirror of
https://github.com/openai/codex.git
synced 2026-10-28 18:58:46 +00:00
Run code mode exclusively through the standalone host (#36217)
## What changed - Move the V8 implementation into a dedicated `codex-code-mode-runtime` crate used by `codex-code-mode-host`, removing the embedded runtime fallback from the Codex process. - Resolve the host executable from the active installation layout and check its availability before selecting tools. - Fall back to direct tools with a one-time warning when optional code mode is unavailable. Keep `code_mode_only` and `disable_in_process_fallback` configurations fail-closed. ## Testing - Cover host discovery for standalone and package layouts, including missing hosts and symlinks. - Verify direct-tool fallback, one-time warnings, and fail-closed code-mode-only behavior. GitOrigin-RevId: 5aa3c6f1db148b2231fc24089a2ee0e2b00dbddb
This commit is contained in:
Channing Conger
committed by
copyberry
copyberry
parent
acd540f158
commit
97576b1794
@@ -0,0 +1,128 @@
|
||||
use std::panic::AssertUnwindSafe;
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures::FutureExt;
|
||||
use tokio::task::JoinSet;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::warn;
|
||||
|
||||
use super::CellHost;
|
||||
use super::CellToolCall;
|
||||
use crate::TaskFailureHandler;
|
||||
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,
|
||||
task_failure_handler: Option<TaskFailureHandler>,
|
||||
) {
|
||||
tasks.spawn(async move {
|
||||
let callback =
|
||||
AssertUnwindSafe(async move { host.notify(call_id, text, cancellation_token).await })
|
||||
.catch_unwind()
|
||||
.await;
|
||||
match callback {
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Err(err)) => warn!("failed to deliver code mode notification: {err}"),
|
||||
Err(_) => report_task_failure(
|
||||
task_failure_handler.as_ref(),
|
||||
"code mode notification task panicked".to_string(),
|
||||
),
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
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,
|
||||
task_failure_handler: Option<TaskFailureHandler>,
|
||||
) {
|
||||
tasks.spawn(async move {
|
||||
let id = invocation.id.clone();
|
||||
let callback =
|
||||
AssertUnwindSafe(async move { host.invoke_tool(invocation, cancellation_token).await })
|
||||
.catch_unwind()
|
||||
.await;
|
||||
let (command, failure_reason) = match callback {
|
||||
Ok(Ok(result)) => (RuntimeCommand::ToolResponse { id, result }, None),
|
||||
Ok(Err(error_text)) => (RuntimeCommand::ToolError { id, error_text }, None),
|
||||
Err(_) => {
|
||||
let failure_reason = "code mode tool task panicked".to_string();
|
||||
(
|
||||
RuntimeCommand::ToolError {
|
||||
id,
|
||||
error_text: failure_reason.clone(),
|
||||
},
|
||||
Some(failure_reason),
|
||||
)
|
||||
}
|
||||
};
|
||||
let _ = runtime_tx.send(command);
|
||||
if let Some(failure_reason) = failure_reason {
|
||||
report_task_failure(task_failure_handler.as_ref(), failure_reason);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub(super) async fn finish_callbacks(
|
||||
cancellation_token: &CancellationToken,
|
||||
notification_tasks: &mut JoinSet<()>,
|
||||
tool_tasks: &mut JoinSet<()>,
|
||||
completion: CallbackCompletion,
|
||||
task_failure_handler: Option<&TaskFailureHandler>,
|
||||
) {
|
||||
if matches!(completion, CallbackCompletion::Cancel) {
|
||||
cancellation_token.cancel();
|
||||
}
|
||||
drain_tasks(notification_tasks, "notification", task_failure_handler).await;
|
||||
cancellation_token.cancel();
|
||||
drain_tasks(tool_tasks, "tool", task_failure_handler).await;
|
||||
}
|
||||
|
||||
pub(super) fn report_task_result(
|
||||
task_result: Option<Result<(), tokio::task::JoinError>>,
|
||||
description: &str,
|
||||
task_failure_handler: Option<&TaskFailureHandler>,
|
||||
) {
|
||||
if let Some(Err(err)) = task_result
|
||||
&& !err.is_cancelled()
|
||||
{
|
||||
report_task_failure(
|
||||
task_failure_handler,
|
||||
format!("code mode {description} task failed: {err}"),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn report_task_failure(task_failure_handler: Option<&TaskFailureHandler>, failure_reason: String) {
|
||||
warn!("{failure_reason}");
|
||||
if let Some(task_failure_handler) = task_failure_handler {
|
||||
task_failure_handler(failure_reason);
|
||||
}
|
||||
}
|
||||
|
||||
async fn drain_tasks(
|
||||
tasks: &mut JoinSet<()>,
|
||||
description: &str,
|
||||
task_failure_handler: Option<&TaskFailureHandler>,
|
||||
) {
|
||||
while let Some(result) = tasks.join_next().await {
|
||||
report_task_result(Some(result), description, task_failure_handler);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "callbacks_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,132 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::mpsc as std_mpsc;
|
||||
use std::time::Duration;
|
||||
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::Value as JsonValue;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::task::JoinSet;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use super::*;
|
||||
use crate::cell_actor::CellState;
|
||||
use crate::cell_actor::CompletionCommit;
|
||||
use crate::runtime::RuntimeCommand;
|
||||
use crate::session_runtime::CellEvent;
|
||||
use crate::session_runtime::ToolKind;
|
||||
use crate::session_runtime::ToolName;
|
||||
|
||||
struct PanickingCallbackHost;
|
||||
|
||||
impl CellHost for PanickingCallbackHost {
|
||||
async fn invoke_tool(
|
||||
&self,
|
||||
_invocation: CellToolCall,
|
||||
_cancellation_token: CancellationToken,
|
||||
) -> Result<JsonValue, String> {
|
||||
panic!("tool callback panic probe");
|
||||
}
|
||||
|
||||
async fn notify(
|
||||
&self,
|
||||
_call_id: String,
|
||||
_text: String,
|
||||
_cancellation_token: CancellationToken,
|
||||
) -> Result<(), String> {
|
||||
panic!("notification callback panic probe");
|
||||
}
|
||||
|
||||
async fn commit_completion(
|
||||
&self,
|
||||
_stored_value_writes: HashMap<String, JsonValue>,
|
||||
_event: CellEvent,
|
||||
_pending_initial_yield_items: Option<Vec<crate::session_runtime::OutputItem>>,
|
||||
_cell_state: Arc<CellState>,
|
||||
) -> CompletionCommit {
|
||||
panic!("unexpected completion commit");
|
||||
}
|
||||
|
||||
async fn closed(&self) {}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_callback_panic_rejects_the_js_promise_and_reports_failure() {
|
||||
let mut tasks = JoinSet::new();
|
||||
let (runtime_tx, runtime_rx) = std_mpsc::channel();
|
||||
let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
|
||||
spawn_tool(
|
||||
&mut tasks,
|
||||
Arc::new(PanickingCallbackHost),
|
||||
CellToolCall {
|
||||
id: "tool-1".to_string(),
|
||||
name: ToolName {
|
||||
name: "panic".to_string(),
|
||||
namespace: None,
|
||||
},
|
||||
kind: ToolKind::Function,
|
||||
input: None,
|
||||
},
|
||||
runtime_tx,
|
||||
CancellationToken::new(),
|
||||
Some(Arc::new(move |reason| {
|
||||
let _ = failure_tx.send(reason);
|
||||
})),
|
||||
);
|
||||
|
||||
tasks
|
||||
.join_next()
|
||||
.await
|
||||
.expect("tool callback task")
|
||||
.expect("tool callback wrapper");
|
||||
let command = runtime_rx
|
||||
.recv_timeout(Duration::from_secs(1))
|
||||
.expect("tool error command");
|
||||
let RuntimeCommand::ToolError { id, error_text } = command else {
|
||||
panic!("expected a tool error command");
|
||||
};
|
||||
assert_eq!(id, "tool-1");
|
||||
assert_eq!(error_text, "code mode tool task panicked");
|
||||
assert_eq!(failure_rx.recv().await, Some(error_text));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn notification_callback_panic_reports_failure() {
|
||||
let mut tasks = JoinSet::new();
|
||||
let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
|
||||
spawn_notification(
|
||||
&mut tasks,
|
||||
Arc::new(PanickingCallbackHost),
|
||||
"notify-1".to_string(),
|
||||
"hello".to_string(),
|
||||
CancellationToken::new(),
|
||||
Some(Arc::new(move |reason| {
|
||||
let _ = failure_tx.send(reason);
|
||||
})),
|
||||
);
|
||||
|
||||
tasks
|
||||
.join_next()
|
||||
.await
|
||||
.expect("notification callback task")
|
||||
.expect("notification callback wrapper");
|
||||
let failure_reason = failure_rx.recv().await.expect("notification failure");
|
||||
assert_eq!(failure_reason, "code mode notification task panicked");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn callback_wrapper_join_error_reports_failure() {
|
||||
let task_result = tokio::spawn(async {
|
||||
panic!("callback wrapper panic probe");
|
||||
})
|
||||
.await;
|
||||
let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
|
||||
let task_failure_handler: TaskFailureHandler = Arc::new(move |reason| {
|
||||
let _ = failure_tx.send(reason);
|
||||
});
|
||||
|
||||
report_task_result(Some(task_result), "tool", Some(&task_failure_handler));
|
||||
|
||||
let failure_reason = failure_rx.recv().await.expect("wrapper failure");
|
||||
assert!(failure_reason.contains("code mode tool task failed"));
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
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 crate::session_runtime::CreateCellRequest as CellRequest;
|
||||
use crate::session_runtime::ImageDetail as CellImageDetail;
|
||||
use crate::session_runtime::OutputItem as CellOutputItem;
|
||||
use crate::session_runtime::ToolKind as 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,
|
||||
}),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputAudio { audio_url } => {
|
||||
CellOutputItem::Audio { audio_url }
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,606 @@
|
||||
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::report_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::CellEventFuture;
|
||||
pub(crate) use self::types::CellHandle;
|
||||
pub(crate) use self::types::CellHost;
|
||||
pub(crate) use self::types::CellState;
|
||||
pub(crate) use self::types::CellToolCall;
|
||||
pub(crate) use self::types::CompletionCommit;
|
||||
use self::types::CompletionDelivery;
|
||||
use self::types::ObservationDelivery;
|
||||
use crate::TaskFailureHandler;
|
||||
use crate::runtime::PendingRuntimeMode;
|
||||
use crate::runtime::RuntimeCommand;
|
||||
use crate::runtime::RuntimeControlCommand;
|
||||
use crate::runtime::RuntimeEvent;
|
||||
use crate::runtime::spawn_runtime;
|
||||
use crate::session_runtime::CellEvent;
|
||||
use crate::session_runtime::CreateCellRequest as CellRequest;
|
||||
use crate::session_runtime::ObserveMode;
|
||||
use crate::session_runtime::OutputItem;
|
||||
use crate::session_runtime::ToolName as CellToolName;
|
||||
|
||||
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,
|
||||
cell_state: Arc<CellState>,
|
||||
task_failure_handler: Option<TaskFailureHandler>,
|
||||
) -> 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,
|
||||
task_failure_handler.clone(),
|
||||
)?;
|
||||
let handle = CellHandle::new(command_tx, Arc::clone(&cell_state));
|
||||
let task = run_cell(
|
||||
host,
|
||||
CellContext {
|
||||
runtime_tx,
|
||||
runtime_control_tx,
|
||||
runtime_terminate_handle,
|
||||
cell_state,
|
||||
},
|
||||
event_rx,
|
||||
command_rx,
|
||||
Observer {
|
||||
mode: initial_observe_mode,
|
||||
response_tx: initial_response_tx,
|
||||
},
|
||||
task_failure_handler,
|
||||
);
|
||||
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,
|
||||
cell_state: Arc<CellState>,
|
||||
}
|
||||
|
||||
struct Observer {
|
||||
mode: ObserveMode,
|
||||
response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
|
||||
}
|
||||
|
||||
async fn run_cell<H: CellHost>(
|
||||
host: Arc<H>,
|
||||
context: CellContext,
|
||||
mut event_rx: mpsc::UnboundedReceiver<RuntimeEvent>,
|
||||
command_rx: mpsc::UnboundedReceiver<CellCommand>,
|
||||
initial_observer: Observer,
|
||||
task_failure_handler: Option<TaskFailureHandler>,
|
||||
) {
|
||||
let CellContext {
|
||||
runtime_tx,
|
||||
runtime_control_tx,
|
||||
runtime_terminate_handle,
|
||||
cell_state,
|
||||
} = context;
|
||||
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 pending_frontier_ready = false;
|
||||
let mut observer = Some(initial_observer);
|
||||
let mut termination = false;
|
||||
let mut runtime_closed = false;
|
||||
let mut runtime_paused = false;
|
||||
let mut runtime_failure_reported = 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();
|
||||
let mut command_rx = Some(command_rx);
|
||||
loop {
|
||||
let yield_deadline_elapsed = yield_timer
|
||||
.as_ref()
|
||||
.is_some_and(|yield_timer| yield_timer.deadline() <= tokio::time::Instant::now());
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = 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,
|
||||
task_failure_handler.as_ref(),
|
||||
).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 Some(CellCommand::Observe { mode, response_tx }) = maybe_command else {
|
||||
cancellation_token.cancel();
|
||||
continue;
|
||||
};
|
||||
if response_tx.is_closed() {
|
||||
continue;
|
||||
}
|
||||
let response_tx = match cell_state.route_observation(mode, 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;
|
||||
}
|
||||
if matches!(mode, ObserveMode::PendingFrontier) && pending_frontier_ready {
|
||||
pending_frontier_ready = false;
|
||||
match send_cell_event(
|
||||
response_tx,
|
||||
CellEvent::Pending {
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
pending_tool_call_ids: std::mem::take(&mut pending_tool_call_ids),
|
||||
},
|
||||
) {
|
||||
Ok(()) => {}
|
||||
Err(CellEvent::Pending {
|
||||
content_items: undelivered_items,
|
||||
pending_tool_call_ids: undelivered_tool_call_ids,
|
||||
}) => {
|
||||
content_items = undelivered_items;
|
||||
pending_tool_call_ids = undelivered_tool_call_ids;
|
||||
pending_frontier_ready = true;
|
||||
}
|
||||
Err(event) => {
|
||||
panic!("pending delivery returned an unexpected event: {event:?}")
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
observer = Some(Observer { mode, response_tx });
|
||||
yield_timer = observer.as_ref().and_then(observer_timer);
|
||||
if runtime_paused && matches!(mode, ObserveMode::YieldAfter(_)) {
|
||||
pending_frontier_ready = false;
|
||||
pending_tool_call_ids.clear();
|
||||
}
|
||||
resume_for_observation(
|
||||
mode,
|
||||
&mut runtime_paused,
|
||||
&runtime_tx,
|
||||
&runtime_control_tx,
|
||||
);
|
||||
}
|
||||
_ = async {
|
||||
if let Some(yield_timer) = yield_timer.as_mut() {
|
||||
yield_timer.await;
|
||||
} else {
|
||||
std::future::pending::<()>().await;
|
||||
}
|
||||
} => {
|
||||
yield_timer = None;
|
||||
restore_undelivered_yield(
|
||||
send_observer_event(
|
||||
observer.take(),
|
||||
CellEvent::Yielded {
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
},
|
||||
),
|
||||
&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 || cancellation_token.is_cancelled() {
|
||||
finish_callbacks(
|
||||
&callback_cancellation_token,
|
||||
&mut notification_tasks,
|
||||
&mut tool_tasks,
|
||||
CallbackCompletion::Cancel,
|
||||
task_failure_handler.as_ref(),
|
||||
).await;
|
||||
finish_termination(
|
||||
&cell_state,
|
||||
observer.take().map(|observer| observer.response_tx),
|
||||
CellEvent::Terminated {
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
},
|
||||
);
|
||||
break;
|
||||
}
|
||||
if !runtime_failure_reported
|
||||
&& let Some(task_failure_handler) = &task_failure_handler
|
||||
{
|
||||
runtime_failure_reported = true;
|
||||
task_failure_handler(
|
||||
"code-mode V8 runtime thread ended unexpectedly".to_string(),
|
||||
);
|
||||
}
|
||||
finish_callbacks(
|
||||
&callback_cancellation_token,
|
||||
&mut notification_tasks,
|
||||
&mut tool_tasks,
|
||||
CallbackCompletion::DrainNotifications,
|
||||
task_failure_handler.as_ref(),
|
||||
)
|
||||
.await;
|
||||
let event = CellEvent::Completed {
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
error_text: Some("exec runtime ended unexpectedly".to_string()),
|
||||
};
|
||||
let rejected_event = match host
|
||||
.commit_completion(
|
||||
HashMap::new(),
|
||||
event,
|
||||
/*pending_initial_yield_items*/ None,
|
||||
Arc::clone(&cell_state),
|
||||
)
|
||||
.await
|
||||
{
|
||||
CompletionCommit::Committed => None,
|
||||
CompletionCommit::Rejected(event) => Some(event),
|
||||
};
|
||||
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: rejected_completion_content(rejected_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;
|
||||
pending_frontier_ready = false;
|
||||
match 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,
|
||||
),
|
||||
},
|
||||
) {
|
||||
Ok(()) => {}
|
||||
Err(CellEvent::Pending {
|
||||
content_items: undelivered_items,
|
||||
pending_tool_call_ids: undelivered_tool_call_ids,
|
||||
}) => {
|
||||
content_items = undelivered_items;
|
||||
pending_tool_call_ids = undelivered_tool_call_ids;
|
||||
pending_frontier_ready = true;
|
||||
}
|
||||
Err(event) => {
|
||||
panic!("pending delivery returned an unexpected event: {event:?}")
|
||||
}
|
||||
}
|
||||
} 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 => {
|
||||
let yield_observer = matches!(
|
||||
observer.as_ref().map(|observer| observer.mode),
|
||||
Some(ObserveMode::YieldAfter(_))
|
||||
);
|
||||
if yield_observer {
|
||||
yield_timer = None;
|
||||
restore_undelivered_yield(
|
||||
send_observer_event(
|
||||
observer.take(),
|
||||
CellEvent::Yielded {
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
},
|
||||
),
|
||||
&mut content_items,
|
||||
);
|
||||
}
|
||||
}
|
||||
RuntimeEvent::Notify { call_id, text } => {
|
||||
spawn_notification(
|
||||
&mut notification_tasks,
|
||||
Arc::clone(&host),
|
||||
call_id,
|
||||
text,
|
||||
callback_cancellation_token.child_token(),
|
||||
task_failure_handler.clone(),
|
||||
);
|
||||
}
|
||||
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(),
|
||||
callback_cancellation_token.child_token(),
|
||||
task_failure_handler.clone(),
|
||||
);
|
||||
}
|
||||
RuntimeEvent::Result { stored_value_writes, error_text } => {
|
||||
runtime_closed = true;
|
||||
yield_timer = None;
|
||||
if termination || cancellation_token.is_cancelled() {
|
||||
finish_callbacks(
|
||||
&callback_cancellation_token,
|
||||
&mut notification_tasks,
|
||||
&mut tool_tasks,
|
||||
CallbackCompletion::Cancel,
|
||||
task_failure_handler.as_ref(),
|
||||
).await;
|
||||
finish_termination(
|
||||
&cell_state,
|
||||
observer.take().map(|observer| observer.response_tx),
|
||||
CellEvent::Terminated {
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
},
|
||||
);
|
||||
break;
|
||||
}
|
||||
finish_callbacks(
|
||||
&callback_cancellation_token,
|
||||
&mut notification_tasks,
|
||||
&mut tool_tasks,
|
||||
CallbackCompletion::DrainNotifications,
|
||||
task_failure_handler.as_ref(),
|
||||
)
|
||||
.await;
|
||||
let event = CellEvent::Completed {
|
||||
content_items: std::mem::take(&mut content_items),
|
||||
error_text,
|
||||
};
|
||||
let rejected_event = match host
|
||||
.commit_completion(
|
||||
stored_value_writes,
|
||||
event,
|
||||
/*pending_initial_yield_items*/ None,
|
||||
Arc::clone(&cell_state),
|
||||
)
|
||||
.await
|
||||
{
|
||||
CompletionCommit::Committed => None,
|
||||
CompletionCommit::Rejected(event) => Some(event),
|
||||
};
|
||||
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: rejected_completion_content(rejected_event),
|
||||
},
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
RuntimeEvent::ThreadPanicked => {
|
||||
runtime_failure_reported = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
task_result = notification_tasks.join_next(), if !notification_tasks.is_empty() => {
|
||||
report_task_result(
|
||||
task_result,
|
||||
"notification",
|
||||
task_failure_handler.as_ref(),
|
||||
);
|
||||
}
|
||||
task_result = tool_tasks.join_next(), if !tool_tasks.is_empty() => {
|
||||
report_task_result(task_result, "tool", task_failure_handler.as_ref());
|
||||
}
|
||||
}
|
||||
}
|
||||
// Reject requests that arrive while asynchronous terminal cleanup runs.
|
||||
cell_state.tombstone();
|
||||
drop(command_rx.take());
|
||||
begin_termination(
|
||||
&runtime_tx,
|
||||
&runtime_control_tx,
|
||||
&runtime_terminate_handle,
|
||||
&cancellation_token,
|
||||
);
|
||||
finish_callbacks(
|
||||
&callback_cancellation_token,
|
||||
&mut notification_tasks,
|
||||
&mut tool_tasks,
|
||||
CallbackCompletion::Cancel,
|
||||
task_failure_handler.as_ref(),
|
||||
)
|
||||
.await;
|
||||
host.closed().await;
|
||||
}
|
||||
|
||||
fn send_observer_event(observer: Option<Observer>, event: CellEvent) -> Result<(), CellEvent> {
|
||||
let Some(observer) = observer else {
|
||||
return Err(event);
|
||||
};
|
||||
send_cell_event(observer.response_tx, event)
|
||||
}
|
||||
|
||||
fn send_cell_event(
|
||||
response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
|
||||
event: CellEvent,
|
||||
) -> Result<(), CellEvent> {
|
||||
match response_tx.send(Ok(event)) {
|
||||
Ok(()) => Ok(()),
|
||||
Err(Ok(event)) => Err(event),
|
||||
Err(Err(error)) => panic!("cell event delivery returned an actor error: {error:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn restore_undelivered_yield(delivery: Result<(), CellEvent>, content_items: &mut Vec<OutputItem>) {
|
||||
match delivery {
|
||||
Ok(()) => {}
|
||||
Err(CellEvent::Yielded {
|
||||
content_items: mut undelivered_items,
|
||||
}) => {
|
||||
undelivered_items.append(content_items);
|
||||
*content_items = undelivered_items;
|
||||
}
|
||||
Err(event) => panic!("yield delivery returned an unexpected event: {event:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn rejected_completion_content(event: Option<CellEvent>) -> Vec<OutputItem> {
|
||||
match event {
|
||||
Some(CellEvent::Completed { content_items, .. }) => content_items,
|
||||
None => Vec::new(),
|
||||
Some(event) => panic!("completion commit rejected an unexpected event: {event:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn finish_termination(
|
||||
cell_state: &CellState,
|
||||
observer_tx: Option<oneshot::Sender<Result<CellEvent, CellError>>>,
|
||||
event: CellEvent,
|
||||
) {
|
||||
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>>> {
|
||||
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,692 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::sync::mpsc as std_mpsc;
|
||||
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::*;
|
||||
use crate::session_runtime::OutputItem;
|
||||
|
||||
struct TestHost;
|
||||
|
||||
#[derive(Default)]
|
||||
struct RecordingHost {
|
||||
notified: AtomicBool,
|
||||
}
|
||||
|
||||
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_completion(
|
||||
&self,
|
||||
_stored_value_writes: HashMap<String, JsonValue>,
|
||||
event: CellEvent,
|
||||
pending_initial_yield_items: Option<Vec<OutputItem>>,
|
||||
cell_state: Arc<CellState>,
|
||||
) -> CompletionCommit {
|
||||
cell_state.commit_completion(event, pending_initial_yield_items, || {})
|
||||
}
|
||||
|
||||
async fn closed(&self) {}
|
||||
}
|
||||
|
||||
impl CellHost for RecordingHost {
|
||||
async fn invoke_tool(
|
||||
&self,
|
||||
_invocation: CellToolCall,
|
||||
_cancellation_token: CancellationToken,
|
||||
) -> Result<JsonValue, String> {
|
||||
Err("unexpected tool call".to_string())
|
||||
}
|
||||
|
||||
async fn notify(
|
||||
&self,
|
||||
_call_id: String,
|
||||
_text: String,
|
||||
_cancellation_token: CancellationToken,
|
||||
) -> Result<(), String> {
|
||||
self.notified.store(true, Ordering::Release);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn commit_completion(
|
||||
&self,
|
||||
_stored_value_writes: HashMap<String, JsonValue>,
|
||||
event: CellEvent,
|
||||
pending_initial_yield_items: Option<Vec<OutputItem>>,
|
||||
cell_state: Arc<CellState>,
|
||||
) -> CompletionCommit {
|
||||
cell_state.commit_completion(event, pending_initial_yield_items, || {})
|
||||
}
|
||||
|
||||
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_control_rx: std_mpsc::Receiver<RuntimeControlCommand>,
|
||||
_runtime_event_rx: mpsc::UnboundedReceiver<RuntimeEvent>,
|
||||
}
|
||||
|
||||
fn spawn_cell_actor_harness(initial_observe_mode: ObserveMode) -> CellActorHarness {
|
||||
spawn_cell_actor_harness_with_host(initial_observe_mode, Arc::new(TestHost))
|
||||
}
|
||||
|
||||
fn spawn_cell_actor_harness_with_host<H: CellHost>(
|
||||
initial_observe_mode: ObserveMode,
|
||||
host: Arc<H>,
|
||||
) -> CellActorHarness {
|
||||
spawn_cell_actor_harness_with_host_and_failure_handler(
|
||||
initial_observe_mode,
|
||||
host,
|
||||
/*task_failure_handler*/ None,
|
||||
)
|
||||
}
|
||||
|
||||
fn spawn_cell_actor_harness_with_host_and_failure_handler<H: CellHost>(
|
||||
initial_observe_mode: ObserveMode,
|
||||
host: Arc<H>,
|
||||
task_failure_handler: Option<TaskFailureHandler>,
|
||||
) -> 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,
|
||||
/*task_failure_handler*/ None,
|
||||
)
|
||||
.unwrap();
|
||||
let (runtime_control_tx, runtime_control_rx) = std_mpsc::channel();
|
||||
let cell_state = Arc::new(CellState::new(CancellationToken::new()));
|
||||
let handle = CellHandle::new(command_tx, Arc::clone(&cell_state));
|
||||
let task = tokio::spawn(run_cell(
|
||||
host,
|
||||
CellContext {
|
||||
runtime_tx,
|
||||
runtime_control_tx,
|
||||
runtime_terminate_handle,
|
||||
cell_state,
|
||||
},
|
||||
event_rx,
|
||||
command_rx,
|
||||
Observer {
|
||||
mode: initial_observe_mode,
|
||||
response_tx: initial_event_tx,
|
||||
},
|
||||
task_failure_handler,
|
||||
));
|
||||
|
||||
CellActorHarness {
|
||||
event_tx,
|
||||
handle,
|
||||
initial_event_rx,
|
||||
task,
|
||||
runtime_control_rx,
|
||||
_runtime_event_rx: runtime_event_rx,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unexpected_runtime_thread_exit_is_reported_to_the_session_owner() {
|
||||
let (failure_tx, mut failure_rx) = mpsc::unbounded_channel();
|
||||
let harness = spawn_cell_actor_harness_with_host_and_failure_handler(
|
||||
ObserveMode::YieldAfter(Duration::from_secs(60)),
|
||||
Arc::new(TestHost),
|
||||
Some(Arc::new(move |reason| {
|
||||
let _ = failure_tx.send(reason);
|
||||
})),
|
||||
);
|
||||
drop(harness.event_tx);
|
||||
|
||||
assert_eq!(
|
||||
tokio::time::timeout(Duration::from_secs(1), failure_rx.recv())
|
||||
.await
|
||||
.expect("runtime failure timeout")
|
||||
.expect("runtime failure"),
|
||||
"code-mode V8 runtime thread ended unexpectedly"
|
||||
);
|
||||
assert!(
|
||||
harness
|
||||
.initial_event_rx
|
||||
.await
|
||||
.expect("initial event")
|
||||
.is_ok()
|
||||
);
|
||||
harness.task.await.expect("cell task");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn runtime_thread_panic_remains_a_cell_error_without_owner_supervision() {
|
||||
let harness = spawn_cell_actor_harness(ObserveMode::YieldAfter(Duration::from_secs(60)));
|
||||
harness
|
||||
.event_tx
|
||||
.send(RuntimeEvent::ThreadPanicked)
|
||||
.expect("runtime panic event");
|
||||
drop(harness.event_tx);
|
||||
|
||||
assert_eq!(
|
||||
harness.initial_event_rx.await.expect("initial event"),
|
||||
Ok(CellEvent::Completed {
|
||||
content_items: Vec::new(),
|
||||
error_text: Some("exec runtime ended unexpectedly".to_string()),
|
||||
})
|
||||
);
|
||||
harness.task.await.expect("cell task");
|
||||
}
|
||||
|
||||
async fn wait_for_notification(host: &RecordingHost) {
|
||||
tokio::time::timeout(Duration::from_secs(1), async {
|
||||
while !host.notified.load(Ordering::Acquire) {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("notification barrier timed out");
|
||||
}
|
||||
|
||||
#[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![OutputItem::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();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn observation_dropped_before_dequeue_does_not_consume_output() {
|
||||
let host = Arc::new(RecordingHost::default());
|
||||
let harness = spawn_cell_actor_harness_with_host(
|
||||
ObserveMode::YieldAfter(Duration::from_secs(60)),
|
||||
Arc::clone(&host),
|
||||
);
|
||||
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
|
||||
assert!(harness.initial_event_rx.await.unwrap().is_ok());
|
||||
|
||||
drop(
|
||||
harness
|
||||
.handle
|
||||
.observe(ObserveMode::YieldAfter(Duration::from_secs(60))),
|
||||
);
|
||||
harness
|
||||
.event_tx
|
||||
.send(RuntimeEvent::ContentItem(
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "survives pre-dequeue cancellation".to_string(),
|
||||
},
|
||||
))
|
||||
.unwrap();
|
||||
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
|
||||
harness
|
||||
.event_tx
|
||||
.send(RuntimeEvent::Notify {
|
||||
call_id: "after-dropped-command".to_string(),
|
||||
text: "barrier".to_string(),
|
||||
})
|
||||
.unwrap();
|
||||
wait_for_notification(&host).await;
|
||||
|
||||
assert_eq!(
|
||||
harness
|
||||
.handle
|
||||
.observe(ObserveMode::YieldAfter(Duration::ZERO))
|
||||
.await,
|
||||
Ok(CellEvent::Yielded {
|
||||
content_items: vec![OutputItem::Text {
|
||||
text: "survives pre-dequeue cancellation".to_string(),
|
||||
}],
|
||||
})
|
||||
);
|
||||
|
||||
let termination = harness.handle.terminate();
|
||||
drop(harness.event_tx);
|
||||
assert_eq!(
|
||||
termination.await,
|
||||
Ok(CellEvent::Terminated {
|
||||
content_items: Vec::new(),
|
||||
})
|
||||
);
|
||||
harness.task.await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dropped_yield_observer_preserves_output_for_the_next_observation() {
|
||||
let host = Arc::new(RecordingHost::default());
|
||||
let harness = spawn_cell_actor_harness_with_host(
|
||||
ObserveMode::YieldAfter(Duration::from_secs(60)),
|
||||
Arc::clone(&host),
|
||||
);
|
||||
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
|
||||
assert!(harness.initial_event_rx.await.unwrap().is_ok());
|
||||
|
||||
let dropped_observation = harness
|
||||
.handle
|
||||
.observe(ObserveMode::YieldAfter(Duration::from_secs(60)));
|
||||
assert_eq!(
|
||||
harness
|
||||
.handle
|
||||
.observe(ObserveMode::YieldAfter(Duration::ZERO))
|
||||
.await,
|
||||
Err(CellError::Busy)
|
||||
);
|
||||
drop(dropped_observation);
|
||||
harness
|
||||
.event_tx
|
||||
.send(RuntimeEvent::ContentItem(
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "survives active cancellation".to_string(),
|
||||
},
|
||||
))
|
||||
.unwrap();
|
||||
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
|
||||
harness
|
||||
.event_tx
|
||||
.send(RuntimeEvent::Notify {
|
||||
call_id: "after-dropped-observer".to_string(),
|
||||
text: "barrier".to_string(),
|
||||
})
|
||||
.unwrap();
|
||||
wait_for_notification(&host).await;
|
||||
|
||||
assert_eq!(
|
||||
harness
|
||||
.handle
|
||||
.observe(ObserveMode::YieldAfter(Duration::ZERO))
|
||||
.await,
|
||||
Ok(CellEvent::Yielded {
|
||||
content_items: vec![OutputItem::Text {
|
||||
text: "survives active cancellation".to_string(),
|
||||
}],
|
||||
})
|
||||
);
|
||||
|
||||
let termination = harness.handle.terminate();
|
||||
drop(harness.event_tx);
|
||||
assert_eq!(
|
||||
termination.await,
|
||||
Ok(CellEvent::Terminated {
|
||||
content_items: Vec::new(),
|
||||
})
|
||||
);
|
||||
harness.task.await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dropped_pending_observer_preserves_the_frontier_for_the_next_observation() {
|
||||
let host = Arc::new(RecordingHost::default());
|
||||
let harness = spawn_cell_actor_harness_with_host(
|
||||
ObserveMode::YieldAfter(Duration::from_secs(60)),
|
||||
Arc::clone(&host),
|
||||
);
|
||||
harness.event_tx.send(RuntimeEvent::YieldRequested).unwrap();
|
||||
assert!(harness.initial_event_rx.await.unwrap().is_ok());
|
||||
|
||||
let dropped_observation = harness.handle.observe(ObserveMode::PendingFrontier);
|
||||
assert_eq!(
|
||||
harness.handle.observe(ObserveMode::PendingFrontier).await,
|
||||
Err(CellError::Busy)
|
||||
);
|
||||
drop(dropped_observation);
|
||||
harness
|
||||
.event_tx
|
||||
.send(RuntimeEvent::ToolCall {
|
||||
id: "tool-1".to_string(),
|
||||
name: codex_protocol::ToolName {
|
||||
name: "echo".to_string(),
|
||||
namespace: None,
|
||||
},
|
||||
kind: codex_code_mode_protocol::CodeModeToolKind::Function,
|
||||
input: Some(serde_json::json!({})),
|
||||
})
|
||||
.unwrap();
|
||||
harness.event_tx.send(RuntimeEvent::Pending).unwrap();
|
||||
harness
|
||||
.event_tx
|
||||
.send(RuntimeEvent::Notify {
|
||||
call_id: "after-dropped-pending".to_string(),
|
||||
text: "barrier".to_string(),
|
||||
})
|
||||
.unwrap();
|
||||
wait_for_notification(&host).await;
|
||||
|
||||
assert_eq!(
|
||||
harness.handle.observe(ObserveMode::PendingFrontier).await,
|
||||
Ok(CellEvent::Pending {
|
||||
content_items: Vec::new(),
|
||||
pending_tool_call_ids: vec!["tool-1".to_string()],
|
||||
})
|
||||
);
|
||||
assert!(matches!(
|
||||
harness.runtime_control_rx.try_recv(),
|
||||
Err(std_mpsc::TryRecvError::Empty)
|
||||
));
|
||||
|
||||
let termination = harness.handle.terminate();
|
||||
drop(harness.event_tx);
|
||||
assert_eq!(
|
||||
termination.await,
|
||||
Ok(CellEvent::Terminated {
|
||||
content_items: Vec::new(),
|
||||
})
|
||||
);
|
||||
harness.task.await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn only_the_first_termination_claims_a_buffered_completion() {
|
||||
let cell_state = CellState::new(CancellationToken::new());
|
||||
let completion = CellEvent::Completed {
|
||||
content_items: Vec::new(),
|
||||
error_text: None,
|
||||
};
|
||||
assert_eq!(
|
||||
cell_state.commit_completion(
|
||||
completion.clone(),
|
||||
/*pending_initial_yield_items*/ None,
|
||||
|| {}
|
||||
),
|
||||
CompletionCommit::Committed
|
||||
);
|
||||
assert!(matches!(
|
||||
cell_state.deliver_completion(/*response_tx*/ None),
|
||||
CompletionDelivery::Buffered
|
||||
));
|
||||
|
||||
let first_termination = cell_state.request_termination();
|
||||
assert_eq!(
|
||||
cell_state.request_termination().await,
|
||||
Err(CellError::AlreadyTerminating)
|
||||
);
|
||||
assert_eq!(first_termination.await, Ok(completion.clone()));
|
||||
assert_eq!(
|
||||
cell_state.finish_termination(CellEvent::Terminated {
|
||||
content_items: Vec::new(),
|
||||
}),
|
||||
Some(completion)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn termination_claim_prevents_stored_value_commit() {
|
||||
let cell_state = CellState::new(CancellationToken::new());
|
||||
let termination = cell_state.request_termination();
|
||||
let mut commit_ran = false;
|
||||
let completion = CellEvent::Completed {
|
||||
content_items: Vec::new(),
|
||||
error_text: None,
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
cell_state.commit_completion(
|
||||
completion.clone(),
|
||||
/*pending_initial_yield_items*/ None,
|
||||
|| commit_ran = true
|
||||
),
|
||||
CompletionCommit::Rejected(completion)
|
||||
);
|
||||
assert!(!commit_ran);
|
||||
|
||||
let terminated = CellEvent::Terminated {
|
||||
content_items: Vec::new(),
|
||||
};
|
||||
assert_eq!(
|
||||
cell_state.finish_termination(terminated.clone()),
|
||||
Some(terminated.clone())
|
||||
);
|
||||
assert_eq!(termination.await, Ok(terminated));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failed_completion_delivery_rebuffers_the_event() {
|
||||
let cell_state = CellState::new(CancellationToken::new());
|
||||
let event = CellEvent::Completed {
|
||||
content_items: Vec::new(),
|
||||
error_text: None,
|
||||
};
|
||||
assert_eq!(
|
||||
cell_state.commit_completion(
|
||||
event.clone(),
|
||||
/*pending_initial_yield_items*/ None,
|
||||
|| {}
|
||||
),
|
||||
CompletionCommit::Committed
|
||||
);
|
||||
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 (response_tx, mut response_rx) = oneshot::channel();
|
||||
assert!(matches!(
|
||||
cell_state.route_observation(ObserveMode::YieldAfter(Duration::ZERO), response_tx),
|
||||
ObservationDelivery::Delivered
|
||||
));
|
||||
assert_eq!(response_rx.try_recv(), Ok(Ok(event)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn buffered_initial_yield_precedes_buffered_completion_for_yield_observer() {
|
||||
let cell_state = CellState::new(CancellationToken::new());
|
||||
let completion = CellEvent::Completed {
|
||||
content_items: vec![OutputItem::Text {
|
||||
text: "after".to_string(),
|
||||
}],
|
||||
error_text: None,
|
||||
};
|
||||
assert_eq!(
|
||||
cell_state.commit_completion(
|
||||
completion.clone(),
|
||||
Some(vec![OutputItem::Text {
|
||||
text: "before".to_string(),
|
||||
}]),
|
||||
|| {}
|
||||
),
|
||||
CompletionCommit::Committed
|
||||
);
|
||||
assert!(matches!(
|
||||
cell_state.deliver_completion(/*response_tx*/ None),
|
||||
CompletionDelivery::Buffered
|
||||
));
|
||||
|
||||
let (response_tx, mut response_rx) = oneshot::channel();
|
||||
assert!(matches!(
|
||||
cell_state.route_observation(ObserveMode::YieldAfter(Duration::ZERO), response_tx),
|
||||
ObservationDelivery::Buffered
|
||||
));
|
||||
assert_eq!(
|
||||
response_rx.try_recv(),
|
||||
Ok(Ok(CellEvent::Yielded {
|
||||
content_items: vec![OutputItem::Text {
|
||||
text: "before".to_string(),
|
||||
}],
|
||||
}))
|
||||
);
|
||||
|
||||
let (response_tx, mut response_rx) = oneshot::channel();
|
||||
assert!(matches!(
|
||||
cell_state.route_observation(ObserveMode::YieldAfter(Duration::ZERO), response_tx),
|
||||
ObservationDelivery::Delivered
|
||||
));
|
||||
assert_eq!(response_rx.try_recv(), Ok(Ok(completion)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_observer_merges_initial_yield_and_completion_output() {
|
||||
let cell_state = CellState::new(CancellationToken::new());
|
||||
assert_eq!(
|
||||
cell_state.commit_completion(
|
||||
CellEvent::Completed {
|
||||
content_items: vec![OutputItem::Text {
|
||||
text: "after".to_string(),
|
||||
}],
|
||||
error_text: None,
|
||||
},
|
||||
Some(vec![OutputItem::Text {
|
||||
text: "before".to_string(),
|
||||
}]),
|
||||
|| {}
|
||||
),
|
||||
CompletionCommit::Committed
|
||||
);
|
||||
assert!(matches!(
|
||||
cell_state.deliver_completion(/*response_tx*/ None),
|
||||
CompletionDelivery::Buffered
|
||||
));
|
||||
|
||||
let (response_tx, mut response_rx) = oneshot::channel();
|
||||
assert!(matches!(
|
||||
cell_state.route_observation(ObserveMode::PendingFrontier, response_tx),
|
||||
ObservationDelivery::Delivered
|
||||
));
|
||||
assert_eq!(
|
||||
response_rx.try_recv(),
|
||||
Ok(Ok(CellEvent::Completed {
|
||||
content_items: vec![
|
||||
OutputItem::Text {
|
||||
text: "before".to_string(),
|
||||
},
|
||||
OutputItem::Text {
|
||||
text: "after".to_string(),
|
||||
},
|
||||
],
|
||||
error_text: None,
|
||||
}))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dropped_pending_observation_preserves_the_initial_yield_boundary() {
|
||||
let cell_state = CellState::new(CancellationToken::new());
|
||||
let completion = CellEvent::Completed {
|
||||
content_items: vec![OutputItem::Text {
|
||||
text: "after".to_string(),
|
||||
}],
|
||||
error_text: None,
|
||||
};
|
||||
assert_eq!(
|
||||
cell_state.commit_completion(
|
||||
completion.clone(),
|
||||
Some(vec![OutputItem::Text {
|
||||
text: "before".to_string(),
|
||||
}]),
|
||||
|| {}
|
||||
),
|
||||
CompletionCommit::Committed
|
||||
);
|
||||
assert!(matches!(
|
||||
cell_state.deliver_completion(/*response_tx*/ None),
|
||||
CompletionDelivery::Buffered
|
||||
));
|
||||
|
||||
let (response_tx, response_rx) = oneshot::channel();
|
||||
drop(response_rx);
|
||||
assert!(matches!(
|
||||
cell_state.route_observation(ObserveMode::PendingFrontier, response_tx),
|
||||
ObservationDelivery::Buffered
|
||||
));
|
||||
|
||||
let (response_tx, mut response_rx) = oneshot::channel();
|
||||
assert!(matches!(
|
||||
cell_state.route_observation(ObserveMode::YieldAfter(Duration::ZERO), response_tx),
|
||||
ObservationDelivery::Buffered
|
||||
));
|
||||
assert_eq!(
|
||||
response_rx.try_recv(),
|
||||
Ok(Ok(CellEvent::Yielded {
|
||||
content_items: vec![OutputItem::Text {
|
||||
text: "before".to_string(),
|
||||
}],
|
||||
}))
|
||||
);
|
||||
|
||||
let (response_tx, mut response_rx) = oneshot::channel();
|
||||
assert!(matches!(
|
||||
cell_state.route_observation(ObserveMode::YieldAfter(Duration::ZERO), response_tx),
|
||||
ObservationDelivery::Delivered
|
||||
));
|
||||
assert_eq!(response_rx.try_recv(), Ok(Ok(completion)));
|
||||
}
|
||||
@@ -0,0 +1,444 @@
|
||||
use std::collections::HashMap;
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use serde_json::Value as JsonValue;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::session_runtime::CellEvent;
|
||||
use crate::session_runtime::ObserveMode;
|
||||
use crate::session_runtime::OutputItem;
|
||||
use crate::session_runtime::ToolKind;
|
||||
use crate::session_runtime::ToolName;
|
||||
|
||||
pub(crate) type CellEventFuture =
|
||||
Pin<Box<dyn Future<Output = Result<CellEvent, CellError>> + Send + 'static>>;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub(crate) enum CellError {
|
||||
Busy,
|
||||
AlreadyTerminating,
|
||||
Closed,
|
||||
}
|
||||
|
||||
pub(crate) struct CellToolCall {
|
||||
pub(crate) id: String,
|
||||
pub(crate) name: ToolName,
|
||||
pub(crate) kind: ToolKind,
|
||||
pub(crate) input: Option<JsonValue>,
|
||||
}
|
||||
|
||||
/// Connects a cell actor to session-owned callbacks and stored values.
|
||||
///
|
||||
/// Implementations should forward callback cancellation to downstream work.
|
||||
/// Implementations 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_completion(
|
||||
&self,
|
||||
stored_value_writes: HashMap<String, JsonValue>,
|
||||
event: CellEvent,
|
||||
pending_initial_yield_items: Option<Vec<OutputItem>>,
|
||||
cell_state: Arc<CellState>,
|
||||
) -> impl Future<Output = CompletionCommit> + Send;
|
||||
|
||||
fn closed(&self) -> impl Future<Output = ()> + Send;
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct CellHandle {
|
||||
command_tx: mpsc::UnboundedSender<CellCommand>,
|
||||
state: Arc<CellState>,
|
||||
}
|
||||
|
||||
impl CellHandle {
|
||||
pub(super) fn new(
|
||||
command_tx: mpsc::UnboundedSender<CellCommand>,
|
||||
state: Arc<CellState>,
|
||||
) -> Self {
|
||||
Self { command_tx, state }
|
||||
}
|
||||
|
||||
pub(crate) fn observe(&self, mode: ObserveMode) -> CellEventFuture {
|
||||
if !self.state.accepting_observations() {
|
||||
return closed_event();
|
||||
}
|
||||
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 {
|
||||
self.state.request_termination()
|
||||
}
|
||||
}
|
||||
|
||||
/// 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,
|
||||
}
|
||||
|
||||
enum CellPhase {
|
||||
Running,
|
||||
Terminating {
|
||||
response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
|
||||
},
|
||||
Completed {
|
||||
// Set only when `yield_control()` races the create-to-first-observe handoff.
|
||||
pending_initial_yield_items: Option<Vec<OutputItem>>,
|
||||
event: CellEvent,
|
||||
},
|
||||
CompletionClaimed(CellEvent),
|
||||
Tombstone,
|
||||
}
|
||||
|
||||
pub(crate) enum CompletionDelivery {
|
||||
Delivered,
|
||||
Buffered,
|
||||
Rejected(Option<oneshot::Sender<Result<CellEvent, CellError>>>),
|
||||
}
|
||||
|
||||
/// Result of atomically publishing a completed cell and its session side effects.
|
||||
#[derive(Debug, PartialEq)]
|
||||
pub(crate) enum CompletionCommit {
|
||||
Committed,
|
||||
Rejected(CellEvent),
|
||||
}
|
||||
|
||||
pub(crate) enum ObservationDelivery {
|
||||
Running(oneshot::Sender<Result<CellEvent, CellError>>),
|
||||
Delivered,
|
||||
Buffered,
|
||||
Closed,
|
||||
}
|
||||
|
||||
impl CellState {
|
||||
pub(crate) fn new(cancellation_token: CancellationToken) -> Self {
|
||||
Self {
|
||||
phase: Mutex::new(CellPhase::Running),
|
||||
cancellation_token,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn accepting_observations(&self) -> bool {
|
||||
let accepting_phase = matches!(
|
||||
*self
|
||||
.phase
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner),
|
||||
CellPhase::Running | CellPhase::Completed { .. }
|
||||
);
|
||||
accepting_phase && !self.cancellation_token.is_cancelled()
|
||||
}
|
||||
|
||||
pub(crate) fn request_termination(&self) -> CellEventFuture {
|
||||
let mut phase = self
|
||||
.phase
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
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)
|
||||
}
|
||||
CellPhase::Terminating { response_tx } => {
|
||||
*phase = CellPhase::Terminating { response_tx };
|
||||
Box::pin(async { Err(CellError::AlreadyTerminating) })
|
||||
}
|
||||
CellPhase::Completed {
|
||||
pending_initial_yield_items,
|
||||
event,
|
||||
} => {
|
||||
let event = prepend_initial_yield(event, pending_initial_yield_items);
|
||||
*phase = CellPhase::CompletionClaimed(event.clone());
|
||||
self.cancellation_token.cancel();
|
||||
ready_event(event)
|
||||
}
|
||||
CellPhase::CompletionClaimed(event) => {
|
||||
*phase = CellPhase::CompletionClaimed(event);
|
||||
Box::pin(async { Err(CellError::AlreadyTerminating) })
|
||||
}
|
||||
CellPhase::Tombstone => closed_event(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn commit_completion(
|
||||
&self,
|
||||
event: CellEvent,
|
||||
pending_initial_yield_items: Option<Vec<OutputItem>>,
|
||||
commit: impl FnOnce(),
|
||||
) -> CompletionCommit {
|
||||
let mut phase = self
|
||||
.phase
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if !matches!(*phase, CellPhase::Running) || self.cancellation_token.is_cancelled() {
|
||||
return CompletionCommit::Rejected(event);
|
||||
}
|
||||
commit();
|
||||
*phase = CellPhase::Completed {
|
||||
pending_initial_yield_items,
|
||||
event,
|
||||
};
|
||||
CompletionCommit::Committed
|
||||
}
|
||||
|
||||
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);
|
||||
let (pending_initial_yield_items, event) =
|
||||
match std::mem::replace(&mut *phase, CellPhase::Tombstone) {
|
||||
CellPhase::Completed {
|
||||
pending_initial_yield_items,
|
||||
event,
|
||||
} => (pending_initial_yield_items, event),
|
||||
previous => {
|
||||
*phase = previous;
|
||||
return CompletionDelivery::Rejected(response_tx);
|
||||
}
|
||||
};
|
||||
let Some(response_tx) = response_tx else {
|
||||
*phase = CellPhase::Completed {
|
||||
pending_initial_yield_items,
|
||||
event,
|
||||
};
|
||||
return CompletionDelivery::Buffered;
|
||||
};
|
||||
match response_tx.send(Ok(event)) {
|
||||
Ok(()) => {
|
||||
self.cancellation_token.cancel();
|
||||
CompletionDelivery::Delivered
|
||||
}
|
||||
Err(Ok(event)) => {
|
||||
*phase = CellPhase::Completed {
|
||||
pending_initial_yield_items,
|
||||
event,
|
||||
};
|
||||
CompletionDelivery::Buffered
|
||||
}
|
||||
Err(Err(error)) => {
|
||||
panic!("completion delivery unexpectedly carried an actor error: {error:?}")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn route_observation(
|
||||
&self,
|
||||
mode: ObserveMode,
|
||||
response_tx: oneshot::Sender<Result<CellEvent, CellError>>,
|
||||
) -> ObservationDelivery {
|
||||
let mut phase = self
|
||||
.phase
|
||||
.lock()
|
||||
.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 {
|
||||
pending_initial_yield_items: Some(content_items),
|
||||
event,
|
||||
} if matches!(mode, ObserveMode::YieldAfter(_)) => {
|
||||
match response_tx.send(Ok(CellEvent::Yielded { content_items })) {
|
||||
Ok(()) => {
|
||||
*phase = CellPhase::Completed {
|
||||
pending_initial_yield_items: None,
|
||||
event,
|
||||
};
|
||||
ObservationDelivery::Buffered
|
||||
}
|
||||
Err(Ok(CellEvent::Yielded { content_items })) => {
|
||||
*phase = CellPhase::Completed {
|
||||
pending_initial_yield_items: Some(content_items),
|
||||
event,
|
||||
};
|
||||
ObservationDelivery::Buffered
|
||||
}
|
||||
Err(Ok(event)) => {
|
||||
panic!("initial yield delivery returned an unexpected event: {event:?}")
|
||||
}
|
||||
Err(Err(error)) => {
|
||||
panic!("initial yield delivery returned an actor error: {error:?}")
|
||||
}
|
||||
}
|
||||
}
|
||||
CellPhase::Completed {
|
||||
pending_initial_yield_items,
|
||||
event,
|
||||
} => {
|
||||
let delivered_event =
|
||||
prepend_initial_yield(event.clone(), pending_initial_yield_items.clone());
|
||||
match response_tx.send(Ok(delivered_event)) {
|
||||
Ok(()) => {
|
||||
self.cancellation_token.cancel();
|
||||
ObservationDelivery::Delivered
|
||||
}
|
||||
Err(Ok(_)) => {
|
||||
*phase = CellPhase::Completed {
|
||||
pending_initial_yield_items,
|
||||
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::CompletionClaimed(event) => {
|
||||
*phase = CellPhase::CompletionClaimed(event);
|
||||
let _ = response_tx.send(Err(CellError::Closed));
|
||||
ObservationDelivery::Closed
|
||||
}
|
||||
CellPhase::Tombstone => {
|
||||
let _ = response_tx.send(Err(CellError::Closed));
|
||||
ObservationDelivery::Closed
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
pending_initial_yield_items,
|
||||
event,
|
||||
} => Some(prepend_initial_yield(event, pending_initial_yield_items)),
|
||||
CellPhase::CompletionClaimed(completed_event) => Some(completed_event),
|
||||
CellPhase::Tombstone => None,
|
||||
};
|
||||
self.cancellation_token.cancel();
|
||||
observer_event
|
||||
}
|
||||
|
||||
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 cancellation_token(&self) -> CancellationToken {
|
||||
self.cancellation_token.clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn prepend_initial_yield(
|
||||
event: CellEvent,
|
||||
pending_initial_yield_items: Option<Vec<OutputItem>>,
|
||||
) -> CellEvent {
|
||||
let Some(mut pending_initial_yield_items) = pending_initial_yield_items else {
|
||||
return event;
|
||||
};
|
||||
match event {
|
||||
CellEvent::Yielded { mut content_items } => {
|
||||
pending_initial_yield_items.append(&mut content_items);
|
||||
CellEvent::Yielded {
|
||||
content_items: pending_initial_yield_items,
|
||||
}
|
||||
}
|
||||
CellEvent::Pending {
|
||||
mut content_items,
|
||||
pending_tool_call_ids,
|
||||
} => {
|
||||
pending_initial_yield_items.append(&mut content_items);
|
||||
CellEvent::Pending {
|
||||
content_items: pending_initial_yield_items,
|
||||
pending_tool_call_ids,
|
||||
}
|
||||
}
|
||||
CellEvent::Completed {
|
||||
mut content_items,
|
||||
error_text,
|
||||
} => {
|
||||
pending_initial_yield_items.append(&mut content_items);
|
||||
CellEvent::Completed {
|
||||
content_items: pending_initial_yield_items,
|
||||
error_text,
|
||||
}
|
||||
}
|
||||
CellEvent::Terminated { mut content_items } => {
|
||||
pending_initial_yield_items.append(&mut content_items);
|
||||
CellEvent::Terminated {
|
||||
content_items: pending_initial_yield_items,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) enum CellCommand {
|
||||
Observe {
|
||||
mode: ObserveMode,
|
||||
response_tx: 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 ready_event(event: CellEvent) -> CellEventFuture {
|
||||
Box::pin(async move { Ok(event) })
|
||||
}
|
||||
|
||||
fn closed_event() -> CellEventFuture {
|
||||
Box::pin(async { Err(CellError::Closed) })
|
||||
}
|
||||
Reference in New Issue
Block a user