diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 94dac45699..a8a5948da8 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1914,6 +1914,7 @@ dependencies = [ "codex-backend-client", "codex-chatgpt", "codex-cloud-config", + "codex-code-mode-runtime", "codex-config", "codex-core", "codex-core-plugins", @@ -2428,12 +2429,24 @@ name = "codex-code-mode" version = "0.0.0" dependencies = [ "codex-protocol", - "deno_core_icudata", "pretty_assertions", "serde", "serde_json", "tokio", "tokio-util", +] + +[[package]] +name = "codex-code-mode-runtime" +version = "0.0.0" +dependencies = [ + "codex-code-mode", + "codex-protocol", + "deno_core_icudata", + "pretty_assertions", + "serde_json", + "tokio", + "tokio-util", "tracing", "v8", ] @@ -2526,6 +2539,7 @@ dependencies = [ "codex-apply-patch", "codex-async-utils", "codex-code-mode", + "codex-code-mode-runtime", "codex-config", "codex-connectors", "codex-core-plugins", @@ -3182,6 +3196,7 @@ version = "0.0.0" dependencies = [ "anyhow", "codex-arg0", + "codex-code-mode-runtime", "codex-config", "codex-core", "codex-exec-server", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 482951fb92..f0bc0a9f20 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -21,6 +21,7 @@ members = [ "install-context", "codex-backend-openapi-models", "code-mode", + "code-mode-runtime", "cloud-config", "cloud-tasks", "cloud-tasks-client", @@ -154,6 +155,7 @@ codex-cloud-config = { path = "cloud-config" } codex-cloud-tasks-client = { path = "cloud-tasks-client" } codex-cloud-tasks-mock-client = { path = "cloud-tasks-mock-client" } codex-code-mode = { path = "code-mode" } +codex-code-mode-runtime = { path = "code-mode-runtime" } codex-config = { path = "config" } codex-connectors = { path = "connectors" } codex-core = { path = "core" } diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index 37cdf219bf..7566ed3dd5 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -33,6 +33,7 @@ axum = { workspace = true, default-features = false, features = [ codex-analytics = { workspace = true } codex-arg0 = { workspace = true } codex-cloud-config = { workspace = true } +codex-code-mode-runtime = { workspace = true } codex-config = { workspace = true } codex-core = { workspace = true } codex-core-plugins = { workspace = true } diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 0996c409bc..374dd8648e 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -187,6 +187,7 @@ mod tests { auth_manager.clone(), SessionSource::Exec, Arc::new(EnvironmentManager::default_for_tests()), + Arc::new(codex_code_mode_runtime::InProcessCodeModeSessionProvider), thread_extensions( guardian_agent_spawner(thread_manager.clone()), Arc::new(NoopExtensionEventSink), diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 84f22809ee..1e45624d2c 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -311,6 +311,7 @@ impl MessageProcessor { auth_manager.clone(), session_source, environment_manager, + Arc::new(codex_code_mode_runtime::InProcessCodeModeSessionProvider), thread_extensions( guardian_agent_spawner(thread_manager.clone()), app_server_extension_event_sink(outgoing.clone()), diff --git a/codex-rs/code-mode-runtime/BUILD.bazel b/codex-rs/code-mode-runtime/BUILD.bazel new file mode 100644 index 0000000000..b34fb8e108 --- /dev/null +++ b/codex-rs/code-mode-runtime/BUILD.bazel @@ -0,0 +1,11 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "code-mode-runtime", + crate_name = "codex_code_mode_runtime", + crate_features = select({ + "@rules_rs//rs/experimental/platforms/constraints:windows_msvc": [], + "//conditions:default": ["sandbox"], + }), + deps_extra = ["@crates//:v8"], +) diff --git a/codex-rs/code-mode-runtime/Cargo.toml b/codex-rs/code-mode-runtime/Cargo.toml new file mode 100644 index 0000000000..b5985577b6 --- /dev/null +++ b/codex-rs/code-mode-runtime/Cargo.toml @@ -0,0 +1,29 @@ +[package] +edition.workspace = true +license.workspace = true +name = "codex-code-mode-runtime" +version.workspace = true + +[lib] +doctest = false +name = "codex_code_mode_runtime" +path = "src/lib.rs" + +[features] +sandbox = ["v8/v8_enable_sandbox"] + +[lints] +workspace = true + +[dependencies] +codex-code-mode = { workspace = true } +codex-protocol = { workspace = true } +deno_core_icudata = { workspace = true } +serde_json = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } +tokio-util = { workspace = true, features = ["rt"] } +tracing = { workspace = true } +v8 = { workspace = true } + +[dev-dependencies] +pretty_assertions = { workspace = true } diff --git a/codex-rs/code-mode-runtime/src/lib.rs b/codex-rs/code-mode-runtime/src/lib.rs new file mode 100644 index 0000000000..0b2613022f --- /dev/null +++ b/codex-rs/code-mode-runtime/src/lib.rs @@ -0,0 +1,6 @@ +mod runtime; +mod service; + +pub use codex_code_mode::*; +pub use service::CodeModeService; +pub use service::InProcessCodeModeSessionProvider; diff --git a/codex-rs/code-mode/src/runtime/callbacks.rs b/codex-rs/code-mode-runtime/src/runtime/callbacks.rs similarity index 99% rename from codex-rs/code-mode/src/runtime/callbacks.rs rename to codex-rs/code-mode-runtime/src/runtime/callbacks.rs index 7d8a28662a..f81f446e35 100644 --- a/codex-rs/code-mode/src/runtime/callbacks.rs +++ b/codex-rs/code-mode-runtime/src/runtime/callbacks.rs @@ -1,4 +1,4 @@ -use crate::response::FunctionCallOutputContentItem; +use crate::FunctionCallOutputContentItem; use super::EXIT_SENTINEL; use super::RuntimeEvent; diff --git a/codex-rs/code-mode/src/runtime/globals.rs b/codex-rs/code-mode-runtime/src/runtime/globals.rs similarity index 100% rename from codex-rs/code-mode/src/runtime/globals.rs rename to codex-rs/code-mode-runtime/src/runtime/globals.rs diff --git a/codex-rs/code-mode/src/runtime/mod.rs b/codex-rs/code-mode-runtime/src/runtime/mod.rs similarity index 78% rename from codex-rs/code-mode/src/runtime/mod.rs rename to codex-rs/code-mode-runtime/src/runtime/mod.rs index 21757328ee..e9657018ad 100644 --- a/codex-rs/code-mode/src/runtime/mod.rs +++ b/codex-rs/code-mode-runtime/src/runtime/mod.rs @@ -10,124 +10,17 @@ use std::sync::mpsc as std_mpsc; use std::thread; use codex_protocol::ToolName; -use serde::Serialize; use serde_json::Value as JsonValue; use tokio::sync::mpsc; -use crate::description::CodeModeToolKind; -use crate::description::EnabledToolMetadata; -use crate::description::ToolDefinition; -use crate::description::enabled_tool_metadata; -use crate::response::FunctionCallOutputContentItem; -use crate::service::CellId; - -pub const DEFAULT_EXEC_YIELD_TIME_MS: u64 = 10_000; -pub const DEFAULT_WAIT_YIELD_TIME_MS: u64 = 10_000; -pub const DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL: usize = 10_000; +use crate::EnabledToolMetadata; +use crate::enabled_tool_metadata; +pub(crate) use codex_code_mode::CodeModeToolKind; +pub(crate) use codex_code_mode::ExecuteRequest; +pub(crate) use codex_code_mode::FunctionCallOutputContentItem; +pub(crate) use codex_code_mode::ToolDefinition; const EXIT_SENTINEL: &str = "__codex_code_mode_exit__"; -#[derive(Clone, Debug)] -pub struct ExecuteRequest { - 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)] -pub struct WaitRequest { - pub cell_id: CellId, - pub yield_time_ms: u64, -} - -#[derive(Clone, Debug)] -pub struct WaitToPendingRequest { - pub cell_id: CellId, -} - -/// Result of waiting on a code-mode cell. -/// -/// The wrapped `RuntimeResponse` is the model-facing wait result. The enum -/// variant carries the extra lifecycle provenance that `RuntimeResponse` cannot: -/// a failed real cell and a missing-cell wait both use -/// `RuntimeResponse::Result { error_text: Some(..), .. }`, but only the former -/// should be treated as a code-cell lifecycle event. -#[derive(Debug, PartialEq)] -pub enum WaitOutcome { - /// The requested code cell was live when the wait command was accepted. - LiveCell(RuntimeResponse), - /// The requested code cell was not live. - MissingCell(RuntimeResponse), -} - -/// Result of executing a code-mode cell until it either completes or reaches a -/// quiescent pending state. -#[derive(Debug, PartialEq)] -pub enum ExecuteToPendingOutcome { - /// The cell is waiting for more runtime input after draining the runtime - /// input queue that was ready at the pending boundary. - Pending { - cell_id: CellId, - content_items: Vec, - /// Runtime tool-call ids emitted before this paused execution frontier - /// sealed. Hosts can use these ids to drain their tool-call transport - /// before surfacing the pending boundary to callers. - pending_tool_call_ids: Vec, - }, - /// The cell reached a terminal runtime response before going pending. - Completed(RuntimeResponse), -} - -/// Result of resuming a live code-mode cell until it completes or becomes -/// quiescent again. -#[derive(Debug, PartialEq)] -pub enum WaitToPendingOutcome { - /// The requested code cell was live when the wait command was accepted. - LiveCell(ExecuteToPendingOutcome), - /// The requested code cell was not live. - MissingCell(RuntimeResponse), -} - -impl From for RuntimeResponse { - fn from(outcome: WaitOutcome) -> Self { - match outcome { - WaitOutcome::LiveCell(response) | WaitOutcome::MissingCell(response) => response, - } - } -} - -#[derive(Debug, PartialEq, Serialize)] -pub enum RuntimeResponse { - Yielded { - cell_id: CellId, - content_items: Vec, - }, - Terminated { - cell_id: CellId, - content_items: Vec, - }, - Result { - cell_id: CellId, - content_items: Vec, - error_text: Option, - }, -} - -/// Nested tool request emitted by one code-mode cell. -/// -/// Code mode owns the per-cell runtime id. Hosts should preserve it for -/// provenance/debugging, but should still assign their own runtime tool call id -/// if their tool-call graph requires globally unique ids. -#[derive(Debug)] -pub struct CodeModeNestedToolCall { - pub cell_id: CellId, - pub runtime_tool_call_id: String, - pub tool_name: ToolName, - pub tool_kind: CodeModeToolKind, - pub input: Option, -} - #[derive(Debug)] pub(crate) enum RuntimeCommand { ToolResponse { id: String, result: JsonValue }, diff --git a/codex-rs/code-mode/src/runtime/module_loader.rs b/codex-rs/code-mode-runtime/src/runtime/module_loader.rs similarity index 100% rename from codex-rs/code-mode/src/runtime/module_loader.rs rename to codex-rs/code-mode-runtime/src/runtime/module_loader.rs diff --git a/codex-rs/code-mode/src/runtime/timers.rs b/codex-rs/code-mode-runtime/src/runtime/timers.rs similarity index 100% rename from codex-rs/code-mode/src/runtime/timers.rs rename to codex-rs/code-mode-runtime/src/runtime/timers.rs diff --git a/codex-rs/code-mode/src/runtime/value.rs b/codex-rs/code-mode-runtime/src/runtime/value.rs similarity index 98% rename from codex-rs/code-mode/src/runtime/value.rs rename to codex-rs/code-mode-runtime/src/runtime/value.rs index 8d76a832d3..28ef7c4215 100644 --- a/codex-rs/code-mode/src/runtime/value.rs +++ b/codex-rs/code-mode-runtime/src/runtime/value.rs @@ -1,8 +1,8 @@ use serde_json::Value as JsonValue; -use crate::response::DEFAULT_IMAGE_DETAIL; -use crate::response::FunctionCallOutputContentItem; -use crate::response::ImageDetail; +use crate::DEFAULT_IMAGE_DETAIL; +use crate::FunctionCallOutputContentItem; +use crate::ImageDetail; const IMAGE_HELPER_EXPECTS_MESSAGE: &str = "image expects a non-empty image URL string, an object with image_url and optional detail, or a raw MCP image block"; const CODEX_IMAGE_DETAIL_META_KEY: &str = "codex/imageDetail"; diff --git a/codex-rs/code-mode-runtime/src/service.rs b/codex-rs/code-mode-runtime/src/service.rs new file mode 100644 index 0000000000..31fc76d6cd --- /dev/null +++ b/codex-rs/code-mode-runtime/src/service.rs @@ -0,0 +1,1715 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use serde_json::Value as JsonValue; +use tokio::sync::Mutex; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio::task::JoinSet; +use tokio_util::sync::CancellationToken; +use tracing::warn; + +use crate::CellId; +use crate::CodeModeNestedToolCall; +use crate::CodeModeSession; +use crate::CodeModeSessionDelegate; +use crate::CodeModeSessionProvider; +use crate::CodeModeSessionProviderFuture; +use crate::CodeModeSessionResultFuture; +use crate::DEFAULT_EXEC_YIELD_TIME_MS; +use crate::ExecuteRequest; +use crate::ExecuteToPendingOutcome; +use crate::FunctionCallOutputContentItem; +use crate::NoopCodeModeSessionDelegate; +use crate::RuntimeResponse; +use crate::StartedCell; +use crate::WaitOutcome; +use crate::WaitRequest; +use crate::WaitToPendingOutcome; +use crate::WaitToPendingRequest; +use crate::runtime::PendingRuntimeMode; +use crate::runtime::RuntimeCommand; +use crate::runtime::RuntimeControlCommand; +use crate::runtime::RuntimeEvent; +use crate::runtime::spawn_runtime; + +#[derive(Default)] +pub struct InProcessCodeModeSessionProvider; + +impl CodeModeSessionProvider for InProcessCodeModeSessionProvider { + fn create_session<'a>( + &'a self, + delegate: Arc, + ) -> CodeModeSessionProviderFuture<'a> { + Box::pin(async move { + let session: Arc = + Arc::new(CodeModeService::with_delegate(delegate)); + Ok(session) + }) + } +} + +#[derive(Clone)] +struct CellHandle { + control_tx: mpsc::UnboundedSender, + runtime_tx: std::sync::mpsc::Sender, + cancellation_token: CancellationToken, +} + +struct Inner { + stored_values: Mutex>, + cells: Mutex>, + delegate: Arc, + shutting_down: AtomicBool, + next_cell_id: AtomicU64, +} + +pub struct CodeModeService { + inner: Arc, +} + +impl CodeModeService { + pub fn new() -> Self { + Self::with_delegate(Arc::new(NoopCodeModeSessionDelegate)) + } + + pub fn with_delegate(delegate: Arc) -> Self { + Self { + inner: Arc::new(Inner { + stored_values: Mutex::new(HashMap::new()), + cells: Mutex::new(HashMap::new()), + delegate, + shutting_down: AtomicBool::new(false), + next_cell_id: AtomicU64::new(1), + }), + } + } + + fn allocate_cell_id(&self) -> CellId { + CellId::new( + self.inner + .next_cell_id + .fetch_add(1, Ordering::Relaxed) + .to_string(), + ) + } + + pub async fn execute(&self, request: ExecuteRequest) -> Result { + if self.inner.shutting_down.load(Ordering::Acquire) { + return Err("code mode session is shutting down".to_string()); + } + let initial_yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS); + let (response_tx, response_rx) = oneshot::channel(); + let cell_id = self.allocate_cell_id(); + self.start_cell( + cell_id.clone(), + request, + CellResponseSender::Runtime(response_tx), + Some(initial_yield_time_ms), + PendingRuntimeMode::Continue, + ) + .await?; + + Ok(StartedCell::new(cell_id, response_rx)) + } + + pub async fn execute_to_pending( + &self, + request: ExecuteRequest, + ) -> Result { + let (response_tx, response_rx) = oneshot::channel(); + let cell_id = self.allocate_cell_id(); + self.start_cell( + cell_id, + request, + CellResponseSender::ExecuteToPending(response_tx), + /*initial_yield_time_ms*/ None, + PendingRuntimeMode::PauseUntilResumed, + ) + .await?; + + response_rx + .await + .map_err(|_| "exec runtime ended unexpectedly".to_string()) + } + + async fn start_cell( + &self, + cell_id: CellId, + request: ExecuteRequest, + initial_response_tx: CellResponseSender, + initial_yield_time_ms: Option, + pending_mode: PendingRuntimeMode, + ) -> Result<(), String> { + let (event_tx, event_rx) = mpsc::unbounded_channel(); + let (control_tx, control_rx) = mpsc::unbounded_channel(); + let stored_values = self.inner.stored_values.lock().await.clone(); + let cancellation_token = CancellationToken::new(); + let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = { + let mut cells = self.inner.cells.lock().await; + if self.inner.shutting_down.load(Ordering::Acquire) { + return Err("code mode session is shutting down".to_string()); + } + if cells.contains_key(&cell_id) { + return Err(format!("exec cell {cell_id} already exists")); + } + + let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = + spawn_runtime(stored_values, request, event_tx, pending_mode)?; + + cells.insert( + cell_id.clone(), + CellHandle { + control_tx, + runtime_tx: runtime_tx.clone(), + cancellation_token: cancellation_token.clone(), + }, + ); + (runtime_tx, runtime_control_tx, runtime_terminate_handle) + }; + + tokio::spawn(run_cell_control( + Arc::clone(&self.inner), + CellControlContext { + cell_id, + runtime_tx, + runtime_control_tx, + pending_mode, + runtime_terminate_handle, + cancellation_token, + }, + event_rx, + control_rx, + initial_response_tx, + initial_yield_time_ms, + )); + + Ok(()) + } + + pub async fn wait(&self, request: WaitRequest) -> Result { + let WaitRequest { + cell_id, + yield_time_ms, + } = request; + let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); + let Some(handle) = handle else { + return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); + }; + let (response_tx, response_rx) = oneshot::channel(); + let control_message = CellControlCommand::Poll { + yield_time_ms, + response_tx, + }; + if handle.control_tx.send(control_message).is_err() { + return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); + } + match response_rx.await { + Ok(response) => Ok(WaitOutcome::LiveCell(response)), + Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))), + } + } + + pub async fn terminate(&self, cell_id: CellId) -> Result { + let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); + let Some(handle) = handle else { + return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); + }; + let (response_tx, response_rx) = oneshot::channel(); + if handle + .control_tx + .send(CellControlCommand::Terminate { response_tx }) + .is_err() + { + return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); + } + match response_rx.await { + Ok(response) => Ok(WaitOutcome::LiveCell(response)), + Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))), + } + } + + pub async fn wait_to_pending( + &self, + request: WaitToPendingRequest, + ) -> Result { + let cell_id = request.cell_id; + let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); + let Some(handle) = handle else { + return Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( + cell_id, + ))); + }; + let (response_tx, response_rx) = oneshot::channel(); + if handle + .control_tx + .send(CellControlCommand::PollToPending { response_tx }) + .is_err() + { + return Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( + cell_id, + ))); + } + match response_rx.await { + Ok(response) => Ok(WaitToPendingOutcome::LiveCell(response)), + Err(_) => Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( + cell_id, + ))), + } + } + + pub async fn shutdown(&self) -> Result<(), String> { + self.inner.shutting_down.store(true, Ordering::Release); + let handles = self + .inner + .cells + .lock() + .await + .values() + .cloned() + .collect::>(); + for handle in handles { + handle.cancellation_token.cancel(); + let (response_tx, _response_rx) = oneshot::channel(); + let _ = handle + .control_tx + .send(CellControlCommand::Terminate { response_tx }); + let _ = handle.runtime_tx.send(RuntimeCommand::Terminate); + } + while !self.inner.cells.lock().await.is_empty() { + tokio::task::yield_now().await; + } + Ok(()) + } +} + +impl Default for CodeModeService { + fn default() -> Self { + Self::new() + } +} + +impl Drop for CodeModeService { + fn drop(&mut self) { + self.inner.shutting_down.store(true, Ordering::Release); + if let Ok(cells) = self.inner.cells.try_lock() { + for handle in cells.values() { + handle.cancellation_token.cancel(); + let (response_tx, _response_rx) = oneshot::channel(); + let _ = handle + .control_tx + .send(CellControlCommand::Terminate { response_tx }); + let _ = handle.runtime_tx.send(RuntimeCommand::Terminate); + } + } + } +} + +impl CodeModeSession for CodeModeService { + fn execute<'a>( + &'a self, + request: ExecuteRequest, + ) -> CodeModeSessionResultFuture<'a, StartedCell> { + Box::pin(CodeModeService::execute(self, request)) + } + + fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(CodeModeService::wait(self, request)) + } + + fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(CodeModeService::terminate(self, cell_id)) + } + + fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { + Box::pin(CodeModeService::shutdown(self)) + } +} + +enum CellControlCommand { + Poll { + yield_time_ms: u64, + response_tx: oneshot::Sender, + }, + PollToPending { + response_tx: oneshot::Sender, + }, + Terminate { + response_tx: oneshot::Sender, + }, +} + +enum CellResponseSender { + Runtime(oneshot::Sender), + ExecuteToPending(oneshot::Sender), +} + +struct PendingResult { + content_items: Vec, + error_text: Option, +} + +struct CellControlContext { + cell_id: CellId, + runtime_tx: std::sync::mpsc::Sender, + runtime_control_tx: std::sync::mpsc::Sender, + pending_mode: PendingRuntimeMode, + runtime_terminate_handle: v8::IsolateHandle, + cancellation_token: CancellationToken, +} + +fn missing_cell_response(cell_id: CellId) -> RuntimeResponse { + RuntimeResponse::Result { + error_text: Some(format!("exec cell {cell_id} not found")), + cell_id, + content_items: Vec::new(), + } +} + +fn pending_result_response(cell_id: &CellId, result: PendingResult) -> RuntimeResponse { + RuntimeResponse::Result { + cell_id: cell_id.clone(), + content_items: result.content_items, + error_text: result.error_text, + } +} + +fn send_terminal_response(response_tx: CellResponseSender, response: RuntimeResponse) { + match response_tx { + CellResponseSender::Runtime(response_tx) => { + let _ = response_tx.send(response); + } + CellResponseSender::ExecuteToPending(response_tx) => { + let _ = response_tx.send(ExecuteToPendingOutcome::Completed(response)); + } + } +} + +fn send_or_buffer_result( + cell_id: &CellId, + result: PendingResult, + response_tx: &mut Option, + pending_result: &mut Option, +) -> bool { + if let Some(response_tx) = response_tx.take() { + let response = pending_result_response(cell_id, result); + send_terminal_response(response_tx, response); + return true; + } + + *pending_result = Some(result); + false +} + +fn send_yield_response( + cell_id: &CellId, + content_items: &mut Vec, + response_tx: &mut Option, +) { + let Some(current_response_tx) = response_tx.take() else { + return; + }; + match current_response_tx { + CellResponseSender::Runtime(response_tx) => { + let _ = response_tx.send(RuntimeResponse::Yielded { + cell_id: cell_id.clone(), + content_items: std::mem::take(content_items), + }); + } + CellResponseSender::ExecuteToPending(execute_to_pending_tx) => { + *response_tx = Some(CellResponseSender::ExecuteToPending(execute_to_pending_tx)); + } + } +} + +async fn run_cell_control( + inner: Arc, + context: CellControlContext, + mut event_rx: mpsc::UnboundedReceiver, + mut control_rx: mpsc::UnboundedReceiver, + initial_response_tx: CellResponseSender, + initial_yield_time_ms: Option, +) { + let CellControlContext { + cell_id, + runtime_tx, + runtime_control_tx, + pending_mode, + runtime_terminate_handle, + cancellation_token, + } = context; + let mut content_items = Vec::new(); + let mut pending_tool_call_ids = Vec::new(); + let mut pending_result: Option = None; + let mut response_tx = Some(initial_response_tx); + let mut termination_requested = false; + let mut runtime_closed = false; + let mut yield_timer: Option>> = None; + let mut notification_tasks = JoinSet::new(); + + loop { + tokio::select! { + maybe_event = async { + if runtime_closed { + std::future::pending::>().await + } else { + event_rx.recv().await + } + } => { + let Some(event) = maybe_event else { + runtime_closed = true; + if termination_requested { + if let Some(response_tx) = response_tx.take() { + let response = RuntimeResponse::Terminated { + cell_id: cell_id.clone(), + content_items: std::mem::take(&mut content_items), + }; + send_terminal_response(response_tx, response); + } + break; + } + if pending_result.is_none() { + let result = PendingResult { + content_items: std::mem::take(&mut content_items), + error_text: Some("exec runtime ended unexpectedly".to_string()), + }; + if send_or_buffer_result( + &cell_id, + result, + &mut response_tx, + &mut pending_result, + ) { + break; + } + } + continue; + }; + match event { + RuntimeEvent::Started => { + yield_timer = initial_yield_time_ms.map(|initial_yield_time_ms| { + Box::pin(tokio::time::sleep(Duration::from_millis(initial_yield_time_ms))) + }); + } + RuntimeEvent::Pending => { + if let Some(current_response_tx) = response_tx.take() { + match current_response_tx { + CellResponseSender::Runtime(runtime_response_tx) => { + response_tx = + Some(CellResponseSender::Runtime(runtime_response_tx)); + } + CellResponseSender::ExecuteToPending(response_tx) => { + let _ = response_tx.send(ExecuteToPendingOutcome::Pending { + cell_id: cell_id.clone(), + content_items: std::mem::take(&mut content_items), + pending_tool_call_ids: std::mem::take( + &mut pending_tool_call_ids, + ), + }); + } + } + } + } + RuntimeEvent::ContentItem(item) => { + content_items.push(item); + } + RuntimeEvent::YieldRequested => { + yield_timer = None; + send_yield_response(&cell_id, &mut content_items, &mut response_tx); + } + RuntimeEvent::Notify { call_id, text } => { + let delegate = Arc::clone(&inner.delegate); + let cell_id = cell_id.clone(); + let cancellation_token = cancellation_token.child_token(); + notification_tasks.spawn(async move { + tokio::select! { + result = delegate.notify( + call_id, + cell_id.clone(), + text, + cancellation_token.clone(), + ) => { + if let Err(err) = result { + warn!( + "failed to deliver code mode notification for cell {cell_id}: {err}" + ); + } + } + _ = cancellation_token.cancelled() => {} + } + }); + } + RuntimeEvent::ToolCall { + id, + name, + kind, + input, + } => { + if pending_mode == PendingRuntimeMode::PauseUntilResumed { + pending_tool_call_ids.push(id.clone()); + } + let tool_call = CodeModeNestedToolCall { + cell_id: cell_id.clone(), + runtime_tool_call_id: id.clone(), + tool_name: name, + tool_kind: kind, + input, + }; + let delegate = Arc::clone(&inner.delegate); + let runtime_tx = runtime_tx.clone(); + let cancellation_token = cancellation_token.child_token(); + tokio::spawn(async move { + let response = tokio::select! { + response = delegate.invoke_tool(tool_call, cancellation_token.clone()) => response, + _ = cancellation_token.cancelled() => return, + }; + let command = match response { + Ok(result) => RuntimeCommand::ToolResponse { id, result }, + Err(error_text) => RuntimeCommand::ToolError { id, error_text }, + }; + let _ = runtime_tx.send(command); + }); + } + RuntimeEvent::Result { + stored_value_writes, + error_text, + } => { + yield_timer = None; + if termination_requested { + if let Some(response_tx) = response_tx.take() { + let response = RuntimeResponse::Terminated { + cell_id: cell_id.clone(), + content_items: std::mem::take(&mut content_items), + }; + send_terminal_response(response_tx, response); + } + break; + } + drain_notification_tasks(&mut notification_tasks).await; + inner + .stored_values + .lock() + .await + .extend(stored_value_writes); + let result = PendingResult { + content_items: std::mem::take(&mut content_items), + error_text, + }; + if send_or_buffer_result( + &cell_id, + result, + &mut response_tx, + &mut pending_result, + ) { + break; + } + } + } + } + task_result = notification_tasks.join_next(), if !notification_tasks.is_empty() => { + if let Some(Err(err)) = task_result + && !err.is_cancelled() + { + warn!("code mode notification task failed: {err}"); + } + } + maybe_command = control_rx.recv() => { + let Some(command) = maybe_command else { + break; + }; + match command { + CellControlCommand::Poll { + yield_time_ms, + response_tx: next_response_tx, + } => { + if let Some(result) = pending_result.take() { + let _ = next_response_tx.send(pending_result_response(&cell_id, result)); + break; + } + response_tx = Some(CellResponseSender::Runtime(next_response_tx)); + yield_timer = Some(Box::pin(tokio::time::sleep(Duration::from_millis(yield_time_ms)))); + resume_paused_runtime(&runtime_control_tx, pending_mode); + } + CellControlCommand::PollToPending { + response_tx: next_response_tx, + } => { + if let Some(result) = pending_result.take() { + let response = pending_result_response(&cell_id, result); + let _ = next_response_tx + .send(ExecuteToPendingOutcome::Completed(response)); + break; + } + response_tx = + Some(CellResponseSender::ExecuteToPending(next_response_tx)); + yield_timer = None; + resume_paused_runtime(&runtime_control_tx, pending_mode); + } + CellControlCommand::Terminate { response_tx: next_response_tx } => { + if let Some(result) = pending_result.take() { + let _ = next_response_tx.send(pending_result_response(&cell_id, result)); + break; + } + + response_tx = Some(CellResponseSender::Runtime(next_response_tx)); + termination_requested = true; + cancellation_token.cancel(); + yield_timer = None; + let _ = runtime_tx.send(RuntimeCommand::Terminate); + terminate_paused_runtime(&runtime_control_tx, pending_mode); + let _ = runtime_terminate_handle.terminate_execution(); + if runtime_closed { + if let Some(response_tx) = response_tx.take() { + let response = RuntimeResponse::Terminated { + cell_id: cell_id.clone(), + content_items: std::mem::take(&mut content_items), + }; + send_terminal_response(response_tx, response); + } + break; + } else { + continue; + } + } + } + } + _ = async { + if let Some(yield_timer) = yield_timer.as_mut() { + yield_timer.await; + } else { + std::future::pending::<()>().await; + } + } => { + yield_timer = None; + send_yield_response(&cell_id, &mut content_items, &mut response_tx); + } + } + } + + let _ = runtime_tx.send(RuntimeCommand::Terminate); + cancellation_token.cancel(); + drain_notification_tasks(&mut notification_tasks).await; + terminate_paused_runtime(&runtime_control_tx, pending_mode); + inner.cells.lock().await.remove(&cell_id); + inner.delegate.cell_closed(&cell_id); +} + +async fn drain_notification_tasks(notification_tasks: &mut JoinSet<()>) { + while let Some(result) = notification_tasks.join_next().await { + if let Err(err) = result + && !err.is_cancelled() + { + warn!("code mode notification task failed: {err}"); + } + } +} + +fn resume_paused_runtime( + runtime_control_tx: &std::sync::mpsc::Sender, + pending_mode: PendingRuntimeMode, +) { + if pending_mode == PendingRuntimeMode::PauseUntilResumed { + let _ = runtime_control_tx.send(RuntimeControlCommand::Resume); + } +} + +fn terminate_paused_runtime( + runtime_control_tx: &std::sync::mpsc::Sender, + pending_mode: PendingRuntimeMode, +) { + if pending_mode == PendingRuntimeMode::PauseUntilResumed { + let _ = runtime_control_tx.send(RuntimeControlCommand::Terminate); + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + use std::sync::atomic::AtomicU64; + use std::sync::atomic::Ordering; + use std::time::Duration; + + use codex_protocol::ToolName; + use pretty_assertions::assert_eq; + use tokio::sync::Mutex; + use tokio::sync::mpsc; + use tokio::sync::oneshot; + + use super::CellControlCommand; + use super::CellControlContext; + use super::CellId; + use super::CellResponseSender; + use super::CodeModeService; + use super::Inner; + use super::NoopCodeModeSessionDelegate; + use super::PendingRuntimeMode; + use super::RuntimeCommand; + use super::RuntimeResponse; + use super::WaitOutcome; + use super::WaitRequest; + use super::WaitToPendingOutcome; + use super::WaitToPendingRequest; + use super::run_cell_control; + use crate::CodeModeToolKind; + use crate::ExecuteToPendingOutcome; + use crate::FunctionCallOutputContentItem; + use crate::ToolDefinition; + use crate::runtime::ExecuteRequest; + use crate::runtime::RuntimeEvent; + use crate::runtime::spawn_runtime; + + fn execute_request(source: &str) -> ExecuteRequest { + ExecuteRequest { + tool_call_id: "call_1".to_string(), + enabled_tools: Vec::new(), + source: source.to_string(), + yield_time_ms: Some(1), + max_output_tokens: None, + } + } + + fn cell_id(value: &str) -> CellId { + CellId::new(value.to_string()) + } + + async fn execute(service: &CodeModeService, request: ExecuteRequest) -> RuntimeResponse { + service + .execute(request) + .await + .unwrap() + .initial_response() + .await + .unwrap() + } + + fn test_inner() -> Arc { + Arc::new(Inner { + stored_values: Mutex::new(HashMap::new()), + cells: Mutex::new(HashMap::new()), + delegate: Arc::new(NoopCodeModeSessionDelegate), + shutting_down: std::sync::atomic::AtomicBool::new(false), + next_cell_id: AtomicU64::new(1), + }) + } + + #[tokio::test] + async fn synchronous_exit_returns_successfully() { + let service = CodeModeService::new(); + + let response = execute( + &service, + ExecuteRequest { + source: r#"text("before"); exit(); text("after");"#.to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "before".to_string(), + }], + error_text: None, + } + ); + } + + #[tokio::test] + async fn stored_values_are_shared_between_cells_but_not_sessions() { + let first_session = CodeModeService::new(); + let second_session = CodeModeService::new(); + + let write_response = execute( + &first_session, + ExecuteRequest { + source: r#"store("key", "visible");"#.to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + let same_session = execute( + &first_session, + ExecuteRequest { + source: r#"text(String(load("key")));"#.to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + let other_session = execute( + &second_session, + ExecuteRequest { + source: r#"text(String(load("key")));"#.to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + write_response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: Vec::new(), + error_text: None, + } + ); + assert_eq!( + same_session, + RuntimeResponse::Result { + cell_id: cell_id("2"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "visible".to_string(), + }], + error_text: None, + } + ); + assert_eq!( + other_session, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "undefined".to_string(), + }], + error_text: None, + } + ); + } + + #[tokio::test] + async fn shutdown_interrupts_cpu_bound_cells() { + let service = CodeModeService::new(); + + let cell = service + .execute(ExecuteRequest { + source: "while (true) {}".to_string(), + ..execute_request("") + }) + .await + .unwrap(); + assert_eq!( + cell.initial_response().await.unwrap(), + RuntimeResponse::Yielded { + cell_id: cell_id("1"), + content_items: Vec::new(), + } + ); + + tokio::time::timeout(Duration::from_secs(1), service.shutdown()) + .await + .unwrap() + .unwrap(); + } + + #[tokio::test] + async fn start_cell_rejects_new_cell_after_shutdown_begins() { + let service = CodeModeService::new(); + service.inner.shutting_down.store(true, Ordering::Release); + let (response_tx, _response_rx) = oneshot::channel(); + + let error = service + .start_cell( + cell_id("late-cell"), + execute_request(""), + CellResponseSender::Runtime(response_tx), + Some(/*initial_yield_time_ms*/ 1), + PendingRuntimeMode::Continue, + ) + .await + .unwrap_err(); + + assert_eq!(error, "code mode session is shutting down".to_string()); + assert!(service.inner.cells.lock().await.is_empty()); + } + + #[tokio::test] + async fn execute_to_pending_returns_completed_for_synchronous_results() { + let service = CodeModeService::new(); + + let response = service + .execute_to_pending(ExecuteRequest { + source: r#"text("done");"#.to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }) + .await + .unwrap(); + + assert_eq!( + response, + ExecuteToPendingOutcome::Completed(RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "done".to_string(), + }], + error_text: None, + }) + ); + } + + #[tokio::test] + async fn execute_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("") + }), + ) + .await + .unwrap() + .unwrap(); + + assert_eq!( + response, + ExecuteToPendingOutcome::Pending { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "before".to_string(), + }], + pending_tool_call_ids: Vec::new(), + } + ); + + let termination = service.terminate(cell_id("1")).await.unwrap(); + + assert_eq!( + termination, + WaitOutcome::LiveCell(RuntimeResponse::Terminated { + cell_id: cell_id("1"), + content_items: Vec::new(), + }) + ); + } + + #[tokio::test] + async fn execute_to_pending_identifies_tool_calls_in_paused_frontier() { + let service = CodeModeService::new(); + + let response = service + .execute_to_pending(ExecuteRequest { + enabled_tools: vec![ToolDefinition { + name: "echo".to_string(), + tool_name: ToolName::plain("echo"), + description: String::new(), + kind: CodeModeToolKind::Function, + input_schema: None, + output_schema: None, + }], + source: r#" +await Promise.all([ + tools.echo({ value: "first" }), + tools.echo({ value: "second" }), +]); +"# + .to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }) + .await + .unwrap(); + + assert_eq!( + response, + ExecuteToPendingOutcome::Pending { + cell_id: cell_id("1"), + content_items: Vec::new(), + pending_tool_call_ids: vec!["tool-1".to_string(), "tool-2".to_string()], + } + ); + + let termination = service.terminate(cell_id("1")).await.unwrap(); + + assert_eq!( + termination, + WaitOutcome::LiveCell(RuntimeResponse::Terminated { + cell_id: cell_id("1"), + content_items: Vec::new(), + }) + ); + } + + #[tokio::test] + async fn execute_to_pending_excludes_delayed_timeout_tool_calls_until_wait() { + let service = CodeModeService::new(); + + let initial_response = service + .execute_to_pending(ExecuteRequest { + enabled_tools: vec![ToolDefinition { + name: "echo".to_string(), + tool_name: ToolName::plain("echo"), + description: String::new(), + kind: CodeModeToolKind::Function, + input_schema: None, + output_schema: None, + }], + source: r#" +setTimeout(() => { + tools.echo({ value: "delayed" }); +}, 1000); +await Promise.all([ + tools.echo({ value: "second" }), + tools.echo({ value: "third" }), +]); +"# + .to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }) + .await + .unwrap(); + + assert_eq!( + initial_response, + ExecuteToPendingOutcome::Pending { + cell_id: cell_id("1"), + content_items: Vec::new(), + pending_tool_call_ids: vec!["tool-1".to_string(), "tool-2".to_string()], + } + ); + + let runtime_tx = service + .inner + .cells + .lock() + .await + .get(&cell_id("1")) + .unwrap() + .runtime_tx + .clone(); + runtime_tx + .send(RuntimeCommand::TimeoutFired { id: 1 }) + .unwrap(); + + let resumed_response = tokio::time::timeout( + Duration::from_secs(1), + service.wait_to_pending(WaitToPendingRequest { + cell_id: cell_id("1"), + }), + ) + .await + .unwrap() + .unwrap(); + + assert_eq!( + resumed_response, + WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending { + cell_id: cell_id("1"), + content_items: Vec::new(), + pending_tool_call_ids: vec!["tool-3".to_string()], + }) + ); + + let termination = service.terminate(cell_id("1")).await.unwrap(); + + assert_eq!( + termination, + WaitOutcome::LiveCell(RuntimeResponse::Terminated { + cell_id: cell_id("1"), + content_items: Vec::new(), + }) + ); + } + + #[tokio::test] + async fn wait_to_pending_returns_after_resumed_runtime_becomes_quiescent_again() { + let service = CodeModeService::new(); + + let initial_response = service + .execute_to_pending(ExecuteRequest { + source: r#" +await new Promise((resolve) => setTimeout(resolve, 60_000)); +text("after"); +await new Promise(() => {}); +"# + .to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }) + .await + .unwrap(); + + assert_eq!( + initial_response, + ExecuteToPendingOutcome::Pending { + cell_id: cell_id("1"), + content_items: Vec::new(), + pending_tool_call_ids: Vec::new(), + } + ); + + let runtime_tx = service + .inner + .cells + .lock() + .await + .get(&cell_id("1")) + .unwrap() + .runtime_tx + .clone(); + runtime_tx + .send(RuntimeCommand::TimeoutFired { id: 1 }) + .unwrap(); + + let resumed_response = tokio::time::timeout( + Duration::from_secs(1), + service.wait_to_pending(WaitToPendingRequest { + cell_id: cell_id("1"), + }), + ) + .await + .unwrap() + .unwrap(); + + assert_eq!( + resumed_response, + WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "after".to_string(), + }], + pending_tool_call_ids: Vec::new(), + }) + ); + + let termination = service.terminate(cell_id("1")).await.unwrap(); + + assert_eq!( + termination, + WaitOutcome::LiveCell(RuntimeResponse::Terminated { + cell_id: cell_id("1"), + content_items: Vec::new(), + }) + ); + } + + #[tokio::test] + async fn wait_to_pending_returns_completed_after_resumed_runtime_finishes() { + let service = CodeModeService::new(); + + let initial_response = service + .execute_to_pending(ExecuteRequest { + source: r#" +await new Promise((resolve) => setTimeout(resolve, 60_000)); +text("done"); +"# + .to_string(), + yield_time_ms: Some(60_000), + ..execute_request("") + }) + .await + .unwrap(); + + assert_eq!( + initial_response, + ExecuteToPendingOutcome::Pending { + cell_id: cell_id("1"), + content_items: Vec::new(), + pending_tool_call_ids: Vec::new(), + } + ); + + let runtime_tx = service + .inner + .cells + .lock() + .await + .get(&cell_id("1")) + .unwrap() + .runtime_tx + .clone(); + runtime_tx + .send(RuntimeCommand::TimeoutFired { id: 1 }) + .unwrap(); + + let resumed_response = tokio::time::timeout( + Duration::from_secs(1), + service.wait_to_pending(WaitToPendingRequest { + cell_id: cell_id("1"), + }), + ) + .await + .unwrap() + .unwrap(); + + 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, + } + )) + ); + } + + #[tokio::test] + async fn v8_console_is_not_exposed_on_global_this() { + let service = CodeModeService::new(); + + let response = execute( + &service, + ExecuteRequest { + source: r#"text(String(Object.hasOwn(globalThis, "console")));"#.to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "false".to_string(), + }], + error_text: None, + } + ); + } + + #[tokio::test] + async fn date_locale_string_formats_with_icu_data() { + let service = CodeModeService::new(); + + let response = execute( + &service, + ExecuteRequest { + source: r#" +const value = new Date("2025-01-02T03:04:05Z") + .toLocaleString("fr-FR", { + weekday: "long", + month: "long", + day: "numeric", + hour: "2-digit", + minute: "2-digit", + second: "2-digit", + hour12: false, + timeZone: "UTC", + }); +text(value); +"# + .to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "jeudi 2 janvier \u{e0} 03:04:05".to_string(), + }], + error_text: None, + } + ); + } + + #[tokio::test] + async fn intl_date_time_format_formats_with_icu_data() { + let service = CodeModeService::new(); + + let response = execute( + &service, + ExecuteRequest { + source: r#" +const formatter = new Intl.DateTimeFormat("fr-FR", { + weekday: "long", + month: "long", + day: "numeric", + hour: "2-digit", + minute: "2-digit", + second: "2-digit", + hour12: false, + timeZone: "UTC", +}); +text(formatter.format(new Date("2025-01-02T03:04:05Z"))); +"# + .to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "jeudi 2 janvier \u{e0} 03:04:05".to_string(), + }], + error_text: None, + } + ); + } + + #[tokio::test] + async fn output_helpers_return_undefined() { + let service = CodeModeService::new(); + + let response = execute( + &service, + ExecuteRequest { + source: r#" +const returnsUndefined = [ + text("first"), + image("https://example.com/image.jpg"), + notify("ping"), +].map((value) => value === undefined); +text(JSON.stringify(returnsUndefined)); +"# + .to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![ + FunctionCallOutputContentItem::InputText { + text: "first".to_string(), + }, + FunctionCallOutputContentItem::InputImage { + image_url: "https://example.com/image.jpg".to_string(), + detail: Some(crate::DEFAULT_IMAGE_DETAIL), + }, + FunctionCallOutputContentItem::InputText { + text: "[true,true,true]".to_string(), + }, + ], + error_text: None, + } + ); + } + + #[tokio::test] + async fn image_helper_accepts_raw_mcp_image_block_with_original_detail() { + let service = CodeModeService::new(); + + let response = execute( + &service, + ExecuteRequest { + source: r#" +image({ + type: "image", + data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==", + mimeType: "image/png", + _meta: { "codex/imageDetail": "original" }, +}); +"# + .to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==".to_string(), + detail: Some(crate::ImageDetail::Original), + }], + error_text: None, + } + ); + } + + #[tokio::test] + async fn image_helper_second_arg_overrides_explicit_object_detail() { + let service = CodeModeService::new(); + + let response = execute( + &service, + ExecuteRequest { + source: r#" +image( + { + image_url: "https://example.com/image.jpg", + detail: "high", + }, + "original", +); +"# + .to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputImage { + image_url: "https://example.com/image.jpg".to_string(), + detail: Some(crate::ImageDetail::Original), + }], + error_text: None, + } + ); + } + + #[tokio::test] + async fn image_helper_second_arg_overrides_raw_mcp_image_detail() { + let service = CodeModeService::new(); + + let response = execute( + &service, + ExecuteRequest { + source: r#" +image( + { + type: "image", + data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==", + mimeType: "image/png", + _meta: { "codex/imageDetail": "original" }, + }, + "high", +); +"# + .to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==".to_string(), + detail: Some(crate::ImageDetail::High), + }], + error_text: None, + } + ); + } + + #[tokio::test] + async fn image_helper_accepts_low_detail() { + let service = CodeModeService::new(); + + let response = execute( + &service, + ExecuteRequest { + source: r#" +image({ + image_url: "https://example.com/image.jpg", + detail: "low", +}); +"# + .to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputImage { + image_url: "https://example.com/image.jpg".to_string(), + detail: Some(crate::ImageDetail::Low), + }], + error_text: None, + } + ); + } + + #[tokio::test] + async fn image_helper_rejects_unsupported_detail() { + let service = CodeModeService::new(); + + let response = execute( + &service, + ExecuteRequest { + source: r#" +image({ + image_url: "https://example.com/image.jpg", + detail: "medium", +}); +"# + .to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: Vec::new(), + error_text: Some( + "image detail must be one of: auto, low, high, original".to_string() + ), + } + ); + } + + #[tokio::test] + async fn image_helper_rejects_raw_mcp_result_container() { + let service = CodeModeService::new(); + + let response = execute( + &service, + ExecuteRequest { + source: r#" +image({ + content: [ + { + type: "image", + data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==", + mimeType: "image/png", + _meta: { "codex/imageDetail": "original" }, + }, + ], + isError: false, +}); +"# + .to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: Vec::new(), + error_text: Some( + "image expects a non-empty image URL string, an object with image_url and optional detail, or a raw MCP image block".to_string(), + ), + } + ); + } + + #[tokio::test] + async fn wait_reports_missing_cell_separately_from_runtime_results() { + let service = CodeModeService::new(); + + let response = service + .wait(WaitRequest { + cell_id: cell_id("missing"), + yield_time_ms: 1, + }) + .await + .unwrap(); + + assert_eq!( + response, + WaitOutcome::MissingCell(RuntimeResponse::Result { + cell_id: cell_id("missing"), + content_items: Vec::new(), + error_text: Some("exec cell missing not found".to_string()), + }) + ); + } + + #[tokio::test] + async fn terminate_waits_for_runtime_shutdown_before_responding() { + let inner = test_inner(); + let (event_tx, event_rx) = mpsc::unbounded_channel(); + let (control_tx, control_rx) = mpsc::unbounded_channel(); + let (initial_response_tx, initial_response_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 { + source: "await new Promise(() => {})".to_string(), + yield_time_ms: None, + ..execute_request("") + }, + runtime_event_tx, + PendingRuntimeMode::Continue, + ) + .unwrap(); + + tokio::spawn(run_cell_control( + inner, + CellControlContext { + cell_id: cell_id("cell-1"), + runtime_tx: runtime_tx.clone(), + runtime_control_tx, + pending_mode: PendingRuntimeMode::Continue, + runtime_terminate_handle, + cancellation_token: tokio_util::sync::CancellationToken::new(), + }, + event_rx, + control_rx, + CellResponseSender::Runtime(initial_response_tx), + Some(/*initial_yield_time_ms*/ 60_000), + )); + + event_tx.send(RuntimeEvent::Started).unwrap(); + event_tx.send(RuntimeEvent::YieldRequested).unwrap(); + assert_eq!( + initial_response_rx.await.unwrap(), + RuntimeResponse::Yielded { + cell_id: cell_id("cell-1"), + content_items: Vec::new(), + } + ); + + let (terminate_response_tx, terminate_response_rx) = oneshot::channel(); + control_tx + .send(CellControlCommand::Terminate { + response_tx: terminate_response_tx, + }) + .unwrap(); + let terminate_response = async { terminate_response_rx.await.unwrap() }; + tokio::pin!(terminate_response); + assert!( + tokio::time::timeout(Duration::from_millis(100), terminate_response.as_mut()) + .await + .is_err() + ); + + drop(event_tx); + + assert_eq!( + terminate_response.await, + RuntimeResponse::Terminated { + cell_id: cell_id("cell-1"), + content_items: Vec::new(), + } + ); + + let _ = runtime_tx.send(RuntimeCommand::Terminate); + } +} diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index 879404d8aa..18bc40cfce 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -9,21 +9,15 @@ doctest = false name = "codex_code_mode" path = "src/lib.rs" -[features] -sandbox = ["v8/v8_enable_sandbox"] - [lints] workspace = true [dependencies] codex-protocol = { workspace = true } -deno_core_icudata = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } -tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } +tokio = { workspace = true, features = ["sync"] } tokio-util = { workspace = true, features = ["rt"] } -tracing = { workspace = true } -v8 = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 37ba1d2a91..a501df4644 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -5,11 +5,15 @@ mod service; pub use description::CODE_MODE_PRAGMA_PREFIX; pub use description::CodeModeToolKind; +#[doc(hidden)] +pub use description::EnabledToolMetadata; pub use description::ToolDefinition; pub use description::ToolNamespaceDescription; pub use description::augment_tool_definition; pub use description::build_exec_tool_description; pub use description::build_wait_tool_description; +#[doc(hidden)] +pub use description::enabled_tool_metadata; pub use description::is_code_mode_nested_tool; pub use description::normalize_code_mode_identifier; pub use description::parse_exec_source; @@ -36,11 +40,11 @@ pub use service::CodeModeSessionDelegate; pub use service::CodeModeSessionProvider; pub use service::CodeModeSessionProviderFuture; pub use service::CodeModeSessionResultFuture; -pub use service::InProcessCodeModeSessionProvider; pub use service::NoopCodeModeSessionDelegate; pub use service::NotificationFuture; pub use service::StartedCell; pub use service::ToolInvocationFuture; +pub use service::UnavailableCodeModeSessionProvider; pub const PUBLIC_TOOL_NAME: &str = "exec"; pub const WAIT_TOOL_NAME: &str = "wait"; diff --git a/codex-rs/code-mode/src/runtime.rs b/codex-rs/code-mode/src/runtime.rs new file mode 100644 index 0000000000..5cdbcac32b --- /dev/null +++ b/codex-rs/code-mode/src/runtime.rs @@ -0,0 +1,114 @@ +use codex_protocol::ToolName; +use serde::Serialize; +use serde_json::Value as JsonValue; + +use crate::CellId; +use crate::CodeModeToolKind; +use crate::FunctionCallOutputContentItem; +use crate::ToolDefinition; + +pub const DEFAULT_EXEC_YIELD_TIME_MS: u64 = 10_000; +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)] +pub struct ExecuteRequest { + 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)] +pub struct WaitRequest { + pub cell_id: CellId, + pub yield_time_ms: u64, +} + +#[derive(Clone, Debug)] +pub struct WaitToPendingRequest { + pub cell_id: CellId, +} + +/// Result of waiting on a code-mode cell. +/// +/// The wrapped `RuntimeResponse` is the model-facing wait result. The enum +/// variant carries the extra lifecycle provenance that `RuntimeResponse` cannot: +/// a failed real cell and a missing-cell wait both use +/// `RuntimeResponse::Result { error_text: Some(..), .. }`, but only the former +/// should be treated as a code-cell lifecycle event. +#[derive(Debug, PartialEq)] +pub enum WaitOutcome { + /// The requested code cell was live when the wait command was accepted. + LiveCell(RuntimeResponse), + /// The requested code cell was not live. + MissingCell(RuntimeResponse), +} + +/// Result of executing a code-mode cell until it either completes or reaches a +/// quiescent pending state. +#[derive(Debug, PartialEq)] +pub enum ExecuteToPendingOutcome { + /// The cell is waiting for more runtime input after draining the runtime + /// input queue that was ready at the pending boundary. + Pending { + cell_id: CellId, + content_items: Vec, + /// Runtime tool-call ids emitted before this paused execution frontier + /// sealed. Hosts can use these ids to drain their tool-call transport + /// before surfacing the pending boundary to callers. + pending_tool_call_ids: Vec, + }, + /// The cell reached a terminal runtime response before going pending. + Completed(RuntimeResponse), +} + +/// Result of resuming a live code-mode cell until it completes or becomes +/// quiescent again. +#[derive(Debug, PartialEq)] +pub enum WaitToPendingOutcome { + /// The requested code cell was live when the wait command was accepted. + LiveCell(ExecuteToPendingOutcome), + /// The requested code cell was not live. + MissingCell(RuntimeResponse), +} + +impl From for RuntimeResponse { + fn from(outcome: WaitOutcome) -> Self { + match outcome { + WaitOutcome::LiveCell(response) | WaitOutcome::MissingCell(response) => response, + } + } +} + +#[derive(Debug, PartialEq, Serialize)] +pub enum RuntimeResponse { + Yielded { + cell_id: CellId, + content_items: Vec, + }, + Terminated { + cell_id: CellId, + content_items: Vec, + }, + Result { + cell_id: CellId, + content_items: Vec, + error_text: Option, + }, +} + +/// Nested tool request emitted by one code-mode cell. +/// +/// Code mode owns the per-cell runtime id. Hosts should preserve it for +/// provenance/debugging, but should still assign their own runtime tool call id +/// if their tool-call graph requires globally unique ids. +#[derive(Debug)] +pub struct CodeModeNestedToolCall { + pub cell_id: CellId, + pub runtime_tool_call_id: String, + pub tool_name: ToolName, + pub tool_kind: CodeModeToolKind, + pub input: Option, +} diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index 6589e6ec5a..966a1df479 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -1,38 +1,19 @@ -use std::collections::HashMap; use std::fmt; use std::future::Future; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::AtomicU64; -use std::sync::atomic::Ordering; -use std::time::Duration; use serde::Deserialize; use serde::Serialize; use serde_json::Value as JsonValue; -use tokio::sync::Mutex; -use tokio::sync::mpsc; use tokio::sync::oneshot; -use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; -use tracing::warn; -use crate::FunctionCallOutputContentItem; -use crate::runtime::CodeModeNestedToolCall; -use crate::runtime::DEFAULT_EXEC_YIELD_TIME_MS; -use crate::runtime::ExecuteRequest; -use crate::runtime::ExecuteToPendingOutcome; -use crate::runtime::PendingRuntimeMode; -use crate::runtime::RuntimeCommand; -use crate::runtime::RuntimeControlCommand; -use crate::runtime::RuntimeEvent; -use crate::runtime::RuntimeResponse; -use crate::runtime::WaitOutcome; -use crate::runtime::WaitRequest; -use crate::runtime::WaitToPendingOutcome; -use crate::runtime::WaitToPendingRequest; -use crate::runtime::spawn_runtime; +use crate::CodeModeNestedToolCall; +use crate::ExecuteRequest; +use crate::RuntimeResponse; +use crate::WaitOutcome; +use crate::WaitRequest; pub type CodeModeSessionResultFuture<'a, T> = Pin> + Send + 'a>>; @@ -73,6 +54,14 @@ pub struct StartedCell { } impl StartedCell { + #[doc(hidden)] + pub fn new(cell_id: CellId, initial_response_rx: oneshot::Receiver) -> Self { + Self { + cell_id, + initial_response_rx, + } + } + pub async fn initial_response(self) -> Result { self.initial_response_rx .await @@ -156,1682 +145,18 @@ pub trait CodeModeSessionProvider: Send + Sync { ) -> CodeModeSessionProviderFuture<'a>; } -#[derive(Default)] -pub struct InProcessCodeModeSessionProvider; +/// Provider used by Codex builds that expose code-mode planning but do not own a runtime. +/// +/// Runtime-owning binaries should inject a concrete provider at their session +/// construction boundary. Calling code mode through this provider returns an +/// explicit error instead of silently changing enabled code-mode behavior. +pub struct UnavailableCodeModeSessionProvider; -impl CodeModeSessionProvider for InProcessCodeModeSessionProvider { +impl CodeModeSessionProvider for UnavailableCodeModeSessionProvider { fn create_session<'a>( &'a self, - delegate: Arc, + _delegate: Arc, ) -> CodeModeSessionProviderFuture<'a> { - Box::pin(async move { - let session: Arc = - Arc::new(CodeModeService::with_delegate(delegate)); - Ok(session) - }) - } -} - -#[derive(Clone)] -struct CellHandle { - control_tx: mpsc::UnboundedSender, - runtime_tx: std::sync::mpsc::Sender, - cancellation_token: CancellationToken, -} - -struct Inner { - stored_values: Mutex>, - cells: Mutex>, - delegate: Arc, - shutting_down: AtomicBool, - next_cell_id: AtomicU64, -} - -pub struct CodeModeService { - inner: Arc, -} - -impl CodeModeService { - pub fn new() -> Self { - Self::with_delegate(Arc::new(NoopCodeModeSessionDelegate)) - } - - pub fn with_delegate(delegate: Arc) -> Self { - Self { - inner: Arc::new(Inner { - stored_values: Mutex::new(HashMap::new()), - cells: Mutex::new(HashMap::new()), - delegate, - shutting_down: AtomicBool::new(false), - next_cell_id: AtomicU64::new(1), - }), - } - } - - fn allocate_cell_id(&self) -> CellId { - CellId::new( - self.inner - .next_cell_id - .fetch_add(1, Ordering::Relaxed) - .to_string(), - ) - } - - pub async fn execute(&self, request: ExecuteRequest) -> Result { - if self.inner.shutting_down.load(Ordering::Acquire) { - return Err("code mode session is shutting down".to_string()); - } - let initial_yield_time_ms = request.yield_time_ms.unwrap_or(DEFAULT_EXEC_YIELD_TIME_MS); - let (response_tx, response_rx) = oneshot::channel(); - let cell_id = self.allocate_cell_id(); - self.start_cell( - cell_id.clone(), - request, - CellResponseSender::Runtime(response_tx), - Some(initial_yield_time_ms), - PendingRuntimeMode::Continue, - ) - .await?; - - Ok(StartedCell { - cell_id, - initial_response_rx: response_rx, - }) - } - - pub async fn execute_to_pending( - &self, - request: ExecuteRequest, - ) -> Result { - let (response_tx, response_rx) = oneshot::channel(); - let cell_id = self.allocate_cell_id(); - self.start_cell( - cell_id, - request, - CellResponseSender::ExecuteToPending(response_tx), - /*initial_yield_time_ms*/ None, - PendingRuntimeMode::PauseUntilResumed, - ) - .await?; - - response_rx - .await - .map_err(|_| "exec runtime ended unexpectedly".to_string()) - } - - async fn start_cell( - &self, - cell_id: CellId, - request: ExecuteRequest, - initial_response_tx: CellResponseSender, - initial_yield_time_ms: Option, - pending_mode: PendingRuntimeMode, - ) -> Result<(), String> { - let (event_tx, event_rx) = mpsc::unbounded_channel(); - let (control_tx, control_rx) = mpsc::unbounded_channel(); - let stored_values = self.inner.stored_values.lock().await.clone(); - let cancellation_token = CancellationToken::new(); - let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = { - let mut cells = self.inner.cells.lock().await; - if self.inner.shutting_down.load(Ordering::Acquire) { - return Err("code mode session is shutting down".to_string()); - } - if cells.contains_key(&cell_id) { - return Err(format!("exec cell {cell_id} already exists")); - } - - let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = - spawn_runtime(stored_values, request, event_tx, pending_mode)?; - - cells.insert( - cell_id.clone(), - CellHandle { - control_tx, - runtime_tx: runtime_tx.clone(), - cancellation_token: cancellation_token.clone(), - }, - ); - (runtime_tx, runtime_control_tx, runtime_terminate_handle) - }; - - tokio::spawn(run_cell_control( - Arc::clone(&self.inner), - CellControlContext { - cell_id, - runtime_tx, - runtime_control_tx, - pending_mode, - runtime_terminate_handle, - cancellation_token, - }, - event_rx, - control_rx, - initial_response_tx, - initial_yield_time_ms, - )); - - Ok(()) - } - - pub async fn wait(&self, request: WaitRequest) -> Result { - let WaitRequest { - cell_id, - yield_time_ms, - } = request; - let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); - let Some(handle) = handle else { - return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); - }; - let (response_tx, response_rx) = oneshot::channel(); - let control_message = CellControlCommand::Poll { - yield_time_ms, - response_tx, - }; - if handle.control_tx.send(control_message).is_err() { - return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); - } - match response_rx.await { - Ok(response) => Ok(WaitOutcome::LiveCell(response)), - Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))), - } - } - - pub async fn terminate(&self, cell_id: CellId) -> Result { - let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); - let Some(handle) = handle else { - return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); - }; - let (response_tx, response_rx) = oneshot::channel(); - if handle - .control_tx - .send(CellControlCommand::Terminate { response_tx }) - .is_err() - { - return Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))); - } - match response_rx.await { - Ok(response) => Ok(WaitOutcome::LiveCell(response)), - Err(_) => Ok(WaitOutcome::MissingCell(missing_cell_response(cell_id))), - } - } - - pub async fn wait_to_pending( - &self, - request: WaitToPendingRequest, - ) -> Result { - let cell_id = request.cell_id; - let handle = self.inner.cells.lock().await.get(&cell_id).cloned(); - let Some(handle) = handle else { - return Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( - cell_id, - ))); - }; - let (response_tx, response_rx) = oneshot::channel(); - if handle - .control_tx - .send(CellControlCommand::PollToPending { response_tx }) - .is_err() - { - return Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( - cell_id, - ))); - } - match response_rx.await { - Ok(response) => Ok(WaitToPendingOutcome::LiveCell(response)), - Err(_) => Ok(WaitToPendingOutcome::MissingCell(missing_cell_response( - cell_id, - ))), - } - } - - pub async fn shutdown(&self) -> Result<(), String> { - self.inner.shutting_down.store(true, Ordering::Release); - let handles = self - .inner - .cells - .lock() - .await - .values() - .cloned() - .collect::>(); - for handle in handles { - handle.cancellation_token.cancel(); - let (response_tx, _response_rx) = oneshot::channel(); - let _ = handle - .control_tx - .send(CellControlCommand::Terminate { response_tx }); - let _ = handle.runtime_tx.send(RuntimeCommand::Terminate); - } - while !self.inner.cells.lock().await.is_empty() { - tokio::task::yield_now().await; - } - Ok(()) - } -} - -impl Default for CodeModeService { - fn default() -> Self { - Self::new() - } -} - -impl Drop for CodeModeService { - fn drop(&mut self) { - self.inner.shutting_down.store(true, Ordering::Release); - if let Ok(cells) = self.inner.cells.try_lock() { - for handle in cells.values() { - handle.cancellation_token.cancel(); - let (response_tx, _response_rx) = oneshot::channel(); - let _ = handle - .control_tx - .send(CellControlCommand::Terminate { response_tx }); - let _ = handle.runtime_tx.send(RuntimeCommand::Terminate); - } - } - } -} - -impl CodeModeSession for CodeModeService { - fn execute<'a>( - &'a self, - request: ExecuteRequest, - ) -> CodeModeSessionResultFuture<'a, StartedCell> { - Box::pin(CodeModeService::execute(self, request)) - } - - fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> { - Box::pin(CodeModeService::wait(self, request)) - } - - fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { - Box::pin(CodeModeService::terminate(self, cell_id)) - } - - fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { - Box::pin(CodeModeService::shutdown(self)) - } -} - -enum CellControlCommand { - Poll { - yield_time_ms: u64, - response_tx: oneshot::Sender, - }, - PollToPending { - response_tx: oneshot::Sender, - }, - Terminate { - response_tx: oneshot::Sender, - }, -} - -enum CellResponseSender { - Runtime(oneshot::Sender), - ExecuteToPending(oneshot::Sender), -} - -struct PendingResult { - content_items: Vec, - error_text: Option, -} - -struct CellControlContext { - cell_id: CellId, - runtime_tx: std::sync::mpsc::Sender, - runtime_control_tx: std::sync::mpsc::Sender, - pending_mode: PendingRuntimeMode, - runtime_terminate_handle: v8::IsolateHandle, - cancellation_token: CancellationToken, -} - -fn missing_cell_response(cell_id: CellId) -> RuntimeResponse { - RuntimeResponse::Result { - error_text: Some(format!("exec cell {cell_id} not found")), - cell_id, - content_items: Vec::new(), - } -} - -fn pending_result_response(cell_id: &CellId, result: PendingResult) -> RuntimeResponse { - RuntimeResponse::Result { - cell_id: cell_id.clone(), - content_items: result.content_items, - error_text: result.error_text, - } -} - -fn send_terminal_response(response_tx: CellResponseSender, response: RuntimeResponse) { - match response_tx { - CellResponseSender::Runtime(response_tx) => { - let _ = response_tx.send(response); - } - CellResponseSender::ExecuteToPending(response_tx) => { - let _ = response_tx.send(ExecuteToPendingOutcome::Completed(response)); - } - } -} - -fn send_or_buffer_result( - cell_id: &CellId, - result: PendingResult, - response_tx: &mut Option, - pending_result: &mut Option, -) -> bool { - if let Some(response_tx) = response_tx.take() { - let response = pending_result_response(cell_id, result); - send_terminal_response(response_tx, response); - return true; - } - - *pending_result = Some(result); - false -} - -fn send_yield_response( - cell_id: &CellId, - content_items: &mut Vec, - response_tx: &mut Option, -) { - let Some(current_response_tx) = response_tx.take() else { - return; - }; - match current_response_tx { - CellResponseSender::Runtime(response_tx) => { - let _ = response_tx.send(RuntimeResponse::Yielded { - cell_id: cell_id.clone(), - content_items: std::mem::take(content_items), - }); - } - CellResponseSender::ExecuteToPending(execute_to_pending_tx) => { - *response_tx = Some(CellResponseSender::ExecuteToPending(execute_to_pending_tx)); - } - } -} - -async fn run_cell_control( - inner: Arc, - context: CellControlContext, - mut event_rx: mpsc::UnboundedReceiver, - mut control_rx: mpsc::UnboundedReceiver, - initial_response_tx: CellResponseSender, - initial_yield_time_ms: Option, -) { - let CellControlContext { - cell_id, - runtime_tx, - runtime_control_tx, - pending_mode, - runtime_terminate_handle, - cancellation_token, - } = context; - let mut content_items = Vec::new(); - let mut pending_tool_call_ids = Vec::new(); - let mut pending_result: Option = None; - let mut response_tx = Some(initial_response_tx); - let mut termination_requested = false; - let mut runtime_closed = false; - let mut yield_timer: Option>> = None; - let mut notification_tasks = JoinSet::new(); - - loop { - tokio::select! { - maybe_event = async { - if runtime_closed { - std::future::pending::>().await - } else { - event_rx.recv().await - } - } => { - let Some(event) = maybe_event else { - runtime_closed = true; - if termination_requested { - if let Some(response_tx) = response_tx.take() { - let response = RuntimeResponse::Terminated { - cell_id: cell_id.clone(), - content_items: std::mem::take(&mut content_items), - }; - send_terminal_response(response_tx, response); - } - break; - } - if pending_result.is_none() { - let result = PendingResult { - content_items: std::mem::take(&mut content_items), - error_text: Some("exec runtime ended unexpectedly".to_string()), - }; - if send_or_buffer_result( - &cell_id, - result, - &mut response_tx, - &mut pending_result, - ) { - break; - } - } - continue; - }; - match event { - RuntimeEvent::Started => { - yield_timer = initial_yield_time_ms.map(|initial_yield_time_ms| { - Box::pin(tokio::time::sleep(Duration::from_millis(initial_yield_time_ms))) - }); - } - RuntimeEvent::Pending => { - if let Some(current_response_tx) = response_tx.take() { - match current_response_tx { - CellResponseSender::Runtime(runtime_response_tx) => { - response_tx = - Some(CellResponseSender::Runtime(runtime_response_tx)); - } - CellResponseSender::ExecuteToPending(response_tx) => { - let _ = response_tx.send(ExecuteToPendingOutcome::Pending { - cell_id: cell_id.clone(), - content_items: std::mem::take(&mut content_items), - pending_tool_call_ids: std::mem::take( - &mut pending_tool_call_ids, - ), - }); - } - } - } - } - RuntimeEvent::ContentItem(item) => { - content_items.push(item); - } - RuntimeEvent::YieldRequested => { - yield_timer = None; - send_yield_response(&cell_id, &mut content_items, &mut response_tx); - } - RuntimeEvent::Notify { call_id, text } => { - let delegate = Arc::clone(&inner.delegate); - let cell_id = cell_id.clone(); - let cancellation_token = cancellation_token.child_token(); - notification_tasks.spawn(async move { - tokio::select! { - result = delegate.notify( - call_id, - cell_id.clone(), - text, - cancellation_token.clone(), - ) => { - if let Err(err) = result { - warn!( - "failed to deliver code mode notification for cell {cell_id}: {err}" - ); - } - } - _ = cancellation_token.cancelled() => {} - } - }); - } - RuntimeEvent::ToolCall { - id, - name, - kind, - input, - } => { - if pending_mode == PendingRuntimeMode::PauseUntilResumed { - pending_tool_call_ids.push(id.clone()); - } - let tool_call = CodeModeNestedToolCall { - cell_id: cell_id.clone(), - runtime_tool_call_id: id.clone(), - tool_name: name, - tool_kind: kind, - input, - }; - let delegate = Arc::clone(&inner.delegate); - let runtime_tx = runtime_tx.clone(); - let cancellation_token = cancellation_token.child_token(); - tokio::spawn(async move { - let response = tokio::select! { - response = delegate.invoke_tool(tool_call, cancellation_token.clone()) => response, - _ = cancellation_token.cancelled() => return, - }; - let command = match response { - Ok(result) => RuntimeCommand::ToolResponse { id, result }, - Err(error_text) => RuntimeCommand::ToolError { id, error_text }, - }; - let _ = runtime_tx.send(command); - }); - } - RuntimeEvent::Result { - stored_value_writes, - error_text, - } => { - yield_timer = None; - if termination_requested { - if let Some(response_tx) = response_tx.take() { - let response = RuntimeResponse::Terminated { - cell_id: cell_id.clone(), - content_items: std::mem::take(&mut content_items), - }; - send_terminal_response(response_tx, response); - } - break; - } - drain_notification_tasks(&mut notification_tasks).await; - inner - .stored_values - .lock() - .await - .extend(stored_value_writes); - let result = PendingResult { - content_items: std::mem::take(&mut content_items), - error_text, - }; - if send_or_buffer_result( - &cell_id, - result, - &mut response_tx, - &mut pending_result, - ) { - break; - } - } - } - } - task_result = notification_tasks.join_next(), if !notification_tasks.is_empty() => { - if let Some(Err(err)) = task_result - && !err.is_cancelled() - { - warn!("code mode notification task failed: {err}"); - } - } - maybe_command = control_rx.recv() => { - let Some(command) = maybe_command else { - break; - }; - match command { - CellControlCommand::Poll { - yield_time_ms, - response_tx: next_response_tx, - } => { - if let Some(result) = pending_result.take() { - let _ = next_response_tx.send(pending_result_response(&cell_id, result)); - break; - } - response_tx = Some(CellResponseSender::Runtime(next_response_tx)); - yield_timer = Some(Box::pin(tokio::time::sleep(Duration::from_millis(yield_time_ms)))); - resume_paused_runtime(&runtime_control_tx, pending_mode); - } - CellControlCommand::PollToPending { - response_tx: next_response_tx, - } => { - if let Some(result) = pending_result.take() { - let response = pending_result_response(&cell_id, result); - let _ = next_response_tx - .send(ExecuteToPendingOutcome::Completed(response)); - break; - } - response_tx = - Some(CellResponseSender::ExecuteToPending(next_response_tx)); - yield_timer = None; - resume_paused_runtime(&runtime_control_tx, pending_mode); - } - CellControlCommand::Terminate { response_tx: next_response_tx } => { - if let Some(result) = pending_result.take() { - let _ = next_response_tx.send(pending_result_response(&cell_id, result)); - break; - } - - response_tx = Some(CellResponseSender::Runtime(next_response_tx)); - termination_requested = true; - cancellation_token.cancel(); - yield_timer = None; - let _ = runtime_tx.send(RuntimeCommand::Terminate); - terminate_paused_runtime(&runtime_control_tx, pending_mode); - let _ = runtime_terminate_handle.terminate_execution(); - if runtime_closed { - if let Some(response_tx) = response_tx.take() { - let response = RuntimeResponse::Terminated { - cell_id: cell_id.clone(), - content_items: std::mem::take(&mut content_items), - }; - send_terminal_response(response_tx, response); - } - break; - } else { - continue; - } - } - } - } - _ = async { - if let Some(yield_timer) = yield_timer.as_mut() { - yield_timer.await; - } else { - std::future::pending::<()>().await; - } - } => { - yield_timer = None; - send_yield_response(&cell_id, &mut content_items, &mut response_tx); - } - } - } - - let _ = runtime_tx.send(RuntimeCommand::Terminate); - cancellation_token.cancel(); - drain_notification_tasks(&mut notification_tasks).await; - terminate_paused_runtime(&runtime_control_tx, pending_mode); - inner.cells.lock().await.remove(&cell_id); - inner.delegate.cell_closed(&cell_id); -} - -async fn drain_notification_tasks(notification_tasks: &mut JoinSet<()>) { - while let Some(result) = notification_tasks.join_next().await { - if let Err(err) = result - && !err.is_cancelled() - { - warn!("code mode notification task failed: {err}"); - } - } -} - -fn resume_paused_runtime( - runtime_control_tx: &std::sync::mpsc::Sender, - pending_mode: PendingRuntimeMode, -) { - if pending_mode == PendingRuntimeMode::PauseUntilResumed { - let _ = runtime_control_tx.send(RuntimeControlCommand::Resume); - } -} - -fn terminate_paused_runtime( - runtime_control_tx: &std::sync::mpsc::Sender, - pending_mode: PendingRuntimeMode, -) { - if pending_mode == PendingRuntimeMode::PauseUntilResumed { - let _ = runtime_control_tx.send(RuntimeControlCommand::Terminate); - } -} - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - use std::sync::Arc; - use std::sync::atomic::AtomicU64; - use std::sync::atomic::Ordering; - use std::time::Duration; - - use codex_protocol::ToolName; - use pretty_assertions::assert_eq; - use tokio::sync::Mutex; - use tokio::sync::mpsc; - use tokio::sync::oneshot; - - use super::CellControlCommand; - use super::CellControlContext; - use super::CellId; - use super::CellResponseSender; - use super::CodeModeService; - use super::Inner; - use super::NoopCodeModeSessionDelegate; - use super::PendingRuntimeMode; - use super::RuntimeCommand; - use super::RuntimeResponse; - use super::WaitOutcome; - use super::WaitRequest; - use super::WaitToPendingOutcome; - use super::WaitToPendingRequest; - use super::run_cell_control; - use crate::CodeModeToolKind; - use crate::FunctionCallOutputContentItem; - use crate::ToolDefinition; - use crate::runtime::ExecuteRequest; - use crate::runtime::ExecuteToPendingOutcome; - use crate::runtime::RuntimeEvent; - use crate::runtime::spawn_runtime; - - fn execute_request(source: &str) -> ExecuteRequest { - ExecuteRequest { - tool_call_id: "call_1".to_string(), - enabled_tools: Vec::new(), - source: source.to_string(), - yield_time_ms: Some(1), - max_output_tokens: None, - } - } - - fn cell_id(value: &str) -> CellId { - CellId::new(value.to_string()) - } - - async fn execute(service: &CodeModeService, request: ExecuteRequest) -> RuntimeResponse { - service - .execute(request) - .await - .unwrap() - .initial_response() - .await - .unwrap() - } - - fn test_inner() -> Arc { - Arc::new(Inner { - stored_values: Mutex::new(HashMap::new()), - cells: Mutex::new(HashMap::new()), - delegate: Arc::new(NoopCodeModeSessionDelegate), - shutting_down: std::sync::atomic::AtomicBool::new(false), - next_cell_id: AtomicU64::new(1), - }) - } - - #[tokio::test] - async fn synchronous_exit_returns_successfully() { - let service = CodeModeService::new(); - - let response = execute( - &service, - ExecuteRequest { - source: r#"text("before"); exit(); text("after");"#.to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "before".to_string(), - }], - error_text: None, - } - ); - } - - #[tokio::test] - async fn stored_values_are_shared_between_cells_but_not_sessions() { - let first_session = CodeModeService::new(); - let second_session = CodeModeService::new(); - - let write_response = execute( - &first_session, - ExecuteRequest { - source: r#"store("key", "visible");"#.to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - - let same_session = execute( - &first_session, - ExecuteRequest { - source: r#"text(String(load("key")));"#.to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - let other_session = execute( - &second_session, - ExecuteRequest { - source: r#"text(String(load("key")));"#.to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - - assert_eq!( - write_response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: Vec::new(), - error_text: None, - } - ); - assert_eq!( - same_session, - RuntimeResponse::Result { - cell_id: cell_id("2"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "visible".to_string(), - }], - error_text: None, - } - ); - assert_eq!( - other_session, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "undefined".to_string(), - }], - error_text: None, - } - ); - } - - #[tokio::test] - async fn shutdown_interrupts_cpu_bound_cells() { - let service = CodeModeService::new(); - - let cell = service - .execute(ExecuteRequest { - source: "while (true) {}".to_string(), - ..execute_request("") - }) - .await - .unwrap(); - assert_eq!( - cell.initial_response().await.unwrap(), - RuntimeResponse::Yielded { - cell_id: cell_id("1"), - content_items: Vec::new(), - } - ); - - tokio::time::timeout(Duration::from_secs(1), service.shutdown()) - .await - .unwrap() - .unwrap(); - } - - #[tokio::test] - async fn start_cell_rejects_new_cell_after_shutdown_begins() { - let service = CodeModeService::new(); - service.inner.shutting_down.store(true, Ordering::Release); - let (response_tx, _response_rx) = oneshot::channel(); - - let error = service - .start_cell( - cell_id("late-cell"), - execute_request(""), - CellResponseSender::Runtime(response_tx), - Some(/*initial_yield_time_ms*/ 1), - PendingRuntimeMode::Continue, - ) - .await - .unwrap_err(); - - assert_eq!(error, "code mode session is shutting down".to_string()); - assert!(service.inner.cells.lock().await.is_empty()); - } - - #[tokio::test] - async fn execute_to_pending_returns_completed_for_synchronous_results() { - let service = CodeModeService::new(); - - let response = service - .execute_to_pending(ExecuteRequest { - source: r#"text("done");"#.to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) - .await - .unwrap(); - - assert_eq!( - response, - ExecuteToPendingOutcome::Completed(RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "done".to_string(), - }], - error_text: None, - }) - ); - } - - #[tokio::test] - async fn execute_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("") - }), - ) - .await - .unwrap() - .unwrap(); - - assert_eq!( - response, - ExecuteToPendingOutcome::Pending { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "before".to_string(), - }], - pending_tool_call_ids: Vec::new(), - } - ); - - let termination = service.terminate(cell_id("1")).await.unwrap(); - - assert_eq!( - termination, - WaitOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: cell_id("1"), - content_items: Vec::new(), - }) - ); - } - - #[tokio::test] - async fn execute_to_pending_identifies_tool_calls_in_paused_frontier() { - let service = CodeModeService::new(); - - let response = service - .execute_to_pending(ExecuteRequest { - enabled_tools: vec![ToolDefinition { - name: "echo".to_string(), - tool_name: ToolName::plain("echo"), - description: String::new(), - kind: CodeModeToolKind::Function, - input_schema: None, - output_schema: None, - }], - source: r#" -await Promise.all([ - tools.echo({ value: "first" }), - tools.echo({ value: "second" }), -]); -"# - .to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) - .await - .unwrap(); - - assert_eq!( - response, - ExecuteToPendingOutcome::Pending { - cell_id: cell_id("1"), - content_items: Vec::new(), - pending_tool_call_ids: vec!["tool-1".to_string(), "tool-2".to_string()], - } - ); - - let termination = service.terminate(cell_id("1")).await.unwrap(); - - assert_eq!( - termination, - WaitOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: cell_id("1"), - content_items: Vec::new(), - }) - ); - } - - #[tokio::test] - async fn execute_to_pending_excludes_delayed_timeout_tool_calls_until_wait() { - let service = CodeModeService::new(); - - let initial_response = service - .execute_to_pending(ExecuteRequest { - enabled_tools: vec![ToolDefinition { - name: "echo".to_string(), - tool_name: ToolName::plain("echo"), - description: String::new(), - kind: CodeModeToolKind::Function, - input_schema: None, - output_schema: None, - }], - source: r#" -setTimeout(() => { - tools.echo({ value: "delayed" }); -}, 1000); -await Promise.all([ - tools.echo({ value: "second" }), - tools.echo({ value: "third" }), -]); -"# - .to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) - .await - .unwrap(); - - assert_eq!( - initial_response, - ExecuteToPendingOutcome::Pending { - cell_id: cell_id("1"), - content_items: Vec::new(), - pending_tool_call_ids: vec!["tool-1".to_string(), "tool-2".to_string()], - } - ); - - let runtime_tx = service - .inner - .cells - .lock() - .await - .get(&cell_id("1")) - .unwrap() - .runtime_tx - .clone(); - runtime_tx - .send(RuntimeCommand::TimeoutFired { id: 1 }) - .unwrap(); - - let resumed_response = tokio::time::timeout( - Duration::from_secs(1), - service.wait_to_pending(WaitToPendingRequest { - cell_id: cell_id("1"), - }), - ) - .await - .unwrap() - .unwrap(); - - assert_eq!( - resumed_response, - WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending { - cell_id: cell_id("1"), - content_items: Vec::new(), - pending_tool_call_ids: vec!["tool-3".to_string()], - }) - ); - - let termination = service.terminate(cell_id("1")).await.unwrap(); - - assert_eq!( - termination, - WaitOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: cell_id("1"), - content_items: Vec::new(), - }) - ); - } - - #[tokio::test] - async fn wait_to_pending_returns_after_resumed_runtime_becomes_quiescent_again() { - let service = CodeModeService::new(); - - let initial_response = service - .execute_to_pending(ExecuteRequest { - source: r#" -await new Promise((resolve) => setTimeout(resolve, 60_000)); -text("after"); -await new Promise(() => {}); -"# - .to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) - .await - .unwrap(); - - assert_eq!( - initial_response, - ExecuteToPendingOutcome::Pending { - cell_id: cell_id("1"), - content_items: Vec::new(), - pending_tool_call_ids: Vec::new(), - } - ); - - let runtime_tx = service - .inner - .cells - .lock() - .await - .get(&cell_id("1")) - .unwrap() - .runtime_tx - .clone(); - runtime_tx - .send(RuntimeCommand::TimeoutFired { id: 1 }) - .unwrap(); - - let resumed_response = tokio::time::timeout( - Duration::from_secs(1), - service.wait_to_pending(WaitToPendingRequest { - cell_id: cell_id("1"), - }), - ) - .await - .unwrap() - .unwrap(); - - assert_eq!( - resumed_response, - WaitToPendingOutcome::LiveCell(ExecuteToPendingOutcome::Pending { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "after".to_string(), - }], - pending_tool_call_ids: Vec::new(), - }) - ); - - let termination = service.terminate(cell_id("1")).await.unwrap(); - - assert_eq!( - termination, - WaitOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: cell_id("1"), - content_items: Vec::new(), - }) - ); - } - - #[tokio::test] - async fn wait_to_pending_returns_completed_after_resumed_runtime_finishes() { - let service = CodeModeService::new(); - - let initial_response = service - .execute_to_pending(ExecuteRequest { - source: r#" -await new Promise((resolve) => setTimeout(resolve, 60_000)); -text("done"); -"# - .to_string(), - yield_time_ms: Some(60_000), - ..execute_request("") - }) - .await - .unwrap(); - - assert_eq!( - initial_response, - ExecuteToPendingOutcome::Pending { - cell_id: cell_id("1"), - content_items: Vec::new(), - pending_tool_call_ids: Vec::new(), - } - ); - - let runtime_tx = service - .inner - .cells - .lock() - .await - .get(&cell_id("1")) - .unwrap() - .runtime_tx - .clone(); - runtime_tx - .send(RuntimeCommand::TimeoutFired { id: 1 }) - .unwrap(); - - let resumed_response = tokio::time::timeout( - Duration::from_secs(1), - service.wait_to_pending(WaitToPendingRequest { - cell_id: cell_id("1"), - }), - ) - .await - .unwrap() - .unwrap(); - - 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, - } - )) - ); - } - - #[tokio::test] - async fn v8_console_is_not_exposed_on_global_this() { - let service = CodeModeService::new(); - - let response = execute( - &service, - ExecuteRequest { - source: r#"text(String(Object.hasOwn(globalThis, "console")));"#.to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "false".to_string(), - }], - error_text: None, - } - ); - } - - #[tokio::test] - async fn date_locale_string_formats_with_icu_data() { - let service = CodeModeService::new(); - - let response = execute( - &service, - ExecuteRequest { - source: r#" -const value = new Date("2025-01-02T03:04:05Z") - .toLocaleString("fr-FR", { - weekday: "long", - month: "long", - day: "numeric", - hour: "2-digit", - minute: "2-digit", - second: "2-digit", - hour12: false, - timeZone: "UTC", - }); -text(value); -"# - .to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "jeudi 2 janvier \u{e0} 03:04:05".to_string(), - }], - error_text: None, - } - ); - } - - #[tokio::test] - async fn intl_date_time_format_formats_with_icu_data() { - let service = CodeModeService::new(); - - let response = execute( - &service, - ExecuteRequest { - source: r#" -const formatter = new Intl.DateTimeFormat("fr-FR", { - weekday: "long", - month: "long", - day: "numeric", - hour: "2-digit", - minute: "2-digit", - second: "2-digit", - hour12: false, - timeZone: "UTC", -}); -text(formatter.format(new Date("2025-01-02T03:04:05Z"))); -"# - .to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "jeudi 2 janvier \u{e0} 03:04:05".to_string(), - }], - error_text: None, - } - ); - } - - #[tokio::test] - async fn output_helpers_return_undefined() { - let service = CodeModeService::new(); - - let response = execute( - &service, - ExecuteRequest { - source: r#" -const returnsUndefined = [ - text("first"), - image("https://example.com/image.jpg"), - notify("ping"), -].map((value) => value === undefined); -text(JSON.stringify(returnsUndefined)); -"# - .to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![ - FunctionCallOutputContentItem::InputText { - text: "first".to_string(), - }, - FunctionCallOutputContentItem::InputImage { - image_url: "https://example.com/image.jpg".to_string(), - detail: Some(crate::DEFAULT_IMAGE_DETAIL), - }, - FunctionCallOutputContentItem::InputText { - text: "[true,true,true]".to_string(), - }, - ], - error_text: None, - } - ); - } - - #[tokio::test] - async fn image_helper_accepts_raw_mcp_image_block_with_original_detail() { - let service = CodeModeService::new(); - - let response = execute( - &service, - ExecuteRequest { - source: r#" -image({ - type: "image", - data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==", - mimeType: "image/png", - _meta: { "codex/imageDetail": "original" }, -}); -"# - .to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputImage { - image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==".to_string(), - detail: Some(crate::ImageDetail::Original), - }], - error_text: None, - } - ); - } - - #[tokio::test] - async fn image_helper_second_arg_overrides_explicit_object_detail() { - let service = CodeModeService::new(); - - let response = execute( - &service, - ExecuteRequest { - source: r#" -image( - { - image_url: "https://example.com/image.jpg", - detail: "high", - }, - "original", -); -"# - .to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputImage { - image_url: "https://example.com/image.jpg".to_string(), - detail: Some(crate::ImageDetail::Original), - }], - error_text: None, - } - ); - } - - #[tokio::test] - async fn image_helper_second_arg_overrides_raw_mcp_image_detail() { - let service = CodeModeService::new(); - - let response = execute( - &service, - ExecuteRequest { - source: r#" -image( - { - type: "image", - data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==", - mimeType: "image/png", - _meta: { "codex/imageDetail": "original" }, - }, - "high", -); -"# - .to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputImage { - image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==".to_string(), - detail: Some(crate::ImageDetail::High), - }], - error_text: None, - } - ); - } - - #[tokio::test] - async fn image_helper_accepts_low_detail() { - let service = CodeModeService::new(); - - let response = execute( - &service, - ExecuteRequest { - source: r#" -image({ - image_url: "https://example.com/image.jpg", - detail: "low", -}); -"# - .to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputImage { - image_url: "https://example.com/image.jpg".to_string(), - detail: Some(crate::ImageDetail::Low), - }], - error_text: None, - } - ); - } - - #[tokio::test] - async fn image_helper_rejects_unsupported_detail() { - let service = CodeModeService::new(); - - let response = execute( - &service, - ExecuteRequest { - source: r#" -image({ - image_url: "https://example.com/image.jpg", - detail: "medium", -}); -"# - .to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: Vec::new(), - error_text: Some( - "image detail must be one of: auto, low, high, original".to_string() - ), - } - ); - } - - #[tokio::test] - async fn image_helper_rejects_raw_mcp_result_container() { - let service = CodeModeService::new(); - - let response = execute( - &service, - ExecuteRequest { - source: r#" -image({ - content: [ - { - type: "image", - data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==", - mimeType: "image/png", - _meta: { "codex/imageDetail": "original" }, - }, - ], - isError: false, -}); -"# - .to_string(), - yield_time_ms: None, - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: Vec::new(), - error_text: Some( - "image expects a non-empty image URL string, an object with image_url and optional detail, or a raw MCP image block".to_string(), - ), - } - ); - } - - #[tokio::test] - async fn wait_reports_missing_cell_separately_from_runtime_results() { - let service = CodeModeService::new(); - - let response = service - .wait(WaitRequest { - cell_id: cell_id("missing"), - yield_time_ms: 1, - }) - .await - .unwrap(); - - assert_eq!( - response, - WaitOutcome::MissingCell(RuntimeResponse::Result { - cell_id: cell_id("missing"), - content_items: Vec::new(), - error_text: Some("exec cell missing not found".to_string()), - }) - ); - } - - #[tokio::test] - async fn terminate_waits_for_runtime_shutdown_before_responding() { - let inner = test_inner(); - let (event_tx, event_rx) = mpsc::unbounded_channel(); - let (control_tx, control_rx) = mpsc::unbounded_channel(); - let (initial_response_tx, initial_response_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 { - source: "await new Promise(() => {})".to_string(), - yield_time_ms: None, - ..execute_request("") - }, - runtime_event_tx, - PendingRuntimeMode::Continue, - ) - .unwrap(); - - tokio::spawn(run_cell_control( - inner, - CellControlContext { - cell_id: cell_id("cell-1"), - runtime_tx: runtime_tx.clone(), - runtime_control_tx, - pending_mode: PendingRuntimeMode::Continue, - runtime_terminate_handle, - cancellation_token: tokio_util::sync::CancellationToken::new(), - }, - event_rx, - control_rx, - CellResponseSender::Runtime(initial_response_tx), - Some(/*initial_yield_time_ms*/ 60_000), - )); - - event_tx.send(RuntimeEvent::Started).unwrap(); - event_tx.send(RuntimeEvent::YieldRequested).unwrap(); - assert_eq!( - initial_response_rx.await.unwrap(), - RuntimeResponse::Yielded { - cell_id: cell_id("cell-1"), - content_items: Vec::new(), - } - ); - - let (terminate_response_tx, terminate_response_rx) = oneshot::channel(); - control_tx - .send(CellControlCommand::Terminate { - response_tx: terminate_response_tx, - }) - .unwrap(); - let terminate_response = async { terminate_response_rx.await.unwrap() }; - tokio::pin!(terminate_response); - assert!( - tokio::time::timeout(Duration::from_millis(100), terminate_response.as_mut()) - .await - .is_err() - ); - - drop(event_tx); - - assert_eq!( - terminate_response.await, - RuntimeResponse::Terminated { - cell_id: cell_id("cell-1"), - content_items: Vec::new(), - } - ); - - let _ = runtime_tx.send(RuntimeCommand::Terminate); + Box::pin(async { Err("code mode runtime is unavailable in this build".to_string()) }) } } diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index e87ee82f30..3051041cbc 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -24,6 +24,7 @@ pub use codex_config::types::TuiKeymap; pub use codex_config::types::TuiNotificationSettings; pub use codex_config::types::TuiPetAnchor; pub use codex_config::types::UriBasedFileOpener; +pub use codex_core::CodeModeSessionProvider; pub use codex_core::CodexThread; pub use codex_core::ForkSnapshot; pub use codex_core::McpManager; diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 4f6ab2e777..4a9cb6692d 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -132,6 +132,7 @@ codex-shell-escalation = { workspace = true } [dev-dependencies] assert_cmd = { workspace = true } assert_matches = { workspace = true } +codex-code-mode-runtime = { workspace = true } codex-otel = { workspace = true } codex-test-binary-support = { workspace = true } codex-utils-cargo-bin = { workspace = true } diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 678a6e5d10..54e0a059ae 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -89,6 +89,7 @@ pub(crate) async fn run_codex_thread_interactive( plugins_manager: Arc::clone(&parent_session.services.plugins_manager), mcp_manager: Arc::clone(&parent_session.services.mcp_manager), extensions: Arc::clone(&parent_session.services.extensions), + code_mode_session_provider: Arc::clone(&parent_session.services.code_mode_session_provider), conversation_history, session_source: SessionSource::SubAgent(subagent_source.clone()), forked_from_thread_id, diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 31ac529556..03dfb9741e 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -104,6 +104,7 @@ pub use client::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER; pub use codex_protocol::config_types::ModelProviderAuthInfo; mod event_mapping; pub mod review_format; +pub use codex_code_mode::CodeModeSessionProvider; pub use codex_prompts as review_prompts; mod thread_manager; pub(crate) mod web_search; diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index f6bf7d0022..309238786e 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -40,7 +40,7 @@ pub async fn build_prompt_input( let thread_store = thread_store_from_config(&config, state_db.clone()); let installation_id = resolve_installation_id(&config.codex_home).await?; - let thread_manager = ThreadManager::new( + let thread_manager = ThreadManager::new_without_code_mode_runtime( &config, Arc::clone(&auth_manager), SessionSource::Exec, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 5a92c9c825..9f32040293 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -314,6 +314,7 @@ use crate::turn_timing::TurnTimingState; use crate::turn_timing::record_turn_ttfm_metric; use crate::unified_exec::UnifiedExecProcessManager; use crate::windows_sandbox::WindowsSandboxLevelExt; +use codex_code_mode::CodeModeSessionProvider; use codex_core_plugins::PluginsManager; use codex_git_utils::get_git_repo_root; use codex_mcp::compute_auth_statuses; @@ -399,6 +400,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) plugins_manager: Arc, pub(crate) mcp_manager: Arc, pub(crate) extensions: Arc>, + pub(crate) code_mode_session_provider: Arc, pub(crate) conversation_history: InitialHistory, pub(crate) session_source: SessionSource, pub(crate) forked_from_thread_id: Option, @@ -483,6 +485,7 @@ impl Codex { plugins_manager, mcp_manager, extensions, + code_mode_session_provider, conversation_history, session_source, forked_from_thread_id, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 05773abb8c..dd86726ca7 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1051,7 +1051,10 @@ impl Session { session_configuration.parent_thread_id, ), ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_session_provider: Arc::clone(&code_mode_session_provider), + code_mode_service: crate::tools::code_mode::CodeModeService::new( + code_mode_session_provider, + ), environment_manager, }; services diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index a3e5558971..54d1781271 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4778,7 +4778,10 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { Session::build_model_client_beta_features_header(config.as_ref()), /*attestation_provider*/ None, ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_session_provider: Arc::new(codex_code_mode::UnavailableCodeModeSessionProvider), + code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::new( + codex_code_mode::UnavailableCodeModeSessionProvider, + )), environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), }; @@ -6868,7 +6871,10 @@ where Session::build_model_client_beta_features_header(config.as_ref()), /*attestation_provider*/ None, ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_session_provider: Arc::new(codex_code_mode::UnavailableCodeModeSessionProvider), + code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::new( + codex_code_mode::UnavailableCodeModeSessionProvider, + )), environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), }; diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index d8b46c9999..8d65053f36 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -713,6 +713,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { plugins_manager, mcp_manager, extensions: codex_extension_api::empty_extension_registry(), + code_mode_session_provider: Arc::new(codex_code_mode::UnavailableCodeModeSessionProvider), conversation_history: InitialHistory::New, session_source: SessionSource::SubAgent(SubAgentSource::Other( GUARDIAN_REVIEWER_NAME.to_string(), diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 59def9c7ab..4beaa84cd0 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -18,6 +18,7 @@ use crate::unified_exec::UnifiedExecProcessManager; use arc_swap::ArcSwap; use arc_swap::ArcSwapOption; use codex_analytics::AnalyticsEventsClient; +use codex_code_mode::CodeModeSessionProvider; use codex_core_plugins::PluginsManager; use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionData; @@ -77,6 +78,7 @@ pub(crate) struct SessionServices { pub(crate) attestation_provider: Option>, /// Session-scoped model client shared across turns. pub(crate) model_client: ModelClient, + pub(crate) code_mode_session_provider: Arc, pub(crate) code_mode_service: CodeModeService, /// Shared process-level environment registry. Sessions carry an `Arc` handle so they can pass /// the same manager through child-thread spawn paths without reconstructing it. diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index d130a35815..40cd61b881 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -19,6 +19,8 @@ use crate::tasks::interrupted_turn_history_marker; use codex_analytics::AnalyticsEventsClient; use codex_app_server_protocol::ThreadHistoryBuilder; use codex_app_server_protocol::TurnStatus; +use codex_code_mode::CodeModeSessionProvider; +use codex_code_mode::UnavailableCodeModeSessionProvider; use codex_core_plugins::PluginsManager; use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionRegistry; @@ -208,6 +210,7 @@ pub(crate) struct ThreadManagerState { mcp_manager: Arc, extensions: Arc>, thread_store: Arc, + code_mode_session_provider: Arc, attestation_provider: Option>, session_source: SessionSource, installation_id: String, @@ -256,6 +259,7 @@ impl ThreadManager { auth_manager: Arc, session_source: SessionSource, environment_manager: Arc, + code_mode_session_provider: Arc, extensions: Arc>, analytics_events_client: Option, thread_store: Arc, @@ -287,6 +291,7 @@ impl ThreadManager { mcp_manager, extensions, thread_store, + code_mode_session_provider, attestation_provider, auth_manager, session_source, @@ -300,6 +305,34 @@ impl ThreadManager { } } + #[allow(clippy::too_many_arguments)] + pub fn new_without_code_mode_runtime( + config: &Config, + auth_manager: Arc, + session_source: SessionSource, + environment_manager: Arc, + extensions: Arc>, + analytics_events_client: Option, + thread_store: Arc, + state_db: Option, + installation_id: String, + attestation_provider: Option>, + ) -> Self { + Self::new( + config, + auth_manager, + session_source, + environment_manager, + Arc::new(UnavailableCodeModeSessionProvider), + extensions, + analytics_events_client, + thread_store, + state_db, + installation_id, + attestation_provider, + ) + } + /// Construct with a dummy AuthManager containing the provided CodexAuth. /// Used for integration tests: should not be used by ordinary business logic. pub(crate) fn with_models_provider_for_tests( @@ -388,6 +421,7 @@ impl ThreadManager { mcp_manager, extensions: empty_extension_registry(), thread_store, + code_mode_session_provider: Arc::new(UnavailableCodeModeSessionProvider), attestation_provider: None, auth_manager, session_source: SessionSource::Exec, @@ -1318,6 +1352,7 @@ impl ThreadManagerState { plugins_manager: Arc::clone(&self.plugins_manager), mcp_manager: Arc::clone(&self.mcp_manager), extensions: Arc::clone(&self.extensions), + code_mode_session_provider: Arc::clone(&self.code_mode_session_provider), conversation_history: initial_history, session_source, forked_from_thread_id, diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 52a0f8eb3b..59e0807c05 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -489,7 +489,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let manager = ThreadManager::new( + let manager = ThreadManager::new_without_code_mode_runtime( &config, auth_manager.clone(), SessionSource::Exec, @@ -605,7 +605,7 @@ async fn explicit_installation_id_skips_codex_home_file() { let installation_id = uuid::Uuid::new_v4().to_string(); let state_db = init_state_db(&config).await; let thread_store = thread_store_from_config(&config, state_db.clone()); - let manager = ThreadManager::new( + let manager = ThreadManager::new_without_code_mode_runtime( &config, auth_manager, SessionSource::Exec, @@ -644,7 +644,7 @@ async fn resume_active_thread_from_rollout_returns_running_thread() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let manager = ThreadManager::new( + let manager = ThreadManager::new_without_code_mode_runtime( &config, auth_manager.clone(), SessionSource::Exec, @@ -701,7 +701,7 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let manager = ThreadManager::new( + let manager = ThreadManager::new_without_code_mode_runtime( &config, auth_manager.clone(), SessionSource::Exec, @@ -765,7 +765,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let state_db = init_state_db(&config).await; let thread_store = thread_store_from_config(&config, state_db.clone()); - let manager = ThreadManager::new( + let manager = ThreadManager::new_without_code_mode_runtime( &config, auth_manager.clone(), SessionSource::Exec, @@ -854,7 +854,7 @@ async fn rollout_path_resume_and_fork_read_history_through_thread_store() { .as_any() .downcast_ref::() .expect("configured in-memory store"); - let manager = ThreadManager::new( + let manager = ThreadManager::new_without_code_mode_runtime( &config, auth_manager.clone(), SessionSource::Exec, @@ -955,7 +955,7 @@ async fn new_uses_active_provider_for_model_refresh() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let manager = ThreadManager::new( + let manager = ThreadManager::new_without_code_mode_runtime( &config, auth_manager, SessionSource::Exec, @@ -1175,7 +1175,7 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let state_db = init_state_db(&config).await; - let manager = ThreadManager::new( + let manager = ThreadManager::new_without_code_mode_runtime( &config, auth_manager.clone(), SessionSource::Exec, @@ -1281,7 +1281,7 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let state_db = init_state_db(&config).await; - let manager = ThreadManager::new( + let manager = ThreadManager::new_without_code_mode_runtime( &config, auth_manager.clone(), SessionSource::Exec, @@ -1377,7 +1377,7 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let state_db = init_state_db(&config).await; - let manager = ThreadManager::new( + let manager = ThreadManager::new_without_code_mode_runtime( &config, auth_manager.clone(), SessionSource::Exec, @@ -1517,7 +1517,7 @@ async fn resumed_thread_keeps_paused_goal_paused() -> anyhow::Result<()> { let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); let state_db = init_state_db(&config).await; - let manager = ThreadManager::new( + let manager = ThreadManager::new_without_code_mode_runtime( &config, auth_manager.clone(), SessionSource::Exec, diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index ade23cf3b7..17d51f101f 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -11,10 +11,12 @@ use std::time::Duration; use codex_code_mode::CellId; use codex_code_mode::CodeModeNestedToolCall; use codex_code_mode::CodeModeSession; +use codex_code_mode::CodeModeSessionProvider; use codex_code_mode::CodeModeToolKind; use codex_code_mode::RuntimeResponse; use codex_protocol::models::FunctionCallOutputContentItem; use serde_json::Value as JsonValue; +use tokio::sync::OnceCell; use tokio_util::sync::CancellationToken; use crate::function_tool::FunctionCallError; @@ -58,17 +60,17 @@ pub(crate) struct ExecContext { } pub(crate) struct CodeModeService { - session: Option>, + session: OnceCell>, + session_provider: Arc, dispatch_broker: Arc, } impl CodeModeService { - pub(crate) fn new() -> Self { + pub(crate) fn new(session_provider: Arc) -> Self { let dispatch_broker = Arc::new(CodeModeDispatchBroker::new()); Self { - session: Some(Arc::new(codex_code_mode::CodeModeService::with_delegate( - dispatch_broker.clone(), - ))), + session: OnceCell::new(), + session_provider, dispatch_broker, } } @@ -77,25 +79,25 @@ impl CodeModeService { &self, request: codex_code_mode::ExecuteRequest, ) -> Result { - self.session()?.execute(request).await + self.session().await?.execute(request).await } pub(crate) async fn wait( &self, request: codex_code_mode::WaitRequest, ) -> Result { - self.session()?.wait(request).await + self.session().await?.wait(request).await } pub(crate) async fn terminate( &self, cell_id: CellId, ) -> Result { - self.session()?.terminate(cell_id).await + self.session().await?.terminate(cell_id).await } pub(crate) async fn shutdown(&self) -> Result<(), String> { - match &self.session { + match self.session.get() { Some(session) => session.shutdown().await, None => Ok(()), } @@ -116,9 +118,7 @@ impl CodeModeService { router: Arc, tracker: SharedTurnDiffTracker, ) -> Option { - if !matches!(turn.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) - || self.session.is_none() - { + if !matches!(turn.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) { return None; } @@ -132,10 +132,13 @@ impl CodeModeService { ) } - fn session(&self) -> Result<&Arc, String> { + async fn session(&self) -> Result<&Arc, String> { self.session - .as_ref() - .ok_or_else(|| "code mode is unavailable".to_string()) + .get_or_try_init(|| { + self.session_provider + .create_session(self.dispatch_broker.clone()) + }) + .await } } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index d4ce84ad55..3a608242d0 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -3982,7 +3982,7 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr .enable(Feature::Sqlite) .expect("test config should allow sqlite"); let state_db = init_state_db(&config).await; - let manager = ThreadManager::new( + let manager = ThreadManager::new_without_code_mode_runtime( &config, AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")), SessionSource::Exec, diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 4fe0609a71..4c7f326e50 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -473,6 +473,7 @@ impl TestCodexBuilder { codex_core::test_support::auth_manager_from_auth(auth.clone()), SessionSource::Exec, Arc::clone(&environment_manager), + Arc::new(codex_code_mode_runtime::InProcessCodeModeSessionProvider), empty_extension_registry(), /*analytics_events_client*/ None, thread_store, diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index c954a954fb..c724f10291 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1131,7 +1131,7 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { let installation_id = resolve_installation_id(&config.codex_home) .await .expect("resolve installation id"); - let thread_manager = ThreadManager::new( + let thread_manager = ThreadManager::new_without_code_mode_runtime( &config, auth_manager, SessionSource::Exec, diff --git a/codex-rs/mcp-server/Cargo.toml b/codex-rs/mcp-server/Cargo.toml index 29b2d6c7d7..ac33149641 100644 --- a/codex-rs/mcp-server/Cargo.toml +++ b/codex-rs/mcp-server/Cargo.toml @@ -19,6 +19,7 @@ workspace = true [dependencies] anyhow = { workspace = true } codex-arg0 = { workspace = true } +codex-code-mode-runtime = { workspace = true } codex-config = { workspace = true } codex-core = { workspace = true } codex-exec-server = { workspace = true } diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index 2f85b35143..bb0e6753af 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -67,6 +67,7 @@ impl MessageProcessor { auth_manager, SessionSource::Mcp, environment_manager, + Arc::new(codex_code_mode_runtime::InProcessCodeModeSessionProvider), empty_extension_registry(), /*analytics_events_client*/ None, codex_core::thread_store_from_config(config.as_ref(), state_db.clone()), diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 9305a6f668..5106f40d76 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -120,7 +120,7 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { .await?, ); let installation_id = resolve_installation_id(&config.codex_home).await?; - let thread_manager = ThreadManager::new( + let thread_manager = ThreadManager::new_without_code_mode_runtime( &config, auth_manager, SessionSource::Exec,