feat: expose Noise exec-server client API

Co-authored-by: Codex noreply@openai.com
This commit is contained in:
viyatb-oai
2026-06-09 16:16:30 -07:00
co-authored by Codex noreply@openai.com
parent 95d74efb58
commit 9dbd96efa5
8 changed files with 304 additions and 28 deletions
+7
View File
@@ -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;
+1
View File
@@ -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]
+32
View File
@@ -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<RemoteExecServerConnectArgs> for ExecServerClientConnectOptions {
}
}
impl From<NoiseRendezvousConnectArgs> 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<StdioExecServerConnectArgs> 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<u64>,
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 {
+108 -1
View File
@@ -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<String>,
}
/// 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", &"<redacted>")
.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<String>,
}
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", &"<redacted>")
.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<NoiseRendezvousConnectArgs, ExecServerError>>;
}
pub type SharedNoiseRendezvousConnectProvider = Arc<dyn NoiseRendezvousConnectProvider>;
/// 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(_) => "<redacted websocket url>".to_string(),
}
}
/// Sends HTTP requests through a runtime-selected transport.
///
/// This is the HTTP capability counterpart to [`crate::ExecBackend`]. Callers
@@ -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<Self, ExecServerError> {
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<Self, ExecServerError> {
+30
View File
@@ -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());
+44 -27
View File
@@ -48,7 +48,7 @@ struct EnvironmentToml {
initialize_timeout_sec: Option<Duration>,
}
#[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]
+4
View File
@@ -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;