agent/codex
2
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
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/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();