Recover code mode after host process failures

codex-rs/code-mode-client/src/connection.rs

@@ -286,6 +286,7 @@ impl ConnectionState {
         if !self.alive.swap(false, Ordering::AcqRel) {
             return;
         }
+        warn!(%reason, "code-mode host connection failed");
         if let Ok(mut failure) = self.failure.lock() {
             *failure = Some(reason.clone());
         }

codex-rs/code-mode-client/src/lib.rs

@@ -132,6 +132,10 @@ struct IpcCodeModeSession {
 }
 
 impl CodeModeSession for IpcCodeModeSession {
+    fn is_alive(&self) -> bool {
+        !self.shutdown.load(Ordering::Acquire) && self.connection.is_alive()
+    }
+
     fn execute<'a>(
         &'a self,
         request: ExecuteRequest,

codex-rs/code-mode-host/src/main.rs

@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::future::Future;
 use std::sync::Arc;
 use std::sync::atomic::AtomicU64;
 use std::sync::atomic::Ordering;
@@ -48,8 +49,8 @@ async fn run() -> Result<(), String> {
     let state = Arc::new(HostState {
         sessions: Mutex::new(HashMap::new()),
         next_session_id: AtomicU64::new(1),
-        peer,
         host_id: std::process::id().to_string(),
+        peer,
     });
 
     let writer = tokio::spawn(async move {
@@ -70,7 +71,7 @@ async fn run() -> Result<(), String> {
         match message {
             ClientMessage::Request { id, request } => {
                 let state = Arc::clone(&state);
-                tokio::spawn(async move {
+                spawn_critical_request_task("request handler", async move {
                     state.handle_request(id, request).await;
                 });
             }
@@ -94,9 +95,25 @@ async fn run() -> Result<(), String> {
     writer.await.map_err(|err| err.to_string())?
 }
 
+fn spawn_critical_request_task(
+    description: &'static str,
+    future: impl Future<Output = ()> + Send + 'static,
+) {
+    let task = tokio::spawn(future);
+    tokio::spawn(async move {
+        if let Err(err) = task.await
+            && err.is_panic()
+        {
+            eprintln!("code-mode host {description} panicked: {err}");
+            std::process::exit(1);
+        }
+    });
+}
+
 struct HostState {
     sessions: Mutex<HashMap<SessionId, Arc<CodeModeService>>>,
     next_session_id: AtomicU64,
+    host_id: String,
     peer: Arc<HostPeer>,
 }
 
@@ -116,7 +133,6 @@ impl HostState {
                         self.host_id.clone(),
                     )),
                 );
-    host_id: String,
                 self.respond(request_id, Ok(HostResponse::SessionCreated { session_id }))
                     .await;
             }
@@ -134,7 +150,7 @@ impl HostState {
                             )
                             .await;
                             let peer = Arc::clone(&self.peer);
-                            tokio::spawn(async move {
+                            spawn_critical_request_task("initial response handler", async move {
                                 let response = started.initial_response().await;
                                 peer.send(HostMessage::InitialResponse {
                                     id: request_id,

codex-rs/code-mode-protocol/src/session.rs

@@ -110,6 +110,14 @@ pub trait CodeModeSessionDelegate: Send + Sync {
 /// must keep those values isolated. Implementations may execute cells
 /// in-process or remotely.
 pub trait CodeModeSession: Send + Sync {
+    /// Returns whether the session can still accept requests.
+    ///
+    /// Remote implementations should return `false` after their underlying
+    /// connection fails so callers can create a fresh session for later work.
+    fn is_alive(&self) -> bool {
+        true
+    }
+
     fn execute<'a>(
         &'a self,
         request: ExecuteRequest,

codex-rs/code-mode/src/service.rs

@@ -353,6 +353,10 @@ impl Drop for CodeModeService {
 }
 
 impl CodeModeSession for CodeModeService {
+    fn is_alive(&self) -> bool {
+        !self.inner.shutting_down.load(Ordering::Acquire)
+    }
+
     fn execute<'a>(
         &'a self,
         request: ExecuteRequest,
@@ -834,11 +838,11 @@ mod tests {
             cells: Mutex::new(HashMap::new()),
             delegate: Arc::new(NoopCodeModeSessionDelegate),
             shutting_down: std::sync::atomic::AtomicBool::new(false),
+            cell_id_prefix: None,
             next_cell_id: AtomicU64::new(1),
         })
     }
 
-            cell_id_prefix: None,
     #[tokio::test]
     async fn synchronous_exit_returns_successfully() {
         let service = CodeModeService::new();
@@ -865,10 +869,6 @@ mod tests {
         );
     }
 
-    #[tokio::test]
-    async fn stored_values_are_shared_between_cells_but_not_sessions() {
-        let first_session = CodeModeService::new();
-        let second_session = CodeModeService::new();
     #[tokio::test]
     async fn cell_ids_include_the_configured_host_prefix() {
         let service = CodeModeService::with_delegate_and_cell_id_prefix(
@@ -890,6 +890,10 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn stored_values_are_shared_between_cells_but_not_sessions() {
+        let first_session = CodeModeService::new();
+        let second_session = CodeModeService::new();
 
         let write_response = execute(
             &first_session,

codex-rs/core/src/tools/code_mode/mod.rs

@@ -16,7 +16,10 @@ use codex_code_mode_protocol::CodeModeToolKind;
 use codex_code_mode_protocol::RuntimeResponse;
 use codex_protocol::models::FunctionCallOutputContentItem;
 use serde_json::Value as JsonValue;
-use tokio::sync::OnceCell;
+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;
@@ -61,18 +64,22 @@ pub(crate) struct ExecContext {
 }
 
 pub(crate) struct CodeModeService {
-    session: OnceCell<Arc<dyn CodeModeSession>>,
+    session: Mutex<Option<Arc<dyn CodeModeSession>>>,
     provider: Arc<dyn CodeModeSessionProvider>,
     dispatch_broker: Arc<CodeModeDispatchBroker>,
+    session_init_permit: Semaphore,
+    shutting_down: AtomicBool,
 }
 
 impl CodeModeService {
     pub(crate) fn new(provider: Arc<dyn CodeModeSessionProvider>) -> Self {
         let dispatch_broker = Arc::new(CodeModeDispatchBroker::new());
         Self {
-            session: OnceCell::new(),
+            session: Mutex::new(None),
             provider,
             dispatch_broker,
+            session_init_permit: Semaphore::new(/*permits*/ 1),
+            shutting_down: AtomicBool::new(false),
         }
     }
 
@@ -80,25 +87,40 @@ impl CodeModeService {
         &self,
         request: codex_code_mode_protocol::ExecuteRequest,
     ) -> Result<codex_code_mode_protocol::StartedCell, String> {
-        self.session().await?.execute(request).await
+        self.session_for_execute().await?.execute(request).await
     }
 
     pub(crate) async fn wait(
         &self,
         request: codex_code_mode_protocol::WaitRequest,
     ) -> Result<codex_code_mode_protocol::WaitOutcome, String> {
-        self.session().await?.wait(request).await
+        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<codex_code_mode_protocol::WaitOutcome, String> {
-        self.session().await?.terminate(cell_id).await
+        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.get() {
+        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(()),
         }
@@ -133,10 +155,43 @@ impl CodeModeService {
         )
     }
 
-    async fn session(&self) -> Result<&Arc<dyn CodeModeSession>, String> {
-        self.session
-            .get_or_try_init(|| self.provider.create_session(self.dispatch_broker.clone()))
+    async fn session_for_execute(&self) -> Result<Arc<dyn CodeModeSession>, 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<Arc<dyn CodeModeSession>> {
+        self.session
+            .lock()
+            .await
+            .as_ref()
+            .filter(|session| session.is_alive())
+            .cloned()
     }
 }
 
@@ -321,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_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() {
@@ -360,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<dyn CodeModeSessionDelegate>,
+        ) -> CodeModeSessionProviderFuture<'a> {
+            let generation = self.sessions_created.fetch_add(1, Ordering::Relaxed) + 1;
+            Box::pin(async move {
+                let session: Arc<dyn CodeModeSession> = 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<dyn CodeModeSessionProvider> = 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);
+    }
 }