Tie remote exec servers to their parent stdin (#35843)

## What changed

- Add `--exit-on-stdin-close` and the `CODEX_EXEC_SERVER_EXIT_ON_STDIN_CLOSE` environment variable as opt-in controls for remote exec servers.
- Gracefully drain active sessions and processes when the parent closes stdin, then flush telemetry before exiting.
- Remove the parent-lifetime environment variable from child process environments.

## Testing

- Cover parent disconnects after signal-listener failures.
- Exercise remote shutdown end to end, including child termination and final telemetry metrics.
- Verify that explicitly disabling the environment variable preserves local exec-server behavior.

GitOrigin-RevId: 63063bc097b54684c370bd545cd32d17c4e55d90
This commit is contained in:
Alexi Christakis
2026-07-28 23:27:46 +00:00
committed by copyberry
parent 9f4c20aadc
commit 166658a34a
9 changed files with 475 additions and 27 deletions
+2
View File
@@ -2365,6 +2365,7 @@ dependencies = [
"codex-utils-path",
"codex-windows-sandbox",
"crossterm",
"futures",
"http 1.4.0",
"insta",
"libc",
@@ -2380,6 +2381,7 @@ dependencies = [
"sys-locale",
"tempfile",
"tokio",
"tokio-tungstenite",
"toml 0.9.11+spec-1.1.0",
"tracing",
"tracing-appender",
+2
View File
@@ -109,10 +109,12 @@ app_test_support = { workspace = true }
assert_cmd = { workspace = true }
assert_matches = { workspace = true }
codex-utils-cargo-bin = { workspace = true }
futures = { workspace = true }
insta = { workspace = true }
predicates = { workspace = true }
pretty_assertions = { workspace = true }
sqlx = { workspace = true }
tokio-tungstenite = { workspace = true }
wiremock = { workspace = true }
zstd = { workspace = true }
+82 -9
View File
@@ -7,6 +7,16 @@ const DEFAULT_ANALYTICS_ENABLED: bool = false;
const DEFAULT_LOG_FILTER: &str = "error,opentelemetry_sdk=off,opentelemetry_otlp=off";
const OTEL_SERVICE_NAME: &str = "codex-exec-server";
pub(crate) enum ParentLifetime {
Independent,
StdinPipe,
}
pub(crate) enum ShutdownBehavior {
Immediate,
Graceful(tokio::sync::oneshot::Sender<()>),
}
pub(crate) fn init(
config: Option<&codex_core::config::Config>,
) -> (impl Send + Sync, codex_exec_server::ExecServerTelemetry) {
@@ -45,10 +55,37 @@ pub(crate) fn init(
(otel, telemetry)
}
pub(crate) async fn run_until_shutdown<F, E>(run: F) -> Result<(), E>
pub(crate) async fn run_until_shutdown<F, E>(
run: F,
parent_lifetime: ParentLifetime,
shutdown_behavior: ShutdownBehavior,
) -> Result<(), E>
where
F: Future<Output = Result<(), E>>,
{
let parent_disconnected = match parent_lifetime {
ParentLifetime::Independent => None,
ParentLifetime::StdinPipe => {
let (sender, receiver) = tokio::sync::oneshot::channel();
std::thread::spawn(move || {
if let Err(error) =
std::io::copy(&mut std::io::stdin().lock(), &mut std::io::sink())
{
tracing::warn!(%error, "Could not read exec-server parent lifetime pipe");
}
let _ = sender.send(());
});
Some(receiver)
}
};
let parent_disconnected = async {
match parent_disconnected {
Some(receiver) => {
let _ = receiver.await;
}
None => std::future::pending().await,
}
};
let shutdown_signal = match shutdown_signal() {
Ok(signal) => Some(signal),
Err(error) => {
@@ -56,23 +93,55 @@ where
None
}
};
tokio::pin!(run);
let shutdown_signal = async {
match shutdown_signal {
Some(signal) => wait_for_shutdown_signal(signal).await,
None => std::future::pending().await,
}
};
run_until_shutdown_with_signals(run, parent_disconnected, shutdown_signal, shutdown_behavior)
.await
}
if let Some(shutdown_signal) = shutdown_signal {
async fn run_until_shutdown_with_signals<F, E, P, S>(
run: F,
parent_disconnected: P,
shutdown_signal: S,
shutdown_behavior: ShutdownBehavior,
) -> Result<(), E>
where
F: Future<Output = Result<(), E>>,
P: Future<Output = ()>,
S: Future<Output = std::io::Result<()>>,
{
tokio::pin!(run, parent_disconnected, shutdown_signal);
let mut signal_enabled = true;
loop {
tokio::select! {
result = &mut run => result,
signal = wait_for_shutdown_signal(shutdown_signal) => {
result = &mut run => return result,
_ = &mut parent_disconnected => {
tracing::info!("Stopping exec-server after its parent closed stdin");
break;
}
signal = &mut shutdown_signal, if signal_enabled => {
match signal {
Ok(()) => Ok(()),
Ok(()) => break,
Err(error) => {
eprintln!("Could not listen for exec-server shutdown signal: {error}");
run.await
signal_enabled = false;
}
}
}
}
} else {
run.await
}
match shutdown_behavior {
ShutdownBehavior::Immediate => Ok(()),
ShutdownBehavior::Graceful(sender) => {
let _ = sender.send(());
run.await
}
}
}
@@ -114,3 +183,7 @@ fn stderr_env_filter() -> EnvFilter {
.or_else(|_| EnvFilter::try_new(DEFAULT_LOG_FILTER))
.unwrap_or_else(|_| EnvFilter::new("error"))
}
#[cfg(test)]
#[path = "exec_server_telemetry_tests.rs"]
mod tests;
@@ -0,0 +1,31 @@
use std::time::Duration;
use super::ShutdownBehavior;
use super::run_until_shutdown_with_signals;
#[tokio::test]
async fn parent_disconnect_still_stops_executor_after_signal_listener_error() {
let (parent_sender, parent_receiver) = tokio::sync::oneshot::channel();
let (shutdown_sender, shutdown_receiver) = tokio::sync::oneshot::channel();
let run = async move {
let _ = shutdown_receiver.await;
Ok::<(), std::io::Error>(())
};
let task = tokio::spawn(run_until_shutdown_with_signals(
run,
async move {
let _ = parent_receiver.await;
},
std::future::ready(Err(std::io::Error::other("signal listener failed"))),
ShutdownBehavior::Graceful(shutdown_sender),
));
tokio::task::yield_now().await;
parent_sender.send(()).expect("parent lifetime receiver");
tokio::time::timeout(Duration::from_secs(1), task)
.await
.expect("parent disconnect should stop the executor")
.expect("shutdown task should not panic")
.expect("executor should shut down successfully");
}
+41 -12
View File
@@ -587,6 +587,14 @@ struct ExecServerCommand {
/// Use Agent Identity auth from CODEX_ACCESS_TOKEN for remote registration.
#[arg(long = "use-agent-identity-auth", requires = "remote")]
use_agent_identity_auth: bool,
/// Exit when the parent-owned standard-input pipe closes.
#[arg(
long = "exit-on-stdin-close",
env = codex_exec_server::CODEX_EXEC_SERVER_EXIT_ON_STDIN_CLOSE_ENV_VAR,
requires_if("true", "remote")
)]
exit_on_stdin_close: bool,
}
#[derive(Debug, clap::Subcommand)]
@@ -1724,9 +1732,26 @@ async fn run_exec_server_command(
remote_config.name = name;
}
let remote_config = remote_config.with_telemetry(telemetry);
exec_server_telemetry::run_until_shutdown(async move {
codex_exec_server::run_remote_environment(remote_config, runtime_paths).await
})
let parent_lifetime = if cmd.exit_on_stdin_close {
exec_server_telemetry::ParentLifetime::StdinPipe
} else {
exec_server_telemetry::ParentLifetime::Independent
};
let (shutdown_sender, shutdown_receiver) = tokio::sync::oneshot::channel();
exec_server_telemetry::run_until_shutdown(
async move {
codex_exec_server::run_remote_environment_until_shutdown(
remote_config,
runtime_paths,
async move {
let _ = shutdown_receiver.await;
},
)
.await
},
parent_lifetime,
exec_server_telemetry::ShutdownBehavior::Graceful(shutdown_sender),
)
.await?;
Ok(())
} else {
@@ -1748,15 +1773,19 @@ async fn run_exec_server_command(
let listen_url = cmd
.listen
.unwrap_or_else(|| codex_exec_server::DEFAULT_LISTEN_URL.to_string());
exec_server_telemetry::run_until_shutdown(async move {
codex_exec_server::run_main_with_telemetry(
&listen_url,
runtime_paths,
telemetry,
http_client_factory,
)
.await
})
exec_server_telemetry::run_until_shutdown(
async move {
codex_exec_server::run_main_with_telemetry(
&listen_url,
runtime_paths,
telemetry,
http_client_factory,
)
.await
},
exec_server_telemetry::ParentLifetime::Independent,
exec_server_telemetry::ShutdownBehavior::Immediate,
)
.await
.map_err(anyhow::Error::from_boxed)
}
+271
View File
@@ -1,3 +1,4 @@
use std::collections::HashMap;
#[cfg(unix)]
use std::io::BufRead as _;
#[cfg(unix)]
@@ -16,7 +17,19 @@ use std::time::Duration;
#[cfg(unix)]
use std::time::Instant;
use anyhow::Context;
use anyhow::Result;
use codex_exec_server::ExecParams;
use codex_exec_server::ExecServerClient;
use codex_exec_server::NoiseChannelIdentity;
use codex_exec_server::NoiseChannelPublicKey;
use codex_exec_server::NoiseRendezvousConnectArgs;
use codex_exec_server::NoiseRendezvousConnectBundle;
use codex_exec_server::ProcessId;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use futures::SinkExt;
use futures::StreamExt;
use predicates::prelude::PredicateBooleanExt;
use predicates::str::contains;
use tempfile::TempDir;
@@ -24,6 +37,9 @@ use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio::net::TcpListener;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::accept_async;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
@@ -74,6 +90,261 @@ fn local_exec_server_ignores_invalid_config_without_strict_config() -> Result<()
Ok(())
}
#[test]
fn local_exec_server_allows_disabled_parent_lifetime_environment_variable() -> Result<()> {
let codex_home = TempDir::new()?;
let mut cmd = codex_command(codex_home.path())?;
cmd.env(
codex_exec_server::CODEX_EXEC_SERVER_EXIT_ON_STDIN_CLOSE_ENV_VAR,
"false",
)
.args(["exec-server", "--listen", "stdio"])
.assert()
.success();
Ok(())
}
#[tokio::test]
async fn remote_exec_server_gracefully_stops_when_parent_stdin_closes() -> Result<()> {
const ENVIRONMENT_ID: &str = "environment-parent-lifetime";
const EXECUTOR_REGISTRATION_ID: &str = "registration-parent-lifetime";
const TEST_TIMEOUT: Duration = Duration::from_secs(30);
let collector = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/metrics"))
.respond_with(ResponseTemplate::new(202))
.mount(&collector)
.await;
let listener = TcpListener::bind("127.0.0.1:0").await?;
let rendezvous_url = format!("ws://{}", listener.local_addr()?);
let registry = MockServer::start().await;
Mock::given(method("POST"))
.and(path(format!(
"/cloud/environment/{ENVIRONMENT_ID}/register"
)))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"environment_id": ENVIRONMENT_ID,
"url": format!("{rendezvous_url}/relay?role=environment"),
"security_profile": "noise_hybrid_ik_v1",
"executor_registration_id": EXECUTOR_REGISTRATION_ID,
})))
.mount(&registry)
.await;
Mock::given(method("POST"))
.and(path(format!(
"/cloud/environment/{ENVIRONMENT_ID}/validate"
)))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"valid": true,
})))
.mount(&registry)
.await;
let codex_home = TempDir::new()?;
let collector_url = collector.uri();
std::fs::write(
codex_home.path().join("config.toml"),
format!(
r#"
[analytics]
enabled = true
[otel]
environment = "test"
metrics_exporter = {{ otlp-http = {{ endpoint = "{collector_url}/v1/metrics", protocol = "json" }} }}
"#
),
)?;
let mut command = tokio::process::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?);
command
.env("CODEX_HOME", codex_home.path())
.env("CODEX_API_KEY", "test-api-key")
.env(
codex_exec_server::CODEX_EXEC_SERVER_EXIT_ON_STDIN_CLOSE_ENV_VAR,
"true",
)
.env("NO_PROXY", "127.0.0.1,localhost")
.env("no_proxy", "127.0.0.1,localhost")
.args([
"exec-server",
"--remote",
&registry.uri(),
"--environment-id",
ENVIRONMENT_ID,
])
.stdin(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
let mut child = command.spawn()?;
let stdin = child
.stdin
.take()
.ok_or_else(|| anyhow::anyhow!("remote exec-server stdin was not piped"))?;
let environment_websocket = accept_parent_lifetime_websocket(&listener, TEST_TIMEOUT).await?;
let executor_public_key = registered_parent_lifetime_executor_public_key(&registry).await?;
let harness_args = NoiseRendezvousConnectArgs {
bundle: NoiseRendezvousConnectBundle {
websocket_url: format!("{rendezvous_url}/relay?role=harness"),
environment_id: ENVIRONMENT_ID.to_string(),
executor_registration_id: EXECUTOR_REGISTRATION_ID.to_string(),
executor_public_key,
harness_key_authorization: "parent-lifetime-harness".to_string(),
},
harness_identity: NoiseChannelIdentity::generate()?,
client_name: "parent-lifetime-test".to_string(),
connect_timeout: TEST_TIMEOUT,
initialize_timeout: TEST_TIMEOUT,
resume_session_id: None,
http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
};
let client_task =
tokio::spawn(async move { ExecServerClient::connect_noise_rendezvous(harness_args).await });
let harness_websocket = accept_parent_lifetime_websocket(&listener, TEST_TIMEOUT).await?;
let relay_task = tokio::spawn(proxy_parent_lifetime_relay(
environment_websocket,
harness_websocket,
));
let client = tokio::time::timeout(TEST_TIMEOUT, client_task)
.await
.context("remote harness did not connect")???;
#[cfg(windows)]
let argv = vec![
"cmd.exe",
"/C",
"if defined CODEX_EXEC_SERVER_EXIT_ON_STDIN_CLOSE (exit /b 1) else ping -n 61 127.0.0.1",
];
#[cfg(not(windows))]
let argv = vec![
"/bin/sh",
"-c",
"[ -z \"${CODEX_EXEC_SERVER_EXIT_ON_STDIN_CLOSE+present}\" ] && exec /bin/sleep 60",
];
let cwd = url::Url::from_directory_path(std::env::current_dir()?)
.map_err(|()| anyhow::anyhow!("could not convert cwd to file URL"))?;
client
.exec(ExecParams {
process_id: ProcessId::from("parent-lifetime-process"),
argv: argv.into_iter().map(str::to_string).collect(),
cwd: cwd.as_str().parse()?,
env_policy: Some(codex_exec_server::ExecEnvPolicy {
inherit: codex_protocol::config_types::ShellEnvironmentPolicyInherit::All,
ignore_default_excludes: false,
exclude: Vec::new(),
r#set: HashMap::new(),
include_only: Vec::new(),
}),
env: HashMap::new(),
tty: false,
pipe_stdin: false,
arg0: None,
sandbox: None,
enforce_managed_network: false,
managed_network: None,
network_proxy: None,
})
.await?;
anyhow::ensure!(
child.try_wait()?.is_none(),
"remote exec-server exited while its parent stdin remained open"
);
drop(stdin);
let output = tokio::time::timeout(TEST_TIMEOUT, child.wait_with_output())
.await
.map_err(|_| anyhow::anyhow!("remote exec-server did not exit after stdin closed"))??;
anyhow::ensure!(
output.status.success(),
"remote exec-server exited with {}; stderr: {}",
output.status,
String::from_utf8_lossy(&output.stderr)
);
relay_task.abort();
let _ = relay_task.await;
let requests = collector
.received_requests()
.await
.ok_or_else(|| anyhow::anyhow!("failed to read OTLP collector requests"))?;
let metrics = requests
.iter()
.filter(|request| request.url.path() == "/v1/metrics")
.map(|request| serde_json::from_slice::<serde_json::Value>(&request.body))
.collect::<serde_json::Result<Vec<_>>>()?;
assert_metric_point(&metrics, "exec_server_processes_active", &[], Some(0));
assert_metric_point(
&metrics,
"exec_server_processes_finished_total",
&[("result", "terminated")],
Some(1),
);
assert_metric_point(
&metrics,
"exec_server_requests_total",
&[("method", "process/start"), ("result", "success")],
Some(1),
);
Ok(())
}
async fn accept_parent_lifetime_websocket(
listener: &TcpListener,
timeout: Duration,
) -> Result<WebSocketStream<tokio::net::TcpStream>> {
let (socket, _) = tokio::time::timeout(timeout, listener.accept())
.await
.context("remote executor or harness did not reach rendezvous")??;
tokio::time::timeout(timeout, accept_async(socket))
.await
.context("fake rendezvous did not complete websocket handshake")?
.map_err(Into::into)
}
async fn registered_parent_lifetime_executor_public_key(
registry: &MockServer,
) -> Result<NoiseChannelPublicKey> {
let requests = registry
.received_requests()
.await
.context("failed to read remote registry requests")?;
let request = requests
.iter()
.find(|request| request.url.path().ends_with("/register"))
.context("remote executor did not register before connecting")?;
let body: serde_json::Value = serde_json::from_slice(&request.body)?;
serde_json::from_value(body["executor_public_key"].clone()).map_err(Into::into)
}
async fn proxy_parent_lifetime_relay(
mut environment: WebSocketStream<tokio::net::TcpStream>,
mut harness: WebSocketStream<tokio::net::TcpStream>,
) -> Result<()> {
loop {
tokio::select! {
message = environment.next() => {
let Some(message) = message else {
break;
};
harness.send(message?).await?;
}
message = harness.next() => {
let Some(message) = message else {
break;
};
environment.send(message?).await?;
}
}
}
Ok(())
}
#[tokio::test]
async fn local_exec_server_flushes_telemetry_on_stdio_disconnect() -> Result<()> {
let collector = MockServer::start().await;
+5
View File
@@ -39,6 +39,10 @@ mod websocket_pong_watchdog;
use codex_exec_server_protocol as protocol;
/// Process-local opt-in for tying a remote executor to its parent's stdin pipe.
pub const CODEX_EXEC_SERVER_EXIT_ON_STDIN_CLOSE_ENV_VAR: &str =
"CODEX_EXEC_SERVER_EXIT_ON_STDIN_CLOSE";
pub use arg0_exec_helper::CODEX_ARG0_EXEC_HELPER_ARG1;
pub use arg0_exec_helper::main as run_arg0_exec_helper_main;
pub use capability_discovery::CapabilityDiscoveryError;
@@ -181,6 +185,7 @@ pub use protocol::WriteResponse;
pub use protocol::WriteStatus;
pub use remote::RemoteEnvironmentConfig;
pub use remote::run_remote_environment;
pub use remote::run_remote_environment_until_shutdown;
pub use resolved_capability::ResolvedSelectedCapabilityRoot;
pub use resolved_capability::SelectedCapabilityRootsStatus;
pub use runtime_paths::ExecServerRuntimePaths;
+9 -6
View File
@@ -624,13 +624,16 @@ impl LocalProcess {
}
fn child_env(params: &ExecParams) -> HashMap<String, String> {
let Some(env_policy) = &params.env_policy else {
return params.env.clone();
let mut env = match &params.env_policy {
Some(env_policy) => {
let policy = shell_environment_policy(env_policy);
let mut env = shell_environment::create_env(&policy, /*thread_id*/ None);
env.extend(params.env.clone());
env
}
None => params.env.clone(),
};
let policy = shell_environment_policy(env_policy);
let mut env = shell_environment::create_env(&policy, /*thread_id*/ None);
env.extend(params.env.clone());
env.remove(crate::CODEX_EXEC_SERVER_EXIT_ON_STDIN_CLOSE_ENV_VAR);
env
}
+32
View File
@@ -514,6 +514,20 @@ pub async fn run_remote_environment(
config: RemoteEnvironmentConfig,
runtime_paths: ExecServerRuntimePaths,
) -> Result<(), ExecServerError> {
run_remote_environment_until_shutdown(config, runtime_paths, std::future::pending()).await
}
/// Serve a remote environment until its owner requests graceful shutdown.
///
/// Active sessions and their processes are drained before this function returns.
pub async fn run_remote_environment_until_shutdown<F>(
config: RemoteEnvironmentConfig,
runtime_paths: ExecServerRuntimePaths,
shutdown: F,
) -> Result<(), ExecServerError>
where
F: std::future::Future<Output = ()>,
{
ensure_rustls_crypto_provider();
let client = EnvironmentRegistryClient::new_with_telemetry(
config.base_url.clone(),
@@ -526,6 +540,24 @@ pub async fn run_remote_environment(
config.telemetry.clone(),
config.http_client_factory.clone(),
);
let result = {
let run = run_remote_environment_connections(config, client, processor.clone());
tokio::pin!(run, shutdown);
tokio::select! {
result = &mut run => result,
_ = &mut shutdown => Ok(()),
}
};
processor.shutdown().await;
result
}
async fn run_remote_environment_connections(
config: RemoteEnvironmentConfig,
client: EnvironmentRegistryClient,
processor: ConnectionProcessor,
) -> Result<(), ExecServerError> {
let identity = NoiseChannelIdentity::generate().map_err(|error| {
ExecServerError::Protocol(format!("failed to generate Noise relay identity: {error}"))
})?;