Reconnect gRPC code-mode sessions after host restarts (#38257)

## What changed

- Reopen a cached code-mode session when its gRPC host stops, while
  serializing concurrent reconnection attempts and coordinating shutdown.
- Scope cell IDs to the new host generation so callbacks remain consistent
  and stale `wait` or `terminate` requests are rejected.
- Accept both `unix://` and `unix:` endpoints for gRPC hosts on Unix systems.

## Testing

- Cover host restart recovery, concurrent execution after reconnection,
  generation-aware callbacks and cell operations, stale cell rejection, and
  Unix socket execution.

GitOrigin-RevId: 548e168fdcef7f7d54bd32262e614886bf7bdd32

codex-rs/code-mode/src/grpc_session/generation.rs

@@ -0,0 +1,124 @@
+use std::sync::Arc;
+
+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::RuntimeResponse;
+use codex_code_mode_protocol::StartedCell;
+use codex_code_mode_protocol::ToolInvocationFuture;
+use codex_code_mode_protocol::WaitOutcome;
+use tokio_util::sync::CancellationToken;
+
+pub(super) struct GenerationDelegate {
+    pub(super) delegate: Arc<dyn CodeModeSessionDelegate>,
+    pub(super) generation: u64,
+}
+
+impl CodeModeSessionDelegate for GenerationDelegate {
+    fn invoke_tool<'a>(
+        &'a self,
+        mut invocation: CodeModeNestedToolCall,
+        cancellation_token: CancellationToken,
+    ) -> ToolInvocationFuture<'a> {
+        invocation.cell_id = public_cell_id(self.generation, &invocation.cell_id);
+        self.delegate.invoke_tool(invocation, cancellation_token)
+    }
+
+    fn notify<'a>(
+        &'a self,
+        call_id: String,
+        cell_id: CellId,
+        text: String,
+        cancellation_token: CancellationToken,
+    ) -> NotificationFuture<'a> {
+        self.delegate.notify(
+            call_id,
+            public_cell_id(self.generation, &cell_id),
+            text,
+            cancellation_token,
+        )
+    }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        self.delegate
+            .cell_closed(&public_cell_id(self.generation, cell_id));
+    }
+}
+
+fn public_cell_id(generation: u64, cell_id: &CellId) -> CellId {
+    if generation == 1 {
+        cell_id.clone()
+    } else {
+        CellId::new(format!("g{generation}:{cell_id}"))
+    }
+}
+
+pub(super) fn remote_cell_id(generation: u64, cell_id: &CellId) -> Result<CellId, String> {
+    if generation == 1 {
+        return Ok(cell_id.clone());
+    }
+
+    let prefix = format!("g{generation}:");
+    cell_id
+        .as_str()
+        .strip_prefix(&prefix)
+        .map(|cell_id| CellId::new(cell_id.to_string()))
+        .ok_or_else(|| "cell belongs to a stale code-mode host generation".to_string())
+}
+
+pub(super) fn public_started_cell(generation: u64, started: StartedCell) -> StartedCell {
+    if generation == 1 {
+        return started;
+    }
+    let cell_id = public_cell_id(generation, &started.cell_id);
+    StartedCell::from_future(cell_id, async move {
+        started
+            .initial_response()
+            .await
+            .map(|response| public_runtime_response(generation, response))
+    })
+}
+
+fn public_runtime_response(generation: u64, response: RuntimeResponse) -> RuntimeResponse {
+    match response {
+        RuntimeResponse::Yielded {
+            cell_id,
+            content_items,
+        } => RuntimeResponse::Yielded {
+            cell_id: public_cell_id(generation, &cell_id),
+            content_items,
+        },
+        RuntimeResponse::Terminated {
+            cell_id,
+            content_items,
+        } => RuntimeResponse::Terminated {
+            cell_id: public_cell_id(generation, &cell_id),
+            content_items,
+        },
+        RuntimeResponse::Result {
+            cell_id,
+            content_items,
+            error_text,
+        } => RuntimeResponse::Result {
+            cell_id: public_cell_id(generation, &cell_id),
+            content_items,
+            error_text,
+        },
+    }
+}
+
+pub(super) fn public_wait_outcome(generation: u64, outcome: WaitOutcome) -> WaitOutcome {
+    match outcome {
+        WaitOutcome::LiveCell(response) => {
+            WaitOutcome::LiveCell(public_runtime_response(generation, response))
+        }
+        WaitOutcome::MissingCell(response) => {
+            WaitOutcome::MissingCell(public_runtime_response(generation, response))
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "generation_tests.rs"]
+mod tests;

codex-rs/code-mode/src/grpc_session/generation_tests.rs

@@ -0,0 +1,230 @@
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeNestedToolCall;
+use codex_code_mode_protocol::CodeModeSessionDelegate;
+use codex_code_mode_protocol::CodeModeToolKind;
+use codex_code_mode_protocol::FunctionCallOutputContentItem;
+use codex_code_mode_protocol::NotificationFuture;
+use codex_code_mode_protocol::RuntimeResponse;
+use codex_code_mode_protocol::StartedCell;
+use codex_code_mode_protocol::ToolInvocationFuture;
+use codex_code_mode_protocol::WaitOutcome;
+use codex_protocol::ToolName;
+use pretty_assertions::assert_eq;
+use serde_json::json;
+use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
+
+use super::GenerationDelegate;
+use super::public_cell_id;
+use super::public_started_cell;
+use super::public_wait_outcome;
+use super::remote_cell_id;
+
+#[derive(Default)]
+struct RecordingDelegate {
+    calls: Mutex<Vec<CodeModeNestedToolCall>>,
+    notifications: Mutex<Vec<(String, CellId, String)>>,
+    closed: Mutex<Vec<CellId>>,
+}
+
+impl CodeModeSessionDelegate for RecordingDelegate {
+    fn invoke_tool<'a>(
+        &'a self,
+        invocation: CodeModeNestedToolCall,
+        _cancellation_token: CancellationToken,
+    ) -> ToolInvocationFuture<'a> {
+        self.calls.lock().expect("calls lock").push(invocation);
+        Box::pin(async { Ok(json!({ "ok": true })) })
+    }
+
+    fn notify<'a>(
+        &'a self,
+        call_id: String,
+        cell_id: CellId,
+        text: String,
+        _cancellation_token: CancellationToken,
+    ) -> NotificationFuture<'a> {
+        self.notifications
+            .lock()
+            .expect("notifications lock")
+            .push((call_id, cell_id, text));
+        Box::pin(async { Ok(()) })
+    }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        self.closed
+            .lock()
+            .expect("closed cells lock")
+            .push(cell_id.clone());
+    }
+}
+
+#[test]
+fn first_generation_preserves_existing_cell_ids() {
+    let cell_id = CellId::new("42".to_string());
+
+    assert_eq!(public_cell_id(/*generation*/ 1, &cell_id), cell_id);
+    assert_eq!(remote_cell_id(/*generation*/ 1, &cell_id), Ok(cell_id));
+}
+
+#[test]
+fn first_generation_preserves_opaque_cell_ids_that_resemble_generation_prefixes() {
+    for value in ["graphics:1", "g2:42"] {
+        let cell_id = CellId::new(value.to_string());
+
+        assert_eq!(public_cell_id(/*generation*/ 1, &cell_id), cell_id);
+        assert_eq!(remote_cell_id(/*generation*/ 1, &cell_id), Ok(cell_id));
+    }
+}
+
+#[test]
+fn later_generations_prefix_public_ids_and_strip_wire_ids() {
+    let wire_id = CellId::new("42".to_string());
+    let public_id = CellId::new("g2:42".to_string());
+
+    assert_eq!(public_cell_id(/*generation*/ 2, &wire_id), public_id);
+    assert_eq!(remote_cell_id(/*generation*/ 2, &public_id), Ok(wire_id));
+}
+
+#[test]
+fn stale_generation_ids_are_rejected_after_reconnection() {
+    for cell_id in [
+        "42".to_string(),
+        "g1:42".to_string(),
+        "g3:42".to_string(),
+        "x".repeat(10_000),
+    ] {
+        assert_eq!(
+            remote_cell_id(/*generation*/ 2, &CellId::new(cell_id)),
+            Err("cell belongs to a stale code-mode host generation".to_string())
+        );
+    }
+}
+
+#[tokio::test]
+async fn reconnect_maps_every_delegate_callback_to_its_generation() {
+    let recording = Arc::new(RecordingDelegate::default());
+    let delegate = GenerationDelegate {
+        delegate: recording.clone(),
+        generation: 2,
+    };
+    let wire_id = CellId::new("42".to_string());
+    let public_id = CellId::new("g2:42".to_string());
+    let invocation = CodeModeNestedToolCall {
+        cell_id: wire_id.clone(),
+        runtime_tool_call_id: "runtime-call".to_string(),
+        tool_name: ToolName::plain("echo"),
+        tool_kind: CodeModeToolKind::Function,
+        input: Some(json!({ "value": true })),
+    };
+
+    assert_eq!(
+        delegate
+            .invoke_tool(invocation.clone(), CancellationToken::new())
+            .await,
+        Ok(json!({ "ok": true }))
+    );
+    assert_eq!(
+        delegate
+            .notify(
+                "outer-call".to_string(),
+                wire_id.clone(),
+                "notice".to_string(),
+                CancellationToken::new(),
+            )
+            .await,
+        Ok(())
+    );
+    delegate.cell_closed(&wire_id);
+
+    assert_eq!(
+        *recording.calls.lock().expect("calls lock"),
+        vec![CodeModeNestedToolCall {
+            cell_id: public_id.clone(),
+            ..invocation
+        }]
+    );
+    assert_eq!(
+        *recording.notifications.lock().expect("notifications lock"),
+        vec![(
+            "outer-call".to_string(),
+            public_id.clone(),
+            "notice".to_string()
+        )]
+    );
+    assert_eq!(
+        *recording.closed.lock().expect("closed cells lock"),
+        vec![public_id]
+    );
+}
+
+#[tokio::test]
+async fn reconnected_execution_maps_started_and_initial_response_ids() {
+    let (response_tx, response_rx) = oneshot::channel();
+    let wire_id = CellId::new("42".to_string());
+    let public_id = CellId::new("g2:42".to_string());
+    let claimed = Arc::new(AtomicBool::new(false));
+    let initial_response_claimed = Arc::clone(&claimed);
+    let response = RuntimeResponse::Result {
+        cell_id: wire_id.clone(),
+        content_items: vec![FunctionCallOutputContentItem::InputText {
+            text: "result".to_string(),
+        }],
+        error_text: None,
+    };
+    let started = StartedCell::from_future(wire_id, async move {
+        initial_response_claimed.store(true, Ordering::Release);
+        response_rx.await.expect("receive initial response")
+    });
+    let started = public_started_cell(/*generation*/ 2, started);
+
+    assert_eq!(started.cell_id, public_id);
+    tokio::task::yield_now().await;
+    assert!(!claimed.load(Ordering::Acquire));
+    response_tx
+        .send(Ok(response))
+        .expect("send initial response");
+    assert_eq!(
+        started.initial_response().await,
+        Ok(RuntimeResponse::Result {
+            cell_id: public_id,
+            content_items: vec![FunctionCallOutputContentItem::InputText {
+                text: "result".to_string(),
+            }],
+            error_text: None,
+        })
+    );
+}
+
+#[test]
+fn reconnected_wait_maps_live_and_missing_outcomes() {
+    let public_id = CellId::new("g2:42".to_string());
+    let yielded = RuntimeResponse::Yielded {
+        cell_id: CellId::new("42".to_string()),
+        content_items: Vec::new(),
+    };
+    let terminated = RuntimeResponse::Terminated {
+        cell_id: CellId::new("42".to_string()),
+        content_items: Vec::new(),
+    };
+
+    assert_eq!(
+        public_wait_outcome(/*generation*/ 2, WaitOutcome::LiveCell(yielded)),
+        WaitOutcome::LiveCell(RuntimeResponse::Yielded {
+            cell_id: public_id.clone(),
+            content_items: Vec::new(),
+        })
+    );
+    assert_eq!(
+        public_wait_outcome(/*generation*/ 2, WaitOutcome::MissingCell(terminated)),
+        WaitOutcome::MissingCell(RuntimeResponse::Terminated {
+            cell_id: public_id,
+            content_items: Vec::new(),
+        })
+    );
+}

