diff --git a/codex-rs/code-mode/src/remote_session.rs b/codex-rs/code-mode/src/remote_session.rs index af7d27845b..ae0a82e6dc 100644 --- a/codex-rs/code-mode/src/remote_session.rs +++ b/codex-rs/code-mode/src/remote_session.rs @@ -22,6 +22,7 @@ use tokio::sync::Semaphore; use tokio::sync::watch; use self::connection::Connection; +use self::connection::ConnectionError; use self::connection::RemoteSession; use self::connection::SessionCleanup; use crate::NoopCodeModeSessionDelegate; @@ -34,30 +35,32 @@ type ShutdownResultReceiver = watch::Receiver>>; /// Creates code-mode sessions backed by one lazily spawned process host. pub struct ProcessOwnedCodeModeSessionProvider { - host_program: PathBuf, - process_host: StdMutex>>, + state: StdMutex, +} + +enum ProviderState { + OwnedProcess(Arc), + InProcess, } impl ProcessOwnedCodeModeSessionProvider { pub fn with_host_program(host_program: PathBuf) -> Self { Self { - host_program, - process_host: StdMutex::new(None), + state: StdMutex::new(ProviderState::OwnedProcess(Arc::new( + OwnedProcessHost::new(host_program), + ))), } } - fn process_host(&self) -> Arc { - let mut process_host = self - .process_host + fn process_host(&self) -> Option> { + match &*self + .state .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if let Some(process_host) = process_host.as_ref() { - return Arc::clone(process_host); + .unwrap_or_else(std::sync::PoisonError::into_inner) + { + ProviderState::OwnedProcess(process_host) => Some(Arc::clone(process_host)), + ProviderState::InProcess => None, } - - let new_process_host = Arc::new(OwnedProcessHost::new(self.host_program.clone())); - *process_host = Some(Arc::clone(&new_process_host)); - new_process_host } } @@ -72,8 +75,28 @@ impl CodeModeSessionProvider for ProcessOwnedCodeModeSessionProvider { &'a self, delegate: Arc, ) -> CodeModeSessionProviderFuture<'a> { - let session = ProcessOwnedCodeModeSession::with_process_host(delegate, self.process_host()); Box::pin(async move { + let Some(process_host) = self.process_host() else { + let session: Arc = + Arc::new(crate::InProcessCodeModeSession::with_delegate(delegate)); + return Ok(session); + }; + + match process_host.connection().await { + Ok(_) => {} + Err(error) if error.host_program_not_found() => { + *self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + ProviderState::InProcess; + let session: Arc = + Arc::new(crate::InProcessCodeModeSession::with_delegate(delegate)); + return Ok(session); + } + Err(error) => return Err(error.to_string()), + } + let session = ProcessOwnedCodeModeSession::with_process_host(delegate, process_host); session.connection().await?; let session: Arc = Arc::new(session); Ok(session) @@ -98,16 +121,14 @@ impl OwnedProcessHost { } } - async fn connection(&self) -> Result, String> { + async fn connection(&self) -> Result, ConnectionError> { if let Some(connection) = self.live_connection() { return Ok(connection); } - let _spawn_permit = self - .spawn_permit - .acquire() - .await - .map_err(|_| "code-mode host spawn coordinator closed".to_string())?; + let _spawn_permit = self.spawn_permit.acquire().await.map_err(|_| { + ConnectionError::Other("code-mode host spawn coordinator closed".into()) + })?; if let Some(connection) = self.live_connection() { return Ok(connection); } @@ -284,7 +305,7 @@ impl SessionInner { cleanup, }) } - Err(err) => Err(err), + Err(err) => Err(err.to_string()), }; { let mut state = self diff --git a/codex-rs/code-mode/src/remote_session/connection.rs b/codex-rs/code-mode/src/remote_session/connection.rs index 7154784def..18ba5e0974 100644 --- a/codex-rs/code-mode/src/remote_session/connection.rs +++ b/codex-rs/code-mode/src/remote_session/connection.rs @@ -1,4 +1,7 @@ +use std::fmt; +use std::io; use std::path::Path; +use std::path::PathBuf; use std::process::Stdio; use std::sync::Arc; use std::sync::atomic::AtomicBool; @@ -46,6 +49,39 @@ mod reader; const IPC_CHANNEL_CAPACITY: usize = 128; const HOST_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10); +pub(super) enum ConnectionError { + Spawn { + host_program: PathBuf, + error: io::Error, + }, + Other(String), +} + +impl ConnectionError { + pub(super) fn host_program_not_found(&self) -> bool { + matches!( + self, + Self::Spawn { error, .. } if error.kind() == io::ErrorKind::NotFound + ) + } +} + +impl fmt::Display for ConnectionError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Spawn { + host_program, + error, + } => write!( + formatter, + "failed to spawn code-mode host {}: {error}", + host_program.display() + ), + Self::Other(message) => formatter.write_str(message), + } + } +} + pub(super) struct Connection { command_tx: mpsc::Sender, execute_claim_tx: mpsc::UnboundedSender, @@ -96,7 +132,7 @@ impl Drop for CallerCancellation { } impl Connection { - pub(super) async fn spawn(host_program: &Path) -> Result { + pub(super) async fn spawn(host_program: &Path) -> Result { let mut command = Command::new(host_program); #[cfg(unix)] command.process_group(0); @@ -106,11 +142,9 @@ impl Connection { .stderr(Stdio::piped()) .kill_on_drop(true) .spawn() - .map_err(|err| { - format!( - "failed to spawn code-mode host {}: {err}", - host_program.display() - ) + .map_err(|error| ConnectionError::Spawn { + host_program: host_program.to_path_buf(), + error, })?; if let Some(stderr) = child.stderr.take() { @@ -132,11 +166,11 @@ impl Connection { let stdin = child .stdin .take() - .ok_or_else(|| "spawned code-mode host has no stdin".to_string())?; + .ok_or_else(|| ConnectionError::Other("spawned code-mode host has no stdin".into()))?; let stdout = child .stdout .take() - .ok_or_else(|| "spawned code-mode host has no stdout".to_string())?; + .ok_or_else(|| ConnectionError::Other("spawned code-mode host has no stdout".into()))?; let mut reader = FramedReader::new(stdout); let mut writer = FramedWriter::new(stdin); let handshake = async { @@ -174,12 +208,14 @@ impl Connection { Ok(result) => result, Err(_) => { kill_and_reap(&mut child).await; - return Err("timed out negotiating with the code-mode host".to_string()); + return Err(ConnectionError::Other( + "timed out negotiating with the code-mode host".into(), + )); } }; if let Err(err) = handshake_result { kill_and_reap(&mut child).await; - return Err(err); + return Err(ConnectionError::Other(err)); } let (command_tx, command_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY); diff --git a/codex-rs/code-mode/src/remote_session_tests.rs b/codex-rs/code-mode/src/remote_session_tests.rs index 03d75e5239..e48ac2d71e 100644 --- a/codex-rs/code-mode/src/remote_session_tests.rs +++ b/codex-rs/code-mode/src/remote_session_tests.rs @@ -3,6 +3,10 @@ use std::path::PathBuf; use std::sync::Arc; use codex_code_mode_protocol::CodeModeSessionProvider; +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::FunctionCallOutputContentItem; +use codex_code_mode_protocol::RuntimeResponse; +use pretty_assertions::assert_eq; use super::ProcessOwnedCodeModeSession; use super::ProcessOwnedCodeModeSessionProvider; @@ -13,8 +17,8 @@ use crate::NoopCodeModeSessionDelegate; fn provider_reuses_its_live_process_host() { let provider = ProcessOwnedCodeModeSessionProvider::default(); - let first = provider.process_host(); - let second = provider.process_host(); + let first = provider.process_host().expect("owned process host"); + let second = provider.process_host().expect("owned process host"); assert!(Arc::ptr_eq(&first, &second)); } @@ -68,18 +72,39 @@ fn host_program_falls_back_to_its_name_when_main_executable_is_unknown() { } #[tokio::test] -async fn provider_reports_host_spawn_failure() { +async fn provider_falls_back_to_in_process_session_when_host_is_missing() { let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( "codex-code-mode-host-does-not-exist".into(), ); - let error = provider + let session = provider .create_session(Arc::new(NoopCodeModeSessionDelegate)) .await - .err() - .expect("session creation should fail"); + .expect("missing host should fall back to an in-process session"); + let response = session + .execute(ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: "text('fallback')".to_string(), + yield_time_ms: None, + max_output_tokens: None, + }) + .await + .expect("execute fallback session") + .initial_response() + .await + .expect("read fallback response"); - assert!(error.contains("failed to spawn code-mode host")); + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: codex_code_mode_protocol::CellId::new("1".to_string()), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "fallback".to_string(), + }], + error_text: None, + } + ); } #[tokio::test] diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index ad746a1de2..bde373de94 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -367,7 +367,9 @@ mod tests { use crate::tools::context::ToolPayload; use codex_code_mode::CodeModeToolKind; use codex_code_mode::ExecuteRequest; + use codex_code_mode::FunctionCallOutputContentItem as CodeModeOutputContentItem; use codex_code_mode::ProcessOwnedCodeModeSessionProvider; + use codex_code_mode::RuntimeResponse; use codex_protocol::models::FunctionCallOutputContentItem; use codex_tools::ToolName; use serde_json::json; @@ -426,26 +428,37 @@ mod tests { } #[tokio::test] - async fn missing_process_host_is_reported_without_failing_service_creation() { + async fn missing_process_host_falls_back_to_in_process_session() { let service = CodeModeService::new(Arc::new( ProcessOwnedCodeModeSessionProvider::with_host_program( "codex-code-mode-host-does-not-exist".into(), ), )); - let error = service + let response = service .execute(ExecuteRequest { tool_call_id: "call-1".to_string(), enabled_tools: Vec::new(), - source: "text('unreachable')".to_string(), + source: "text('fallback')".to_string(), yield_time_ms: None, max_output_tokens: None, }) .await - .err() - .expect("missing host should reject execution"); + .expect("missing host should fall back to an in-process session") + .initial_response() + .await + .expect("read fallback response"); - assert!(error.contains("failed to spawn code-mode host")); - service.shutdown().await.expect("shutdown unused service"); + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: codex_code_mode::CellId::new("1".to_string()), + content_items: vec![CodeModeOutputContentItem::InputText { + text: "fallback".to_string(), + }], + error_text: None, + } + ); + service.shutdown().await.expect("shutdown service"); } } diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 52af3824d5..6edc44e297 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -230,7 +230,7 @@ async fn run_code_mode_turn_with_builder( } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn missing_process_host_returns_a_tool_error() -> Result<()> { +async fn missing_process_host_falls_back_to_in_process_code_mode() -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; @@ -244,16 +244,15 @@ async fn missing_process_host_returns_a_tool_error() -> Result<()> { .expect("code mode should be enabled"); }); let (_test, follow_up_mock) = - run_code_mode_turn_with_builder(&server, "Run code mode", "text('unreachable')", builder) + run_code_mode_turn_with_builder(&server, "Run code mode", "text('fallback')", builder) .await?; - let output = follow_up_mock - .single_request() - .custom_tool_call_output("call-1"); - assert!( - output["output"] - .as_str() - .is_some_and(|output| output.contains("failed to spawn code-mode host")) + assert_eq!( + text_item( + &custom_tool_output_items(&follow_up_mock.single_request(), "call-1"), + /*index*/ 1, + ), + "fallback" ); Ok(())