refactor: keep direct Noise client independently landable

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 6a163833f9
commit 2488836aca
8 changed files with 57 additions and 170 deletions
-7
View File
@@ -47,14 +47,7 @@ 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;
-30
View File
@@ -23,8 +23,6 @@ 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;
@@ -123,16 +121,6 @@ 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 {
@@ -155,23 +143,6 @@ 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,
@@ -266,7 +237,6 @@ impl LazyRemoteExecServerClient {
if matches!(
&self.transport_params,
ExecServerTransportParams::WebSocketUrl { .. }
| ExecServerTransportParams::NoiseRendezvous { .. }
) =>
{
ExecServerClient::connect_for_transport(self.transport_params.clone()).await?
+11 -48
View File
@@ -1,6 +1,5 @@
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use futures::future::BoxFuture;
@@ -34,6 +33,11 @@ pub struct RemoteExecServerConnectArgs {
}
/// Registry-authorized material for one Noise rendezvous connection attempt.
///
/// Treat this as an atomic, single-use bundle. The URL authorization, executor
/// registration, pinned executor key, and harness-key authorization describe one
/// physical connection attempt and must not be mixed with values from another
/// registry response.
pub struct NoiseRendezvousConnectBundle {
pub websocket_url: String,
pub environment_id: String,
@@ -58,6 +62,10 @@ impl std::fmt::Debug for NoiseRendezvousConnectBundle {
}
/// Connection arguments for an authenticated Noise rendezvous exec-server.
///
/// `harness_identity` identifies the logical harness endpoint and may be reused
/// across reconnects. In contrast, callers must supply a fresh
/// [`NoiseRendezvousConnectBundle`] for each physical connection attempt.
pub struct NoiseRendezvousConnectArgs {
pub bundle: NoiseRendezvousConnectBundle,
pub harness_identity: NoiseChannelIdentity,
@@ -80,20 +88,6 @@ impl std::fmt::Debug for NoiseRendezvousConnectArgs {
}
}
/// 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 {
@@ -113,16 +107,13 @@ pub(crate) struct StdioExecServerCommand {
}
/// Parameters used to connect to a remote exec-server environment.
#[derive(Clone)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum ExecServerTransportParams {
WebSocketUrl {
websocket_url: String,
connect_timeout: Duration,
initialize_timeout: Duration,
},
NoiseRendezvous {
provider: SharedNoiseRendezvousConnectProvider,
},
#[allow(dead_code)]
StdioCommand {
command: StdioExecServerCommand,
@@ -130,35 +121,6 @@ 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 {
@@ -169,6 +131,7 @@ impl ExecServerTransportParams {
}
}
/// Removes URL query and fragment data before a rendezvous URL reaches logs or errors.
pub(crate) fn redacted_websocket_url(websocket_url: &str) -> String {
match url::Url::parse(websocket_url) {
Ok(mut url) => {
+6 -9
View File
@@ -44,15 +44,6 @@ 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,
@@ -94,6 +85,12 @@ impl ExecServerClient {
Self::connect(connection, args.into()).await
}
/// Connects to one exec-server through an authenticated, encrypted rendezvous stream.
///
/// This method pins the executor's Noise public key from `args`, completes
/// the encrypted channel before starting JSON-RPC, and uses the rendezvous
/// websocket only as a ciphertext transport. Callers are responsible for
/// obtaining a fresh atomic connect bundle for every physical connection.
pub async fn connect_noise_rendezvous(
args: NoiseRendezvousConnectArgs,
) -> Result<Self, ExecServerError> {
-30
View File
@@ -9,7 +9,6 @@ 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;
@@ -283,34 +282,6 @@ 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.
@@ -449,7 +420,6 @@ 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());
+27 -44
View File
@@ -48,7 +48,7 @@ struct EnvironmentToml {
initialize_timeout_sec: Option<Duration>,
}
#[derive(Clone, Debug)]
#[derive(Clone, Debug, PartialEq, Eq)]
struct TomlEnvironmentProvider {
default: EnvironmentDefault,
include_local: bool,
@@ -577,26 +577,18 @@ mod tests {
)
.expect("provider");
let ExecServerTransportParams::StdioCommand {
command,
initialize_timeout,
} = &provider.environments[0].1
else {
panic!("expected stdio transport");
};
assert_eq!(
command,
&StdioExecServerCommand {
program: "ssh".to_string(),
args: Vec::new(),
env: HashMap::new(),
cwd: Some(config_dir.path().join("workspace")),
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,
}
);
assert_eq!(
*initialize_timeout,
DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT
);
}
#[test]
@@ -622,35 +614,26 @@ 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!(
command,
&StdioExecServerCommand {
program: "ssh".to_string(),
args: Vec::new(),
env: HashMap::new(),
cwd: None,
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),
}
);
assert_eq!(*initialize_timeout, Duration::from_secs(56));
}
#[test]
-2
View File
@@ -35,9 +35,7 @@ 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;
@@ -2,10 +2,23 @@ mod harness;
mod message_framing;
mod ordered_ciphertext;
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
use crate::ExecServerError;
pub(crate) use harness::noise_harness_connection_from_websocket;
// This bounds allocation in tungstenite before protobuf and Noise record
// validation run. It comfortably fits one maximum Noise record plus metadata.
const MAX_NOISE_RELAY_WEBSOCKET_MESSAGE_SIZE: usize = 256 * 1024;
/// Return the websocket limits required by every Noise relay endpoint.
pub(crate) fn noise_relay_websocket_config() -> WebSocketConfig {
WebSocketConfig::default()
.max_frame_size(Some(MAX_NOISE_RELAY_WEBSOCKET_MESSAGE_SIZE))
.max_message_size(Some(MAX_NOISE_RELAY_WEBSOCKET_MESSAGE_SIZE))
}
fn take_next_sequence(next_seq: &mut u32) -> Result<u32, ExecServerError> {
// Never wrap: relay sequence is the explicit ordering key for an implicit
// Noise nonce. Reusing zero after u32::MAX would be ambiguous and unsafe.