mirror of
https://github.com/openai/codex.git
synced 2026-11-05 20:10:33 +00:00
feat: expose Noise exec-server client API
Co-authored-by: Codex noreply@openai.com
This commit is contained in:
co-authored by
Codex noreply@openai.com
parent
95d74efb58
commit
9dbd96efa5
@@ -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;
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user