agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

[1/3] core: add remote environment connection lifecycle (#28674)

Browse Source 
## Why

Remote environments can be registered before their exec-server is first
used. Starting the connection at registration time uses that startup
window, while sharing one startup result prevents background work and
capability calls from opening competing connections.

Keep initial startup simple: each environment makes one connection
attempt using its configured transport timeout. A failed initial attempt
is final for that environment, while an environment that disconnects
after connecting can still recover on a later operation.

## What changed

- Start URL and Noise environments in the background when they are added
to `EnvironmentManager`. Provider snapshots are fully validated before
connection work begins.
- Share one initial connection attempt and its saved result across
metadata, process, filesystem, and HTTP callers.
- Keep configured stdio environments lazy until first use so
registration does not launch a process.
- Tie background startup work to the environment lifetime so replacing
or dropping an environment cancels unfinished work.
- After an established client disconnects, share one fresh connection
attempt across concurrent callers. A failed attempt fails the current
operation without permanently preventing a later attempt.
- Store the shared lazy client directly on `Environment` and expose
small methods for starting, observing, and awaiting startup.

## Test plan

- `just test -p codex-exec-server`
- `just test -p codex-app-server
turn_start_resolves_sticky_thread_local_environment_and_turn_overrides`
This commit is contained in:
sayan-oai
2026-06-18 21:50:15 -07:00
committed by GitHub
Unverified
parent 4e6bc42266
commit 41988e6a24
3 changed files with 469 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/exec-server/src/client.rs
+295 -28
View File
@@ -14,9 +14,10 @@ use futures::FutureExt;
use futures::future::BoxFuture;
use serde_json::Value;
use tokio::sync::Mutex;
use tokio::sync::Semaphore;
use tokio::sync::OnceCell;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio_util::task::AbortOnDropHandle;
use tokio::time::timeout;
use tracing::debug;
@@ -227,54 +228,132 @@ impl Drop for ActiveProcessStart {
}
}
type ConnectionResult = Result<ExecServerClient, Arc<ExecServerError>>;
type ConnectionAttempt = OnceCell<ConnectionResult>;
#[derive(Clone)]
pub(crate) struct LazyRemoteExecServerClient {
transport_params: ExecServerTransportParams,
client: Arc<StdMutex<Option<ExecServerClient>>>,
connect_lock: Arc<Semaphore>,
// Saves the first startup result so callers share it and failures remain final.
startup: Arc<ConnectionAttempt>,
// The latest successful client, replaced whenever reconnecting succeeds.
current_client: Arc<StdMutex<Option<ExecServerClient>>>,
reconnect: Arc<StdMutex<Option<Arc<ConnectionAttempt>>>>,
}
impl LazyRemoteExecServerClient {
pub(crate) fn new(transport_params: ExecServerTransportParams) -> Self {
Self {
transport_params,
client: Arc::new(StdMutex::new(None)),
connect_lock: Arc::new(Semaphore::new(/*permits*/ 1)),
startup: Arc::new(ConnectionAttempt::new()),
current_client: Arc::new(StdMutex::new(None)),
reconnect: Arc::new(StdMutex::new(None)),
}
}
pub(crate) fn start_connecting(&self) -> Option<AbortOnDropHandle<()>> {
// Stdio starts a process, so keep it lazy until the environment is used.
if matches!(
self.transport_params,
ExecServerTransportParams::StdioCommand { .. }
) {
return None;
}
let client = self.clone();
Some(AbortOnDropHandle::new(tokio::spawn(async move {
if let Err(error) = client.wait_until_ready().await {
debug!(%error, "exec-server environment startup failed");
}
})))
}
pub(crate) fn startup_finished(&self) -> bool {
self.startup.get().is_some()
}
pub(crate) async fn wait_until_ready(&self) -> Result<(), ExecServerError> {
self.initial_client().await.map(drop)
}
pub(crate) async fn get(&self) -> Result<ExecServerClient, ExecServerError> {
if let Some(client) = self.connected_client() {
return Ok(client);
}
let _connect_permit = self.connect_lock.acquire().await.map_err(|_| {
ExecServerError::Protocol("exec-server connect lock closed".to_string())
})?;
if let Some(client) = self.connected_client() {
return Ok(client);
}
let next_client = match self.cached_client() {
Some(_client)
if matches!(
&self.transport_params,
ExecServerTransportParams::WebSocketUrl { .. }
| ExecServerTransportParams::NoiseRendezvous { .. }
) =>
{
ExecServerClient::connect_for_transport(self.transport_params.clone()).await?
let Some(cached_client) = self.cached_client() else {
let client = self.initial_client().await?;
if !client.is_disconnected() || !self.can_reconnect() {
return Ok(client);
}
Some(client) => return Ok(client),
None => ExecServerClient::connect_for_transport(self.transport_params.clone()).await?,
return self.reconnect().await;
};
let mut cached_client = self
.client
if !self.can_reconnect() {
return Ok(cached_client);
}
self.reconnect().await
}
async fn initial_client(&self) -> Result<ExecServerClient, ExecServerError> {
// The first caller starts the work; every other caller waits for that same result.
let result = self
.startup
.get_or_init(|| connect_once(self.transport_params.clone()))
.await;
match result {
Ok(client) => {
let mut current_client = self
.current_client
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if current_client.is_none() {
*current_client = Some(client.clone());
}
Ok(client.clone())
}
Err(error) => Err(ExecServerError::ConnectionAttempt(Arc::clone(error))),
}
}
async fn reconnect(&self) -> Result<ExecServerClient, ExecServerError> {
// Callers handling the same outage share one reconnect attempt.
let attempt = {
let mut reconnect = self
.reconnect
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(client) = self.connected_client() {
return Ok(client);
}
reconnect
.get_or_insert_with(|| Arc::new(ConnectionAttempt::new()))
.clone()
};
let result = attempt
.get_or_init(|| async {
let result = connect_once(self.transport_params.clone()).await;
if let Ok(client) = &result {
*self
.current_client
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(client.clone());
}
result
})
.await;
let mut reconnect = self
.reconnect
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
*cached_client = Some(next_client.clone());
Ok(next_client)
// Forget only this completed attempt so a later operation can retry after failure.
if reconnect
.as_ref()
.is_some_and(|current| Arc::ptr_eq(current, &attempt))
{
*reconnect = None;
}
result.clone().map_err(ExecServerError::ConnectionAttempt)
}
fn connected_client(&self) -> Option<ExecServerClient> {
@@ -283,11 +362,25 @@ impl LazyRemoteExecServerClient {
}
fn cached_client(&self) -> Option<ExecServerClient> {
self.client
self.current_client
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
}
fn can_reconnect(&self) -> bool {
matches!(
self.transport_params,
ExecServerTransportParams::WebSocketUrl { .. }
| ExecServerTransportParams::NoiseRendezvous { .. }
)
}
}
async fn connect_once(transport_params: ExecServerTransportParams) -> ConnectionResult {
ExecServerClient::connect_for_transport(transport_params)
.await
.map_err(Arc::new)
}
impl HttpClient for LazyRemoteExecServerClient {
@@ -353,6 +446,8 @@ pub enum ExecServerError {
EnvironmentRegistryAuth(String),
#[error("environment registry request failed: {0}")]
EnvironmentRegistryRequest(#[from] reqwest::Error),
#[error("exec-server connection attempt failed: {0}")]
ConnectionAttempt(#[source] Arc<ExecServerError>),
}
impl ExecServerClient {
@@ -1112,6 +1207,7 @@ mod tests {
use tokio::net::TcpStream;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::sync::watch;
use tokio::time::Duration;
#[cfg(unix)]
use tokio::time::sleep;
@@ -1919,6 +2015,177 @@ mod tests {
server.await.expect("server task should finish");
}
#[tokio::test]
async fn initial_connection_is_shared_by_all_waiters() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let websocket_url = format!(
"ws://{}",
listener.local_addr().expect("listener should have address")
);
let server = tokio::spawn(async move {
let mut connection = accept_websocket(&listener).await;
complete_websocket_initialize(
&mut connection,
"startup-session",
/*expected_resume_session_id*/ None,
)
.await;
timeout(Duration::from_secs(1), connection.next())
.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),
});
assert!(!client.startup_finished());
let _startup_task = client.start_connecting();
let (ready, first, second) =
tokio::join!(client.wait_until_ready(), client.get(), client.get());
ready.expect("background startup should finish");
let first = first.expect("first waiter should receive the client");
let second = second.expect("second waiter should receive the same client");
assert!(client.startup_finished());
assert_eq!(first.session_id().as_deref(), Some("startup-session"));
assert!(Arc::ptr_eq(&first.inner, &second.inner));
drop(first);
drop(second);
drop(client);
server.await.expect("server task should finish");
}
#[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,
},
initialize_timeout: Duration::from_secs(1),
});
assert!(client.start_connecting().is_none());
assert!(!client.startup_finished());
let first = match client.get().await {
Ok(_) => panic!("missing executable should fail"),
Err(error) => error,
};
assert!(client.startup_finished());
let second = match client.get().await {
Ok(_) => panic!("burned environment should stay failed"),
Err(error) => error,
};
let (
super::ExecServerError::ConnectionAttempt(first),
super::ExecServerError::ConnectionAttempt(second),
) = (first, second)
else {
panic!("expected saved connection failures");
};
assert!(Arc::ptr_eq(&first, &second));
}
#[tokio::test]
async fn failed_reconnect_does_not_burn_environment() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let websocket_url = format!(
"ws://{}",
listener.local_addr().expect("listener should have address")
);
let (replacement_initialized_tx, replacement_initialized_rx) = oneshot::channel();
let (allow_replacement_tx, allow_replacement_rx) = watch::channel(false);
let server = tokio::spawn(async move {
let mut first = accept_websocket(&listener).await;
complete_websocket_initialize(
&mut first,
"startup-session",
/*expected_resume_session_id*/ None,
)
.await;
first
.close(None)
.await
.expect("startup websocket should close");
let successful_reconnect = loop {
let (stream, _) = listener.accept().await.expect("reconnect should arrive");
if *allow_replacement_rx.borrow() {
break stream;
}
let mut failed_reconnect = stream;
failed_reconnect
.write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n")
.await
.expect("failed handshake response should write");
};
let mut successful_reconnect = accept_async(successful_reconnect)
.await
.expect("replacement websocket handshake should succeed");
complete_websocket_initialize(
&mut successful_reconnect,
"replacement-session",
/*expected_resume_session_id*/ None,
)
.await;
replacement_initialized_tx
.send(())
.expect("replacement initialization should be observed");
timeout(Duration::from_secs(1), successful_reconnect.next())
.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 initial = client.get().await.expect("startup should connect");
timeout(Duration::from_secs(1), async {
while !initial.is_disconnected() {
tokio::task::yield_now().await;
}
})
.await
.expect("client should observe disconnect");
let failed_reconnect = match client.get().await {
Ok(_) => panic!("first lazy reconnect should fail"),
Err(error) => error,
};
assert!(matches!(
failed_reconnect,
super::ExecServerError::ConnectionAttempt(_)
));
allow_replacement_tx
.send(true)
.expect("server should allow a fresh client");
let replacement = client.get().await.expect("later reconnect should succeed");
assert_eq!(
replacement.session_id().as_deref(),
Some("replacement-session")
);
replacement_initialized_rx
.await
.expect("server should observe replacement initialization");
drop(initial);
drop(replacement);
drop(client);
server.await.expect("server task should finish");
}
#[tokio::test]
async fn wake_notifications_do_not_block_other_sessions() {
let (client_stdin, server_reader) = duplex(1 << 20);
codex-rs/exec-server/src/environment.rs
+174 -54
View File
@@ -1,10 +1,8 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::RwLock;
use futures::FutureExt;
use futures::future::BoxFuture;
use crate::ExecServerError;
use crate::ExecServerRuntimePaths;
use crate::ExecutorFileSystem;
@@ -29,6 +27,7 @@ use crate::remote::NoiseRendezvousEnvironmentConfig;
use crate::remote_file_system::RemoteFileSystem;
use crate::remote_process::RemoteProcess;
use codex_shell_command::shell_detect::DetectedShell;
use tokio_util::task::AbortOnDropHandle;
pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL";
pub const CODEX_EXEC_SERVER_NOISE_REGISTRY_URL_ENV_VAR: &str =
@@ -50,9 +49,9 @@ pub const CODEX_EXEC_SERVER_NOISE_CHATGPT_ACCOUNT_ID_ENV_VAR: &str =
/// use `default_environment().is_some()` as the signal for model-facing
/// shell/filesystem tool availability.
///
/// Remote environments create remote filesystem and execution backends that
/// lazy-connect to the configured exec-server on first use. The remote
/// transport is not opened when the manager or environment is constructed.
/// Remote environments begin connecting when added to the manager. Their
/// filesystem and execution backends share that startup result and reconnect
/// after later disconnects as needed.
#[derive(Debug)]
pub struct EnvironmentManager {
default_environment: Option<String>,
@@ -224,6 +223,10 @@ impl EnvironmentManager {
Some(environment_id)
}
};
// The snapshot is valid; start connecting its remote environments in the background.
for environment in environment_map.values() {
environment.start_connecting();
}
Ok(Self {
default_environment,
environments: RwLock::new(environment_map),
@@ -307,12 +310,15 @@ impl EnvironmentManager {
"remote environment requires an exec-server url".to_string(),
));
};
let environment =
Environment::remote_inner(exec_server_url, self.local_runtime_paths.clone());
let environment = Arc::new(Environment::remote_inner(
exec_server_url,
self.local_runtime_paths.clone(),
));
environment.start_connecting();
self.environments
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(environment_id, Arc::new(environment));
.insert(environment_id, environment);
Ok(())
}
@@ -336,14 +342,15 @@ impl EnvironmentManager {
"failed to generate Noise harness identity: {error}"
))
})?;
let environment = Environment::remote_with_transport(
let environment = Arc::new(Environment::remote_with_transport(
ExecServerTransportParams::NoiseRendezvous { provider, identity },
self.local_runtime_paths.clone(),
);
));
environment.start_connecting();
self.environments
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(environment_id, Arc::new(environment));
.insert(environment_id, environment);
Ok(())
}
}
@@ -402,50 +409,22 @@ fn optional_environment_value(name: &str) -> Option<String> {
#[derive(Clone)]
pub struct Environment {
exec_server_url: Option<String>,
remote_transport: Option<ExecServerTransportParams>,
info_provider: Arc<dyn EnvironmentInfoProvider>,
remote_client: Option<LazyRemoteExecServerClient>,
// Dropping the environment stops unfinished background startup work.
startup_task: Arc<Mutex<Option<AbortOnDropHandle<()>>>>,
exec_backend: Arc<dyn ExecBackend>,
filesystem: Arc<dyn ExecutorFileSystem>,
http_client: Arc<dyn HttpClient>,
local_runtime_paths: Option<ExecServerRuntimePaths>,
}
/// Provides environment metadata from either a local environment or a remote exec-server.
trait EnvironmentInfoProvider: Send + Sync {
fn info(&self) -> BoxFuture<'_, Result<EnvironmentInfo, ExecServerError>>;
}
struct LocalEnvironmentInfoProvider;
impl EnvironmentInfoProvider for LocalEnvironmentInfoProvider {
fn info(&self) -> BoxFuture<'_, Result<EnvironmentInfo, ExecServerError>> {
std::future::ready(Ok(EnvironmentInfo::local())).boxed()
}
}
struct RemoteEnvironmentInfoProvider {
client: LazyRemoteExecServerClient,
}
impl RemoteEnvironmentInfoProvider {
fn new(client: LazyRemoteExecServerClient) -> Self {
Self { client }
}
}
impl EnvironmentInfoProvider for RemoteEnvironmentInfoProvider {
fn info(&self) -> BoxFuture<'_, Result<EnvironmentInfo, ExecServerError>> {
async move { self.client.environment_info().await }.boxed()
}
}
impl Environment {
/// Builds a test-only local environment without configured sandbox helper paths.
pub fn default_for_tests() -> Self {
Self {
exec_server_url: None,
remote_transport: None,
info_provider: Arc::new(LocalEnvironmentInfoProvider),
remote_client: None,
startup_task: Arc::new(Mutex::new(None)),
exec_backend: Arc::new(LocalProcess::default()),
filesystem: Arc::new(LocalFileSystem::unsandboxed()),
http_client: Arc::new(ReqwestHttpClient),
@@ -501,8 +480,8 @@ impl Environment {
pub(crate) fn local(local_runtime_paths: ExecServerRuntimePaths) -> Self {
Self {
exec_server_url: None,
remote_transport: None,
info_provider: Arc::new(LocalEnvironmentInfoProvider),
remote_client: None,
startup_task: Arc::new(Mutex::new(None)),
exec_backend: Arc::new(LocalProcess::default()),
filesystem: Arc::new(LocalFileSystem::with_runtime_paths(
local_runtime_paths.clone(),
@@ -534,15 +513,15 @@ impl Environment {
ExecServerTransportParams::NoiseRendezvous { .. } => None,
ExecServerTransportParams::StdioCommand { .. } => None,
};
let client = LazyRemoteExecServerClient::new(remote_transport.clone());
let client = LazyRemoteExecServerClient::new(remote_transport);
let exec_backend: Arc<dyn ExecBackend> = Arc::new(RemoteProcess::new(client.clone()));
let filesystem: Arc<dyn ExecutorFileSystem> =
Arc::new(RemoteFileSystem::new(client.clone()));
Self {
exec_server_url,
remote_transport: Some(remote_transport),
info_provider: Arc::new(RemoteEnvironmentInfoProvider::new(client.clone())),
remote_client: Some(client.clone()),
startup_task: Arc::new(Mutex::new(None)),
exec_backend,
filesystem,
http_client: Arc::new(client),
@@ -551,7 +530,7 @@ impl Environment {
}
pub fn is_remote(&self) -> bool {
self.remote_transport.is_some()
self.remote_client.is_some()
}
/// Returns the remote exec-server URL when this environment is remote.
@@ -565,7 +544,40 @@ impl Environment {
/// Returns environment information from the selected execution/filesystem environment.
pub async fn info(&self) -> Result<EnvironmentInfo, ExecServerError> {
self.info_provider.info().await
match &self.remote_client {
Some(client) => client.environment_info().await,
None => Ok(EnvironmentInfo::local()),
}
}
/// Starts connecting a remote environment without waiting for it.
/// Requires an active Tokio runtime when background startup is supported.
pub fn start_connecting(&self) {
let Some(client) = &self.remote_client else {
return;
};
let mut startup_task = self
.startup_task
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if startup_task.is_none() {
*startup_task = client.start_connecting();
}
}
/// Returns whether initial startup has either succeeded or permanently failed.
pub fn startup_finished(&self) -> bool {
self.remote_client
.as_ref()
.is_none_or(LazyRemoteExecServerClient::startup_finished)
}
/// Waits for initial startup. A failed startup is never attempted again.
pub async fn wait_until_ready(&self) -> Result<(), ExecServerError> {
match &self.remote_client {
Some(client) => client.wait_until_ready().await,
None => Ok(()),
}
}
pub fn get_exec_backend(&self) -> Arc<dyn ExecBackend> {
@@ -600,7 +612,9 @@ impl From<DetectedShell> for ShellInfo {
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use super::Environment;
use super::EnvironmentManager;
@@ -609,10 +623,14 @@ mod tests {
use super::noise_environment_config_from_values;
use crate::ExecServerRuntimePaths;
use crate::ProcessId;
use crate::client_api::ExecServerTransportParams;
use crate::client_api::StdioExecServerCommand;
use crate::environment_provider::EnvironmentDefault;
use crate::environment_provider::EnvironmentProviderSnapshot;
use codex_utils_path_uri::PathUri;
use pretty_assertions::assert_eq;
use tokio::net::TcpListener;
use tokio::time::timeout;
fn test_runtime_paths() -> ExecServerRuntimePaths {
ExecServerRuntimePaths::new(
@@ -626,8 +644,8 @@ mod tests {
assert!(manager.try_local_environment().is_none());
}
#[test]
fn noise_environment_config_selects_remote_as_default() {
#[tokio::test]
async fn noise_environment_config_selects_remote_as_default() {
let config = noise_environment_config_from_values(
Some("http://registry.example/api".to_string()),
Some("environment-requested".to_string()),
@@ -982,6 +1000,108 @@ mod tests {
assert!(!Arc::ptr_eq(&first, &second));
}
#[tokio::test]
async fn environment_manager_starts_remote_environment_when_upserted() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind websocket listener");
let manager = EnvironmentManager::without_environments();
manager
.upsert_environment(
"executor-a".to_string(),
format!("ws://{}", listener.local_addr().expect("listener address")),
)
.expect("remote environment");
timeout(Duration::from_secs(5), listener.accept())
.await
.expect("environment should start connecting when registered")
.expect("accept connection");
}
#[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),
},
/*local_runtime_paths*/ None,
);
let manager = EnvironmentManager::from_snapshot(
EnvironmentProviderSnapshot {
environments: vec![("stdio".to_string(), environment)],
default: EnvironmentDefault::Disabled,
include_local: false,
},
/*local_runtime_paths*/ None,
)
.expect("environment manager");
let environment = manager.get_environment("stdio").expect("stdio environment");
assert!(!environment.startup_finished());
assert!(environment.wait_until_ready().await.is_err());
assert!(environment.startup_finished());
}
#[tokio::test]
async fn replacing_environment_stops_its_startup_task() {
let first_listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind first websocket listener");
let second_listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind second websocket listener");
let manager = EnvironmentManager::without_environments();
manager
.upsert_environment(
"executor-a".to_string(),
format!(
"ws://{}",
first_listener.local_addr().expect("first listener address")
),
)
.expect("first remote environment");
let environment = manager
.get_environment("executor-a")
.expect("first remote environment");
let startup_abort = environment
.startup_task
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.as_ref()
.expect("startup task")
.abort_handle();
assert!(!startup_abort.is_finished());
drop(environment);
manager
.upsert_environment(
"executor-a".to_string(),
format!(
"ws://{}",
second_listener
.local_addr()
.expect("second listener address")
),
)
.expect("replacement remote environment");
timeout(Duration::from_secs(1), async {
while !startup_abort.is_finished() {
tokio::task::yield_now().await;
}
})
.await
.expect("replacing the environment should cancel its startup task");
}
#[tokio::test]
async fn environment_manager_rejects_empty_remote_environment_url() {
let manager = EnvironmentManager::without_environments();
codex-rs/exec-server/tests/relay.rs
-65
View File
@@ -6,8 +6,6 @@ mod relay_proto;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Duration;
use anyhow::Context;
@@ -15,25 +13,20 @@ use anyhow::Result;
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD;
use codex_api::AuthProvider;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::ExecParams;
use codex_exec_server::ExecResponse;
use codex_exec_server::ExecServerClient;
use codex_exec_server::ExecServerError;
use codex_exec_server::ExecServerRuntimePaths;
use codex_exec_server::FsReadFileParams;
use codex_exec_server::NoiseChannelIdentity;
use codex_exec_server::NoiseChannelPublicKey;
use codex_exec_server::NoiseRendezvousConnectArgs;
use codex_exec_server::NoiseRendezvousConnectBundle;
use codex_exec_server::NoiseRendezvousConnectProvider;
use codex_exec_server::ProcessId;
use codex_exec_server::RemoteEnvironmentConfig;
use codex_utils_path_uri::PathUri;
use futures::FutureExt;
use futures::SinkExt;
use futures::StreamExt;
use futures::future::BoxFuture;
use http::HeaderMap;
use http::HeaderValue;
use pretty_assertions::assert_eq;
@@ -72,68 +65,10 @@ impl AuthProvider for StaticRegistryAuthProvider {
}
}
struct FailingNoiseConnectProvider {
attempts: Arc<AtomicUsize>,
}
impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider {
fn connect_bundle(
&self,
_: NoiseChannelPublicKey,
) -> BoxFuture<'_, Result<NoiseRendezvousConnectBundle, ExecServerError>> {
self.attempts.fetch_add(1, Ordering::SeqCst);
async {
Err(ExecServerError::Protocol(
"test registry connect failure".to_string(),
))
}
.boxed()
}
}
fn static_registry_auth_provider() -> codex_api::SharedAuthProvider {
Arc::new(StaticRegistryAuthProvider)
}
#[tokio::test]
async fn noise_environment_refreshes_bundle_for_each_connection_attempt() -> Result<()> {
let attempts = Arc::new(AtomicUsize::new(0));
let manager = EnvironmentManager::without_environments();
manager.upsert_noise_environment(
ENVIRONMENT_ID.to_string(),
Arc::new(FailingNoiseConnectProvider {
attempts: Arc::clone(&attempts),
}),
)?;
let backend = manager
.get_environment(ENVIRONMENT_ID)
.context("Noise environment should be materialized")?
.get_exec_backend();
for attempt in 1..=2 {
let result = backend
.start(ExecParams {
process_id: ProcessId::new(format!("proc-{attempt}")),
argv: vec!["true".to_string()],
cwd: PathUri::from_path(std::env::current_dir()?)?,
env_policy: None,
env: HashMap::new(),
tty: false,
pipe_stdin: false,
arg0: None,
})
.await;
assert!(matches!(
result,
Err(ExecServerError::Protocol(ref message))
if message == "test registry connect failure"
));
}
assert_eq!(attempts.load(Ordering::SeqCst), 2);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;