feat: use ProcessId in exec-server (#15866)

Use a full struct for the ProcessId to increase readability and make it
easier in the future to make it evolve if needed

Author: jif-oai

codex-rs/core/src/unified_exec/process_manager.rs

@@ -603,7 +603,7 @@ impl UnifiedExecProcessManager {
             let started = environment
                 .get_exec_backend()
                 .start(codex_exec_server::ExecParams {
-                    process_id: exec_server_process_id(process_id),
+                    process_id: exec_server_process_id(process_id).into(),
                     argv: env.command.clone(),
                     cwd: env.cwd.clone(),
                     env: env.env.clone(),

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

@@ -113,7 +113,7 @@ struct Inner {
     // process on the connection. Keep a local process_id -> session registry so
     // we can turn those connection-global notifications into process wakeups
     // without making notifications the source of truth for output delivery.
-    sessions: ArcSwap<HashMap<String, Arc<SessionState>>>,
+    sessions: ArcSwap<HashMap<ProcessId, Arc<SessionState>>>,
     // ArcSwap makes reads cheap on the hot notification path, but writes still
     // need serialization so concurrent register/remove operations do not
     // overwrite each other's copy-on-write updates.
@@ -225,29 +225,32 @@ impl ExecServerClient {
 
     pub async fn write(
         &self,
-        process_id: &str,
+        process_id: &ProcessId,
         chunk: Vec<u8>,
     ) -> Result<WriteResponse, ExecServerError> {
         self.inner
             .client
             .call(
                 EXEC_WRITE_METHOD,
                 &WriteParams {
-                    process_id: process_id.to_string(),
+                    process_id: process_id.clone(),
                     chunk: chunk.into(),
                 },
             )
             .await
             .map_err(Into::into)
     }
 
-    pub async fn terminate(&self, process_id: &str) -> Result<TerminateResponse, ExecServerError> {
+    pub async fn terminate(
+        &self,
+        process_id: &ProcessId,
+    ) -> Result<TerminateResponse, ExecServerError> {
         self.inner
             .client
             .call(
                 EXEC_TERMINATE_METHOD,
                 &TerminateParams {
-                    process_id: process_id.to_string(),
+                    process_id: process_id.clone(),
                 },
             )
             .await
@@ -330,20 +333,20 @@ impl ExecServerClient {
 
     pub(crate) async fn register_session(
         &self,
-        process_id: &str,
+        process_id: &ProcessId,
     ) -> Result<Session, ExecServerError> {
         let state = Arc::new(SessionState::new());
         self.inner
             .insert_session(process_id, Arc::clone(&state))
             .await?;
         Ok(Session {
             client: self.clone(),
-            process_id: process_id.to_string().into(),
+            process_id: process_id.clone(),
             state,
         })
     }
 
-    pub(crate) async fn unregister_session(&self, process_id: &str) {
+    pub(crate) async fn unregister_session(&self, process_id: &ProcessId) {
         self.inner.remove_session(process_id).await;
     }
 
@@ -487,7 +490,7 @@ impl Session {
         match self
             .client
             .read(ReadParams {
-                process_id: self.process_id.to_string(),
+                process_id: self.process_id.clone(),
                 after_seq,
                 max_bytes,
                 wait_ms,
@@ -519,13 +522,13 @@ impl Session {
 }
 
 impl Inner {
-    fn get_session(&self, process_id: &str) -> Option<Arc<SessionState>> {
+    fn get_session(&self, process_id: &ProcessId) -> Option<Arc<SessionState>> {
         self.sessions.load().get(process_id).cloned()
     }
 
     async fn insert_session(
         &self,
-        process_id: &str,
+        process_id: &ProcessId,
         session: Arc<SessionState>,
     ) -> Result<(), ExecServerError> {
         let _sessions_write_guard = self.sessions_write_lock.lock().await;
@@ -536,12 +539,12 @@ impl Inner {
             )));
         }
         let mut next_sessions = sessions.as_ref().clone();
-        next_sessions.insert(process_id.to_string(), session);
+        next_sessions.insert(process_id.clone(), session);
         self.sessions.store(Arc::new(next_sessions));
         Ok(())
     }
 
-    async fn remove_session(&self, process_id: &str) -> Option<Arc<SessionState>> {
+    async fn remove_session(&self, process_id: &ProcessId) -> Option<Arc<SessionState>> {
         let _sessions_write_guard = self.sessions_write_lock.lock().await;
         let sessions = self.sessions.load();
         let session = sessions.get(process_id).cloned();
@@ -552,7 +555,7 @@ impl Inner {
         session
     }
 
-    async fn take_all_sessions(&self) -> HashMap<String, Arc<SessionState>> {
+    async fn take_all_sessions(&self) -> HashMap<ProcessId, Arc<SessionState>> {
         let _sessions_write_guard = self.sessions_write_lock.lock().await;
         let sessions = self.sessions.load();
         let drained_sessions = sessions.as_ref().clone();
@@ -640,6 +643,7 @@ mod tests {
 
     use super::ExecServerClient;
     use super::ExecServerClientConnectOptions;
+    use crate::ProcessId;
     use crate::connection::JsonRpcConnection;
     use crate::protocol::EXEC_EXITED_METHOD;
     use crate::protocol::EXEC_OUTPUT_DELTA_METHOD;
@@ -718,12 +722,14 @@ mod tests {
         .await
         .expect("client should connect");
 
+        let noisy_process_id = ProcessId::from("noisy");
+        let quiet_process_id = ProcessId::from("quiet");
         let _noisy_session = client
-            .register_session("noisy")
+            .register_session(&noisy_process_id)
             .await
             .expect("noisy session should register");
         let quiet_session = client
-            .register_session("quiet")
+            .register_session(&quiet_process_id)
             .await
             .expect("quiet session should register");
         let mut quiet_wake_rx = quiet_session.subscribe_wake();
@@ -734,7 +740,7 @@ mod tests {
                     method: EXEC_OUTPUT_DELTA_METHOD.to_string(),
                     params: Some(
                         serde_json::to_value(ExecOutputDeltaNotification {
-                            process_id: "noisy".to_string(),
+                            process_id: noisy_process_id.clone(),
                             seq,
                             stream: ExecOutputStream::Stdout,
                             chunk: b"x".to_vec().into(),
@@ -751,7 +757,7 @@ mod tests {
                 method: EXEC_EXITED_METHOD.to_string(),
                 params: Some(
                     serde_json::to_value(ExecExitedNotification {
-                        process_id: "quiet".to_string(),
+                        process_id: quiet_process_id,
                         seq: 1,
                         exit_code: 17,
                     })

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

@@ -159,6 +159,7 @@ mod tests {
 
     use super::Environment;
     use super::EnvironmentManager;
+    use crate::ProcessId;
     use pretty_assertions::assert_eq;
 
     #[tokio::test]
@@ -195,7 +196,7 @@ mod tests {
         let response = environment
             .get_exec_backend()
             .start(crate::ExecParams {
-                process_id: "default-env-proc".to_string(),
+                process_id: ProcessId::from("default-env-proc"),
                 argv: vec!["true".to_string()],
                 cwd: std::env::current_dir().expect("read current dir"),
                 env: Default::default(),

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

@@ -6,6 +6,7 @@ mod file_system;
 mod local_file_system;
 mod local_process;
 mod process;
+mod process_id;
 mod protocol;
 mod remote_file_system;
 mod remote_process;
@@ -43,8 +44,8 @@ pub use file_system::ReadDirectoryEntry;
 pub use file_system::RemoveOptions;
 pub use process::ExecBackend;
 pub use process::ExecProcess;
-pub use process::ProcessId;
 pub use process::StartedExecProcess;
+pub use process_id::ProcessId;
 pub use protocol::ExecClosedNotification;
 pub use protocol::ExecExitedNotification;
 pub use protocol::ExecOutputDeltaNotification;

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

@@ -75,7 +75,7 @@ enum ProcessEntry {
 
 struct Inner {
     notifications: RpcNotificationSender,
-    processes: Mutex<HashMap<String, ProcessEntry>>,
+    processes: Mutex<HashMap<ProcessId, ProcessEntry>>,
     initialize_requested: AtomicBool,
     initialized: AtomicBool,
 }
@@ -420,7 +420,7 @@ impl ExecBackend for LocalProcess {
             .map_err(map_handler_error)?;
         Ok(StartedExecProcess {
             process: Arc::new(LocalExecProcess {
-                process_id: response.process_id.into(),
+                process_id: response.process_id,
                 backend: self.clone(),
                 wake_tx,
             }),
@@ -461,13 +461,13 @@ impl ExecProcess for LocalExecProcess {
 impl LocalProcess {
     async fn read(
         &self,
-        process_id: &str,
+        process_id: &ProcessId,
         after_seq: Option<u64>,
         max_bytes: Option<usize>,
         wait_ms: Option<u64>,
     ) -> Result<ReadResponse, ExecServerError> {
         self.exec_read(ReadParams {
-            process_id: process_id.to_string(),
+            process_id: process_id.clone(),
             after_seq,
             max_bytes,
             wait_ms,
@@ -478,20 +478,20 @@ impl LocalProcess {
 
     async fn write(
         &self,
-        process_id: &str,
+        process_id: &ProcessId,
         chunk: Vec<u8>,
     ) -> Result<WriteResponse, ExecServerError> {
         self.exec_write(WriteParams {
-            process_id: process_id.to_string(),
+            process_id: process_id.clone(),
             chunk: chunk.into(),
         })
         .await
         .map_err(map_handler_error)
     }
 
-    async fn terminate(&self, process_id: &str) -> Result<(), ExecServerError> {
+    async fn terminate(&self, process_id: &ProcessId) -> Result<(), ExecServerError> {
         self.terminate_process(TerminateParams {
-            process_id: process_id.to_string(),
+            process_id: process_id.clone(),
         })
         .await
         .map_err(map_handler_error)?;
@@ -507,7 +507,7 @@ fn map_handler_error(error: JSONRPCErrorError) -> ExecServerError {
 }
 
 async fn stream_output(
-    process_id: String,
+    process_id: ProcessId,
     stream: ExecOutputStream,
     mut receiver: tokio::sync::mpsc::Receiver<Vec<u8>>,
     inner: Arc<Inner>,
@@ -560,7 +560,7 @@ async fn stream_output(
 }
 
 async fn watch_exit(
-    process_id: String,
+    process_id: ProcessId,
     exit_rx: tokio::sync::oneshot::Receiver<i32>,
     inner: Arc<Inner>,
     output_notify: Arc<Notify>,
@@ -605,7 +605,7 @@ async fn watch_exit(
     }
 }
 
-async fn finish_output_stream(process_id: String, inner: Arc<Inner>) {
+async fn finish_output_stream(process_id: ProcessId, inner: Arc<Inner>) {
     {
         let mut processes = inner.processes.lock().await;
         let Some(ProcessEntry::Running(process)) = processes.get_mut(&process_id) else {
@@ -620,7 +620,7 @@ async fn finish_output_stream(process_id: String, inner: Arc<Inner>) {
     maybe_emit_closed(process_id, inner).await;
 }
 
-async fn maybe_emit_closed(process_id: String, inner: Arc<Inner>) {
+async fn maybe_emit_closed(process_id: ProcessId, inner: Arc<Inner>) {
     let notification = {
         let mut processes = inner.processes.lock().await;
         let Some(ProcessEntry::Running(process)) = processes.get_mut(&process_id) else {

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

@@ -1,58 +1,18 @@
-use std::fmt;
-use std::ops::Deref;
 use std::sync::Arc;
 
 use async_trait::async_trait;
 use tokio::sync::watch;
 
 use crate::ExecServerError;
+use crate::ProcessId;
 use crate::protocol::ExecParams;
 use crate::protocol::ReadResponse;
 use crate::protocol::WriteResponse;
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-pub struct ProcessId(String);
-
 pub struct StartedExecProcess {
     pub process: Arc<dyn ExecProcess>,
 }
 
-impl ProcessId {
-    pub fn as_str(&self) -> &str {
-        &self.0
-    }
-
-    pub fn into_inner(self) -> String {
-        self.0
-    }
-}
-
-impl Deref for ProcessId {
-    type Target = str;
-
-    fn deref(&self) -> &Self::Target {
-        self.as_str()
-    }
-}
-
-impl AsRef<str> for ProcessId {
-    fn as_ref(&self) -> &str {
-        self.as_str()
-    }
-}
-
-impl fmt::Display for ProcessId {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        self.0.fmt(f)
-    }
-}
-
-impl From<String> for ProcessId {
-    fn from(value: String) -> Self {
-        Self(value)
-    }
-}
-
 #[async_trait]
 pub trait ExecProcess: Send + Sync {
     fn process_id(&self) -> &ProcessId;