feat: move exec-server ownership (openai#16344)

This introduces session-scoped ownership for exec-server so ws
disconnects no longer immediately kill running remote exec processes,
and it prepares the protocol for reconnect-based resume.
- add session_id / resume_session_id to the exec-server initialize
handshake
  - move process ownership under a shared session registry
- detach sessions on websocket disconnect and expire them after a TTL
instead of killing processes immediately (we will resume based on this)
- allow a new connection to resume an existing session and take over
notifications/ownership
- I use UUID to make them not predictable as we don't have auth for now
- make detached-session expiry authoritative at resume time so teardown
wins at the TTL boundary
- reject long-poll process/read calls that get resumed out from under an
older attachment

---------

Co-authored-by: Codex <noreply@openai.com>

Author: jif-oai

codex-rs/Cargo.lock

@@ -40,6 +40,7 @@ tokio = { workspace = true, features = [
 ] }
 tokio-tungstenite = { workspace = true }
 tracing = { workspace = true }
+uuid = { workspace = true, features = ["v4"] }
 
 [dev-dependencies]
 anyhow = { workspace = true }

codex-rs/exec-server/Cargo.toml

@@ -71,6 +71,7 @@ impl Default for ExecServerClientConnectOptions {
         Self {
             client_name: "codex-core".to_string(),
             initialize_timeout: INITIALIZE_TIMEOUT,
+            resume_session_id: None,
         }
     }
 }
@@ -80,6 +81,7 @@ impl From<RemoteExecServerConnectArgs> for ExecServerClientConnectOptions {
         Self {
             client_name: value.client_name,
             initialize_timeout: value.initialize_timeout,
+            resume_session_id: value.resume_session_id,
         }
     }
 }
@@ -91,6 +93,7 @@ impl RemoteExecServerConnectArgs {
             client_name,
             connect_timeout: CONNECT_TIMEOUT,
             initialize_timeout: INITIALIZE_TIMEOUT,
+            resume_session_id: None,
         }
     }
 }
@@ -118,6 +121,7 @@ struct Inner {
     // need serialization so concurrent register/remove operations do not
     // overwrite each other's copy-on-write updates.
     sessions_write_lock: Mutex<()>,
+    session_id: std::sync::RwLock<Option<String>>,
     reader_task: tokio::task::JoinHandle<()>,
 }
 
@@ -190,14 +194,29 @@ impl ExecServerClient {
         let ExecServerClientConnectOptions {
             client_name,
             initialize_timeout,
+            resume_session_id,
         } = options;
 
         timeout(initialize_timeout, async {
-            let response = self
+            let response: InitializeResponse = self
                 .inner
                 .client
-                .call(INITIALIZE_METHOD, &InitializeParams { client_name })
+                .call(
+                    INITIALIZE_METHOD,
+                    &InitializeParams {
+                        client_name,
+                        resume_session_id,
+                    },
+                )
                 .await?;
+            {
+                let mut session_id = self
+                    .inner
+                    .session_id
+                    .write()
+                    .unwrap_or_else(std::sync::PoisonError::into_inner);
+                *session_id = Some(response.session_id.clone());
+            }
             self.notify_initialized().await?;
             Ok(response)
         })
@@ -350,6 +369,14 @@ impl ExecServerClient {
         self.inner.remove_session(process_id).await;
     }
 
+    pub fn session_id(&self) -> Option<String> {
+        self.inner
+            .session_id
+            .read()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .clone()
+    }
+
     async fn connect(
         connection: JsonRpcConnection,
         options: ExecServerClientConnectOptions,
@@ -388,6 +415,7 @@ impl ExecServerClient {
                 client: rpc_client,
                 sessions: ArcSwap::from_pointee(HashMap::new()),
                 sessions_write_lock: Mutex::new(()),
+                session_id: std::sync::RwLock::new(None),
                 reader_task,
             }
         });
@@ -693,8 +721,10 @@ mod tests {
                 &mut server_writer,
                 JSONRPCMessage::Response(JSONRPCResponse {
                     id: request.id,
-                    result: serde_json::to_value(InitializeResponse {})
-                        .expect("initialize response should serialize"),
+                    result: serde_json::to_value(InitializeResponse {
+                        session_id: "session-1".to_string(),
+                    })
+                    .expect("initialize response should serialize"),
                 }),
             )
             .await;

