diff --git a/codex-rs/code-mode/src/grpc_session/mod.rs b/codex-rs/code-mode/src/grpc_session/mod.rs index 6d3be6eaed..a2c09487e6 100644 --- a/codex-rs/code-mode/src/grpc_session/mod.rs +++ b/codex-rs/code-mode/src/grpc_session/mod.rs @@ -94,6 +94,24 @@ impl GrpcCodeModeSessionProvider { Self::from_transport(SharedTransport::with_channel(channel)) } + /// Uses a caller-owned transport to control connection establishment and recovery. + pub fn with_transport(transport: T) -> Self + where + T: tower::Service< + tonic::codegen::http::Request, + Response = tonic::codegen::http::Response, + Error = tonic::Status, + > + Clone + + Send + + Sync + + 'static, + T::Future: Send + 'static, + { + Self::from_transport(SharedTransport::with_transport( + tower::util::BoxCloneSyncService::new(transport), + )) + } + fn from_transport(transport: SharedTransport) -> Self { Self { transport: Arc::new(transport), diff --git a/codex-rs/code-mode/src/grpc_session/transport.rs b/codex-rs/code-mode/src/grpc_session/transport.rs index f16a3fef88..286a122d5f 100644 --- a/codex-rs/code-mode/src/grpc_session/transport.rs +++ b/codex-rs/code-mode/src/grpc_session/transport.rs @@ -48,6 +48,15 @@ impl SharedTransport { ) } + pub(super) fn with_transport(transport: GrpcTransport) -> Self { + Self::Connected( + CodeModeHostClient::new(transport) + .accept_compressed(tonic::codec::CompressionEncoding::Zstd) + .max_decoding_message_size(MAX_FRAME_BYTES) + .max_encoding_message_size(MAX_FRAME_BYTES), + ) + } + pub(super) async fn client(&self) -> Result { let (endpoint, http_client_factory, client) = match self { Self::Url { diff --git a/codex-rs/core/src/code_mode_host.rs b/codex-rs/core/src/code_mode_host.rs index de328875cb..83c39a9a24 100644 --- a/codex-rs/core/src/code_mode_host.rs +++ b/codex-rs/core/src/code_mode_host.rs @@ -1,13 +1,18 @@ //! Owns the local gRPC code-mode client and host process. Sessions share a lazy -//! HTTP/2 channel; reconnecting spawns a fresh host, and dropping its I/O kills and -//! reaps the child. Session recovery and stale cell IDs use the shared gRPC client. +//! HTTP/2 connection. Only session admission can replace a dead host; other RPCs +//! use the current generation. Dropping its I/O kills and reaps the child. use std::io; use std::path::Path; use std::path::PathBuf; use std::pin::Pin; use std::process::Stdio; +use std::sync::Arc; +use std::sync::Mutex; use std::sync::OnceLock; +use std::sync::PoisonError; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; use std::time::Duration; @@ -29,12 +34,17 @@ use tokio::process::Child; use tokio::process::ChildStdin; use tokio::process::ChildStdout; use tokio::process::Command; +use tonic::Status; +use tonic::body::Body; +use tonic::transport::Channel; use tonic::transport::Endpoint; +use tower::ServiceExt; use tower::service_fn; pub(crate) struct ProcessOwnedGrpcCodeModeSessionProvider { host_program: PathBuf, provider: OnceLock, + current: Arc>>, } impl ProcessOwnedGrpcCodeModeSessionProvider { @@ -42,6 +52,7 @@ impl ProcessOwnedGrpcCodeModeSessionProvider { Self { host_program, provider: OnceLock::new(), + current: Arc::new(Mutex::new(None)), } } } @@ -63,21 +74,90 @@ impl CodeModeSessionProvider for ProcessOwnedGrpcCodeModeSessionProvider { Box::pin(async move { self.availability()?; let provider = self.provider.get_or_init(|| { - let host_program = self.host_program.clone(); - // The URI is an HTTP/2 authority only; the connector exclusively uses pipes. - // A live but unresponsive child can leave both pipes open. Let HTTP/2 - // detect that and drop the I/O, which reaps the child on reconnect. + let transport = Arc::new(HostTransport { + host_program: self.host_program.clone(), + current: self.current.clone(), + }); + GrpcCodeModeSessionProvider::with_transport(service_fn(move |request| { + let transport = transport.clone(); + async move { transport.send(request).await } + })) + }); + provider.create_session_with_limits(limits).await + }) + } +} + +// Every channel gets exactly one child. Tonic's internal reconnect loop cannot +// spawn a process, including if the child dies between channel selection and +// dispatch. Session admission selects a new generation during recovery. +struct HostTransport { + host_program: PathBuf, + current: Arc>>, +} + +struct HostChannel { + channel: Channel, + alive: Arc, +} + +impl HostTransport { + async fn send(&self, request: http::Request) -> Result, Status> { + self.channel(request.uri().path())? + .oneshot(request) + .await + .map_err(|error| { + // Channel errors are transport failures, not host RPC statuses. + // In particular a child exiting during handshake can surface as + // BrokenPipe, which generic tonic conversion calls Unknown and + // would prevent the existing bounded OpenSession retry policy. + let mut status = Status::unavailable("code-mode host transport failed"); + status.set_source(Arc::new(error)); + status + }) + } + + fn channel(&self, rpc_path: &str) -> Result { + Ok({ + let mut current = self.current.lock().unwrap_or_else(PoisonError::into_inner); + if let Some(host) = current + .as_ref() + .filter(|host| host.alive.load(Ordering::Acquire)) + { + host.channel.clone() + } else { + if rpc_path != "/codex.code_mode.v1.CodeModeHost/OpenSession" { + return Err(Status::unavailable("code-mode host is not running")); + } + let connection = HostConnection::spawn(&self.host_program) + .map_err(|error| Status::unavailable(error.to_string()))?; + let alive = connection.alive.clone(); + let connection = Arc::new(Mutex::new(Some(connection))); + // Detect a live but unresponsive child even while the connection is idle. let channel = Endpoint::from_static("http://code-mode-host") .http2_keep_alive_interval(Duration::from_secs(/*secs*/ 30)) .keep_alive_timeout(Duration::from_secs(/*secs*/ 10)) .keep_alive_while_idle(/*enabled*/ true) .connect_with_connector_lazy(service_fn(move |_| { - let host_program = host_program.clone(); - async move { HostConnection::spawn(&host_program).map(TokioIo::new) } + let connection = connection + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take(); + async move { + connection.map(TokioIo::new).ok_or_else(|| { + io::Error::new( + io::ErrorKind::NotConnected, + "code-mode host connection ended", + ) + }) + } })); - GrpcCodeModeSessionProvider::with_channel(channel) - }); - provider.create_session_with_limits(limits).await + *current = Some(HostChannel { + channel: channel.clone(), + alive, + }); + channel + } }) } } @@ -85,6 +165,7 @@ impl CodeModeSessionProvider for ProcessOwnedGrpcCodeModeSessionProvider { struct HostConnection { io: tokio::io::Join, child: Option, + alive: Arc, } impl HostConnection { @@ -119,6 +200,7 @@ impl HostConnection { Ok(Self { io: tokio::io::join(stdout, stdin), child: Some(child), + alive: Arc::new(AtomicBool::new(true)), }) } } @@ -153,6 +235,7 @@ impl AsyncWrite for HostConnection { impl Drop for HostConnection { fn drop(&mut self) { + self.alive.store(false, Ordering::Release); if let Some(mut child) = self.child.take() { let _ = child.start_kill(); if let Ok(runtime) = tokio::runtime::Handle::try_current() { diff --git a/codex-rs/core/src/code_mode_host_lifecycle_tests.rs b/codex-rs/core/src/code_mode_host_lifecycle_tests.rs new file mode 100644 index 0000000000..fff3e7ebc0 --- /dev/null +++ b/codex-rs/core/src/code_mode_host_lifecycle_tests.rs @@ -0,0 +1,155 @@ +#![cfg(test)] +//! Deterministic child failures around channel selection and session recovery. + +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use codex_code_mode::CodeModeSessionProvider; +use codex_code_mode::FunctionCallOutputContentItem; +use codex_code_mode::GrpcCodeModeSessionProvider; +use core_test_support::code_mode_host::CodeModeHostRecorder; +use futures::future::try_join_all; +use pretty_assertions::assert_eq; +use tokio::sync::Notify; +use tokio::sync::Semaphore; +use tokio::time::timeout; +use tonic::Status; +use tonic::body::Body; +use tower::ServiceExt; +use tower::service_fn; + +use super::super::HostTransport; +use super::super::ProcessOwnedGrpcCodeModeSessionProvider; +use super::TEST_TIMEOUT; +use super::execute; + +/// Failed initial HTTP/2 setup can retry admission, with exactly one replacement. +#[tokio::test] +async fn child_exit_during_initial_connection_retries_admission() -> Result<()> { + let recorder = CodeModeHostRecorder::new()?; + let directory = tempfile::tempdir()?; + let wrapper = directory.path().join("exit-first"); + let attempts = directory.path().join("attempts"); + let marker = directory.path().join("failed-once"); + let program = recorder.program(); + let program = shlex::try_quote(program.to_str().context("host program")?)?; + let quoted_attempts = shlex::try_quote(attempts.to_str().context("attempts")?)?; + let marker = shlex::try_quote(marker.to_str().context("marker")?)?; + codex_utils_cargo_bin::write_executable( + &wrapper, + &format!( + "#!/bin/sh\nprintf x >> {quoted_attempts}\nif [ ! -e {marker} ]; then : > {marker}; exit 0; fi\nexec {program} \"$@\"\n" + ), + )?; + let provider = ProcessOwnedGrpcCodeModeSessionProvider::with_host_program(wrapper); + let session = timeout(TEST_TIMEOUT, provider.create_session()) + .await? + .map_err(anyhow::Error::msg)?; + assert_eq!( + execute(&session, "text('ready');").await?, + vec![FunctionCallOutputContentItem::InputText { + text: "ready".to_string() + }] + ); + assert_eq!(std::fs::read_to_string(attempts)?, "xx"); + assert_eq!(recorder.invocations()?.len(), 1); + session.shutdown().await.map_err(anyhow::Error::msg)?; + Ok(()) +} + +/// A selected generation cannot reconnect; concurrent sessions recover without replay. +#[tokio::test] +async fn child_exit_after_channel_selection_does_not_replay_execution() -> Result<()> { + let recorder = CodeModeHostRecorder::new()?; + let transport = Arc::new(HostTransport { + host_program: recorder.program(), + current: Arc::new(Mutex::new(None)), + }); + let armed = Arc::new(AtomicBool::new(false)); + let selected = Arc::new(Notify::new()); + let dispatch = Arc::new(Semaphore::new(0)); + let provider = GrpcCodeModeSessionProvider::with_transport(service_fn({ + let armed = armed.clone(); + let selected = selected.clone(); + let dispatch = dispatch.clone(); + move |request: http::Request| { + let transport = transport.clone(); + let armed = armed.clone(); + let selected = selected.clone(); + let dispatch = dispatch.clone(); + async move { + // Exercise the production selector, pausing exactly between + // selecting a live generation and dispatching to that channel. + let channel = transport.channel(request.uri().path())?; + if request.uri().path() == "/codex.code_mode.v1.CodeModeHost/Execute" + && armed.swap(false, Ordering::SeqCst) + { + selected.notify_one(); + dispatch.acquire().await.unwrap().forget(); + } + channel + .oneshot(request) + .await + .map_err(|error| Status::from_error(Box::new(error))) + } + } + })); + let sessions = timeout( + TEST_TIMEOUT, + try_join_all((0..4).map(|_| provider.create_session())), + ) + .await? + .map_err(anyhow::Error::msg)?; + armed.store(true, Ordering::SeqCst); + let pending = tokio::spawn({ + let session = sessions[0].clone(); + async move { execute(&session, "store('replayed', true);").await } + }); + timeout(TEST_TIMEOUT, selected.notified()).await?; + let pid = recorder.pid()?; + assert_eq!(unsafe { libc::kill(pid as libc::pid_t, libc::SIGKILL) }, 0); + timeout(TEST_TIMEOUT, async { + while unsafe { + libc::kill(pid as libc::pid_t, /*sig*/ 0) + } == 0 + { + tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await; + } + }) + .await + .context("child was not reaped")?; + assert_eq!(recorder.invocations()?.len(), 1); + dispatch.add_permits(1); + assert!(timeout(TEST_TIMEOUT, pending).await??.is_err()); + assert_eq!(recorder.invocations()?.len(), 1); + + let recovery = try_join_all(sessions.iter().map(|session| async move { + let error = execute(session, "store('replayed', true);") + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("Stored values and running cells were lost"), + "{error}" + ); + assert_eq!( + execute(session, "text(load('replayed') ?? 'clean');").await?, + vec![FunctionCallOutputContentItem::InputText { + text: "clean".to_string() + }] + ); + Ok::<_, anyhow::Error>(()) + })); + timeout(TEST_TIMEOUT, recovery).await??; + assert_eq!(recorder.invocations()?.len(), 2); + for session in sessions { + session.shutdown().await.map_err(anyhow::Error::msg)?; + } + Ok(()) +} diff --git a/codex-rs/core/src/code_mode_host_tests.rs b/codex-rs/core/src/code_mode_host_tests.rs index b850c3c947..2088380bb7 100644 --- a/codex-rs/core/src/code_mode_host_tests.rs +++ b/codex-rs/core/src/code_mode_host_tests.rs @@ -21,6 +21,10 @@ use super::ProcessOwnedGrpcCodeModeSessionProvider; const TEST_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 20); +#[cfg(unix)] +#[path = "code_mode_host_lifecycle_tests.rs"] +mod lifecycle; + async fn execute( session: &Arc, source: &str, @@ -92,6 +96,23 @@ async fn grpc_stdio_sessions_share_a_host_and_isolate_state() -> Result<()> { vec!["--listen grpc+stdio://".to_string()], "both sessions must share one gRPC host process" ); + drop(first); + drop(second); + drop(provider); + #[cfg(unix)] + { + let pid = recorder.pid()?; + timeout(TEST_TIMEOUT, async { + while unsafe { + libc::kill(pid as libc::pid_t, /*sig*/ 0) + } == 0 + { + tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await; + } + }) + .await + .context("host was not reaped after provider teardown")?; + } Ok(()) } diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 534d831d13..92a06da783 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -3,6 +3,10 @@ #[path = "code_mode_dynamic_tools_tests.rs"] mod dynamic_tools; +#[cfg(unix)] +#[path = "code_mode_host_recovery_tests.rs"] +mod host_recovery; + use anyhow::Context; use anyhow::Result; use base64::Engine; diff --git a/codex-rs/core/tests/suite/code_mode_host_recovery_tests.rs b/codex-rs/core/tests/suite/code_mode_host_recovery_tests.rs new file mode 100644 index 0000000000..4fe19bb7b1 --- /dev/null +++ b/codex-rs/core/tests/suite/code_mode_host_recovery_tests.rs @@ -0,0 +1,156 @@ +//! Host loss must reach the model without replaying completed external effects. + +use super::*; +use codex_extension_api::ToolCallSource; +use core_test_support::test_codex::TestCodexHarness; +use pretty_assertions::assert_eq; + +struct PauseCompletedCommand { + gate: Mutex, oneshot::Receiver<()>)>>, +} + +impl ToolLifecycleContributor for PauseCompletedCommand { + fn on_tool_finish<'a>(&'a self, input: ToolFinishInput<'a>) -> ToolLifecycleFuture<'a> { + Box::pin(async move { + if input.tool_name.name != "exec_command" + || !matches!(input.source, ToolCallSource::CodeMode { .. }) + { + return; + } + let gate = self.gate.lock().unwrap().take(); + if let Some((reached, release)) = gate { + let _ = reached.send(input.outcome); + // Keep the host executing until the test kills it: the external + // command completed, but its result has not reached JavaScript. + let _ = release.await; + } + }) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn grpc_host_loss_reports_reset_without_replaying_side_effects() -> Result<()> { + skip_if_no_network!(Ok(())); + + let recorder = CodeModeHostRecorder::new()?; + let (reached_tx, reached_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + let mut extensions = ExtensionRegistryBuilder::::new(); + extensions.tool_lifecycle_contributor(Arc::new(PauseCompletedCommand { + gate: Mutex::new(Some((reached_tx, release_rx))), + })); + let builder = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_code_mode_host_program(recorder.program()) + .with_config(|config| { + config.features.enable(Feature::CodeModeOnly).unwrap(); + config.features.enable(Feature::CodeModeHostGrpc).unwrap(); + }); + let harness = TestCodexHarness::with_auto_env_builder(builder).await?; + let test = harness.test(); + let server = harness.server(); + let (append, read) = match core_test_support::test_target_os() { + core_test_support::TestTargetOs::Windows => ( + "[IO.File]::AppendAllText((Join-Path (Get-Location) 'effects.txt'), 'x')", + "[Console]::Write([IO.File]::ReadAllText((Join-Path (Get-Location) 'effects.txt')))", + ), + core_test_support::TestTargetOs::Linux | core_test_support::TestTargetOs::MacOs => { + ("printf x >> effects.txt", "cat effects.txt") + } + }; + let interrupted = format!( + r#"// @exec: {{"yield_time_ms": 60000}} +store("before-crash", true); +await tools.exec_command({{ cmd: {append:?} }}); +text("unexpected success");"#, + ); + let fresh = format!( + r#"const result = await tools.exec_command({{ cmd: {read:?} }}); +text(JSON.stringify({{ previous: load("before-crash") ?? null, output: result.output }}));"#, + ); + let mut requests = Vec::new(); + for (call_id, source) in [ + ("crashed", &interrupted), + ("reset", &interrupted), + ("fresh", &fresh), + ] { + requests.push( + responses::mount_sse_once( + server, + sse(vec![ + ev_response_created(call_id), + ev_custom_tool_call(call_id, "exec", source), + ev_completed(call_id), + ]), + ) + .await, + ); + } + let completion = responses::mount_sse_once( + server, + sse(vec![ + ev_assistant_message("done", "recovered"), + ev_completed("done"), + ]), + ) + .await; + + let crash = async { + assert_eq!( + reached_rx.await?, + ToolCallOutcome::Completed { success: true } + ); + assert_eq!(harness.read_file_text("effects.txt").await?, "x"); + assert_eq!( + recorder.invocations()?, + vec!["--listen grpc+stdio://".to_string()] + ); + let pid = recorder.pid()?; + assert_eq!(unsafe { libc::kill(pid as libc::pid_t, libc::SIGKILL) }, 0); + let _ = release_tx.send(()); + Ok::<_, anyhow::Error>(()) + }; + let (turn_result, crash_result) = + tokio::time::timeout(Duration::from_secs(/*secs*/ 30), async { + tokio::join!( + test.submit_turn("Run a command and recover if the host fails"), + crash + ) + }) + .await + .context("host-loss recovery stalled")?; + crash_result?; + turn_result?; + + // The inference requests must expose the failure and the next admission's + // reset warning. The API serializes these errors as text, without a success flag. + let (failure, _) = + custom_tool_output_body_and_success(&requests[1].single_request(), "crashed"); + assert!( + failure.contains("gRPC code-mode") || failure.contains("exec runtime ended unexpectedly"), + "{failure}" + ); + let (reset, _) = custom_tool_output_body_and_success(&requests[2].single_request(), "reset"); + assert!( + reset.contains("The exec runtime was replaced. Stored values and running cells were lost."), + "{reset}" + ); + assert!( + reset.contains("Prior external tool operations may have completed; inspect their results before retrying."), + "{reset}" + ); + let (output, success) = + custom_tool_output_body_and_success(&completion.single_request(), "fresh"); + assert_ne!(success, Some(false), "{output}"); + assert_eq!( + serde_json::from_str::(&output)?, + serde_json::json!({ "previous": null, "output": "x" }) + ); + assert_eq!(harness.read_file_text("effects.txt").await?, "x"); + assert_eq!( + recorder.invocations()?, + vec!["--listen grpc+stdio://".to_string(); 2] + ); + test.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/rmcp-client/tests/streamable_http_remote.rs b/codex-rs/rmcp-client/tests/streamable_http_remote.rs index 7df871976b..88462ff5cb 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_remote.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_remote.rs @@ -76,7 +76,7 @@ async fn stalled_handshake_unblocks_remote_executor( let http_server = MockServer::start().await; Mock::given(method("POST")) .and(body_partial_json(json!({ "method": expected_method }))) - .respond_with(ResponseTemplate::new(200).set_delay(Duration::from_secs(10))) + .respond_with(ResponseTemplate::new(200).set_delay(Duration::from_secs(60))) .expect(1) .mount(&http_server) .await; @@ -102,7 +102,9 @@ async fn stalled_handshake_unblocks_remote_executor( let error = client .initialize( params, - Some(Duration::from_millis(500)), + // Allow the request to reach the remote HTTP server on loaded CI + // workers before testing cancellation of its stalled response. + Some(Duration::from_secs(5)), Box::new(|_, _| { async { Ok(ElicitationResponse {