mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +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:
committed by
copyberry
parent
9f4c20aadc
commit
166658a34a
Generated
+2
@@ -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",
|
||||
|
||||
@@ -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 }
|
||||
|
||||
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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(®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]
|
||||
async fn local_exec_server_flushes_telemetry_on_stdio_disconnect() -> Result<()> {
|
||||
let collector = MockServer::start().await;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -624,13 +624,16 @@ impl LocalProcess {
|
||||
}
|
||||
|
||||
fn child_env(params: &ExecParams) -> HashMap<String, String> {
|
||||
let Some(env_policy) = ¶ms.env_policy else {
|
||||
return params.env.clone();
|
||||
let mut env = match ¶ms.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
|
||||
}
|
||||
|
||||
|
||||
@@ -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}"))
|
||||
})?;
|
||||
|
||||
Reference in New Issue
Block a user