From b0cf57bdeb4d78527f52cb815e69ea13cfd514ba Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Fri, 12 Jun 2026 00:02:33 +0000 Subject: [PATCH] code-mode: add opt-in IPC consumer --- codex-rs/Cargo.lock | 3 + codex-rs/app-server/BUILD.bazel | 1 + .../tests/suite/v2/imagegen_extension.rs | 14 +- codex-rs/core/BUILD.bazel | 1 + codex-rs/core/Cargo.toml | 2 + 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 | 10 +- .../src/tools/code_mode/execute_handler.rs | 11 +- .../core/src/tools/code_mode/execute_spec.rs | 16 +- codex-rs/core/src/tools/code_mode/mod.rs | 256 ++++++++++++++++-- .../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 | 19 +- .../src/tools/tool_dispatch_trace_tests.rs | 76 ++++++ codex-rs/core/tests/common/Cargo.toml | 1 + codex-rs/core/tests/common/test_codex.rs | 12 + codex-rs/core/tests/suite/code_mode.rs | 15 +- .../tests/suite/model_runtime_selectors.rs | 8 +- 26 files changed, 470 insertions(+), 113 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index bb5c8a6ba5..00505796e3 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2621,6 +2621,8 @@ dependencies = [ "codex-apply-patch", "codex-async-utils", "codex-code-mode", + "codex-code-mode-client", + "codex-code-mode-protocol", "codex-config", "codex-connectors", "codex-context-fragments", @@ -4624,6 +4626,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 6765141bdc..1dd4fde7a7 100644 --- a/codex-rs/app-server/BUILD.bazel +++ b/codex-rs/app-server/BUILD.bazel @@ -16,6 +16,7 @@ codex_rust_crate( }, extra_binaries = [ "//codex-rs/bwrap:bwrap", + "//codex-rs/code-mode-host:codex-code-mode-host", ], test_tags = ["no-sandbox"], ) diff --git a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs index 025c6ec186..3d53e38243 100644 --- a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs +++ b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs @@ -290,9 +290,19 @@ generatedImage(result); ChatGptAuthFixture::new("access-chatgpt"), AuthCredentialsStoreMode::File, )?; + let code_mode_host = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + let code_mode_host = code_mode_host + .to_str() + .context("code-mode host path is not valid UTF-8")?; - let mut mcp = - TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + let mut mcp = TestAppServer::new_with_env( + codex_home.path(), + &[ + ("OPENAI_API_KEY", None), + ("CODEX_CODE_MODE_HOST_PATH", Some(code_mode_host)), + ], + ) + .await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; start_image_generation_turn(&mut mcp).await?; timeout( diff --git a/codex-rs/core/BUILD.bazel b/codex-rs/core/BUILD.bazel index 478699c0a1..4fb5ab0725 100644 --- a/codex-rs/core/BUILD.bazel +++ b/codex-rs/core/BUILD.bazel @@ -41,6 +41,7 @@ codex_rust_crate( unit_test_timeout = "long", 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 d21245c63e..cf37e1021d 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -29,6 +29,7 @@ codex-app-server-protocol = { workspace = true } codex-apply-patch = { workspace = true } codex-async-utils = { workspace = true } codex-code-mode = { workspace = true } +codex-code-mode-protocol = { workspace = true } codex-connectors = { workspace = true } codex-context-fragments = { workspace = true } codex-config = { workspace = true } @@ -133,6 +134,7 @@ codex-shell-escalation = { workspace = true } [dev-dependencies] assert_cmd = { workspace = true } assert_matches = { workspace = true } +codex-code-mode-client = { workspace = true } codex-image-generation-extension = { workspace = true } codex-home = { workspace = true } codex-otel = { workspace = true } diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 135909bf00..76180f5520 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -113,6 +113,7 @@ pub(crate) async fn run_codex_thread_interactive( analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), attestation_provider: parent_session.services.attestation_provider.clone(), + code_mode_session_provider: Arc::clone(&parent_session.services.code_mode_session_provider), inherited_multi_agent_version: Some(MultiAgentVersion::Disabled), })) .or_cancel(&cancel_token) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 5e93352fd8..ff2f38d885 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -432,6 +432,8 @@ pub(crate) struct CodexSpawnArgs { pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, + pub(crate) code_mode_session_provider: + Arc, pub(crate) inherited_multi_agent_version: Option, } @@ -514,6 +516,7 @@ impl Codex { analytics_events_client, thread_store, attestation_provider, + code_mode_session_provider, inherited_multi_agent_version, } = args; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); @@ -667,6 +670,7 @@ impl Codex { thread_store, parent_rollout_thread_trace, attestation_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 df60870a6e..5764206a85 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -492,6 +492,7 @@ impl Session { thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, attestation_provider: Option>, + code_mode_session_provider: Arc, multi_agent_version: Option, ) -> anyhow::Result> { debug!( @@ -1031,7 +1032,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, environment_manager, }; let (out_of_band_elicitation_paused, _out_of_band_elicitation_paused_rx) = diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 214489a18a..749128eb9c 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4835,6 +4835,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + Arc::new(codex_code_mode_client::IpcCodeModeSessionProvider::default()), Some(config.multi_agent_version_from_features()), ) .await; @@ -5001,7 +5002,12 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { Session::build_model_client_beta_features_header(config.as_ref()), /*attestation_provider*/ None, ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::new( + codex_code_mode_client::IpcCodeModeSessionProvider::default(), + )), + code_mode_session_provider: Arc::new( + codex_code_mode_client::IpcCodeModeSessionProvider::default(), + ), environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), }; @@ -5173,6 +5179,7 @@ async fn make_session_with_config_and_rx( )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + Arc::new(codex_code_mode_client::IpcCodeModeSessionProvider::default()), Some(config.multi_agent_version_from_features()), ) .await?; @@ -5282,6 +5289,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, + Arc::new(codex_code_mode_client::IpcCodeModeSessionProvider::default()), Some(config.multi_agent_version_from_features()), ) .await?; @@ -7000,7 +7008,12 @@ where Session::build_model_client_beta_features_header(config.as_ref()), /*attestation_provider*/ None, ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::new( + codex_code_mode_client::IpcCodeModeSessionProvider::default(), + )), + code_mode_session_provider: Arc::new( + codex_code_mode_client::IpcCodeModeSessionProvider::default(), + ), environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), }; diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 0d3c9c033b..397dc5c5fa 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -736,6 +736,9 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { analytics_events_client: None, thread_store, attestation_provider: None, + code_mode_session_provider: Arc::new( + codex_code_mode_client::IpcCodeModeSessionProvider::default(), + ), inherited_multi_agent_version: None, }) .await diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index d8d91e85a8..74382a0d02 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -19,6 +19,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_exec_server::EnvironmentManager; use codex_extension_api::ExtensionData; @@ -79,6 +80,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, /// Shared process-level environment registry. Sessions carry an `Arc` handle so they can pass /// the same manager through child-thread spawn paths without reconstructing it. pub(crate) environment_manager: Arc, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index c332622acb..6fb725af7b 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::InProcessCodeModeSessionProvider; +use codex_code_mode_protocol::CodeModeSessionProvider; use codex_core_plugins::PluginsManager; use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionDataInit; @@ -218,6 +220,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, } @@ -304,6 +307,7 @@ impl ThreadManager { installation_id, analytics_events_client, state_db, + code_mode_session_provider: Arc::new(InProcessCodeModeSessionProvider), ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -409,6 +413,7 @@ impl ThreadManager { installation_id, analytics_events_client: None, state_db, + code_mode_session_provider: Arc::new(InProcessCodeModeSessionProvider), ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -420,6 +425,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() } @@ -1430,6 +1447,7 @@ impl ThreadManagerState { analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), attestation_provider: self.attestation_provider.clone(), + code_mode_session_provider: Arc::clone(&self.code_mode_session_provider), inherited_multi_agent_version: multi_agent_version, })) .await?; diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index 06bb71a280..cc56caeab9 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/codex-rs/core/src/tools/code_mode/delegate.rs @@ -2,11 +2,11 @@ use std::collections::HashMap; 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; diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 2e46c1bd6b..ea24be4e15 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -33,8 +33,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); @@ -43,7 +43,7 @@ impl CodeModeExecuteHandler { .session .services .code_mode_service - .execute(codex_code_mode::ExecuteRequest { + .execute(codex_code_mode_protocol::ExecuteRequest { tool_call_id: call_id.clone(), enabled_tools, source: args.code.clone(), @@ -78,7 +78,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 ade23cf3b7..9224ba45a4 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -8,13 +8,18 @@ pub(crate) mod wait_spec; use std::sync::Arc; 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::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 std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use tokio::sync::Mutex; +use tokio::sync::Semaphore; use tokio_util::sync::CancellationToken; use crate::function_tool::FunctionCallError; @@ -42,9 +47,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 { @@ -58,50 +64,69 @@ pub(crate) struct ExecContext { } pub(crate) struct CodeModeService { - session: Option>, + session: Mutex>>, + provider: Arc, dispatch_broker: Arc, + session_init_permit: Semaphore, + shutting_down: AtomicBool, } 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 execute( &self, - request: codex_code_mode::ExecuteRequest, - ) -> Result { - self.session()?.execute(request).await + request: codex_code_mode_protocol::ExecuteRequest, + ) -> Result { + self.session_for_execute().await?.execute(request).await } pub(crate) async fn wait( &self, - request: codex_code_mode::WaitRequest, - ) -> Result { - self.session()?.wait(request).await + request: codex_code_mode_protocol::WaitRequest, + ) -> Result { + self.current_session() + .await + .ok_or_else(|| "code mode session is unavailable".to_string())? + .wait(request) + .await } pub(crate) async fn terminate( &self, cell_id: CellId, - ) -> Result { - self.session()?.terminate(cell_id).await + ) -> Result { + self.current_session() + .await + .ok_or_else(|| "code mode session is unavailable".to_string())? + .terminate(cell_id) + .await } 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: &codex_code_mode_protocol::CellId) { self.dispatch_broker.mark_cell_ready_for_dispatch(cell_id); } @@ -116,9 +141,7 @@ impl CodeModeService { router: Arc, tracker: SharedTurnDiffTracker, ) -> Option { - if !matches!(turn.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) - || self.session.is_none() - { + if !matches!(turn.tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) { return None; } @@ -132,10 +155,43 @@ impl CodeModeService { ) } - fn session(&self) -> Result<&Arc, String> { + async fn session_for_execute(&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() } } @@ -320,11 +376,30 @@ fn build_freeform_tool_payload( #[cfg(test)] mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use super::CodeModeService; use super::build_nested_tool_payload; use crate::tools::context::ToolPayload; - use codex_code_mode::CodeModeToolKind; + use codex_code_mode_protocol::CellId; + 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::ExecuteRequest; + use codex_code_mode_protocol::RuntimeResponse; + use codex_code_mode_protocol::StartedCell; + use codex_code_mode_protocol::WaitOutcome; + use codex_code_mode_protocol::WaitRequest; use codex_tools::ToolName; + use pretty_assertions::assert_eq; use serde_json::json; + use tokio::sync::oneshot; #[test] fn build_nested_tool_payload_uses_function_kind() { @@ -359,4 +434,127 @@ mod tests { other => panic!("expected freeform payload, got {other:?}"), } } + + 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 execute<'a>( + &'a self, + _request: ExecuteRequest, + ) -> CodeModeSessionResultFuture<'a, StartedCell> { + Box::pin(async move { + if self.generation == 1 { + self.alive.store(false, Ordering::Release); + return Err("host crashed".to_string()); + } + let cell_id = CellId::new(format!("host{}_1", self.generation)); + let (response_tx, response_rx) = oneshot::channel(); + response_tx + .send(RuntimeResponse::Result { + cell_id: cell_id.clone(), + content_items: Vec::new(), + error_text: None, + }) + .expect("test response receiver should be live"); + Ok(StartedCell::new(cell_id, response_rx)) + }) + } + + fn wait<'a>( + &'a self, + _request: WaitRequest, + ) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(async { panic!("wait should not be sent to a failed session") }) + } + + fn terminate<'a>( + &'a self, + _cell_id: CellId, + ) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + 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 execute_request() -> ExecuteRequest { + ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: "text('hello')".to_string(), + yield_time_ms: None, + max_output_tokens: None, + } + } + + #[tokio::test] + async fn exec_replaces_failed_session_but_wait_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.execute(execute_request()).await.err().as_deref(), + Some("host crashed") + ); + assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1); + + assert_eq!( + service + .wait(WaitRequest { + cell_id: CellId::new("host1_1".to_string()), + yield_time_ms: 1, + }) + .await + .err() + .as_deref(), + Some("code mode session is unavailable") + ); + assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1); + + assert_eq!( + service + .terminate(CellId::new("host1_1".to_string())) + .await + .err() + .as_deref(), + Some("code mode session is unavailable") + ); + assert_eq!(provider.sessions_created.load(Ordering::Relaxed), 1); + + let started = service.execute(execute_request()).await.unwrap(); + assert_eq!(started.cell_id, CellId::new("host2_1".to_string())); + 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 953391a573..0db26296ec 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 - .wait(codex_code_mode::WaitRequest { + .wait(codex_code_mode_protocol::WaitRequest { cell_id, yield_time_ms: args.yield_time_ms, }) .await } .map_err(FunctionCallError::RespondToModel)?; - if let codex_code_mode::WaitOutcome::LiveCell(response) = &wait_response - && !matches!(response, codex_code_mode::RuntimeResponse::Yielded { .. }) + if let codex_code_mode_protocol::WaitOutcome::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 79c86b94d1..292192bc32 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -242,7 +242,7 @@ fn spec_for_model_request( 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 { @@ -408,9 +408,9 @@ fn is_hidden_by_code_mode_only( ) -> bool { 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 { @@ -533,7 +533,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 { @@ -542,7 +542,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(), }); @@ -883,8 +883,8 @@ fn append_extension_tool_executors( turn_context.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) @@ -983,9 +983,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); @@ -997,8 +997,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 2f7f23947b..7c4cd69985 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -876,8 +876,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( @@ -895,8 +895,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"), @@ -926,7 +926,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!( @@ -1071,8 +1072,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, ]); } @@ -1307,8 +1308,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, // Multi-agent v2 tools. "spawn_agent", "send_message", diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index b09996872e..dba0127dc6 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -3,6 +3,17 @@ use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use codex_code_mode_protocol::CellId; +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::ExecuteRequest; +use codex_code_mode_protocol::RuntimeResponse; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::WaitOutcome; +use codex_code_mode_protocol::WaitRequest; use codex_protocol::protocol::SessionSource; use codex_rollout_trace::ExecutionStatus; use codex_rollout_trace::ThreadStartedTraceMetadata; @@ -58,6 +69,51 @@ impl ToolExecutor for TestHandler { impl CoreToolRuntime for TestHandler {} +struct MissingCellSessionProvider; + +impl CodeModeSessionProvider for MissingCellSessionProvider { + fn create_session<'a>( + &'a self, + _delegate: Arc, + ) -> CodeModeSessionProviderFuture<'a> { + Box::pin(async { + let session: Arc = Arc::new(MissingCellSession); + Ok(session) + }) + } +} + +struct MissingCellSession; + +impl CodeModeSession for MissingCellSession { + fn execute<'a>( + &'a self, + _request: ExecuteRequest, + ) -> CodeModeSessionResultFuture<'a, StartedCell> { + Box::pin(async { Err("execute is not supported by this test session".to_string()) }) + } + + fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(async move { Ok(missing_cell_outcome(request.cell_id)) }) + } + + fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(async move { Ok(missing_cell_outcome(cell_id)) }) + } + + fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { + Box::pin(async { Ok(()) }) + } +} + +fn missing_cell_outcome(cell_id: CellId) -> WaitOutcome { + WaitOutcome::MissingCell(RuntimeResponse::Result { + error_text: Some(format!("exec cell {cell_id} not found")), + cell_id, + content_items: Vec::new(), + }) +} + #[tokio::test] async fn dispatch_lifecycle_trace_records_direct_and_code_mode_requesters() -> anyhow::Result<()> { let temp = TempDir::new()?; @@ -213,6 +269,26 @@ async fn dispatch_lifecycle_trace_records_incompatible_payload_failures() -> any async fn missing_code_mode_wait_traces_only_the_wait_tool_call() -> anyhow::Result<()> { let temp = TempDir::new()?; let (mut session, turn) = make_session_and_context().await; + let provider: Arc = Arc::new(MissingCellSessionProvider); + session.services.code_mode_service = + crate::tools::code_mode::CodeModeService::new(Arc::clone(&provider)); + session.services.code_mode_session_provider = provider; + assert_eq!( + session + .services + .code_mode_service + .execute(ExecuteRequest { + tool_call_id: "setup-call".to_string(), + enabled_tools: Vec::new(), + source: String::new(), + yield_time_ms: None, + max_output_tokens: None, + }) + .await + .err() + .as_deref(), + Some("execute is not supported by this test session") + ); attach_test_trace(&mut session, &turn, temp.path())?; let registry = ToolRegistry::with_handler_for_test(Arc::new(CodeModeWaitHandler)); diff --git a/codex-rs/core/tests/common/Cargo.toml b/codex-rs/core/tests/common/Cargo.toml index 76f06df1ed..956162f25e 100644 --- a/codex-rs/core/tests/common/Cargo.toml +++ b/codex-rs/core/tests/common/Cargo.toml @@ -17,6 +17,7 @@ assert_cmd = { workspace = true } base64 = { workspace = true } codex-arg0 = { workspace = true } codex-config = { workspace = true } +codex-code-mode-client = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-exec-server = { workspace = true } diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 1582b44d7d..cee430f496 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::IpcCodeModeSessionProvider; use codex_config::CloudConfigBundleLoader; use codex_core::CodexThread; use codex_core::ThreadManager; @@ -567,6 +569,16 @@ impl TestCodexBuilder { installation_id, /*attestation_provider*/ None, ); + let code_mode_session_provider = + match codex_utils_cargo_bin::cargo_bin("codex-code-mode-host") { + Ok(program) => IpcCodeModeSessionProvider::new(CodeModeHostCommand { + program, + args: Vec::new(), + }), + Err(_) => IpcCodeModeSessionProvider::default(), + }; + let thread_manager = + thread_manager.with_code_mode_session_provider(Arc::new(code_mode_session_provider)); let thread_manager = Arc::new(thread_manager); let user_shell_override = self.user_shell_override.clone(); diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 0f954bac7b..11d78986b9 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -1921,7 +1921,7 @@ text("after terminate"); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn code_mode_wait_returns_error_for_unknown_session() -> Result<()> { +async fn code_mode_wait_returns_error_without_a_session() -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; @@ -1964,17 +1964,10 @@ async fn code_mode_wait_returns_error_for_unknown_session() -> Result<()> { assert_ne!(success, Some(true)); let items = function_tool_output_items(&request, "call-1"); - assert_eq!(items.len(), 2); - assert_regex_match( - concat!( - r"(?s)\A", - r"Script failed\nWall time \d+\.\d seconds\nOutput:\n\z" - ), - text_item(&items, /*index*/ 0), - ); + assert_eq!(items.len(), 1); assert_eq!( - text_item(&items, /*index*/ 1), - "Script error:\nexec cell 999999 not found" + text_item(&items, /*index*/ 0), + "code mode session is unavailable" ); Ok(()) diff --git a/codex-rs/core/tests/suite/model_runtime_selectors.rs b/codex-rs/core/tests/suite/model_runtime_selectors.rs index 38c95394ea..f14f998c05 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(), // Hosted Responses tools. "web_search".to_string(), "image_generation".to_string(),