diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index f6f0e55831..061bc7769c 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2517,7 +2517,6 @@ dependencies = [ "pretty_assertions", "serde", "serde_json", - "tokio", "tokio-util", ] diff --git a/codex-rs/code-mode-protocol/Cargo.toml b/codex-rs/code-mode-protocol/Cargo.toml index f6280fd091..f3ef06fd32 100644 --- a/codex-rs/code-mode-protocol/Cargo.toml +++ b/codex-rs/code-mode-protocol/Cargo.toml @@ -16,9 +16,7 @@ workspace = true codex-protocol = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } -tokio = { workspace = true, features = ["io-util", "sync"] } tokio-util = { workspace = true, features = ["rt"] } [dev-dependencies] pretty_assertions = { workspace = true } -tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/code-mode-protocol/src/lib.rs b/codex-rs/code-mode-protocol/src/lib.rs index b18fff4b6c..ec7e03c50b 100644 --- a/codex-rs/code-mode-protocol/src/lib.rs +++ b/codex-rs/code-mode-protocol/src/lib.rs @@ -20,17 +20,17 @@ pub use description::render_json_schema_to_typescript; pub use response::DEFAULT_IMAGE_DETAIL; pub use response::FunctionCallOutputContentItem; pub use response::ImageDetail; +pub use runtime::CellOutcome; pub use runtime::CodeModeNestedToolCall; +pub use runtime::CreateCellRequest; pub use runtime::DEFAULT_EXEC_YIELD_TIME_MS; pub use runtime::DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL; pub use runtime::DEFAULT_WAIT_YIELD_TIME_MS; -pub use runtime::ExecuteRequest; -pub use runtime::ExecuteToPendingOutcome; +pub use runtime::ObserveRequest; +pub use runtime::ObserveToPendingOutcome; +pub use runtime::ObserveToPendingRequest; +pub use runtime::PendingOutcome; pub use runtime::RuntimeResponse; -pub use runtime::WaitOutcome; -pub use runtime::WaitRequest; -pub use runtime::WaitToPendingOutcome; -pub use runtime::WaitToPendingRequest; pub use session::CellId; pub use session::CodeModeSession; pub use session::CodeModeSessionDelegate; @@ -38,7 +38,6 @@ pub use session::CodeModeSessionProvider; pub use session::CodeModeSessionProviderFuture; pub use session::CodeModeSessionResultFuture; pub use session::NotificationFuture; -pub use session::StartedCell; pub use session::ToolInvocationFuture; pub const PUBLIC_TOOL_NAME: &str = "exec"; diff --git a/codex-rs/code-mode-protocol/src/runtime.rs b/codex-rs/code-mode-protocol/src/runtime.rs index 147822063f..54f402dfef 100644 --- a/codex-rs/code-mode-protocol/src/runtime.rs +++ b/codex-rs/code-mode-protocol/src/runtime.rs @@ -13,33 +13,31 @@ pub const DEFAULT_WAIT_YIELD_TIME_MS: u64 = 10_000; pub const DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL: usize = 10_000; #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -pub struct ExecuteRequest { +pub struct CreateCellRequest { pub tool_call_id: String, pub enabled_tools: Vec, pub source: String, - pub yield_time_ms: Option, - pub max_output_tokens: Option, } #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -pub struct WaitRequest { +pub struct ObserveRequest { pub cell_id: CellId, pub yield_time_ms: u64, } #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] -pub struct WaitToPendingRequest { +pub struct ObserveToPendingRequest { pub cell_id: CellId, } #[derive(Debug, Deserialize, PartialEq, Serialize)] -pub enum WaitOutcome { +pub enum CellOutcome { LiveCell(RuntimeResponse), MissingCell(RuntimeResponse), } #[derive(Debug, Deserialize, PartialEq, Serialize)] -pub enum ExecuteToPendingOutcome { +pub enum PendingOutcome { Pending { cell_id: CellId, content_items: Vec, @@ -49,15 +47,15 @@ pub enum ExecuteToPendingOutcome { } #[derive(Debug, Deserialize, PartialEq, Serialize)] -pub enum WaitToPendingOutcome { - LiveCell(ExecuteToPendingOutcome), +pub enum ObserveToPendingOutcome { + LiveCell(PendingOutcome), MissingCell(RuntimeResponse), } -impl From for RuntimeResponse { - fn from(outcome: WaitOutcome) -> Self { +impl From for RuntimeResponse { + fn from(outcome: CellOutcome) -> Self { match outcome { - WaitOutcome::LiveCell(response) | WaitOutcome::MissingCell(response) => response, + CellOutcome::LiveCell(response) | CellOutcome::MissingCell(response) => response, } } } @@ -87,3 +85,7 @@ pub struct CodeModeNestedToolCall { pub tool_kind: CodeModeToolKind, pub input: Option, } + +#[cfg(test)] +#[path = "runtime_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-protocol/src/runtime_tests.rs b/codex-rs/code-mode-protocol/src/runtime_tests.rs new file mode 100644 index 0000000000..1f24e92204 --- /dev/null +++ b/codex-rs/code-mode-protocol/src/runtime_tests.rs @@ -0,0 +1,301 @@ +use std::fmt::Debug; + +use codex_protocol::ToolName; +use pretty_assertions::assert_eq; +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::Value as JsonValue; +use serde_json::json; + +use super::*; +use crate::ImageDetail; + +fn assert_json_round_trip(value: &T, expected: JsonValue) +where + T: Debug + DeserializeOwned + PartialEq + Serialize, +{ + assert_eq!(serde_json::to_value(value).unwrap(), expected); + assert_eq!(&serde_json::from_value::(expected).unwrap(), value); +} + +fn text(value: &str) -> FunctionCallOutputContentItem { + FunctionCallOutputContentItem::InputText { + text: value.to_string(), + } +} + +fn image(image_url: &str, detail: Option) -> FunctionCallOutputContentItem { + FunctionCallOutputContentItem::InputImage { + image_url: image_url.to_string(), + detail, + } +} + +#[test] +fn requests_round_trip_with_exact_tool_and_namespace_fields() { + let create = CreateCellRequest { + tool_call_id: "response-call-7".to_string(), + enabled_tools: vec![ + ToolDefinition { + name: "mcp__search__query".to_string(), + tool_name: ToolName::namespaced("search", "query"), + description: "Search indexed documents".to_string(), + kind: CodeModeToolKind::Function, + input_schema: Some(json!({ + "type": "object", + "properties": {"q": {"type": "string"}}, + })), + output_schema: None, + }, + ToolDefinition { + name: "apply_patch".to_string(), + tool_name: ToolName::plain("apply_patch"), + description: "Apply a patch".to_string(), + kind: CodeModeToolKind::Freeform, + input_schema: None, + output_schema: Some(json!({"type": "string"})), + }, + ], + source: "await tools.mcp__search__query({ q: 'rust' });".to_string(), + }; + assert_json_round_trip( + &create, + json!({ + "tool_call_id": "response-call-7", + "enabled_tools": [ + { + "name": "mcp__search__query", + "tool_name": {"name": "query", "namespace": "search"}, + "description": "Search indexed documents", + "kind": "function", + "input_schema": { + "type": "object", + "properties": {"q": {"type": "string"}}, + }, + "output_schema": null, + }, + { + "name": "apply_patch", + "tool_name": {"name": "apply_patch", "namespace": null}, + "description": "Apply a patch", + "kind": "freeform", + "input_schema": null, + "output_schema": {"type": "string"}, + }, + ], + "source": "await tools.mcp__search__query({ q: 'rust' });", + }), + ); + + assert_json_round_trip( + &ObserveRequest { + cell_id: CellId::new("cell-a7".to_string()), + yield_time_ms: 250, + }, + json!({"cell_id": "cell-a7", "yield_time_ms": 250}), + ); + assert_json_round_trip( + &ObserveToPendingRequest { + cell_id: CellId::new("cell-b9".to_string()), + }, + json!({"cell_id": "cell-b9"}), + ); +} + +#[test] +fn nested_tool_calls_round_trip_with_optional_input_and_tool_kind() { + assert_json_round_trip( + &CodeModeNestedToolCall { + cell_id: CellId::new("cell-a7".to_string()), + runtime_tool_call_id: "runtime-call-1".to_string(), + tool_name: ToolName::namespaced("search", "query"), + tool_kind: CodeModeToolKind::Function, + input: Some(json!({"q": "rust"})), + }, + json!({ + "cell_id": "cell-a7", + "runtime_tool_call_id": "runtime-call-1", + "tool_name": {"name": "query", "namespace": "search"}, + "tool_kind": "function", + "input": {"q": "rust"}, + }), + ); + assert_json_round_trip( + &CodeModeNestedToolCall { + cell_id: CellId::new("cell-b9".to_string()), + runtime_tool_call_id: "runtime-call-2".to_string(), + tool_name: ToolName::plain("apply_patch"), + tool_kind: CodeModeToolKind::Freeform, + input: None, + }, + json!({ + "cell_id": "cell-b9", + "runtime_tool_call_id": "runtime-call-2", + "tool_name": {"name": "apply_patch", "namespace": null}, + "tool_kind": "freeform", + "input": null, + }), + ); +} + +#[test] +fn runtime_responses_round_trip_with_every_output_and_terminal_variant() { + let yielded = RuntimeResponse::Yielded { + cell_id: CellId::new("cell-a7".to_string()), + content_items: vec![ + text("before"), + image("data:image/png;base64,auto", Some(ImageDetail::Auto)), + image("data:image/png;base64,low", Some(ImageDetail::Low)), + image("data:image/png;base64,high", Some(ImageDetail::High)), + image( + "data:image/png;base64,original", + Some(ImageDetail::Original), + ), + image("data:image/png;base64,default", None), + ], + }; + assert_json_round_trip( + &yielded, + json!({ + "Yielded": { + "cell_id": "cell-a7", + "content_items": [ + {"type": "input_text", "text": "before"}, + {"type": "input_image", "image_url": "data:image/png;base64,auto", "detail": "auto"}, + {"type": "input_image", "image_url": "data:image/png;base64,low", "detail": "low"}, + {"type": "input_image", "image_url": "data:image/png;base64,high", "detail": "high"}, + {"type": "input_image", "image_url": "data:image/png;base64,original", "detail": "original"}, + {"type": "input_image", "image_url": "data:image/png;base64,default"}, + ], + }, + }), + ); + assert_json_round_trip( + &RuntimeResponse::Terminated { + cell_id: CellId::new("cell-b9".to_string()), + content_items: vec![text("partial")], + }, + json!({ + "Terminated": { + "cell_id": "cell-b9", + "content_items": [{"type": "input_text", "text": "partial"}], + }, + }), + ); + assert_json_round_trip( + &RuntimeResponse::Result { + cell_id: CellId::new("cell-c3".to_string()), + content_items: vec![text("failed")], + error_text: Some("tool failed".to_string()), + }, + json!({ + "Result": { + "cell_id": "cell-c3", + "content_items": [{"type": "input_text", "text": "failed"}], + "error_text": "tool failed", + }, + }), + ); + assert_json_round_trip( + &RuntimeResponse::Result { + cell_id: CellId::new("cell-d4".to_string()), + content_items: Vec::new(), + error_text: None, + }, + json!({ + "Result": { + "cell_id": "cell-d4", + "content_items": [], + "error_text": null, + }, + }), + ); +} + +#[test] +fn observation_outcomes_round_trip_with_live_missing_pending_and_completed_cells() { + let result = RuntimeResponse::Result { + cell_id: CellId::new("cell-a7".to_string()), + content_items: vec![text("done")], + error_text: None, + }; + assert_json_round_trip( + &CellOutcome::LiveCell(result.clone()), + json!({ + "LiveCell": { + "Result": { + "cell_id": "cell-a7", + "content_items": [{"type": "input_text", "text": "done"}], + "error_text": null, + }, + }, + }), + ); + assert_json_round_trip( + &CellOutcome::MissingCell(RuntimeResponse::Terminated { + cell_id: CellId::new("missing".to_string()), + content_items: Vec::new(), + }), + json!({ + "MissingCell": { + "Terminated": {"cell_id": "missing", "content_items": []}, + }, + }), + ); + + let pending = PendingOutcome::Pending { + cell_id: CellId::new("cell-b9".to_string()), + content_items: vec![text("waiting")], + pending_tool_call_ids: vec!["runtime-call-1".to_string()], + }; + assert_json_round_trip( + &pending, + json!({ + "Pending": { + "cell_id": "cell-b9", + "content_items": [{"type": "input_text", "text": "waiting"}], + "pending_tool_call_ids": ["runtime-call-1"], + }, + }), + ); + assert_json_round_trip( + &PendingOutcome::Completed(result), + json!({ + "Completed": { + "Result": { + "cell_id": "cell-a7", + "content_items": [{"type": "input_text", "text": "done"}], + "error_text": null, + }, + }, + }), + ); + assert_json_round_trip( + &ObserveToPendingOutcome::LiveCell(pending), + json!({ + "LiveCell": { + "Pending": { + "cell_id": "cell-b9", + "content_items": [{"type": "input_text", "text": "waiting"}], + "pending_tool_call_ids": ["runtime-call-1"], + }, + }, + }), + ); + assert_json_round_trip( + &ObserveToPendingOutcome::MissingCell(RuntimeResponse::Result { + cell_id: CellId::new("missing".to_string()), + content_items: Vec::new(), + error_text: Some("cell not found".to_string()), + }), + json!({ + "MissingCell": { + "Result": { + "cell_id": "missing", + "content_items": [], + "error_text": "cell not found", + }, + }, + }), + ); +} diff --git a/codex-rs/code-mode-protocol/src/session.rs b/codex-rs/code-mode-protocol/src/session.rs index 2ebbe6e8f1..a836e27943 100644 --- a/codex-rs/code-mode-protocol/src/session.rs +++ b/codex-rs/code-mode-protocol/src/session.rs @@ -6,14 +6,12 @@ use std::sync::Arc; use serde::Deserialize; use serde::Serialize; use serde_json::Value as JsonValue; -use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; +use crate::CellOutcome; use crate::CodeModeNestedToolCall; -use crate::ExecuteRequest; -use crate::RuntimeResponse; -use crate::WaitOutcome; -use crate::WaitRequest; +use crate::CreateCellRequest; +use crate::ObserveRequest; pub type CodeModeSessionResultFuture<'a, T> = Pin> + Send + 'a>>; @@ -48,42 +46,6 @@ impl fmt::Display for CellId { } } -pub struct StartedCell { - pub cell_id: CellId, - initial_response: CodeModeSessionResultFuture<'static, RuntimeResponse>, -} - -impl StartedCell { - pub fn new(cell_id: CellId, initial_response_rx: oneshot::Receiver) -> Self { - Self { - cell_id, - initial_response: Box::pin(async move { - initial_response_rx - .await - .map_err(|_| "exec runtime ended unexpectedly".to_string()) - }), - } - } - - pub fn from_result_receiver( - cell_id: CellId, - initial_response_rx: oneshot::Receiver>, - ) -> Self { - Self { - cell_id, - initial_response: Box::pin(async move { - initial_response_rx - .await - .map_err(|_| "exec runtime ended unexpectedly".to_string())? - }), - } - } - - pub async fn initial_response(self) -> Result { - self.initial_response.await - } -} - /// Host callbacks used by a code-mode session while cells are executing. pub trait CodeModeSessionDelegate: Send + Sync { fn invoke_tool<'a>( @@ -104,11 +66,13 @@ pub trait CodeModeSessionDelegate: Send + Sync { fn cell_closed(&self, cell_id: &CellId); } -/// A durable code-mode session owned by one Codex thread. +/// A stateful code-mode session owned by one Codex thread. /// /// Cells executed in the same session share stored values. Separate sessions /// must keep those values isolated. Implementations may execute cells -/// in-process or remotely. +/// in-process or remotely. Implementations should surface lost connections or +/// protocol desynchronization, but do not need to preserve cells or stored +/// values across process failure or restart. pub trait CodeModeSession: Send + Sync { /// Returns whether the session can still accept requests. /// @@ -116,14 +80,17 @@ pub trait CodeModeSession: Send + Sync { /// connection fails so callers can create a fresh session for later work. fn is_alive(&self) -> bool; - fn execute<'a>( + fn create_cell<'a>( &'a self, - request: ExecuteRequest, - ) -> CodeModeSessionResultFuture<'a, StartedCell>; + request: CreateCellRequest, + ) -> CodeModeSessionResultFuture<'a, CellId>; - fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome>; + fn observe<'a>( + &'a self, + request: ObserveRequest, + ) -> CodeModeSessionResultFuture<'a, CellOutcome>; - fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>; + fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, CellOutcome>; fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()>; } @@ -138,7 +105,3 @@ pub trait CodeModeSessionProvider: Send + Sync { delegate: Arc, ) -> CodeModeSessionProviderFuture<'a>; } - -#[cfg(test)] -#[path = "session_tests.rs"] -mod tests; diff --git a/codex-rs/code-mode-protocol/src/session_tests.rs b/codex-rs/code-mode-protocol/src/session_tests.rs deleted file mode 100644 index d0f6491105..0000000000 --- a/codex-rs/code-mode-protocol/src/session_tests.rs +++ /dev/null @@ -1,19 +0,0 @@ -use pretty_assertions::assert_eq; -use tokio::sync::oneshot; - -use super::CellId; -use super::StartedCell; - -#[tokio::test] -async fn started_cell_preserves_remote_initial_response_errors() { - let (response_tx, response_rx) = oneshot::channel(); - response_tx - .send(Err("remote runtime failed".to_string())) - .expect("initial response receiver should be open"); - let started = StartedCell::from_result_receiver(CellId::new("1".to_string()), response_rx); - - assert_eq!( - started.initial_response().await, - Err("remote runtime failed".to_string()) - ); -} diff --git a/codex-rs/code-mode/src/cell_actor/conversions.rs b/codex-rs/code-mode/src/cell_actor/conversions.rs index 7fb2a25c63..324002a79a 100644 --- a/codex-rs/code-mode/src/cell_actor/conversions.rs +++ b/codex-rs/code-mode/src/cell_actor/conversions.rs @@ -1,5 +1,5 @@ use codex_code_mode_protocol::CodeModeToolKind; -use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::CreateCellRequest; use codex_code_mode_protocol::FunctionCallOutputContentItem; use codex_code_mode_protocol::ImageDetail; use codex_code_mode_protocol::ToolDefinition; @@ -10,8 +10,8 @@ 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 { +pub(super) fn runtime_request(request: CellRequest) -> CreateCellRequest { + CreateCellRequest { tool_call_id: request.tool_call_id, enabled_tools: request .enabled_tools @@ -32,8 +32,6 @@ pub(super) fn runtime_request(request: CellRequest) -> ExecuteRequest { }) .collect(), source: request.source, - yield_time_ms: None, - max_output_tokens: None, } } diff --git a/codex-rs/code-mode/src/cell_actor/tests.rs b/codex-rs/code-mode/src/cell_actor/tests.rs index b5bf89d753..93fc3e4771 100644 --- a/codex-rs/code-mode/src/cell_actor/tests.rs +++ b/codex-rs/code-mode/src/cell_actor/tests.rs @@ -5,7 +5,7 @@ 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::CreateCellRequest; use codex_code_mode_protocol::FunctionCallOutputContentItem; use pretty_assertions::assert_eq; use serde_json::Value as JsonValue; @@ -109,12 +109,10 @@ fn spawn_cell_actor_harness_with_host(host: Arc) -> CellActorHar let (runtime_event_tx, runtime_event_rx) = mpsc::unbounded_channel(); let (runtime_tx, runtime_pause_tx, runtime_terminate_handle) = spawn_runtime( HashMap::new(), - ExecuteRequest { + CreateCellRequest { 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, diff --git a/codex-rs/code-mode/src/runtime/mod.rs b/codex-rs/code-mode/src/runtime/mod.rs index bfc54f71be..ce6a27c9ae 100644 --- a/codex-rs/code-mode/src/runtime/mod.rs +++ b/codex-rs/code-mode/src/runtime/mod.rs @@ -10,8 +10,8 @@ use std::sync::mpsc as std_mpsc; use std::thread; use codex_code_mode_protocol::CodeModeToolKind; +use codex_code_mode_protocol::CreateCellRequest; use codex_code_mode_protocol::EnabledToolMetadata; -use codex_code_mode_protocol::ExecuteRequest; use codex_code_mode_protocol::FunctionCallOutputContentItem; use codex_code_mode_protocol::enabled_tool_metadata; use codex_protocol::ToolName; @@ -67,7 +67,7 @@ pub(crate) enum RuntimeEvent { pub(crate) fn spawn_runtime( stored_values: HashMap, - request: ExecuteRequest, + request: CreateCellRequest, event_tx: mpsc::UnboundedSender, pending_mode: PendingRuntimeMode, ) -> Result< @@ -330,7 +330,7 @@ mod tests { use pretty_assertions::assert_eq; use tokio::sync::mpsc; - use super::ExecuteRequest; + use super::CreateCellRequest; use super::PendingRuntimeMode; use super::RuntimeCommand; use super::RuntimeControlCommand; @@ -338,13 +338,11 @@ mod tests { use super::spawn_runtime; use crate::FunctionCallOutputContentItem; - fn execute_request(source: &str) -> ExecuteRequest { - ExecuteRequest { + fn execute_request(source: &str) -> CreateCellRequest { + CreateCellRequest { tool_call_id: "call_1".to_string(), enabled_tools: Vec::new(), source: source.to_string(), - yield_time_ms: Some(1), - max_output_tokens: None, } } diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index 24baa0c457..0f3181d864 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use std::time::Duration; use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CellOutcome; use codex_code_mode_protocol::CodeModeNestedToolCall; use codex_code_mode_protocol::CodeModeSession; use codex_code_mode_protocol::CodeModeSessionDelegate; @@ -9,21 +10,17 @@ use codex_code_mode_protocol::CodeModeSessionProvider; use codex_code_mode_protocol::CodeModeSessionProviderFuture; use codex_code_mode_protocol::CodeModeSessionResultFuture; use codex_code_mode_protocol::CodeModeToolKind; -use codex_code_mode_protocol::DEFAULT_EXEC_YIELD_TIME_MS; -use codex_code_mode_protocol::ExecuteRequest; -use codex_code_mode_protocol::ExecuteToPendingOutcome; +use codex_code_mode_protocol::CreateCellRequest; use codex_code_mode_protocol::FunctionCallOutputContentItem; use codex_code_mode_protocol::ImageDetail; use codex_code_mode_protocol::NotificationFuture; +use codex_code_mode_protocol::ObserveRequest; +use codex_code_mode_protocol::ObserveToPendingOutcome; +use codex_code_mode_protocol::ObserveToPendingRequest; +use codex_code_mode_protocol::PendingOutcome; use codex_code_mode_protocol::RuntimeResponse; -use codex_code_mode_protocol::StartedCell; use codex_code_mode_protocol::ToolInvocationFuture; -use codex_code_mode_protocol::WaitOutcome; -use codex_code_mode_protocol::WaitRequest; -use codex_code_mode_protocol::WaitToPendingOutcome; -use codex_code_mode_protocol::WaitToPendingRequest; use serde_json::Value as JsonValue; -use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; use crate::session_runtime as runtime; @@ -87,62 +84,24 @@ impl CodeModeService { } } - pub async fn execute(&self, request: ExecuteRequest) -> Result { - let yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS); + pub async fn create_cell(&self, request: CreateCellRequest) -> Result { let runtime_cell_id = self .runtime .create_cell(runtime_request(request)) .await .map_err(|error| error.to_string())?; - let pending_event = self - .runtime - .begin_observe( - &runtime_cell_id, - runtime::ObserveMode::YieldAfter(Duration::from_millis(yield_time_ms)), - ) - .await - .map_err(|error| error.to_string())?; - let cell_id = protocol_cell_id(&runtime_cell_id); - let response_cell_id = cell_id.clone(); - let (response_tx, response_rx) = oneshot::channel(); - tokio::spawn(async move { - let response = pending_event - .event() - .await - .map_err(|error| error.to_string()) - .and_then(|event| runtime_response(&response_cell_id, event)); - let _ = response_tx.send(response); - }); - Ok(StartedCell::from_result_receiver(cell_id, response_rx)) + Ok(protocol_cell_id(&runtime_cell_id)) } - pub async fn execute_to_pending( + pub async fn observe(&self, request: ObserveRequest) -> Result { + self.begin_observe(request).await.await + } + + async fn begin_observe( &self, - request: ExecuteRequest, - ) -> Result { - let runtime_cell_id = self - .runtime - .create_cell(runtime_request(request)) - .await - .map_err(|error| error.to_string())?; - let cell_id = protocol_cell_id(&runtime_cell_id); - let event = self - .runtime - .observe(&runtime_cell_id, runtime::ObserveMode::PendingFrontier) - .await - .map_err(|error| error.to_string())?; - pending_outcome(&cell_id, event) - } - - pub async fn wait(&self, request: WaitRequest) -> Result { - self.begin_wait(request).await.await - } - - async fn begin_wait( - &self, - request: WaitRequest, - ) -> CodeModeSessionResultFuture<'static, WaitOutcome> { - let WaitRequest { + request: ObserveRequest, + ) -> CodeModeSessionResultFuture<'static, CellOutcome> { + let ObserveRequest { cell_id, yield_time_ms, } = request; @@ -157,34 +116,34 @@ impl CodeModeService { { Ok(pending_event) => Box::pin(async move { match pending_event.event().await { - Ok(event) => Ok(WaitOutcome::LiveCell(runtime_response(&cell_id, event)?)), + Ok(event) => Ok(CellOutcome::LiveCell(runtime_response(&cell_id, event)?)), Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => { - Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))) + Ok(CellOutcome::MissingCell(missing_cell_response(cell_id))) } Err(error) => Err(error.to_string()), } }), Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => { - missing_wait(cell_id) + missing_observation(cell_id) } Err(error) => Box::pin(async move { Err(error.to_string()) }), } } - pub async fn terminate(&self, cell_id: CellId) -> Result { + pub async fn terminate(&self, cell_id: CellId) -> Result { match self.runtime.terminate(&runtime_cell_id(&cell_id)).await { - Ok(event) => Ok(WaitOutcome::LiveCell(runtime_response(&cell_id, event)?)), + Ok(event) => Ok(CellOutcome::LiveCell(runtime_response(&cell_id, event)?)), Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => { - Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))) + Ok(CellOutcome::MissingCell(missing_cell_response(cell_id))) } Err(error) => Err(error.to_string()), } } - pub async fn wait_to_pending( + pub async fn observe_to_pending( &self, - request: WaitToPendingRequest, - ) -> Result { + request: ObserveToPendingRequest, + ) -> Result { let cell_id = request.cell_id; match self .runtime @@ -194,11 +153,11 @@ impl CodeModeService { ) .await { - Ok(event) => Ok(WaitToPendingOutcome::LiveCell(pending_outcome( + Ok(event) => Ok(ObserveToPendingOutcome::LiveCell(pending_outcome( &cell_id, event, )?)), Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => Ok( - WaitToPendingOutcome::MissingCell(missing_cell_response(cell_id)), + ObserveToPendingOutcome::MissingCell(missing_cell_response(cell_id)), ), Err(error) => Err(error.to_string()), } @@ -223,18 +182,21 @@ impl CodeModeSession for CodeModeService { self.runtime.is_alive() } - fn execute<'a>( + fn create_cell<'a>( &'a self, - request: ExecuteRequest, - ) -> CodeModeSessionResultFuture<'a, StartedCell> { - Box::pin(CodeModeService::execute(self, request)) + request: CreateCellRequest, + ) -> CodeModeSessionResultFuture<'a, CellId> { + Box::pin(CodeModeService::create_cell(self, request)) } - fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> { - Box::pin(CodeModeService::wait(self, request)) + fn observe<'a>( + &'a self, + request: ObserveRequest, + ) -> CodeModeSessionResultFuture<'a, CellOutcome> { + Box::pin(CodeModeService::observe(self, request)) } - fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, CellOutcome> { Box::pin(CodeModeService::terminate(self, cell_id)) } @@ -295,7 +257,7 @@ impl runtime::SessionRuntimeDelegate for ProtocolDelegate { } } -fn runtime_request(request: ExecuteRequest) -> runtime::CreateCellRequest { +fn runtime_request(request: CreateCellRequest) -> runtime::CreateCellRequest { runtime::CreateCellRequest { tool_call_id: request.tool_call_id, enabled_tools: request @@ -326,22 +288,17 @@ fn protocol_cell_id(cell_id: &runtime::CellId) -> CellId { CellId::new(cell_id.as_str().to_string()) } -fn pending_outcome( - cell_id: &CellId, - event: runtime::CellEvent, -) -> Result { +fn pending_outcome(cell_id: &CellId, event: runtime::CellEvent) -> Result { match event { runtime::CellEvent::Pending { content_items, pending_tool_call_ids, - } => Ok(ExecuteToPendingOutcome::Pending { + } => Ok(PendingOutcome::Pending { cell_id: cell_id.clone(), content_items: content_items.into_iter().map(output_item).collect(), pending_tool_call_ids, }), - event => Ok(ExecuteToPendingOutcome::Completed(runtime_response( - cell_id, event, - )?)), + event => Ok(PendingOutcome::Completed(runtime_response(cell_id, event)?)), } } @@ -397,8 +354,8 @@ fn missing_cell_response(cell_id: CellId) -> RuntimeResponse { } } -fn missing_wait(cell_id: CellId) -> CodeModeSessionResultFuture<'static, WaitOutcome> { - Box::pin(async move { Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))) }) +fn missing_observation(cell_id: CellId) -> CodeModeSessionResultFuture<'static, CellOutcome> { + Box::pin(async move { Ok(CellOutcome::MissingCell(missing_cell_response(cell_id))) }) } #[cfg(test)] diff --git a/codex-rs/code-mode/src/service_contract_tests.rs b/codex-rs/code-mode/src/service_contract_tests.rs index 17f4789d33..bbeff7f346 100644 --- a/codex-rs/code-mode/src/service_contract_tests.rs +++ b/codex-rs/code-mode/src/service_contract_tests.rs @@ -178,13 +178,45 @@ fn cell_id(value: &str) -> CellId { CellId::new(value.to_string()) } -fn execute_request(source: &str) -> ExecuteRequest { - ExecuteRequest { +fn execute_request(source: &str) -> CreateCellRequest { + CreateCellRequest { tool_call_id: "call-1".to_string(), enabled_tools: Vec::new(), source: source.to_string(), - yield_time_ms: Some(1), - max_output_tokens: None, + } +} + +async fn execute(service: &CodeModeService, request: CreateCellRequest) -> RuntimeResponse { + execute_with_yield_time(service, request, /*yield_time_ms*/ 1).await +} + +async fn execute_with_yield_time( + service: &CodeModeService, + request: CreateCellRequest, + yield_time_ms: u64, +) -> RuntimeResponse { + let cell_id = service.create_cell(request).await.unwrap(); + service + .observe(ObserveRequest { + cell_id, + yield_time_ms, + }) + .await + .unwrap() + .into() +} + +async fn create_and_observe_to_pending( + service: &CodeModeService, + request: CreateCellRequest, +) -> Result { + let cell_id = service.create_cell(request).await?; + match service + .observe_to_pending(ObserveToPendingRequest { cell_id }) + .await? + { + ObserveToPendingOutcome::LiveCell(outcome) => Ok(outcome), + ObserveToPendingOutcome::MissingCell(response) => Ok(PendingOutcome::Completed(response)), } } @@ -209,17 +241,18 @@ async fn next_event(events_rx: &mut mpsc::UnboundedReceiver) -> D #[tokio::test] async fn yields_and_resumes() { let service = CodeModeService::new(); - let cell = service - .execute(ExecuteRequest { + let cell = execute_with_yield_time( + &service, + CreateCellRequest { source: r#"text("before"); yield_control(); text("after");"#.to_string(), - yield_time_ms: Some(60_000), ..execute_request("") - }) - .await - .unwrap(); + }, + /*yield_time_ms*/ 60_000, + ) + .await; assert_eq!( - cell.initial_response().await.unwrap(), + cell, RuntimeResponse::Yielded { cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { @@ -229,13 +262,13 @@ async fn yields_and_resumes() { ); assert_eq!( service - .wait(WaitRequest { + .observe(ObserveRequest { cell_id: cell_id("1"), yield_time_ms: 60_000, }) .await .unwrap(), - WaitOutcome::LiveCell(RuntimeResponse::Result { + CellOutcome::LiveCell(RuntimeResponse::Result { cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "after".to_string(), @@ -251,20 +284,21 @@ async fn returns_and_resumes_from_the_pending_frontier() { let service = CodeModeService::with_delegate(delegate.clone()); assert_eq!( - service - .execute_to_pending(ExecuteRequest { + create_and_observe_to_pending( + &service, + CreateCellRequest { enabled_tools: vec![blocking_tool()], source: r#" await tools.block({}); text("after"); "# .to_string(), - yield_time_ms: Some(60_000), ..execute_request("") - }) - .await - .unwrap(), - ExecuteToPendingOutcome::Pending { + }, + ) + .await + .unwrap(), + PendingOutcome::Pending { cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: vec!["tool-1".to_string()], @@ -276,53 +310,51 @@ text("after"); assert_eq!( service - .wait_to_pending(WaitToPendingRequest { + .observe_to_pending(ObserveToPendingRequest { cell_id: cell_id("1"), }) .await .unwrap(), - WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Completed( - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "after".to_string(), - }], - error_text: None, - } - )) + ObserveToPendingOutcome::LiveCell(PendingOutcome::Completed(RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "after".to_string(), + }], + error_text: None, + })) ); } #[tokio::test] async fn observed_natural_completion_wins_over_termination() { let service = CodeModeService::new(); - let cell = service - .execute(execute_request( - r#"yield_control(); store("finished", true); text("done");"#, - )) - .await - .unwrap(); + let cell = execute( + &service, + execute_request(r#"yield_control(); store("finished", true); text("done");"#), + ) + .await; assert_eq!( - cell.initial_response().await.unwrap(), + cell, RuntimeResponse::Yielded { cell_id: cell_id("1"), content_items: Vec::new(), } ); - tokio::time::timeout(Duration::from_secs(1), async { + tokio::time::timeout(Duration::from_secs(10), async { loop { - let response = service - .execute(ExecuteRequest { - yield_time_ms: Some(60_000), + // Observe the synchronous probe through completion so scheduler load + // cannot turn its short yield deadline into a false lifecycle failure. + let response = create_and_observe_to_pending( + &service, + CreateCellRequest { ..execute_request(r#"text(String(load("finished")));"#) - }) - .await - .unwrap() - .initial_response() - .await - .unwrap(); - let RuntimeResponse::Result { content_items, .. } = response else { + }, + ) + .await + .unwrap(); + let PendingOutcome::Completed(RuntimeResponse::Result { content_items, .. }) = response + else { panic!("expected stored-value probe to complete"); }; if content_items @@ -339,7 +371,7 @@ async fn observed_natural_completion_wins_over_termination() { .unwrap(); assert_eq!( service.terminate(cell_id("1")).await.unwrap(), - WaitOutcome::LiveCell(RuntimeResponse::Result { + CellOutcome::LiveCell(RuntimeResponse::Result { cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), @@ -353,19 +385,18 @@ async fn observed_natural_completion_wins_over_termination() { async fn termination_discards_pending_callbacks_before_responding() { let (delegate, mut events_rx) = BlockingDelegate::new(); let service = CodeModeService::with_delegate(delegate.clone()); - let cell = service - .execute(execute_request( - r#"notify("pending"); await new Promise(() => {});"#, - )) - .await - .unwrap(); + let cell = execute( + &service, + execute_request(r#"notify("pending"); await new Promise(() => {});"#), + ) + .await; assert_eq!( next_event(&mut events_rx).await, DelegateEvent::NotificationStarted ); assert_eq!( - cell.initial_response().await.unwrap(), + cell, RuntimeResponse::Yielded { cell_id: cell_id("1"), content_items: Vec::new(), @@ -373,7 +404,7 @@ async fn termination_discards_pending_callbacks_before_responding() { ); assert_eq!( service.terminate(cell_id("1")).await.unwrap(), - WaitOutcome::LiveCell(RuntimeResponse::Terminated { + CellOutcome::LiveCell(RuntimeResponse::Terminated { cell_id: cell_id("1"), content_items: Vec::new(), }) @@ -388,15 +419,14 @@ async fn termination_discards_pending_callbacks_before_responding() { async fn termination_discards_stored_writes_before_the_next_cell_can_load_them() { let (delegate, mut events_rx) = BlockingDelegate::new(); let service = CodeModeService::with_delegate(delegate.clone()); - let cell = service - .execute(ExecuteRequest { + let created_cell_id = service + .create_cell(CreateCellRequest { enabled_tools: vec![blocking_tool()], source: r#" store("candidate", "leaked"); await tools.block({}); "# .to_string(), - yield_time_ms: Some(60_000), ..execute_request("") }) .await @@ -405,8 +435,8 @@ await tools.block({}); // Reaching the delegate proves that the store ran before execution became gated. assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted); assert_eq!( - service.terminate(cell.cell_id).await.unwrap(), - WaitOutcome::LiveCell(RuntimeResponse::Terminated { + service.terminate(created_cell_id).await.unwrap(), + CellOutcome::LiveCell(RuntimeResponse::Terminated { cell_id: cell_id("1"), content_items: Vec::new(), }) @@ -418,16 +448,12 @@ await tools.block({}); ); assert_eq!( - service - .execute(ExecuteRequest { - yield_time_ms: Some(60_000), - ..execute_request(r#"text(String(load("candidate")));"#) - }) - .await - .unwrap() - .initial_response() - .await - .unwrap(), + execute_with_yield_time( + &service, + execute_request(r#"text(String(load("candidate")));"#), + /*yield_time_ms*/ 60_000, + ) + .await, RuntimeResponse::Result { cell_id: cell_id("2"), content_items: vec![FunctionCallOutputContentItem::InputText { @@ -443,7 +469,7 @@ async fn shutdown_does_not_await_notifications_during_natural_completion() { let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new(); let service = Arc::new(CodeModeService::with_delegate(delegate)); service - .execute(execute_request(r#"notify("pending");"#)) + .create_cell(execute_request(r#"notify("pending");"#)) .await .unwrap(); @@ -469,8 +495,8 @@ async fn shutdown_does_not_await_notifications_during_natural_completion() { async fn shutdown_does_not_await_a_non_cooperative_nested_tool() { let (delegate, mut events_rx) = NeverResolvingToolDelegate::new(); let service = CodeModeService::with_delegate(delegate); - let _started = service - .execute(ExecuteRequest { + let _cell_id = service + .create_cell(CreateCellRequest { enabled_tools: vec![blocking_tool()], source: r#"await tools.block({});"#.to_string(), ..execute_request("") @@ -493,19 +519,18 @@ async fn shutdown_does_not_await_a_non_cooperative_nested_tool() { async fn termination_does_not_await_a_non_cooperative_notification() { let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new(); let service = CodeModeService::with_delegate(delegate); - let cell = service - .execute(execute_request( - r#"notify("pending"); await new Promise(() => {});"#, - )) - .await - .unwrap(); + let cell = execute( + &service, + execute_request(r#"notify("pending"); await new Promise(() => {});"#), + ) + .await; assert_eq!( next_event(&mut events_rx).await, DelegateEvent::NotificationStarted ); assert_eq!( - cell.initial_response().await.unwrap(), + cell, RuntimeResponse::Yielded { cell_id: cell_id("1"), content_items: Vec::new(), @@ -520,7 +545,7 @@ async fn termination_does_not_await_a_non_cooperative_notification() { .await .expect("termination should not await a non-cooperative notification") .unwrap(), - WaitOutcome::LiveCell(RuntimeResponse::Terminated { + CellOutcome::LiveCell(RuntimeResponse::Terminated { cell_id: cell_id("1"), content_items: Vec::new(), }) @@ -534,13 +559,10 @@ async fn termination_does_not_await_a_non_cooperative_notification() { #[tokio::test] async fn second_observer_is_rejected_without_displacing_the_first() { let service = CodeModeService::new(); - let cell = service - .execute(execute_request("await new Promise(() => {});")) - .await - .unwrap(); + let cell = execute(&service, execute_request("await new Promise(() => {});")).await; assert_eq!( - cell.initial_response().await.unwrap(), + cell, RuntimeResponse::Yielded { cell_id: cell_id("1"), content_items: Vec::new(), @@ -548,14 +570,14 @@ async fn second_observer_is_rejected_without_displacing_the_first() { ); let first_observer = service - .begin_wait(WaitRequest { + .begin_observe(ObserveRequest { cell_id: cell_id("1"), yield_time_ms: 60_000, }) .await; assert_eq!( service - .wait(WaitRequest { + .observe(ObserveRequest { cell_id: cell_id("1"), yield_time_ms: 60_000, }) @@ -570,11 +592,11 @@ async fn second_observer_is_rejected_without_displacing_the_first() { }; assert_eq!( service.terminate(cell_id("1")).await.unwrap(), - WaitOutcome::LiveCell(terminated.clone()) + CellOutcome::LiveCell(terminated.clone()) ); assert_eq!( first_observer.await.unwrap(), - WaitOutcome::LiveCell(terminated) + CellOutcome::LiveCell(terminated) ); } @@ -582,8 +604,8 @@ async fn second_observer_is_rejected_without_displacing_the_first() { async fn natural_completion_cleans_up_callbacks_before_responding() { let (delegate, mut events_rx) = BlockingDelegate::new(); let service = CodeModeService::with_delegate(delegate.clone()); - let cell = service - .execute(ExecuteRequest { + let created_cell_id = service + .create_cell(CreateCellRequest { enabled_tools: vec![blocking_tool()], source: concat!( "tools.block({});", @@ -591,7 +613,6 @@ async fn natural_completion_cleans_up_callbacks_before_responding() { "text('done');", ) .to_string(), - yield_time_ms: Some(60_000), ..execute_request("") }) .await @@ -599,14 +620,20 @@ async fn natural_completion_cleans_up_callbacks_before_responding() { assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted); assert_eq!( - cell.initial_response().await.unwrap(), - RuntimeResponse::Result { + service + .observe(ObserveRequest { + cell_id: created_cell_id, + yield_time_ms: 60_000, + }) + .await + .unwrap(), + CellOutcome::LiveCell(RuntimeResponse::Result { cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), }], error_text: None, - } + }) ); assert!(delegate.tool_future_dropped.load(Ordering::Acquire)); } diff --git a/codex-rs/code-mode/src/service_tests.rs b/codex-rs/code-mode/src/service_tests.rs index 222701a7bc..db821481cc 100644 --- a/codex-rs/code-mode/src/service_tests.rs +++ b/codex-rs/code-mode/src/service_tests.rs @@ -2,20 +2,20 @@ use std::sync::Arc; use std::time::Duration; use super::CellId; +use super::CellOutcome; use super::CodeModeNestedToolCall; use super::CodeModeService; use super::CodeModeSessionDelegate; use super::NotificationFuture; +use super::ObserveRequest; +use super::ObserveToPendingOutcome; +use super::ObserveToPendingRequest; use super::RuntimeResponse; use super::ToolInvocationFuture; -use super::WaitOutcome; -use super::WaitRequest; -use super::WaitToPendingOutcome; -use super::WaitToPendingRequest; use crate::CodeModeToolKind; -use crate::ExecuteRequest; -use crate::ExecuteToPendingOutcome; +use crate::CreateCellRequest; use crate::FunctionCallOutputContentItem; +use crate::PendingOutcome; use crate::ToolDefinition; use codex_protocol::ToolName; use pretty_assertions::assert_eq; @@ -61,13 +61,11 @@ impl CodeModeSessionDelegate for ReleasableToolDelegate { fn cell_closed(&self, _cell_id: &CellId) {} } -fn execute_request(source: &str) -> ExecuteRequest { - ExecuteRequest { +fn execute_request(source: &str) -> CreateCellRequest { + CreateCellRequest { tool_call_id: "call_1".to_string(), enabled_tools: Vec::new(), source: source.to_string(), - yield_time_ms: Some(1), - max_output_tokens: None, } } @@ -86,14 +84,38 @@ fn echo_tool() -> ToolDefinition { } } -async fn execute(service: &CodeModeService, request: ExecuteRequest) -> RuntimeResponse { +async fn execute(service: &CodeModeService, request: CreateCellRequest) -> RuntimeResponse { + execute_with_yield_time(service, request, /*yield_time_ms*/ 10_000).await +} + +async fn execute_with_yield_time( + service: &CodeModeService, + request: CreateCellRequest, + yield_time_ms: u64, +) -> RuntimeResponse { + let cell_id = service.create_cell(request).await.unwrap(); service - .execute(request) - .await - .unwrap() - .initial_response() + .observe(ObserveRequest { + cell_id, + yield_time_ms, + }) .await .unwrap() + .into() +} + +async fn create_and_observe_to_pending( + service: &CodeModeService, + request: CreateCellRequest, +) -> Result { + let cell_id = service.create_cell(request).await?; + match service + .observe_to_pending(ObserveToPendingRequest { cell_id }) + .await? + { + ObserveToPendingOutcome::LiveCell(outcome) => Ok(outcome), + ObserveToPendingOutcome::MissingCell(response) => Ok(PendingOutcome::Completed(response)), + } } #[tokio::test] @@ -102,9 +124,8 @@ async fn synchronous_exit_returns_successfully() { let response = execute( &service, - ExecuteRequest { + CreateCellRequest { source: r#"text("before"); exit(); text("after");"#.to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -129,9 +150,8 @@ async fn stored_values_are_shared_between_cells_but_not_sessions() { let write_response = execute( &first_session, - ExecuteRequest { + CreateCellRequest { source: r#"store("key", "visible");"#.to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -139,18 +159,16 @@ async fn stored_values_are_shared_between_cells_but_not_sessions() { let same_session = execute( &first_session, - ExecuteRequest { + CreateCellRequest { source: r#"text(String(load("key")));"#.to_string(), - yield_time_ms: None, ..execute_request("") }, ) .await; let other_session = execute( &second_session, - ExecuteRequest { + CreateCellRequest { source: r#"text(String(load("key")));"#.to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -190,15 +208,17 @@ async fn stored_values_are_shared_between_cells_but_not_sessions() { async fn shutdown_interrupts_cpu_bound_cells() { let service = CodeModeService::new(); - let cell = service - .execute(ExecuteRequest { + let response = execute_with_yield_time( + &service, + CreateCellRequest { source: "while (true) {}".to_string(), ..execute_request("") - }) - .await - .unwrap(); + }, + /*yield_time_ms*/ 1, + ) + .await; assert_eq!( - cell.initial_response().await.unwrap(), + response, RuntimeResponse::Yielded { cell_id: cell_id("1"), content_items: Vec::new(), @@ -217,7 +237,7 @@ async fn start_cell_rejects_new_cell_after_shutdown_begins() { service.shutdown().await.unwrap(); let error = service - .execute(execute_request("text('late');")) + .create_cell(execute_request("text('late');")) .await .err() .unwrap(); @@ -226,21 +246,22 @@ async fn start_cell_rejects_new_cell_after_shutdown_begins() { } #[tokio::test] -async fn execute_to_pending_returns_completed_for_synchronous_results() { +async fn create_and_observe_to_pending_returns_completed_for_synchronous_results() { let service = CodeModeService::new(); - let response = service - .execute_to_pending(ExecuteRequest { + let response = create_and_observe_to_pending( + &service, + CreateCellRequest { source: r#"text("done");"#.to_string(), - yield_time_ms: Some(60_000), ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await + .unwrap(); assert_eq!( response, - ExecuteToPendingOutcome::Completed(RuntimeResponse::Result { + PendingOutcome::Completed(RuntimeResponse::Result { cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "done".to_string(), @@ -251,16 +272,18 @@ async fn execute_to_pending_returns_completed_for_synchronous_results() { } #[tokio::test] -async fn execute_to_pending_returns_once_the_runtime_is_quiescent() { +async fn create_and_observe_to_pending_returns_once_the_runtime_is_quiescent() { let service = CodeModeService::new(); let response = tokio::time::timeout( Duration::from_secs(1), - service.execute_to_pending(ExecuteRequest { - source: r#"text("before"); await new Promise(() => {});"#.to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }), + create_and_observe_to_pending( + &service, + CreateCellRequest { + source: r#"text("before"); await new Promise(() => {});"#.to_string(), + ..execute_request("") + }, + ), ) .await .unwrap() @@ -268,7 +291,7 @@ async fn execute_to_pending_returns_once_the_runtime_is_quiescent() { assert_eq!( response, - ExecuteToPendingOutcome::Pending { + PendingOutcome::Pending { cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "before".to_string(), @@ -281,7 +304,7 @@ async fn execute_to_pending_returns_once_the_runtime_is_quiescent() { assert_eq!( termination, - WaitOutcome::LiveCell(RuntimeResponse::Terminated { + CellOutcome::LiveCell(RuntimeResponse::Terminated { cell_id: cell_id("1"), content_items: Vec::new(), }) @@ -289,11 +312,12 @@ async fn execute_to_pending_returns_once_the_runtime_is_quiescent() { } #[tokio::test] -async fn execute_to_pending_identifies_tool_calls_in_paused_frontier() { +async fn create_and_observe_to_pending_identifies_tool_calls_in_paused_frontier() { let service = CodeModeService::new(); - let response = service - .execute_to_pending(ExecuteRequest { + let response = create_and_observe_to_pending( + &service, + CreateCellRequest { enabled_tools: vec![echo_tool()], source: r#" await Promise.all([ @@ -302,15 +326,15 @@ await Promise.all([ ]); "# .to_string(), - yield_time_ms: Some(60_000), ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await + .unwrap(); assert_eq!( response, - ExecuteToPendingOutcome::Pending { + PendingOutcome::Pending { cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: vec!["tool-1".to_string(), "tool-2".to_string()], @@ -321,7 +345,7 @@ await Promise.all([ assert_eq!( termination, - WaitOutcome::LiveCell(RuntimeResponse::Terminated { + CellOutcome::LiveCell(RuntimeResponse::Terminated { cell_id: cell_id("1"), content_items: Vec::new(), }) @@ -329,11 +353,12 @@ await Promise.all([ } #[tokio::test] -async fn execute_to_pending_excludes_delayed_timeout_tool_calls_until_wait() { +async fn create_and_observe_to_pending_excludes_delayed_timeout_tool_calls_until_observe() { let service = CodeModeService::new(); - let initial_response = service - .execute_to_pending(ExecuteRequest { + let initial_response = create_and_observe_to_pending( + &service, + CreateCellRequest { enabled_tools: vec![echo_tool()], source: r#" setTimeout(() => { @@ -345,15 +370,15 @@ await Promise.all([ ]); "# .to_string(), - yield_time_ms: Some(60_000), ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await + .unwrap(); assert_eq!( initial_response, - ExecuteToPendingOutcome::Pending { + PendingOutcome::Pending { cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: vec!["tool-1".to_string(), "tool-2".to_string()], @@ -367,13 +392,13 @@ await Promise.all([ let resumed_response = tokio::time::timeout(Duration::from_secs(5), async { loop { let response = service - .wait_to_pending(WaitToPendingRequest { + .observe_to_pending(ObserveToPendingRequest { cell_id: cell_id("1"), }) .await?; if !matches!( &response, - WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending { + ObserveToPendingOutcome::LiveCell(PendingOutcome::Pending { pending_tool_call_ids, .. }) if pending_tool_call_ids.is_empty() @@ -388,7 +413,7 @@ await Promise.all([ assert_eq!( resumed_response, - WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending { + ObserveToPendingOutcome::LiveCell(PendingOutcome::Pending { cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: vec!["tool-3".to_string()], @@ -399,7 +424,7 @@ await Promise.all([ assert_eq!( termination, - WaitOutcome::LiveCell(RuntimeResponse::Terminated { + CellOutcome::LiveCell(RuntimeResponse::Terminated { cell_id: cell_id("1"), content_items: Vec::new(), }) @@ -407,12 +432,13 @@ await Promise.all([ } #[tokio::test] -async fn wait_to_pending_returns_after_resumed_runtime_becomes_quiescent_again() { +async fn observe_to_pending_returns_after_resumed_runtime_becomes_quiescent_again() { let delegate = Arc::new(ReleasableToolDelegate::default()); let service = CodeModeService::with_delegate(delegate.clone()); - let initial_response = service - .execute_to_pending(ExecuteRequest { + let initial_response = create_and_observe_to_pending( + &service, + CreateCellRequest { enabled_tools: vec![echo_tool()], source: r#" await tools.echo({}); @@ -420,15 +446,15 @@ text("after"); await new Promise(() => {}); "# .to_string(), - yield_time_ms: Some(60_000), ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await + .unwrap(); assert_eq!( initial_response, - ExecuteToPendingOutcome::Pending { + PendingOutcome::Pending { cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: vec!["tool-1".to_string()], @@ -439,7 +465,7 @@ await new Promise(() => {}); let resumed_response = tokio::time::timeout( Duration::from_secs(1), - service.wait_to_pending(WaitToPendingRequest { + service.observe_to_pending(ObserveToPendingRequest { cell_id: cell_id("1"), }), ) @@ -449,7 +475,7 @@ await new Promise(() => {}); assert_eq!( resumed_response, - WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending { + ObserveToPendingOutcome::LiveCell(PendingOutcome::Pending { cell_id: cell_id("1"), content_items: vec![FunctionCallOutputContentItem::InputText { text: "after".to_string(), @@ -462,7 +488,7 @@ await new Promise(() => {}); assert_eq!( termination, - WaitOutcome::LiveCell(RuntimeResponse::Terminated { + CellOutcome::LiveCell(RuntimeResponse::Terminated { cell_id: cell_id("1"), content_items: Vec::new(), }) @@ -470,27 +496,28 @@ await new Promise(() => {}); } #[tokio::test] -async fn wait_to_pending_returns_completed_after_resumed_runtime_finishes() { +async fn observe_to_pending_returns_completed_after_resumed_runtime_finishes() { let delegate = Arc::new(ReleasableToolDelegate::default()); let service = CodeModeService::with_delegate(delegate.clone()); - let initial_response = service - .execute_to_pending(ExecuteRequest { + let initial_response = create_and_observe_to_pending( + &service, + CreateCellRequest { enabled_tools: vec![echo_tool()], source: r#" await tools.echo({}); text("done"); "# .to_string(), - yield_time_ms: Some(60_000), ..execute_request("") - }) - .await - .unwrap(); + }, + ) + .await + .unwrap(); assert_eq!( initial_response, - ExecuteToPendingOutcome::Pending { + PendingOutcome::Pending { cell_id: cell_id("1"), content_items: Vec::new(), pending_tool_call_ids: vec!["tool-1".to_string()], @@ -501,7 +528,7 @@ text("done"); let resumed_response = tokio::time::timeout( Duration::from_secs(1), - service.wait_to_pending(WaitToPendingRequest { + service.observe_to_pending(ObserveToPendingRequest { cell_id: cell_id("1"), }), ) @@ -511,15 +538,13 @@ text("done"); assert_eq!( resumed_response, - WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Completed( - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "done".to_string(), - }], - error_text: None, - } - )) + ObserveToPendingOutcome::LiveCell(PendingOutcome::Completed(RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "done".to_string(), + }], + error_text: None, + })) ); } @@ -529,9 +554,8 @@ async fn v8_console_is_not_exposed_on_global_this() { let response = execute( &service, - ExecuteRequest { + CreateCellRequest { source: r#"text(String(Object.hasOwn(globalThis, "console")));"#.to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -555,7 +579,7 @@ async fn date_locale_string_formats_with_icu_data() { let response = execute( &service, - ExecuteRequest { + CreateCellRequest { source: r#" const value = new Date("2025-01-02T03:04:05Z") .toLocaleString("fr-FR", { @@ -571,7 +595,6 @@ const value = new Date("2025-01-02T03:04:05Z") text(value); "# .to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -595,7 +618,7 @@ async fn intl_date_time_format_formats_with_icu_data() { let response = execute( &service, - ExecuteRequest { + CreateCellRequest { source: r#" const formatter = new Intl.DateTimeFormat("fr-FR", { weekday: "long", @@ -610,7 +633,6 @@ const formatter = new Intl.DateTimeFormat("fr-FR", { text(formatter.format(new Date("2025-01-02T03:04:05Z"))); "# .to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -634,7 +656,7 @@ async fn output_helpers_return_undefined() { let response = execute( &service, - ExecuteRequest { + CreateCellRequest { source: r#" const returnsUndefined = [ text("first"), @@ -644,7 +666,6 @@ const returnsUndefined = [ text(JSON.stringify(returnsUndefined)); "# .to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -677,7 +698,7 @@ async fn image_helper_accepts_raw_mcp_image_block_with_original_detail() { let response = execute( &service, - ExecuteRequest { + CreateCellRequest { source: r#" image({ type: "image", @@ -687,7 +708,6 @@ image({ }); "# .to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -712,7 +732,7 @@ async fn generated_image_helper_appends_image_and_output_hint() { let response = execute( &service, - ExecuteRequest { + CreateCellRequest { source: r#" generatedImage({ image_url: "data:image/png;base64,AAA", @@ -720,7 +740,6 @@ generatedImage({ }); "# .to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -750,7 +769,7 @@ async fn image_helper_second_arg_overrides_explicit_object_detail() { let response = execute( &service, - ExecuteRequest { + CreateCellRequest { source: r#" image( { @@ -761,7 +780,6 @@ image( ); "# .to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -786,7 +804,7 @@ async fn image_helper_second_arg_overrides_raw_mcp_image_detail() { let response = execute( &service, - ExecuteRequest { + CreateCellRequest { source: r#" image( { @@ -799,7 +817,6 @@ image( ); "# .to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -824,7 +841,7 @@ async fn image_helper_accepts_low_detail() { let response = execute( &service, - ExecuteRequest { + CreateCellRequest { source: r#" image({ image_url: "data:image/png;base64,AAA", @@ -832,7 +849,6 @@ image({ }); "# .to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -865,9 +881,8 @@ async fn image_helpers_reject_remote_urls() { let response = execute( &service, - ExecuteRequest { + CreateCellRequest { source, - yield_time_ms: None, ..execute_request("") }, ) @@ -893,7 +908,7 @@ async fn image_helper_rejects_unsupported_detail() { let response = execute( &service, - ExecuteRequest { + CreateCellRequest { source: r#" image({ image_url: "data:image/png;base64,AAA", @@ -901,7 +916,6 @@ image({ }); "# .to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -923,7 +937,7 @@ async fn image_helper_rejects_raw_mcp_result_container() { let response = execute( &service, - ExecuteRequest { + CreateCellRequest { source: r#" image({ content: [ @@ -938,7 +952,6 @@ image({ }); "# .to_string(), - yield_time_ms: None, ..execute_request("") }, ) @@ -957,11 +970,11 @@ image({ } #[tokio::test] -async fn wait_reports_missing_cell_separately_from_runtime_results() { +async fn observe_reports_missing_cell_separately_from_runtime_results() { let service = CodeModeService::new(); let response = service - .wait(WaitRequest { + .observe(ObserveRequest { cell_id: cell_id("missing"), yield_time_ms: 1, }) @@ -970,7 +983,7 @@ async fn wait_reports_missing_cell_separately_from_runtime_results() { assert_eq!( response, - WaitOutcome::MissingCell(RuntimeResponse::Result { + CellOutcome::MissingCell(RuntimeResponse::Result { cell_id: cell_id("missing"), content_items: Vec::new(), error_text: Some("exec cell missing not found".to_string()), diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 2e46c1bd6b..8410e08c33 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -9,6 +9,7 @@ use codex_tools::ToolName; use codex_tools::ToolSpec; use super::ExecContext; +use super::InitialObservationGuard; use super::PUBLIC_TOOL_NAME; use super::handle_runtime_response; use super::is_exec_tool_name; @@ -39,20 +40,17 @@ impl CodeModeExecuteHandler { let enabled_tools = codex_tools::collect_code_mode_tool_definitions(&self.nested_tool_specs); let started_at = std::time::Instant::now(); - let started_cell = exec + let cell_id = exec .session .services .code_mode_service - .execute(codex_code_mode::ExecuteRequest { + .create_cell(codex_code_mode::CreateCellRequest { tool_call_id: call_id.clone(), enabled_tools, source: args.code.clone(), - yield_time_ms: args.yield_time_ms, - max_output_tokens: args.max_output_tokens, }) .await .map_err(FunctionCallError::RespondToModel)?; - let cell_id = started_cell.cell_id.clone(); let runtime_cell_id = cell_id.to_string(); let code_cell_trace = exec .session @@ -64,14 +62,34 @@ impl CodeModeExecuteHandler { call_id.as_str(), args.code.as_str(), ); + let mut initial_observation_guard = InitialObservationGuard::new( + &exec.session.services.code_mode_service, + cell_id.clone(), + code_cell_trace.clone(), + ) + .map_err(FunctionCallError::RespondToModel)?; exec.session .services .code_mode_service .mark_cell_ready_for_dispatch(&cell_id); - let response = started_cell - .initial_response() - .await - .map_err(FunctionCallError::RespondToModel)?; + let observed = exec + .session + .services + .code_mode_service + .observe(codex_code_mode::ObserveRequest { + cell_id: cell_id.clone(), + yield_time_ms: args + .yield_time_ms + .unwrap_or(codex_code_mode::DEFAULT_EXEC_YIELD_TIME_MS), + }) + .await; + let response: codex_code_mode::RuntimeResponse = match observed { + Ok(outcome) => outcome.into(), + Err(error) => { + initial_observation_guard.finish_with_error(&error); + return Err(FunctionCallError::RespondToModel(error)); + } + }; // Record the raw runtime boundary. The model-visible custom-tool output // is produced by `handle_runtime_response` and later linked through // `CodeCell.output_item_ids` in the reduced trace. @@ -85,6 +103,7 @@ impl CodeModeExecuteHandler { .code_mode_service .finish_cell_dispatch(&cell_id); } + initial_observation_guard.disarm(); handle_runtime_response(&exec, response, args.max_output_tokens, started_at) .await .map_err(FunctionCallError::RespondToModel) diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 09918033e6..ee5527ea4c 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -63,6 +63,63 @@ pub(crate) struct CodeModeService { dispatch_broker: Arc, } +struct InitialObservationGuard { + session: Arc, + dispatch_broker: Arc, + cell_id: Option, + code_cell_trace: codex_rollout_trace::CodeCellTraceContext, +} + +impl InitialObservationGuard { + fn new( + service: &CodeModeService, + cell_id: CellId, + code_cell_trace: codex_rollout_trace::CodeCellTraceContext, + ) -> Result { + Ok(Self { + session: Arc::clone(service.session()?), + dispatch_broker: Arc::clone(&service.dispatch_broker), + cell_id: Some(cell_id), + code_cell_trace, + }) + } + + fn finish_with_error(&mut self, error_text: &str) { + let Some(cell_id) = self.cell_id.take() else { + return; + }; + let response = RuntimeResponse::Result { + cell_id: cell_id.clone(), + content_items: Vec::new(), + error_text: Some(error_text.to_string()), + }; + self.code_cell_trace.record_initial_response(&response); + self.code_cell_trace.record_ended(&response); + self.dispatch_broker.close_cell(&cell_id); + + let session = Arc::clone(&self.session); + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + tracing::warn!("unable to terminate code-mode cell {cell_id}: no Tokio runtime"); + return; + }; + let _termination_task = runtime.spawn(async move { + if let Err(error) = session.terminate(cell_id.clone()).await { + tracing::warn!("failed to terminate code-mode cell {cell_id}: {error}"); + } + }); + } + + fn disarm(&mut self) { + self.cell_id = None; + } +} + +impl Drop for InitialObservationGuard { + fn drop(&mut self) { + self.finish_with_error("initial code-mode observation cancelled"); + } +} + impl CodeModeService { pub(crate) fn new() -> Self { let dispatch_broker = Arc::new(CodeModeDispatchBroker::new()); @@ -74,24 +131,24 @@ impl CodeModeService { } } - pub(crate) async fn execute( + pub(crate) async fn create_cell( &self, - request: codex_code_mode::ExecuteRequest, - ) -> Result { - self.session()?.execute(request).await + request: codex_code_mode::CreateCellRequest, + ) -> Result { + self.session()?.create_cell(request).await } - pub(crate) async fn wait( + pub(crate) async fn observe( &self, - request: codex_code_mode::WaitRequest, - ) -> Result { - self.session()?.wait(request).await + request: codex_code_mode::ObserveRequest, + ) -> Result { + self.session()?.observe(request).await } pub(crate) async fn terminate( &self, cell_id: CellId, - ) -> Result { + ) -> Result { self.session()?.terminate(cell_id).await } @@ -322,13 +379,71 @@ fn build_freeform_tool_payload( #[cfg(test)] mod tests { + use super::CodeModeDispatchBroker; + use super::InitialObservationGuard; use super::build_nested_tool_payload; use super::truncate_code_mode_result; use crate::tools::context::ToolPayload; + use codex_code_mode::CellId; + use codex_code_mode::CellOutcome; + use codex_code_mode::CodeModeSession; + use codex_code_mode::CodeModeSessionResultFuture; use codex_code_mode::CodeModeToolKind; + use codex_code_mode::CreateCellRequest; + use codex_code_mode::ObserveRequest; + use codex_code_mode::RuntimeResponse; use codex_protocol::models::FunctionCallOutputContentItem; use codex_tools::ToolName; use serde_json::json; + use std::sync::Arc; + use std::sync::Mutex; + + struct RecordingCodeModeSession { + terminated_cells: Mutex>, + termination: tokio::sync::Notify, + } + + impl CodeModeSession for RecordingCodeModeSession { + fn is_alive(&self) -> bool { + true + } + + fn create_cell<'a>( + &'a self, + _request: CreateCellRequest, + ) -> CodeModeSessionResultFuture<'a, CellId> { + Box::pin(async { panic!("unexpected cell creation") }) + } + + fn observe<'a>( + &'a self, + _request: ObserveRequest, + ) -> CodeModeSessionResultFuture<'a, CellOutcome> { + Box::pin(async { panic!("unexpected cell observation") }) + } + + fn terminate<'a>( + &'a self, + cell_id: CellId, + ) -> CodeModeSessionResultFuture<'a, CellOutcome> { + Box::pin(async move { + self.terminated_cells + .lock() + .expect("terminated-cell lock should not be poisoned") + .push(cell_id.clone()); + self.termination.notify_one(); + Ok(CellOutcome::MissingCell(RuntimeResponse::Result { + cell_id, + content_items: Vec::new(), + error_text: Some("terminated by test".to_string()), + })) + }) + } + + fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { + Box::pin(async { Ok(()) }) + } + } #[test] fn build_nested_tool_payload_uses_function_kind() { @@ -382,4 +497,39 @@ mod tests { }] ); } + + #[tokio::test] + async fn dropping_initial_observation_guard_terminates_the_cell() { + let session = Arc::new(RecordingCodeModeSession { + terminated_cells: Mutex::new(Vec::new()), + termination: tokio::sync::Notify::new(), + }); + let session_trait: Arc = session.clone(); + let dispatch_broker = Arc::new(CodeModeDispatchBroker::new()); + let cell_id = CellId::new("cell-1".to_string()); + dispatch_broker.mark_cell_ready_for_dispatch(&cell_id); + let code_cell_trace = codex_rollout_trace::ThreadTraceContext::disabled() + .code_cell_trace_context("turn-1", cell_id.as_str()); + + drop(InitialObservationGuard { + session: session_trait, + dispatch_broker, + cell_id: Some(cell_id.clone()), + code_cell_trace, + }); + + tokio::time::timeout( + std::time::Duration::from_secs(1), + session.termination.notified(), + ) + .await + .expect("guard should schedule cell termination"); + assert_eq!( + *session + .terminated_cells + .lock() + .expect("terminated-cell lock should not be poisoned"), + vec![cell_id] + ); + } } diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index 953391a573..a9b985f155 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -89,14 +89,14 @@ impl CodeModeWaitHandler { exec.session .services .code_mode_service - .wait(codex_code_mode::WaitRequest { + .observe(codex_code_mode::ObserveRequest { cell_id, yield_time_ms: args.yield_time_ms, }) .await } .map_err(FunctionCallError::RespondToModel)?; - if let codex_code_mode::WaitOutcome::LiveCell(response) = &wait_response + if let codex_code_mode::CellOutcome::LiveCell(response) = &wait_response && !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) { // Only a live-cell wait can close a CodeCell. A missing diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 0a765a0e26..584bb7d588 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -2530,6 +2530,131 @@ text({ json: true }); Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_exec_observes_a_synchronously_completed_cell() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let (_test, second_mock) = run_code_mode_turn( + &server, + "use exec to return immediately", + r#"text("created before observation");"#, + ) + .await?; + + let request = second_mock.single_request(); + let (output, success) = custom_tool_output_body_and_success(&request, "call-1"); + assert_ne!( + success, + Some(false), + "exec call failed unexpectedly: {output}" + ); + assert_eq!(output, "created before observation"); + + Ok(()) +} + +#[cfg_attr(windows, ignore = "no exec_command on Windows")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn interrupting_the_initial_exec_observation_terminates_the_created_cell() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let mut builder = test_codex().with_config(|config| { + let _ = config.features.enable(Feature::CodeMode); + }); + let test = builder.build(&server).await?; + let started_marker = test.workspace_path("code-mode-initial-observation-started"); + let started_marker_quoted = shlex::try_join([started_marker.to_string_lossy().as_ref()])?; + let start_command = format!("printf started > {started_marker_quoted}"); + let code = format!( + r#"// @exec: {{"yield_time_ms": 60000}} +await tools.exec_command({{ cmd: {start_command:?} }}); +store("initial_observation_cancel_probe", "leaked"); +await new Promise((resolve) => setTimeout(resolve, 2000)); +"# + ); + + responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_custom_tool_call("call-1", "exec", &code), + ev_completed("resp-1"), + ]), + ) + .await; + + let cwd = test.config.cwd.clone(); + let (sandbox_policy, permission_profile) = + turn_permission_fields(PermissionProfile::Disabled, cwd.as_path()); + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "start an exec whose initial observation will be interrupted".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: codex_protocol::protocol::ThreadSettingsOverrides { + approval_policy: Some(AskForApproval::Never), + sandbox_policy: Some(sandbox_policy), + permission_profile, + ..Default::default() + }, + }) + .await?; + + tokio::time::timeout(Duration::from_secs(5), async { + while !started_marker.exists() { + tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await; + } + }) + .await?; + test.codex.submit(Op::Interrupt).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnAborted(_)) + }) + .await; + + // The sentinel is already pending when the observation is interrupted. Without the guard + // around create + initial observe, the detached cell finishes this timer and commits it after + // the interrupted turn has gone away. + tokio::time::sleep(Duration::from_millis(/*millis*/ 2_250)).await; + + responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-2"), + ev_custom_tool_call( + "call-2", + "exec", + r#"text(String(load("initial_observation_cancel_probe")));"#, + ), + ev_completed("resp-2"), + ]), + ) + .await; + let probe_completion = responses::mount_sse_once( + &server, + sse(vec![ + ev_assistant_message("msg-2", "done"), + ev_completed("resp-3"), + ]), + ) + .await; + + test.submit_turn("check whether the interrupted cell committed state") + .await?; + let request = probe_completion.single_request(); + let (output, success) = custom_tool_output_body_and_success(&request, "call-2"); + assert_ne!(success, Some(false), "probe exec failed: {output}"); + assert_eq!(output, "undefined"); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_can_resume_after_set_timeout() -> Result<()> { skip_if_no_network!(Ok(()));