Restrict code-mode host replacement to session admission (#57782)

## Why

Automatic gRPC reconnection could spawn a replacement host while dispatching an RPC for an existing session. Host loss must remain visible so recovery can report lost state without replaying execution or completed external effects.

## What changed

- Bind each local host channel to one child process and allow only `OpenSession` to replace a dead host.
- Add caller-owned transport support to `GrpcCodeModeSessionProvider` and map local transport failures to `Unavailable` so initial connection failures use the existing bounded admission retries.

## Testing

Add regression coverage for child exit during initial connection and between channel selection and dispatch, concurrent session recovery, and child reaping after teardown. An integration test kills the host after a command completes and verifies that recovery reports the reset, clears stored state, and leaves the command's external effect unchanged.

Increase the remote MCP stalled-handshake test's initialization timeout and mock response delay to let requests reach the server on loaded CI workers.

GitOrigin-RevId: 4ffd810e4cfcfb7a56e606816aa867bc206a09b6

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<T>(transport: T) -> Self
+    where
+        T: tower::Service<
+                tonic::codegen::http::Request<tonic::body::Body>,
+                Response = tonic::codegen::http::Response<tonic::body::Body>,
+                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),

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<GrpcClient, String> {
         let (endpoint, http_client_factory, client) = match self {
             Self::Url {

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<GrpcCodeModeSessionProvider>,
+    current: Arc<Mutex<Option<HostChannel>>>,
 }
 
 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<Mutex<Option<HostChannel>>>,
+}
+
+struct HostChannel {
+    channel: Channel,
+    alive: Arc<AtomicBool>,
+}
+
+impl HostTransport {
+    async fn send(&self, request: http::Request<Body>) -> Result<http::Response<Body>, 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<Channel, Status> {
+        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<ChildStdout, ChildStdin>,
     child: Option<Child>,
+    alive: Arc<AtomicBool>,
 }
 
 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() {

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<Body>| {
+            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(())
+}

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<dyn CodeModeSession>,
     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(())
 }
 

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;

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<Option<(oneshot::Sender<ToolCallOutcome>, 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::<Config>::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::<Value>(&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(())
+}

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 {