codex-rs/code-mode/src/grpc_session/mod.rs

@@ -39,12 +39,16 @@ mod callbacks;
 mod completion;
 mod conversion;
 mod deadline;
+mod generation;
 mod operations;
+mod reconnect;
 mod state;
 mod transport;
 
 type GrpcClient = CodeModeHostClient<GrpcTransport>;
 
+const SHUTDOWN_ERROR: &str = "code mode session is shutting down";
+
 /// Creates code-mode sessions over an HTTP/2 gRPC connection.
 #[derive(Clone)]
 pub struct GrpcCodeModeSessionProvider {
@@ -52,7 +56,7 @@ pub struct GrpcCodeModeSessionProvider {
 }
 
 impl GrpcCodeModeSessionProvider {
-    /// Connects lazily to an `http://` or `https://` gRPC endpoint.
+    /// Connects lazily to an `http://`, `https://`, or `unix://` gRPC endpoint.
     pub fn new(endpoint: impl Into<String>) -> Self {
         Self::with_http_client_factory(
             endpoint,
@@ -166,9 +170,13 @@ impl CodeModeSessionProvider for GrpcCodeModeSessionProvider {
         limits: CodeModeSessionCellExecutionLimits,
     ) -> CodeModeSessionProviderFuture<'a> {
         Box::pin(async move {
-            self.open_binding(delegate, limits)
-                .await
-                .map(|session| session as _)
+            let session = Arc::new(reconnect::ReconnectableSession::new(
+                self.clone(),
+                delegate,
+                limits,
+            ));
+            session.initialize().await?;
+            Ok(session as _)
         })
     }
 }
@@ -246,7 +254,7 @@ impl SessionInner {
             .unwrap_or_else(PoisonError::into_inner)
             .require_open()?;
         if self.shutdown_requested.load(Ordering::Acquire) {
-            return Err("code mode session is shutting down".to_string());
+            return Err(SHUTDOWN_ERROR.to_string());
         }
         Ok(())
     }

codex-rs/code-mode/src/grpc_session/reconnect.rs

@@ -0,0 +1,237 @@
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::PoisonError;
+use std::sync::atomic::AtomicU64;
+use std::sync::atomic::Ordering;
+
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeSession;
+use codex_code_mode_protocol::CodeModeSessionCellExecutionLimits;
+use codex_code_mode_protocol::CodeModeSessionDelegate;
+use codex_code_mode_protocol::CodeModeSessionResultFuture;
+use codex_code_mode_protocol::ExecuteRequest;
+use codex_code_mode_protocol::StartedCell;
+use codex_code_mode_protocol::WaitOutcome;
+use codex_code_mode_protocol::WaitRequest;
+use tokio::sync::Semaphore;
+use tokio::sync::watch;
+use tokio_util::sync::CancellationToken;
+
+use super::GrpcCodeModeSession;
+use super::GrpcCodeModeSessionProvider;
+use super::SHUTDOWN_ERROR;
+use super::generation;
+use super::generation::GenerationDelegate;
+use crate::remote_session::ShutdownResultReceiver;
+use crate::remote_session::wait_for_watch;
+
+pub(super) struct ReconnectableSession {
+    inner: Arc<ReconnectInner>,
+}
+
+struct ReconnectInner {
+    provider: GrpcCodeModeSessionProvider,
+    delegate: Arc<dyn CodeModeSessionDelegate>,
+    limits: CodeModeSessionCellExecutionLimits,
+    binding: Mutex<Option<SessionBinding>>,
+    opening_permit: Semaphore,
+    next_generation: AtomicU64,
+    shutdown_requested: CancellationToken,
+    shutdown_result: Mutex<Option<ShutdownResultReceiver>>,
+}
+
+#[derive(Clone)]
+struct SessionBinding {
+    session: Arc<GrpcCodeModeSession>,
+    generation: u64,
+}
+
+impl ReconnectableSession {
+    pub(super) fn new(
+        provider: GrpcCodeModeSessionProvider,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+        limits: CodeModeSessionCellExecutionLimits,
+    ) -> Self {
+        Self {
+            inner: Arc::new(ReconnectInner {
+                provider,
+                delegate,
+                limits,
+                binding: Mutex::new(None),
+                opening_permit: Semaphore::new(/*permits*/ 1),
+                next_generation: AtomicU64::new(1),
+                shutdown_requested: CancellationToken::new(),
+                shutdown_result: Mutex::new(None),
+            }),
+        }
+    }
+
+    pub(super) async fn initialize(&self) -> Result<(), String> {
+        self.inner.get_or_open_binding().await.map(|_| ())
+    }
+}
+
+impl CodeModeSession for ReconnectableSession {
+    fn execute<'a>(
+        &'a self,
+        request: ExecuteRequest,
+    ) -> CodeModeSessionResultFuture<'a, StartedCell> {
+        Box::pin(async move {
+            let binding = self.inner.get_or_open_binding().await?;
+            let started = binding.session.execute(request).await?;
+            Ok(generation::public_started_cell(binding.generation, started))
+        })
+    }
+
+    fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
+        Box::pin(async move {
+            let binding = self.inner.get_or_open_binding().await?;
+            let request = WaitRequest {
+                cell_id: generation::remote_cell_id(binding.generation, &request.cell_id)?,
+                yield_time_ms: request.yield_time_ms,
+            };
+            let outcome = binding.session.wait(request).await?;
+            Ok(generation::public_wait_outcome(binding.generation, outcome))
+        })
+    }
+
+    fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
+        Box::pin(async move {
+            let binding = self.inner.get_or_open_binding().await?;
+            let cell_id = generation::remote_cell_id(binding.generation, &cell_id)?;
+            let outcome = binding.session.terminate(cell_id).await?;
+            Ok(generation::public_wait_outcome(binding.generation, outcome))
+        })
+    }
+
+    fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> {
+        Box::pin(wait_for_watch(self.inner.request_shutdown()))
+    }
+}
+
+impl Drop for ReconnectableSession {
+    fn drop(&mut self) {
+        if tokio::runtime::Handle::try_current().is_ok() {
+            self.inner.request_shutdown();
+        }
+    }
+}
+
+impl ReconnectInner {
+    async fn get_or_open_binding(&self) -> Result<SessionBinding, String> {
+        if self.shutdown_requested.is_cancelled() {
+            return Err(SHUTDOWN_ERROR.to_string());
+        }
+        if let Some(binding) = self.live_binding() {
+            return Ok(binding);
+        }
+
+        let _opening_permit = tokio::select! {
+            biased;
+            _ = self.shutdown_requested.cancelled() => {
+                return Err(SHUTDOWN_ERROR.to_string());
+            }
+            permit = self.opening_permit.acquire() => permit
+                .map_err(|_| "gRPC code-mode session opening coordinator closed".to_string())?,
+        };
+        if self.shutdown_requested.is_cancelled() {
+            return Err(SHUTDOWN_ERROR.to_string());
+        }
+        if let Some(binding) = self.live_binding() {
+            return Ok(binding);
+        }
+
+        let previous_binding = self
+            .binding
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .clone();
+        if let Some(binding) = previous_binding {
+            wait_for_watch(binding.session.inner.request_shutdown()).await?;
+        }
+
+        let generation = self.next_generation.fetch_add(1, Ordering::Relaxed);
+        let delegate = Arc::new(GenerationDelegate {
+            delegate: Arc::clone(&self.delegate),
+            generation,
+        });
+        let session = tokio::select! {
+            biased;
+            _ = self.shutdown_requested.cancelled() => {
+                return Err(SHUTDOWN_ERROR.to_string());
+            }
+            session = self.provider.open_binding(delegate, self.limits.clone()) => session?,
+        };
+        let binding = SessionBinding {
+            session,
+            generation,
+        };
+        let published = {
+            let mut current = self.binding.lock().unwrap_or_else(PoisonError::into_inner);
+            if self.shutdown_requested.is_cancelled() {
+                false
+            } else {
+                *current = Some(binding.clone());
+                true
+            }
+        };
+        if !published {
+            let _ = wait_for_watch(binding.session.inner.request_shutdown()).await;
+            return Err(SHUTDOWN_ERROR.to_string());
+        }
+        Ok(binding)
+    }
+
+    fn live_binding(&self) -> Option<SessionBinding> {
+        self.binding
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .as_ref()
+            .filter(|binding| !binding.session.inner.stopped.is_cancelled())
+            .cloned()
+    }
+
+    fn request_shutdown(self: &Arc<Self>) -> ShutdownResultReceiver {
+        {
+            let binding = self.binding.lock().unwrap_or_else(PoisonError::into_inner);
+            self.shutdown_requested.cancel();
+            if let Some(binding) = binding.as_ref() {
+                binding.session.inner.request_shutdown();
+            }
+        }
+        let mut result = self
+            .shutdown_result
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner);
+        if let Some(receiver) = result.as_ref() {
+            return receiver.clone();
+        }
+
+        let (sender, receiver) = watch::channel(None);
+        *result = Some(receiver.clone());
+        let inner = Arc::clone(self);
+        tokio::spawn(async move {
+            let opening_permit = match inner.opening_permit.acquire().await {
+                Ok(permit) => permit,
+                Err(_) => {
+                    sender.send_replace(Some(Err(
+                        "gRPC code-mode session opening coordinator closed".to_string(),
+                    )));
+                    return;
+                }
+            };
+            let binding = inner
+                .binding
+                .lock()
+                .unwrap_or_else(PoisonError::into_inner)
+                .take();
+            drop(opening_permit);
+            let result = match binding {
+                Some(binding) => wait_for_watch(binding.session.inner.request_shutdown()).await,
+                None => Ok(()),
+            };
+            sender.send_replace(Some(result));
+        });
+        receiver
+    }
+}

codex-rs/code-mode/src/grpc_session/transport.rs

@@ -10,6 +10,7 @@ use tonic::codegen::http::Request;
 use tonic::codegen::http::Response;
 use tonic::codegen::http::Uri;
 use tonic::transport::Channel;
+use tonic::transport::Endpoint;
 use tower::ServiceExt;
 use tower::service_fn;
 use tower::util::BoxCloneSyncService;
@@ -53,6 +54,15 @@ impl SharedTransport {
         self.client
             .get_or_try_init(|| async {
                 let client = match &self.endpoint {
+                    TransportEndpoint::Url { endpoint, .. } if endpoint.starts_with("unix:") => {
+                        let channel = Endpoint::from_shared(endpoint.clone())
+                            .map_err(|error| {
+                                format!("invalid gRPC code-mode Unix socket endpoint: {error}")
+                            })?
+                            .connect_lazy();
+                        let transport = channel.map_err(io::Error::other);
+                        CodeModeHostClient::new(BoxCloneSyncService::new(transport))
+                    }
                     TransportEndpoint::Url {
                         endpoint,
                         http_client_factory,