From 4d2807023aae3ee317d39a9868fb187618699df2 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Tue, 15 Sep 2026 23:20:01 +0000 Subject: [PATCH] Record interrupted turns in managed daemon recovery snapshots (#45807) ## Why Managed daemon recovery snapshots previously saved only loaded thread IDs, without identifying active turns or preserving their turn-specific options. ## What changed - Capture regular, uncanceled turns after their input is recorded, saving the turn ID, output schema, service tier, and cyber access program alongside persisted thread IDs. - Store interruption metadata atomically in the existing candidate array format so older servers can still read thread candidates. - Begin snapshotting once admitted operations drain, while turns may still be running. Run snapshot collection and thread listener attachment independently of the event loop to keep forced shutdown responsive. ## Testing Add coverage for running, completed, canceled, and compacting turns; recovery readiness for automatic and user turns; admitted resumes during shutdown; forced shutdown with a blocked rollout writer and child listener attachment; and legacy candidate-array compatibility. GitOrigin-RevId: ed46342c3a5c71b09c48fa9acece2f15ae748e2f --- codex-rs/app-server-transport/Cargo.toml | 2 +- .../src/daemon_recovery.rs | 72 +++++- .../app-server/src/daemon_thread_recovery.rs | 11 +- codex-rs/app-server/src/lib.rs | 41 ++-- codex-rs/app-server/src/message_processor.rs | 6 +- .../src/request_processors/daemon_snapshot.rs | 22 +- .../tests/suite/v2/daemon_update_recovery.rs | 210 +++++++++++++++++- codex-rs/core/src/codex_thread.rs | 5 + codex-rs/core/src/session/daemon_recovery.rs | 39 ++++ codex-rs/core/src/session/mod.rs | 1 + codex-rs/core/src/session/turn.rs | 3 + .../core/tests/suite/turn_input_submission.rs | 33 +++ 12 files changed, 407 insertions(+), 38 deletions(-) create mode 100644 codex-rs/core/src/session/daemon_recovery.rs diff --git a/codex-rs/app-server-transport/Cargo.toml b/codex-rs/app-server-transport/Cargo.toml index aa8150de07..ef56ca54a0 100644 --- a/codex-rs/app-server-transport/Cargo.toml +++ b/codex-rs/app-server-transport/Cargo.toml @@ -27,6 +27,7 @@ codex-app-server-protocol = { workspace = true } codex-core = { workspace = true } codex-login = { workspace = true } codex-model-provider = { workspace = true } +codex-protocol = { workspace = true } codex-state = { workspace = true } codex-uds = { workspace = true } codex-utils-absolute-path = { workspace = true } @@ -61,7 +62,6 @@ signal-hook = { workspace = true } [dev-dependencies] chrono = { workspace = true } codex-config = { workspace = true } -codex-protocol = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["test-util"] } diff --git a/codex-rs/app-server-transport/src/daemon_recovery.rs b/codex-rs/app-server-transport/src/daemon_recovery.rs index 85fb7a826f..9b3b590d18 100644 --- a/codex-rs/app-server-transport/src/daemon_recovery.rs +++ b/codex-rs/app-server-transport/src/daemon_recovery.rs @@ -1,21 +1,79 @@ //! Shared on-disk candidate set for managed daemon restarts. +use std::collections::BTreeMap; use std::collections::BTreeSet; use std::io; use std::path::Path; use codex_core::path_utils::write_atomically; +use serde::Deserialize; +use serde::Serialize; + +#[derive(Debug, Default, Deserialize, Serialize, PartialEq, Eq)] +pub struct RecoverySnapshot { + #[serde(skip)] + pub loaded: BTreeSet, + pub interrupted: BTreeMap, +} + +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)] +pub struct InterruptedTurn { + pub turn_id: String, + pub output_schema: Option, + pub service_tier: Option, + pub cyber_access_program: Option, +} + +// Old servers accept the array and skip this non-thread entry during best-effort +// restoration. Keeping metadata in the same atomic file avoids stale sidecars. +const INTERRUPTION_PREFIX: &str = "codex-interrupted-v1:"; + +pub fn read_snapshot(path: &Path) -> io::Result { + let mut loaded: BTreeSet = 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 => return Ok(RecoverySnapshot::default()), + Err(err) => return Err(err), + }; + let mut snapshot = RecoverySnapshot::default(); + loaded.retain(|entry| { + if let Some(metadata) = entry.strip_prefix(INTERRUPTION_PREFIX) { + if let Ok(saved) = serde_json::from_str::(metadata) { + snapshot = saved; + } + false + } else { + true + } + }); + snapshot.interrupted.retain(|id, _| loaded.contains(id)); + snapshot.loaded = loaded; + Ok(snapshot) +} + pub fn read_candidates(path: &Path) -> io::Result> { - 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), - } + Ok(read_snapshot(path)?.loaded) } pub fn write_candidates(path: &Path, candidates: &BTreeSet) -> io::Result<()> { - write_atomically( + write_snapshot( path, - &serde_json::to_string(candidates).map_err(io::Error::other)?, + &RecoverySnapshot { + loaded: candidates.clone(), + ..Default::default() + }, + ) +} + +pub fn write_snapshot(path: &Path, snapshot: &RecoverySnapshot) -> io::Result<()> { + let mut saved = snapshot.loaded.clone(); + if !snapshot.interrupted.is_empty() { + saved.insert(format!( + "{INTERRUPTION_PREFIX}{}", + serde_json::to_string(snapshot).map_err(io::Error::other)? + )); + } + write_atomically( + path, + &serde_json::to_string(&saved).map_err(io::Error::other)?, ) } diff --git a/codex-rs/app-server/src/daemon_thread_recovery.rs b/codex-rs/app-server/src/daemon_thread_recovery.rs index 0bf6e2b1aa..7460cc0499 100644 --- a/codex-rs/app-server/src/daemon_thread_recovery.rs +++ b/codex-rs/app-server/src/daemon_thread_recovery.rs @@ -2,22 +2,21 @@ //! 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; use std::path::PathBuf; use codex_app_server_transport::daemon_recovery; -pub(crate) async fn snapshot(path: PathBuf, loaded: Vec) -> io::Result<()> { +pub(crate) async fn snapshot( + path: PathBuf, + saved: daemon_recovery::RecoverySnapshot, +) -> 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::>(), - ); + let result = daemon_recovery::write_snapshot(&path, &saved); if result.is_err() && let Err(err) = std::fs::remove_file(&path) && err.kind() != io::ErrorKind::NotFound diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 36dd9dbca8..68c48a2613 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -984,6 +984,7 @@ pub async fn run_main_with_transport_options( let mut active_admissions_rx = processor.turn_admission.subscribe_active(); let mut connections = HashMap::::new(); let mut connection_cleanup_tasks = ConnectionCleanupTasks::new(); + let mut thread_listener_tasks = tokio::task::JoinSet::new(); let mut remote_control_status_rx = remote_control_handle.status_receiver(); let mut remote_control_status = remote_control_status_rx.borrow().clone(); let transport_shutdown_token = transport_shutdown_token.clone(); @@ -1007,12 +1008,16 @@ pub async fn run_main_with_transport_options( let mut listen_for_threads = true; // Keep force signals and daemon control events responsive while saving. 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 processor = Arc::clone(&processor); + let recovery_file = recovery_file.clone(); + // Run independently: snapshot locks can require the event loop to make progress. + let task = tokio::spawn(async move { + let saved = processor.daemon_recovery_snapshot().await; + if let Err(err) = daemon_thread_recovery::snapshot(recovery_file, saved).await { + warn!("failed to save threads during daemon shutdown: {err}"); + } + }); + let _ = tokio_util::task::AbortOnDropHandle::new(task).await; }); let mut snapshot_finished = !managed_daemon; let mut clients_disconnected = false; @@ -1049,7 +1054,7 @@ pub async fn run_main_with_transport_options( } tokio::select! { - _ = &mut snapshot, if ready_to_exit && !snapshot_finished => { + _ = &mut snapshot, if shutdown_state.requested() && active_admissions == 0 && !snapshot_finished => { snapshot_finished = true; } shutdown_signal_result = shutdown_signal(), if graceful_signal_restart_enabled && !shutdown_state.forced() => { @@ -1237,6 +1242,11 @@ pub async fn run_main_with_transport_options( } } _ = connection_cleanup_tasks.reap_next() => {} + result = thread_listener_tasks.join_next(), if !thread_listener_tasks.is_empty() => { + if let Some(Err(err)) = result { + warn!("thread listener attachment failed: {err}"); + } + } changed = remote_control_status_rx.changed() => { if changed.is_err() { continue; @@ -1260,12 +1270,16 @@ pub async fn run_main_with_transport_options( initialized_connection_ids.push(*connection_id); } } - processor - .try_attach_thread_listener( - thread_id, - initialized_connection_ids, - ) - .await; + let processor = Arc::clone(&processor); + // Attachment can wait on snapshot locks; keep force shutdown responsive. + thread_listener_tasks.spawn(async move { + processor + .try_attach_thread_listener( + thread_id, + initialized_connection_ids, + ) + .await; + }); } Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => { // TODO(jif) handle lag. @@ -1286,6 +1300,7 @@ pub async fn run_main_with_transport_options( task.abort(); } drop(snapshot); + drop(thread_listener_tasks); if !shutdown_state.forced() { futures::future::join_all(connections.iter().map( |(&connection_id, connection_state)| { diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 2c5bba1729..17b6d2d4e5 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -757,8 +757,10 @@ impl MessageProcessor { self.thread_processor.thread_created_receiver() } - pub(crate) async fn daemon_recovery_candidates(&self) -> Vec { - self.thread_processor.daemon_recovery_candidates().await + pub(crate) async fn daemon_recovery_snapshot( + &self, + ) -> codex_app_server_transport::daemon_recovery::RecoverySnapshot { + self.thread_processor.daemon_recovery_snapshot().await } pub(crate) async fn restore_daemon_threads( diff --git a/codex-rs/app-server/src/request_processors/daemon_snapshot.rs b/codex-rs/app-server/src/request_processors/daemon_snapshot.rs index a25554923d..e1238f4c06 100644 --- a/codex-rs/app-server/src/request_processors/daemon_snapshot.rs +++ b/codex-rs/app-server/src/request_processors/daemon_snapshot.rs @@ -2,6 +2,8 @@ //! Only threads successfully persisted through the existing thread store become candidates. use super::ThreadRequestProcessor; +use codex_app_server_transport::daemon_recovery::InterruptedTurn; +use codex_app_server_transport::daemon_recovery::RecoverySnapshot; use codex_thread_store::PersistContext; use tracing::warn; @@ -10,8 +12,8 @@ impl ThreadRequestProcessor { clippy::await_holding_invalid_type, reason = "snapshot selection must be serialized against pending unloads" )] - pub(crate) async fn daemon_recovery_candidates(&self) -> Vec { - let mut ids = Vec::new(); + pub(crate) async fn daemon_recovery_snapshot(&self) -> RecoverySnapshot { + let mut snapshot = RecoverySnapshot::default(); 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) { @@ -25,6 +27,7 @@ impl ThreadRequestProcessor { && config.parent_thread_id.is_none() && !config.session_source.is_non_root_agent() { + let interrupted = thread.interrupted_turn().await; if let Err(err) = self .thread_store .persist_thread(thread_id, PersistContext::Standard) @@ -33,9 +36,20 @@ impl ThreadRequestProcessor { warn!(%thread_id, %err, "skipping daemon restore for thread that could not be persisted"); continue; } - ids.push(thread_id.to_string()); + snapshot.loaded.insert(thread_id.to_string()); + if let Some((turn_id, options)) = interrupted { + snapshot.interrupted.insert( + thread_id.to_string(), + InterruptedTurn { + turn_id, + output_schema: options.final_output_json_schema, + service_tier: options.service_tier, + cyber_access_program: options.cyber_access_program, + }, + ); + } } } - ids + snapshot } } 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 a4585e3fcf..1aefa8f132 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 @@ -187,14 +187,13 @@ async fn managed_restart_resumes_loaded_threads_and_goal_without_client() -> Res } #[tokio::test] -async fn managed_force_shutdown_exits_without_snapshotting_active_work() -> Result<()> { +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 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?; @@ -204,7 +203,6 @@ async fn managed_force_shutdown_exits_without_snapshotting_active_work() -> Resu 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(()) } @@ -305,16 +303,105 @@ async fn managed_shutdown_skips_nonpersistent_threads_and_tolerates_save_failure Ok(()) } +#[tokio::test] +async fn managed_shutdown_preserves_admitted_resume() -> Result<()> { + let home = TempDir::new()?; + let mock = wiremock::MockServer::start().await; + let (mcp, control) = + core_test_support::apps_test_server::AppsTestServer::mount_with_startup_control(&mock) + .await?; + let release = control.hold_next_successful_initialize(); + create_config_toml(home.path(), "http://127.0.0.1:1", "never")?; + let config = home.path().join("config.toml"); + std::fs::write( + &config, + format!( + "{}\n[mcp_servers.stalled]\nurl = {:?}\nrequired = true\nstartup_timeout_sec = 120\nhttp_headers = {{ Authorization = \"Bearer synthetic-test-token\" }}\n", + std::fs::read_to_string(&config)?, + format!("{}/api/codex/ps/mcp", mcp.chatgpt_base_url), + ), + )?; + let id = app_test_support::create_fake_rollout( + home.path(), + "2026-09-01T12-00-00", + "2026-09-01T12:00:00Z", + "Saved task", + Some("mock_provider"), + /*git_info*/ None, + )?; + let socket = home.path().join("control/server.sock"); + let mut server = spawn_server(home.path(), &socket)?; + let mut client = connect_default_daemon_client(&socket).await?; + client + .send(Message::Text( + json!({"id":2,"method":"thread/resume","params":{"threadId":id}}) + .to_string() + .into(), + )) + .await?; + // Required MCP startup holds the admitted resume before its runtime is registered. + timeout(DEFAULT_READ_TIMEOUT, async { + while control.initialize_attempts() == 0 { + sleep(Duration::from_millis(25)).await; + } + }) + .await + .context("required MCP initialization did not begin")?; + request_shutdown(&server, &socket).await?; + assert_still_running(&mut server, "shutdown must wait for the admitted resume").await; + assert!(!daemon_recovery_file_path(home.path()).exists()); + release + .send(()) + .expect("release required MCP initialization"); + wait_success(&mut server).await?; + assert_eq!( + daemon_recovery::read_candidates(&daemon_recovery_file_path(home.path()))?, + [id].into(), + ); + Ok(()) +} + #[cfg(unix)] #[tokio::test] async fn managed_force_shutdown_exits_with_blocked_rollout_writer() -> Result<()> { + use core_test_support::responses; + let home = TempDir::new()?; - let mock = wiremock::MockServer::start().await; - create_config_toml(home.path(), &mock.uri(), "never")?; + let (release_spawn, spawn_gate) = oneshot::channel(); + let (_release_parent, parent_gate) = oneshot::channel(); + let (_release_child, child_gate) = oneshot::channel(); + let (mock, _completions) = start_streaming_sse_server(vec![ + vec![StreamingSseChunk { + gate: Some(spawn_gate), + body: responses::sse(vec![ + responses::ev_response_created("spawn-child"), + responses::ev_function_call_with_namespace( + "spawn-child", + "collaboration", + "spawn_agent", + r#"{"task_name":"child","message":"Wait here","fork_turns":"none"}"#, + ), + responses::ev_completed("spawn-child"), + ]), + }], + vec![stream_chunk(Some(parent_gate), "Done")?], + vec![stream_chunk(Some(child_gate), "Done")?], + ]) + .await; + create_config_toml(home.path(), mock.uri(), "never")?; + let config_path = home.path().join("config.toml"); + let config = std::fs::read_to_string(&config_path)?; + std::fs::write( + config_path, + format!("{config}\n[features.multi_agent_v2]\nenabled = true\n"), + )?; 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 parent = start_thread(&mut client, /*id*/ 3, json!({})).await?; + start_turn(&mut client, /*id*/ 4, &parent.thread.id).await?; + wait_for_requests(&mock, /*count*/ 1).await?; let rollout = thread .thread .path @@ -343,6 +430,10 @@ async fn managed_force_shutdown_exits_with_blocked_rollout_writer() -> Result<() }) .await .context("rollout writer did not open the blocked input")?; + // The running parent can create a child after admissions drain. Its listener + // needs the unload lock held by the blocked snapshot. + release_spawn.send(()).expect("parent is waiting to spawn"); + wait_for_requests(&mock, /*count*/ 3).await?; assert_still_running(&mut server, "graceful shutdown must wait for the writer").await; request_shutdown(&server, &socket_path).await?; wait_success(&mut server).await?; @@ -368,6 +459,115 @@ async fn assert_still_running(server: &mut Child, message: &str) { ); } +#[test_case::test_case("running")] +#[test_case::test_case("completed")] +#[test_case::test_case("canceled")] +#[test_case::test_case("compacting")] +#[tokio::test] +async fn managed_shutdown_records_interrupted_turn(outcome: &str) -> Result<()> { + let home = TempDir::new()?; + let (release, gate) = oneshot::channel(); + let (_release_compaction, compaction_gate) = oneshot::channel(); + let (mock, _completions) = start_streaming_sse_server(vec![ + vec![stream_chunk(Some(gate), "Original")?], + vec![stream_chunk(Some(compaction_gate), "Compacted")?], + ]) + .await; + create_config_toml(home.path(), mock.uri(), "never")?; + let socket = home.path().join("control/server.sock"); + let mut server = spawn_server(home.path(), &socket)?; + let mut client = connect_default_daemon_client(&socket).await?; + let thread = start_thread(&mut client, /*id*/ 2, json!({})).await?; + let id = thread.thread.id; + let schema = json!({"type":"object","properties":{},"additionalProperties":false}); + let turn = request( + &mut client, + /*id*/ 3, + "turn/start", + json!({ + "threadId":id,"input":[{"type":"text","text":"Do the work"}],"outputSchema":schema + }), + ) + .await?; + wait_for_requests(&mock, /*count*/ 1).await?; + match outcome { + "completed" => { + release.send(()).expect("original request is waiting"); + timeout(DEFAULT_READ_TIMEOUT, async { + loop { + let Message::Text(text) = client.next().await.context("socket closed")?? else { + continue; + }; + if let JSONRPCMessage::Notification(notification) = serde_json::from_str(&text)? + && notification.method == "turn/completed" + { + break; + } + } + Ok::<(), anyhow::Error>(()) + }) + .await??; + } + "canceled" => { + request( + &mut client, + /*id*/ 4, + "turn/interrupt", + json!({"threadId":id,"turnId":turn["turn"]["id"]}), + ) + .await?; + } + "compacting" => { + request( + &mut client, + /*id*/ 4, + "thread/compact/start", + json!({"threadId":id}), + ) + .await?; + wait_for_requests(&mock, /*count*/ 2).await?; + } + "running" => {} + _ => unreachable!(), + } + request_shutdown(&server, &socket).await?; + let path = daemon_recovery_file_path(home.path()); + timeout(DEFAULT_READ_TIMEOUT, async { + while !path.exists() { + sleep(Duration::from_millis(25)).await; + } + }) + .await?; + let expected = if outcome == "running" { + [( + id.clone(), + daemon_recovery::InterruptedTurn { + turn_id: turn["turn"]["id"].as_str().context("turn ID")?.to_string(), + output_schema: Some(schema), + service_tier: Some("default".into()), + cyber_access_program: None, + }, + )] + .into() + } else { + Default::default() + }; + assert_eq!( + daemon_recovery::read_snapshot(&path)?, + daemon_recovery::RecoverySnapshot { + loaded: [id.clone()].into(), + interrupted: expected, + } + ); + // Existing binaries can still read the candidate array and ignore the metadata entry. + let legacy: std::collections::BTreeSet = serde_json::from_slice(&std::fs::read(path)?)?; + assert!(legacy.contains(&id)); + if server.try_wait()?.is_none() { + server.kill().await?; + } + Ok(()) +} + fn stream_chunk(gate: Option>, body: &str) -> Result { Ok(StreamingSseChunk { gate, diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 76c5ad7df9..a2f4e9bf63 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -497,6 +497,11 @@ impl CodexThread { self.session.inject_if_running(items).await } + /// Captures a regular turn only after its input is recorded. The caller must flush the rollout. + pub async fn interrupted_turn(&self) -> Option<(String, TurnStartOptions)> { + self.session.interrupted_turn().await + } + /// Returns the trusted root when the expected turn is currently active. pub async fn active_turn_root(&self, expected_turn_id: &str) -> Option { let active = self.session.active_turn.lock().await; diff --git a/codex-rs/core/src/session/daemon_recovery.rs b/codex-rs/core/src/session/daemon_recovery.rs new file mode 100644 index 0000000000..8e12ee026b --- /dev/null +++ b/codex-rs/core/src/session/daemon_recovery.rs @@ -0,0 +1,39 @@ +//! Captures recorded, uncanceled regular turns for managed daemon recovery. +//! Callers must flush the rollout after capture before persisting the snapshot. + +use super::Session; + +/// Turn input and turn-start injections have entered the rollout writer. +pub(super) struct RecordedTurnInput; + +impl Session { + /// Captures a regular turn only after its input is recorded. The caller must flush the rollout. + pub(crate) async fn interrupted_turn(&self) -> Option<(String, crate::TurnStartOptions)> { + let active = self.active_turn.lock().await; + let task = active.as_ref()?.task.as_ref()?; + if task.kind != crate::state::TaskKind::Regular || task.cancellation_token.is_cancelled() { + return None; + } + let context = &task.turn_context; + context.extension_data.get::()?; + Some(( + context.sub_id.clone(), + crate::TurnStartOptions { + final_output_json_schema: context.final_output_json_schema.clone(), + service_tier: Some( + context + .current_settings + .load() + .service_tier + .clone() + .unwrap_or_else(|| { + codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE + .to_string() + }), + ), + cyber_access_program: context.cyber_access_program, + ..Default::default() + }, + )) + } +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index cc78d189f2..9252c8952d 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -227,6 +227,7 @@ use codex_protocol::exec_output::StreamOutput; mod code_mode_warning; pub(crate) mod context_window; +mod daemon_recovery; mod environment; mod extension_interruption; pub(crate) mod extension_metrics; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 8ab44553be..1f80548e4e 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -35,6 +35,7 @@ use crate::responses_retry::ResponsesStreamRetryState; use crate::responses_retry::handle_retryable_response_stream_error; use crate::session::PreviousTurnSettings; use crate::session::TurnInput; +use crate::session::daemon_recovery::RecordedTurnInput; use crate::session::session::Session; use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; @@ -445,6 +446,8 @@ pub(crate) async fn run_turn( break; } + // Input and turn-start injections are recorded before recovery can continue this turn. + turn_context.extension_data.insert(RecordedTurnInput); let window_id = sess.current_window_id().await; super::rollout_budget::maybe_record_reminder( sess.as_ref(), diff --git a/codex-rs/core/tests/suite/turn_input_submission.rs b/codex-rs/core/tests/suite/turn_input_submission.rs index 5a4c2ce586..ff1d56c003 100644 --- a/codex-rs/core/tests/suite/turn_input_submission.rs +++ b/codex-rs/core/tests/suite/turn_input_submission.rs @@ -832,3 +832,36 @@ async fn start_or_steer_turn_requires_matching_active_output_schema() { assert!(!second_request.contains("rejected steer")); server.shutdown().await; } + +#[test_case(Vec::new(); "automatic")] +#[test_case(vec![UserInput::Text { text: "Do the work".into(), text_elements: Vec::new() }]; "user")] +#[tokio::test] +async fn sampling_is_ready_for_daemon_recovery(input: Vec) -> anyhow::Result<()> { + let (release, gate) = oneshot::channel(); + let (server, _completions) = start_streaming_sse_server(vec![vec![StreamingSseChunk { + gate: Some(gate), + body: responses::sse_completed("automatic"), + }]]) + .await; + let test = test_codex().build_with_streaming_server(&server).await?; + let StartIfIdleSubmission::Started { turn_id } = test + .codex + .start_turn_if_idle(TurnInputRequest::user_input(input)) + .await? + else { + panic!("sampling should start"); + }; + timeout( + Duration::from_secs(5), + server.wait_for_request_count(/*count*/ 1), + ) + .await?; + let active = test.codex.interrupted_turn().await; + assert_eq!(active.map(|(id, _)| id), Some(turn_id)); + release.send(()).expect("sampling is waiting"); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + Ok(()) +}