diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 5a16b4768..2d992507b 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/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(), diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index b19bca74e..d73f6cf95 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/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>>, + sessions: ArcSwap>>, // 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, ) -> Result { 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 { + pub async fn terminate( + &self, + process_id: &ProcessId, + ) -> Result { 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 { 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> { + fn get_session(&self, process_id: &ProcessId) -> Option> { self.sessions.load().get(process_id).cloned() } async fn insert_session( &self, - process_id: &str, + process_id: &ProcessId, session: Arc, ) -> 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> { + async fn remove_session(&self, process_id: &ProcessId) -> Option> { 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> { + async fn take_all_sessions(&self) -> HashMap> { 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, }) diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index de576aa25..0518612c3 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/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(), diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 544340c4c..77834c614 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/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; diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index 591beedf3..7c7f0af8c 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -75,7 +75,7 @@ enum ProcessEntry { struct Inner { notifications: RpcNotificationSender, - processes: Mutex>, + processes: Mutex>, 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, max_bytes: Option, wait_ms: Option, ) -> Result { 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, ) -> Result { 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>, inner: Arc, @@ -560,7 +560,7 @@ async fn stream_output( } async fn watch_exit( - process_id: String, + process_id: ProcessId, exit_rx: tokio::sync::oneshot::Receiver, inner: Arc, output_notify: Arc, @@ -605,7 +605,7 @@ async fn watch_exit( } } -async fn finish_output_stream(process_id: String, inner: Arc) { +async fn finish_output_stream(process_id: ProcessId, inner: Arc) { { 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) { maybe_emit_closed(process_id, inner).await; } -async fn maybe_emit_closed(process_id: String, inner: Arc) { +async fn maybe_emit_closed(process_id: ProcessId, inner: Arc) { let notification = { let mut processes = inner.processes.lock().await; let Some(ProcessEntry::Running(process)) = processes.get_mut(&process_id) else { diff --git a/codex-rs/exec-server/src/process.rs b/codex-rs/exec-server/src/process.rs index f96b90ad3..e455c77bd 100644 --- a/codex-rs/exec-server/src/process.rs +++ b/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, } -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 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 for ProcessId { - fn from(value: String) -> Self { - Self(value) - } -} - #[async_trait] pub trait ExecProcess: Send + Sync { fn process_id(&self) -> &ProcessId; diff --git a/codex-rs/exec-server/src/process_id.rs b/codex-rs/exec-server/src/process_id.rs new file mode 100644 index 000000000..f25c81009 --- /dev/null +++ b/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) -> 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 for ProcessId { + fn borrow(&self) -> &str { + self.as_str() + } +} + +impl AsRef 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 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 for String { + fn from(value: ProcessId) -> Self { + value.0 + } +} diff --git a/codex-rs/exec-server/src/protocol.rs b/codex-rs/exec-server/src/protocol.rs index 5d074689e..0a0d0a4e8 100644 --- a/codex-rs/exec-server/src/protocol.rs +++ b/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, pub cwd: PathBuf, pub env: HashMap, @@ -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, pub max_bytes: Option, pub wait_ms: Option, @@ -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, } diff --git a/codex-rs/exec-server/src/server/handler/tests.rs b/codex-rs/exec-server/src/server/handler/tests.rs index 5b6c9074f..0bd78fffb 100644 --- a/codex-rs/exec-server/src/server/handler/tests.rs +++ b/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"); diff --git a/codex-rs/exec-server/tests/exec_process.rs b/codex-rs/exec-server/tests/exec_process.rs index 425974a15..689869a6b 100644 --- a/codex-rs/exec-server/tests/exec_process.rs +++ b/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(), diff --git a/codex-rs/exec-server/tests/process.rs b/codex-rs/exec-server/tests/process.rs index 4926e6088..c210b5a9c 100644 --- a/codex-rs/exec-server/tests/process.rs +++ b/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") } );