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

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,7 +225,7 @@ impl ExecServerClient {
 
     pub async fn write(
         &self,
-        process_id: &str,
+        process_id: &ProcessId,
         chunk: Vec<u8>,
     ) -> Result<WriteResponse, ExecServerError> {
         self.inner
@@ -233,7 +233,7 @@ impl ExecServerClient {
             .call(
                 EXEC_WRITE_METHOD,
                 &WriteParams {
-                    process_id: process_id.to_string(),
+                    process_id: process_id.clone(),
                     chunk: chunk.into(),
                 },
             )
@@ -241,13 +241,16 @@ impl ExecServerClient {
             .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,7 +333,7 @@ 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
@@ -338,12 +341,12 @@ impl ExecServerClient {
             .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;

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

@@ -0,0 +1,74 @@
+use std::borrow::Borrow;
+use std::fmt;
+use std::ops::Deref;
+
+use serde::Deserialize;
+use serde::Serialize;
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
+#[serde(transparent)]
+pub struct ProcessId(String);
+
+impl ProcessId {
+    pub fn new(value: impl Into<String>) -> Self {
+        Self(value.into())
+    }
+
+    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 Borrow<str> for ProcessId {
+    fn borrow(&self) -> &str {
+        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)
+    }
+}
+
+impl From<&str> for ProcessId {
+    fn from(value: &str) -> Self {
+        Self(value.to_string())
+    }
+}
+
+impl From<&String> for ProcessId {
+    fn from(value: &String) -> Self {
+        Self(value.clone())
+    }
+}
+
+impl From<ProcessId> for String {
+    fn from(value: ProcessId) -> Self {
+        value.0
+    }
+}

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

@@ -5,6 +5,8 @@ use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
 use serde::Deserialize;
 use serde::Serialize;
 
+use crate::ProcessId;
+
 pub const INITIALIZE_METHOD: &str = "initialize";
 pub const INITIALIZED_METHOD: &str = "initialized";
 pub const EXEC_METHOD: &str = "process/start";
@@ -53,7 +55,7 @@ pub struct InitializeResponse {}
 pub struct ExecParams {
     /// Client-chosen logical process handle scoped to this connection/session.
     /// This is a protocol key, not an OS pid.
-    pub process_id: String,
+    pub process_id: ProcessId,
     pub argv: Vec<String>,
     pub cwd: PathBuf,
     pub env: HashMap<String, String>,
@@ -64,13 +66,13 @@ pub struct ExecParams {
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ExecResponse {
-    pub process_id: String,
+    pub process_id: ProcessId,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ReadParams {
-    pub process_id: String,
+    pub process_id: ProcessId,
     pub after_seq: Option<u64>,
     pub max_bytes: Option<usize>,
     pub wait_ms: Option<u64>,
@@ -98,7 +100,7 @@ pub struct ReadResponse {
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct WriteParams {
-    pub process_id: String,
+    pub process_id: ProcessId,
     pub chunk: ByteChunk,
 }
 
@@ -120,7 +122,7 @@ pub struct WriteResponse {
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct TerminateParams {
-    pub process_id: String,
+    pub process_id: ProcessId,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -140,7 +142,7 @@ pub enum ExecOutputStream {
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ExecOutputDeltaNotification {
-    pub process_id: String,
+    pub process_id: ProcessId,
     pub seq: u64,
     pub stream: ExecOutputStream,
     pub chunk: ByteChunk,
@@ -149,7 +151,7 @@ pub struct ExecOutputDeltaNotification {
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ExecExitedNotification {
-    pub process_id: String,
+    pub process_id: ProcessId,
     pub seq: u64,
     pub exit_code: i32,
 }
@@ -157,7 +159,7 @@ pub struct ExecExitedNotification {
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ExecClosedNotification {
-    pub process_id: String,
+    pub process_id: ProcessId,
     pub seq: u64,
 }
 

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

@@ -6,6 +6,7 @@ use pretty_assertions::assert_eq;
 use tokio::sync::mpsc;
 
 use super::ExecServerHandler;
+use crate::ProcessId;
 use crate::protocol::ExecParams;
 use crate::protocol::InitializeResponse;
 use crate::protocol::TerminateParams;
@@ -18,7 +19,7 @@ fn exec_params(process_id: &str) -> ExecParams {
         env.insert("PATH".to_string(), path.to_string_lossy().into_owned());
     }
     ExecParams {
-        process_id: process_id.to_string(),
+        process_id: ProcessId::from(process_id),
         argv: vec![
             "bash".to_string(),
             "-lc".to_string(),
@@ -84,7 +85,7 @@ async fn terminate_reports_false_after_process_exit() {
     loop {
         let response = handler
             .terminate(TerminateParams {
-                process_id: "proc-1".to_string(),
+                process_id: ProcessId::from("proc-1"),
             })
             .await
             .expect("terminate response");

codex-rs/exec-server/tests/exec_process.rs

@@ -9,6 +9,7 @@ use codex_exec_server::Environment;
 use codex_exec_server::ExecBackend;
 use codex_exec_server::ExecParams;
 use codex_exec_server::ExecProcess;
+use codex_exec_server::ProcessId;
 use codex_exec_server::ReadResponse;
 use codex_exec_server::StartedExecProcess;
 use pretty_assertions::assert_eq;
@@ -47,7 +48,7 @@ async fn assert_exec_process_starts_and_exits(use_remote: bool) -> Result<()> {
     let session = context
         .backend
         .start(ExecParams {
-            process_id: "proc-1".to_string(),
+            process_id: ProcessId::from("proc-1"),
             argv: vec!["true".to_string()],
             cwd: std::env::current_dir()?,
             env: Default::default(),
@@ -119,7 +120,7 @@ async fn assert_exec_process_streams_output(use_remote: bool) -> Result<()> {
     let session = context
         .backend
         .start(ExecParams {
-            process_id: process_id.clone(),
+            process_id: process_id.clone().into(),
             argv: vec![
                 "/bin/sh".to_string(),
                 "-c".to_string(),
@@ -148,7 +149,7 @@ async fn assert_exec_process_write_then_read(use_remote: bool) -> Result<()> {
     let session = context
         .backend
         .start(ExecParams {
-            process_id: process_id.clone(),
+            process_id: process_id.clone().into(),
             argv: vec![
                 "/usr/bin/python3".to_string(),
                 "-c".to_string(),
@@ -184,7 +185,7 @@ async fn assert_exec_process_preserves_queued_events_before_subscribe(
     let session = context
         .backend
         .start(ExecParams {
-            process_id: "proc-queued".to_string(),
+            process_id: ProcessId::from("proc-queued"),
             argv: vec![
                 "/bin/sh".to_string(),
                 "-c".to_string(),
@@ -214,7 +215,7 @@ async fn remote_exec_process_reports_transport_disconnect() -> Result<()> {
     let session = context
         .backend
         .start(ExecParams {
-            process_id: "proc-disconnect".to_string(),
+            process_id: ProcessId::from("proc-disconnect"),
             argv: vec![
                 "/bin/sh".to_string(),
                 "-c".to_string(),

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

@@ -6,6 +6,7 @@ use codex_app_server_protocol::JSONRPCMessage;
 use codex_app_server_protocol::JSONRPCResponse;
 use codex_exec_server::ExecResponse;
 use codex_exec_server::InitializeParams;
+use codex_exec_server::ProcessId;
 use common::exec_server::exec_server;
 use pretty_assertions::assert_eq;
 
@@ -62,7 +63,7 @@ async fn exec_server_starts_process_over_websocket() -> anyhow::Result<()> {
     assert_eq!(
         process_start_response,
         ExecResponse {
-            process_id: "proc-1".to_string()
+            process_id: ProcessId::from("proc-1")
         }
     );