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/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;