From afef87eb077656851f40af8c5f33d388f91eb6dc Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Tue, 16 Jun 2026 05:11:03 +0000 Subject: [PATCH] code-mode: cut production over to stdio host --- codex-rs/Cargo.lock | 8 +- codex-rs/app-server/BUILD.bazel | 1 + .../tests/common/test_app_server.rs | 3 + .../code-mode-host/tests/provider_contract.rs | 49 +- codex-rs/code-mode/src/cell_actor/tests.rs | 36 +- codex-rs/code-mode/src/lib.rs | 5 - codex-rs/code-mode/src/runtime/mod.rs | 2 +- codex-rs/code-mode/src/service.rs | 369 ----- .../code-mode/src/service_contract_tests.rs | 830 ---------- codex-rs/code-mode/src/service_tests.rs | 1408 ----------------- codex-rs/code-mode/src/session_runtime/mod.rs | 9 - codex-rs/core/BUILD.bazel | 1 + codex-rs/core/Cargo.toml | 3 +- codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/session/mod.rs | 4 + codex-rs/core/src/session/session.rs | 6 +- codex-rs/core/src/session/tests.rs | 17 +- .../core/src/session/tests/guardian_tests.rs | 3 + codex-rs/core/src/state/service.rs | 2 + codex-rs/core/src/thread_manager.rs | 18 + codex-rs/core/src/tools/code_mode/delegate.rs | 90 +- .../src/tools/code_mode/delegate_tests.rs | 79 +- .../src/tools/code_mode/execute_handler.rs | 18 +- .../core/src/tools/code_mode/execute_spec.rs | 16 +- codex-rs/core/src/tools/code_mode/mod.rs | 301 +++- .../src/tools/code_mode/response_adapter.rs | 25 +- .../core/src/tools/code_mode/wait_handler.rs | 21 +- .../core/src/tools/code_mode/wait_spec.rs | 12 +- codex-rs/core/src/tools/spec_plan.rs | 29 +- codex-rs/core/src/tools/spec_plan_tests.rs | 33 +- codex-rs/core/tests/common/Cargo.toml | 1 + codex-rs/core/tests/common/test_codex.rs | 10 + .../tests/suite/model_runtime_selectors.rs | 8 +- codex-rs/rollout-trace/Cargo.toml | 2 +- codex-rs/rollout-trace/src/code_cell.rs | 2 +- codex-rs/rollout-trace/src/tool_dispatch.rs | 6 +- codex-rs/tools/Cargo.toml | 2 +- codex-rs/tools/src/code_mode.rs | 19 +- codex-rs/tools/src/code_mode_tests.rs | 8 +- scripts/run-nextest.py | 7 + 40 files changed, 598 insertions(+), 2866 deletions(-) delete mode 100644 codex-rs/code-mode/src/service.rs delete mode 100644 codex-rs/code-mode/src/service_contract_tests.rs delete mode 100644 codex-rs/code-mode/src/service_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 7c24c3dc9b..e14a907e46 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2641,7 +2641,8 @@ dependencies = [ "codex-app-server-protocol", "codex-apply-patch", "codex-async-utils", - "codex-code-mode", + "codex-code-mode-client", + "codex-code-mode-protocol", "codex-config", "codex-connectors", "codex-context-fragments", @@ -3800,7 +3801,7 @@ name = "codex-rollout-trace" version = "0.0.0" dependencies = [ "anyhow", - "codex-code-mode", + "codex-code-mode-protocol", "codex-protocol", "http 1.4.0", "pretty_assertions", @@ -4016,7 +4017,7 @@ name = "codex-tools" version = "0.0.0" dependencies = [ "codex-app-server-protocol", - "codex-code-mode", + "codex-code-mode-protocol", "codex-features", "codex-file-system", "codex-protocol", @@ -4675,6 +4676,7 @@ dependencies = [ "assert_cmd", "base64 0.22.1", "codex-arg0", + "codex-code-mode-client", "codex-config", "codex-core", "codex-exec-server", diff --git a/codex-rs/app-server/BUILD.bazel b/codex-rs/app-server/BUILD.bazel index 534e7209cd..26a102043c 100644 --- a/codex-rs/app-server/BUILD.bazel +++ b/codex-rs/app-server/BUILD.bazel @@ -5,6 +5,7 @@ codex_rust_crate( crate_name = "codex_app_server", extra_binaries = [ "//codex-rs/bwrap:bwrap", + "//codex-rs/code-mode-host:codex-code-mode-host", ], extra_binaries_non_windows = [ "//codex-rs/cli:codex", diff --git a/codex-rs/app-server/tests/common/test_app_server.rs b/codex-rs/app-server/tests/common/test_app_server.rs index e940d9a42b..4994157fcb 100644 --- a/codex-rs/app-server/tests/common/test_app_server.rs +++ b/codex-rs/app-server/tests/common/test_app_server.rs @@ -224,6 +224,9 @@ impl TestAppServer { cmd.current_dir(codex_home); cmd.env("CODEX_HOME", codex_home); cmd.env("RUST_LOG", "warn"); + let code_mode_host = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host") + .context("should find binary for codex-code-mode-host")?; + cmd.env("CODEX_CODE_MODE_HOST_PATH", code_mode_host); // Keep integration tests isolated from host managed configuration. cmd.env( "CODEX_APP_SERVER_MANAGED_CONFIG_PATH", diff --git a/codex-rs/code-mode-host/tests/provider_contract.rs b/codex-rs/code-mode-host/tests/provider_contract.rs index b7f2766d8a..a08d450d13 100644 --- a/codex-rs/code-mode-host/tests/provider_contract.rs +++ b/codex-rs/code-mode-host/tests/provider_contract.rs @@ -3,23 +3,22 @@ use std::sync::Arc; use std::time::Duration; -use codex_code_mode::CellId; -use codex_code_mode::CellOutcome; -use codex_code_mode::CodeModeNestedToolCall; -use codex_code_mode::CodeModeSession; -use codex_code_mode::CodeModeSessionDelegate; -use codex_code_mode::CodeModeSessionProvider; -use codex_code_mode::CodeModeToolKind; -use codex_code_mode::CreateCellRequest; -use codex_code_mode::FunctionCallOutputContentItem; -use codex_code_mode::InProcessCodeModeSessionProvider; -use codex_code_mode::NotificationFuture; -use codex_code_mode::ObserveRequest; -use codex_code_mode::RuntimeResponse; -use codex_code_mode::ToolDefinition; -use codex_code_mode::ToolInvocationFuture; use codex_code_mode_client::CodeModeHostCommand; use codex_code_mode_client::StdioCodeModeSessionProvider; +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CellOutcome; +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeSession; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::CodeModeSessionProvider; +use codex_code_mode_protocol::CodeModeToolKind; +use codex_code_mode_protocol::CreateCellRequest; +use codex_code_mode_protocol::FunctionCallOutputContentItem; +use codex_code_mode_protocol::NotificationFuture; +use codex_code_mode_protocol::ObserveRequest; +use codex_code_mode_protocol::RuntimeResponse; +use codex_code_mode_protocol::ToolDefinition; +use codex_code_mode_protocol::ToolInvocationFuture; use codex_protocol::ToolName; use pretty_assertions::assert_eq; use tokio::sync::mpsc; @@ -33,19 +32,13 @@ struct ProviderCase { fn provider_cases() -> Vec { let host_program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("resolve host binary"); - vec![ - ProviderCase { - name: "in-process", - provider: Arc::new(InProcessCodeModeSessionProvider), - }, - ProviderCase { - name: "stdio", - provider: Arc::new(StdioCodeModeSessionProvider::new(CodeModeHostCommand { - program: host_program, - args: Vec::new(), - })), - }, - ] + vec![ProviderCase { + name: "stdio", + provider: Arc::new(StdioCodeModeSessionProvider::new(CodeModeHostCommand { + program: host_program, + args: Vec::new(), + })), + }] } #[derive(Debug, PartialEq)] diff --git a/codex-rs/code-mode/src/cell_actor/tests.rs b/codex-rs/code-mode/src/cell_actor/tests.rs index 66fcf3fbad..6bb499101d 100644 --- a/codex-rs/code-mode/src/cell_actor/tests.rs +++ b/codex-rs/code-mode/src/cell_actor/tests.rs @@ -119,7 +119,6 @@ fn spawn_cell_actor_harness_with_host(host: Arc) -> CellActorHar source: "await new Promise(() => {});".to_string(), }, runtime_event_tx, - PendingRuntimeMode::PauseUntilResumed, ) .unwrap(); let (runtime_control_tx, runtime_control_rx) = std_mpsc::channel(); @@ -1102,3 +1101,38 @@ fn dropped_pending_observation_preserves_the_initial_yield_boundary() { )); assert_eq!(response_rx.try_recv(), Ok(Ok(completion))); } + +#[tokio::test] +async fn repeated_termination_is_rejected_before_the_first_response() { + let harness = spawn_cell_actor_harness(); + let initial_observation = harness + .handle + .observe(ObserveMode::YieldAfter(Duration::from_secs(60))); + assert_eq!( + harness + .handle + .observe(ObserveMode::YieldAfter(Duration::from_secs(60))) + .await, + Err(CellError::Busy) + ); + let first_termination = harness.handle.terminate(); + + assert_eq!( + harness.handle.terminate().await, + Err(CellError::AlreadyTerminating) + ); + + harness + .event_tx + .send(RuntimeEvent::Result { + stored_value_writes: HashMap::new(), + error_text: None, + }) + .unwrap(); + let terminated = Ok(CellEvent::Terminated { + content_items: Vec::new(), + }); + assert_eq!(first_termination.await, terminated.clone()); + assert_eq!(initial_observation.await, terminated); + harness.task.await.unwrap(); +} diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 079e8962d6..21d5d2a16d 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -1,10 +1,5 @@ mod cell_actor; mod runtime; -mod service; pub mod session_runtime; -pub use codex_code_mode_protocol::*; -pub use service::CodeModeService; -pub use service::InProcessCodeModeSessionProvider; -pub use service::NoopCodeModeSessionDelegate; pub use session_runtime::SessionRuntime; diff --git a/codex-rs/code-mode/src/runtime/mod.rs b/codex-rs/code-mode/src/runtime/mod.rs index ca2536f402..3c7305f4ba 100644 --- a/codex-rs/code-mode/src/runtime/mod.rs +++ b/codex-rs/code-mode/src/runtime/mod.rs @@ -318,7 +318,7 @@ mod tests { use super::RuntimeControlCommand; use super::RuntimeEvent; use super::spawn_runtime; - use crate::FunctionCallOutputContentItem; + use codex_code_mode_protocol::FunctionCallOutputContentItem; fn execute_request(source: &str) -> CreateCellRequest { CreateCellRequest { diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs deleted file mode 100644 index 14f600d970..0000000000 --- a/codex-rs/code-mode/src/service.rs +++ /dev/null @@ -1,369 +0,0 @@ -use std::sync::Arc; -use std::time::Duration; - -use codex_code_mode_protocol::CellId; -use codex_code_mode_protocol::CellOutcome; -use codex_code_mode_protocol::CodeModeNestedToolCall; -use codex_code_mode_protocol::CodeModeSession; -use codex_code_mode_protocol::CodeModeSessionDelegate; -use codex_code_mode_protocol::CodeModeSessionProvider; -use codex_code_mode_protocol::CodeModeSessionProviderFuture; -use codex_code_mode_protocol::CodeModeSessionResultFuture; -use codex_code_mode_protocol::CodeModeToolKind; -use codex_code_mode_protocol::CreateCellRequest; -use codex_code_mode_protocol::FunctionCallOutputContentItem; -use codex_code_mode_protocol::ImageDetail; -use codex_code_mode_protocol::NotificationFuture; -use codex_code_mode_protocol::ObserveRequest; -use codex_code_mode_protocol::ObserveToPendingOutcome; -use codex_code_mode_protocol::ObserveToPendingRequest; -use codex_code_mode_protocol::PendingOutcome; -use codex_code_mode_protocol::RuntimeResponse; -use codex_code_mode_protocol::ToolInvocationFuture; -use serde_json::Value as JsonValue; -use tokio_util::sync::CancellationToken; - -use crate::session_runtime as runtime; -use crate::session_runtime::SessionRuntime; - -pub struct NoopCodeModeSessionDelegate; - -impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate { - fn invoke_tool<'a>( - &'a self, - _invocation: CodeModeNestedToolCall, - cancellation_token: CancellationToken, - ) -> ToolInvocationFuture<'a> { - Box::pin(async move { - cancellation_token.cancelled().await; - Err("code mode nested tools are unavailable".to_string()) - }) - } - - fn notify<'a>( - &'a self, - _call_id: String, - _cell_id: CellId, - _text: String, - _cancellation_token: CancellationToken, - ) -> NotificationFuture<'a> { - Box::pin(async { Ok(()) }) - } - - fn cell_closed(&self, _cell_id: &CellId) {} -} - -#[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) - }) - } -} - -pub struct CodeModeService { - runtime: SessionRuntime, -} - -impl CodeModeService { - pub fn new() -> Self { - Self::with_delegate(Arc::new(NoopCodeModeSessionDelegate)) - } - - pub fn with_delegate(delegate: Arc) -> Self { - #[cfg(not(test))] - let runtime = SessionRuntime::new(Arc::new(ProtocolDelegate { delegate })); - #[cfg(test)] - let runtime = SessionRuntime::new_for_test(Arc::new(ProtocolDelegate { delegate })); - Self { runtime } - } - - pub async fn create_cell(&self, request: CreateCellRequest) -> Result { - let runtime_cell_id = self - .runtime - .create_cell(runtime_request(request)) - .await - .map_err(|error| error.to_string())?; - Ok(protocol_cell_id(&runtime_cell_id)) - } - - pub async fn observe(&self, request: ObserveRequest) -> Result { - self.begin_observe(request).await.await - } - - async fn begin_observe( - &self, - request: ObserveRequest, - ) -> CodeModeSessionResultFuture<'static, CellOutcome> { - let ObserveRequest { - cell_id, - yield_time_ms, - } = request; - let runtime_cell_id = runtime_cell_id(&cell_id); - match self - .runtime - .begin_observe( - &runtime_cell_id, - runtime::ObserveMode::YieldAfter(Duration::from_millis(yield_time_ms)), - ) - .await - { - Ok(pending_event) => Box::pin(async move { - match pending_event.event().await { - Ok(event) => Ok(CellOutcome::LiveCell(runtime_response(&cell_id, event)?)), - Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => { - Ok(CellOutcome::MissingCell(missing_cell_response(cell_id))) - } - Err(error) => Err(error.to_string()), - } - }), - Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => { - missing_observation(cell_id) - } - Err(error) => Box::pin(async move { Err(error.to_string()) }), - } - } - - pub async fn terminate(&self, cell_id: CellId) -> Result { - match self.runtime.terminate(&runtime_cell_id(&cell_id)).await { - Ok(event) => Ok(CellOutcome::LiveCell(runtime_response(&cell_id, event)?)), - Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => { - Ok(CellOutcome::MissingCell(missing_cell_response(cell_id))) - } - Err(error) => Err(error.to_string()), - } - } - - pub async fn observe_to_pending( - &self, - request: ObserveToPendingRequest, - ) -> Result { - let cell_id = request.cell_id; - match self - .runtime - .observe( - &runtime_cell_id(&cell_id), - runtime::ObserveMode::PendingFrontier, - ) - .await - { - Ok(event) => Ok(ObserveToPendingOutcome::LiveCell(pending_outcome( - &cell_id, event, - )?)), - Err(runtime::Error::MissingCell(_) | runtime::Error::ClosedCell(_)) => Ok( - ObserveToPendingOutcome::MissingCell(missing_cell_response(cell_id)), - ), - Err(error) => Err(error.to_string()), - } - } - - pub async fn shutdown(&self) -> Result<(), String> { - self.runtime - .shutdown() - .await - .map_err(|error| error.to_string()) - } -} - -impl Default for CodeModeService { - fn default() -> Self { - Self::new() - } -} - -impl CodeModeSession for CodeModeService { - fn is_alive(&self) -> bool { - self.runtime.is_alive() - } - - fn create_cell<'a>( - &'a self, - request: CreateCellRequest, - ) -> CodeModeSessionResultFuture<'a, CellId> { - Box::pin(CodeModeService::create_cell(self, request)) - } - - fn observe<'a>( - &'a self, - request: ObserveRequest, - ) -> CodeModeSessionResultFuture<'a, CellOutcome> { - Box::pin(CodeModeService::observe(self, request)) - } - - fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, CellOutcome> { - Box::pin(CodeModeService::terminate(self, cell_id)) - } - - fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { - Box::pin(CodeModeService::shutdown(self)) - } -} - -struct ProtocolDelegate { - delegate: Arc, -} - -impl runtime::SessionRuntimeDelegate for ProtocolDelegate { - async fn invoke_tool( - &self, - invocation: runtime::NestedToolCall, - cancellation_token: CancellationToken, - ) -> Result { - self.delegate - .invoke_tool( - CodeModeNestedToolCall { - cell_id: protocol_cell_id(&invocation.cell_id), - runtime_tool_call_id: invocation.runtime_tool_call_id, - tool_name: codex_protocol::ToolName { - name: invocation.tool_name.name, - namespace: invocation.tool_name.namespace, - }, - tool_kind: match invocation.tool_kind { - runtime::ToolKind::Function => CodeModeToolKind::Function, - runtime::ToolKind::Freeform => CodeModeToolKind::Freeform, - }, - input: invocation.input, - }, - cancellation_token, - ) - .await - } - - async fn notify( - &self, - call_id: String, - cell_id: runtime::CellId, - text: String, - cancellation_token: CancellationToken, - ) -> Result<(), String> { - self.delegate - .notify( - call_id, - protocol_cell_id(&cell_id), - text, - cancellation_token, - ) - .await - } - - fn cell_closed(&self, cell_id: &runtime::CellId) { - self.delegate.cell_closed(&protocol_cell_id(cell_id)); - } -} - -fn runtime_request(request: CreateCellRequest) -> runtime::CreateCellRequest { - runtime::CreateCellRequest { - tool_call_id: request.tool_call_id, - enabled_tools: request - .enabled_tools - .into_iter() - .map(|definition| runtime::ToolDefinition { - name: definition.name, - tool_name: runtime::ToolName { - name: definition.tool_name.name, - namespace: definition.tool_name.namespace, - }, - description: definition.description, - kind: match definition.kind { - CodeModeToolKind::Function => runtime::ToolKind::Function, - CodeModeToolKind::Freeform => runtime::ToolKind::Freeform, - }, - }) - .collect(), - source: request.source, - } -} - -fn runtime_cell_id(cell_id: &CellId) -> runtime::CellId { - runtime::CellId::new(cell_id.as_str()) -} - -fn protocol_cell_id(cell_id: &runtime::CellId) -> CellId { - CellId::new(cell_id.as_str().to_string()) -} - -fn pending_outcome(cell_id: &CellId, event: runtime::CellEvent) -> Result { - match event { - runtime::CellEvent::Pending { - content_items, - pending_tool_call_ids, - } => Ok(PendingOutcome::Pending { - cell_id: cell_id.clone(), - content_items: content_items.into_iter().map(output_item).collect(), - pending_tool_call_ids, - }), - event => Ok(PendingOutcome::Completed(runtime_response(cell_id, event)?)), - } -} - -fn runtime_response( - cell_id: &CellId, - event: runtime::CellEvent, -) -> Result { - match event { - runtime::CellEvent::Yielded { content_items } => Ok(RuntimeResponse::Yielded { - cell_id: cell_id.clone(), - content_items: content_items.into_iter().map(output_item).collect(), - }), - runtime::CellEvent::Completed { - content_items, - error_text, - } => Ok(RuntimeResponse::Result { - cell_id: cell_id.clone(), - content_items: content_items.into_iter().map(output_item).collect(), - error_text, - }), - runtime::CellEvent::Terminated { content_items } => Ok(RuntimeResponse::Terminated { - cell_id: cell_id.clone(), - content_items: content_items.into_iter().map(output_item).collect(), - }), - runtime::CellEvent::Pending { .. } => { - Err("cell returned a pending frontier unexpectedly".to_string()) - } - } -} - -fn output_item(item: runtime::OutputItem) -> FunctionCallOutputContentItem { - match item { - runtime::OutputItem::Text { text } => FunctionCallOutputContentItem::InputText { text }, - runtime::OutputItem::Image { image_url, detail } => { - FunctionCallOutputContentItem::InputImage { - image_url, - detail: detail.map(|detail| match detail { - runtime::ImageDetail::Auto => ImageDetail::Auto, - runtime::ImageDetail::Low => ImageDetail::Low, - runtime::ImageDetail::High => ImageDetail::High, - runtime::ImageDetail::Original => ImageDetail::Original, - }), - } - } - } -} - -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 missing_observation(cell_id: CellId) -> CodeModeSessionResultFuture<'static, CellOutcome> { - Box::pin(async move { Ok(CellOutcome::MissingCell(missing_cell_response(cell_id))) }) -} - -#[cfg(test)] -#[path = "service_tests.rs"] -mod tests; - -#[cfg(test)] -#[path = "service_contract_tests.rs"] -mod contract_tests; diff --git a/codex-rs/code-mode/src/service_contract_tests.rs b/codex-rs/code-mode/src/service_contract_tests.rs deleted file mode 100644 index 3e8e665633..0000000000 --- a/codex-rs/code-mode/src/service_contract_tests.rs +++ /dev/null @@ -1,830 +0,0 @@ -use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; -use std::time::Duration; - -use codex_protocol::ToolName; -use pretty_assertions::assert_eq; -use tokio::sync::Notify; -use tokio::sync::Semaphore; -use tokio::sync::mpsc; -use tokio_util::sync::CancellationToken; - -use super::*; -use crate::CodeModeToolKind; -use crate::ToolDefinition; - -#[derive(Debug, PartialEq)] -enum DelegateEvent { - NotificationStarted, - NotificationFinished, - ToolStarted, - CellClosed(CellId), -} - -fn record_cell_closed(events_tx: &mpsc::UnboundedSender, cell_id: &CellId) { - let _ = events_tx.send(DelegateEvent::CellClosed(cell_id.clone())); -} - -struct BlockingDelegate { - events_tx: mpsc::UnboundedSender, - tool_future_dropped: AtomicBool, - tool_release: Notify, -} - -struct DropFlag<'a>(&'a AtomicBool); - -impl Drop for DropFlag<'_> { - fn drop(&mut self) { - self.0.store(true, Ordering::Release); - } -} - -struct NeverResolvingNotificationDelegate { - events_tx: mpsc::UnboundedSender, -} - -struct ReleasableNotificationDelegate { - events_tx: mpsc::UnboundedSender, - notification_release: Semaphore, -} - -struct NeverResolvingToolDelegate { - events_tx: mpsc::UnboundedSender, -} - -impl NeverResolvingNotificationDelegate { - fn new() -> (Arc, mpsc::UnboundedReceiver) { - let (events_tx, events_rx) = mpsc::unbounded_channel(); - (Arc::new(Self { events_tx }), events_rx) - } -} - -impl ReleasableNotificationDelegate { - fn new() -> (Arc, mpsc::UnboundedReceiver) { - let (events_tx, events_rx) = mpsc::unbounded_channel(); - ( - Arc::new(Self { - events_tx, - notification_release: Semaphore::new(/*permits*/ 0), - }), - events_rx, - ) - } - - fn release_notification(&self) { - self.notification_release.add_permits(/*permits*/ 1); - } -} - -impl NeverResolvingToolDelegate { - fn new() -> (Arc, mpsc::UnboundedReceiver) { - let (events_tx, events_rx) = mpsc::unbounded_channel(); - (Arc::new(Self { events_tx }), events_rx) - } -} - -impl CodeModeSessionDelegate for NeverResolvingNotificationDelegate { - fn invoke_tool<'a>( - &'a self, - _invocation: CodeModeNestedToolCall, - _cancellation_token: CancellationToken, - ) -> ToolInvocationFuture<'a> { - Box::pin(async { Err("unexpected tool call".to_string()) }) - } - - fn notify<'a>( - &'a self, - _call_id: String, - _cell_id: CellId, - _text: String, - _cancellation_token: CancellationToken, - ) -> NotificationFuture<'a> { - Box::pin(async move { - let _ = self.events_tx.send(DelegateEvent::NotificationStarted); - std::future::pending().await - }) - } - - fn cell_closed(&self, cell_id: &CellId) { - record_cell_closed(&self.events_tx, cell_id); - } -} - -impl CodeModeSessionDelegate for ReleasableNotificationDelegate { - fn invoke_tool<'a>( - &'a self, - _invocation: CodeModeNestedToolCall, - cancellation_token: CancellationToken, - ) -> ToolInvocationFuture<'a> { - Box::pin(async move { - cancellation_token.cancelled().await; - Err("cancelled".to_string()) - }) - } - - fn notify<'a>( - &'a self, - _call_id: String, - _cell_id: CellId, - _text: String, - cancellation_token: CancellationToken, - ) -> NotificationFuture<'a> { - Box::pin(async move { - let _ = self.events_tx.send(DelegateEvent::NotificationStarted); - tokio::select! { - _ = self.notification_release.acquire() => { - let _ = self.events_tx.send(DelegateEvent::NotificationFinished); - Ok(()) - } - _ = cancellation_token.cancelled() => { - Err("cancelled".to_string()) - } - } - }) - } - - fn cell_closed(&self, cell_id: &CellId) { - record_cell_closed(&self.events_tx, cell_id); - } -} - -impl CodeModeSessionDelegate for NeverResolvingToolDelegate { - fn invoke_tool<'a>( - &'a self, - _invocation: CodeModeNestedToolCall, - _cancellation_token: CancellationToken, - ) -> ToolInvocationFuture<'a> { - Box::pin(async move { - let _ = self.events_tx.send(DelegateEvent::ToolStarted); - std::future::pending().await - }) - } - - fn notify<'a>( - &'a self, - _call_id: String, - _cell_id: CellId, - _text: String, - _cancellation_token: CancellationToken, - ) -> NotificationFuture<'a> { - Box::pin(async { Ok(()) }) - } - - fn cell_closed(&self, cell_id: &CellId) { - record_cell_closed(&self.events_tx, cell_id); - } -} - -impl BlockingDelegate { - fn new() -> (Arc, mpsc::UnboundedReceiver) { - let (events_tx, events_rx) = mpsc::unbounded_channel(); - ( - Arc::new(Self { - events_tx, - tool_future_dropped: AtomicBool::new(false), - tool_release: Notify::new(), - }), - events_rx, - ) - } - - fn release_tool(&self) { - self.tool_release.notify_one(); - } -} - -impl CodeModeSessionDelegate for BlockingDelegate { - fn invoke_tool<'a>( - &'a self, - _invocation: CodeModeNestedToolCall, - cancellation_token: CancellationToken, - ) -> ToolInvocationFuture<'a> { - let drop_flag = DropFlag(&self.tool_future_dropped); - Box::pin(async move { - let _drop_flag = drop_flag; - let _ = self.events_tx.send(DelegateEvent::ToolStarted); - tokio::select! { - _ = self.tool_release.notified() => { - Ok(serde_json::Value::Null) - } - _ = cancellation_token.cancelled() => { - Err("cancelled".to_string()) - } - } - }) - } - - fn notify<'a>( - &'a self, - _call_id: String, - _cell_id: CellId, - _text: String, - cancellation_token: CancellationToken, - ) -> NotificationFuture<'a> { - Box::pin(async move { - let _ = self.events_tx.send(DelegateEvent::NotificationStarted); - cancellation_token.cancelled().await; - Err("cancelled".to_string()) - }) - } - - fn cell_closed(&self, cell_id: &CellId) { - record_cell_closed(&self.events_tx, cell_id); - } -} - -fn cell_id(value: &str) -> CellId { - CellId::new(value.to_string()) -} - -fn execute_request(source: &str) -> CreateCellRequest { - CreateCellRequest { - tool_call_id: "call-1".to_string(), - enabled_tools: Vec::new(), - source: source.to_string(), - } -} - -async fn execute(service: &CodeModeService, request: CreateCellRequest) -> RuntimeResponse { - execute_with_yield_time(service, request, /*yield_time_ms*/ 1).await -} - -async fn execute_with_yield_time( - service: &CodeModeService, - request: CreateCellRequest, - yield_time_ms: u64, -) -> RuntimeResponse { - let cell_id = service.create_cell(request).await.unwrap(); - service - .observe(ObserveRequest { - cell_id, - yield_time_ms, - }) - .await - .unwrap() - .into() -} - -async fn create_and_observe_to_pending( - service: &CodeModeService, - request: CreateCellRequest, -) -> Result { - let cell_id = service.create_cell(request).await?; - match service - .observe_to_pending(ObserveToPendingRequest { cell_id }) - .await? - { - ObserveToPendingOutcome::LiveCell(outcome) => Ok(outcome), - ObserveToPendingOutcome::MissingCell(response) => Ok(PendingOutcome::Completed(response)), - } -} - -fn blocking_tool() -> ToolDefinition { - ToolDefinition { - name: "block".to_string(), - tool_name: ToolName::plain("block"), - description: String::new(), - kind: CodeModeToolKind::Function, - input_schema: None, - output_schema: None, - } -} - -async fn next_event(events_rx: &mut mpsc::UnboundedReceiver) -> DelegateEvent { - tokio::time::timeout(Duration::from_secs(2), events_rx.recv()) - .await - .expect("delegate event timeout") - .expect("delegate event channel closed") -} - -#[tokio::test] -async fn yields_and_resumes() { - let service = CodeModeService::new(); - let cell = execute_with_yield_time( - &service, - CreateCellRequest { - source: r#"text("before"); yield_control(); text("after");"#.to_string(), - ..execute_request("") - }, - /*yield_time_ms*/ 60_000, - ) - .await; - - assert_eq!( - cell, - RuntimeResponse::Yielded { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "before".to_string(), - }], - } - ); - assert_eq!( - service - .observe(ObserveRequest { - cell_id: cell_id("1"), - yield_time_ms: 60_000, - }) - .await - .unwrap(), - CellOutcome::LiveCell(RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "after".to_string(), - }], - error_text: None, - }) - ); -} - -#[tokio::test] -async fn yield_before_first_observation_preserves_its_output_boundary() { - let (delegate, mut events_rx) = BlockingDelegate::new(); - let service = CodeModeService::with_delegate(delegate.clone()); - let cell_id = service - .create_cell(CreateCellRequest { - enabled_tools: vec![blocking_tool()], - source: r#" -text("before"); -yield_control(); -text("after"); -await tools.block({}); -"# - .to_string(), - ..execute_request("") - }) - .await - .unwrap(); - assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted); - - assert_eq!( - service - .observe(ObserveRequest { - cell_id: cell_id.clone(), - yield_time_ms: 60_000, - }) - .await - .unwrap(), - CellOutcome::LiveCell(RuntimeResponse::Yielded { - cell_id: cell_id.clone(), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "before".to_string(), - }], - }) - ); - - delegate.release_tool(); - assert_eq!( - service - .observe(ObserveRequest { - cell_id: cell_id.clone(), - yield_time_ms: 60_000, - }) - .await - .unwrap(), - CellOutcome::LiveCell(RuntimeResponse::Result { - cell_id, - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "after".to_string(), - }], - error_text: None, - }) - ); -} - -#[tokio::test] -async fn background_completion_notifies_the_delegate_without_another_observation() { - let (delegate, mut events_rx) = BlockingDelegate::new(); - let service = CodeModeService::with_delegate(delegate); - let created_cell_id = service - .create_cell(execute_request( - r#"await new Promise(resolve => setTimeout(resolve, 100)); text("done");"#, - )) - .await - .unwrap(); - assert_eq!( - tokio::time::timeout( - Duration::from_secs(2), - service.observe(ObserveRequest { - cell_id: created_cell_id.clone(), - yield_time_ms: 1, - }), - ) - .await - .expect("initial observation should yield while the cell is still running") - .unwrap(), - CellOutcome::LiveCell(RuntimeResponse::Yielded { - cell_id: created_cell_id.clone(), - content_items: Vec::new(), - }) - ); - - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::CellClosed(created_cell_id) - ); -} - -#[tokio::test] -async fn returns_and_resumes_from_the_pending_frontier() { - let (delegate, mut events_rx) = BlockingDelegate::new(); - let service = CodeModeService::with_delegate(delegate.clone()); - - assert_eq!( - create_and_observe_to_pending( - &service, - CreateCellRequest { - enabled_tools: vec![blocking_tool()], - source: r#" -await tools.block({}); -text("after"); -"# - .to_string(), - ..execute_request("") - }, - ) - .await - .unwrap(), - PendingOutcome::Pending { - cell_id: cell_id("1"), - content_items: Vec::new(), - pending_tool_call_ids: vec!["tool-1".to_string()], - } - ); - - assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted); - delegate.release_tool(); - - assert_eq!( - service - .observe_to_pending(ObserveToPendingRequest { - cell_id: cell_id("1"), - }) - .await - .unwrap(), - ObserveToPendingOutcome::LiveCell(PendingOutcome::Completed(RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "after".to_string(), - }], - error_text: None, - })) - ); -} - -#[tokio::test] -async fn observed_natural_completion_wins_over_termination() { - let service = CodeModeService::new(); - let cell = execute( - &service, - execute_request(r#"yield_control(); store("finished", true); text("done");"#), - ) - .await; - - assert_eq!( - cell, - RuntimeResponse::Yielded { - cell_id: cell_id("1"), - content_items: Vec::new(), - } - ); - tokio::time::timeout(Duration::from_secs(10), async { - loop { - // Observe the synchronous probe through completion so scheduler load - // cannot turn its short yield deadline into a false lifecycle failure. - let response = create_and_observe_to_pending( - &service, - CreateCellRequest { - ..execute_request(r#"text(String(load("finished")));"#) - }, - ) - .await - .unwrap(); - let PendingOutcome::Completed(RuntimeResponse::Result { content_items, .. }) = response - else { - panic!("expected stored-value probe to complete"); - }; - if content_items - == vec![FunctionCallOutputContentItem::InputText { - text: "true".to_string(), - }] - { - break; - } - tokio::task::yield_now().await; - } - }) - .await - .unwrap(); - assert_eq!( - service.terminate(cell_id("1")).await.unwrap(), - CellOutcome::LiveCell(RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "done".to_string(), - }], - error_text: None, - }) - ); -} - -#[tokio::test] -async fn termination_discards_pending_callbacks_before_responding() { - let (delegate, mut events_rx) = BlockingDelegate::new(); - let service = CodeModeService::with_delegate(delegate.clone()); - let cell = execute( - &service, - execute_request(r#"notify("pending"); await new Promise(() => {});"#), - ) - .await; - - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::NotificationStarted - ); - assert_eq!( - cell, - RuntimeResponse::Yielded { - cell_id: cell_id("1"), - content_items: Vec::new(), - } - ); - assert_eq!( - service.terminate(cell_id("1")).await.unwrap(), - CellOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: cell_id("1"), - content_items: Vec::new(), - }) - ); -} - -#[tokio::test] -async fn termination_discards_stored_writes_before_the_next_cell_can_load_them() { - let (delegate, mut events_rx) = BlockingDelegate::new(); - let service = CodeModeService::with_delegate(delegate.clone()); - let created_cell_id = service - .create_cell(CreateCellRequest { - enabled_tools: vec![blocking_tool()], - source: r#" -store("candidate", "leaked"); -await tools.block({}); -"# - .to_string(), - ..execute_request("") - }) - .await - .unwrap(); - - // Reaching the delegate proves that the store ran before execution became gated. - assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted); - assert_eq!( - service.terminate(created_cell_id).await.unwrap(), - CellOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: cell_id("1"), - content_items: Vec::new(), - }) - ); - delegate.release_tool(); - - assert_eq!( - execute_with_yield_time( - &service, - execute_request(r#"text(String(load("candidate")));"#), - /*yield_time_ms*/ 60_000, - ) - .await, - RuntimeResponse::Result { - cell_id: cell_id("2"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "undefined".to_string(), - }], - error_text: None, - } - ); -} - -#[tokio::test] -async fn shutdown_does_not_await_notifications_during_natural_completion() { - let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new(); - let service = Arc::new(CodeModeService::with_delegate(delegate)); - service - .create_cell(execute_request(r#"notify("pending");"#)) - .await - .unwrap(); - - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::NotificationStarted - ); - - let shutdown_service = Arc::clone(&service); - tokio::time::timeout(Duration::from_millis(/*millis*/ 100), async move { - shutdown_service.shutdown().await - }) - .await - .expect("shutdown should not await a non-cooperative notification") - .unwrap(); -} - -#[tokio::test] -async fn shutdown_does_not_await_a_non_cooperative_nested_tool() { - let (delegate, mut events_rx) = NeverResolvingToolDelegate::new(); - let service = CodeModeService::with_delegate(delegate); - let _cell_id = service - .create_cell(CreateCellRequest { - enabled_tools: vec![blocking_tool()], - source: r#"await tools.block({});"#.to_string(), - ..execute_request("") - }) - .await - .unwrap(); - - assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted); - tokio::time::timeout(Duration::from_millis(/*millis*/ 100), service.shutdown()) - .await - .expect("shutdown should not await a non-cooperative nested tool") - .unwrap(); -} - -#[tokio::test] -async fn termination_does_not_await_a_non_cooperative_notification() { - let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new(); - let service = CodeModeService::with_delegate(delegate); - let cell = execute( - &service, - execute_request(r#"notify("pending"); await new Promise(() => {});"#), - ) - .await; - - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::NotificationStarted - ); - assert_eq!( - cell, - RuntimeResponse::Yielded { - cell_id: cell_id("1"), - content_items: Vec::new(), - } - ); - - assert_eq!( - tokio::time::timeout( - Duration::from_millis(/*millis*/ 100), - service.terminate(cell_id("1")), - ) - .await - .expect("termination should not await a non-cooperative notification") - .unwrap(), - CellOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: cell_id("1"), - content_items: Vec::new(), - }) - ); -} - -#[tokio::test] -async fn create_cell_returns_before_natural_completion() { - let (delegate, mut events_rx) = ReleasableNotificationDelegate::new(); - let service = CodeModeService::with_delegate(delegate.clone()); - let created_cell_id = service - .create_cell(execute_request(r#"notify("pending");"#)) - .await - .unwrap(); - assert_eq!(created_cell_id, cell_id("1")); - let mut observation = Box::pin(service.observe(ObserveRequest { - cell_id: created_cell_id, - yield_time_ms: 60_000, - })); - - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::NotificationStarted - ); - std::future::poll_fn(|context| match observation.as_mut().poll(context) { - std::task::Poll::Pending => std::task::Poll::Ready(()), - std::task::Poll::Ready(result) => { - panic!("observation returned while the notification was blocked: {result:?}") - } - }) - .await; - - delegate.release_notification(); - - assert_eq!( - observation.await.unwrap(), - CellOutcome::LiveCell(RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: Vec::new(), - error_text: None, - }) - ); - assert_eq!( - next_event(&mut events_rx).await, - DelegateEvent::NotificationFinished - ); -} - -#[tokio::test] -async fn created_cell_can_be_terminated_before_observation() { - let service = CodeModeService::new(); - let created_cell_id = service - .create_cell(CreateCellRequest { - source: "await new Promise(() => {});".to_string(), - ..execute_request("") - }) - .await - .unwrap(); - - assert_eq!(created_cell_id, cell_id("1")); - assert_eq!( - service.terminate(created_cell_id.clone()).await.unwrap(), - CellOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: created_cell_id, - content_items: Vec::new(), - }) - ); -} - -#[tokio::test] -async fn second_observer_is_rejected_without_displacing_the_first() { - let service = CodeModeService::new(); - let cell = execute(&service, execute_request("await new Promise(() => {});")).await; - - assert_eq!( - cell, - RuntimeResponse::Yielded { - cell_id: cell_id("1"), - content_items: Vec::new(), - } - ); - - let first_observer = service - .begin_observe(ObserveRequest { - cell_id: cell_id("1"), - yield_time_ms: 60_000, - }) - .await; - assert_eq!( - service - .observe(ObserveRequest { - cell_id: cell_id("1"), - yield_time_ms: 60_000, - }) - .await - .unwrap_err(), - "exec cell 1 already has an active observer" - ); - - let terminated = RuntimeResponse::Terminated { - cell_id: cell_id("1"), - content_items: Vec::new(), - }; - assert_eq!( - service.terminate(cell_id("1")).await.unwrap(), - CellOutcome::LiveCell(terminated.clone()) - ); - assert_eq!( - first_observer.await.unwrap(), - CellOutcome::LiveCell(terminated) - ); -} - -#[tokio::test] -async fn natural_completion_cleans_up_callbacks_before_responding() { - let (delegate, mut events_rx) = BlockingDelegate::new(); - let service = CodeModeService::with_delegate(delegate.clone()); - let created_cell_id = service - .create_cell(CreateCellRequest { - enabled_tools: vec![blocking_tool()], - source: concat!( - "tools.block({});", - "await new Promise(resolve => setTimeout(resolve, 100));", - "text('done');", - ) - .to_string(), - ..execute_request("") - }) - .await - .unwrap(); - - assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted); - assert_eq!( - service - .observe(ObserveRequest { - cell_id: created_cell_id, - yield_time_ms: 60_000, - }) - .await - .unwrap(), - CellOutcome::LiveCell(RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "done".to_string(), - }], - error_text: None, - }) - ); - assert!(delegate.tool_future_dropped.load(Ordering::Acquire)); -} diff --git a/codex-rs/code-mode/src/service_tests.rs b/codex-rs/code-mode/src/service_tests.rs deleted file mode 100644 index 6cfe8b9107..0000000000 --- a/codex-rs/code-mode/src/service_tests.rs +++ /dev/null @@ -1,1408 +0,0 @@ -use std::collections::VecDeque; -use std::sync::Arc; -use std::sync::Mutex; -use std::time::Duration; - -use super::CellId; -use super::CellOutcome; -use super::CodeModeNestedToolCall; -use super::CodeModeService; -use super::CodeModeSessionDelegate; -use super::NotificationFuture; -use super::ObserveRequest; -use super::ObserveToPendingOutcome; -use super::ObserveToPendingRequest; -use super::ProtocolDelegate; -use super::RuntimeResponse; -use super::ToolInvocationFuture; -use super::missing_cell_response; -use super::pending_outcome; -use super::protocol_cell_id; -use super::runtime; -use super::runtime_cell_id; -use super::runtime_request; -use super::runtime_response; -use crate::CodeModeToolKind; -use crate::CreateCellRequest; -use crate::FunctionCallOutputContentItem; -use crate::PendingOutcome; -use crate::ToolDefinition; -use codex_protocol::ToolName; -use pretty_assertions::assert_eq; -use serde_json::Value as JsonValue; -use serde_json::json; -use tokio::sync::Notify; -use tokio_util::sync::CancellationToken; - -#[derive(Default)] -struct ReleasableToolDelegate { - tool_release: Notify, -} - -impl ReleasableToolDelegate { - fn release_tool(&self) { - self.tool_release.notify_one(); - } -} - -impl CodeModeSessionDelegate for ReleasableToolDelegate { - fn invoke_tool<'a>( - &'a self, - _invocation: CodeModeNestedToolCall, - cancellation_token: CancellationToken, - ) -> ToolInvocationFuture<'a> { - Box::pin(async move { - tokio::select! { - _ = self.tool_release.notified() => Ok(JsonValue::Null), - _ = cancellation_token.cancelled() => Err("cancelled".to_string()), - } - }) - } - - fn notify<'a>( - &'a self, - _call_id: String, - _cell_id: CellId, - _text: String, - _cancellation_token: CancellationToken, - ) -> NotificationFuture<'a> { - Box::pin(async { Ok(()) }) - } - - fn cell_closed(&self, _cell_id: &CellId) {} -} - -#[derive(Debug, PartialEq)] -enum RecordedDelegateCall { - Tool { - invocation: CodeModeNestedToolCall, - cancellation_requested: bool, - }, - Notification { - call_id: String, - cell_id: CellId, - text: String, - cancellation_requested: bool, - }, - CellClosed { - cell_id: CellId, - }, -} - -struct RecordingDelegate { - calls: Mutex>, - tool_results: Mutex>>, - notification_results: Mutex>>, -} - -impl RecordingDelegate { - fn new( - tool_results: Vec>, - notification_results: Vec>, - ) -> Self { - Self { - calls: Mutex::new(Vec::new()), - tool_results: Mutex::new(tool_results.into()), - notification_results: Mutex::new(notification_results.into()), - } - } - - fn take_calls(&self) -> Vec { - std::mem::take(&mut *self.calls.lock().unwrap()) - } -} - -impl CodeModeSessionDelegate for RecordingDelegate { - fn invoke_tool<'a>( - &'a self, - invocation: CodeModeNestedToolCall, - cancellation_token: CancellationToken, - ) -> ToolInvocationFuture<'a> { - Box::pin(async move { - self.calls.lock().unwrap().push(RecordedDelegateCall::Tool { - invocation, - cancellation_requested: cancellation_token.is_cancelled(), - }); - self.tool_results - .lock() - .unwrap() - .pop_front() - .expect("test must provide one result per tool call") - }) - } - - fn notify<'a>( - &'a self, - call_id: String, - cell_id: CellId, - text: String, - cancellation_token: CancellationToken, - ) -> NotificationFuture<'a> { - Box::pin(async move { - self.calls - .lock() - .unwrap() - .push(RecordedDelegateCall::Notification { - call_id, - cell_id, - text, - cancellation_requested: cancellation_token.is_cancelled(), - }); - self.notification_results - .lock() - .unwrap() - .pop_front() - .expect("test must provide one result per notification") - }) - } - - fn cell_closed(&self, cell_id: &CellId) { - self.calls - .lock() - .unwrap() - .push(RecordedDelegateCall::CellClosed { - cell_id: cell_id.clone(), - }); - } -} - -fn execute_request(source: &str) -> CreateCellRequest { - CreateCellRequest { - tool_call_id: "call_1".to_string(), - enabled_tools: Vec::new(), - source: source.to_string(), - } -} - -fn cell_id(value: &str) -> CellId { - CellId::new(value.to_string()) -} - -fn echo_tool() -> ToolDefinition { - ToolDefinition { - name: "echo".to_string(), - tool_name: ToolName::plain("echo"), - description: String::new(), - kind: CodeModeToolKind::Function, - input_schema: None, - output_schema: None, - } -} - -async fn execute(service: &CodeModeService, request: CreateCellRequest) -> RuntimeResponse { - execute_with_yield_time(service, request, /*yield_time_ms*/ 10_000).await -} - -async fn execute_with_yield_time( - service: &CodeModeService, - request: CreateCellRequest, - yield_time_ms: u64, -) -> RuntimeResponse { - let cell_id = service.create_cell(request).await.unwrap(); - service - .observe(ObserveRequest { - cell_id, - yield_time_ms, - }) - .await - .unwrap() - .into() -} - -async fn create_and_observe_to_pending( - service: &CodeModeService, - request: CreateCellRequest, -) -> Result { - let cell_id = service.create_cell(request).await?; - match service - .observe_to_pending(ObserveToPendingRequest { cell_id }) - .await? - { - ObserveToPendingOutcome::LiveCell(outcome) => Ok(outcome), - ObserveToPendingOutcome::MissingCell(response) => Ok(PendingOutcome::Completed(response)), - } -} - -#[tokio::test] -async fn synchronous_exit_returns_successfully() { - let service = CodeModeService::new(); - - let response = execute( - &service, - CreateCellRequest { - source: r#"text("before"); exit(); text("after");"#.to_string(), - ..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, - CreateCellRequest { - source: r#"store("key", "visible");"#.to_string(), - ..execute_request("") - }, - ) - .await; - - let same_session = execute( - &first_session, - CreateCellRequest { - source: r#"text(String(load("key")));"#.to_string(), - ..execute_request("") - }, - ) - .await; - let other_session = execute( - &second_session, - CreateCellRequest { - source: r#"text(String(load("key")));"#.to_string(), - ..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 response = execute_with_yield_time( - &service, - CreateCellRequest { - source: "while (true) {}".to_string(), - ..execute_request("") - }, - /*yield_time_ms*/ 1, - ) - .await; - assert_eq!( - response, - 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.shutdown().await.unwrap(); - - let error = service - .create_cell(execute_request("text('late');")) - .await - .err() - .unwrap(); - - assert_eq!(error, "code mode session is shutting down".to_string()); -} - -#[tokio::test] -async fn create_and_observe_to_pending_returns_completed_for_synchronous_results() { - let service = CodeModeService::new(); - - let response = create_and_observe_to_pending( - &service, - CreateCellRequest { - source: r#"text("done");"#.to_string(), - ..execute_request("") - }, - ) - .await - .unwrap(); - - assert_eq!( - response, - PendingOutcome::Completed(RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "done".to_string(), - }], - error_text: None, - }) - ); -} - -#[tokio::test] -async fn create_and_observe_to_pending_returns_once_the_runtime_is_quiescent() { - let service = CodeModeService::new(); - - let response = tokio::time::timeout( - Duration::from_secs(1), - create_and_observe_to_pending( - &service, - CreateCellRequest { - source: r#"text("before"); await new Promise(() => {});"#.to_string(), - ..execute_request("") - }, - ), - ) - .await - .unwrap() - .unwrap(); - - assert_eq!( - response, - PendingOutcome::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, - CellOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: cell_id("1"), - content_items: Vec::new(), - }) - ); -} - -#[tokio::test] -async fn create_and_observe_to_pending_identifies_tool_calls_in_paused_frontier() { - let service = CodeModeService::new(); - - let response = create_and_observe_to_pending( - &service, - CreateCellRequest { - enabled_tools: vec![echo_tool()], - source: r#" -await Promise.all([ - tools.echo({ value: "first" }), - tools.echo({ value: "second" }), -]); -"# - .to_string(), - ..execute_request("") - }, - ) - .await - .unwrap(); - - assert_eq!( - response, - PendingOutcome::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, - CellOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: cell_id("1"), - content_items: Vec::new(), - }) - ); -} - -#[tokio::test] -async fn create_and_observe_to_pending_excludes_delayed_timeout_tool_calls_until_observe() { - let service = CodeModeService::new(); - - let initial_response = create_and_observe_to_pending( - &service, - CreateCellRequest { - enabled_tools: vec![echo_tool()], - source: r#" -setTimeout(() => { - tools.echo({ value: "delayed" }); -}, 1000); -await Promise.all([ - tools.echo({ value: "second" }), - tools.echo({ value: "third" }), -]); -"# - .to_string(), - ..execute_request("") - }, - ) - .await - .unwrap(); - - assert_eq!( - initial_response, - PendingOutcome::Pending { - cell_id: cell_id("1"), - content_items: Vec::new(), - pending_tool_call_ids: vec!["tool-1".to_string(), "tool-2".to_string()], - } - ); - - tokio::time::sleep(Duration::from_millis(1100)).await; - - // Resuming can expose an empty quiescent frontier before the expired timer - // callback is dequeued, so keep observing until its tool call is visible. - let resumed_response = tokio::time::timeout(Duration::from_secs(5), async { - loop { - let response = service - .observe_to_pending(ObserveToPendingRequest { - cell_id: cell_id("1"), - }) - .await?; - if !matches!( - &response, - ObserveToPendingOutcome::LiveCell(PendingOutcome::Pending { - pending_tool_call_ids, - .. - }) if pending_tool_call_ids.is_empty() - ) { - break Ok::<_, String>(response); - } - } - }) - .await - .unwrap() - .unwrap(); - - assert_eq!( - resumed_response, - ObserveToPendingOutcome::LiveCell(PendingOutcome::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, - CellOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: cell_id("1"), - content_items: Vec::new(), - }) - ); -} - -#[tokio::test] -async fn observe_to_pending_returns_after_resumed_runtime_becomes_quiescent_again() { - let delegate = Arc::new(ReleasableToolDelegate::default()); - let service = CodeModeService::with_delegate(delegate.clone()); - - let initial_response = create_and_observe_to_pending( - &service, - CreateCellRequest { - enabled_tools: vec![echo_tool()], - source: r#" -await tools.echo({}); -text("after"); -await new Promise(() => {}); -"# - .to_string(), - ..execute_request("") - }, - ) - .await - .unwrap(); - - assert_eq!( - initial_response, - PendingOutcome::Pending { - cell_id: cell_id("1"), - content_items: Vec::new(), - pending_tool_call_ids: vec!["tool-1".to_string()], - } - ); - - delegate.release_tool(); - - let resumed_response = tokio::time::timeout( - Duration::from_secs(1), - service.observe_to_pending(ObserveToPendingRequest { - cell_id: cell_id("1"), - }), - ) - .await - .unwrap() - .unwrap(); - - assert_eq!( - resumed_response, - ObserveToPendingOutcome::LiveCell(PendingOutcome::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, - CellOutcome::LiveCell(RuntimeResponse::Terminated { - cell_id: cell_id("1"), - content_items: Vec::new(), - }) - ); -} - -#[tokio::test] -async fn observe_to_pending_returns_completed_after_resumed_runtime_finishes() { - let delegate = Arc::new(ReleasableToolDelegate::default()); - let service = CodeModeService::with_delegate(delegate.clone()); - - let initial_response = create_and_observe_to_pending( - &service, - CreateCellRequest { - enabled_tools: vec![echo_tool()], - source: r#" -await tools.echo({}); -text("done"); -"# - .to_string(), - ..execute_request("") - }, - ) - .await - .unwrap(); - - assert_eq!( - initial_response, - PendingOutcome::Pending { - cell_id: cell_id("1"), - content_items: Vec::new(), - pending_tool_call_ids: vec!["tool-1".to_string()], - } - ); - - delegate.release_tool(); - - let resumed_response = tokio::time::timeout( - Duration::from_secs(1), - service.observe_to_pending(ObserveToPendingRequest { - cell_id: cell_id("1"), - }), - ) - .await - .unwrap() - .unwrap(); - - assert_eq!( - resumed_response, - ObserveToPendingOutcome::LiveCell(PendingOutcome::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, - CreateCellRequest { - source: r#"text(String(Object.hasOwn(globalThis, "console")));"#.to_string(), - ..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, - CreateCellRequest { - 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(), - ..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, - CreateCellRequest { - 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(), - ..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, - CreateCellRequest { - source: r#" -const returnsUndefined = [ - text("first"), - image("data:image/png;base64,AAA"), - notify("ping"), -].map((value) => value === undefined); -text(JSON.stringify(returnsUndefined)); -"# - .to_string(), - ..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: "data:image/png;base64,AAA".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, - CreateCellRequest { - source: r#" -image({ - type: "image", - data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==", - mimeType: "image/png", - _meta: { "codex/imageDetail": "original" }, -}); -"# - .to_string(), - ..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 generated_image_helper_appends_image_and_output_hint() { - let service = CodeModeService::new(); - - let response = execute( - &service, - CreateCellRequest { - source: r#" -generatedImage({ - image_url: "data:image/png;base64,AAA", - output_hint: "generated image save hint", -}); -"# - .to_string(), - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![ - FunctionCallOutputContentItem::InputImage { - image_url: "data:image/png;base64,AAA".to_string(), - detail: Some(crate::DEFAULT_IMAGE_DETAIL), - }, - FunctionCallOutputContentItem::InputText { - text: "generated image save hint".to_string(), - }, - ], - error_text: None, - } - ); -} - -#[tokio::test] -async fn image_helper_second_arg_overrides_explicit_object_detail() { - let service = CodeModeService::new(); - - let response = execute( - &service, - CreateCellRequest { - source: r#" -image( - { - image_url: "data:image/png;base64,AAA", - detail: "high", - }, - "original", -); -"# - .to_string(), - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputImage { - image_url: "data:image/png;base64,AAA".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, - CreateCellRequest { - source: r#" -image( - { - type: "image", - data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==", - mimeType: "image/png", - _meta: { "codex/imageDetail": "original" }, - }, - "high", -); -"# - .to_string(), - ..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, - CreateCellRequest { - source: r#" -image({ - image_url: "data:image/png;base64,AAA", - detail: "low", -}); -"# - .to_string(), - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: vec![FunctionCallOutputContentItem::InputImage { - image_url: "data:image/png;base64,AAA".to_string(), - detail: Some(crate::ImageDetail::Low), - }], - error_text: None, - } - ); -} - -#[tokio::test] -async fn image_helpers_reject_remote_urls() { - for image_url in [ - "http://example.com/image.jpg", - "https://example.com/image.jpg", - ] { - for source in [ - format!("image({image_url:?});"), - format!("generatedImage({{ image_url: {image_url:?} }});"), - ] { - let service = CodeModeService::new(); - - let response = execute( - &service, - CreateCellRequest { - source, - ..execute_request("") - }, - ) - .await; - - assert_eq!( - response, - RuntimeResponse::Result { - cell_id: cell_id("1"), - content_items: Vec::new(), - error_text: Some( - "Tool call failed: remote image URLs are not supported in tool outputs. Pass a base64 data URI instead".to_string(), - ), - } - ); - } - } -} - -#[tokio::test] -async fn image_helper_rejects_unsupported_detail() { - let service = CodeModeService::new(); - - let response = execute( - &service, - CreateCellRequest { - source: r#" -image({ - image_url: "data:image/png;base64,AAA", - detail: "medium", -}); -"# - .to_string(), - ..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, - CreateCellRequest { - source: r#" -image({ - content: [ - { - type: "image", - data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==", - mimeType: "image/png", - _meta: { "codex/imageDetail": "original" }, - }, - ], - isError: false, -}); -"# - .to_string(), - ..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 observe_reports_missing_cell_separately_from_runtime_results() { - let service = CodeModeService::new(); - - let response = service - .observe(ObserveRequest { - cell_id: cell_id("missing"), - yield_time_ms: 1, - }) - .await - .unwrap(); - - assert_eq!( - response, - CellOutcome::MissingCell(RuntimeResponse::Result { - cell_id: cell_id("missing"), - content_items: Vec::new(), - error_text: Some("exec cell missing not found".to_string()), - }) - ); -} - -#[test] -fn protocol_requests_map_to_runtime_requests_field_for_field() { - let request = CreateCellRequest { - tool_call_id: "response-call-7".to_string(), - enabled_tools: vec![ - ToolDefinition { - name: "mcp__search__query".to_string(), - tool_name: ToolName::namespaced("search", "query"), - description: "Search indexed documents".to_string(), - kind: CodeModeToolKind::Function, - input_schema: Some(json!({"type": "object"})), - output_schema: None, - }, - ToolDefinition { - name: "apply_patch".to_string(), - tool_name: ToolName::plain("apply_patch"), - description: "Apply a patch".to_string(), - kind: CodeModeToolKind::Freeform, - input_schema: None, - output_schema: Some(json!({"type": "string"})), - }, - ], - source: "await tools.mcp__search__query({ q: 'rust' });".to_string(), - }; - - assert_eq!( - runtime_request(request), - runtime::CreateCellRequest { - tool_call_id: "response-call-7".to_string(), - enabled_tools: vec![ - runtime::ToolDefinition { - name: "mcp__search__query".to_string(), - tool_name: runtime::ToolName { - name: "query".to_string(), - namespace: Some("search".to_string()), - }, - description: "Search indexed documents".to_string(), - kind: runtime::ToolKind::Function, - }, - runtime::ToolDefinition { - name: "apply_patch".to_string(), - tool_name: runtime::ToolName { - name: "apply_patch".to_string(), - namespace: None, - }, - description: "Apply a patch".to_string(), - kind: runtime::ToolKind::Freeform, - }, - ], - source: "await tools.mcp__search__query({ q: 'rust' });".to_string(), - } - ); - - let protocol_id = cell_id("cell-a7"); - assert_eq!( - protocol_cell_id(&runtime_cell_id(&protocol_id)), - protocol_id - ); -} - -#[tokio::test] -async fn protocol_delegate_maps_callbacks_cancellation_and_errors_field_for_field() { - let delegate = Arc::new(RecordingDelegate::new( - vec![ - Ok(json!({"matches": 3})), - Err("freeform tool failed".to_string()), - ], - vec![Err("notification failed".to_string())], - )); - let adapter = ProtocolDelegate { - delegate: delegate.clone(), - }; - let cancelled = CancellationToken::new(); - cancelled.cancel(); - - assert_eq!( - runtime::SessionRuntimeDelegate::invoke_tool( - &adapter, - runtime::NestedToolCall { - cell_id: runtime::CellId::new("cell-a7"), - runtime_tool_call_id: "runtime-call-1".to_string(), - tool_name: runtime::ToolName { - name: "query".to_string(), - namespace: Some("search".to_string()), - }, - tool_kind: runtime::ToolKind::Function, - input: Some(json!({"q": "rust"})), - }, - CancellationToken::new(), - ) - .await, - Ok(json!({"matches": 3})) - ); - assert_eq!( - runtime::SessionRuntimeDelegate::invoke_tool( - &adapter, - runtime::NestedToolCall { - cell_id: runtime::CellId::new("cell-b9"), - runtime_tool_call_id: "runtime-call-2".to_string(), - tool_name: runtime::ToolName { - name: "apply_patch".to_string(), - namespace: None, - }, - tool_kind: runtime::ToolKind::Freeform, - input: None, - }, - cancelled.clone(), - ) - .await, - Err("freeform tool failed".to_string()) - ); - assert_eq!( - runtime::SessionRuntimeDelegate::notify( - &adapter, - "notification-1".to_string(), - runtime::CellId::new("cell-c3"), - "progress".to_string(), - cancelled, - ) - .await, - Err("notification failed".to_string()) - ); - runtime::SessionRuntimeDelegate::cell_closed(&adapter, &runtime::CellId::new("cell-d4")); - - assert_eq!( - delegate.take_calls(), - vec![ - RecordedDelegateCall::Tool { - invocation: CodeModeNestedToolCall { - cell_id: cell_id("cell-a7"), - runtime_tool_call_id: "runtime-call-1".to_string(), - tool_name: ToolName::namespaced("search", "query"), - tool_kind: CodeModeToolKind::Function, - input: Some(json!({"q": "rust"})), - }, - cancellation_requested: false, - }, - RecordedDelegateCall::Tool { - invocation: CodeModeNestedToolCall { - cell_id: cell_id("cell-b9"), - runtime_tool_call_id: "runtime-call-2".to_string(), - tool_name: ToolName::plain("apply_patch"), - tool_kind: CodeModeToolKind::Freeform, - input: None, - }, - cancellation_requested: true, - }, - RecordedDelegateCall::Notification { - call_id: "notification-1".to_string(), - cell_id: cell_id("cell-c3"), - text: "progress".to_string(), - cancellation_requested: true, - }, - RecordedDelegateCall::CellClosed { - cell_id: cell_id("cell-d4"), - }, - ] - ); -} - -#[test] -fn runtime_events_map_to_protocol_outcomes_field_for_field() { - let output_items = vec![ - runtime::OutputItem::Text { - text: "before".to_string(), - }, - runtime::OutputItem::Image { - image_url: "data:image/png;base64,auto".to_string(), - detail: Some(runtime::ImageDetail::Auto), - }, - runtime::OutputItem::Image { - image_url: "data:image/png;base64,low".to_string(), - detail: Some(runtime::ImageDetail::Low), - }, - runtime::OutputItem::Image { - image_url: "data:image/png;base64,high".to_string(), - detail: Some(runtime::ImageDetail::High), - }, - runtime::OutputItem::Image { - image_url: "data:image/png;base64,original".to_string(), - detail: Some(runtime::ImageDetail::Original), - }, - runtime::OutputItem::Image { - image_url: "data:image/png;base64,default".to_string(), - detail: None, - }, - ]; - let expected_output_items = vec![ - FunctionCallOutputContentItem::InputText { - text: "before".to_string(), - }, - FunctionCallOutputContentItem::InputImage { - image_url: "data:image/png;base64,auto".to_string(), - detail: Some(crate::ImageDetail::Auto), - }, - FunctionCallOutputContentItem::InputImage { - image_url: "data:image/png;base64,low".to_string(), - detail: Some(crate::ImageDetail::Low), - }, - FunctionCallOutputContentItem::InputImage { - image_url: "data:image/png;base64,high".to_string(), - detail: Some(crate::ImageDetail::High), - }, - FunctionCallOutputContentItem::InputImage { - image_url: "data:image/png;base64,original".to_string(), - detail: Some(crate::ImageDetail::Original), - }, - FunctionCallOutputContentItem::InputImage { - image_url: "data:image/png;base64,default".to_string(), - detail: None, - }, - ]; - - assert_eq!( - runtime_response( - &cell_id("cell-a7"), - runtime::CellEvent::Yielded { - content_items: output_items, - }, - ), - Ok(RuntimeResponse::Yielded { - cell_id: cell_id("cell-a7"), - content_items: expected_output_items, - }) - ); - assert_eq!( - runtime_response( - &cell_id("cell-b9"), - runtime::CellEvent::Completed { - content_items: vec![runtime::OutputItem::Text { - text: "failed".to_string(), - }], - error_text: Some("tool failed".to_string()), - }, - ), - Ok(RuntimeResponse::Result { - cell_id: cell_id("cell-b9"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "failed".to_string(), - }], - error_text: Some("tool failed".to_string()), - }) - ); - assert_eq!( - runtime_response( - &cell_id("cell-c3"), - runtime::CellEvent::Terminated { - content_items: vec![runtime::OutputItem::Text { - text: "partial".to_string(), - }], - }, - ), - Ok(RuntimeResponse::Terminated { - cell_id: cell_id("cell-c3"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "partial".to_string(), - }], - }) - ); - assert_eq!( - pending_outcome( - &cell_id("cell-d4"), - runtime::CellEvent::Pending { - content_items: vec![runtime::OutputItem::Text { - text: "waiting".to_string(), - }], - pending_tool_call_ids: vec!["runtime-call-1".to_string()], - }, - ), - Ok(PendingOutcome::Pending { - cell_id: cell_id("cell-d4"), - content_items: vec![FunctionCallOutputContentItem::InputText { - text: "waiting".to_string(), - }], - pending_tool_call_ids: vec!["runtime-call-1".to_string()], - }) - ); - assert_eq!( - pending_outcome( - &cell_id("cell-e5"), - runtime::CellEvent::Completed { - content_items: Vec::new(), - error_text: None, - }, - ), - Ok(PendingOutcome::Completed(RuntimeResponse::Result { - cell_id: cell_id("cell-e5"), - content_items: Vec::new(), - error_text: None, - })) - ); - assert_eq!( - runtime_response( - &cell_id("cell-f6"), - runtime::CellEvent::Pending { - content_items: Vec::new(), - pending_tool_call_ids: Vec::new(), - }, - ), - Err("cell returned a pending frontier unexpectedly".to_string()) - ); - assert_eq!( - missing_cell_response(cell_id("missing")), - RuntimeResponse::Result { - cell_id: cell_id("missing"), - content_items: Vec::new(), - error_text: Some("exec cell missing not found".to_string()), - } - ); -} diff --git a/codex-rs/code-mode/src/session_runtime/mod.rs b/codex-rs/code-mode/src/session_runtime/mod.rs index 693569da4e..37aaf9880e 100644 --- a/codex-rs/code-mode/src/session_runtime/mod.rs +++ b/codex-rs/code-mode/src/session_runtime/mod.rs @@ -58,8 +58,6 @@ struct Inner { enum CellIdNamespace { Runtime(uuid::Uuid), - #[cfg(test)] - Unscoped, } impl SessionRuntime { @@ -67,11 +65,6 @@ impl SessionRuntime { Self::with_cell_id_namespace(delegate, CellIdNamespace::Runtime(uuid::Uuid::new_v4())) } - #[cfg(test)] - pub(crate) fn new_for_test(delegate: Arc) -> Self { - Self::with_cell_id_namespace(delegate, CellIdNamespace::Unscoped) - } - fn with_cell_id_namespace(delegate: Arc, cell_id_namespace: CellIdNamespace) -> Self { Self { inner: Arc::new(Inner { @@ -167,8 +160,6 @@ impl SessionRuntime { } encoded } - #[cfg(test)] - CellIdNamespace::Unscoped => sequence.to_string(), }; CellId::new(value) } diff --git a/codex-rs/core/BUILD.bazel b/codex-rs/core/BUILD.bazel index a55a98e186..a3f0583b80 100644 --- a/codex-rs/core/BUILD.bazel +++ b/codex-rs/core/BUILD.bazel @@ -14,6 +14,7 @@ codex_rust_crate( crate_name = "codex_core", extra_binaries = [ "//codex-rs/bwrap:bwrap", + "//codex-rs/code-mode-host:codex-code-mode-host", "//codex-rs/linux-sandbox:codex-linux-sandbox", "//codex-rs/rmcp-client:test_stdio_server", "//codex-rs/rmcp-client:test_streamable_http_server", diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index d16e8d9cbe..f90a691675 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -28,7 +28,8 @@ codex-api = { workspace = true } codex-app-server-protocol = { workspace = true } codex-apply-patch = { workspace = true } codex-async-utils = { workspace = true } -codex-code-mode = { workspace = true } +codex-code-mode-client = { workspace = true } +codex-code-mode-protocol = { workspace = true } codex-connectors = { workspace = true } codex-context-fragments = { workspace = true } codex-config = { workspace = true } diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 9f5b867f75..35a7126955 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -121,6 +121,7 @@ pub(crate) async fn run_codex_thread_interactive( thread_store: Arc::clone(&parent_session.services.thread_store), attestation_provider: parent_session.services.attestation_provider.clone(), external_time_provider: Some(Arc::clone(&parent_session.services.time_provider)), + code_mode_session_provider: Arc::clone(&parent_session.services.code_mode_session_provider), inherited_multi_agent_version: Some(MultiAgentVersion::Disabled), initial_multi_agent_mode: None, })) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 0792334b30..8d304db67a 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -444,6 +444,8 @@ pub(crate) struct CodexSpawnArgs { pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, pub(crate) external_time_provider: Option>, + pub(crate) code_mode_session_provider: + Arc, pub(crate) inherited_multi_agent_version: Option, pub(crate) initial_multi_agent_mode: Option, } @@ -529,6 +531,7 @@ impl Codex { thread_store, attestation_provider, external_time_provider, + code_mode_session_provider, inherited_multi_agent_version, initial_multi_agent_mode, } = args; @@ -680,6 +683,7 @@ impl Codex { parent_rollout_thread_trace, attestation_provider, external_time_provider, + code_mode_session_provider, multi_agent_version, )) .await diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 0c07f5e951..3c1dc460a2 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -501,6 +501,7 @@ impl Session { parent_rollout_thread_trace: ThreadTraceContext, attestation_provider: Option>, external_time_provider: Option>, + code_mode_session_provider: Arc, multi_agent_version: Option, ) -> anyhow::Result> { debug!( @@ -1055,7 +1056,10 @@ impl Session { session_configuration.parent_thread_id, ), ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::clone( + &code_mode_session_provider, + )), + code_mode_session_provider, tool_search_handler_cache: Default::default(), turn_environments: Arc::clone(&turn_environments), }; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index bca6be6a74..62fbe330d6 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4924,6 +4924,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, /*external_time_provider*/ None, + Arc::new(codex_code_mode_client::StdioCodeModeSessionProvider::default()), Some(config.multi_agent_version_from_features()), ) .await; @@ -5104,7 +5105,12 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { /*item_ids_enabled*/ config.features.enabled(Feature::ItemIds), /*attestation_provider*/ None, ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::new( + codex_code_mode_client::StdioCodeModeSessionProvider::default(), + )), + code_mode_session_provider: Arc::new( + codex_code_mode_client::StdioCodeModeSessionProvider::default(), + ), tool_search_handler_cache: Default::default(), turn_environments: Arc::clone(&turn_environments), }; @@ -5284,6 +5290,7 @@ async fn make_session_with_config_and_rx( codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, /*external_time_provider*/ None, + Arc::new(codex_code_mode_client::StdioCodeModeSessionProvider::default()), Some(config.multi_agent_version_from_features()), ) .await?; @@ -5398,6 +5405,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, /*external_time_provider*/ None, + Arc::new(codex_code_mode_client::StdioCodeModeSessionProvider::default()), Some(config.multi_agent_version_from_features()), ) .await?; @@ -7163,7 +7171,12 @@ where /*item_ids_enabled*/ config.features.enabled(Feature::ItemIds), /*attestation_provider*/ None, ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::new( + codex_code_mode_client::StdioCodeModeSessionProvider::default(), + )), + code_mode_session_provider: Arc::new( + codex_code_mode_client::StdioCodeModeSessionProvider::default(), + ), tool_search_handler_cache: Default::default(), turn_environments: Arc::clone(&turn_environments), }; diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 7011e6d0f5..62432e60b1 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -737,6 +737,9 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { thread_store, attestation_provider: None, external_time_provider: None, + code_mode_session_provider: Arc::new( + codex_code_mode_client::StdioCodeModeSessionProvider::default(), + ), inherited_multi_agent_version: None, initial_multi_agent_mode: None, }) diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 853b96e921..35d2211159 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -23,6 +23,7 @@ use anyhow::Result; use arc_swap::ArcSwap; use arc_swap::ArcSwapOption; use codex_analytics::AnalyticsEventsClient; +use codex_code_mode_protocol::CodeModeSessionProvider; use codex_core_plugins::PluginsManager; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionDataInit; @@ -84,6 +85,7 @@ pub(crate) struct SessionServices { /// Session-scoped model client shared across turns. pub(crate) model_client: ModelClient, pub(crate) code_mode_service: CodeModeService, + pub(crate) code_mode_session_provider: Arc, pub(crate) tool_search_handler_cache: ToolSearchHandlerCache, pub(crate) turn_environments: Arc, } diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 3a0e374750..414a17d569 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_client::StdioCodeModeSessionProvider; +use codex_code_mode_protocol::CodeModeSessionProvider; use codex_core_plugins::PluginsManager; use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionDataInit; @@ -223,6 +225,7 @@ pub(crate) struct ThreadManagerState { installation_id: String, analytics_events_client: Option, state_db: Option, + code_mode_session_provider: Arc, // Captures submitted ops for testing purpose when test mode is enabled. ops_log: Option, } @@ -311,6 +314,7 @@ impl ThreadManager { installation_id, analytics_events_client, state_db, + code_mode_session_provider: Arc::new(StdioCodeModeSessionProvider::default()), ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -417,6 +421,7 @@ impl ThreadManager { installation_id, analytics_events_client: None, state_db, + code_mode_session_provider: Arc::new(StdioCodeModeSessionProvider::default()), ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -428,6 +433,18 @@ impl ThreadManager { self.state.session_source.clone() } + #[doc(hidden)] + pub fn with_code_mode_session_provider( + mut self, + provider: Arc, + ) -> Self { + let Some(state) = Arc::get_mut(&mut self.state) else { + panic!("code-mode provider must be configured before sharing the thread manager"); + }; + state.code_mode_session_provider = provider; + self + } + pub fn auth_manager(&self) -> Arc { self.state.auth_manager.clone() } @@ -1499,6 +1516,7 @@ impl ThreadManagerState { thread_store: Arc::clone(&self.thread_store), attestation_provider: self.attestation_provider.clone(), external_time_provider: self.external_time_provider.clone(), + code_mode_session_provider: Arc::clone(&self.code_mode_session_provider), inherited_multi_agent_version: multi_agent_version, initial_multi_agent_mode, })) diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index 993c30eaa3..b491225a0c 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/codex-rs/core/src/tools/code_mode/delegate.rs @@ -1,12 +1,13 @@ use std::collections::HashMap; +use std::collections::HashSet; use std::sync::Arc; use std::sync::Mutex; -use codex_code_mode::CellId; -use codex_code_mode::CodeModeNestedToolCall; -use codex_code_mode::CodeModeSessionDelegate; -use codex_code_mode::NotificationFuture; -use codex_code_mode::ToolInvocationFuture; +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::NotificationFuture; +use codex_code_mode_protocol::ToolInvocationFuture; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseItem; use serde_json::Value as JsonValue; @@ -24,7 +25,13 @@ use crate::tools::parallel::ToolCallRuntime; pub(super) struct CodeModeDispatchBroker { dispatch_tx: async_channel::Sender, dispatch_rx: async_channel::Receiver, - dispatch_gates: Arc>>>, + dispatch_gates: Arc>, +} + +#[derive(Default)] +struct DispatchGateState { + gates: HashMap>, + terminal_before_ready: HashSet, } impl CodeModeDispatchBroker { @@ -33,18 +40,22 @@ impl CodeModeDispatchBroker { Self { dispatch_tx, dispatch_rx, - dispatch_gates: Arc::new(Mutex::new(HashMap::new())), + dispatch_gates: Arc::new(Mutex::new(DispatchGateState::default())), } } pub(super) fn mark_cell_ready_for_dispatch(&self, cell_id: &CellId) { - dispatch_gate(&self.dispatch_gates, cell_id).send_replace(true); + mark_dispatch_gate_ready(&self.dispatch_gates, cell_id); } pub(super) fn close_cell(&self, cell_id: &CellId) { remove_dispatch_gate(&self.dispatch_gates, cell_id); } + fn record_cell_closed(&self, cell_id: &CellId) { + record_terminal_dispatch_gate(&self.dispatch_gates, cell_id); + } + pub(super) fn start_turn_worker( &self, exec: ExecContext, @@ -130,39 +141,74 @@ impl CodeModeDispatchBroker { } fn dispatch_gate( - dispatch_gates: &Mutex>>, + dispatch_gates: &Mutex, cell_id: &CellId, -) -> watch::Sender { +) -> Option> { let mut dispatch_gates = match dispatch_gates.lock() { Ok(dispatch_gates) => dispatch_gates, Err(poisoned) => poisoned.into_inner(), }; - dispatch_gates - .entry(cell_id.clone()) - .or_insert_with(|| watch::channel(false).0) - .clone() + if dispatch_gates.terminal_before_ready.contains(cell_id) { + return None; + } + Some( + dispatch_gates + .gates + .entry(cell_id.clone()) + .or_insert_with(|| watch::channel(false).0) + .clone(), + ) } -fn remove_dispatch_gate( - dispatch_gates: &Mutex>>, - cell_id: &CellId, -) { +fn mark_dispatch_gate_ready(dispatch_gates: &Mutex, cell_id: &CellId) { let mut dispatch_gates = match dispatch_gates.lock() { Ok(dispatch_gates) => dispatch_gates, Err(poisoned) => poisoned.into_inner(), }; - dispatch_gates.remove(cell_id); + if dispatch_gates.terminal_before_ready.remove(cell_id) { + return; + } + dispatch_gates + .gates + .entry(cell_id.clone()) + .or_insert_with(|| watch::channel(false).0) + .send_replace(true); +} + +fn remove_dispatch_gate(dispatch_gates: &Mutex, cell_id: &CellId) { + let mut dispatch_gates = match dispatch_gates.lock() { + Ok(dispatch_gates) => dispatch_gates, + Err(poisoned) => poisoned.into_inner(), + }; + dispatch_gates.gates.remove(cell_id); +} + +fn record_terminal_dispatch_gate(dispatch_gates: &Mutex, cell_id: &CellId) { + let mut dispatch_gates = match dispatch_gates.lock() { + Ok(dispatch_gates) => dispatch_gates, + Err(poisoned) => poisoned.into_inner(), + }; + let ready = dispatch_gates + .gates + .remove(cell_id) + .is_some_and(|gate| *gate.borrow()); + if !ready { + dispatch_gates.terminal_before_ready.insert(cell_id.clone()); + } } async fn wait_until_cell_ready_for_dispatch( - dispatch_gates: &Mutex>>, + dispatch_gates: &Mutex, cell_id: &CellId, cancellation_token: &CancellationToken, ) -> bool { if cancellation_token.is_cancelled() { return false; } - let mut ready_rx = dispatch_gate(dispatch_gates, cell_id).subscribe(); + let Some(dispatch_gate) = dispatch_gate(dispatch_gates, cell_id) else { + return false; + }; + let mut ready_rx = dispatch_gate.subscribe(); loop { if *ready_rx.borrow_and_update() { return true; @@ -240,7 +286,7 @@ impl CodeModeSessionDelegate for CodeModeDispatchBroker { } fn cell_closed(&self, cell_id: &CellId) { - self.close_cell(cell_id); + self.record_cell_closed(cell_id); } } diff --git a/codex-rs/core/src/tools/code_mode/delegate_tests.rs b/codex-rs/core/src/tools/code_mode/delegate_tests.rs index 23e153d4bf..a9e23c0ab0 100644 --- a/codex-rs/core/src/tools/code_mode/delegate_tests.rs +++ b/codex-rs/core/src/tools/code_mode/delegate_tests.rs @@ -1,66 +1,43 @@ -use std::sync::Arc; -use std::time::Duration; - -use codex_code_mode::CellId; -use codex_code_mode::CellOutcome; -use codex_code_mode::CodeModeService; -use codex_code_mode::CodeModeSessionDelegate; -use codex_code_mode::CreateCellRequest; -use codex_code_mode::ObserveRequest; -use codex_code_mode::RuntimeResponse; +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeSessionDelegate; use super::CodeModeDispatchBroker; #[test] -fn terminal_notification_removes_the_cell_dispatch_gate() { +fn terminal_notification_removes_a_ready_cell_dispatch_gate() { let broker = CodeModeDispatchBroker::new(); let cell_id = CellId::new("cell-a7".to_string()); broker.mark_cell_ready_for_dispatch(&cell_id); - assert!(broker.dispatch_gates.lock().unwrap().contains_key(&cell_id)); + assert!( + broker + .dispatch_gates + .lock() + .unwrap() + .gates + .contains_key(&cell_id) + ); CodeModeSessionDelegate::cell_closed(&broker, &cell_id); - assert!(!broker.dispatch_gates.lock().unwrap().contains_key(&cell_id)); + assert!( + !broker + .dispatch_gates + .lock() + .unwrap() + .gates + .contains_key(&cell_id) + ); } -#[tokio::test] -async fn background_completion_removes_the_dispatch_gate_without_another_observation() { - let broker = Arc::new(CodeModeDispatchBroker::new()); - let service = CodeModeService::with_delegate(broker.clone()); - let cell_id = service - .create_cell(CreateCellRequest { - tool_call_id: "call-1".to_string(), - enabled_tools: Vec::new(), - source: concat!( - "await new Promise(resolve => setTimeout(resolve, 100));", - "text('done');", - ) - .to_string(), - }) - .await - .unwrap(); +#[test] +fn terminal_notification_before_readiness_does_not_leave_a_dispatch_gate() { + let broker = CodeModeDispatchBroker::new(); + let cell_id = CellId::new("cell-b8".to_string()); + + CodeModeSessionDelegate::cell_closed(&broker, &cell_id); broker.mark_cell_ready_for_dispatch(&cell_id); - assert_eq!( - service - .observe(ObserveRequest { - cell_id: cell_id.clone(), - yield_time_ms: 1, - }) - .await - .unwrap(), - CellOutcome::LiveCell(RuntimeResponse::Yielded { - cell_id: cell_id.clone(), - content_items: Vec::new(), - }) - ); - tokio::time::timeout(Duration::from_secs(/*secs*/ 2), async { - while broker.dispatch_gates.lock().unwrap().contains_key(&cell_id) { - tokio::task::yield_now().await; - } - }) - .await - .expect("terminal notification should remove the dispatch gate"); - - service.shutdown().await.unwrap(); + let dispatch_gates = broker.dispatch_gates.lock().unwrap(); + assert!(!dispatch_gates.gates.contains_key(&cell_id)); + assert!(!dispatch_gates.terminal_before_ready.contains(&cell_id)); } diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 8410e08c33..939f1d9a55 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -34,8 +34,8 @@ impl CodeModeExecuteHandler { call_id: String, code: String, ) -> Result { - let args = - codex_code_mode::parse_exec_source(&code).map_err(FunctionCallError::RespondToModel)?; + let args = codex_code_mode_protocol::parse_exec_source(&code) + .map_err(FunctionCallError::RespondToModel)?; let exec = ExecContext { session, turn }; let enabled_tools = codex_tools::collect_code_mode_tool_definitions(&self.nested_tool_specs); @@ -44,7 +44,7 @@ impl CodeModeExecuteHandler { .session .services .code_mode_service - .create_cell(codex_code_mode::CreateCellRequest { + .create_cell(codex_code_mode_protocol::CreateCellRequest { tool_call_id: call_id.clone(), enabled_tools, source: args.code.clone(), @@ -67,6 +67,7 @@ impl CodeModeExecuteHandler { cell_id.clone(), code_cell_trace.clone(), ) + .await .map_err(FunctionCallError::RespondToModel)?; exec.session .services @@ -76,14 +77,14 @@ impl CodeModeExecuteHandler { .session .services .code_mode_service - .observe(codex_code_mode::ObserveRequest { + .observe(codex_code_mode_protocol::ObserveRequest { cell_id: cell_id.clone(), yield_time_ms: args .yield_time_ms - .unwrap_or(codex_code_mode::DEFAULT_EXEC_YIELD_TIME_MS), + .unwrap_or(codex_code_mode_protocol::DEFAULT_EXEC_YIELD_TIME_MS), }) .await; - let response: codex_code_mode::RuntimeResponse = match observed { + let response: codex_code_mode_protocol::RuntimeResponse = match observed { Ok(outcome) => outcome.into(), Err(error) => { initial_observation_guard.finish_with_error(&error); @@ -96,7 +97,10 @@ impl CodeModeExecuteHandler { code_cell_trace.record_initial_response(&response); // Yielded cells keep running, so terminal lifecycle is only emitted // here when the first response also ended the runtime. - if !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) { + if !matches!( + response, + codex_code_mode_protocol::RuntimeResponse::Yielded { .. } + ) { code_cell_trace.record_ended(&response); exec.session .services diff --git a/codex-rs/core/src/tools/code_mode/execute_spec.rs b/codex-rs/core/src/tools/code_mode/execute_spec.rs index 0a858bd206..404525bea6 100644 --- a/codex-rs/core/src/tools/code_mode/execute_spec.rs +++ b/codex-rs/core/src/tools/code_mode/execute_spec.rs @@ -1,4 +1,4 @@ -use codex_code_mode::ToolDefinition as CodeModeToolDefinition; +use codex_code_mode_protocol::ToolDefinition as CodeModeToolDefinition; use codex_tools::FreeformTool; use codex_tools::FreeformToolFormat; use codex_tools::ToolSpec; @@ -6,7 +6,7 @@ use std::collections::BTreeMap; pub(crate) fn create_code_mode_tool( enabled_tools: &[CodeModeToolDefinition], - namespace_descriptions: &BTreeMap, + namespace_descriptions: &BTreeMap, code_mode_only: bool, deferred_tools_available: bool, ) -> ToolSpec { @@ -21,8 +21,8 @@ SOURCE: /[\s\S]+/ "#; ToolSpec::Freeform(FreeformTool { - name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(), - description: codex_code_mode::build_exec_tool_description( + name: codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(), + description: codex_code_mode_protocol::build_exec_tool_description( enabled_tools, namespace_descriptions, code_mode_only, @@ -44,11 +44,11 @@ mod tests { #[test] fn create_code_mode_tool_matches_expected_spec() { - let enabled_tools = vec![codex_code_mode::ToolDefinition { + let enabled_tools = vec![codex_code_mode_protocol::ToolDefinition { name: "update_plan".to_string(), tool_name: ToolName::plain("update_plan"), description: "Update the plan".to_string(), - kind: codex_code_mode::CodeModeToolKind::Function, + kind: codex_code_mode_protocol::CodeModeToolKind::Function, input_schema: None, output_schema: None, }]; @@ -61,8 +61,8 @@ mod tests { /*deferred_tools_available*/ false, ), ToolSpec::Freeform(FreeformTool { - name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(), - description: codex_code_mode::build_exec_tool_description( + name: codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(), + description: codex_code_mode_protocol::build_exec_tool_description( &enabled_tools, &BTreeMap::new(), /*code_mode_only*/ true, diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index ee5527ea4c..b6028cc792 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -6,15 +6,21 @@ mod wait_handler; pub(crate) mod wait_spec; use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::time::Duration; -use codex_code_mode::CellId; -use codex_code_mode::CodeModeNestedToolCall; -use codex_code_mode::CodeModeSession; -use codex_code_mode::CodeModeToolKind; -use codex_code_mode::RuntimeResponse; +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CellOutcome; +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeSession; +use codex_code_mode_protocol::CodeModeSessionProvider; +use codex_code_mode_protocol::CodeModeToolKind; +use codex_code_mode_protocol::RuntimeResponse; use codex_protocol::models::FunctionCallOutputContentItem; use serde_json::Value as JsonValue; +use tokio::sync::Mutex; +use tokio::sync::Semaphore; use tokio_util::sync::CancellationToken; use crate::function_tool::FunctionCallError; @@ -43,9 +49,10 @@ pub(crate) use execute_handler::CodeModeExecuteHandler; use response_adapter::into_function_call_output_content_items; pub(crate) use wait_handler::CodeModeWaitHandler; -pub(crate) const PUBLIC_TOOL_NAME: &str = codex_code_mode::PUBLIC_TOOL_NAME; -pub(crate) const WAIT_TOOL_NAME: &str = codex_code_mode::WAIT_TOOL_NAME; -pub(crate) const DEFAULT_WAIT_YIELD_TIME_MS: u64 = codex_code_mode::DEFAULT_WAIT_YIELD_TIME_MS; +pub(crate) const PUBLIC_TOOL_NAME: &str = codex_code_mode_protocol::PUBLIC_TOOL_NAME; +pub(crate) const WAIT_TOOL_NAME: &str = codex_code_mode_protocol::WAIT_TOOL_NAME; +pub(crate) const DEFAULT_WAIT_YIELD_TIME_MS: u64 = + codex_code_mode_protocol::DEFAULT_WAIT_YIELD_TIME_MS; /// Returns true for the un-namespaced code-mode `exec` tool. pub(crate) fn is_exec_tool_name(tool_name: &ToolName) -> bool { @@ -59,8 +66,11 @@ pub(crate) struct ExecContext { } pub(crate) struct CodeModeService { - session: Option>, + session: Mutex>>, + provider: Arc, dispatch_broker: Arc, + session_init_permit: Semaphore, + shutting_down: AtomicBool, } struct InitialObservationGuard { @@ -71,13 +81,19 @@ struct InitialObservationGuard { } impl InitialObservationGuard { - fn new( + async fn new( service: &CodeModeService, cell_id: CellId, code_cell_trace: codex_rollout_trace::CodeCellTraceContext, ) -> Result { + let session = service + .session + .lock() + .await + .clone() + .ok_or_else(|| "code mode is unavailable".to_string())?; Ok(Self { - session: Arc::clone(service.session()?), + session, dispatch_broker: Arc::clone(&service.dispatch_broker), cell_id: Some(cell_id), code_cell_trace, @@ -95,10 +111,10 @@ impl InitialObservationGuard { }; self.code_cell_trace.record_initial_response(&response); self.code_cell_trace.record_ended(&response); - self.dispatch_broker.close_cell(&cell_id); - let session = Arc::clone(&self.session); + let dispatch_broker = Arc::clone(&self.dispatch_broker); let Ok(runtime) = tokio::runtime::Handle::try_current() else { + dispatch_broker.close_cell(&cell_id); tracing::warn!("unable to terminate code-mode cell {cell_id}: no Tokio runtime"); return; }; @@ -106,6 +122,7 @@ impl InitialObservationGuard { if let Err(error) = session.terminate(cell_id.clone()).await { tracing::warn!("failed to terminate code-mode cell {cell_id}: {error}"); } + dispatch_broker.close_cell(&cell_id); }); } @@ -121,45 +138,56 @@ impl Drop for InitialObservationGuard { } impl CodeModeService { - pub(crate) fn new() -> Self { + pub(crate) fn new(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: Mutex::new(None), + provider, dispatch_broker, + session_init_permit: Semaphore::new(/*permits*/ 1), + shutting_down: AtomicBool::new(false), } } pub(crate) async fn create_cell( &self, - request: codex_code_mode::CreateCellRequest, + request: codex_code_mode_protocol::CreateCellRequest, ) -> Result { - self.session()?.create_cell(request).await + self.session_for_create().await?.create_cell(request).await } pub(crate) async fn observe( &self, - request: codex_code_mode::ObserveRequest, - ) -> Result { - self.session()?.observe(request).await + request: codex_code_mode_protocol::ObserveRequest, + ) -> Result { + match self.current_session().await { + Some(session) => session.observe(request).await, + None => Ok(missing_cell_outcome(request.cell_id)), + } } - pub(crate) async fn terminate( - &self, - cell_id: CellId, - ) -> Result { - self.session()?.terminate(cell_id).await + pub(crate) async fn terminate(&self, cell_id: CellId) -> Result { + match self.current_session().await { + Some(session) => session.terminate(cell_id).await, + None => Ok(missing_cell_outcome(cell_id)), + } } pub(crate) async fn shutdown(&self) -> Result<(), String> { - match &self.session { + self.shutting_down.store(true, Ordering::Release); + let _permit = self + .session_init_permit + .acquire() + .await + .map_err(|_| "code mode session initializer closed".to_string())?; + let session = self.session.lock().await.clone(); + match session { Some(session) => session.shutdown().await, None => Ok(()), } } - pub(crate) fn mark_cell_ready_for_dispatch(&self, cell_id: &codex_code_mode::CellId) { + pub(crate) fn mark_cell_ready_for_dispatch(&self, cell_id: &CellId) { self.dispatch_broker.mark_cell_ready_for_dispatch(cell_id); } @@ -175,9 +203,7 @@ impl CodeModeService { tracker: SharedTurnDiffTracker, ) -> Option { let tool_mode = effective_tool_mode(turn); - if !matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) - || self.session.is_none() - { + if !matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) { return None; } @@ -191,10 +217,43 @@ impl CodeModeService { ) } - fn session(&self) -> Result<&Arc, String> { + async fn session_for_create(&self) -> Result, String> { + if self.shutting_down.load(Ordering::Acquire) { + return Err("code mode session is shutting down".to_string()); + } + if let Some(session) = self.current_session().await { + return Ok(session); + } + let _permit = self + .session_init_permit + .acquire() + .await + .map_err(|_| "code mode session initializer closed".to_string())?; + if self.shutting_down.load(Ordering::Acquire) { + return Err("code mode session is shutting down".to_string()); + } + if let Some(session) = self.current_session().await { + return Ok(session); + } + let session = self + .provider + .create_session(self.dispatch_broker.clone()) + .await?; + if self.shutting_down.load(Ordering::Acquire) { + let _ = session.shutdown().await; + return Err("code mode session is shutting down".to_string()); + } + *self.session.lock().await = Some(Arc::clone(&session)); + Ok(session) + } + + async fn current_session(&self) -> Option> { self.session + .lock() + .await .as_ref() - .ok_or_else(|| "code mode is unavailable".to_string()) + .filter(|session| session.is_alive()) + .cloned() } } @@ -367,6 +426,14 @@ fn serialize_function_tool_arguments( } } +fn missing_cell_outcome(cell_id: CellId) -> CellOutcome { + CellOutcome::MissingCell(RuntimeResponse::Result { + error_text: Some(format!("exec cell {cell_id} not found")), + cell_id, + content_items: Vec::new(), + }) +} + fn build_freeform_tool_payload( tool_name: &ToolName, input: Option, @@ -379,24 +446,33 @@ fn build_freeform_tool_payload( #[cfg(test)] mod tests { + use std::sync::Arc; + use std::sync::Mutex; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use super::CodeModeDispatchBroker; + use super::CodeModeService; use super::InitialObservationGuard; use super::build_nested_tool_payload; use super::truncate_code_mode_result; use crate::tools::context::ToolPayload; - use codex_code_mode::CellId; - use codex_code_mode::CellOutcome; - use codex_code_mode::CodeModeSession; - use codex_code_mode::CodeModeSessionResultFuture; - use codex_code_mode::CodeModeToolKind; - use codex_code_mode::CreateCellRequest; - use codex_code_mode::ObserveRequest; - use codex_code_mode::RuntimeResponse; + use codex_code_mode_protocol::CellId; + use codex_code_mode_protocol::CellOutcome; + use codex_code_mode_protocol::CodeModeSession; + use codex_code_mode_protocol::CodeModeSessionDelegate; + use codex_code_mode_protocol::CodeModeSessionProvider; + use codex_code_mode_protocol::CodeModeSessionProviderFuture; + use codex_code_mode_protocol::CodeModeSessionResultFuture; + use codex_code_mode_protocol::CodeModeToolKind; + use codex_code_mode_protocol::CreateCellRequest; + use codex_code_mode_protocol::ObserveRequest; + use codex_code_mode_protocol::RuntimeResponse; use codex_protocol::models::FunctionCallOutputContentItem; use codex_tools::ToolName; + use pretty_assertions::assert_eq; use serde_json::json; - use std::sync::Arc; - use std::sync::Mutex; struct RecordingCodeModeSession { terminated_cells: Mutex>, @@ -532,4 +608,141 @@ mod tests { vec![cell_id] ); } + + struct RecoveringSessionProvider { + sessions_created: AtomicUsize, + } + + impl CodeModeSessionProvider for RecoveringSessionProvider { + fn create_session<'a>( + &'a self, + _delegate: Arc, + ) -> CodeModeSessionProviderFuture<'a> { + let generation = self.sessions_created.fetch_add(1, Ordering::Relaxed) + 1; + Box::pin(async move { + let session: Arc = Arc::new(RecoveringSession { + generation, + alive: AtomicBool::new(true), + }); + Ok(session) + }) + } + } + + struct RecoveringSession { + generation: usize, + alive: AtomicBool, + } + + impl CodeModeSession for RecoveringSession { + fn is_alive(&self) -> bool { + self.alive.load(Ordering::Acquire) + } + + fn create_cell<'a>( + &'a self, + _request: CreateCellRequest, + ) -> CodeModeSessionResultFuture<'a, CellId> { + Box::pin(async move { + if self.generation == 1 { + self.alive.store(false, Ordering::Release); + return Err("host crashed".to_string()); + } + Ok(CellId::new(format!("cell-{}", self.generation))) + }) + } + + fn observe<'a>( + &'a self, + request: ObserveRequest, + ) -> CodeModeSessionResultFuture<'a, CellOutcome> { + Box::pin(async move { + Ok(CellOutcome::LiveCell(RuntimeResponse::Result { + cell_id: request.cell_id, + content_items: Vec::new(), + error_text: None, + })) + }) + } + + fn terminate<'a>( + &'a self, + _cell_id: CellId, + ) -> CodeModeSessionResultFuture<'a, CellOutcome> { + Box::pin(async { panic!("terminate should not be sent to a failed session") }) + } + + fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { + Box::pin(async { Ok(()) }) + } + } + + fn create_cell_request() -> CreateCellRequest { + CreateCellRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: "text('hello')".to_string(), + } + } + + #[tokio::test] + async fn create_replaces_failed_session_but_observe_does_not() { + let provider = Arc::new(RecoveringSessionProvider { + sessions_created: AtomicUsize::new(0), + }); + let provider_trait: Arc = provider.clone(); + let service = CodeModeService::new(provider_trait); + + assert_eq!( + service + .create_cell(create_cell_request()) + .await + .err() + .as_deref(), + Some("host crashed") + ); + assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1); + + assert_eq!( + service + .observe(ObserveRequest { + cell_id: CellId::new("1".to_string()), + yield_time_ms: 1, + }) + .await, + Ok(CellOutcome::MissingCell(RuntimeResponse::Result { + cell_id: CellId::new("1".to_string()), + content_items: Vec::new(), + error_text: Some("exec cell 1 not found".to_string()), + })) + ); + assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1); + + assert_eq!( + service.terminate(CellId::new("1".to_string())).await, + Ok(CellOutcome::MissingCell(RuntimeResponse::Result { + cell_id: CellId::new("1".to_string()), + content_items: Vec::new(), + error_text: Some("exec cell 1 not found".to_string()), + })) + ); + assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1); + + let cell_id = service.create_cell(create_cell_request()).await.unwrap(); + assert_eq!(cell_id, CellId::new("cell-2".to_string())); + assert_eq!( + service + .observe(ObserveRequest { + cell_id: cell_id.clone(), + yield_time_ms: 1, + }) + .await, + Ok(CellOutcome::LiveCell(RuntimeResponse::Result { + cell_id, + content_items: Vec::new(), + error_text: None, + })) + ); + assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 2); + } } diff --git a/codex-rs/core/src/tools/code_mode/response_adapter.rs b/codex-rs/core/src/tools/code_mode/response_adapter.rs index e20cf6a071..37782f6f74 100644 --- a/codex-rs/core/src/tools/code_mode/response_adapter.rs +++ b/codex-rs/core/src/tools/code_mode/response_adapter.rs @@ -1,4 +1,4 @@ -use codex_code_mode::ImageDetail as CodeModeImageDetail; +use codex_code_mode_protocol::ImageDetail as CodeModeImageDetail; use codex_protocol::models::DEFAULT_IMAGE_DETAIL; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::ImageDetail; @@ -8,7 +8,7 @@ trait IntoProtocol { } pub(super) fn into_function_call_output_content_items( - items: Vec, + items: Vec, ) -> Vec { items.into_iter().map(IntoProtocol::into_protocol).collect() } @@ -26,22 +26,23 @@ impl IntoProtocol for CodeModeImageDetail { } impl IntoProtocol - for codex_code_mode::FunctionCallOutputContentItem + for codex_code_mode_protocol::FunctionCallOutputContentItem { fn into_protocol(self) -> FunctionCallOutputContentItem { let value = self; match value { - codex_code_mode::FunctionCallOutputContentItem::InputText { text } => { + codex_code_mode_protocol::FunctionCallOutputContentItem::InputText { text } => { FunctionCallOutputContentItem::InputText { text } } - codex_code_mode::FunctionCallOutputContentItem::InputImage { image_url, detail } => { - FunctionCallOutputContentItem::InputImage { - image_url, - detail: detail - .map(IntoProtocol::into_protocol) - .or(Some(DEFAULT_IMAGE_DETAIL)), - } - } + codex_code_mode_protocol::FunctionCallOutputContentItem::InputImage { + image_url, + detail, + } => FunctionCallOutputContentItem::InputImage { + image_url, + detail: detail + .map(IntoProtocol::into_protocol) + .or(Some(DEFAULT_IMAGE_DETAIL)), + }, } } } diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index a9b985f155..80b6e60fd8 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -78,7 +78,7 @@ impl CodeModeWaitHandler { let args: ExecWaitArgs = parse_arguments(&arguments)?; let exec = ExecContext { session, turn }; let started_at = std::time::Instant::now(); - let cell_id = codex_code_mode::CellId::new(args.cell_id); + let cell_id = codex_code_mode_protocol::CellId::new(args.cell_id); let wait_response = if args.terminate { exec.session .services @@ -89,23 +89,30 @@ impl CodeModeWaitHandler { exec.session .services .code_mode_service - .observe(codex_code_mode::ObserveRequest { + .observe(codex_code_mode_protocol::ObserveRequest { cell_id, yield_time_ms: args.yield_time_ms, }) .await } .map_err(FunctionCallError::RespondToModel)?; - if let codex_code_mode::CellOutcome::LiveCell(response) = &wait_response - && !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) + if let codex_code_mode_protocol::CellOutcome::LiveCell(response) = &wait_response + && !matches!( + response, + codex_code_mode_protocol::RuntimeResponse::Yielded { .. } + ) { // Only a live-cell wait can close a CodeCell. A missing // cell is still an ordinary `wait` tool result, but there // is no runtime object for the reducer to complete. let runtime_cell_id = match response { - codex_code_mode::RuntimeResponse::Yielded { cell_id, .. } - | codex_code_mode::RuntimeResponse::Terminated { cell_id, .. } - | codex_code_mode::RuntimeResponse::Result { cell_id, .. } => cell_id, + codex_code_mode_protocol::RuntimeResponse::Yielded { cell_id, .. } + | codex_code_mode_protocol::RuntimeResponse::Terminated { + cell_id, .. + } + | codex_code_mode_protocol::RuntimeResponse::Result { cell_id, .. } => { + cell_id + } }; exec.session .services diff --git a/codex-rs/core/src/tools/code_mode/wait_spec.rs b/codex-rs/core/src/tools/code_mode/wait_spec.rs index 72bb09084c..8232705137 100644 --- a/codex-rs/core/src/tools/code_mode/wait_spec.rs +++ b/codex-rs/core/src/tools/code_mode/wait_spec.rs @@ -30,11 +30,11 @@ pub(crate) fn create_wait_tool() -> ToolSpec { ]); ToolSpec::Function(ResponsesApiTool { - name: codex_code_mode::WAIT_TOOL_NAME.to_string(), + name: codex_code_mode_protocol::WAIT_TOOL_NAME.to_string(), description: format!( "Waits on a yielded `{}` cell and returns new output or completion.\n{}", - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::build_wait_tool_description().trim() + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::build_wait_tool_description().trim() ), strict: false, parameters: JsonSchema::object( @@ -57,11 +57,11 @@ mod tests { assert_eq!( create_wait_tool(), ToolSpec::Function(ResponsesApiTool { - name: codex_code_mode::WAIT_TOOL_NAME.to_string(), + name: codex_code_mode_protocol::WAIT_TOOL_NAME.to_string(), description: format!( "Waits on a yielded `{}` cell and returns new output or completion.\n{}", - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::build_wait_tool_description().trim() + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::build_wait_tool_description().trim() ), strict: false, defer_loading: None, diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index af7acf17a0..a378279293 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -279,7 +279,7 @@ fn spec_for_model_request( if matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) && exposure != ToolExposure::DirectModelOnly && !is_excluded_from_code_mode(turn_context, tool_name) - && codex_code_mode::is_code_mode_nested_tool(spec.name()) + && codex_code_mode_protocol::is_code_mode_nested_tool(spec.name()) { codex_tools::augment_tool_spec_for_code_mode(spec) } else { @@ -457,9 +457,9 @@ fn is_hidden_by_code_mode_only( let tool_mode = effective_tool_mode(turn_context); tool_mode == ToolMode::CodeModeOnly && exposure != ToolExposure::DirectModelOnly - && codex_code_mode::is_code_mode_nested_tool(&codex_tools::code_mode_name_for_tool_name( - tool_name, - )) + && codex_code_mode_protocol::is_code_mode_nested_tool( + &codex_tools::code_mode_name_for_tool_name(tool_name), + ) } fn is_excluded_from_code_mode(turn_context: &TurnContext, tool_name: &ToolName) -> bool { @@ -581,7 +581,7 @@ fn merge_into_namespaces(specs: Vec) -> Vec { fn code_mode_namespace_descriptions( specs: &[ToolSpec], -) -> BTreeMap { +) -> BTreeMap { let mut namespace_descriptions = BTreeMap::new(); for spec in specs { let ToolSpec::Namespace(namespace) = spec else { @@ -590,7 +590,7 @@ fn code_mode_namespace_descriptions( let entry = namespace_descriptions .entry(namespace.name.clone()) - .or_insert_with(|| codex_code_mode::ToolNamespaceDescription { + .or_insert_with(|| codex_code_mode_protocol::ToolNamespaceDescription { name: namespace.name.clone(), description: namespace.description.clone(), }); @@ -986,8 +986,8 @@ fn append_extension_tool_executors( .collect::>(); let tool_mode = effective_tool_mode(turn_context); if matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) { - reserved_tool_names.insert(ToolName::plain(codex_code_mode::PUBLIC_TOOL_NAME)); - reserved_tool_names.insert(ToolName::plain(codex_code_mode::WAIT_TOOL_NAME)); + reserved_tool_names.insert(ToolName::plain(codex_code_mode_protocol::PUBLIC_TOOL_NAME)); + reserved_tool_names.insert(ToolName::plain(codex_code_mode_protocol::WAIT_TOOL_NAME)); } if search_tool_enabled(turn_context) && namespace_tools_enabled(turn_context) @@ -1086,9 +1086,9 @@ impl CoreToolRuntime for MultiAgentV2NamespaceOverride { } fn compare_code_mode_tools( - left: &codex_code_mode::ToolDefinition, - right: &codex_code_mode::ToolDefinition, - namespace_descriptions: &BTreeMap, + left: &codex_code_mode_protocol::ToolDefinition, + right: &codex_code_mode_protocol::ToolDefinition, + namespace_descriptions: &BTreeMap, ) -> std::cmp::Ordering { let left_namespace = code_mode_namespace_name(left, namespace_descriptions); let right_namespace = code_mode_namespace_name(right, namespace_descriptions); @@ -1100,8 +1100,11 @@ fn compare_code_mode_tools( } fn code_mode_namespace_name<'a>( - tool: &codex_code_mode::ToolDefinition, - namespace_descriptions: &'a BTreeMap, + tool: &codex_code_mode_protocol::ToolDefinition, + namespace_descriptions: &'a BTreeMap< + String, + codex_code_mode_protocol::ToolNamespaceDescription, + >, ) -> Option<&'a str> { tool.tool_name .namespace diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 5d70fb8f17..caf5cb1836 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -453,8 +453,8 @@ async fn request_user_input_stays_direct_in_code_mode_only() { plan.assert_visible_contains(&[ "request_user_input", - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::WAIT_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::WAIT_TOOL_NAME, ]); plan.assert_registered_contains(&["request_user_input"]); assert_eq!( @@ -462,7 +462,8 @@ async fn request_user_input_stays_direct_in_code_mode_only() { ToolExposure::DirectModelOnly ); - let ToolSpec::Freeform(exec) = plan.visible_spec(codex_code_mode::PUBLIC_TOOL_NAME) else { + let ToolSpec::Freeform(exec) = plan.visible_spec(codex_code_mode_protocol::PUBLIC_TOOL_NAME) + else { panic!("expected code mode exec tool"); }; assert!(!exec.description.contains("request_user_input")); @@ -983,8 +984,8 @@ async fn code_mode_only_exposes_code_executor_and_hides_nested_tools() { &["lookup".to_string()] ); plain.assert_visible_lacks(&[ - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::WAIT_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::WAIT_TOOL_NAME, ]); let code_mode_only = probe_with( @@ -1002,8 +1003,8 @@ async fn code_mode_only_exposes_code_executor_and_hides_nested_tools() { ) .await; code_mode_only.assert_visible_contains(&[ - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::WAIT_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::WAIT_TOOL_NAME, ]); assert_eq!( code_mode_only.namespace_function_names("codex_app"), @@ -1033,8 +1034,8 @@ async fn code_mode_only_exposes_configured_dynamic_namespace_directly() { .await; plan.assert_visible_contains(&[ - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::WAIT_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::WAIT_TOOL_NAME, "direct_only", ]); plan.assert_visible_lacks(&["tool_search"]); @@ -1047,7 +1048,8 @@ async fn code_mode_only_exposes_configured_dynamic_namespace_directly() { }; let ResponsesApiNamespaceTool::Function(tool) = &namespace.tools[0]; assert_eq!(tool.defer_loading, None); - let ToolSpec::Freeform(exec) = plan.visible_spec(codex_code_mode::PUBLIC_TOOL_NAME) else { + let ToolSpec::Freeform(exec) = plan.visible_spec(codex_code_mode_protocol::PUBLIC_TOOL_NAME) + else { panic!("expected code mode exec tool"); }; assert!(!exec.description.contains("direct_only_lookup(args:")); @@ -1075,7 +1077,8 @@ async fn excluded_deferred_namespaces_do_not_enable_nested_tool_guidance() { ) .await; - let ToolSpec::Freeform(exec) = plan.visible_spec(codex_code_mode::PUBLIC_TOOL_NAME) else { + let ToolSpec::Freeform(exec) = plan.visible_spec(codex_code_mode_protocol::PUBLIC_TOOL_NAME) + else { panic!("expected code mode exec tool"); }; assert!( @@ -1222,8 +1225,8 @@ async fn tool_mode_selector_overrides_feature_flags() { }) .await; direct.assert_visible_lacks(&[ - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::WAIT_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::WAIT_TOOL_NAME, ]); } @@ -1460,8 +1463,8 @@ async fn hosted_tools_follow_provider_auth_model_and_config_gates() { code_mode_only.visible_names, vec![ // Code-mode entrypoints. - codex_code_mode::PUBLIC_TOOL_NAME, - codex_code_mode::WAIT_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::WAIT_TOOL_NAME, "request_user_input", // Multi-agent v2 tools. "spawn_agent", diff --git a/codex-rs/core/tests/common/Cargo.toml b/codex-rs/core/tests/common/Cargo.toml index 76f06df1ed..88e96925af 100644 --- a/codex-rs/core/tests/common/Cargo.toml +++ b/codex-rs/core/tests/common/Cargo.toml @@ -16,6 +16,7 @@ anyhow = { workspace = true } assert_cmd = { workspace = true } base64 = { workspace = true } codex-arg0 = { workspace = true } +codex-code-mode-client = { workspace = true } codex-config = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index d843e0097b..6e9cfa59dd 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -13,6 +13,8 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; use anyhow::anyhow; +use codex_code_mode_client::CodeModeHostCommand; +use codex_code_mode_client::StdioCodeModeSessionProvider; use codex_config::CloudConfigBundleLoader; use codex_core::CodexThread; use codex_core::StartThreadOptions; @@ -591,6 +593,14 @@ impl TestCodexBuilder { /*attestation_provider*/ None, /*external_time_provider*/ self.external_time_provider.clone(), ); + let program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host") + .context("resolve code-mode host for stdio core tests")?; + let thread_manager = thread_manager.with_code_mode_session_provider(Arc::new( + StdioCodeModeSessionProvider::new(CodeModeHostCommand { + program, + args: Vec::new(), + }), + )); let thread_manager = Arc::new(thread_manager); let user_shell_override = self.user_shell_override.clone(); diff --git a/codex-rs/core/tests/suite/model_runtime_selectors.rs b/codex-rs/core/tests/suite/model_runtime_selectors.rs index 3ee758372d..0f3008422c 100644 --- a/codex-rs/core/tests/suite/model_runtime_selectors.rs +++ b/codex-rs/core/tests/suite/model_runtime_selectors.rs @@ -158,8 +158,8 @@ async fn remote_tool_mode_selector_overrides_feature_flags() -> Result<()> { assert!( direct_tools .iter() - .all(|name| name != codex_code_mode::PUBLIC_TOOL_NAME - && name != codex_code_mode::WAIT_TOOL_NAME), + .all(|name| name != codex_code_mode_protocol::PUBLIC_TOOL_NAME + && name != codex_code_mode_protocol::WAIT_TOOL_NAME), "direct mode should override enabled code mode flags: {direct_tools:?}" ); @@ -171,8 +171,8 @@ async fn remote_tool_mode_selector_overrides_feature_flags() -> Result<()> { tool_names(&code_mode_only_body), vec![ // Code-mode entrypoints. - codex_code_mode::PUBLIC_TOOL_NAME.to_string(), - codex_code_mode::WAIT_TOOL_NAME.to_string(), + codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(), + codex_code_mode_protocol::WAIT_TOOL_NAME.to_string(), "request_user_input".to_string(), // Hosted Responses tools. "web_search".to_string(), diff --git a/codex-rs/rollout-trace/Cargo.toml b/codex-rs/rollout-trace/Cargo.toml index b368c9acc5..1335fecc20 100644 --- a/codex-rs/rollout-trace/Cargo.toml +++ b/codex-rs/rollout-trace/Cargo.toml @@ -14,7 +14,7 @@ workspace = true [dependencies] anyhow = { workspace = true } -codex-code-mode = { workspace = true } +codex-code-mode-protocol = { workspace = true } codex-protocol = { workspace = true } http = { workspace = true } serde = { workspace = true, features = ["derive"] } diff --git a/codex-rs/rollout-trace/src/code_cell.rs b/codex-rs/rollout-trace/src/code_cell.rs index 5f2603b70a..e215ad41e6 100644 --- a/codex-rs/rollout-trace/src/code_cell.rs +++ b/codex-rs/rollout-trace/src/code_cell.rs @@ -7,7 +7,7 @@ use std::sync::Arc; -use codex_code_mode::RuntimeResponse; +use codex_code_mode_protocol::RuntimeResponse; use serde::Serialize; use tracing::warn; diff --git a/codex-rs/rollout-trace/src/tool_dispatch.rs b/codex-rs/rollout-trace/src/tool_dispatch.rs index 36c6651b10..28a5c003e0 100644 --- a/codex-rs/rollout-trace/src/tool_dispatch.rs +++ b/codex-rs/rollout-trace/src/tool_dispatch.rs @@ -194,7 +194,7 @@ impl ToolDispatchTraceContext { fn suppresses_tool_dispatch_trace(invocation: &ToolDispatchInvocation) -> bool { matches!(invocation.payload, ToolDispatchPayload::Custom { .. }) && invocation.tool_namespace.is_none() - && invocation.tool_name == codex_code_mode::PUBLIC_TOOL_NAME + && invocation.tool_name == codex_code_mode_protocol::PUBLIC_TOOL_NAME } fn record_started(context: &EnabledToolDispatchTraceContext, invocation: ToolDispatchInvocation) { @@ -394,7 +394,7 @@ mod tests { #[test] fn suppresses_only_noncanonical_dispatch_boundaries() { assert!(suppresses_tool_dispatch_trace(&invocation( - codex_code_mode::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, /*tool_namespace*/ None, ToolDispatchRequester::Model { model_visible_call_id: "call-exec".to_string(), @@ -414,7 +414,7 @@ mod tests { }, ))); assert!(!suppresses_tool_dispatch_trace(&invocation( - codex_code_mode::PUBLIC_TOOL_NAME, + codex_code_mode_protocol::PUBLIC_TOOL_NAME, Some("mcp__server".to_string()), ToolDispatchRequester::Model { model_visible_call_id: "call-namespaced".to_string(), diff --git a/codex-rs/tools/Cargo.toml b/codex-rs/tools/Cargo.toml index 7d06acb0c0..3257f1691d 100644 --- a/codex-rs/tools/Cargo.toml +++ b/codex-rs/tools/Cargo.toml @@ -9,7 +9,7 @@ workspace = true [dependencies] codex-app-server-protocol = { workspace = true } -codex-code-mode = { workspace = true } +codex-code-mode-protocol = { workspace = true } codex-features = { workspace = true } codex-file-system = { workspace = true } codex-protocol = { workspace = true } diff --git a/codex-rs/tools/src/code_mode.rs b/codex-rs/tools/src/code_mode.rs index ae1848948f..1e532494ba 100644 --- a/codex-rs/tools/src/code_mode.rs +++ b/codex-rs/tools/src/code_mode.rs @@ -1,8 +1,8 @@ use crate::ResponsesApiNamespaceTool; use crate::ToolName; use crate::ToolSpec; -use codex_code_mode::CodeModeToolKind; -use codex_code_mode::ToolDefinition as CodeModeToolDefinition; +use codex_code_mode_protocol::CodeModeToolKind; +use codex_code_mode_protocol::ToolDefinition as CodeModeToolDefinition; /// Augment tool descriptions with code-mode-specific exec samples. pub fn augment_tool_spec_for_code_mode(spec: ToolSpec) -> ToolSpec { @@ -40,7 +40,8 @@ pub fn augment_tool_spec_for_code_mode(spec: ToolSpec) -> ToolSpec { output_schema: tool.output_schema.clone(), }; tool.description = - codex_code_mode::augment_tool_definition(definition).description; + codex_code_mode_protocol::augment_tool_definition(definition) + .description; } } } @@ -54,8 +55,8 @@ pub fn augment_tool_spec_for_code_mode(spec: ToolSpec) -> ToolSpec { /// including the code-mode-specific description sample. pub fn tool_spec_to_code_mode_tool_definition(spec: &ToolSpec) -> Option { let definition = code_mode_tool_definition_for_spec(spec)?; - codex_code_mode::is_code_mode_nested_tool(&definition.name) - .then(|| codex_code_mode::augment_tool_definition(definition)) + codex_code_mode_protocol::is_code_mode_nested_tool(&definition.name) + .then(|| codex_code_mode_protocol::augment_tool_definition(definition)) } pub fn collect_code_mode_tool_definitions<'a>( @@ -76,8 +77,8 @@ pub fn collect_code_mode_tool_definitions<'a>( } definitions }) - .filter(|definition| codex_code_mode::is_code_mode_nested_tool(&definition.name)) - .map(codex_code_mode::augment_tool_definition) + .filter(|definition| codex_code_mode_protocol::is_code_mode_nested_tool(&definition.name)) + .map(codex_code_mode_protocol::augment_tool_definition) .collect::>(); tool_definitions.sort_by(|left, right| left.name.cmp(&right.name)); tool_definitions.dedup_by(|left, right| left.name == right.name); @@ -90,7 +91,7 @@ pub fn collect_code_mode_exec_prompt_tool_definitions<'a>( let mut tool_definitions = specs .into_iter() .flat_map(code_mode_tool_definitions_for_spec) - .filter(|definition| codex_code_mode::is_code_mode_nested_tool(&definition.name)) + .filter(|definition| codex_code_mode_protocol::is_code_mode_nested_tool(&definition.name)) .collect::>(); tool_definitions.sort_by(|left, right| left.name.cmp(&right.name)); tool_definitions.dedup_by(|left, right| left.name == right.name); @@ -99,7 +100,7 @@ pub fn collect_code_mode_exec_prompt_tool_definitions<'a>( fn augmented_description_for_spec(spec: &ToolSpec) -> Option { code_mode_tool_definition_for_spec(spec) - .map(codex_code_mode::augment_tool_definition) + .map(codex_code_mode_protocol::augment_tool_definition) .map(|definition| definition.description) } diff --git a/codex-rs/tools/src/code_mode_tests.rs b/codex-rs/tools/src/code_mode_tests.rs index c4c4c7ce26..c5c7b7b3f1 100644 --- a/codex-rs/tools/src/code_mode_tests.rs +++ b/codex-rs/tools/src/code_mode_tests.rs @@ -69,7 +69,7 @@ declare const tools: { lookup_order(args: { order_id: string; }): Promise<{ ok: fn augment_tool_spec_for_code_mode_preserves_exec_tool_description() { assert_eq!( augment_tool_spec_for_code_mode(ToolSpec::Freeform(FreeformTool { - name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(), + name: codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(), description: "Run code".to_string(), format: FreeformToolFormat { r#type: "grammar".to_string(), @@ -78,7 +78,7 @@ fn augment_tool_spec_for_code_mode_preserves_exec_tool_description() { }, })), ToolSpec::Freeform(FreeformTool { - name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(), + name: codex_code_mode_protocol::PUBLIC_TOOL_NAME.to_string(), description: "Run code".to_string(), format: FreeformToolFormat { r#type: "grammar".to_string(), @@ -103,7 +103,7 @@ fn tool_spec_to_code_mode_tool_definition_returns_augmented_nested_tools() { assert_eq!( tool_spec_to_code_mode_tool_definition(&spec), - Some(codex_code_mode::ToolDefinition { + Some(codex_code_mode_protocol::ToolDefinition { name: "apply_patch".to_string(), tool_name: ToolName::plain("apply_patch"), description: r#"Apply a patch @@ -113,7 +113,7 @@ exec tool declaration: declare const tools: { apply_patch(input: string): Promise; }; ```"# .to_string(), - kind: codex_code_mode::CodeModeToolKind::Freeform, + kind: codex_code_mode_protocol::CodeModeToolKind::Freeform, input_schema: None, output_schema: None, }) diff --git a/scripts/run-nextest.py b/scripts/run-nextest.py index 6a6d0679ba..c996c47143 100644 --- a/scripts/run-nextest.py +++ b/scripts/run-nextest.py @@ -54,6 +54,13 @@ def main() -> int: check=True, env=env, ) + target_dir = Path(env.get("CARGO_TARGET_DIR", CODEX_RS_ROOT / "target")) + if not target_dir.is_absolute(): + target_dir = CODEX_RS_ROOT / target_dir + host_name = ( + "codex-code-mode-host.exe" if os.name == "nt" else "codex-code-mode-host" + ) + env["CODEX_CODE_MODE_HOST_PATH"] = str(target_dir / "debug" / host_name) return subprocess.run( [cargo, "nextest", "run", "--no-fail-fast", *sys.argv[1:]], cwd=CODEX_RS_ROOT,