From 9dbd96efa52b52878ed31841e22f015d50a49a6f Mon Sep 17 00:00:00 2001 From: viyatb-oai Date: Wed, 3 Jun 2026 12:25:08 -0700 Subject: [PATCH] feat: expose Noise exec-server client API Co-authored-by: Codex noreply@openai.com --- codex-rs/core-api/src/lib.rs | 7 ++ codex-rs/exec-server/Cargo.toml | 1 + codex-rs/exec-server/src/client.rs | 32 ++++++ codex-rs/exec-server/src/client_api.rs | 109 ++++++++++++++++++- codex-rs/exec-server/src/client_transport.rs | 78 +++++++++++++ codex-rs/exec-server/src/environment.rs | 30 +++++ codex-rs/exec-server/src/environment_toml.rs | 71 +++++++----- codex-rs/exec-server/src/lib.rs | 4 + 8 files changed, 304 insertions(+), 28 deletions(-) diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index 47452eb032..8d242e1b0d 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -47,7 +47,14 @@ pub use codex_core::resolve_installation_id; pub use codex_core::skills::SkillsManager; pub use codex_core::thread_store_from_config; pub use codex_exec_server::EnvironmentManager; +pub use codex_exec_server::ExecServerError; pub use codex_exec_server::ExecServerRuntimePaths; +pub use codex_exec_server::NoiseChannelIdentity; +pub use codex_exec_server::NoiseChannelPublicKey; +pub use codex_exec_server::NoiseRendezvousConnectArgs; +pub use codex_exec_server::NoiseRendezvousConnectBundle; +pub use codex_exec_server::NoiseRendezvousConnectProvider; +pub use codex_exec_server::SharedNoiseRendezvousConnectProvider; pub use codex_extension_api::empty_extension_registry; pub use codex_features::Feature; pub use codex_features::Features; diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index e4cdf800fe..e8821ea3f9 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/codex-rs/exec-server/Cargo.toml @@ -49,6 +49,7 @@ tokio = { workspace = true, features = [ tokio-util = { workspace = true, features = ["rt"] } tokio-tungstenite = { workspace = true } tracing = { workspace = true } +url = { workspace = true } uuid = { workspace = true, features = ["v4"] } [dev-dependencies] diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 26b3e46e88..6d05accfc4 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -23,6 +23,8 @@ use crate::ProcessId; use crate::client_api::ExecServerClientConnectOptions; use crate::client_api::ExecServerTransportParams; use crate::client_api::HttpClient; +use crate::client_api::NoiseRendezvousConnectArgs; +use crate::client_api::NoiseRendezvousConnectBundle; use crate::client_api::RemoteExecServerConnectArgs; use crate::client_api::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; @@ -121,6 +123,16 @@ impl From for ExecServerClientConnectOptions { } } +impl From for ExecServerClientConnectOptions { + fn from(value: NoiseRendezvousConnectArgs) -> Self { + Self { + client_name: value.client_name, + initialize_timeout: value.initialize_timeout, + resume_session_id: value.resume_session_id, + } + } +} + impl From for ExecServerClientConnectOptions { fn from(value: StdioExecServerConnectArgs) -> Self { Self { @@ -143,6 +155,23 @@ impl RemoteExecServerConnectArgs { } } +impl NoiseRendezvousConnectArgs { + pub fn new( + bundle: NoiseRendezvousConnectBundle, + harness_identity: crate::NoiseChannelIdentity, + client_name: String, + ) -> Self { + Self { + bundle, + harness_identity, + client_name, + connect_timeout: CONNECT_TIMEOUT, + initialize_timeout: INITIALIZE_TIMEOUT, + resume_session_id: None, + } + } +} + pub(crate) struct SessionState { wake_tx: watch::Sender, events: ExecProcessEventLog, @@ -237,6 +266,7 @@ impl LazyRemoteExecServerClient { if matches!( &self.transport_params, ExecServerTransportParams::WebSocketUrl { .. } + | ExecServerTransportParams::NoiseRendezvous { .. } ) => { ExecServerClient::connect_for_transport(self.transport_params.clone()).await? @@ -329,6 +359,8 @@ pub enum ExecServerError { EnvironmentRegistryAuth(String), #[error("environment registry request failed: {0}")] EnvironmentRegistryRequest(#[from] reqwest::Error), + #[error(transparent)] + NoiseChannel(#[from] crate::NoiseChannelError), } impl ExecServerClient { diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index 899863723f..a6251a8757 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::path::PathBuf; +use std::sync::Arc; use std::time::Duration; use futures::future::BoxFuture; @@ -8,6 +9,8 @@ use crate::ExecServerError; use crate::HttpRequestParams; use crate::HttpRequestResponse; use crate::HttpResponseBodyStream; +use crate::NoiseChannelIdentity; +use crate::NoiseChannelPublicKey; pub(crate) const DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); pub(crate) const DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10); @@ -30,6 +33,67 @@ pub struct RemoteExecServerConnectArgs { pub resume_session_id: Option, } +/// Registry-authorized material for one Noise rendezvous connection attempt. +pub struct NoiseRendezvousConnectBundle { + pub websocket_url: String, + pub environment_id: String, + pub executor_registration_id: String, + pub executor_public_key: NoiseChannelPublicKey, + pub harness_key_authorization: String, +} + +impl std::fmt::Debug for NoiseRendezvousConnectBundle { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("NoiseRendezvousConnectBundle") + .field( + "websocket_url", + &redacted_websocket_url(&self.websocket_url), + ) + .field("environment_id", &self.environment_id) + .field("executor_registration_id", &self.executor_registration_id) + .field("executor_public_key", &self.executor_public_key) + .field("harness_key_authorization", &"") + .finish() + } +} + +/// Connection arguments for an authenticated Noise rendezvous exec-server. +pub struct NoiseRendezvousConnectArgs { + pub bundle: NoiseRendezvousConnectBundle, + pub harness_identity: NoiseChannelIdentity, + pub client_name: String, + pub connect_timeout: Duration, + pub initialize_timeout: Duration, + pub resume_session_id: Option, +} + +impl std::fmt::Debug for NoiseRendezvousConnectArgs { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("NoiseRendezvousConnectArgs") + .field("bundle", &self.bundle) + .field("harness_identity", &"") + .field("client_name", &self.client_name) + .field("connect_timeout", &self.connect_timeout) + .field("initialize_timeout", &self.initialize_timeout) + .field("resume_session_id", &self.resume_session_id) + .finish() + } +} + +/// Supplies fresh registry-authorized material for Noise rendezvous connections. +/// +/// Implementations must preserve one endpoint-local harness identity while +/// refreshing short-lived registry material for every physical connection attempt. +pub trait NoiseRendezvousConnectProvider: Send + Sync { + /// Environment ID this provider is authorized to connect to. + fn environment_id(&self) -> &str; + + /// Returns a fresh atomic bundle for one physical connection attempt. + fn connect_args(&self) -> BoxFuture<'_, Result>; +} + +pub type SharedNoiseRendezvousConnectProvider = Arc; + /// Stdio connection arguments for a command-backed exec-server. #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct StdioExecServerConnectArgs { @@ -49,13 +113,16 @@ pub(crate) struct StdioExecServerCommand { } /// Parameters used to connect to a remote exec-server environment. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Clone)] pub(crate) enum ExecServerTransportParams { WebSocketUrl { websocket_url: String, connect_timeout: Duration, initialize_timeout: Duration, }, + NoiseRendezvous { + provider: SharedNoiseRendezvousConnectProvider, + }, #[allow(dead_code)] StdioCommand { command: StdioExecServerCommand, @@ -63,6 +130,35 @@ pub(crate) enum ExecServerTransportParams { }, } +impl std::fmt::Debug for ExecServerTransportParams { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::WebSocketUrl { + websocket_url, + connect_timeout, + initialize_timeout, + } => f + .debug_struct("WebSocketUrl") + .field("websocket_url", websocket_url) + .field("connect_timeout", connect_timeout) + .field("initialize_timeout", initialize_timeout) + .finish(), + Self::NoiseRendezvous { provider } => f + .debug_struct("NoiseRendezvous") + .field("environment_id", &provider.environment_id()) + .finish(), + Self::StdioCommand { + command, + initialize_timeout, + } => f + .debug_struct("StdioCommand") + .field("command", command) + .field("initialize_timeout", initialize_timeout) + .finish(), + } + } +} + impl ExecServerTransportParams { pub(crate) fn websocket_url(websocket_url: String) -> Self { Self::WebSocketUrl { @@ -73,6 +169,17 @@ impl ExecServerTransportParams { } } +pub(crate) fn redacted_websocket_url(websocket_url: &str) -> String { + match url::Url::parse(websocket_url) { + Ok(mut url) => { + url.set_query(None); + url.set_fragment(None); + url.to_string() + } + Err(_) => "".to_string(), + } +} + /// Sends HTTP requests through a runtime-selected transport. /// /// This is the HTTP capability counterpart to [`crate::ExecBackend`]. Callers diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 4bdc09a80e..ade4c0689d 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -4,6 +4,7 @@ use tokio::io::BufReader; use tokio::process::Command; use tokio::time::timeout; use tokio_tungstenite::connect_async; +use tokio_tungstenite::connect_async_with_config; use tracing::debug; use tracing::warn; @@ -11,10 +12,15 @@ use codex_utils_rustls_provider::ensure_rustls_crypto_provider; use crate::ExecServerClient; use crate::ExecServerError; +use crate::client_api::NoiseRendezvousConnectArgs; +use crate::client_api::NoiseRendezvousConnectBundle; use crate::client_api::RemoteExecServerConnectArgs; use crate::client_api::StdioExecServerCommand; use crate::client_api::StdioExecServerConnectArgs; +use crate::client_api::redacted_websocket_url; use crate::connection::JsonRpcConnection; +use crate::noise_relay::noise_harness_connection_from_websocket; +use crate::noise_relay::noise_relay_websocket_config; use crate::relay::harness_connection_from_websocket; const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment"; @@ -38,6 +44,15 @@ impl ExecServerClient { }) .await } + crate::client_api::ExecServerTransportParams::NoiseRendezvous { provider } => { + let args = provider.connect_args().await?; + if args.bundle.environment_id != provider.environment_id() { + return Err(ExecServerError::Protocol( + "Noise rendezvous provider returned a different environment id".to_string(), + )); + } + Self::connect_noise_rendezvous(args).await + } crate::client_api::ExecServerTransportParams::StdioCommand { command, initialize_timeout, @@ -79,6 +94,69 @@ impl ExecServerClient { Self::connect(connection, args.into()).await } + pub async fn connect_noise_rendezvous( + args: NoiseRendezvousConnectArgs, + ) -> Result { + ensure_rustls_crypto_provider(); + // This connect call owns the complete registry-issued bundle. Move each + // sensitive value into the transport task exactly once rather than + // leaving extra copies of the harness authorization or endpoint identity + // alive in `args` after the handshake starts. + let NoiseRendezvousConnectArgs { + bundle, + harness_identity, + client_name, + connect_timeout, + initialize_timeout, + resume_session_id, + } = args; + let NoiseRendezvousConnectBundle { + websocket_url, + environment_id, + executor_registration_id, + executor_public_key, + harness_key_authorization, + } = bundle; + let diagnostic_url = redacted_websocket_url(&websocket_url); + let (stream, _) = timeout( + connect_timeout, + connect_async_with_config( + websocket_url.as_str(), + Some(noise_relay_websocket_config()), + /*disable_nagle*/ false, + ), + ) + .await + .map_err(|_| ExecServerError::WebSocketConnectTimeout { + url: diagnostic_url.clone(), + timeout: connect_timeout, + })? + .map_err(|source| ExecServerError::WebSocketConnect { + url: diagnostic_url.clone(), + source, + })?; + + let connection_label = format!("Noise exec-server rendezvous websocket {diagnostic_url}"); + let connection = noise_harness_connection_from_websocket( + stream, + connection_label, + environment_id, + executor_registration_id, + harness_identity, + executor_public_key, + harness_key_authorization, + ); + Self::connect( + connection, + crate::client_api::ExecServerClientConnectOptions { + client_name, + initialize_timeout, + resume_session_id, + }, + ) + .await + } + pub(crate) async fn connect_stdio_command( args: StdioExecServerConnectArgs, ) -> Result { diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 1dfd49d3fd..2ad29d60f1 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -9,6 +9,7 @@ use crate::ExecServerError; use crate::ExecServerRuntimePaths; use crate::ExecutorFileSystem; use crate::HttpClient; +use crate::SharedNoiseRendezvousConnectProvider; use crate::client::LazyRemoteExecServerClient; use crate::client::http_client::ReqwestHttpClient; use crate::client_api::ExecServerTransportParams; @@ -282,6 +283,34 @@ impl EnvironmentManager { .insert(environment_id, Arc::new(environment)); Ok(()) } + + /// Adds or replaces a named remote environment that connects through an + /// authenticated, end-to-end encrypted rendezvous stream. + pub fn upsert_noise_environment( + &self, + environment_id: String, + provider: SharedNoiseRendezvousConnectProvider, + ) -> Result<(), ExecServerError> { + if environment_id.is_empty() { + return Err(ExecServerError::Protocol( + "environment id cannot be empty".to_string(), + )); + } + if environment_id != provider.environment_id() { + return Err(ExecServerError::Protocol( + "Noise environment id does not match connection provider".to_string(), + )); + } + let environment = Environment::remote_with_transport( + ExecServerTransportParams::NoiseRendezvous { provider }, + self.local_runtime_paths.clone(), + ); + self.environments + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(environment_id, Arc::new(environment)); + Ok(()) + } } /// Concrete execution/filesystem environment selected for a session. @@ -420,6 +449,7 @@ impl Environment { websocket_url: exec_server_url, .. } => Some(exec_server_url.clone()), + ExecServerTransportParams::NoiseRendezvous { .. } => None, ExecServerTransportParams::StdioCommand { .. } => None, }; let client = LazyRemoteExecServerClient::new(remote_transport.clone()); diff --git a/codex-rs/exec-server/src/environment_toml.rs b/codex-rs/exec-server/src/environment_toml.rs index 26c178b5b8..daa024b4f5 100644 --- a/codex-rs/exec-server/src/environment_toml.rs +++ b/codex-rs/exec-server/src/environment_toml.rs @@ -48,7 +48,7 @@ struct EnvironmentToml { initialize_timeout_sec: Option, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug)] struct TomlEnvironmentProvider { default: EnvironmentDefault, include_local: bool, @@ -577,18 +577,26 @@ mod tests { ) .expect("provider"); + let ExecServerTransportParams::StdioCommand { + command, + initialize_timeout, + } = &provider.environments[0].1 + else { + panic!("expected stdio transport"); + }; assert_eq!( - provider.environments[0].1, - ExecServerTransportParams::StdioCommand { - command: StdioExecServerCommand { - program: "ssh".to_string(), - args: Vec::new(), - env: HashMap::new(), - cwd: Some(config_dir.path().join("workspace")), - }, - initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT, + command, + &StdioExecServerCommand { + program: "ssh".to_string(), + args: Vec::new(), + env: HashMap::new(), + cwd: Some(config_dir.path().join("workspace")), } ); + assert_eq!( + *initialize_timeout, + DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT + ); } #[test] @@ -614,26 +622,35 @@ mod tests { }) .expect("provider"); + let ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout, + initialize_timeout, + } = &provider.environments[0].1 + else { + panic!("expected websocket transport"); + }; + assert_eq!(websocket_url, "ws://127.0.0.1:8765"); + assert_eq!(*connect_timeout, Duration::from_secs(12)); + assert_eq!(*initialize_timeout, Duration::from_secs(34)); + + let ExecServerTransportParams::StdioCommand { + command, + initialize_timeout, + } = &provider.environments[1].1 + else { + panic!("expected stdio transport"); + }; assert_eq!( - provider.environments[0].1, - ExecServerTransportParams::WebSocketUrl { - websocket_url: "ws://127.0.0.1:8765".to_string(), - connect_timeout: Duration::from_secs(12), - initialize_timeout: Duration::from_secs(34), - } - ); - assert_eq!( - provider.environments[1].1, - ExecServerTransportParams::StdioCommand { - command: StdioExecServerCommand { - program: "ssh".to_string(), - args: Vec::new(), - env: HashMap::new(), - cwd: None, - }, - initialize_timeout: Duration::from_secs(56), + command, + &StdioExecServerCommand { + program: "ssh".to_string(), + args: Vec::new(), + env: HashMap::new(), + cwd: None, } ); + assert_eq!(*initialize_timeout, Duration::from_secs(56)); } #[test] diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index c786c1a073..f60cac4f24 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -33,7 +33,11 @@ pub use client::http_client::HttpResponseBodyStream; pub use client::http_client::ReqwestHttpClient; pub use client_api::ExecServerClientConnectOptions; pub use client_api::HttpClient; +pub use client_api::NoiseRendezvousConnectArgs; +pub use client_api::NoiseRendezvousConnectBundle; +pub use client_api::NoiseRendezvousConnectProvider; pub use client_api::RemoteExecServerConnectArgs; +pub use client_api::SharedNoiseRendezvousConnectProvider; pub use codex_file_system::CopyOptions; pub use codex_file_system::CreateDirectoryOptions; pub use codex_file_system::ExecutorFileSystem;