Replace Windows app-server shutdown files with socket requests (#43308)

## What changed

Route managed Windows app-server shutdown through `/daemon/shutdown` on the local control socket. Require the server's PID and acknowledge it before triggering shutdown through the existing drain logic. Enable the endpoint only for managed Windows launches via `CODEX_DAEMON_SHUTDOWN_SOCKET`.

Watch running-turn changes whenever shutdown is requested so socket-triggered shutdown can finish draining. Keep the updater's file-based shutdown mechanism and retain forced termination as the manager's timeout fallback.

## Testing

Add a transport test covering rejection of unmanaged shutdown requests and mismatched PIDs, plus acknowledgment before emitting the shutdown event.

GitOrigin-RevId: 3e95dc1914f61e8a131f9b6fd56b6c63688d1744
This commit is contained in:
Eric Traut
2026-09-07 00:00:23 +00:00
committed by copyberry
parent a51da75131
commit 455318c202
16 changed files with 246 additions and 57 deletions
+1
View File
@@ -2312,6 +2312,7 @@ dependencies = [
"tempfile",
"tokio",
"tokio-tungstenite",
"tracing",
"windows-sys 0.52.0",
]
+1
View File
@@ -34,6 +34,7 @@ tokio = { workspace = true, features = [
"time",
] }
tokio-tungstenite = { workspace = true }
tracing = { workspace = true }
[dev-dependencies]
pretty_assertions = { workspace = true }
+23 -5
View File
@@ -138,6 +138,8 @@ impl PidBackend {
}
let pid = record.pid;
let started_at = tokio::time::Instant::now();
let deadline = started_at + STOP_TIMEOUT;
#[cfg(unix)]
self.terminate_process(pid)?;
#[cfg(windows)]
@@ -148,13 +150,29 @@ impl PidBackend {
if process.start_time()? != record.process_start_time {
continue;
}
fs::write(self.pid_file.with_extension("shutdown"), pid.to_string())
.await
.context("failed to request daemon shutdown")?;
match self.command_kind {
PidCommandKind::AppServer { .. } => {
let codex_home = self
.pid_file
.parent()
.and_then(Path::parent)
.context("daemon pid path has no Codex home")?;
let socket_path =
codex_app_server_transport::app_server_control_socket_path(codex_home)?;
if let Err(err) =
crate::client::request_shutdown(socket_path.as_path(), pid).await
{
tracing::warn!(%pid, %err, "managed app-server shutdown request failed; waiting for force deadline");
}
}
PidCommandKind::UpdateLoop => {
fs::write(self.pid_file.with_extension("shutdown"), pid.to_string())
.await
.context("failed to request updater shutdown")?;
}
}
process
};
let started_at = tokio::time::Instant::now();
let deadline = tokio::time::Instant::now() + STOP_TIMEOUT;
let mut forced = false;
loop {
#[cfg(unix)]
@@ -174,16 +174,25 @@ impl PidBackend {
Err(err) => return Err(err).context("failed to clear updater readiness"),
}
}
let shutdown_file = self.pid_file.with_extension("shutdown");
match fs::remove_file(&shutdown_file).await {
Ok(()) => {}
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
Err(err) => return Err(err).context("failed to clear daemon shutdown request"),
match self.command_kind {
PidCommandKind::AppServer { .. } => {
command.env(codex_app_server_transport::DAEMON_SHUTDOWN_SOCKET_ENV, "1");
}
PidCommandKind::UpdateLoop => {
let shutdown_file = self.pid_file.with_extension("shutdown");
match fs::remove_file(&shutdown_file).await {
Ok(()) => {}
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
Err(err) => {
return Err(err).context("failed to clear updater shutdown request");
}
}
command.env(
codex_app_server_transport::DAEMON_SHUTDOWN_FILE_ENV,
shutdown_file,
);
}
}
command.env(
codex_app_server_transport::DAEMON_SHUTDOWN_FILE_ENV,
shutdown_file,
);
}
let child = match command.spawn() {
+27 -1
View File
@@ -61,15 +61,41 @@ async fn probe_inner(socket_path: &Path) -> Result<ProbeInfo> {
}
pub(crate) async fn connect(socket_path: &Path) -> Result<WebSocketStream<UnixStream>> {
connect_at(socket_path, "ws://localhost/").await
}
async fn connect_at(socket_path: &Path, url: &str) -> Result<WebSocketStream<UnixStream>> {
let stream = UnixStream::connect(socket_path)
.await
.with_context(|| format!("failed to connect to {}", socket_path.display()))?;
let (websocket, _response) = client_async("ws://localhost/", stream)
let (websocket, _response) = client_async(url, stream)
.await
.with_context(|| format!("failed to upgrade {}", socket_path.display()))?;
Ok(websocket)
}
#[cfg(windows)]
pub(crate) async fn request_shutdown(socket_path: &Path, pid: u32) -> Result<()> {
timeout(CONTROL_SOCKET_RESPONSE_TIMEOUT, async {
let mut websocket = connect_at(socket_path, "ws://localhost/daemon/shutdown").await?;
websocket
.send(Message::Text(pid.to_string().into()))
.await?;
let reply = websocket
.next()
.await
.context("shutdown socket closed without acknowledgment")??;
anyhow::ensure!(
matches!(reply, Message::Text(ack) if ack == pid.to_string()),
"shutdown acknowledgment did not match the managed process {pid}"
);
websocket.close(None).await?;
Ok(())
})
.await
.context("timed out waiting for managed app-server shutdown acknowledgment")?
}
pub(crate) async fn initialize<S>(
websocket: &mut WebSocketStream<S>,
experimental_api: bool,
@@ -1,6 +1,6 @@
//! Detached Windows processes cannot receive console signals. The daemon passes
//! a request path inside its user-only state directory. Only the addressed PID
//! consumes a request, triggering the SIGTERM drain without changing the wire API.
//! The detached Windows updater has no control socket, so its owner requests
//! termination through a file in its private state directory. Managed app-server
//! shutdown instead uses the protected local socket.
use std::io;
use std::io::Read;
@@ -13,8 +13,7 @@ use std::time::Duration;
#[cfg(windows)]
pub const DAEMON_SHUTDOWN_FILE_ENV: &str = "CODEX_DAEMON_SHUTDOWN_FILE";
/// Waits for and consumes one daemon shutdown request. Without a managed launch
/// environment this stays pending, leaving normal console shutdown unchanged.
/// Waits for and consumes one updater shutdown request.
#[cfg(windows)]
pub async fn daemon_shutdown_signal() -> io::Result<()> {
let Some(path) = std::env::var_os(DAEMON_SHUTDOWN_FILE_ENV).map(PathBuf::from) else {
+3
View File
@@ -4,6 +4,8 @@ mod daemon_shutdown;
pub use daemon_shutdown::DAEMON_SHUTDOWN_FILE_ENV;
#[cfg(windows)]
pub use daemon_shutdown::daemon_shutdown_signal;
/// Only managed app-server launches accept the local socket shutdown request.
pub const DAEMON_SHUTDOWN_SOCKET_ENV: &str = "CODEX_DAEMON_SHUTDOWN_SOCKET";
mod outgoing_message;
mod transport;
@@ -17,6 +19,7 @@ pub use transport::AppServerTransport;
pub use transport::AppServerTransportParseError;
pub use transport::CHANNEL_CAPACITY;
pub use transport::ConnectionOrigin;
pub use transport::DaemonShutdownAccess;
pub use transport::REMOTE_CONTROL_DISABLED_ENV_VAR;
pub use transport::RemoteControlDisabledByRequirements;
pub use transport::RemoteControlEnableError;
@@ -43,6 +43,7 @@ pub use remote_control::start_remote_control;
pub use remote_control::take_remote_control_disabled_env;
pub use stdio::start_stdio_connection;
pub use unix_socket::AppServerStartupLock;
pub use unix_socket::DaemonShutdownAccess;
pub use unix_socket::acquire_app_server_startup_lock;
pub use unix_socket::prepare_control_socket_path;
pub use unix_socket::start_control_socket_acceptor;
@@ -170,6 +171,8 @@ impl FromStr for AppServerTransport {
#[derive(Debug)]
pub enum TransportEvent {
/// Accepted on the managed local control socket, outside JSON-RPC.
DaemonShutdown,
ConnectionOpened {
connection_id: ConnectionId,
origin: ConnectionOrigin,
@@ -423,6 +423,7 @@ fn transport_event_name(event: &TransportEvent) -> &'static str {
TransportEvent::ConnectionOpened { .. } => "connection_opened",
TransportEvent::ConnectionClosed { .. } => "connection_closed",
TransportEvent::IncomingMessage { .. } => "incoming_message",
TransportEvent::DaemonShutdown => "daemon_shutdown",
}
}
@@ -10,11 +10,15 @@ use crate::transport::websocket::run_websocket_connection;
use codex_uds::UnixListener;
use codex_uds::UnixStream;
use codex_utils_absolute_path::AbsolutePathBuf;
use futures::SinkExt;
use futures::StreamExt;
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tokio::time::Duration;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::accept_hdr_async;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::http::Response;
use tokio_tungstenite::tungstenite::http::StatusCode;
use tokio_util::sync::CancellationToken;
use tracing::error;
use tracing::info;
@@ -23,10 +27,17 @@ use tracing::warn;
#[cfg(unix)]
const CONTROL_SOCKET_MODE: u32 = 0o600;
#[derive(Clone, Copy)]
pub enum DaemonShutdownAccess {
Disabled,
Managed,
}
pub async fn start_control_socket_acceptor(
socket_path: AbsolutePathBuf,
transport_event_tx: mpsc::Sender<TransportEvent>,
shutdown_token: CancellationToken,
daemon_shutdown_access: DaemonShutdownAccess,
) -> IoResult<JoinHandle<()>> {
#[cfg(windows)]
let (socket_path, directory_guard) = {
@@ -54,6 +65,7 @@ pub async fn start_control_socket_acceptor(
transport_event_tx,
shutdown_token,
socket_guard,
daemon_shutdown_access,
)))
}
@@ -62,6 +74,7 @@ async fn run_control_socket_acceptor(
transport_event_tx: mpsc::Sender<TransportEvent>,
shutdown_token: CancellationToken,
socket_guard: ControlSocketFileGuard,
daemon_shutdown_access: DaemonShutdownAccess,
) {
let _socket_guard = socket_guard;
loop {
@@ -90,13 +103,33 @@ async fn run_control_socket_acceptor(
let transport_event_tx = transport_event_tx.clone();
tokio::spawn(async move {
let websocket_stream = match accept_async(stream).await {
let mut shutdown_request = false;
let websocket_stream = match accept_hdr_async(
stream,
|request: &tokio_tungstenite::tungstenite::handshake::server::Request, response| {
if request.uri().path() == "/daemon/shutdown" {
if !matches!(daemon_shutdown_access, DaemonShutdownAccess::Managed) {
let mut rejection = Response::new(Some("unmanaged server".to_string()));
*rejection.status_mut() = StatusCode::FORBIDDEN;
return Err(rejection);
}
shutdown_request = true;
}
Ok(response)
},
)
.await
{
Ok(websocket_stream) => websocket_stream,
Err(err) => {
warn!("failed to upgrade control socket websocket connection: {err}");
return;
}
};
if shutdown_request {
run_daemon_shutdown(websocket_stream, transport_event_tx).await;
return;
}
let (websocket_writer, websocket_reader) = websocket_stream.split();
run_websocket_connection(websocket_writer, websocket_reader, transport_event_tx).await;
});
@@ -104,6 +137,24 @@ async fn run_control_socket_acceptor(
info!("control socket acceptor shutting down");
}
async fn run_daemon_shutdown(
mut websocket: tokio_tungstenite::WebSocketStream<UnixStream>,
transport_event_tx: mpsc::Sender<TransportEvent>,
) {
let pid = std::process::id().to_string();
if !matches!(websocket.next().await, Some(Ok(Message::Text(request))) if request == pid) {
return;
}
if websocket.send(Message::Text(pid.into())).await.is_err() {
return;
}
// Let the manager receive the acknowledgment before the main loop closes connections.
let _ = tokio::time::timeout(Duration::from_secs(2), websocket.next()).await;
let _ = transport_event_tx
.send(TransportEvent::DaemonShutdown)
.await;
}
pub async fn prepare_control_socket_path(socket_path: &Path) -> IoResult<()> {
if let Some(parent) = socket_path.parent() {
codex_uds::prepare_private_socket_directory(parent).await?;
@@ -1,5 +1,6 @@
use super::AppServerTransport;
use super::CHANNEL_CAPACITY;
use super::DaemonShutdownAccess;
use super::TransportEvent;
use super::acquire_app_server_startup_lock;
use super::app_server_control_socket_path;
@@ -64,6 +65,7 @@ async fn control_socket_acceptor_upgrades_and_forwards_websocket_text_messages_a
socket_path.clone(),
transport_event_tx,
shutdown_token.clone(),
DaemonShutdownAccess::Disabled,
)
.await
.expect("control socket acceptor should start");
@@ -141,6 +143,87 @@ async fn control_socket_acceptor_upgrades_and_forwards_websocket_text_messages_a
assert_socket_path_removed(socket_path.as_path());
}
#[tokio::test]
async fn shutdown_is_only_accepted_on_managed_local_socket_for_its_own_pid() {
let temp_dir = tempfile::TempDir::new().expect("temp dir");
let socket_path = test_socket_path(temp_dir.path());
let (tx, mut rx) = mpsc::channel(CHANNEL_CAPACITY);
let shutdown = CancellationToken::new();
let acceptor = start_control_socket_acceptor(
socket_path.clone(),
tx,
shutdown.clone(),
DaemonShutdownAccess::Disabled,
)
.await
.expect("acceptor");
let stream = connect_to_socket(socket_path.as_path())
.await
.expect("connect");
assert!(
client_async("ws://localhost/daemon/shutdown", stream)
.await
.is_err()
);
assert!(rx.try_recv().is_err());
shutdown.cancel();
acceptor.await.expect("acceptor shutdown");
let (tx, mut rx) = mpsc::channel(CHANNEL_CAPACITY);
let shutdown = CancellationToken::new();
let acceptor = start_control_socket_acceptor(
socket_path.clone(),
tx,
shutdown.clone(),
DaemonShutdownAccess::Managed,
)
.await
.expect("managed acceptor");
let stream = connect_to_socket(socket_path.as_path())
.await
.expect("connect");
let (mut websocket, _) = client_async("ws://localhost/daemon/shutdown", stream)
.await
.expect("upgrade");
websocket
.send(WebSocketMessage::Text("0".into()))
.await
.expect("wrong pid");
assert!(!matches!(
websocket.next().await,
Some(Ok(WebSocketMessage::Text(_)))
));
assert!(rx.try_recv().is_err());
let stream = connect_to_socket(socket_path.as_path())
.await
.expect("connect");
let (mut websocket, _) = client_async("ws://localhost/daemon/shutdown", stream)
.await
.expect("upgrade");
let pid = std::process::id().to_string();
websocket
.send(WebSocketMessage::Text(pid.clone().into()))
.await
.expect("request");
assert_eq!(
websocket.next().await.expect("ack").expect("ack frame"),
WebSocketMessage::Text(pid.into())
);
assert!(
rx.try_recv().is_err(),
"server must wait until the ack is received"
);
websocket.close(None).await.expect("confirm receipt");
assert!(matches!(
timeout(Duration::from_secs(2), rx.recv()).await,
Ok(Some(TransportEvent::DaemonShutdown))
));
shutdown.cancel();
acceptor.await.expect("acceptor shutdown");
}
#[tokio::test]
async fn app_server_startup_lock_serializes_waiters() {
let temp_dir = tempfile::TempDir::new().expect("temp dir");
@@ -177,6 +260,7 @@ async fn control_socket_file_is_private_after_bind() {
socket_path.clone(),
transport_event_tx,
shutdown_token.clone(),
DaemonShutdownAccess::Disabled,
)
.await
.expect("control socket acceptor should start");
@@ -199,9 +283,14 @@ async fn control_socket_pins_directory_until_shutdown() {
let moved = temp_dir.path().join("moved");
let (tx, _rx) = mpsc::channel::<TransportEvent>(CHANNEL_CAPACITY);
let shutdown = CancellationToken::new();
let acceptor = start_control_socket_acceptor(socket_path.clone(), tx, shutdown.clone())
.await
.expect("acceptor");
let acceptor = start_control_socket_acceptor(
socket_path.clone(),
tx,
shutdown.clone(),
DaemonShutdownAccess::Disabled,
)
.await
.expect("acceptor");
assert!(std::fs::rename(directory, &moved).is_err());
shutdown.cancel();
acceptor.await.expect("shutdown");
+15 -15
View File
@@ -35,6 +35,7 @@ use crate::plugin_config_reload::PluginStartupConfig;
use crate::transport::CHANNEL_CAPACITY;
use crate::transport::ConnectionOrigin;
use crate::transport::ConnectionState;
use crate::transport::DaemonShutdownAccess;
use crate::transport::OutboundConnectionState;
use crate::transport::RemoteControlPolicy;
use crate::transport::RemoteControlStartConfig;
@@ -226,20 +227,8 @@ async fn shutdown_signal() -> IoResult<ShutdownSignal> {
std::future::pending::<()>().await;
}
};
let daemon_signal = async {
let result = codex_app_server_transport::daemon_shutdown_signal().await;
// The processor retries listener errors; the updater instead needs
// immediate errors from its nonblocking shutdown probe.
if result.is_err() {
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
}
result
};
tokio::select! {
_ = console_signal => Ok(ShutdownSignal::Forceable),
result = daemon_signal =>
result.map(|_| ShutdownSignal::Forceable),
}
console_signal.await;
Ok(ShutdownSignal::Forceable)
}
}
@@ -761,6 +750,14 @@ pub async fn run_main_with_transport_options(
socket_path.clone(),
transport_event_tx.clone(),
transport_shutdown_token.clone(),
if cfg!(windows)
&& std::env::var_os(codex_app_server_transport::DAEMON_SHUTDOWN_SOCKET_ENV)
.is_some()
{
DaemonShutdownAccess::Managed
} else {
DaemonShutdownAccess::Disabled
},
)
.await?;
transport_accept_handles.push(accept_handle);
@@ -989,7 +986,7 @@ pub async fn run_main_with_transport_options(
let running_turn_count = *running_turn_count_rx.borrow();
shutdown_state.on_signal(signal, connections.len(), running_turn_count);
}
changed = running_turn_count_rx.changed(), if graceful_signal_restart_enabled && shutdown_state.requested() => {
changed = running_turn_count_rx.changed(), if shutdown_state.requested() => {
if changed.is_err() {
warn!("running-turn watcher closed during graceful restart drain");
}
@@ -999,6 +996,9 @@ pub async fn run_main_with_transport_options(
break "transport_channel_closed";
};
match event {
TransportEvent::DaemonShutdown => {
shutdown_state.on_signal(ShutdownSignal::Forceable, connections.len(), *running_turn_count_rx.borrow());
}
TransportEvent::ConnectionOpened {
connection_id,
origin,
+1
View File
@@ -16,6 +16,7 @@ pub use codex_app_server_transport::AppServerTransport;
pub(crate) use codex_app_server_transport::CHANNEL_CAPACITY;
pub(crate) use codex_app_server_transport::ConnectionId;
pub(crate) use codex_app_server_transport::ConnectionOrigin;
pub(crate) use codex_app_server_transport::DaemonShutdownAccess;
pub(crate) use codex_app_server_transport::OutgoingMessage;
pub(crate) use codex_app_server_transport::QueuedOutgoingMessage;
pub(crate) use codex_app_server_transport::RemoteControlEnableError;
@@ -508,11 +508,6 @@ pub(super) async fn spawn_websocket_server_with_args(
.stderr(Stdio::piped())
.env("CODEX_HOME", codex_home)
.env("RUST_LOG", "warn");
#[cfg(windows)]
cmd.env(
codex_app_server_transport::DAEMON_SHUTDOWN_FILE_ENV,
codex_home.join("daemon.shutdown"),
);
let mut process = cmd
.kill_on_drop(true)
@@ -1,5 +1,4 @@
//! Signal-driven shutdown tests; Windows uses the detached daemon control request
//! for SIGTERM-equivalent graceful and forced shutdown coverage.
//! Signal-driven shutdown tests for Unix websocket servers.
use super::connection_handling_websocket::DEFAULT_READ_TIMEOUT;
use super::connection_handling_websocket::WsClient;
@@ -89,6 +88,7 @@ async fn websocket_transport_second_ctrl_c_forces_exit_while_turn_running() -> R
Ok(())
}
#[cfg(unix)]
#[tokio::test]
async fn websocket_transport_sigterm_waits_for_running_turn_before_exit() -> Result<()> {
let GracefulCtrlCFixture {
@@ -114,6 +114,7 @@ async fn websocket_transport_sigterm_waits_for_running_turn_before_exit() -> Res
Ok(())
}
#[cfg(unix)]
#[tokio::test]
async fn websocket_transport_second_sigterm_forces_exit_while_turn_running() -> Result<()> {
let GracefulCtrlCFixture {
@@ -337,12 +338,3 @@ async fn expect_websocket_disconnect(stream: &mut WsClient) -> Result<()> {
}
}
}
#[cfg(windows)]
fn send_sigterm(process: &Child, home: &std::path::Path) -> Result<()> {
std::fs::write(
home.join("daemon.shutdown"),
process.id().context("server pid")?.to_string(),
)?;
Ok(())
}
+1 -1
View File
@@ -19,7 +19,7 @@ mod config_requirements_browser_use;
mod config_requirements_in_app_browser;
mod config_rpc;
mod connection_handling_websocket;
#[cfg(any(unix, windows))]
#[cfg(unix)]
mod connection_handling_websocket_unix;
#[cfg(unix)]
mod curated_mcp_sync;