feat: move exec-server ownership (#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>

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

@@ -1,3 +1,8 @@
+use std::sync::Arc;
+use std::sync::Mutex as StdMutex;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+
 use codex_app_server_protocol::JSONRPCErrorError;
 
 use crate::protocol::ExecParams;
@@ -16,6 +21,7 @@ use crate::protocol::FsRemoveParams;
 use crate::protocol::FsRemoveResponse;
 use crate::protocol::FsWriteFileParams;
 use crate::protocol::FsWriteFileResponse;
+use crate::protocol::InitializeParams;
 use crate::protocol::InitializeResponse;
 use crate::protocol::ReadParams;
 use crate::protocol::ReadResponse;
@@ -24,65 +30,126 @@ use crate::protocol::TerminateResponse;
 use crate::protocol::WriteParams;
 use crate::protocol::WriteResponse;
 use crate::rpc::RpcNotificationSender;
+use crate::rpc::invalid_request;
 use crate::server::file_system_handler::FileSystemHandler;
-use crate::server::process_handler::ProcessHandler;
+use crate::server::session_registry::SessionHandle;
+use crate::server::session_registry::SessionRegistry;
 
-#[derive(Clone)]
 pub(crate) struct ExecServerHandler {
-    process: ProcessHandler,
+    session_registry: Arc<SessionRegistry>,
+    notifications: RpcNotificationSender,
+    session: StdMutex<Option<SessionHandle>>,
     file_system: FileSystemHandler,
+    initialize_requested: AtomicBool,
+    initialized: AtomicBool,
 }
 
 impl ExecServerHandler {
-    pub(crate) fn new(notifications: RpcNotificationSender) -> Self {
+    pub(crate) fn new(
+        session_registry: Arc<SessionRegistry>,
+        notifications: RpcNotificationSender,
+    ) -> Self {
         Self {
-            process: ProcessHandler::new(notifications),
+            session_registry,
+            notifications,
+            session: StdMutex::new(None),
             file_system: FileSystemHandler::default(),
+            initialize_requested: AtomicBool::new(false),
+            initialized: AtomicBool::new(false),
         }
     }
 
     pub(crate) async fn shutdown(&self) {
-        self.process.shutdown().await;
+        if let Some(session) = self.session() {
+            session.detach().await;
+        }
     }
 
-    pub(crate) fn initialize(&self) -> Result<InitializeResponse, JSONRPCErrorError> {
-        self.process.initialize()
+    pub(crate) fn is_session_attached(&self) -> bool {
+        self.session()
+            .is_none_or(|session| session.is_session_attached())
+    }
+
+    pub(crate) async fn initialize(
+        &self,
+        params: InitializeParams,
+    ) -> Result<InitializeResponse, JSONRPCErrorError> {
+        if self.initialize_requested.swap(true, Ordering::SeqCst) {
+            return Err(invalid_request(
+                "initialize may only be sent once per connection".to_string(),
+            ));
+        }
+
+        let session = match self
+            .session_registry
+            .attach(params.resume_session_id.clone(), self.notifications.clone())
+            .await
+        {
+            Ok(session) => session,
+            Err(error) => {
+                self.initialize_requested.store(false, Ordering::SeqCst);
+                return Err(error);
+            }
+        };
+        let session_id = session.session_id().to_string();
+        tracing::debug!(
+            session_id,
+            connection_id = %session.connection_id(),
+            "exec-server session attached"
+        );
+        *self
+            .session
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(session);
+        Ok(InitializeResponse { session_id })
     }
 
     pub(crate) fn initialized(&self) -> Result<(), String> {
-        self.process.initialized()
+        if !self.initialize_requested.load(Ordering::SeqCst) {
+            return Err("received `initialized` notification before `initialize`".into());
+        }
+        self.require_session_attached()
+            .map_err(|error| error.message)?;
+        self.initialized.store(true, Ordering::SeqCst);
+        Ok(())
     }
 
     pub(crate) async fn exec(&self, params: ExecParams) -> Result<ExecResponse, JSONRPCErrorError> {
-        self.process.exec(params).await
+        let session = self.require_initialized_for("exec")?;
+        session.process().exec(params).await
     }
 
     pub(crate) async fn exec_read(
         &self,
         params: ReadParams,
     ) -> Result<ReadResponse, JSONRPCErrorError> {
-        self.process.exec_read(params).await
+        let session = self.require_initialized_for("exec")?;
+        let response = session.process().exec_read(params).await?;
+        self.require_session_attached()?;
+        Ok(response)
     }
 
     pub(crate) async fn exec_write(
         &self,
         params: WriteParams,
     ) -> Result<WriteResponse, JSONRPCErrorError> {
-        self.process.exec_write(params).await
+        let session = self.require_initialized_for("exec")?;
+        session.process().exec_write(params).await
     }
 
     pub(crate) async fn terminate(
         &self,
         params: TerminateParams,
     ) -> Result<TerminateResponse, JSONRPCErrorError> {
-        self.process.terminate(params).await
+        let session = self.require_initialized_for("exec")?;
+        session.process().terminate(params).await
     }
 
     pub(crate) async fn fs_read_file(
         &self,
         params: FsReadFileParams,
     ) -> Result<FsReadFileResponse, JSONRPCErrorError> {
-        self.process.require_initialized_for("filesystem")?;
+        self.require_initialized_for("filesystem")?;
         self.file_system.read_file(params).await
     }
 
@@ -90,7 +157,7 @@ impl ExecServerHandler {
         &self,
         params: FsWriteFileParams,
     ) -> Result<FsWriteFileResponse, JSONRPCErrorError> {
-        self.process.require_initialized_for("filesystem")?;
+        self.require_initialized_for("filesystem")?;
         self.file_system.write_file(params).await
     }
 
@@ -98,7 +165,7 @@ impl ExecServerHandler {
         &self,
         params: FsCreateDirectoryParams,
     ) -> Result<FsCreateDirectoryResponse, JSONRPCErrorError> {
-        self.process.require_initialized_for("filesystem")?;
+        self.require_initialized_for("filesystem")?;
         self.file_system.create_directory(params).await
     }
 
@@ -106,7 +173,7 @@ impl ExecServerHandler {
         &self,
         params: FsGetMetadataParams,
     ) -> Result<FsGetMetadataResponse, JSONRPCErrorError> {
-        self.process.require_initialized_for("filesystem")?;
+        self.require_initialized_for("filesystem")?;
         self.file_system.get_metadata(params).await
     }
 
@@ -114,7 +181,7 @@ impl ExecServerHandler {
         &self,
         params: FsReadDirectoryParams,
     ) -> Result<FsReadDirectoryResponse, JSONRPCErrorError> {
-        self.process.require_initialized_for("filesystem")?;
+        self.require_initialized_for("filesystem")?;
         self.file_system.read_directory(params).await
     }
 
@@ -122,7 +189,7 @@ impl ExecServerHandler {
         &self,
         params: FsRemoveParams,
     ) -> Result<FsRemoveResponse, JSONRPCErrorError> {
-        self.process.require_initialized_for("filesystem")?;
+        self.require_initialized_for("filesystem")?;
         self.file_system.remove(params).await
     }
 
@@ -130,9 +197,49 @@ impl ExecServerHandler {
         &self,
         params: FsCopyParams,
     ) -> Result<FsCopyResponse, JSONRPCErrorError> {
-        self.process.require_initialized_for("filesystem")?;
+        self.require_initialized_for("filesystem")?;
         self.file_system.copy(params).await
     }
+
+    fn require_initialized_for(
+        &self,
+        method_family: &str,
+    ) -> Result<SessionHandle, JSONRPCErrorError> {
+        if !self.initialize_requested.load(Ordering::SeqCst) {
+            return Err(invalid_request(format!(
+                "client must call initialize before using {method_family} methods"
+            )));
+        }
+        let session = self.require_session_attached()?;
+        if !self.initialized.load(Ordering::SeqCst) {
+            return Err(invalid_request(format!(
+                "client must send initialized before using {method_family} methods"
+            )));
+        }
+        Ok(session)
+    }
+
+    fn require_session_attached(&self) -> Result<SessionHandle, JSONRPCErrorError> {
+        let Some(session) = self.session() else {
+            return Err(invalid_request(
+                "client must call initialize before using methods".to_string(),
+            ));
+        };
+        if session.is_session_attached() {
+            return Ok(session);
+        }
+
+        Err(invalid_request(
+            "session has been resumed by another connection".to_string(),
+        ))
+    }
+
+    fn session(&self) -> Option<SessionHandle> {
+        self.session
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .clone()
+    }
 }
 
 #[cfg(test)]