openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 17:28:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Route exec-server WebSockets through configured proxies (#35056)

Browse Source 
## Why

Remote environment connections need to honor Codex's effective outbound proxy policy, including when a rendezvous connection reconnects.

## What changed

- Pass the configured `HttpClientFactory` into remote environment transports and use `WebSocketConnector` for exec-server and rendezvous WebSockets.
- Resolve proxy routes asynchronously so these connections can use the configured system proxy.
- Add connector options that preserve Tungstenite's default TLS behavior and enable `TCP_NODELAY` for latency-sensitive rendezvous traffic.

## Testing

- Verify prepared remote environments connect through a configured system proxy.
- Verify initial and reconnected encrypted relay peers use the system proxy.
- Cover default TLS selection and opt-in `TCP_NODELAY` behavior in the WebSocket client.

GitOrigin-RevId: 8a8da2116e37cb3a891269d0c0b037986fecdd3c
This commit is contained in:
Celia Chen
2026-07-24 00:16:42 +00:00
committed by copyberry
parent fb4e6ba2f4
commit 94ebae725e
24 changed files with 727 additions and 226 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+1
View File
@@ -2929,6 +2929,7 @@ dependencies = [
"codex-utils-path-uri",
"codex-utils-pty",
"codex-utils-rustls-provider",
"codex-websocket-client",
"ctor 0.6.3",
"futures",
"http 1.4.0",
codex-rs/exec-server/Cargo.toml
+1
View File
@@ -28,6 +28,7 @@ codex-utils-absolute-path = { workspace = true }
codex-utils-path-uri = { workspace = true }
codex-utils-pty = { workspace = true }
codex-utils-rustls-provider = { workspace = true }
codex-websocket-client = { workspace = true }
futures = { workspace = true }
http = { workspace = true }
reqwest = { workspace = true, features = ["json", "rustls-tls", "stream"] }
codex-rs/exec-server/src/client.rs
+73 -34
View File
@@ -112,6 +112,7 @@ use crate::protocol::WriteParams;
use crate::protocol::WriteResponse;
use crate::rpc::RpcCallError;
use crate::rpc::RpcClient;
use codex_http_client::HttpClientFactory;
pub(crate) mod http_client;
#[path = "client_recovery.rs"]
@@ -157,13 +158,18 @@ impl From<StdioExecServerConnectArgs> for ExecServerClientConnectOptions {
}
impl RemoteExecServerConnectArgs {
pub fn new(websocket_url: String, client_name: String) -> Self {
pub fn new(
websocket_url: String,
client_name: String,
http_client_factory: HttpClientFactory,
) -> Self {
Self {
websocket_url,
client_name,
connect_timeout: CONNECT_TIMEOUT,
initialize_timeout: INITIALIZE_TIMEOUT,
resume_session_id: None,
http_client_factory,
}
}
}
@@ -288,7 +294,8 @@ type ConnectionAttempt = OnceCell<ConnectionResult>;
#[derive(Clone)]
pub(crate) struct LazyRemoteExecServerClient {
transport_params: ExecServerTransportParams,
pub(crate) transport_params: ExecServerTransportParams,
http_client_factory: HttpClientFactory,
recovery_policy: RecoveryPolicy,
// Saves the first startup result so callers share it and failures remain final.
startup: Arc<ConnectionAttempt>,
@@ -299,9 +306,13 @@ pub(crate) struct LazyRemoteExecServerClient {
}
impl LazyRemoteExecServerClient {
pub(crate) fn new(transport_params: ExecServerTransportParams) -> Self {
pub(crate) fn new(
transport_params: ExecServerTransportParams,
http_client_factory: HttpClientFactory,
) -> Self {
Self {
transport_params,
http_client_factory,
recovery_policy: RecoveryPolicy::Wait,
startup: Arc::new(ConnectionAttempt::new()),
current_client: Arc::new(StdMutex::new(None)),
@@ -499,9 +510,12 @@ impl LazyRemoteExecServerClient {
}
async fn connect_once(&self) -> ConnectionResult {
let result = ExecServerClient::connect_for_transport(self.transport_params.clone())
.await
.map_err(Arc::new);
let result = ExecServerClient::connect_for_transport(
self.transport_params.clone(),
self.http_client_factory.clone(),
)
.await
.map_err(Arc::new);
if let Ok(client) = &result {
client
.attach_environment_connection_state(self.environment_connection_state_tx.clone());
@@ -547,6 +561,8 @@ pub enum ExecServerError {
#[source]
source: tokio_tungstenite::tungstenite::Error,
},
#[error("failed to configure exec-server websocket: {0}")]
WebSocketConfiguration(String),
#[error("timed out waiting for exec-server initialize handshake after {timeout:?}")]
InitializeTimedOut { timeout: Duration },
#[error("exec-server transport closed")]
@@ -1493,6 +1509,8 @@ mod tests {
use codex_exec_server_protocol::JSONRPCMessage;
use codex_exec_server_protocol::JSONRPCNotification;
use codex_exec_server_protocol::JSONRPCResponse;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use codex_utils_path_uri::PathUri;
use futures::SinkExt;
use futures::StreamExt;
@@ -1809,6 +1827,9 @@ mod tests {
},
initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
},
codex_http_client::HttpClientFactory::new(
codex_http_client::OutboundProxyPolicy::ReqwestDefault,
),
)
.await
.expect("stdio transport should connect");
@@ -2221,11 +2242,14 @@ mod tests {
finish_rx.await.expect("test should finish");
});
let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl {
websocket_url,
connect_timeout: Duration::from_secs(1),
initialize_timeout: Duration::from_secs(1),
});
let client = LazyRemoteExecServerClient::new(
ExecServerTransportParams::WebSocketUrl {
websocket_url,
connect_timeout: Duration::from_secs(1),
initialize_timeout: Duration::from_secs(1),
},
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
);
let stable_client = client.get().await.expect("client should connect");
timeout(Duration::from_secs(1), resumed_rx)
.await
@@ -2328,11 +2352,14 @@ mod tests {
finish_rx.await.expect("test should finish");
});
let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl {
websocket_url,
connect_timeout: Duration::from_secs(1),
initialize_timeout: Duration::from_secs(1),
});
let client = LazyRemoteExecServerClient::new(
ExecServerTransportParams::WebSocketUrl {
websocket_url,
connect_timeout: Duration::from_secs(1),
initialize_timeout: Duration::from_secs(1),
},
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
);
let stable_client = client.get().await.expect("client should connect");
let session = stable_client
.register_session(&ProcessId::from("proc-write"))
@@ -2427,6 +2454,9 @@ mod tests {
connect_timeout: Duration::from_secs(1),
initialize_timeout: Duration::from_secs(1),
resume_session_id: Some("session-1".to_string()),
http_client_factory: codex_http_client::HttpClientFactory::new(
codex_http_client::OutboundProxyPolicy::ReqwestDefault,
),
}),
)
.await
@@ -2463,11 +2493,14 @@ mod tests {
.await
.expect("client should close after the test");
});
let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl {
websocket_url,
connect_timeout: Duration::from_secs(1),
initialize_timeout: Duration::from_secs(1),
});
let client = LazyRemoteExecServerClient::new(
ExecServerTransportParams::WebSocketUrl {
websocket_url,
connect_timeout: Duration::from_secs(1),
initialize_timeout: Duration::from_secs(1),
},
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
);
assert!(!client.startup_finished());
let _startup_task = client.start_connecting();
@@ -2489,15 +2522,18 @@ mod tests {
#[tokio::test]
async fn terminal_stdio_startup_failure_is_remembered() {
let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::StdioCommand {
command: StdioExecServerCommand {
program: "codex-missing-exec-server-for-test".to_string(),
args: Vec::new(),
env: HashMap::new(),
cwd: None,
let client = LazyRemoteExecServerClient::new(
ExecServerTransportParams::StdioCommand {
command: StdioExecServerCommand {
program: "codex-missing-exec-server-for-test".to_string(),
args: Vec::new(),
env: HashMap::new(),
cwd: None,
},
initialize_timeout: Duration::from_secs(1),
},
initialize_timeout: Duration::from_secs(1),
});
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
);
assert!(client.start_connecting().is_none());
assert!(!client.startup_finished());
@@ -2576,11 +2612,14 @@ mod tests {
.await
.expect("client should close after the test");
});
let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl {
websocket_url,
connect_timeout: Duration::from_secs(1),
initialize_timeout: Duration::from_secs(1),
});
let client = LazyRemoteExecServerClient::new(
ExecServerTransportParams::WebSocketUrl {
websocket_url,
connect_timeout: Duration::from_secs(1),
initialize_timeout: Duration::from_secs(1),
},
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
);
let initial = client.get().await.expect("startup should connect");
timeout(Duration::from_secs(1), async {
codex-rs/exec-server/src/client_api.rs
+3
View File
@@ -3,6 +3,7 @@ use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use codex_http_client::HttpClientFactory;
use futures::future::BoxFuture;
use futures::future::Shared;
use tokio::sync::oneshot;
@@ -33,6 +34,7 @@ pub struct RemoteExecServerConnectArgs {
pub connect_timeout: Duration,
pub initialize_timeout: Duration,
pub resume_session_id: Option<String>,
pub http_client_factory: HttpClientFactory,
}
/// Registry-authorized material for one Noise rendezvous connection attempt.
@@ -61,6 +63,7 @@ pub struct NoiseRendezvousConnectArgs {
pub connect_timeout: Duration,
pub initialize_timeout: Duration,
pub resume_session_id: Option<String>,
pub http_client_factory: HttpClientFactory,
}
/// Supplies fresh registry-authorized material for Noise rendezvous connections.
codex-rs/exec-server/src/client_transport.rs
+49 -17
View File
@@ -6,13 +6,14 @@ use tokio::io::AsyncBufReadExt;
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 tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tracing::debug;
use tracing::warn;
use codex_http_client::HttpClientFactory;
use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
use codex_websocket_client::WebSocketConnector;
use codex_websocket_client::WebSocketTlsMode;
use crate::ExecServerClient;
use crate::ExecServerError;
@@ -50,6 +51,7 @@ pub(crate) enum ExecServerReconnectStrategy {
client_name: String,
connect_timeout: Duration,
initialize_timeout: Duration,
http_client_factory: HttpClientFactory,
},
}
@@ -71,6 +73,7 @@ impl ExecServerReconnectStrategy {
client_name,
connect_timeout,
initialize_timeout,
http_client_factory,
} => {
let bundle = provider.connect_bundle(identity.public_key()).await?;
ExecServerClient::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs {
@@ -80,6 +83,7 @@ impl ExecServerReconnectStrategy {
connect_timeout: *connect_timeout,
initialize_timeout: *initialize_timeout,
resume_session_id: Some(session_id.to_string()),
http_client_factory: http_client_factory.clone(),
})
.await
}
@@ -93,6 +97,7 @@ impl ExecServerClient {
/// and authorization without replacing the harness identity.
pub(crate) async fn connect_for_transport(
transport_params: ExecServerTransportParams,
http_client_factory: HttpClientFactory,
) -> Result<Self, ExecServerError> {
let (transport_params, deferred_readiness) = match transport_params {
ExecServerTransportParams::Deferred(deferred) => {
@@ -130,9 +135,14 @@ impl ExecServerClient {
client_name: ENVIRONMENT_CLIENT_NAME.to_string(),
connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
http_client_factory: http_client_factory.clone(),
};
let (connection, options) =
Self::open_initial_noise_rendezvous_connection(&provider, &identity).await?;
let (connection, options) = Self::open_initial_noise_rendezvous_connection(
&provider,
&identity,
http_client_factory,
)
.await?;
return Self::connect_with_recovery(connection, options, Some(reconnect_strategy))
.await;
}
@@ -155,6 +165,7 @@ impl ExecServerClient {
connect_timeout,
initialize_timeout,
resume_session_id: None,
http_client_factory,
})
.await
}
@@ -162,6 +173,7 @@ impl ExecServerClient {
async fn open_initial_noise_rendezvous_connection(
provider: &Arc<dyn NoiseRendezvousConnectProvider>,
identity: &NoiseChannelIdentity,
http_client_factory: HttpClientFactory,
) -> Result<(JsonRpcConnection, ExecServerClientConnectOptions), ExecServerError> {
let open_connection = |bundle: NoiseRendezvousConnectBundle| {
Self::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs {
@@ -171,6 +183,7 @@ impl ExecServerClient {
connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,
resume_session_id: None,
http_client_factory: http_client_factory.clone(),
})
};
let bundle = provider.connect_bundle(identity.public_key()).await?;
@@ -212,16 +225,34 @@ impl ExecServerClient {
ensure_rustls_crypto_provider();
let websocket_url = args.websocket_url.clone();
let connect_timeout = args.connect_timeout;
let (stream, _) = timeout(connect_timeout, connect_async(websocket_url.as_str()))
.await
.map_err(|_| ExecServerError::WebSocketConnectTimeout {
url: websocket_url.clone(),
timeout: connect_timeout,
})?
let request = websocket_url
.as_str()
.into_client_request()
.map_err(|source| ExecServerError::WebSocketConnect {
url: websocket_url.clone(),
source,
})?;
let connector = WebSocketConnector::new_with_tls_mode(
&args.http_client_factory,
WebSocketTlsMode::TungsteniteDefault,
)
.map_err(|error| ExecServerError::WebSocketConfiguration(error.to_string()))?;
let (stream, _) = timeout(
connect_timeout,
connector.connect(
request,
tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default(),
),
)
.await
.map_err(|_| ExecServerError::WebSocketConnectTimeout {
url: websocket_url.clone(),
timeout: connect_timeout,
})?
.map_err(|source| ExecServerError::WebSocketConnect {
url: websocket_url.clone(),
source,
})?;
let connection_label = format!("exec-server websocket {websocket_url}");
let connection = if is_rendezvous_harness_url(&websocket_url) {
@@ -267,6 +298,7 @@ impl ExecServerClient {
connect_timeout,
initialize_timeout,
resume_session_id,
http_client_factory,
} = args;
let NoiseRendezvousConnectBundle {
websocket_url,
@@ -292,13 +324,13 @@ impl ExecServerClient {
.extend(current_trace_context_headers());
let (stream, _) = timeout(
connect_timeout,
connect_async_with_config(
request,
Some(noise_relay_websocket_config()),
// Rendezvous sends small, latency-sensitive frames, so avoid Nagle's coalescing delay.
/*disable_nagle*/
true,
),
WebSocketConnector::new_with_tls_mode(
&http_client_factory,
WebSocketTlsMode::TungsteniteDefault,
)
.map_err(|error| ExecServerError::WebSocketConfiguration(error.to_string()))?
.with_tcp_nodelay()
.connect(request, noise_relay_websocket_config()),
)
.await
.map_err(|_| ExecServerError::WebSocketConnectTimeout {
codex-rs/exec-server/src/client_transport_tests.rs
+8 -2
View File
@@ -99,8 +99,14 @@ async fn initial_noise_connection_refreshes_bundle_after_unauthorized_handshake(
let provider: Arc<dyn NoiseRendezvousConnectProvider> = sequence.clone();
let identity = NoiseChannelIdentity::generate()?;
let _connection =
ExecServerClient::open_initial_noise_rendezvous_connection(&provider, &identity).await?;
let _connection = ExecServerClient::open_initial_noise_rendezvous_connection(
&provider,
&identity,
codex_http_client::HttpClientFactory::new(
codex_http_client::OutboundProxyPolicy::ReqwestDefault,
),
)
.await?;
assert_eq!(
sequence.returned_urls(),
codex-rs/exec-server/src/connection.rs
+4 -2
View File
@@ -18,6 +18,7 @@ use tokio::process::Child;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::time::timeout;
#[cfg(test)]
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::tungstenite::Message;
use tracing::debug;
@@ -385,9 +386,10 @@ impl JsonRpcConnection {
}
}
pub(crate) fn from_websocket<S>(stream: WebSocketStream<S>, connection_label: String) -> Self
pub(crate) fn from_websocket<T, E>(stream: T, connection_label: String) -> Self
where
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
T: Sink<Message, Error = E> + Stream<Item = Result<Message, E>> + Unpin + Send + 'static,
E: std::fmt::Display + Send + 'static,
{
Self::from_websocket_stream(stream, connection_label, /*ping_interval*/ None)
}
codex-rs/exec-server/src/environment.rs
+72 -69
View File
@@ -287,7 +287,7 @@ impl EnvironmentManager {
} else {
None
};
for (id, environment) in environments {
for (id, transport) in environments {
if id.is_empty() {
return Err(ExecServerError::Protocol(
"environment id cannot be empty".to_string(),
@@ -298,6 +298,11 @@ impl EnvironmentManager {
"environment id `{LOCAL_ENVIRONMENT_ID}` is reserved for EnvironmentManager"
)));
}
let environment = Environment::remote_with_transport(
transport,
/*local_runtime_paths*/ None,
http_client_factory.clone(),
);
if environment_map
.insert(id.clone(), Arc::new(environment))
.is_some()
@@ -408,6 +413,7 @@ impl EnvironmentManager {
connect_timeout.unwrap_or(DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT),
),
self.local_runtime_paths.clone(),
self.http_client_factory.clone(),
));
self.insert_environment(environment_id, environment);
Ok(())
@@ -429,6 +435,7 @@ impl EnvironmentManager {
transport: ExecServerTransportParams::NoiseRendezvous { provider, identity },
})),
self.local_runtime_paths.clone(),
self.http_client_factory.clone(),
);
environment.ready_info = Some(Arc::clone(&ready_info));
let environment = Arc::new(environment);
@@ -455,6 +462,7 @@ impl EnvironmentManager {
let environment = Arc::new(Environment::remote_with_transport(
ExecServerTransportParams::NoiseRendezvous { provider, identity },
self.local_runtime_paths.clone(),
self.http_client_factory.clone(),
));
self.insert_environment(environment_id, environment);
Ok(())
@@ -639,17 +647,26 @@ impl std::fmt::Debug for Environment {
}
impl Environment {
/// Builds an environment from the raw `CODEX_EXEC_SERVER_URL` value.
/// Builds an environment using the caller's effective outbound HTTP policy.
pub fn create(
exec_server_url: Option<String>,
local_runtime_paths: ExecServerRuntimePaths,
http_client_factory: HttpClientFactory,
) -> Result<Self, ExecServerError> {
Self::create_inner(exec_server_url, Some(local_runtime_paths))
Self::create_inner(
exec_server_url,
Some(local_runtime_paths),
http_client_factory,
)
}
/// Builds a test-only environment without configured sandbox helper paths.
pub fn create_for_tests(exec_server_url: Option<String>) -> Result<Self, ExecServerError> {
Self::create_inner(exec_server_url, /*local_runtime_paths*/ None)
Self::create_inner(
exec_server_url,
/*local_runtime_paths*/ None,
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
)
}
/// Builds an environment from the raw `CODEX_EXEC_SERVER_URL` value and
@@ -657,6 +674,7 @@ impl Environment {
fn create_inner(
exec_server_url: Option<String>,
local_runtime_paths: Option<ExecServerRuntimePaths>,
http_client_factory: HttpClientFactory,
) -> Result<Self, ExecServerError> {
let (exec_server_url, disabled) = normalize_exec_server_url(exec_server_url);
if disabled {
@@ -666,13 +684,16 @@ impl Environment {
}
Ok(match exec_server_url {
Some(exec_server_url) => Self::remote_inner(exec_server_url, local_runtime_paths),
None => match local_runtime_paths {
Some(local_runtime_paths) => Self::local(
local_runtime_paths,
// This legacy constructor has no resolved application proxy configuration.
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
Some(exec_server_url) => Self::remote_with_transport(
ExecServerTransportParams::websocket_url(
exec_server_url,
DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
),
local_runtime_paths,
http_client_factory,
),
None => match local_runtime_paths {
Some(local_runtime_paths) => Self::local(local_runtime_paths, http_client_factory),
None => Self::default_for_tests(),
},
})
@@ -697,24 +718,12 @@ impl Environment {
}
}
pub(crate) fn remote_inner(
exec_server_url: String,
local_runtime_paths: Option<ExecServerRuntimePaths>,
) -> Self {
Self::remote_with_transport(
ExecServerTransportParams::websocket_url(
exec_server_url,
DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
),
local_runtime_paths,
)
}
pub(crate) fn remote_with_transport(
remote_transport: ExecServerTransportParams,
local_runtime_paths: Option<ExecServerRuntimePaths>,
http_client_factory: HttpClientFactory,
) -> Self {
let client = LazyRemoteExecServerClient::new(remote_transport);
let client = LazyRemoteExecServerClient::new(remote_transport, http_client_factory);
let exec_backend: Arc<dyn ExecBackend> = Arc::new(RemoteProcess::new(client.clone()));
let filesystem: Arc<dyn ExecutorFileSystem> =
Arc::new(RemoteFileSystem::new(client.clone()));
@@ -883,6 +892,7 @@ mod tests {
use super::noise_environment_config_from_values;
use crate::ExecServerRuntimePaths;
use crate::ProcessId;
use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
use crate::client_api::ExecServerTransportParams;
use crate::client_api::StdioExecServerCommand;
use crate::environment_provider::EnvironmentDefault;
@@ -898,6 +908,13 @@ mod tests {
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault)
}
fn prepared_websocket_environment() -> ExecServerTransportParams {
ExecServerTransportParams::websocket_url(
"ws://127.0.0.1:8765".to_string(),
DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
)
}
fn test_runtime_paths() -> ExecServerRuntimePaths {
ExecServerRuntimePaths::new(
std::env::current_exe().expect("current exe"),
@@ -960,8 +977,12 @@ mod tests {
#[tokio::test]
async fn create_local_environment_does_not_connect() {
let environment = Environment::create(/*exec_server_url*/ None, test_runtime_paths())
.expect("create environment");
let environment = Environment::create(
/*exec_server_url*/ None,
test_runtime_paths(),
legacy_http_client_factory(),
)
.expect("create environment");
assert!(!environment.is_remote());
assert!(environment.info().await.is_ok());
@@ -1050,8 +1071,7 @@ mod tests {
let snapshot = EnvironmentProviderSnapshot {
environments: vec![(
REMOTE_ENVIRONMENT_ID.to_string(),
Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string()))
.expect("remote environment"),
prepared_websocket_environment(),
)],
default: EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string()),
include_local: false,
@@ -1080,7 +1100,7 @@ mod tests {
#[tokio::test]
async fn environment_manager_rejects_empty_environment_id() {
let snapshot = EnvironmentProviderSnapshot {
environments: vec![("".to_string(), Environment::default_for_tests())],
environments: vec![("".to_string(), prepared_websocket_environment())],
default: EnvironmentDefault::Disabled,
include_local: false,
};
@@ -1102,7 +1122,7 @@ mod tests {
let snapshot = EnvironmentProviderSnapshot {
environments: vec![(
LOCAL_ENVIRONMENT_ID.to_string(),
Environment::default_for_tests(),
prepared_websocket_environment(),
)],
default: EnvironmentDefault::Disabled,
include_local: false,
@@ -1123,11 +1143,7 @@ mod tests {
#[tokio::test]
async fn environment_manager_uses_explicit_provider_default() {
let snapshot = EnvironmentProviderSnapshot {
environments: vec![(
"devbox".to_string(),
Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string()))
.expect("remote environment"),
)],
environments: vec![("devbox".to_string(), prepared_websocket_environment())],
default: EnvironmentDefault::EnvironmentId("devbox".to_string()),
include_local: true,
};
@@ -1149,11 +1165,7 @@ mod tests {
#[tokio::test]
async fn environment_manager_disables_provider_default() {
let snapshot = EnvironmentProviderSnapshot {
environments: vec![(
"devbox".to_string(),
Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string()))
.expect("remote environment"),
)],
environments: vec![("devbox".to_string(), prepared_websocket_environment())],
default: EnvironmentDefault::Disabled,
include_local: true,
};
@@ -1177,11 +1189,7 @@ mod tests {
#[tokio::test]
async fn environment_manager_rejects_unknown_provider_default() {
let snapshot = EnvironmentProviderSnapshot {
environments: vec![(
"devbox".to_string(),
Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string()))
.expect("remote environment"),
)],
environments: vec![("devbox".to_string(), prepared_websocket_environment())],
default: EnvironmentDefault::EnvironmentId("missing".to_string()),
include_local: true,
};
@@ -1356,6 +1364,7 @@ mod tests {
initialize_timeout: Duration::from_secs(1),
},
/*local_runtime_paths*/ None,
legacy_http_client_factory(),
);
assert_eq!(
@@ -1367,21 +1376,18 @@ mod tests {
#[tokio::test]
async fn environment_manager_leaves_stdio_environment_lazy() {
let environment = Environment::remote_with_transport(
ExecServerTransportParams::StdioCommand {
command: StdioExecServerCommand {
program: "codex-missing-exec-server-for-test".to_string(),
args: Vec::new(),
env: HashMap::new(),
cwd: None,
},
initialize_timeout: Duration::from_secs(1),
let transport = ExecServerTransportParams::StdioCommand {
command: StdioExecServerCommand {
program: "codex-missing-exec-server-for-test".to_string(),
args: Vec::new(),
env: HashMap::new(),
cwd: None,
},
/*local_runtime_paths*/ None,
);
initialize_timeout: Duration::from_secs(1),
};
let manager = EnvironmentManager::from_snapshot(
EnvironmentProviderSnapshot {
environments: vec![("stdio".to_string(), environment)],
environments: vec![("stdio".to_string(), transport)],
default: EnvironmentDefault::Disabled,
include_local: false,
},
@@ -1401,21 +1407,18 @@ mod tests {
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
let environment = Environment::remote_with_transport(
ExecServerTransportParams::StdioCommand {
command: StdioExecServerCommand {
program: "codex-missing-exec-server-for-test".to_string(),
args: Vec::new(),
env: HashMap::new(),
cwd: None,
},
initialize_timeout: Duration::from_secs(1),
let transport = ExecServerTransportParams::StdioCommand {
command: StdioExecServerCommand {
program: "codex-missing-exec-server-for-test".to_string(),
args: Vec::new(),
env: HashMap::new(),
cwd: None,
},
/*local_runtime_paths*/ None,
);
initialize_timeout: Duration::from_secs(1),
};
let manager = EnvironmentManager::from_snapshot(
EnvironmentProviderSnapshot {
environments: vec![("stdio".to_string(), environment)],
environments: vec![("stdio".to_string(), transport)],
default: EnvironmentDefault::Disabled,
include_local: false,
},
codex-rs/exec-server/src/environment_bootstrap.rs
+1 -2
View File
@@ -33,8 +33,7 @@ impl PreparedEnvironmentManager {
EnvironmentDefault::EnvironmentId(default_id) => snapshot
.environments
.iter()
.find(|(environment_id, _)| environment_id == default_id)
.is_some_and(|(_, environment)| environment.is_remote()),
.any(|(environment_id, _)| environment_id == default_id),
},
}
}
codex-rs/exec-server/src/environment_bootstrap_tests.rs
+29 -21
View File
@@ -5,35 +5,37 @@ use pretty_assertions::assert_eq;
use super::PreparedEnvironmentManager;
use super::PreparedEnvironmentSource;
use crate::DefaultEnvironmentProvider;
use crate::Environment;
use crate::EnvironmentConnectionState;
use crate::ExecServerRuntimePaths;
use crate::LOCAL_ENVIRONMENT_ID;
use crate::REMOTE_ENVIRONMENT_ID;
use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
use crate::client_api::ExecServerTransportParams;
use crate::environment_provider::EnvironmentDefault;
use crate::environment_provider::EnvironmentProviderSnapshot;
use crate::remote::NoiseRendezvousEnvironmentConfig;
#[test]
fn prepared_remote_environment_is_detected_without_starting_a_connection() {
let environment = Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string()))
.expect("remote environment");
let connection_state = environment
.subscribe_connection_state()
.expect("remote connection state");
fn prepared_remote_environment_is_detected_without_constructing_a_connection() {
let prepared = PreparedEnvironmentManager {
source: PreparedEnvironmentSource::Snapshot(EnvironmentProviderSnapshot {
environments: vec![(REMOTE_ENVIRONMENT_ID.to_string(), environment)],
environments: vec![(
REMOTE_ENVIRONMENT_ID.to_string(),
ExecServerTransportParams::websocket_url(
"ws://username:password@executor.example/private?token=secret".to_string(),
DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
),
)],
default: EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string()),
include_local: false,
}),
};
assert!(prepared.default_environment_is_remote());
assert_eq!(
*connection_state.borrow(),
EnvironmentConnectionState::Disconnected
);
let debug = format!("{prepared:?}");
assert!(!debug.contains("username"));
assert!(!debug.contains("password"));
assert!(!debug.contains("executor.example"));
assert!(!debug.contains("secret"));
}
#[test]
@@ -130,21 +132,27 @@ fn prepared_local_and_disabled_environments_are_not_remote() {
async fn prepared_environment_manager_builds_with_the_explicit_http_policy() {
let prepared = PreparedEnvironmentManager {
source: PreparedEnvironmentSource::Snapshot(
DefaultEnvironmentProvider::new(/*exec_server_url*/ None).snapshot_inner(),
DefaultEnvironmentProvider::new(Some("ws://127.0.0.1:8765".to_string()))
.snapshot_inner(),
),
};
let runtime_paths = ExecServerRuntimePaths::new(
std::env::current_exe().expect("current exe"),
/*codex_linux_sandbox_exe*/ None,
)
.expect("runtime paths");
let manager = prepared
.build(
Some(runtime_paths),
/*local_runtime_paths*/ None,
HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy),
)
.expect("environment manager");
assert_eq!(
manager.default_environment_id(),
Some(REMOTE_ENVIRONMENT_ID)
);
assert!(
manager
.default_environment()
.expect("remote environment")
.is_remote()
);
assert_eq!(
manager.http_client_factory().outbound_proxy_policy(),
OutboundProxyPolicy::RespectSystemProxy
codex-rs/exec-server/src/environment_provider.rs
+29 -11
View File
@@ -1,19 +1,20 @@
use std::future::Future;
use std::pin::Pin;
use crate::Environment;
use crate::ExecServerError;
use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;
use crate::client_api::ExecServerTransportParams;
use crate::environment::CODEX_EXEC_SERVER_URL_ENV_VAR;
use crate::environment::LOCAL_ENVIRONMENT_ID;
use crate::environment::REMOTE_ENVIRONMENT_ID;
/// Lists the concrete environments available to Codex.
/// Lists the remote environment transports available to Codex.
///
/// Implementations own a startup snapshot containing both the available
/// environment list in configured order and the default environment
/// selection. Providers should only return provider-owned remote environments;
/// `include_local` controls whether `EnvironmentManager` should add the local
/// environment to the snapshot.
/// environment transport list in configured order and the default environment
/// selection. Providers return transport descriptions before the effective HTTP
/// policy is available; `include_local` controls whether `EnvironmentManager`
/// should add the local environment when the snapshot is built.
pub trait EnvironmentProvider: Send + Sync {
/// Returns the provider-owned environment startup snapshot.
fn snapshot(&self) -> EnvironmentProviderFuture<'_>;
@@ -22,13 +23,24 @@ pub trait EnvironmentProvider: Send + Sync {
pub type EnvironmentProviderFuture<'a> =
Pin<Box<dyn Future<Output = Result<EnvironmentProviderSnapshot, ExecServerError>> + Send + 'a>>;
#[derive(Clone, Debug)]
#[derive(Clone)]
pub struct EnvironmentProviderSnapshot {
pub environments: Vec<(String, Environment)>,
pub(crate) environments: Vec<(String, ExecServerTransportParams)>,
pub default: EnvironmentDefault,
pub include_local: bool,
}
impl std::fmt::Debug for EnvironmentProviderSnapshot {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let environment_ids: Vec<_> = self.environments.iter().map(|(id, _)| id).collect();
f.debug_struct("EnvironmentProviderSnapshot")
.field("environments", &environment_ids)
.field("default", &self.default)
.field("include_local", &self.include_local)
.finish()
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum EnvironmentDefault {
Disabled,
@@ -59,7 +71,10 @@ impl DefaultEnvironmentProvider {
if let Some(exec_server_url) = exec_server_url {
environments.push((
REMOTE_ENVIRONMENT_ID.to_string(),
Environment::remote_inner(exec_server_url, /*local_runtime_paths*/ None),
ExecServerTransportParams::websocket_url(
exec_server_url,
DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
),
));
}
@@ -175,8 +190,11 @@ mod tests {
assert!(!include_local);
assert!(!environments.contains_key(LOCAL_ENVIRONMENT_ID));
let remote_environment = &environments[REMOTE_ENVIRONMENT_ID];
assert!(remote_environment.is_remote());
assert!(matches!(
&environments[REMOTE_ENVIRONMENT_ID],
ExecServerTransportParams::WebSocketUrl { websocket_url, .. }
if websocket_url == "ws://127.0.0.1:8765"
));
assert_eq!(
default,
EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string())
codex-rs/exec-server/src/environment_toml.rs
+9 -15
View File
@@ -8,7 +8,6 @@ use serde::Deserialize;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use crate::DefaultEnvironmentProvider;
use crate::Environment;
use crate::EnvironmentProvider;
use crate::EnvironmentProviderFuture;
use crate::ExecServerError;
@@ -94,19 +93,8 @@ impl TomlEnvironmentProvider {
}
async fn snapshot(&self) -> Result<EnvironmentProviderSnapshot, ExecServerError> {
let mut environments = Vec::with_capacity(self.environments.len());
for (id, transport_params) in &self.environments {
environments.push((
id.clone(),
Environment::remote_with_transport(
transport_params.clone(),
/*local_runtime_paths*/ None,
),
));
}
Ok(EnvironmentProviderSnapshot {
environments,
environments: self.environments.clone(),
default: self.default.clone(),
include_local: self.include_local,
})
@@ -389,8 +377,14 @@ mod tests {
assert!(include_local);
assert!(!environments.contains_key(LOCAL_ENVIRONMENT_ID));
assert!(environments["devbox"].is_remote());
assert!(environments["ssh-dev"].is_remote());
assert!(matches!(
&environments["devbox"],
ExecServerTransportParams::WebSocketUrl { .. }
));
assert!(matches!(
&environments["ssh-dev"],
ExecServerTransportParams::StdioCommand { .. }
));
assert_eq!(
default,
EnvironmentDefault::EnvironmentId("ssh-dev".to_string())
codex-rs/exec-server/src/relay.rs
+4 -6
View File
@@ -7,13 +7,10 @@ use futures::SinkExt;
use futures::Stream;
use futures::StreamExt;
use prost::Message as ProstMessage;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::task::JoinSet;
use tokio::time::timeout;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::tungstenite::Message;
use tracing::debug;
use tracing::info;
@@ -465,8 +462,8 @@ pub(crate) trait HarnessKeyValidator: Send + Sync {
/// Parsing the first Noise message authenticates the harness key. Only a
/// successful registry check turns that pending handshake into a virtual stream.
#[tracing::instrument(level = "debug", skip_all, fields(noise_side = "executor"))]
pub(crate) async fn run_multiplexed_environment<S, V>(
stream: WebSocketStream<S>,
pub(crate) async fn run_multiplexed_environment<T, E, V>(
stream: T,
processor: ConnectionProcessor,
environment_id: String,
executor_registration_id: String,
@@ -474,7 +471,8 @@ pub(crate) async fn run_multiplexed_environment<S, V>(
validator: V,
) -> RendezvousDisconnectReason
where
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
T: Sink<Message, Error = E> + Stream<Item = Result<Message, E>> + Unpin + Send + 'static,
E: std::fmt::Display + Send + 'static,
V: HarnessKeyValidator + Clone + 'static,
{
debug!(
codex-rs/exec-server/src/remote.rs
+21 -16
View File
@@ -15,15 +15,15 @@ use http::HeaderValue;
use reqwest::StatusCode;
use serde::Deserialize;
use tokio::time::sleep;
use tokio_tungstenite::MaybeTlsStream;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::connect_async_with_config;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tracing::debug;
use tracing::info;
use tracing::warn;
use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
use codex_websocket_client::WebSocketConnection;
use codex_websocket_client::WebSocketConnector;
use codex_websocket_client::WebSocketTlsMode;
use crate::EnvironmentRegistryConnectRequest;
use crate::EnvironmentRegistryConnectResponse;
@@ -514,7 +514,13 @@ pub async fn run_remote_environment(
.await?;
loop {
match connect_rendezvous(&response.url, &config.telemetry).await {
match connect_rendezvous(
&response.url,
&config.telemetry,
&config.http_client_factory,
)
.await
{
Ok(websocket) => {
backoff = Duration::from_secs(1);
let executor_registration_id = response.executor_registration_id.clone();
@@ -587,25 +593,24 @@ pub async fn run_remote_environment(
async fn connect_rendezvous(
url: &str,
telemetry: &ExecServerTelemetry,
) -> Result<
WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>,
tokio_tungstenite::tungstenite::Error,
> {
http_client_factory: &HttpClientFactory,
) -> Result<WebSocketConnection, tokio_tungstenite::tungstenite::Error> {
let started_at = Instant::now();
let result = async {
let mut request = url.into_client_request()?;
request
.headers_mut()
.extend(current_trace_context_headers());
connect_async_with_config(
request,
Some(noise_relay_websocket_config()),
// Rendezvous sends small, latency-sensitive frames, so avoid Nagle's coalescing delay.
/*disable_nagle*/
true,
let connector = WebSocketConnector::new_with_tls_mode(
http_client_factory,
WebSocketTlsMode::TungsteniteDefault,
)
.await
.map(|(websocket, _)| websocket)
.map_err(|error| tokio_tungstenite::tungstenite::Error::Io(std::io::Error::other(error)))?;
connector
.with_tcp_nodelay()
.connect(request, noise_relay_websocket_config())
.await
.map(|(websocket, _)| websocket)
}
.await;
let result_name = if result.is_ok() { "success" } else { "error" };
codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs
+3
View File
@@ -2,6 +2,8 @@
use codex_exec_server_protocol::JSONRPCMessage;
use codex_exec_server_protocol::JSONRPCResponse;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use codex_protocol::models::PermissionProfile;
use codex_protocol::permissions::FileSystemAccessMode;
use codex_protocol::permissions::FileSystemPath;
@@ -41,6 +43,7 @@ async fn remote_file_system_sends_path_and_sandbox_cwd_uris_without_native_conve
websocket_url,
DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,
),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
));
let paths = vec![
PathUri::parse("file:///C:/Users/Alice/src/main.rs").expect("valid drive URI"),
codex-rs/exec-server/src/server/handler/tests.rs
+10 -6
View File
@@ -83,6 +83,10 @@ fn test_runtime_paths() -> ExecServerRuntimePaths {
.expect("runtime paths")
}
fn test_http_client_factory() -> HttpClientFactory {
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault)
}
async fn initialized_handler() -> Arc<ExecServerHandler> {
let (outgoing_tx, _outgoing_rx) = mpsc::channel(16);
let registry = SessionRegistry::new(crate::ExecServerTelemetry::default());
@@ -90,7 +94,7 @@ async fn initialized_handler() -> Arc<ExecServerHandler> {
registry,
RpcNotificationSender::new(outgoing_tx),
test_runtime_paths(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
test_http_client_factory(),
));
let initialize_response = handler
.initialize(InitializeParams {
@@ -169,7 +173,7 @@ async fn long_poll_read_fails_after_session_resume() {
Arc::clone(&registry),
RpcNotificationSender::new(first_tx),
test_runtime_paths(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
test_http_client_factory(),
));
let initialize_response = first_handler
.initialize(InitializeParams {
@@ -210,7 +214,7 @@ async fn long_poll_read_fails_after_session_resume() {
registry,
RpcNotificationSender::new(second_tx),
test_runtime_paths(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
test_http_client_factory(),
));
second_handler
.initialize(InitializeParams {
@@ -244,7 +248,7 @@ async fn active_session_resume_is_rejected() {
Arc::clone(&registry),
RpcNotificationSender::new(first_tx),
test_runtime_paths(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
test_http_client_factory(),
));
let initialize_response = first_handler
.initialize(InitializeParams {
@@ -259,7 +263,7 @@ async fn active_session_resume_is_rejected() {
registry,
RpcNotificationSender::new(second_tx),
test_runtime_paths(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
test_http_client_factory(),
));
let err = second_handler
.initialize(InitializeParams {
@@ -288,7 +292,7 @@ async fn output_and_exit_are_retained_after_notification_receiver_closes() {
SessionRegistry::new(crate::ExecServerTelemetry::default()),
RpcNotificationSender::new(outgoing_tx),
test_runtime_paths(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
test_http_client_factory(),
));
handler
.initialize(InitializeParams {
codex-rs/exec-server/tests/environment.rs
+82
View File
@@ -6,13 +6,95 @@ use anyhow::Context;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::REMOTE_ENVIRONMENT_ID;
use codex_exec_server::SelectedCapabilityRootsStatus;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use codex_http_client::cache_system_proxy_route_for_test;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_utils_path_uri::PathUri;
use common::exec_server::exec_server;
use pretty_assertions::assert_eq;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpListener;
use tokio::net::TcpStream;
use tokio::sync::oneshot;
use tokio::time::sleep;
use tokio::time::timeout;
use tokio_util::task::AbortOnDropHandle;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial(remote_exec_server)]
async fn prepared_remote_environment_uses_configured_system_proxy() -> anyhow::Result<()> {
let server = exec_server().await?;
let upstream = server
.websocket_url()
.strip_prefix("ws://")
.context("exec-server websocket should use ws://")?
.to_string();
let proxy_listener = TcpListener::bind("127.0.0.1:0").await?;
let proxy_url = format!("http://{}", proxy_listener.local_addr()?);
let websocket_url = "ws://exec-server-system-proxy.invalid:8765/";
let proxy_resolution_url = "http://exec-server-system-proxy.invalid:8765/";
cache_system_proxy_route_for_test(proxy_resolution_url, proxy_url);
let (request_tx, request_rx) = oneshot::channel();
let _proxy_task = AbortOnDropHandle::new(tokio::spawn(async move {
let (mut client, _) = proxy_listener.accept().await?;
let mut request = Vec::new();
let mut byte = [0_u8; 1];
while !request.ends_with(b"\r\n\r\n") {
client.read_exact(&mut byte).await?;
request.push(byte[0]);
}
let request_line = String::from_utf8(request)?
.lines()
.next()
.context("system proxy should receive a CONNECT request")?
.to_string();
request_tx
.send(request_line)
.map_err(|_| anyhow::anyhow!("system proxy request receiver was dropped"))?;
let mut target = TcpStream::connect(upstream).await?;
client
.write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
.await?;
tokio::io::copy_bidirectional(&mut client, &mut target).await?;
Ok::<(), anyhow::Error>(())
}));
let codex_home = tempfile::tempdir()?;
std::fs::write(
codex_home.path().join("environments.toml"),
format!(
"default = \"{REMOTE_ENVIRONMENT_ID}\"\ninclude_local = false\n\n[[environments]]\nid = \"{REMOTE_ENVIRONMENT_ID}\"\nurl = \"{websocket_url}\"\n"
),
)?;
let prepared = EnvironmentManager::prepare_from_codex_home(codex_home.path()).await?;
assert!(prepared.default_environment_is_remote());
let manager = prepared.build(
/*local_runtime_paths*/ None,
HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy),
)?;
let request_line = timeout(Duration::from_secs(5), request_rx)
.await
.context("prepared environment did not connect through the system proxy")??;
assert_eq!(
request_line,
"CONNECT exec-server-system-proxy.invalid:8765 HTTP/1.1"
);
let environment = manager
.default_environment()
.context("prepared remote environment")?;
timeout(Duration::from_secs(5), environment.info())
.await
.context("prepared remote environment did not initialize through the system proxy")??;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial(remote_exec_server)]
codex-rs/exec-server/tests/file_stream.rs
+5
View File
@@ -11,6 +11,8 @@ use codex_exec_server::FsOpenParams;
use codex_exec_server::FsReadBlockParams;
use codex_exec_server::FsReadBlockResponse;
use codex_exec_server::RemoteExecServerConnectArgs;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use codex_protocol::models::PermissionProfile;
use codex_protocol::permissions::FileSystemAccessMode;
use codex_protocol::permissions::FileSystemPath;
@@ -234,6 +236,7 @@ async fn read_block_supports_non_sequential_offsets_and_lengths() -> Result<()>
let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new(
server.websocket_url().to_string(),
"file-stream-protocol-test".to_string(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
))
.await?;
let tmp = TempDir::new()?;
@@ -296,6 +299,7 @@ async fn open_enforces_the_per_connection_limit_and_close_releases_capacity() ->
let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new(
server.websocket_url().to_string(),
"file-stream-protocol-test".to_string(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
))
.await?;
let tmp = TempDir::new()?;
@@ -356,6 +360,7 @@ async fn open_rejects_handle_ids_longer_than_32_bytes() -> Result<()> {
let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new(
server.websocket_url().to_string(),
"file-stream-protocol-test".to_string(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
))
.await?;
let tmp = TempDir::new()?;
codex-rs/exec-server/tests/http_client.rs
+3
View File
@@ -19,6 +19,8 @@ use codex_exec_server_protocol::JSONRPCRequest;
use codex_exec_server_protocol::JSONRPCResponse;
use codex_exec_server_protocol::MAX_HTTP_BODY_DELTA_BYTES;
use codex_exec_server_protocol::RequestId;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use futures::SinkExt;
use futures::StreamExt;
use pretty_assertions::assert_eq;
@@ -1369,6 +1371,7 @@ impl ScriptedExecServer {
ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new(
self.websocket_url.clone(),
CLIENT_NAME.to_string(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
))
.await
.context("client should connect to fake exec-server")
codex-rs/exec-server/tests/relay.rs
+82 -6
View File
@@ -16,6 +16,7 @@ use base64::Engine as _;
use base64::engine::general_purpose::STANDARD;
use codex_api::AuthProvider;
use codex_exec_server::EnvironmentConnectionState;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::EnvironmentReadyInfo;
use codex_exec_server::ExecParams;
use codex_exec_server::ExecResponse;
@@ -30,9 +31,9 @@ use codex_exec_server::NoiseRendezvousConnectBundle;
use codex_exec_server::NoiseRendezvousConnectProvider;
use codex_exec_server::ProcessId;
use codex_exec_server::RemoteEnvironmentConfig;
use codex_exec_server_test_support::environment_manager_without_environments;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use codex_http_client::cache_system_proxy_route_for_test;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_utils_path_uri::PathUri;
@@ -46,13 +47,18 @@ use prost::Message as ProstMessage;
use relay_proto::RelayMessageFrame;
use relay_proto::relay_message_frame;
use tempfile::TempDir;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpListener;
use tokio::net::TcpStream;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::task::JoinSet;
use tokio::time::timeout;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::tungstenite::Message;
use tokio_util::task::AbortOnDropHandle;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
@@ -114,7 +120,46 @@ impl NoiseRendezvousConnectProvider for FreshBundleNoiseConnectProvider {
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let rendezvous_url = format!("ws://{}", listener.local_addr()?);
let rendezvous_address = listener.local_addr()?;
let environment_rendezvous_url =
"ws://environment-noise-relay-system-proxy.invalid:8765/relay?role=environment";
let harness_rendezvous_url =
"ws://harness-noise-relay-system-proxy.invalid:8765/relay?role=harness";
let proxy_listener = TcpListener::bind("127.0.0.1:0").await?;
let proxy_url = format!("http://{}", proxy_listener.local_addr()?);
for rendezvous_url in [environment_rendezvous_url, harness_rendezvous_url] {
let proxy_resolution_url = rendezvous_url.replacen("ws://", "http://", /*count*/ 1);
cache_system_proxy_route_for_test(&proxy_resolution_url, proxy_url.clone());
}
let (proxy_request_tx, mut proxy_request_rx) = mpsc::unbounded_channel();
let _proxy_task = AbortOnDropHandle::new(tokio::spawn(async move {
let mut proxy_connections = JoinSet::new();
while let Ok((mut client, _)) = proxy_listener.accept().await {
let proxy_request_tx = proxy_request_tx.clone();
proxy_connections.spawn(async move {
let mut request = Vec::new();
let mut byte = [0_u8; 1];
while !request.ends_with(b"\r\n\r\n") {
client.read_exact(&mut byte).await?;
request.push(byte[0]);
}
let request_line = String::from_utf8(request)?
.lines()
.next()
.context("system proxy should receive a CONNECT request")?
.to_string();
proxy_request_tx
.send(request_line)
.map_err(|_| anyhow::anyhow!("system proxy request receiver was dropped"))?;
let mut target = TcpStream::connect(rendezvous_address).await?;
client
.write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
.await?;
tokio::io::copy_bidirectional(&mut client, &mut target).await?;
Ok::<(), anyhow::Error>(())
});
}
}));
let registry = MockServer::start().await;
Mock::given(method("POST"))
.and(path(format!(
@@ -122,7 +167,7 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle()
)))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"environment_id": ENVIRONMENT_ID,
"url": format!("{rendezvous_url}/relay?role=environment"),
"url": environment_rendezvous_url,
"security_profile": "noise_hybrid_ik_v1",
"executor_registration_id": EXECUTOR_REGISTRATION_ID,
})))
@@ -142,11 +187,12 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle()
let (codex_exe, codex_linux_sandbox_exe) = common::current_test_binary_helper_paths()?;
let runtime_paths = ExecServerRuntimePaths::new(codex_exe, codex_linux_sandbox_exe)?;
let http_client_factory = HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy);
let config = RemoteEnvironmentConfig::new(
registry.uri(),
ENVIRONMENT_ID.to_string(),
static_registry_auth_provider(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
http_client_factory.clone(),
)?;
let remote_environment = tokio::spawn(codex_exec_server::run_remote_environment(
config,
@@ -154,12 +200,19 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle()
));
let environment_websocket = accept_websocket(&listener, "environment").await?;
let environment_proxy_request =
"CONNECT environment-noise-relay-system-proxy.invalid:8765 HTTP/1.1";
let harness_proxy_request = "CONNECT harness-noise-relay-system-proxy.invalid:8765 HTTP/1.1";
assert_eq!(
timeout(TEST_TIMEOUT, proxy_request_rx.recv()).await?,
Some(environment_proxy_request.to_string())
);
let provider = Arc::new(FreshBundleNoiseConnectProvider {
websocket_url: format!("{rendezvous_url}/relay?role=harness"),
websocket_url: harness_rendezvous_url.to_string(),
executor_public_key: registered_executor_public_key(&registry).await?,
calls: AtomicUsize::new(0),
});
let manager = environment_manager_without_environments();
let manager = EnvironmentManager::without_environments(http_client_factory);
let registration = manager
.register_deferred_noise_environment(ENVIRONMENT_ID.to_string(), provider.clone())?;
let environment = manager
@@ -181,6 +234,10 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle()
selected_capability_roots: selected_capability_roots.clone(),
}))?;
let harness_websocket = accept_websocket(&listener, "harness").await?;
assert_eq!(
timeout(TEST_TIMEOUT, proxy_request_rx.recv()).await?,
Some(harness_proxy_request.to_string())
);
let first_relay = tokio::spawn(proxy_relay_frames(
environment_websocket,
harness_websocket,
@@ -207,6 +264,22 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle()
);
let first_reconnected_websocket = accept_websocket(&listener, "reconnected peer").await?;
let second_reconnected_websocket = accept_websocket(&listener, "reconnected peer").await?;
let mut reconnect_proxy_requests = vec![
timeout(TEST_TIMEOUT, proxy_request_rx.recv())
.await?
.context("first reconnected peer should use the system proxy")?,
timeout(TEST_TIMEOUT, proxy_request_rx.recv())
.await?
.context("second reconnected peer should use the system proxy")?,
];
reconnect_proxy_requests.sort();
assert_eq!(
reconnect_proxy_requests,
vec![
environment_proxy_request.to_string(),
harness_proxy_request.to_string(),
]
);
let second_relay = tokio::spawn(proxy_relay_frames(
first_reconnected_websocket,
second_reconnected_websocket,
@@ -300,6 +373,9 @@ async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> {
connect_timeout: TEST_TIMEOUT,
initialize_timeout: TEST_TIMEOUT,
resume_session_id: None,
http_client_factory: codex_http_client::HttpClientFactory::new(
codex_http_client::OutboundProxyPolicy::ReqwestDefault,
),
};
let client_task =
tokio::spawn(async move { ExecServerClient::connect_noise_rendezvous(client_args).await });
codex-rs/rmcp-client/tests/streamable_http_test_support.rs
+3
View File
@@ -23,6 +23,8 @@ use codex_exec_server::Environment;
use codex_exec_server::ExecServerClient;
use codex_exec_server::HttpClient;
use codex_exec_server::RemoteExecServerConnectArgs;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use codex_rmcp_client::ElicitationAction;
use codex_rmcp_client::ElicitationResponse;
use codex_rmcp_client::RmcpClient;
@@ -345,6 +347,7 @@ pub(crate) async fn spawn_exec_server() -> anyhow::Result<ExecServerProcess> {
let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new(
websocket_url,
"rmcp-client-remote-http-test".to_string(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
))
.await?;
codex-rs/websocket-client/src/dialer.rs
+24 -11
View File
@@ -6,6 +6,7 @@ use std::sync::Arc;
use std::time::Duration;
use codex_http_client::OutboundProxyRoute;
use codex_http_client::build_rustls_client_config_with_custom_ca;
use futures::StreamExt;
use futures::stream::FuturesUnordered;
use rustls::ClientConfig;
@@ -28,15 +29,18 @@ use tokio_tungstenite::tungstenite::proxy::ProxyConfig;
use crate::AsyncIo;
use crate::ConnectionInner;
use crate::TcpNodelay;
const HAPPY_EYEBALLS_DELAY: Duration = Duration::from_millis(250);
pub(crate) async fn connect(
request: Request,
config: WebSocketConfig,
tls_config: Arc<ClientConfig>,
tls_config: Option<Arc<ClientConfig>>,
proxy_route: OutboundProxyRoute,
tcp_nodelay: TcpNodelay,
) -> Result<(ConnectionInner, Response), WebSocketError> {
let disable_nagle = tcp_nodelay == TcpNodelay::Enabled;
let proxy_url = match proxy_route {
OutboundProxyRoute::TransportDefault => {
// The workspace enables tokio-tungstenite's `proxy` feature, so its default dialer
@@ -44,8 +48,8 @@ pub(crate) async fn connect(
let (stream, response) = connect_async_tls_with_config(
request,
Some(config),
false, // Preserve Tungstenite's recommended Nagle default.
Some(Connector::Rustls(tls_config)),
disable_nagle,
tls_config.map(Connector::Rustls),
)
.await?;
return Ok((ConnectionInner::TransportDefault(stream), response));
@@ -65,8 +69,8 @@ pub(crate) async fn connect(
match connect_async_tls_with_config(
request.clone(),
Some(config),
false, // Preserve Tungstenite's recommended Nagle default.
Some(Connector::Rustls(Arc::clone(&tls_config))),
disable_nagle,
tls_config.clone().map(Connector::Rustls),
)
.await
{
@@ -84,7 +88,7 @@ pub(crate) async fn connect(
let host = websocket_host(&request)?;
let port = websocket_port(&request)?;
Box::new(
connect_tcp(host_port(host, port))
connect_tcp(host_port(host, port), tcp_nodelay)
.await
.map_err(WebSocketError::Io)?,
)
@@ -93,13 +97,18 @@ pub(crate) async fn connect(
let proxy = ProxyEndpoint::parse(&url)?;
let host = websocket_host(&request)?;
let port = websocket_port(&request)?;
let stream = connect_tcp(proxy.config.authority())
let stream = connect_tcp(proxy.config.authority(), tcp_nodelay)
.await
.map_err(WebSocketError::Io)?;
let stream: Box<dyn AsyncIo> = if proxy.tls {
let proxy_tls_config = match &tls_config {
Some(tls_config) => Arc::clone(tls_config),
None => build_rustls_client_config_with_custom_ca()
.map_err(|error| WebSocketError::Io(error.into()))?,
};
let server_name = ServerName::try_from(proxy.config.host.clone())
.map_err(|_| WebSocketError::Tls(TlsError::InvalidDnsName))?;
let stream = TlsConnector::from(Arc::clone(&tls_config))
let stream = TlsConnector::from(proxy_tls_config)
.connect(server_name, stream)
.await
.map_err(WebSocketError::Io)?;
@@ -115,7 +124,7 @@ pub(crate) async fn connect(
request,
stream,
Some(config),
Some(Connector::Rustls(tls_config)),
tls_config.map(Connector::Rustls),
)
.await?;
Ok((ConnectionInner::Routed(stream), response))
@@ -183,9 +192,13 @@ fn host_port(host: &str, port: u16) -> String {
}
}
async fn connect_tcp(address: String) -> io::Result<TcpStream> {
async fn connect_tcp(address: String, tcp_nodelay: TcpNodelay) -> io::Result<TcpStream> {
let addresses = tokio::net::lookup_host(address).await?.collect::<Vec<_>>();
connect_happy_eyeballs(addresses, TcpStream::connect).await
let stream = connect_happy_eyeballs(addresses, TcpStream::connect).await?;
if tcp_nodelay == TcpNodelay::Enabled {
stream.set_nodelay(/*nodelay*/ true)?;
}
Ok(stream)
}
async fn connect_happy_eyeballs<T, F, Fut>(
codex-rs/websocket-client/src/dialer_tests.rs
+158 -3
View File
@@ -24,6 +24,7 @@ use tokio::net::TcpListener;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use tokio_rustls::TlsAcceptor;
use tokio_tungstenite::MaybeTlsStream;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
@@ -34,6 +35,7 @@ use super::*;
use crate::AsyncIo;
use crate::WebSocketConnection;
use crate::WebSocketConnector;
use crate::WebSocketTlsMode;
#[tokio::test]
async fn public_connector_uses_factory_and_exposes_stream_and_sink() {
@@ -48,6 +50,7 @@ async fn public_connector_uses_factory_and_exposes_stream_and_sink() {
.connect(request, WebSocketConfig::default())
.await
.expect("websocket handshake should succeed");
assert!(!websocket_tcp_nodelay(&websocket));
let expected = Message::Text("hello".into());
websocket
.send(expected.clone())
@@ -63,6 +66,120 @@ async fn public_connector_uses_factory_and_exposes_stream_and_sink() {
target_task.await.expect("target task should finish");
}
#[tokio::test]
async fn public_connector_enables_tcp_nodelay_when_requested() {
let (target_addr, target_task) = start_echo_websocket_server(/*acceptor*/ None).await;
let request = format!("ws://localhost:{}/v1/responses", target_addr.port())
.into_client_request()
.expect("websocket request should build");
let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault);
let connector = WebSocketConnector::new(&factory)
.expect("connector should build")
.with_tcp_nodelay();
let (mut websocket, _) = connector
.connect(request, WebSocketConfig::default())
.await
.expect("websocket handshake should succeed");
assert!(websocket_tcp_nodelay(&websocket));
let expected = Message::Text("latency-sensitive".into());
websocket
.send(expected.clone())
.await
.expect("websocket should send");
assert_eq!(
websocket
.next()
.await
.expect("websocket should receive a message")
.expect("websocket message should be valid"),
expected
);
target_task.await.expect("target task should finish");
}
#[tokio::test]
async fn tungstenite_default_tls_mode_ignores_invalid_custom_ca_in_a_subprocess() {
let (target_addr, target_task) = start_echo_websocket_server(/*acceptor*/ None).await;
let target_url = format!("ws://127.0.0.1:{}/v1/responses", target_addr.port());
let executable = std::env::current_exe().expect("test executable should be available");
let output = tokio::task::spawn_blocking(move || {
let mut command = Command::new(executable);
command.args([
"--exact",
"dialer::tests::tungstenite_default_tls_mode_subprocess_probe",
"--nocapture",
]);
for key in [
"HTTP_PROXY",
"http_proxy",
"HTTPS_PROXY",
"https_proxy",
"ALL_PROXY",
"all_proxy",
"NO_PROXY",
"no_proxy",
"SSL_CERT_FILE",
] {
command.env_remove(key);
}
command
.env(
"CODEX_CA_CERTIFICATE",
"/codex-websocket-client-nonexistent-custom-ca.pem",
)
.env("CODEX_WEBSOCKET_DEFAULT_TLS_PROBE_URL", target_url)
.output()
.expect("WebSocket default-TLS subprocess should run")
})
.await
.expect("WebSocket default-TLS subprocess should join");
assert!(
output.status.success(),
"WebSocket default-TLS subprocess failed\nstdout:\n{}\nstderr:\n{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
target_task.await.expect("target task should finish");
}
#[tokio::test]
async fn tungstenite_default_tls_mode_subprocess_probe() {
let Ok(url) = std::env::var("CODEX_WEBSOCKET_DEFAULT_TLS_PROBE_URL") else {
return;
};
let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault);
assert!(
WebSocketConnector::new(&factory).is_err(),
"explicit Codex TLS should reject the invalid custom CA"
);
let connector =
WebSocketConnector::new_with_tls_mode(&factory, WebSocketTlsMode::TungsteniteDefault)
.expect("Tungstenite-default TLS should ignore custom CA configuration");
let request = url
.into_client_request()
.expect("websocket request should build");
let (mut websocket, _) = connector
.connect(request, WebSocketConfig::default())
.await
.expect("WebSocket should connect without constructing Codex TLS");
let expected = Message::Text("Tungstenite default TLS".into());
websocket
.send(expected.clone())
.await
.expect("WebSocket should send");
assert_eq!(
websocket
.next()
.await
.expect("WebSocket should receive a message")
.expect("WebSocket message should be valid"),
expected
);
}
#[tokio::test]
async fn direct_route_connects_secure_websocket() {
let (tls_config, acceptor, _) = test_tls_configs();
@@ -74,8 +191,9 @@ async fn direct_route_connects_secure_websocket() {
let (inner, _) = connect(
request,
WebSocketConfig::default(),
tls_config,
Some(tls_config),
OutboundProxyRoute::Direct,
TcpNodelay::Enabled,
)
.await
.expect("direct websocket handshake should succeed");
@@ -157,11 +275,12 @@ async fn no_proxy_subprocess_probe() {
let (inner, _) = connect(
request,
WebSocketConfig::default(),
tls_config,
Some(tls_config),
OutboundProxyRoute::Proxy {
url: proxy_url,
no_proxy: Some(no_proxy),
},
TcpNodelay::Enabled,
)
.await
.expect("websocket handshake should succeed");
@@ -238,6 +357,41 @@ async fn happy_eyeballs_does_not_wait_for_stalled_preferred_family() {
assert_eq!(connected, reachable);
}
#[tokio::test]
async fn routed_tcp_connections_only_enable_nodelay_when_requested() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let address = listener
.local_addr()
.expect("listener should have an address")
.to_string();
let default_stream = connect_tcp(address.clone(), TcpNodelay::Default)
.await
.expect("default TCP connection should succeed");
assert!(!default_stream.nodelay().expect("TCP_NODELAY should read"));
let low_latency_stream = connect_tcp(address, TcpNodelay::Enabled)
.await
.expect("low-latency TCP connection should succeed");
assert!(
low_latency_stream
.nodelay()
.expect("TCP_NODELAY should read")
);
}
fn websocket_tcp_nodelay(websocket: &WebSocketConnection) -> bool {
let ConnectionInner::TransportDefault(stream) = &websocket.inner else {
panic!("default connector should use Tungstenite's transport");
};
let MaybeTlsStream::Plain(stream) = stream.get_ref() else {
panic!("test websocket should use a plain TCP stream");
};
stream.nodelay().expect("TCP_NODELAY should read")
}
async fn start_echo_websocket_server(
acceptor: Option<TlsAcceptor>,
) -> (SocketAddr, JoinHandle<()>) {
@@ -453,11 +607,12 @@ async fn assert_proxy_tunnels_secure_websocket(proxy_tls: bool) {
let (inner, _) = connect(
request,
WebSocketConfig::default(),
tls_config,
Some(tls_config),
OutboundProxyRoute::Proxy {
url: format!("{proxy_scheme}://localhost:{}", proxy_addr.port()),
no_proxy: None,
},
TcpNodelay::Enabled,
)
.await
.expect("proxied websocket handshake should succeed");
codex-rs/websocket-client/src/lib.rs
+53 -5
View File
@@ -32,7 +32,23 @@ use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
#[derive(Clone)]
pub struct WebSocketConnector {
http_client_factory: HttpClientFactory,
tls_config: Arc<ClientConfig>,
tls_config: Option<Arc<ClientConfig>>,
tcp_nodelay: TcpNodelay,
}
/// Selects whether WebSocket TLS follows Codex custom-CA policy or Tungstenite defaults.
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum WebSocketTlsMode {
/// Build an explicit TLS configuration from native roots and configured Codex custom CAs.
ExplicitCodexTls,
/// Let Tungstenite build its default TLS configuration when the target requires TLS.
TungsteniteDefault,
}
#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) enum TcpNodelay {
Default,
Enabled,
}
impl WebSocketConnector {
@@ -40,12 +56,36 @@ impl WebSocketConnector {
pub fn new(
http_client_factory: &HttpClientFactory,
) -> Result<Self, BuildCustomCaTransportError> {
Self::new_with_tls_mode(http_client_factory, WebSocketTlsMode::ExplicitCodexTls)
}
/// Creates a connector with explicit Codex TLS or the transport's existing TLS defaults.
///
/// HTTPS proxy connections still build Codex TLS configuration when they establish their
/// proxy tunnel; default-mode target connections otherwise remain entirely with Tungstenite.
pub fn new_with_tls_mode(
http_client_factory: &HttpClientFactory,
tls_mode: WebSocketTlsMode,
) -> Result<Self, BuildCustomCaTransportError> {
let tls_config = match tls_mode {
WebSocketTlsMode::ExplicitCodexTls => {
Some(build_rustls_client_config_with_custom_ca()?)
}
WebSocketTlsMode::TungsteniteDefault => None,
};
Ok(Self {
http_client_factory: http_client_factory.clone(),
tls_config: build_rustls_client_config_with_custom_ca()?,
tls_config,
tcp_nodelay: TcpNodelay::Default,
})
}
/// Disables Nagle's algorithm for latency-sensitive WebSocket connections.
pub fn with_tcp_nodelay(mut self) -> Self {
self.tcp_nodelay = TcpNodelay::Enabled;
self
}
/// Connects a WebSocket after resolving the request destination through the configured proxy
/// policy.
pub async fn connect(
@@ -55,9 +95,17 @@ impl WebSocketConnector {
) -> Result<(WebSocketConnection, Response), WebSocketError> {
let proxy_route = self
.http_client_factory
.resolve_proxy_route(&request.uri().to_string());
let (inner, response) =
dialer::connect(request, config, Arc::clone(&self.tls_config), proxy_route).await?;
.resolve_proxy_route_async(request.uri().to_string())
.await
.map_err(WebSocketError::Io)?;
let (inner, response) = dialer::connect(
request,
config,
self.tls_config.clone(),
proxy_route,
self.tcp_nodelay,
)
.await?;
Ok((WebSocketConnection { inner }, response))
}
}