diff --git a/codex-rs/app-server/src/daemon_thread_recovery.rs b/codex-rs/app-server/src/daemon_thread_recovery.rs index fbe866ece6..0bf6e2b1aa 100644 --- a/codex-rs/app-server/src/daemon_thread_recovery.rs +++ b/codex-rs/app-server/src/daemon_thread_recovery.rs @@ -1,4 +1,6 @@ -//! Persists the candidate set for a planned managed-daemon replacement. +//! Saves persistent root thread IDs and restores them through the shared internal resume path. +//! Recovery uses normal cold-resume semantics without delaying readiness for runtime loading. +//! Already-loaded runtimes remain owned by their current clients. use std::collections::BTreeSet; use std::io; @@ -26,3 +28,24 @@ pub(crate) async fn snapshot(path: PathBuf, loaded: Vec) -> io::Result<( })?; result_rx.await.map_err(io::Error::other)? } + +/// Consume the handoff before serving requests, then restore runtimes in the background. +pub(crate) async fn start_recovery( + path: PathBuf, + processor: std::sync::Arc, +) -> io::Result> { + let candidates = tokio::task::spawn_blocking(move || { + let candidates = daemon_recovery::read_candidates(&path); + // Even malformed or temporarily unreadable snapshots belong to this generation only. + match std::fs::remove_file(&path) { + Ok(()) => candidates, + Err(err) if err.kind() == io::ErrorKind::NotFound => candidates, + Err(err) => Err(err), + } + }) + .await + .map_err(io::Error::other)??; + Ok(tokio::spawn(async move { + processor.restore_daemon_threads(candidates).await; + })) +} diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 42b8238b33..daf87590ed 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -985,14 +985,22 @@ 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 recovery_task = if managed_daemon { + match daemon_thread_recovery::start_recovery( + recovery_file.clone(), + Arc::clone(&processor), + ) + .await + { + Ok(task) => Some(task), + Err(err) => { + warn!("failed to consume daemon recovery snapshot: {err}"); + None + } + } + } else { + None + }; let mut listen_for_threads = true; // Keep force signals and daemon control events responsive while saving. let mut snapshot = Box::pin(async { @@ -1016,6 +1024,9 @@ pub async fn run_main_with_transport_options( ShutdownAction::Finish ); if ready_to_exit { + if let Some(task) = &recovery_task { + task.abort(); + } let finished = snapshot_finished || shutdown_state.forced(); if finished { transport_shutdown_token.cancel(); @@ -1275,6 +1286,9 @@ pub async fn run_main_with_transport_options( } }; + if let Some(task) = recovery_task { + task.abort(); + } drop(snapshot); if !shutdown_state.forced() { futures::future::join_all(connections.iter().map( diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 77c29040d9..040192c75b 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -46,6 +46,7 @@ use crate::request_processors::SearchRequestProcessor; use crate::request_processors::ThreadGoalRequestProcessor; use crate::request_processors::ThreadQueueRequestProcessor; use crate::request_processors::ThreadRequestProcessor; +use crate::request_processors::ThreadResumeTarget; use crate::request_processors::TurnRequestProcessor; use crate::request_processors::WindowsSandboxRequestProcessor; use crate::request_processors::read_server_diagnostics; @@ -763,6 +764,34 @@ impl MessageProcessor { self.thread_processor.daemon_recovery_candidates().await } + pub(crate) async fn restore_daemon_threads( + &self, + candidates: std::collections::BTreeSet, + ) { + for thread_id in candidates { + let Ok(_permit) = self.turn_admission.admit() else { + break; + }; + if let Err(err) = self + .thread_processor + .thread_resume( + ThreadResumeTarget::DaemonRecovery, + codex_app_server_protocol::ThreadResumeParams { + thread_id, + exclude_turns: true, + ..Default::default() + }, + /*app_server_client_name*/ None, + /*app_server_client_version*/ None, + ClientMcpExtensions::default(), + ) + .await + { + tracing::warn!(code = err.code, "failed to restore saved daemon thread"); + } + } + } + pub(crate) async fn send_initialize_notifications_to_connection( &self, connection_id: ConnectionId, @@ -1265,7 +1294,7 @@ impl MessageProcessor { ClientRequest::ThreadResume { params, .. } => { self.thread_processor .thread_resume( - request_id.clone(), + ThreadResumeTarget::Client(request_id.clone()), params, app_server_client_name.clone(), client_version.clone(), diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 9795665a98..a79759d9d0 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -591,6 +591,7 @@ pub(crate) use remote_control_processor::RemoteControlRequestProcessor; pub(crate) use search::SearchRequestProcessor; pub(crate) use thread_goal_processor::ThreadGoalRequestProcessor; pub(crate) use thread_processor::ThreadRequestProcessor; +pub(crate) use thread_processor::ThreadResumeTarget; pub(crate) use thread_queue_processor::ThreadQueueRequestProcessor; pub(crate) use turn_processor::TurnRequestProcessor; pub(crate) use windows_sandbox_processor::WindowsSandboxRequestProcessor; diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index e17e0e38d7..172177cdac 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -460,6 +460,12 @@ pub(crate) struct ThreadRequestProcessor { pub(super) initial_config_warnings: Arc>, } +/// Whether resume attaches a client or restores a cold runtime during daemon startup. +pub(crate) enum ThreadResumeTarget { + Client(ConnectionRequestId), + DaemonRecovery, +} + /// Outcome of trying to satisfy a resume request from an already loaded thread. enum RunningThreadResumeResult { /// The request was delegated to the loaded thread. @@ -547,7 +553,7 @@ impl ThreadRequestProcessor { pub(crate) async fn thread_resume( &self, - request_id: ConnectionRequestId, + target: ThreadResumeTarget, params: ThreadResumeParams, app_server_client_name: Option, app_server_client_version: Option, @@ -556,7 +562,7 @@ impl ThreadRequestProcessor { let mut prepared_config = None; // Keep the resume future off the request handler's stack. while Box::pin(self.thread_resume_inner( - request_id.clone(), + &target, ¶ms, app_server_client_name.clone(), app_server_client_version.clone(), @@ -3634,7 +3640,7 @@ impl ThreadRequestProcessor { async fn thread_resume_inner( &self, - request_id: ConnectionRequestId, + target: &ThreadResumeTarget, params: &ThreadResumeParams, app_server_client_name: Option, app_server_client_version: Option, @@ -3648,56 +3654,47 @@ impl ThreadRequestProcessor { .await .contains(&thread_id) { - self.outgoing - .send_error( - request_id, - invalid_request(format!( - "thread {thread_id} is closing; retry thread/resume after the thread is closed" - )), - ) - .await; - return Ok(ControlFlow::Break(())); + return Err(invalid_request(format!( + "thread {thread_id} is closing; retry thread/resume after the thread is closed" + ))); } if params.sandbox.is_some() && params.permissions.is_some() { - self.outgoing - .send_error( - request_id, - invalid_request("`permissions` cannot be combined with `sandbox`"), - ) - .await; - return Ok(ControlFlow::Break(())); + return Err(invalid_request( + "`permissions` cannot be combined with `sandbox`", + )); } let redact_resume_payloads = should_redact_thread_resume_payloads(app_server_client_name.as_deref()); - let _thread_list_state_permit = match self.acquire_thread_list_state_permit().await { - Ok(permit) => permit, - Err(error) => { - self.outgoing.send_error(request_id, error).await; - return Ok(ControlFlow::Break(())); - } - }; - let stored_thread_from_running_probe = match self - .resume_running_thread( - &request_id, - params, - app_server_client_name.clone(), - app_server_client_version.clone(), - /*cold_resume_history*/ None, - ) - .await - { - Ok(RunningThreadResumeResult::Handled(completion)) => { - // The listener may need this permit to finish the response. - drop(_thread_list_state_permit); - let _ = completion.await; - return Ok(ControlFlow::Break(())); - } - Ok(RunningThreadResumeResult::NotRunning(stored_thread)) => stored_thread, - Err(error) => { - self.outgoing.send_error(request_id, error).await; - return Ok(ControlFlow::Break(())); + let _thread_list_state_permit = self.acquire_thread_list_state_permit().await?; + let stored_thread_from_running_probe = match target { + ThreadResumeTarget::Client(request_id) => match self + .resume_running_thread( + request_id, + params, + app_server_client_name.clone(), + app_server_client_version.clone(), + /*cold_resume_history*/ None, + ) + .await? + { + RunningThreadResumeResult::Handled(completion) => { + // The listener may need this permit to finish the response. + drop(_thread_list_state_permit); + let _ = completion.await; + return Ok(ControlFlow::Break(())); + } + RunningThreadResumeResult::NotRunning(stored_thread) => stored_thread, + }, + ThreadResumeTarget::DaemonRecovery => { + let thread_id = ThreadId::from_string(¶ms.thread_id) + .map_err(|err| invalid_request(format!("invalid thread id: {err}")))?; + // Recheck under the same permit as client resume, including after config loading. + if self.thread_manager.get_thread(thread_id).await.is_ok() { + return Ok(ControlFlow::Break(())); + } + None } }; @@ -3747,13 +3744,7 @@ impl ThreadRequestProcessor { Err(error) => Err(error), } }; - let (thread_history, resume_source_thread) = match resume_result { - Ok(value) => value, - Err(error) => { - self.outgoing.send_error(request_id, error).await; - return Ok(ControlFlow::Break(())); - } - }; + let (thread_history, resume_source_thread) = resume_result?; if let InitialHistory::Resumed(resumed) = &thread_history && self .pending_thread_unloads @@ -3770,7 +3761,11 @@ impl ThreadRequestProcessor { matches!(thread.history_mode, ThreadHistoryMode::Paginated).then_some(thread.thread_id) }); let paginated_resume = paginated_thread_id.is_some(); - if paginated_resume && include_turns && prepared_config.is_none() { + if paginated_resume + && include_turns + && prepared_config.is_none() + && let ThreadResumeTarget::Client(request_id) = target + { self.send_deprecation_notice( request_id.connection_id, PAGINATED_FULL_HISTORY_DEPRECATION_SUMMARY, @@ -3783,6 +3778,9 @@ impl ThreadRequestProcessor { && let Some((source, _)) = thread_history.get_resumed_session_sources() && !can_accept_direct_input(thread_history.get_multi_agent_version(), &source) { + let ThreadResumeTarget::Client(request_id) = target else { + return Ok(ControlFlow::Break(())); + }; let child_thread_id = resumed_history.conversation_id; self.thread_manager .ensure_multi_agent_v2_child_loaded(child_thread_id) @@ -3808,7 +3806,7 @@ impl ThreadRequestProcessor { }; return match self .resume_running_thread( - &request_id, + request_id, &attach_params, app_server_client_name, app_server_client_version, @@ -3913,20 +3911,9 @@ impl ThreadRequestProcessor { .as_mut() .and_then(|overrides| overrides.remove("approval_policy")) { - let approval_policy = match serde_json::from_value(value) { - Ok(approval_policy) => approval_policy, - Err(err) => { - self.outgoing - .send_error( - request_id, - invalid_params(format!( - "invalid `approval_policy` config override: {err}" - )), - ) - .await; - return Ok(ControlFlow::Break(())); - } - }; + let approval_policy = serde_json::from_value(value).map_err(|err| { + invalid_params(format!("invalid `approval_policy` config override: {err}")) + })?; typesafe_overrides.approval_policy = Some(approval_policy); } let has_explicit_model_resume_override = @@ -3959,19 +3946,11 @@ impl ThreadRequestProcessor { _ => { // Config loading can call back into Desktop; release the permit during host work. drop(_thread_list_state_permit); - let config = match self + let config = self .config_manager .load_for_cwd(request_overrides, typesafe_overrides, history_cwd) .await - { - Ok(config) => config, - Err(err) => { - self.outgoing - .send_error(request_id, config_load_error(&err)) - .await; - return Ok(ControlFlow::Break(())); - } - }; + .map_err(|err| config_load_error(&err))?; *prepared_config = Some(PreparedResumeConfig { state: config_state, config, @@ -3992,7 +3971,12 @@ impl ThreadRequestProcessor { config, thread_history, self.auth_manager.clone(), - self.request_trace_context(&request_id).await, + match target { + ThreadResumeTarget::Client(request_id) => { + self.request_trace_context(request_id).await + } + ThreadResumeTarget::DaemonRecovery => None, + }, client_mcp_extensions, ) .await @@ -4003,6 +3987,21 @@ impl ThreadRequestProcessor { session_configured, .. }) => { + let ThreadResumeTarget::Client(request_id) = target else { + // Observe lifecycle events without attaching a client subscription. + self.thread_watch_manager + .upsert_thread(&thread_id.to_string()) + .await; + // Invoke idle work before arming subscriber-based unloading. + codex_thread + .emit_thread_idle_lifecycle_if_idle(ThreadIdleCause::Completed) + .await; + let state = self.thread_state_manager.thread_state(thread_id).await; + self.ensure_listener_task_running(thread_id, Arc::clone(&codex_thread), state) + .await?; + return Ok(ControlFlow::Break(())); + }; + let request_id = request_id.clone(); if let Err(err) = Self::set_app_server_client_info( codex_thread.as_ref(), app_server_client_name, @@ -4213,7 +4212,7 @@ impl ThreadRequestProcessor { CodexErrorDetails::InvalidRequest(message) => invalid_request(message.clone()), _ => internal_error(format!("error resuming thread: {err}")), }; - self.outgoing.send_error(request_id, error).await; + return Err(error); } } Ok(ControlFlow::Break(())) diff --git a/codex-rs/app-server/tests/suite/v2/daemon_update_recovery.rs b/codex-rs/app-server/tests/suite/v2/daemon_update_recovery.rs index 2abd5efe23..a4585e3fcf 100644 --- a/codex-rs/app-server/tests/suite/v2/daemon_update_recovery.rs +++ b/codex-rs/app-server/tests/suite/v2/daemon_update_recovery.rs @@ -35,15 +35,20 @@ use tokio_tungstenite::client_async; use tokio_tungstenite::tungstenite::Message; #[tokio::test] -async fn managed_shutdown_saves_loaded_thread_after_active_turn() -> Result<()> { +async fn managed_restart_resumes_loaded_threads_and_goal_without_client() -> 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; + let (release_recovered, recovered_gate) = oneshot::channel(); + let (mock, _completions) = start_streaming_sse_server(vec![ + vec![stream_chunk(Some(response_gate), "Done")?], + vec![stream_chunk(Some(recovered_gate), "Recovered")?], + ]) + .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())?; + std::fs::create_dir_all(recovery_file.parent().context("snapshot parent")?)?; + std::fs::write(&recovery_file, "invalid JSON")?; let mut server = spawn_server(home.path(), &socket_path)?; let mut client = connect_daemon_client( &socket_path, @@ -83,16 +88,101 @@ async fn managed_shutdown_saves_loaded_thread_after_active_turn() -> Result<()> ] { request(&mut client, id, method, params).await?; } + let idle = start_thread(&mut client, /*id*/ 7, json!({})).await?; start_turn(&mut client, /*id*/ 5, &thread.thread.id).await?; wait_for_requests(&mock, /*count*/ 1).await?; + request( + &mut client, + /*id*/ 6, + "thread/goal/set", + json!({ + "threadId":thread.thread.id,"objective":"continue after restart" + }), + ) + .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?; assert_eq!( daemon_recovery::read_candidates(&recovery_file)?, - [thread.thread.id].into_iter().collect() + [thread.thread.id.clone(), idle.thread.id.clone()] + .into_iter() + .collect() ); + let mut saved = daemon_recovery::read_candidates(&recovery_file)?; + saved.insert("!missing".into()); + daemon_recovery::write_candidates(&recovery_file, &saved)?; + let mut successor = spawn_server(home.path(), &socket_path)?; + // No client reconnects until the goal has made an inference request. + wait_for_requests(&mock, /*count*/ 2).await?; + let requests = mock.requests().await; + let recovered: serde_json::Value = serde_json::from_slice(&requests[1])?; + assert_eq!(recovered["model"], "gpt-5.2"); + assert!(recovered["input"] + .as_array() + .context("recovered conversation")? + .iter() + .any(|item| item["role"] == "assistant" && item["content"].to_string().contains("Done"))); + assert!( + recovered["tools"].to_string().contains("lookup_ticket"), + "{recovered}" + ); + assert!(!recovery_file.exists()); + let mut reconnected = 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?; + let resumed = request( + &mut reconnected, + /*id*/ 2, + "thread/resume", + json!({ + "threadId":thread.thread.id,"excludeTurns":true + }), + ) + .await?; + assert_eq!(resumed["thread"]["id"], thread.thread.id); + timeout(DEFAULT_READ_TIMEOUT, async { + loop { + let loaded = request( + &mut reconnected, + /*id*/ 4, + "thread/loaded/list", + json!({}), + ) + .await?; + if loaded["data"] + .as_array() + .context("loaded IDs")? + .contains(&json!(idle.thread.id)) + { + return Ok::<(), anyhow::Error>(()); + } + sleep(Duration::from_millis(25)).await; + } + }) + .await??; + request( + &mut reconnected, + /*id*/ 3, + "thread/goal/clear", + json!({"threadId":thread.thread.id}), + ) + .await?; + release_recovered + .send(()) + .expect("recovered goal is waiting"); + request_shutdown(&successor, &socket_path).await?; + wait_success(&mut successor).await?; Ok(()) }