codex-rs/exec-server/src/client.rs

@@ -5,6 +5,7 @@ use std::time::Duration;
 pub struct ExecServerClientConnectOptions {
     pub client_name: String,
     pub initialize_timeout: Duration,
+    pub resume_session_id: Option<String>,
 }
 
 /// WebSocket connection arguments for a remote exec-server.
@@ -14,4 +15,5 @@ pub struct RemoteExecServerConnectArgs {
     pub client_name: String,
     pub connect_timeout: Duration,
     pub initialize_timeout: Duration,
+    pub resume_session_id: Option<String>,
 }

codex-rs/exec-server/src/client_api.rs

@@ -4,6 +4,7 @@ use futures::StreamExt;
 use tokio::io::AsyncRead;
 use tokio::io::AsyncWrite;
 use tokio::sync::mpsc;
+use tokio::sync::watch;
 use tokio_tungstenite::WebSocketStream;
 use tokio_tungstenite::tungstenite::Message;
 
@@ -28,6 +29,7 @@ pub(crate) enum JsonRpcConnectionEvent {
 pub(crate) struct JsonRpcConnection {
     outgoing_tx: mpsc::Sender<JSONRPCMessage>,
     incoming_rx: mpsc::Receiver<JsonRpcConnectionEvent>,
+    disconnected_rx: watch::Receiver<bool>,
     task_handles: Vec<tokio::task::JoinHandle<()>>,
 }
 
@@ -40,9 +42,11 @@ impl JsonRpcConnection {
     {
         let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY);
         let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY);
+        let (disconnected_tx, disconnected_rx) = watch::channel(false);
 
         let reader_label = connection_label.clone();
         let incoming_tx_for_reader = incoming_tx.clone();
+        let disconnected_tx_for_reader = disconnected_tx.clone();
         let reader_task = tokio::spawn(async move {
             let mut lines = BufReader::new(reader).lines();
             loop {
@@ -73,12 +77,18 @@ impl JsonRpcConnection {
                         }
                     }
                     Ok(None) => {
-                        send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await;
+                        send_disconnected(
+                            &incoming_tx_for_reader,
+                            &disconnected_tx_for_reader,
+                            /*reason*/ None,
+                        )
+                        .await;
                         break;
                     }
                     Err(err) => {
                         send_disconnected(
                             &incoming_tx_for_reader,
+                            &disconnected_tx_for_reader,
                             Some(format!(
                                 "failed to read JSON-RPC message from {reader_label}: {err}"
                             )),
@@ -96,6 +106,7 @@ impl JsonRpcConnection {
                 if let Err(err) = write_jsonrpc_line_message(&mut writer, &message).await {
                     send_disconnected(
                         &incoming_tx,
+                        &disconnected_tx,
                         Some(format!(
                             "failed to write JSON-RPC message to {connection_label}: {err}"
                         )),
@@ -109,6 +120,7 @@ impl JsonRpcConnection {
         Self {
             outgoing_tx,
             incoming_rx,
+            disconnected_rx,
             task_handles: vec![reader_task, writer_task],
         }
     }
@@ -119,10 +131,12 @@ impl JsonRpcConnection {
     {
         let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY);
         let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY);
+        let (disconnected_tx, disconnected_rx) = watch::channel(false);
         let (mut websocket_writer, mut websocket_reader) = stream.split();
 
         let reader_label = connection_label.clone();
         let incoming_tx_for_reader = incoming_tx.clone();
+        let disconnected_tx_for_reader = disconnected_tx.clone();
         let reader_task = tokio::spawn(async move {
             loop {
                 match websocket_reader.next().await {
@@ -171,14 +185,20 @@ impl JsonRpcConnection {
                         }
                     }
                     Some(Ok(Message::Close(_))) => {
-                        send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await;
+                        send_disconnected(
+                            &incoming_tx_for_reader,
+                            &disconnected_tx_for_reader,
+                            /*reason*/ None,
+                        )
+                        .await;
                         break;
                     }
                     Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) => {}
                     Some(Ok(_)) => {}
                     Some(Err(err)) => {
                         send_disconnected(
                             &incoming_tx_for_reader,
+                            &disconnected_tx_for_reader,
                             Some(format!(
                                 "failed to read websocket JSON-RPC message from {reader_label}: {err}"
                             )),
@@ -187,7 +207,12 @@ impl JsonRpcConnection {
                         break;
                     }
                     None => {
-                        send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await;
+                        send_disconnected(
+                            &incoming_tx_for_reader,
+                            &disconnected_tx_for_reader,
+                            /*reason*/ None,
+                        )
+                        .await;
                         break;
                     }
                 }
@@ -202,6 +227,7 @@ impl JsonRpcConnection {
                         {
                             send_disconnected(
                                 &incoming_tx,
+                                &disconnected_tx,
                                 Some(format!(
                                     "failed to write websocket JSON-RPC message to {connection_label}: {err}"
                                 )),
@@ -213,6 +239,7 @@ impl JsonRpcConnection {
                     Err(err) => {
                         send_disconnected(
                             &incoming_tx,
+                            &disconnected_tx,
                             Some(format!(
                                 "failed to serialize JSON-RPC message for {connection_label}: {err}"
                             )),
@@ -227,6 +254,7 @@ impl JsonRpcConnection {
         Self {
             outgoing_tx,
             incoming_rx,
+            disconnected_rx,
             task_handles: vec![reader_task, writer_task],
         }
     }
@@ -236,16 +264,24 @@ impl JsonRpcConnection {
     ) -> (
         mpsc::Sender<JSONRPCMessage>,
         mpsc::Receiver<JsonRpcConnectionEvent>,
+        watch::Receiver<bool>,
         Vec<tokio::task::JoinHandle<()>>,
     ) {
-        (self.outgoing_tx, self.incoming_rx, self.task_handles)
+        (
+            self.outgoing_tx,
+            self.incoming_rx,
+            self.disconnected_rx,
+            self.task_handles,
+        )
     }
 }
 
 async fn send_disconnected(
     incoming_tx: &mpsc::Sender<JsonRpcConnectionEvent>,
+    disconnected_tx: &watch::Sender<bool>,
     reason: Option<String>,
 ) {
+    let _ = disconnected_tx.send(true);
     let _ = incoming_tx
         .send(JsonRpcConnectionEvent::Disconnected { reason })
         .await;

codex-rs/exec-server/src/connection.rs

@@ -105,18 +105,10 @@ pub struct Environment {
 
 impl Default for Environment {
     fn default() -> Self {
-        let local_process = LocalProcess::default();
-        if let Err(err) = local_process.initialize() {
-            panic!("default local process initialization should succeed: {err:?}");
-        }
-        if let Err(err) = local_process.initialized() {
-            panic!("default local process should accept initialized notification: {err}");
-        }
-
         Self {
             exec_server_url: None,
             remote_exec_server_client: None,
-            exec_backend: Arc::new(local_process),
+            exec_backend: Arc::new(LocalProcess::default()),
         }
     }
 }
@@ -146,31 +138,20 @@ impl Environment {
                     client_name: "codex-environment".to_string(),
                     connect_timeout: std::time::Duration::from_secs(5),
                     initialize_timeout: std::time::Duration::from_secs(5),
+                    resume_session_id: None,
                 })
                 .await?,
             )
         } else {
             None
         };
 
-        let exec_backend: Arc<dyn ExecBackend> = match remote_exec_server_client.clone() {
-            Some(client) => Arc::new(RemoteProcess::new(client)),
-            None if exec_server_url.is_some() => {
-                return Err(ExecServerError::Protocol(
-                    "remote mode should have an exec-server client".to_string(),
-                ));
-            }
-            None => {
-                let local_process = LocalProcess::default();
-                local_process
-                    .initialize()
-                    .map_err(|err| ExecServerError::Protocol(err.message))?;
-                local_process
-                    .initialized()
-                    .map_err(ExecServerError::Protocol)?;
-                Arc::new(local_process)
-            }
-        };
+        let exec_backend: Arc<dyn ExecBackend> =
+            if let Some(client) = remote_exec_server_client.clone() {
+                Arc::new(RemoteProcess::new(client))
+            } else {
+                Arc::new(LocalProcess::default())
+            };
 
         Ok(Self {
             exec_server_url,