diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 4e282c8fd3..c4526e13bd 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -10,6 +10,7 @@ use arc_swap::ArcSwap; use codex_app_server_protocol::JSONRPCNotification; use serde_json::Value; use tokio::sync::Mutex; +use tokio::sync::OnceCell; use tokio::sync::mpsc; use tokio::sync::watch; @@ -174,6 +175,37 @@ pub struct ExecServerClient { inner: Arc, } +#[derive(Clone, Debug)] +pub(crate) struct LazyRemoteExecServerClient { + websocket_url: String, + client: Arc>, +} + +impl LazyRemoteExecServerClient { + pub(crate) fn new(websocket_url: String) -> Self { + Self { + websocket_url, + client: Arc::new(OnceCell::new()), + } + } + + pub(crate) async fn get(&self) -> Result { + self.client + .get_or_try_init(|| async { + ExecServerClient::connect_websocket(RemoteExecServerConnectArgs { + websocket_url: self.websocket_url.clone(), + client_name: "codex-environment".to_string(), + connect_timeout: Duration::from_secs(5), + initialize_timeout: Duration::from_secs(5), + resume_session_id: None, + }) + .await + }) + .await + .cloned() + } +} + #[derive(Debug, thiserror::Error)] pub enum ExecServerError { #[error("failed to spawn exec-server: {0}")] diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 634b2e7c79..2452144001 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -1,13 +1,9 @@ use std::collections::HashMap; use std::sync::Arc; -use std::time::Duration; -use tokio::sync::OnceCell; - -use crate::ExecServerClient; use crate::ExecServerError; use crate::ExecServerRuntimePaths; -use crate::RemoteExecServerConnectArgs; +use crate::client::LazyRemoteExecServerClient; use crate::file_system::ExecutorFileSystem; use crate::local_file_system::LocalFileSystem; use crate::local_process::LocalProcess; @@ -37,37 +33,6 @@ pub struct EnvironmentManagerArgs { pub local_runtime_paths: Option, } -#[derive(Clone, Debug)] -pub(crate) struct LazyRemoteExecServerClient { - websocket_url: String, - client: Arc>, -} - -impl LazyRemoteExecServerClient { - fn new(websocket_url: String) -> Self { - Self { - websocket_url, - client: Arc::new(OnceCell::new()), - } - } - - pub(crate) async fn get(&self) -> Result { - self.client - .get_or_try_init(|| async { - ExecServerClient::connect_websocket(RemoteExecServerConnectArgs { - websocket_url: self.websocket_url.clone(), - client_name: "codex-environment".to_string(), - connect_timeout: Duration::from_secs(5), - initialize_timeout: Duration::from_secs(5), - resume_session_id: None, - }) - .await - }) - .await - .cloned() - } -} - impl Default for EnvironmentManager { fn default() -> Self { Self::new(EnvironmentManagerArgs::default()) diff --git a/codex-rs/exec-server/src/remote_file_system.rs b/codex-rs/exec-server/src/remote_file_system.rs index 1b4150fe5e..02a5e28836 100644 --- a/codex-rs/exec-server/src/remote_file_system.rs +++ b/codex-rs/exec-server/src/remote_file_system.rs @@ -14,7 +14,7 @@ use crate::FileSystemResult; use crate::FileSystemSandboxContext; use crate::ReadDirectoryEntry; use crate::RemoveOptions; -use crate::environment::LazyRemoteExecServerClient; +use crate::client::LazyRemoteExecServerClient; use crate::protocol::FsCopyParams; use crate::protocol::FsCreateDirectoryParams; use crate::protocol::FsGetMetadataParams; diff --git a/codex-rs/exec-server/src/remote_process.rs b/codex-rs/exec-server/src/remote_process.rs index 19828d9d69..d8d06735cd 100644 --- a/codex-rs/exec-server/src/remote_process.rs +++ b/codex-rs/exec-server/src/remote_process.rs @@ -9,8 +9,8 @@ use crate::ExecProcess; use crate::ExecProcessEventReceiver; use crate::ExecServerError; use crate::StartedExecProcess; +use crate::client::LazyRemoteExecServerClient; use crate::client::Session; -use crate::environment::LazyRemoteExecServerClient; use crate::protocol::ExecParams; use crate::protocol::ReadResponse; use crate::protocol::WriteResponse;