Persist loaded threads before managed daemon shutdown (#44283)

## Why

Loaded threads, including idle threads whose rollout files are still deferred, need to survive a managed daemon restart. Shutdown must also remain forceable when rollout I/O is blocked.

## What changed

- Add a hidden `--managed-daemon` option for app-server Unix socket transports. After active turns and admitted requests drain, save loaded persistent root threads before exiting.
- Skip ephemeral threads, non-root agent threads, and threads pending unload. Log persistence failures and continue saving other threads.
- Apply shutdown admission checks to thread and turn settings updates, thread deletion, and archiving.
- Keep force signals and daemon shutdown requests responsive during persistence. Return `AppServerExit::Forced` so executables can exit without waiting for runtime teardown.

## Testing

Add integration coverage for resuming active and idle threads after restart, forcing shutdown during active work, and forcing shutdown with a blocked rollout writer. Extend shutdown rejection coverage to deletion and settings updates.

GitOrigin-RevId: 4344e97d39f9f80c5d84c41c17300fd6aa99b4cc

codex-rs/Cargo.lock

@@ -2226,6 +2226,7 @@ dependencies = [
  "codex-state",
  "codex-thread-store",
  "codex-tools",
+ "codex-uds",
  "codex-user-verification",
  "codex-utils-absolute-path",
  "codex-utils-cargo-bin",

codex-rs/app-server/Cargo.toml

@@ -123,6 +123,7 @@ tikv-jemallocator = { workspace = true }
 codex-windows-sandbox = { workspace = true }
 
 [dev-dependencies]
+codex-uds = { workspace = true }
 app_test_support = { workspace = true }
 axum = { workspace = true, default-features = false, features = [
     "http1",

codex-rs/app-server/src/lib.rs

@@ -445,6 +445,15 @@ pub async fn run_main(
         AppServerRuntimeOptions::default(),
     )
     .await
+    .map(|_| ())
+}
+
+/// Reports whether shutdown finished or the managed daemon must exit without waiting for I/O.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum AppServerExit {
+    Graceful,
+    /// Executables should exit before dropping their runtime; libraries must not exit their host.
+    Forced,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -459,6 +468,7 @@ pub struct AppServerRuntimeOptions {
     pub plugin_startup_tasks: PluginStartupTasks,
     pub remote_control_startup_mode: RemoteControlStartupMode,
     pub install_shutdown_signal_handler: bool,
+    pub managed_daemon: bool,
 }
 
 impl Default for AppServerRuntimeOptions {
@@ -468,6 +478,7 @@ impl Default for AppServerRuntimeOptions {
             plugin_startup_tasks: PluginStartupTasks::Start,
             remote_control_startup_mode: RemoteControlStartupMode::ResolvePersisted,
             install_shutdown_signal_handler: true,
+            managed_daemon: false,
         }
     }
 }
@@ -483,7 +494,7 @@ pub async fn run_main_with_transport_options(
     session_source: SessionSource,
     auth: AppServerWebsocketAuthSettings,
     runtime_options: AppServerRuntimeOptions,
-) -> IoResult<()> {
+) -> IoResult<AppServerExit> {
     let loader_overrides = loader_overrides_with_test_user_config_file(
         loader_overrides,
         test_user_config_file_from_env(),
@@ -741,6 +752,8 @@ pub async fn run_main_with_transport_options(
     let single_client_mode = matches!(&transport, AppServerTransport::Stdio);
     let graceful_signal_restart_enabled =
         runtime_options.install_shutdown_signal_handler && !single_client_mode;
+    let managed_daemon = matches!(&transport, AppServerTransport::UnixSocket { .. })
+        && runtime_options.managed_daemon;
     let mut app_server_client_name_rx = None;
 
     match &transport {
@@ -970,24 +983,43 @@ pub async fn run_main_with_transport_options(
         let transport_shutdown_token = transport_shutdown_token.clone();
         async move {
             let mut listen_for_threads = true;
+            // Keep force signals and daemon control events responsive while saving.
+            let mut snapshot = Box::pin(processor.persist_daemon_threads());
+            let mut snapshot_finished = !managed_daemon;
+            let mut clients_disconnected = false;
             let mut shutdown_state = ShutdownState::default();
             let exit_reason = loop {
                 // Sample submissions first: one can publish a running turn before
                 // releasing its permit, and shutdown must observe that new turn.
                 let active_admissions = *active_admissions_rx.borrow_and_update();
                 let running_turn_count = *running_turn_count_rx.borrow_and_update();
-                if matches!(
+                let ready_to_exit = matches!(
                     shutdown_state.update(running_turn_count, active_admissions, connections.len()),
                     ShutdownAction::Finish
-                ) {
+                );
+                if ready_to_exit {
+                    let finished = snapshot_finished || shutdown_state.forced();
+                    if finished {
                         transport_shutdown_token.cancel();
+                    }
+                    if managed_daemon && shutdown_state.forced() {
+                        break "forced_shutdown_requested";
+                    }
+                    if !clients_disconnected {
                         let _ = outbound_control_tx
                             .send(OutboundControlEvent::DisconnectAll)
                             .await;
+                        clients_disconnected = true;
+                    }
+                    if finished {
                         break "shutdown_requested";
                     }
+                }
 
                 tokio::select! {
+                    _ = &mut snapshot, if ready_to_exit && !snapshot_finished => {
+                        snapshot_finished = true;
+                    }
                     shutdown_signal_result = shutdown_signal(), if graceful_signal_restart_enabled && !shutdown_state.forced() => {
                         let signal = match shutdown_signal_result {
                             Ok(signal) => signal,
@@ -1013,6 +1045,12 @@ pub async fn run_main_with_transport_options(
                         let Some(event) = event else {
                             break "transport_channel_closed";
                         };
+                        if ready_to_exit && !matches!(event, TransportEvent::DaemonShutdown) {
+                            if let TransportEvent::ConnectionOpened { disconnect_sender: Some(token), .. } = event {
+                                token.cancel();
+                            }
+                            continue;
+                        }
                         match event {
                             TransportEvent::DaemonShutdown => {
                                 shutdown_state.on_signal(ShutdownSignal::Forceable, connections.len(), *running_turn_count_rx.borrow(), &processor.turn_admission);
@@ -1188,7 +1226,7 @@ pub async fn run_main_with_transport_options(
                             .send_server_notification(notification)
                             .await;
                     }
-                    created = thread_created_rx.recv(), if listen_for_threads => {
+                    created = thread_created_rx.recv(), if listen_for_threads && !ready_to_exit => {
                         match created {
                             Ok(thread_id) => {
                                 let mut initialized_connection_ids = Vec::new();
@@ -1219,6 +1257,7 @@ pub async fn run_main_with_transport_options(
                 }
             };
 
+            drop(snapshot);
             if !shutdown_state.forced() {
                 futures::future::join_all(connections.iter().map(
                     |(&connection_id, connection_state)| {
@@ -1238,12 +1277,19 @@ pub async fn run_main_with_transport_options(
                 shutdown_forced = shutdown_state.forced(),
                 "processor task exited"
             );
+            if managed_daemon && shutdown_state.forced() {
+                AppServerExit::Forced
+            } else {
+                AppServerExit::Graceful
+            }
         }
     });
 
     drop(transport_event_tx);
 
-    let _ = processor_handle.await;
+    if matches!(processor_handle.await, Ok(AppServerExit::Forced)) {
+        return Ok(AppServerExit::Forced);
+    }
     let _ = outbound_handle.await;
 
     transport_shutdown_token.cancel();
@@ -1252,7 +1298,7 @@ pub async fn run_main_with_transport_options(
         let _ = handle.await;
     }
 
-    Ok(())
+    Ok(AppServerExit::Graceful)
 }
 
 struct SqliteRecoveryNotice {

codex-rs/app-server/src/main.rs

@@ -70,6 +70,10 @@ struct AppServerArgs {
     /// Enable remote control for this app-server process without changing persistence.
     #[arg(long = "remote-control", hide = true)]
     remote_control: bool,
+
+    /// Save loaded threads during managed daemon shutdown.
+    #[arg(long, hide = true)]
+    managed_daemon: bool,
 }
 
 fn main() -> anyhow::Result<()> {
@@ -85,6 +89,7 @@ fn main() -> anyhow::Result<()> {
             #[cfg(debug_assertions)]
             disable_plugin_startup_tasks_for_tests,
             remote_control,
+            managed_daemon,
         } = AppServerArgs::parse();
         let loader_overrides = if disable_managed_config_from_debug_env() {
             LoaderOverrides::without_managed_config_for_tests()
@@ -97,6 +102,7 @@ fn main() -> anyhow::Result<()> {
         let auth = auth.try_into_settings()?;
         let mut runtime_options = AppServerRuntimeOptions {
             code_mode_host_transport: code_mode_host.into(),
+            managed_daemon,
             ..Default::default()
         };
         #[cfg(debug_assertions)]
@@ -110,7 +116,7 @@ fn main() -> anyhow::Result<()> {
                 (false, false) => codex_app_server::RemoteControlStartupMode::ResolvePersisted,
             };
 
-        run_main_with_transport_options(
+        let exit = run_main_with_transport_options(
             arg0_paths,
             config_overrides,
             loader_overrides,
@@ -122,6 +128,10 @@ fn main() -> anyhow::Result<()> {
             runtime_options,
         )
         .await?;
+        if exit == codex_app_server::AppServerExit::Forced {
+            // Runtime teardown can wait forever for blocked rollout I/O.
+            std::process::exit(0);
+        }
         Ok(())
     })
 }

codex-rs/app-server/src/message_processor.rs

@@ -759,6 +759,10 @@ impl MessageProcessor {
         self.thread_processor.thread_created_receiver()
     }
 
+    pub(crate) async fn persist_daemon_threads(&self) {
+        self.thread_processor.persist_daemon_threads().await
+    }
+
     pub(crate) async fn send_initialize_notifications_to_connection(
         &self,
         connection_id: ConnectionId,
@@ -951,7 +955,11 @@ impl MessageProcessor {
             | ClientRequest::ThreadFork { .. }
             | ClientRequest::ThreadResume { .. }
             | ClientRequest::ThreadRollback { .. }
-            | ClientRequest::ThreadRevert { .. } => (Some(self.turn_admission.admit()?), false),
+            | ClientRequest::ThreadRevert { .. }
+            | ClientRequest::ThreadSettingsUpdate { .. }
+            | ClientRequest::TurnSettingsUpdate { .. }
+            | ClientRequest::ThreadDelete { .. }
+            | ClientRequest::ThreadArchive { .. } => (Some(self.turn_admission.admit()?), false),
             ClientRequest::TurnStart { .. }
             | ClientRequest::TurnSteer { .. }
             | ClientRequest::ReviewStart { .. }

codex-rs/app-server/src/request_processors/daemon_snapshot.rs

@@ -0,0 +1,35 @@
+//! Persists loaded persistent root threads before managed daemon shutdown.
+//! Failures are logged so other threads can still be saved.
+
+use super::ThreadRequestProcessor;
+use codex_thread_store::PersistContext;
+use tracing::warn;
+
+impl ThreadRequestProcessor {
+    #[expect(
+        clippy::await_holding_invalid_type,
+        reason = "snapshot selection must be serialized against pending unloads"
+    )]
+    pub(crate) async fn persist_daemon_threads(&self) {
+        for thread_id in self.thread_manager.list_thread_ids().await {
+            let pending_unloads = self.pending_thread_unloads.lock().await;
+            if pending_unloads.contains(&thread_id) {
+                continue;
+            }
+            let Ok(thread) = self.thread_manager.get_thread(thread_id).await else {
+                continue;
+            };
+            let config = thread.config_snapshot().await;
+            if !config.ephemeral
+                && config.parent_thread_id.is_none()
+                && !config.session_source.is_non_root_agent()
+                && let Err(err) = self
+                    .thread_store
+                    .persist_thread(thread_id, PersistContext::Standard)
+                    .await
+            {
+                warn!(%thread_id, %err, "failed to persist thread during daemon shutdown");
+            }
+        }
+    }
+}

codex-rs/app-server/src/request_processors/thread_processor.rs

@@ -1,3 +1,6 @@
+#[path = "daemon_snapshot.rs"]
+mod daemon_snapshot;
+
 use super::persisted_resume_settings::PersistedResumeSettings;
 use super::persisted_resume_settings::latest_persisted_resume_settings;
 use super::thread_enrichment::enrich_loaded_threads;

codex-rs/app-server/tests/suite/v2/connection_handling_websocket_unix.rs

@@ -88,6 +88,15 @@ async fn websocket_transport_ctrl_c_waits_for_running_turn_before_exit() -> Resu
         ("thread/start", json!({})),
         ("thread/fork", json!({"threadId": thread_id})),
         ("thread/resume", json!({"threadId": thread_id})),
+        ("thread/delete", json!({"threadId": thread_id})),
+        (
+            "thread/settings/update",
+            json!({"threadId": thread_id, "model": "other"}),
+        ),
+        (
+            "turn/settings/update",
+            json!({"threadId": thread_id, "turnId": "unused", "model": "other"}),
+        ),
         (
             "thread/rollback",
             json!({"threadId": thread_id, "numTurns": 1}),

codex-rs/app-server/tests/suite/v2/daemon_update_recovery.rs

@@ -0,0 +1,324 @@
+//! Best-effort thread persistence during managed app-server shutdown.
+
+use super::connection_handling_websocket::DEFAULT_READ_TIMEOUT;
+use super::connection_handling_websocket::create_config_toml;
+use anyhow::Context;
+use anyhow::Result;
+use app_test_support::DISABLE_PLUGIN_STARTUP_TASKS_ARG;
+use app_test_support::create_final_assistant_message_sse_response;
+use codex_app_server_protocol::InitializeCapabilities;
+use codex_app_server_protocol::JSONRPCMessage;
+use codex_app_server_protocol::RequestId;
+use codex_app_server_protocol::ThreadStartResponse;
+use codex_uds::UnixStream;
+use core_test_support::streaming_sse::StreamingSseChunk;
+use core_test_support::streaming_sse::start_streaming_sse_server;
+use futures::SinkExt;
+use futures::StreamExt;
+use pretty_assertions::assert_eq;
+use serde_json::json;
+use std::path::Path;
+#[cfg(unix)]
+use std::process::Command as StdCommand;
+use std::process::Stdio;
+use tempfile::TempDir;
+use tokio::process::Child;
+use tokio::process::Command;
+use tokio::sync::oneshot;
+use tokio::time::Duration;
+use tokio::time::sleep;
+use tokio::time::timeout;
+use tokio_tungstenite::WebSocketStream;
+use tokio_tungstenite::client_async;
+use tokio_tungstenite::tungstenite::Message;
+
+#[tokio::test]
+async fn managed_shutdown_persists_loaded_threads_after_active_turn() -> Result<()> {
+    let home = TempDir::new()?;
+    let (release_response, response_gate) = oneshot::channel();
+    let (mock, _completions) =
+        start_streaming_sse_server(vec![vec![stream_chunk(Some(response_gate), "Done")?]]).await;
+    create_config_toml(home.path(), mock.uri(), "never")?;
+    let socket_path = home.path().join("control/server.sock");
+    let mut server = spawn_server(home.path(), &socket_path)?;
+    let mut client = connect_default_daemon_client(&socket_path).await?;
+    let active = start_thread(&mut client, /*id*/ 2, json!({})).await?;
+    let idle = start_thread(&mut client, /*id*/ 3, json!({})).await?;
+    start_turn(&mut client, /*id*/ 4, &active.thread.id).await?;
+    wait_for_requests(&mock, /*count*/ 1).await?;
+    request_shutdown(&server, &socket_path).await?;
+    assert_still_running(&mut server, "active turn must finish first").await;
+    release_response.send(()).expect("response is waiting");
+    wait_success(&mut server).await?;
+
+    let mut successor = spawn_server(home.path(), &socket_path)?;
+    let mut reconnected = connect_default_daemon_client(&socket_path).await?;
+    for (id, thread) in [(2, &active.thread), (3, &idle.thread)] {
+        let resumed = request(
+            &mut reconnected,
+            id,
+            "thread/resume",
+            json!({"threadId":thread.id}),
+        )
+        .await?;
+        assert_eq!(resumed["thread"]["id"], thread.id);
+        if thread.id == active.thread.id {
+            assert!(
+                resumed["thread"]["turns"][0]["items"]
+                    .as_array()
+                    .context("recovered turn items")?
+                    .iter()
+                    .any(|item| item["type"] == "agentMessage" && item["text"] == "Done")
+            );
+        }
+    }
+    request_shutdown(&successor, &socket_path).await?;
+    wait_success(&mut successor).await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn managed_force_shutdown_exits_with_active_work() -> Result<()> {
+    let home = TempDir::new()?;
+    let (_release_response, response_gate) = oneshot::channel();
+    let (mock, _completions) =
+        start_streaming_sse_server(vec![vec![stream_chunk(Some(response_gate), "Done")?]]).await;
+    create_config_toml(home.path(), mock.uri(), "never")?;
+    let socket_path = home.path().join("control/server.sock");
+    let mut server = spawn_server(home.path(), &socket_path)?;
+    let mut client = connect_default_daemon_client(&socket_path).await?;
+    let thread = start_thread(&mut client, /*id*/ 2, json!({})).await?;
+    start_turn(&mut client, /*id*/ 3, &thread.thread.id).await?;
+    wait_for_requests(&mock, /*count*/ 1).await?;
+    request_shutdown(&server, &socket_path).await?;
+    assert_still_running(&mut server, "graceful shutdown must wait").await;
+    request_shutdown(&server, &socket_path).await?;
+    wait_success(&mut server).await?;
+    Ok(())
+}
+
+#[cfg(unix)]
+#[tokio::test]
+async fn managed_force_shutdown_exits_with_blocked_rollout_writer() -> Result<()> {
+    let home = TempDir::new()?;
+    let mock = wiremock::MockServer::start().await;
+    create_config_toml(home.path(), &mock.uri(), "never")?;
+    let socket_path = home.path().join("control/server.sock");
+    let mut server = spawn_server(home.path(), &socket_path)?;
+    let mut client = connect_default_daemon_client(&socket_path).await?;
+    let thread = start_thread(&mut client, /*id*/ 2, json!({})).await?;
+    let rollout = thread
+        .thread
+        .path
+        .context("persistent thread rollout path")?;
+    assert!(!rollout.exists(), "idle rollout should still be deferred");
+    let compressed = rollout.with_extension("jsonl.zst");
+    std::fs::create_dir_all(compressed.parent().context("rollout directory")?)?;
+    assert!(
+        StdCommand::new("mkfifo")
+            .arg(&compressed)
+            .status()?
+            .success()
+    );
+
+    request_shutdown(&server, &socket_path).await?;
+    // Opening the write end succeeds only once the real rollout writer opens
+    // the compressed input. Keep it empty and open to block that writer's read.
+    let _blocked_input = timeout(Duration::from_secs(10), async {
+        loop {
+            if let Ok(sender) = tokio::net::unix::pipe::OpenOptions::new().open_sender(&compressed)
+            {
+                return sender;
+            }
+            sleep(Duration::from_millis(25)).await;
+        }
+    })
+    .await
+    .context("rollout writer did not open the blocked input")?;
+    assert_still_running(&mut server, "graceful shutdown must wait for the writer").await;
+    request_shutdown(&server, &socket_path).await?;
+    wait_success(&mut server).await?;
+    Ok(())
+}
+
+async fn wait_success(server: &mut Child) -> Result<()> {
+    anyhow::ensure!(
+        timeout(Duration::from_secs(10), server.wait())
+            .await??
+            .success(),
+        "server did not exit successfully"
+    );
+    Ok(())
+}
+
+async fn assert_still_running(server: &mut Child, message: &str) {
+    assert!(
+        timeout(Duration::from_millis(300), server.wait())
+            .await
+            .is_err(),
+        "{message}"
+    );
+}
+
+fn stream_chunk(gate: Option<oneshot::Receiver<()>>, body: &str) -> Result<StreamingSseChunk> {
+    Ok(StreamingSseChunk {
+        gate,
+        body: create_final_assistant_message_sse_response(body)?,
+    })
+}
+
+async fn wait_for_requests(
+    mock: &core_test_support::streaming_sse::StreamingSseServer,
+    count: usize,
+) -> Result<()> {
+    timeout(Duration::from_secs(5), mock.wait_for_request_count(count)).await?;
+    Ok(())
+}
+
+async fn start_turn(
+    websocket: &mut WebSocketStream<UnixStream>,
+    id: i64,
+    thread_id: impl serde::Serialize,
+) -> Result<()> {
+    request(
+        websocket,
+        id,
+        "turn/start",
+        json!({"threadId":thread_id,"input":[{"type":"text","text":"Hello"}]}),
+    )
+    .await?;
+    Ok(())
+}
+
+async fn start_thread(
+    websocket: &mut WebSocketStream<UnixStream>,
+    id: i64,
+    params: serde_json::Value,
+) -> Result<ThreadStartResponse> {
+    Ok(serde_json::from_value(
+        request(websocket, id, "thread/start", params).await?,
+    )?)
+}
+
+async fn connect_default_daemon_client(socket_path: &Path) -> Result<WebSocketStream<UnixStream>> {
+    connect_daemon_client(socket_path, InitializeCapabilities::default()).await
+}
+
+async fn connect_daemon_client(
+    socket_path: &Path,
+    capabilities: InitializeCapabilities,
+) -> Result<WebSocketStream<UnixStream>> {
+    connect_initialized(socket_path, capabilities, "daemon_recovery_test", "0.1.0").await
+}
+
+fn spawn_server(home: &Path, socket_path: &Path) -> Result<Child> {
+    let binary = codex_utils_cargo_bin::cargo_bin("codex-app-server")?;
+    Ok(Command::new(binary)
+        .args(["--listen", &format!("unix://{}", socket_path.display())])
+        .arg(DISABLE_PLUGIN_STARTUP_TASKS_ARG)
+        .env("CODEX_HOME", home)
+        .arg("--managed-daemon")
+        .env(
+            codex_app_server_transport::DAEMON_SHUTDOWN_SOCKET_ENV,
+            socket_path,
+        )
+        .stdin(Stdio::null())
+        .stdout(Stdio::null())
+        .stderr(Stdio::null())
+        .kill_on_drop(true)
+        .spawn()?)
+}
+
+#[cfg(unix)]
+async fn request_shutdown(server: &Child, _socket_path: &Path) -> Result<()> {
+    let pid = server.id().context("server pid")?;
+    let status = StdCommand::new("kill")
+        .arg("-TERM")
+        .arg(pid.to_string())
+        .status()?;
+    anyhow::ensure!(status.success(), "failed to signal server {pid}: {status}");
+    Ok(())
+}
+
+#[cfg(windows)]
+async fn request_shutdown(server: &Child, socket_path: &Path) -> Result<()> {
+    let pid = server.id().context("server pid")?.to_string();
+    timeout(DEFAULT_READ_TIMEOUT, async {
+        let stream = UnixStream::connect(socket_path).await?;
+        let (mut websocket, _) = client_async("ws://localhost/daemon/shutdown", stream).await?;
+        websocket.send(Message::Text(pid.clone().into())).await?;
+        let reply = websocket
+            .next()
+            .await
+            .context("missing shutdown acknowledgment")??;
+        assert_eq!(reply, Message::Text(pid.into()));
+        websocket.close(None).await?;
+        Ok(())
+    })
+    .await?
+}
+
+async fn connect_initialized(
+    socket_path: &Path,
+    capabilities: InitializeCapabilities,
+    client_name: &str,
+    client_version: &str,
+) -> Result<WebSocketStream<UnixStream>> {
+    let mut websocket = timeout(DEFAULT_READ_TIMEOUT, async {
+        loop {
+            if let Ok(stream) = UnixStream::connect(socket_path).await
+                && let Ok((websocket, _)) = client_async("ws://localhost/", stream).await
+            {
+                return websocket;
+            }
+            sleep(Duration::from_millis(25)).await;
+        }
+    })
+    .await
+    .context("server did not open its control socket")?;
+    request(
+        &mut websocket,
+        /*id*/ 1,
+        "initialize",
+        json!({"clientInfo":{"name":client_name,"version":client_version},"capabilities":capabilities}),
+    )
+    .await?;
+    websocket
+        .send(Message::Text(
+            json!({"method":"initialized"}).to_string().into(),
+        ))
+        .await?;
+    Ok(websocket)
+}
+
+async fn request(
+    websocket: &mut WebSocketStream<UnixStream>,
+    id: i64,
+    method: &str,
+    params: serde_json::Value,
+) -> Result<serde_json::Value> {
+    websocket
+        .send(Message::Text(
+            json!({"id":id,"method":method,"params":params})
+                .to_string()
+                .into(),
+        ))
+        .await?;
+    timeout(DEFAULT_READ_TIMEOUT, async {
+        loop {
+            let frame = websocket.next().await.context("socket closed")??;
+            let Message::Text(text) = frame else { continue };
+            match serde_json::from_str::<JSONRPCMessage>(&text)? {
+                JSONRPCMessage::Response(response) if response.id == RequestId::Integer(id) => {
+                    return Ok(response.result);
+                }
+                JSONRPCMessage::Error(error) if error.id == RequestId::Integer(id) => {
+                    anyhow::bail!("{}: {}", method, error.error.message);
+                }
+                _ => {}
+            }
+        }
+    })
+    .await
+    .context("timed out waiting for app-server response")?
+}

codex-rs/app-server/tests/suite/v2/mod.rs

@@ -25,6 +25,7 @@ mod connection_handling_websocket_unix;
 mod curated_mcp_sync;
 mod current_time;
 mod cyber_access_program;
+mod daemon_update_recovery;
 mod daybreak_access;
 mod dynamic_tools;
 mod environment_add;

codex-rs/cli/src/main.rs

@@ -577,6 +577,10 @@ struct AppServerCommand {
     #[arg(long = "remote-control", hide = true)]
     remote_control: bool,
 
+    /// Save loaded threads during managed daemon shutdown.
+    #[arg(long, hide = true)]
+    managed_daemon: bool,
+
     /// Controls whether analytics are enabled by default.
     ///
     /// Analytics are disabled by default for app-server. Users have to explicitly opt in
@@ -1319,6 +1323,7 @@ async fn cli_main(
                 listen,
                 stdio,
                 remote_control,
+                managed_daemon,
                 analytics_default_enabled,
                 auth,
             } = app_server_cli;
@@ -1339,6 +1344,7 @@ async fn cli_main(
                     let auth = auth.try_into_settings()?;
                     let runtime_options = codex_app_server::AppServerRuntimeOptions {
                         code_mode_host_transport: code_mode_host.into(),
+                        managed_daemon,
                         remote_control_startup_mode: match (remote_control, remote_control_disabled)
                         {
                             (true, _) => {
@@ -1353,7 +1359,7 @@ async fn cli_main(
                         },
                         ..Default::default()
                     };
-                    codex_app_server::run_main_with_transport_options(
+                    let exit = codex_app_server::run_main_with_transport_options(
                         arg0_paths.clone(),
                         root_config_overrides,
                         LoaderOverrides::default(),
@@ -1365,6 +1371,10 @@ async fn cli_main(
                         runtime_options,
                     )
                     .await?;
+                    if exit == codex_app_server::AppServerExit::Forced {
+                        // Runtime teardown can wait forever for blocked rollout I/O.
+                        std::process::exit(0);
+                    }
                 }
                 Some(AppServerSubcommand::Daemon(daemon_cli)) => match daemon_cli.subcommand {
                     AppServerDaemonSubcommand::Start => {

codex-rs/cli/src/remote_control_cmd.rs

@@ -129,7 +129,7 @@ async fn run_foreground_remote_control(
         ..Default::default()
     };
     let (stop_rx, stop_signal_task) = foreground_stop_signal();
-    let mut app_server_task = tokio::spawn(codex_app_server::run_main_with_transport_options(
+    let app_server = codex_app_server::run_main_with_transport_options(
         arg0_paths,
         root_config_overrides,
         LoaderOverrides::default(),
@@ -139,7 +139,8 @@ async fn run_foreground_remote_control(
         SessionSource::VSCode,
         AppServerWebsocketAuthSettings::default(),
         runtime_options,
-    ));
+    );
+    let mut app_server_task = tokio::spawn(async move { app_server.await.map(|_| ()) });
 
     let summary = match wait_for_foreground_remote_control_start(
         &mut app_server_task,