Record thread recovery candidates on managed daemon shutdown (#44299)

## What changed

- Atomically save successfully persisted, loaded root thread IDs to `app-server-daemon/loaded-threads.json` during graceful managed shutdown, excluding ephemeral threads and threads pending unload.
- Clear stale recovery state at startup and around daemon lifecycle operations, including explicit stops and fresh starts.
- Probe support for `--managed-daemon` before launching managed app servers, and clear recovery state best-effort when launching older binaries.
- Keep forced shutdown responsive while saving the snapshot and log snapshot failures without preventing shutdown.

## Testing

Add coverage for snapshots after active turns finish, forced shutdown without a snapshot, excluded threads, snapshot write failures, and recovery-state cleanup for stops, fresh starts, and legacy launches.

GitOrigin-RevId: 422dc32b85355510f7e4c57b7540597b47d9f0b2
This commit is contained in:
Eric Traut
2026-09-09 20:39:33 +00:00
committed by copyberry
parent 742472c525
commit 7c88f037d9
12 changed files with 398 additions and 41 deletions
@@ -93,6 +93,40 @@ impl PidBackend {
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::from(stderr_log.into_std().await));
// Older or pinned managed binaries may predate this optional startup flag.
let managed_app_server =
matches!(self.command_kind, super::PidCommandKind::AppServer { .. });
if managed_app_server
&& matches!(
tokio::time::timeout(
std::time::Duration::from_secs(5),
Command::new(&codex_bin)
.args(["app-server", "--managed-daemon", "--help"])
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.kill_on_drop(true)
.status(),
).await,
Ok(Ok(status)) if status.success()
)
{
command.arg("--managed-daemon");
} else if managed_app_server {
let codex_home = self
.pid_file
.parent()
.and_then(std::path::Path::parent)
.context("daemon pid path has no Codex home")?;
let recovery_file = codex_app_server_transport::daemon_recovery_file_path(codex_home);
match fs::remove_file(&recovery_file).await {
Ok(()) => {}
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
Err(err) => {
tracing::warn!(path = %recovery_file.display(), %err, "failed to clear daemon recovery state before legacy launch");
}
}
}
if let Some((key, value)) = self.command_env() {
command.env(key, value);
}
@@ -155,6 +155,36 @@ async fn start_retries_stale_empty_pid_file_under_its_own_lock() {
);
}
#[cfg(unix)]
#[tokio::test]
async fn legacy_launch_clears_recovery_best_effort() {
for snapshot_is_directory in [false, true] {
let home = TempDir::new().expect("temp dir");
let state_dir = home.path().join("app-server-daemon");
std::fs::create_dir_all(&state_dir).expect("state dir");
let recovery_file = codex_app_server_transport::daemon_recovery_file_path(home.path());
if snapshot_is_directory {
std::fs::create_dir(&recovery_file).expect("invalid snapshot directory");
} else {
std::fs::write(&recovery_file, "{}").expect("pending snapshot");
}
let backend = PidBackend::new(
home.path().join("missing-codex"),
state_dir.join("app-server.pid"),
/*remote_control_enabled*/ false,
);
let error = backend.start().await.expect_err("missing binary");
assert!(
error
.to_string()
.starts_with("failed to spawn detached app-server process using "),
"{error:#}"
);
assert_eq!(recovery_file.exists(), snapshot_is_directory);
}
}
#[tokio::test]
async fn stale_record_cleanup_preserves_replacement_record() {
let temp_dir = TempDir::new().expect("temp dir");
+77 -1
View File
@@ -7,6 +7,7 @@ mod client;
mod managed_install;
mod remote_control_client;
mod settings;
mod thread_recovery;
mod update_loop;
use std::path::Path;
@@ -314,6 +315,15 @@ impl Daemon {
})
}
fn recovery_file(&self) -> Result<PathBuf> {
Ok(codex_app_server_transport::daemon_recovery_file_path(
self.settings_file
.parent()
.and_then(Path::parent)
.context("daemon settings path has no Codex home")?,
))
}
async fn run(&self, command: LifecycleCommand) -> Result<LifecycleOutput> {
match command {
LifecycleCommand::Start => {
@@ -326,7 +336,11 @@ impl Daemon {
}
LifecycleCommand::Stop => {
let _operation_lock = self.acquire_operation_lock().await?;
self.stop().await
let output = self.stop().await?;
if let Err(err) = thread_recovery::discard_pending(self) {
eprintln!("warning: failed to clear saved threads after daemon stop: {err}");
}
Ok(output)
}
LifecycleCommand::Version => self.version().await,
}
@@ -350,6 +364,10 @@ impl Daemon {
self.wait_until_ready().await?,
)
} else {
// A fresh start must ignore snapshots left by older stop clients.
if let Err(err) = thread_recovery::discard_pending(self) {
eprintln!("warning: failed to clear stale daemon recovery before start: {err}");
}
self.ensure_managed_codex_bin()?;
let pid = self.start_managed_backend(&settings).await?;
(
@@ -386,6 +404,9 @@ impl Daemon {
self.ensure_managed_codex_bin()?;
if let Some(backend) = self.running_backend_instance(&settings).await? {
if let Err(err) = thread_recovery::discard_pending(self) {
eprintln!("warning: failed to clear stale daemon recovery before restart: {err}");
}
backend
.stop_with_grace(settings.shutdown_grace_seconds)
.await?;
@@ -456,6 +477,11 @@ impl Daemon {
RestartDecision::Restart => {
#[cfg(windows)]
backend::windows::ensure_detached_launch(managed_codex_bin)?;
if let Err(err) = thread_recovery::discard_pending(self) {
eprintln!(
"warning: failed to clear stale daemon recovery before update: {err}"
);
}
backend
.stop_with_grace(settings.shutdown_grace_seconds)
.await?;
@@ -654,6 +680,11 @@ impl Daemon {
settings.save(&self.settings_file).await?;
let app_server_version = if let Some(backend) = backend {
if let Err(err) = thread_recovery::discard_pending(self) {
eprintln!(
"warning: failed to clear stale recovery before remote-control restart: {err}"
);
}
backend
.stop_with_grace(settings.shutdown_grace_seconds)
.await?;
@@ -695,6 +726,9 @@ impl Daemon {
.stop()
.await?;
if let Some(backend) = self.running_backend_instance(&settings).await? {
if let Err(err) = thread_recovery::discard_pending(self) {
eprintln!("warning: failed to clear stale daemon recovery before bootstrap: {err}");
}
backend
.stop_with_grace(settings.shutdown_grace_seconds)
.await?;
@@ -1176,6 +1210,48 @@ mod tests {
);
}
#[tokio::test]
async fn stop_and_fresh_start_discard_pending_thread_restore() {
let home = TempDir::new().expect("home");
let state = home.path().join("app-server-daemon");
codex_uds::prepare_private_socket_directory(&state)
.await
.expect("private state directory");
let daemon = Daemon {
socket_path: home.path().join("server.sock"),
pid_file: state.join("server.pid"),
update_pid_file: state.join("updater.pid"),
operation_lock_file: state.join("daemon.lock"),
settings_file: state.join("settings.json"),
managed_codex_bin: home.path().join("codex"),
};
codex_app_server_transport::daemon_recovery::write_candidates(
&daemon.recovery_file().expect("recovery path"),
&["thread".to_string()].into_iter().collect(),
)
.expect("saved threads");
assert_eq!(
daemon
.run(super::LifecycleCommand::Stop)
.await
.expect("stop")
.status,
LifecycleStatus::NotRunning
);
assert!(!daemon.recovery_file().expect("recovery path").exists());
// Simulate an older stop client leaving a snapshot behind.
std::fs::write(
daemon.recovery_file().expect("recovery path"),
r#"["thread"]"#,
)
.expect("legacy saved threads");
daemon
.run(super::LifecycleCommand::Start)
.await
.expect_err("missing backend binary");
assert!(!daemon.recovery_file().expect("recovery path").exists());
}
#[cfg(unix)]
#[tokio::test]
async fn managed_local_backend_counts_as_bootstrapped_without_updater() {
@@ -0,0 +1,16 @@
//! Removes stale snapshots before planned daemon replacements.
use std::io::ErrorKind;
use anyhow::Context;
use anyhow::Result;
use crate::Daemon;
pub(crate) fn discard_pending(daemon: &Daemon) -> Result<()> {
match std::fs::remove_file(daemon.recovery_file()?) {
Ok(()) => Ok(()),
Err(err) if err.kind() == ErrorKind::NotFound => Ok(()),
Err(err) => Err(err).context("failed to clear daemon recovery file"),
}
}
@@ -0,0 +1,21 @@
//! Shared on-disk candidate set for managed daemon restarts.
use std::collections::BTreeSet;
use std::io;
use std::path::Path;
use codex_core::path_utils::write_atomically;
pub fn read_candidates(path: &Path) -> io::Result<BTreeSet<String>> {
match std::fs::read(path) {
Ok(contents) => serde_json::from_slice(&contents).map_err(io::Error::other),
Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(BTreeSet::new()),
Err(err) => Err(err),
}
}
pub fn write_candidates(path: &Path, candidates: &BTreeSet<String>) -> io::Result<()> {
write_atomically(
path,
&serde_json::to_string(candidates).map_err(io::Error::other)?,
)
}
+2
View File
@@ -1,3 +1,4 @@
pub mod daemon_recovery;
#[cfg(any(windows, test))]
mod daemon_shutdown;
#[cfg(windows)]
@@ -33,6 +34,7 @@ pub use transport::acquire_app_server_startup_lock;
pub use transport::app_server_control_socket_path;
pub use transport::app_server_startup_lock_path;
pub use transport::auth;
pub use transport::daemon_recovery_file_path;
pub use transport::prepare_control_socket_path;
pub use transport::start_control_socket_acceptor;
pub use transport::start_remote_control;
@@ -11,6 +11,7 @@ use codex_core::config::find_codex_home;
use codex_utils_absolute_path::AbsolutePathBuf;
use std::net::SocketAddr;
use std::path::Path;
use std::path::PathBuf;
use std::str::FromStr;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
@@ -55,6 +56,13 @@ const OVERLOADED_ERROR_CODE: i64 = -32001;
const APP_SERVER_CONTROL_SOCKET_DIR_NAME: &str = "app-server-control";
const APP_SERVER_CONTROL_SOCKET_FILE_NAME: &str = "app-server-control.sock";
const APP_SERVER_STARTUP_LOCK_FILE_NAME: &str = "app-server-startup.lock";
const DAEMON_RECOVERY_FILE_NAME: &str = "loaded-threads.json";
pub fn daemon_recovery_file_path(codex_home: &Path) -> PathBuf {
codex_home
.join("app-server-daemon")
.join(DAEMON_RECOVERY_FILE_NAME)
}
pub fn app_server_control_socket_path(codex_home: &Path) -> std::io::Result<AbsolutePathBuf> {
AbsolutePathBuf::from_absolute_path(
@@ -0,0 +1,28 @@
//! Persists the candidate set for a planned managed-daemon replacement.
use std::collections::BTreeSet;
use std::io;
use std::path::PathBuf;
use codex_app_server_transport::daemon_recovery;
pub(crate) async fn snapshot(path: PathBuf, loaded: Vec<String>) -> io::Result<()> {
// A forced exit must not wait for file I/O in Tokio's blocking pool.
let (result_tx, result_rx) = tokio::sync::oneshot::channel();
std::thread::Builder::new()
.name("daemon-snapshot".into())
.spawn(move || {
let result = daemon_recovery::write_candidates(
&path,
&loaded.into_iter().collect::<BTreeSet<_>>(),
);
if result.is_err()
&& let Err(err) = std::fs::remove_file(&path)
&& err.kind() != io::ErrorKind::NotFound
{
tracing::warn!("failed to clear stale daemon recovery file: {err}");
}
let _ = result_tx.send(result);
})?;
result_rx.await.map_err(io::Error::other)?
}
+19 -1
View File
@@ -55,6 +55,7 @@ use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::ServerNotification;
use codex_app_server_protocol::TextPosition as AppTextPosition;
use codex_app_server_protocol::TextRange as AppTextRange;
use codex_app_server_transport::daemon_recovery_file_path;
use codex_config::ConfigLayerSource;
use codex_config::ConfigLoadError;
use codex_config::TextRange as CoreTextRange;
@@ -106,6 +107,7 @@ mod config_manager_service;
mod connection_cleanup;
mod connection_rpc_gate;
mod current_time;
mod daemon_thread_recovery;
mod dynamic_tools;
mod effective_plugin_change;
mod error_code;
@@ -937,6 +939,7 @@ pub async fn run_main_with_transport_options(
info!("outbound router task exited (channel closed)");
});
let recovery_file = daemon_recovery_file_path(&config.codex_home);
let processor_handle = tokio::spawn({
let auth_manager = Arc::clone(&auth_manager);
let analytics_events_client =
@@ -982,9 +985,24 @@ pub async fn run_main_with_transport_options(
let mut remote_control_status = remote_control_status_rx.borrow().clone();
let transport_shutdown_token = transport_shutdown_token.clone();
async move {
// This process must never hand an earlier generation's candidates to
// its successor, including when the snapshot writer cannot start.
if managed_daemon
&& let Err(err) = tokio::fs::remove_file(&recovery_file).await
&& err.kind() != std::io::ErrorKind::NotFound
{
warn!("failed to clear stale daemon recovery file: {err}");
}
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 = Box::pin(async {
let loaded = processor.daemon_recovery_candidates().await;
if let Err(err) =
daemon_thread_recovery::snapshot(recovery_file.clone(), loaded).await
{
warn!("failed to save loaded threads during daemon shutdown: {err}");
}
});
let mut snapshot_finished = !managed_daemon;
let mut clients_disconnected = false;
let mut shutdown_state = ShutdownState::default();
+2 -2
View File
@@ -759,8 +759,8 @@ 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 daemon_recovery_candidates(&self) -> Vec<String> {
self.thread_processor.daemon_recovery_candidates().await
}
pub(crate) async fn send_initialize_notifications_to_connection(
@@ -1,5 +1,5 @@
//! Persists loaded persistent root threads before managed daemon shutdown.
//! Failures are logged so other threads can still be saved.
//! Selects persistent root threads for managed daemon recovery using normal resume semantics.
//! Only threads successfully persisted through the existing thread store become candidates.
use super::ThreadRequestProcessor;
use codex_thread_store::PersistContext;
@@ -10,7 +10,8 @@ impl ThreadRequestProcessor {
clippy::await_holding_invalid_type,
reason = "snapshot selection must be serialized against pending unloads"
)]
pub(crate) async fn persist_daemon_threads(&self) {
pub(crate) async fn daemon_recovery_candidates(&self) -> Vec<String> {
let mut ids = Vec::new();
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) {
@@ -23,13 +24,18 @@ impl ThreadRequestProcessor {
if !config.ephemeral
&& config.parent_thread_id.is_none()
&& !config.session_source.is_non_root_agent()
&& let Err(err) = self
{
if let Err(err) = self
.thread_store
.persist_thread(thread_id, PersistContext::Standard)
.await
{
warn!(%thread_id, %err, "failed to persist thread during daemon shutdown");
{
warn!(%thread_id, %err, "skipping daemon restore for thread that could not be persisted");
continue;
}
ids.push(thread_id.to_string());
}
}
ids
}
}
@@ -1,4 +1,4 @@
//! Best-effort thread persistence during managed app-server shutdown.
//! Best-effort loaded-thread snapshots during managed app-server shutdown.
use super::connection_handling_websocket::DEFAULT_READ_TIMEOUT;
use super::connection_handling_websocket::create_config_toml;
@@ -10,6 +10,8 @@ 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_app_server_transport::daemon_recovery;
use codex_app_server_transport::daemon_recovery_file_path;
use codex_uds::UnixStream;
use core_test_support::streaming_sse::StreamingSseChunk;
use core_test_support::streaming_sse::start_streaming_sse_server;
@@ -33,58 +35,76 @@ use tokio_tungstenite::client_async;
use tokio_tungstenite::tungstenite::Message;
#[tokio::test]
async fn managed_shutdown_persists_loaded_threads_after_active_turn() -> Result<()> {
async fn managed_shutdown_saves_loaded_thread_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 recovery_file = daemon_recovery_file_path(home.path());
daemon_recovery::write_candidates(&recovery_file, &["stale".into()].into_iter().collect())?;
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?;
let mut client = connect_daemon_client(
&socket_path,
InitializeCapabilities {
experimental_api: true,
extensions: Some(std::collections::HashMap::from([(
"io.modelcontextprotocol/ui".into(),
json!({"mimeTypes":["text/html"]}),
)])),
..Default::default()
},
)
.await?;
assert!(!recovery_file.exists());
let thread = start_thread(
&mut client,
/*id*/ 2,
json!({
"model":"gpt-5.2",
"developerInstructions":"Session-only instructions",
"dynamicTools":[{"name":"lookup_ticket","description":"Look up a ticket",
"inputSchema":{"type":"object","properties":{}}}]
}),
)
.await?;
for (id, method, params) in [
(
3,
"skills/extraRoots/set",
json!({"extraRoots":[home.path()]}),
),
(
4,
"experimentalFeature/enablement/set",
json!({"enablement":{"mcp_2026_07_28":true}}),
),
] {
request(&mut client, id, method, params).await?;
}
start_turn(&mut client, /*id*/ 5, &thread.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?;
assert_eq!(
daemon_recovery::read_candidates(&recovery_file)?,
[thread.thread.id].into_iter().collect()
);
Ok(())
}
#[tokio::test]
async fn managed_force_shutdown_exits_with_active_work() -> Result<()> {
async fn managed_force_shutdown_exits_without_snapshotting_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 recovery_file = daemon_recovery_file_path(home.path());
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?;
@@ -94,6 +114,104 @@ async fn managed_force_shutdown_exits_with_active_work() -> Result<()> {
assert_still_running(&mut server, "graceful shutdown must wait").await;
request_shutdown(&server, &socket_path).await?;
wait_success(&mut server).await?;
assert!(!recovery_file.exists());
Ok(())
}
enum SnapshotScenario {
WriteFailure,
Ephemeral,
Archived,
Child,
Parented,
}
#[test_case::test_case(SnapshotScenario::WriteFailure)]
#[test_case::test_case(SnapshotScenario::Ephemeral)]
#[test_case::test_case(SnapshotScenario::Archived)]
#[test_case::test_case(SnapshotScenario::Child)]
#[test_case::test_case(SnapshotScenario::Parented)]
#[tokio::test]
async fn managed_shutdown_skips_nonpersistent_threads_and_tolerates_save_failure(
scenario: SnapshotScenario,
) -> Result<()> {
let home = TempDir::new()?;
create_config_toml(home.path(), "http://127.0.0.1:1", "never")?;
let source = if matches!(scenario, SnapshotScenario::Child) {
codex_protocol::protocol::SessionSource::SubAgent(
codex_protocol::protocol::SubAgentSource::ThreadSpawn {
parent_thread_id: codex_protocol::ThreadId::new(),
depth: 1,
agent_path: None,
agent_nickname: None,
agent_role: None,
},
)
} else {
codex_protocol::protocol::SessionSource::Cli
};
let stored = if matches!(scenario, SnapshotScenario::Parented) {
app_test_support::create_fake_parented_rollout_with_source(
home.path(),
"2026-09-01T12-00-00",
"2026-09-01T12:00:00Z",
"Saved task",
Some("mock_provider"),
/*git_info*/ None,
source,
codex_protocol::SessionId::new(),
codex_protocol::ThreadId::new(),
)?
} else {
app_test_support::create_fake_rollout_with_source(
home.path(),
"2026-09-01T12-00-00",
"2026-09-01T12:00:00Z",
"Saved task",
Some("mock_provider"),
/*git_info*/ None,
source,
)?
};
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?;
if matches!(scenario, SnapshotScenario::Ephemeral) {
request(
&mut client,
/*id*/ 2,
"thread/start",
json!({"ephemeral":true}),
)
.await?;
} else {
request(
&mut client,
/*id*/ 2,
"thread/resume",
json!({"threadId":stored}),
)
.await?;
}
let path = daemon_recovery_file_path(home.path());
if matches!(scenario, SnapshotScenario::WriteFailure) {
std::fs::create_dir_all(&path)?;
} else if matches!(scenario, SnapshotScenario::Archived) {
request(
&mut client,
/*id*/ 3,
"thread/archive",
json!({"threadId":stored}),
)
.await?;
}
request_shutdown(&server, &socket_path).await?;
wait_success(&mut server).await?;
if matches!(scenario, SnapshotScenario::WriteFailure) {
assert!(path.is_dir());
} else {
assert_eq!(daemon_recovery::read_candidates(&path)?, Default::default());
}
Ok(())
}