mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +00:00
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:
1 parent
9f4c20aadc
commit
166658a34a
9 files changed
+468
-20
No files matched your search
Generated
+2
@@ -2365,6 +2365,7 @@ dependencies = [
|
|||||||
"codex-utils-path",
|
"codex-utils-path",
|
||||||
"codex-windows-sandbox",
|
"codex-windows-sandbox",
|
||||||
"crossterm",
|
"crossterm",
|
||||||
|
"futures",
|
||||||
"http 1.4.0",
|
"http 1.4.0",
|
||||||
"insta",
|
"insta",
|
||||||
"libc",
|
"libc",
|
||||||
@@ -2380,6 +2381,7 @@ dependencies = [
|
|||||||
"sys-locale",
|
"sys-locale",
|
||||||
"tempfile",
|
"tempfile",
|
||||||
"tokio",
|
"tokio",
|
||||||
|
"tokio-tungstenite",
|
||||||
"toml 0.9.11+spec-1.1.0",
|
"toml 0.9.11+spec-1.1.0",
|
||||||
"tracing",
|
"tracing",
|
||||||
"tracing-appender",
|
"tracing-appender",
|
||||||
|
|||||||
@@ -109,10 +109,12 @@ app_test_support = { workspace = true }
|
|||||||
assert_cmd = { workspace = true }
|
assert_cmd = { workspace = true }
|
||||||
assert_matches = { workspace = true }
|
assert_matches = { workspace = true }
|
||||||
codex-utils-cargo-bin = { workspace = true }
|
codex-utils-cargo-bin = { workspace = true }
|
||||||
|
futures = { workspace = true }
|
||||||
insta = { workspace = true }
|
insta = { workspace = true }
|
||||||
predicates = { workspace = true }
|
predicates = { workspace = true }
|
||||||
pretty_assertions = { workspace = true }
|
pretty_assertions = { workspace = true }
|
||||||
sqlx = { workspace = true }
|
sqlx = { workspace = true }
|
||||||
|
tokio-tungstenite = { workspace = true }
|
||||||
wiremock = { workspace = true }
|
wiremock = { workspace = true }
|
||||||
zstd = { workspace = true }
|
zstd = { workspace = true }
|
||||||
|
|
||||||
|
|||||||
@@ -7,6 +7,16 @@ const DEFAULT_ANALYTICS_ENABLED: bool = false;
|
|||||||
const DEFAULT_LOG_FILTER: &str = "error,opentelemetry_sdk=off,opentelemetry_otlp=off";
|
const DEFAULT_LOG_FILTER: &str = "error,opentelemetry_sdk=off,opentelemetry_otlp=off";
|
||||||
const OTEL_SERVICE_NAME: &str = "codex-exec-server";
|
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(
|
pub(crate) fn init(
|
||||||
config: Option<&codex_core::config::Config>,
|
config: Option<&codex_core::config::Config>,
|
||||||
) -> (impl Send + Sync, codex_exec_server::ExecServerTelemetry) {
|
) -> (impl Send + Sync, codex_exec_server::ExecServerTelemetry) {
|
||||||
@@ -45,10 +55,37 @@ pub(crate) fn init(
|
|||||||
(otel, telemetry)
|
(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
|
where
|
||||||
F: Future<Output = Result<(), E>>,
|
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() {
|
let shutdown_signal = match shutdown_signal() {
|
||||||
Ok(signal) => Some(signal),
|
Ok(signal) => Some(signal),
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
@@ -56,25 +93,57 @@ where
|
|||||||
None
|
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! {
|
tokio::select! {
|
||||||
result = &mut run => result,
|
result = &mut run => return result,
|
||||||
signal = wait_for_shutdown_signal(shutdown_signal) => {
|
_ = &mut parent_disconnected => {
|
||||||
|
tracing::info!("Stopping exec-server after its parent closed stdin");
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
signal = &mut shutdown_signal, if signal_enabled => {
|
||||||
match signal {
|
match signal {
|
||||||
Ok(()) => Ok(()),
|
Ok(()) => break,
|
||||||
Err(error) => {
|
Err(error) => {
|
||||||
eprintln!("Could not listen for exec-server shutdown signal: {error}");
|
eprintln!("Could not listen for exec-server shutdown signal: {error}");
|
||||||
|
signal_enabled = false;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
match shutdown_behavior {
|
||||||
|
ShutdownBehavior::Immediate => Ok(()),
|
||||||
|
ShutdownBehavior::Graceful(sender) => {
|
||||||
|
let _ = sender.send(());
|
||||||
run.await
|
run.await
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
|
||||||
} else {
|
|
||||||
run.await
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(unix)]
|
#[cfg(unix)]
|
||||||
struct ShutdownSignal {
|
struct ShutdownSignal {
|
||||||
@@ -114,3 +183,7 @@ fn stderr_env_filter() -> EnvFilter {
|
|||||||
.or_else(|_| EnvFilter::try_new(DEFAULT_LOG_FILTER))
|
.or_else(|_| EnvFilter::try_new(DEFAULT_LOG_FILTER))
|
||||||
.unwrap_or_else(|_| EnvFilter::new("error"))
|
.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");
|
||||||
|
}
|
||||||
@@ -587,6 +587,14 @@ struct ExecServerCommand {
|
|||||||
/// Use Agent Identity auth from CODEX_ACCESS_TOKEN for remote registration.
|
/// Use Agent Identity auth from CODEX_ACCESS_TOKEN for remote registration.
|
||||||
#[arg(long = "use-agent-identity-auth", requires = "remote")]
|
#[arg(long = "use-agent-identity-auth", requires = "remote")]
|
||||||
use_agent_identity_auth: bool,
|
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)]
|
#[derive(Debug, clap::Subcommand)]
|
||||||
@@ -1724,9 +1732,26 @@ async fn run_exec_server_command(
|
|||||||
remote_config.name = name;
|
remote_config.name = name;
|
||||||
}
|
}
|
||||||
let remote_config = remote_config.with_telemetry(telemetry);
|
let remote_config = remote_config.with_telemetry(telemetry);
|
||||||
exec_server_telemetry::run_until_shutdown(async move {
|
let parent_lifetime = if cmd.exit_on_stdin_close {
|
||||||
codex_exec_server::run_remote_environment(remote_config, runtime_paths).await
|
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?;
|
.await?;
|
||||||
Ok(())
|
Ok(())
|
||||||
} else {
|
} else {
|
||||||
@@ -1748,7 +1773,8 @@ async fn run_exec_server_command(
|
|||||||
let listen_url = cmd
|
let listen_url = cmd
|
||||||
.listen
|
.listen
|
||||||
.unwrap_or_else(|| codex_exec_server::DEFAULT_LISTEN_URL.to_string());
|
.unwrap_or_else(|| codex_exec_server::DEFAULT_LISTEN_URL.to_string());
|
||||||
exec_server_telemetry::run_until_shutdown(async move {
|
exec_server_telemetry::run_until_shutdown(
|
||||||
|
async move {
|
||||||
codex_exec_server::run_main_with_telemetry(
|
codex_exec_server::run_main_with_telemetry(
|
||||||
&listen_url,
|
&listen_url,
|
||||||
runtime_paths,
|
runtime_paths,
|
||||||
@@ -1756,7 +1782,10 @@ async fn run_exec_server_command(
|
|||||||
http_client_factory,
|
http_client_factory,
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
})
|
},
|
||||||
|
exec_server_telemetry::ParentLifetime::Independent,
|
||||||
|
exec_server_telemetry::ShutdownBehavior::Immediate,
|
||||||
|
)
|
||||||
.await
|
.await
|
||||||
.map_err(anyhow::Error::from_boxed)
|
.map_err(anyhow::Error::from_boxed)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,3 +1,4 @@
|
|||||||
|
use std::collections::HashMap;
|
||||||
#[cfg(unix)]
|
#[cfg(unix)]
|
||||||
use std::io::BufRead as _;
|
use std::io::BufRead as _;
|
||||||
#[cfg(unix)]
|
#[cfg(unix)]
|
||||||
@@ -16,7 +17,19 @@ use std::time::Duration;
|
|||||||
#[cfg(unix)]
|
#[cfg(unix)]
|
||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
|
|
||||||
|
use anyhow::Context;
|
||||||
use anyhow::Result;
|
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::prelude::PredicateBooleanExt;
|
||||||
use predicates::str::contains;
|
use predicates::str::contains;
|
||||||
use tempfile::TempDir;
|
use tempfile::TempDir;
|
||||||
@@ -24,6 +37,9 @@ use tokio::io::AsyncBufReadExt;
|
|||||||
use tokio::io::AsyncReadExt;
|
use tokio::io::AsyncReadExt;
|
||||||
use tokio::io::AsyncWriteExt;
|
use tokio::io::AsyncWriteExt;
|
||||||
use tokio::io::BufReader;
|
use tokio::io::BufReader;
|
||||||
|
use tokio::net::TcpListener;
|
||||||
|
use tokio_tungstenite::WebSocketStream;
|
||||||
|
use tokio_tungstenite::accept_async;
|
||||||
use wiremock::Mock;
|
use wiremock::Mock;
|
||||||
use wiremock::MockServer;
|
use wiremock::MockServer;
|
||||||
use wiremock::ResponseTemplate;
|
use wiremock::ResponseTemplate;
|
||||||
@@ -74,6 +90,261 @@ fn local_exec_server_ignores_invalid_config_without_strict_config() -> Result<()
|
|||||||
Ok(())
|
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(®istry)
|
||||||
|
.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(®istry)
|
||||||
|
.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",
|
||||||
|
®istry.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(®istry).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]
|
#[tokio::test]
|
||||||
async fn local_exec_server_flushes_telemetry_on_stdio_disconnect() -> Result<()> {
|
async fn local_exec_server_flushes_telemetry_on_stdio_disconnect() -> Result<()> {
|
||||||
let collector = MockServer::start().await;
|
let collector = MockServer::start().await;
|
||||||
|
|||||||
@@ -39,6 +39,10 @@ mod websocket_pong_watchdog;
|
|||||||
|
|
||||||
use codex_exec_server_protocol as protocol;
|
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::CODEX_ARG0_EXEC_HELPER_ARG1;
|
||||||
pub use arg0_exec_helper::main as run_arg0_exec_helper_main;
|
pub use arg0_exec_helper::main as run_arg0_exec_helper_main;
|
||||||
pub use capability_discovery::CapabilityDiscoveryError;
|
pub use capability_discovery::CapabilityDiscoveryError;
|
||||||
@@ -181,6 +185,7 @@ pub use protocol::WriteResponse;
|
|||||||
pub use protocol::WriteStatus;
|
pub use protocol::WriteStatus;
|
||||||
pub use remote::RemoteEnvironmentConfig;
|
pub use remote::RemoteEnvironmentConfig;
|
||||||
pub use remote::run_remote_environment;
|
pub use remote::run_remote_environment;
|
||||||
|
pub use remote::run_remote_environment_until_shutdown;
|
||||||
pub use resolved_capability::ResolvedSelectedCapabilityRoot;
|
pub use resolved_capability::ResolvedSelectedCapabilityRoot;
|
||||||
pub use resolved_capability::SelectedCapabilityRootsStatus;
|
pub use resolved_capability::SelectedCapabilityRootsStatus;
|
||||||
pub use runtime_paths::ExecServerRuntimePaths;
|
pub use runtime_paths::ExecServerRuntimePaths;
|
||||||
|
|||||||
@@ -624,15 +624,18 @@ impl LocalProcess {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn child_env(params: &ExecParams) -> HashMap<String, String> {
|
fn child_env(params: &ExecParams) -> HashMap<String, String> {
|
||||||
let Some(env_policy) = ¶ms.env_policy else {
|
let mut env = match ¶ms.env_policy {
|
||||||
return params.env.clone();
|
Some(env_policy) => {
|
||||||
};
|
|
||||||
|
|
||||||
let policy = shell_environment_policy(env_policy);
|
let policy = shell_environment_policy(env_policy);
|
||||||
let mut env = shell_environment::create_env(&policy, /*thread_id*/ None);
|
let mut env = shell_environment::create_env(&policy, /*thread_id*/ None);
|
||||||
env.extend(params.env.clone());
|
env.extend(params.env.clone());
|
||||||
env
|
env
|
||||||
}
|
}
|
||||||
|
None => params.env.clone(),
|
||||||
|
};
|
||||||
|
env.remove(crate::CODEX_EXEC_SERVER_EXIT_ON_STDIN_CLOSE_ENV_VAR);
|
||||||
|
env
|
||||||
|
}
|
||||||
|
|
||||||
fn shell_environment_policy(env_policy: &ExecEnvPolicy) -> ShellEnvironmentPolicy {
|
fn shell_environment_policy(env_policy: &ExecEnvPolicy) -> ShellEnvironmentPolicy {
|
||||||
ShellEnvironmentPolicy {
|
ShellEnvironmentPolicy {
|
||||||
|
|||||||
@@ -514,6 +514,20 @@ pub async fn run_remote_environment(
|
|||||||
config: RemoteEnvironmentConfig,
|
config: RemoteEnvironmentConfig,
|
||||||
runtime_paths: ExecServerRuntimePaths,
|
runtime_paths: ExecServerRuntimePaths,
|
||||||
) -> Result<(), ExecServerError> {
|
) -> 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();
|
ensure_rustls_crypto_provider();
|
||||||
let client = EnvironmentRegistryClient::new_with_telemetry(
|
let client = EnvironmentRegistryClient::new_with_telemetry(
|
||||||
config.base_url.clone(),
|
config.base_url.clone(),
|
||||||
@@ -526,6 +540,24 @@ pub async fn run_remote_environment(
|
|||||||
config.telemetry.clone(),
|
config.telemetry.clone(),
|
||||||
config.http_client_factory.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| {
|
let identity = NoiseChannelIdentity::generate().map_err(|error| {
|
||||||
ExecServerError::Protocol(format!("failed to generate Noise relay identity: {error}"))
|
ExecServerError::Protocol(format!("failed to generate Noise relay identity: {error}"))
|
||||||
})?;
|
})?;
|
||||||
|
|||||||
Reference in new issue
Block a user