From 49530f6570bbc4d3db12c54eea8d45cbedb796c5 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Fri, 7 Aug 2026 16:53:27 +0300 Subject: [PATCH 01/20] Fix crash-safe scheduled occurrences Agent: cossus --- .../src/message_processor_schedule_tests.rs | 32 + .../request_processors/thread_lifecycle.rs | 30 +- .../thread_lifecycle/scheduled_runs.rs | 205 ++ .../thread_schedule_runtime.rs | 969 ++------- .../thread_schedule_runtime/occurrence.rs | 32 + .../occurrence/execution.rs | 511 +++++ .../occurrence/terminal.rs | 269 +++ .../occurrence_tests.rs | 604 ++++++ codex-rs/core/src/codex_thread.rs | 38 + codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/session/inject.rs | 119 +- codex-rs/core/src/session/input_queue.rs | 28 + codex-rs/core/src/session/tests.rs | 46 + codex-rs/core/src/tasks/mod.rs | 160 +- .../core/src/tools/handlers/loop_control.rs | 87 +- .../src/tools/handlers/schedule_control.rs | 42 + .../0063_thread_schedule_occurrence_state.sql | 198 ++ codex-rs/state/src/lib.rs | 2 + codex-rs/state/src/runtime.rs | 2 + codex-rs/state/src/runtime/schedules.rs | 1881 ++++------------- .../state/src/runtime/schedules/occurrence.rs | 159 ++ .../src/runtime/schedules/occurrence/claim.rs | 564 +++++ .../runtime/schedules/occurrence/finish.rs | 730 +++++++ .../src/runtime/schedules/occurrence/start.rs | 251 +++ .../src/runtime/schedules/occurrence_tests.rs | 719 +++++++ 25 files changed, 5321 insertions(+), 2358 deletions(-) create mode 100644 codex-rs/app-server/src/request_processors/thread_lifecycle/scheduled_runs.rs create mode 100644 codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence.rs create mode 100644 codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/execution.rs create mode 100644 codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs create mode 100644 codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence_tests.rs create mode 100644 codex-rs/state/migrations/0063_thread_schedule_occurrence_state.sql create mode 100644 codex-rs/state/src/runtime/schedules/occurrence.rs create mode 100644 codex-rs/state/src/runtime/schedules/occurrence/claim.rs create mode 100644 codex-rs/state/src/runtime/schedules/occurrence/finish.rs create mode 100644 codex-rs/state/src/runtime/schedules/occurrence/start.rs create mode 100644 codex-rs/state/src/runtime/schedules/occurrence_tests.rs diff --git a/codex-rs/app-server/src/message_processor_schedule_tests.rs b/codex-rs/app-server/src/message_processor_schedule_tests.rs index ac964bb0e..bc62cb088 100644 --- a/codex-rs/app-server/src/message_processor_schedule_tests.rs +++ b/codex-rs/app-server/src/message_processor_schedule_tests.rs @@ -440,6 +440,38 @@ impl ScheduleHarness { }) .await? .expect("schedule should claim for seeded failure"); + let turn_id = claim + .run + .turn_id + .as_deref() + .expect("claimed occurrence should have a stable turn id"); + self.state_db + .thread_schedules() + .enqueue_thread_schedule_run(codex_state::ThreadScheduleRunEnqueueParams { + schedule_id, + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + goal_id: None, + auth_profile_recorded: false, + auth_profile: None, + turn_input: "seeded app-server schedule failure", + now, + }) + .await? + .expect("seeded failure occurrence should enqueue"); + self.state_db + .thread_schedules() + .mark_thread_schedule_run_started(codex_state::ThreadScheduleRunStartParams { + schedule_id, + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + turn_id, + goal_id: None, + now, + lease_duration: std::time::Duration::from_secs(300), + }) + .await? + .expect("seeded failure occurrence should start"); self.state_db .thread_schedules() .fail_thread_schedule_run( diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs index 3ac34b588..985c49847 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -2,6 +2,8 @@ use super::*; use crate::request_processors::thread_goal_processor::api_thread_goal_from_state; use crate::request_processors::thread_goal_processor::api_thread_goal_plan_from_state_for_thread; +mod scheduled_runs; + pub(super) const THREAD_UNLOADING_DELAY: Duration = Duration::from_secs(30 * 60); #[derive(Clone)] @@ -321,25 +323,17 @@ pub(super) async fn ensure_listener_task_running( // Track the event before emitting any typed translations // so thread-local state such as raw event opt-in stays // synchronized with the conversation. - let terminal_event = matches!( - event.msg, - EventMsg::TurnComplete(_) | EventMsg::TurnAborted(_) | EventMsg::Error(_) - ); - let (raw_events_enabled, tracked_scheduled_run, turn_error) = { + let scheduled_runs::TrackedScheduledEvent { + raw_events_enabled, + terminal_event, + scheduled_run: tracked_scheduled_run, + turn_error, + } = { let mut thread_state = thread_state.lock().await; - thread_state.track_current_turn_event(&event.id, &event.msg); - let tracked_scheduled_run = if terminal_event { - thread_state.take_scheduled_run(&event.id) - } else { - None - }; - let turn_error = terminal_event - .then(|| thread_state.turn_summary.last_error.clone()) - .flatten(); - ( - thread_state.experimental_raw_events, - tracked_scheduled_run, - turn_error, + scheduled_runs::track_scheduled_event( + &mut thread_state, + event.id.as_str(), + &event.msg, ) }; let terminal_scheduled_run = match ( diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle/scheduled_runs.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle/scheduled_runs.rs new file mode 100644 index 000000000..e2c2c3acf --- /dev/null +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle/scheduled_runs.rs @@ -0,0 +1,205 @@ +//! Scheduled-run tracking owned by the live thread event listener. + +use super::*; + +pub(super) struct TrackedScheduledEvent { + pub(super) raw_events_enabled: bool, + pub(super) terminal_event: bool, + pub(super) scheduled_run: Option, + pub(super) turn_error: Option, +} + +pub(super) fn track_scheduled_event( + thread_state: &mut ThreadState, + turn_id: &str, + event: &EventMsg, +) -> TrackedScheduledEvent { + let terminal_event = matches!(event, EventMsg::TurnComplete(_) | EventMsg::TurnAborted(_)) + || matches!(event, EventMsg::Error(error) if error.affects_turn_status()); + thread_state.track_current_turn_event(turn_id, event); + let scheduled_run = terminal_event + .then(|| thread_state.take_scheduled_run(turn_id)) + .flatten(); + let turn_error = terminal_event + .then(|| thread_state.turn_summary.last_error.clone()) + .flatten(); + TrackedScheduledEvent { + raw_events_enabled: thread_state.experimental_raw_events, + terminal_event, + scheduled_run, + turn_error, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use codex_protocol::protocol::CodexErrorInfo; + use codex_protocol::protocol::ErrorEvent; + use codex_protocol::protocol::TurnCompleteEvent; + use pretty_assertions::assert_eq; + + #[tokio::test] + async fn non_affecting_error_keeps_scheduled_run_running_until_turn_complete() { + let temp_dir = tempfile::tempdir().expect("temp dir should be created"); + let state_db = codex_state::StateRuntime::init( + temp_dir.path().to_path_buf(), + "fallback-provider".to_string(), + ) + .await + .expect("state db should initialize"); + let thread_id = ThreadId::new(); + let now = chrono::DateTime::::from_timestamp(1_700_000_000, 0) + .expect("test timestamp should be valid"); + let mut builder = ThreadMetadataBuilder::new( + thread_id, + temp_dir.path().join("thread.jsonl"), + now, + SessionSource::Cli, + ); + builder.cwd = temp_dir.path().join("workspace"); + state_db + .upsert_thread(&builder.build("fallback-provider")) + .await + .expect("thread metadata should persist"); + let schedule = state_db + .thread_schedules() + .create_thread_schedule(codex_state::ThreadScheduleCreateParams { + thread_id, + prompt: "finish after a non-terminal error".to_string(), + prompt_source: codex_state::ThreadSchedulePromptSource::Inline, + schedule: codex_state::ThreadScheduleSpec::Interval( + codex_state::ThreadScheduleInterval { + amount: 5, + unit: codex_state::ThreadScheduleIntervalUnit::Minutes, + }, + ), + timezone: "UTC".to_string(), + status: codex_state::ThreadScheduleStatus::Active, + next_run_at: Some(now), + expires_at: None, + }) + .await + .expect("schedule should create"); + let claim = state_db + .thread_schedules() + .claim_due_thread_schedule(now, "lease-live", Duration::from_secs(300)) + .await + .expect("schedule claim should succeed") + .expect("schedule should be due"); + let turn_id = claim + .run + .turn_id + .as_deref() + .expect("claimed occurrence should reserve a turn") + .to_string(); + state_db + .thread_schedules() + .enqueue_thread_schedule_run(codex_state::ThreadScheduleRunEnqueueParams { + schedule_id: schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + goal_id: None, + auth_profile_recorded: true, + auth_profile: None, + turn_input: "scheduled input", + now, + }) + .await + .expect("occurrence should enqueue") + .expect("owned occurrence should enqueue"); + let running = state_db + .thread_schedules() + .mark_thread_schedule_run_started(codex_state::ThreadScheduleRunStartParams { + schedule_id: schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + turn_id: turn_id.as_str(), + goal_id: None, + now, + lease_duration: Duration::from_secs(300), + }) + .await + .expect("occurrence should start") + .expect("owned occurrence should materialize one run"); + + let mut thread_state = ThreadState::default(); + thread_state.track_scheduled_run( + turn_id.clone(), + crate::thread_state::ScheduledThreadScheduleRun { + schedule_id: schedule.schedule_id.clone(), + run_id: running.run_id.clone(), + lease_id: running.lease_id.clone(), + goal_id: None, + state_db: state_db.clone(), + }, + ); + let non_terminal_error = EventMsg::Error(ErrorEvent { + message: "rollback request failed".to_string(), + codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed), + }); + let tracked = + track_scheduled_event(&mut thread_state, turn_id.as_str(), &non_terminal_error); + assert!(!tracked.terminal_event); + assert!(tracked.scheduled_run.is_none()); + assert!(thread_state.has_scheduled_run(turn_id.as_str())); + assert_eq!( + codex_state::ThreadScheduleRunStatus::Running, + state_db + .thread_schedules() + .get_thread_schedule_run(running.run_id.as_str()) + .await + .expect("running row should load") + .expect("running row should remain") + .status + ); + + let complete = EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: turn_id.clone(), + last_agent_message: Some("finished after rollback warning".to_string()), + completed_at: Some(1_700_000_001), + duration_ms: Some(1_000), + time_to_first_token_ms: Some(100), + }); + let tracked = track_scheduled_event(&mut thread_state, turn_id.as_str(), &complete); + assert!(tracked.terminal_event); + assert!(!thread_state.has_scheduled_run(turn_id.as_str())); + let scheduled_run = tracked + .scheduled_run + .expect("terminal event should take the tracked run once"); + let (outgoing_tx, _outgoing_rx) = mpsc::channel(/*buffer*/ 8); + let outgoing = Arc::new(OutgoingMessageSender::new( + outgoing_tx, + codex_analytics::AnalyticsEventsClient::disabled(), + )); + super::super::thread_schedule_runtime::finish_scheduled_run_after_turn( + thread_id, + scheduled_run, + &complete, + tracked.turn_error, + &outgoing, + ) + .await; + assert_eq!( + codex_state::ThreadScheduleRunStatus::Completed, + state_db + .thread_schedules() + .get_thread_schedule_run(running.run_id.as_str()) + .await + .expect("completed row should load") + .expect("completed row should remain") + .status + ); + assert!( + super::super::thread_schedule_runtime::recover_scheduled_run_for_terminal_turn( + &state_db, + thread_id, + turn_id.as_str(), + ) + .await + .expect("completed run recovery should not fail") + .is_none(), + "a completed run must not be finalized twice" + ); + } +} diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs index 278ca500c..6bc70c64f 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs @@ -13,6 +13,22 @@ use croner::Cron; use std::fmt::Write as _; use std::str::FromStr; +mod occurrence; +#[cfg(test)] +use occurrence::PersistedScheduledTurnTerminal; +#[cfg(test)] +use occurrence::ScheduledTurnFinish; +pub(super) use occurrence::default_thread_schedule_expires_at; +pub(super) use occurrence::finish_scheduled_run_after_turn; +use occurrence::next_thread_schedule_run_after_completion; +pub(super) use occurrence::next_thread_schedule_run_at; +pub(super) use occurrence::normalize_schedule_timezone; +#[cfg(test)] +use occurrence::persisted_scheduled_turn_terminal; +pub(super) use occurrence::recover_scheduled_run_for_terminal_turn; +#[cfg(test)] +use occurrence::scheduled_turn_finish; + const SCHEDULE_POLL_INTERVAL: Duration = Duration::from_secs(10); const SCHEDULE_LEASE_DURATION: Duration = Duration::from_secs(10 * 60); const SCHEDULE_LEASE_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5 * 60); @@ -173,281 +189,6 @@ impl ThreadScheduleRuntime { } } - async fn execute_claim( - &self, - state_db: StateDbHandle, - claim: codex_state::ThreadScheduleClaim, - ) { - let thread_id = claim.schedule.thread_id; - self.emit_schedule_run_updated(thread_id, claim.run.clone()) - .await; - - let result = match self - .resolve_claim_prompt(&state_db, thread_id, &claim.schedule) - .await - { - Ok(prompt) => { - let scheduled_goal_objective = - scheduled_goal_objective(&prompt).map(str::to_string); - self.submit_claimed_schedule( - thread_id, - state_db.clone(), - &claim, - prompt, - scheduled_goal_objective, - ) - .await - } - Err(error) => Err(ScheduleSubmitError { - error, - goal_id: None, - }), - }; - if let Err(ScheduleSubmitError { error, goal_id }) = result { - warn!( - schedule_id = %claim.schedule.schedule_id, - thread_id = %thread_id, - "failed to submit scheduled thread run: {error}" - ); - if let Some(wait) = error.downcast_ref::() { - self.defer_claimed_run_for_usage_profile_wait(state_db, claim, wait.clone()) - .await; - return; - } - if let Some(deferral) = error.downcast_ref::() { - self.defer_claimed_run(state_db, claim, deferral.clone()) - .await; - return; - } - self.fail_claimed_run_after_submit_error( - state_db, - claim, - goal_id, - schedule_submit_error(&error), - ) - .await; - } - } - - async fn submit_claimed_schedule( - &self, - thread_id: ThreadId, - state_db: StateDbHandle, - claim: &codex_state::ThreadScheduleClaim, - prompt: String, - scheduled_goal_objective: Option, - ) -> Result<(), ScheduleSubmitError> { - let claim_auth_profile = self - .claim_auth_profile(&state_db, thread_id, &claim.schedule) - .await; - let broker_decision = super::usage_profile_broker::resolve_dispatch_auth_profile( - &self.auth_manager, - &self.config, - claim_auth_profile.clone(), - ) - .await; - let claim_auth_profile = match schedule_auth_profile_after_broker_decision( - claim_auth_profile, - broker_decision, - self.config.usage_self_heal.reset_retry_buffer_secs, - Utc::now(), - ) { - Ok(resolved) => resolved, - Err(wait) => { - return Err(ScheduleSubmitError { - error: anyhow::Error::new(wait), - goal_id: None, - }); - } - }; - let thread = self - .load_or_resume_thread(thread_id, claim_auth_profile.clone()) - .await - .map_err(|error| ScheduleSubmitError { - error, - goal_id: None, - })?; - self.ensure_schedule_listener(thread_id, thread.clone()) - .await - .map_err(|error| ScheduleSubmitError { - error, - goal_id: None, - })?; - let thread_state = self.thread_state_manager.thread_state(thread_id).await; - let listener_command_tx = { - let thread_state = thread_state.lock().await; - thread_state.listener_command_tx() - }; - let (turn_prompt, scheduled_goal_id) = - if let Some(objective) = scheduled_goal_objective.as_deref() { - let scheduled_goal_id = self - .prepare_scheduled_goal( - thread_id, - &state_db, - objective, - listener_command_tx.clone(), - ) - .await - .map_err(|error| { - let goal_id = error - .downcast_ref::() - .map(|held| held.goal_id.clone()); - ScheduleSubmitError { error, goal_id } - })?; - ( - scheduled_goal_thread_prompt( - objective, - claim.run.run_id.as_str(), - claim.run.scheduled_for, - &claim.schedule, - ), - Some(scheduled_goal_id), - ) - } else { - ( - scheduled_thread_prompt( - &prompt, - &claim.schedule, - claim.run.run_id.as_str(), - claim.run.scheduled_for, - ), - None, - ) - }; - let thread_settings = scheduled_thread_settings_from_snapshot( - thread.config_snapshot().await, - claim_auth_profile, - ); - let turn_id = Uuid::now_v7().to_string(); - - let run_start = state_db - .thread_schedules() - .mark_thread_schedule_run_started(codex_state::ThreadScheduleRunStartParams { - schedule_id: claim.schedule.schedule_id.as_str(), - run_id: claim.run.run_id.as_str(), - lease_id: claim.run.lease_id.as_str(), - turn_id: turn_id.as_str(), - goal_id: scheduled_goal_id.as_deref(), - now: Utc::now(), - lease_duration: SCHEDULE_LEASE_DURATION, - }) - .await; - let run = match run_start { - Ok(Some(run)) => run, - Ok(None) => { - return Err(ScheduleSubmitError { - error: anyhow::anyhow!( - "claimed schedule run {} no longer owns the current unexpired lease", - claim.run.run_id - ), - goal_id: scheduled_goal_id, - }); - } - Err(error) => { - return Err(ScheduleSubmitError { - error, - goal_id: scheduled_goal_id, - }); - } - }; - let ownership_lost = match self.start_lease_heartbeat(state_db.clone(), &run).await { - Ok(Some(ownership_lost)) => ownership_lost, - Ok(None) => { - return Err(ScheduleSubmitError { - error: anyhow::anyhow!( - "claimed schedule run {} lost lease ownership before dispatch readiness", - claim.run.run_id - ), - goal_id: scheduled_goal_id, - }); - } - Err(error) => { - return Err(ScheduleSubmitError { - error, - goal_id: scheduled_goal_id, - }); - } - }; - { - let mut thread_state = thread_state.lock().await; - thread_state.track_scheduled_run( - turn_id.clone(), - crate::thread_state::ScheduledThreadScheduleRun { - schedule_id: run.schedule_id.clone(), - run_id: run.run_id.clone(), - lease_id: run.lease_id.clone(), - goal_id: run.goal_id.clone(), - state_db: state_db.clone(), - }, - ); - } - let start_result = match submit_scheduled_turn_if_owned( - &state_db, - codex_state::ThreadScheduleRunLeaseParams { - schedule_id: run.schedule_id.as_str(), - run_id: run.run_id.as_str(), - lease_id: run.lease_id.as_str(), - now: Utc::now(), - lease_duration: SCHEDULE_LEASE_DURATION, - }, - &ownership_lost, - thread.try_start_user_input_turn_if_idle( - turn_id.clone(), - vec![CoreInputItem::Text { - text: turn_prompt, - text_elements: Vec::new(), - }], - Default::default(), - thread_settings, - ), - ) - .await - { - Ok(Some(start_result)) => start_result, - Ok(None) => { - thread_state - .lock() - .await - .take_scheduled_run(turn_id.as_str()); - return Err(ScheduleSubmitError { - error: anyhow::anyhow!( - "claimed schedule run {} lost lease ownership before turn submission", - claim.run.run_id - ), - goal_id: scheduled_goal_id, - }); - } - Err(error) => { - thread_state - .lock() - .await - .take_scheduled_run(turn_id.as_str()); - return Err(ScheduleSubmitError { - error, - goal_id: scheduled_goal_id, - }); - } - }; - if let Err(err) = start_result { - thread_state - .lock() - .await - .take_scheduled_run(turn_id.as_str()); - if let Some(deferral) = schedule_deferral_for_idle_rejection(&err, Utc::now()) { - return Err(ScheduleSubmitError { - error: anyhow::Error::new(deferral), - goal_id: scheduled_goal_id, - }); - } - return Err(ScheduleSubmitError { - error: anyhow::anyhow!("failed to start scheduled prompt: {err}"), - goal_id: scheduled_goal_id, - }); - } - self.emit_schedule_run_updated(thread_id, run).await; - Ok(()) - } - async fn prepare_scheduled_goal( &self, thread_id: ThreadId, @@ -727,6 +468,62 @@ impl ThreadScheduleRuntime { goal_id: Option, error: String, ) { + let completed_at = Utc::now(); + match finish_scheduled_run_state( + &state_db, + &claim.schedule.schedule_id, + &claim.run.run_id, + &claim.run.lease_id, + goal_id.as_deref(), + Some(error.clone()), + completed_at, + ) + .await + { + Ok(Some((schedule, run))) => { + self.emit_schedule_updated(claim.schedule.thread_id, schedule) + .await; + self.emit_schedule_run_updated(claim.schedule.thread_id, run) + .await; + return; + } + Ok(None) => {} + Err(err) => { + warn!( + schedule_id = %claim.schedule.schedule_id, + "failed to finish started scheduled occurrence after submit error: {err}" + ); + return; + } + } + if matches!( + claim.occurrence_state, + codex_state::ThreadScheduleOccurrenceState::WaitingIdle + | codex_state::ThreadScheduleOccurrenceState::Enqueued + ) { + match state_db + .thread_schedules() + .fail_thread_schedule_occurrence_before_start( + claim.schedule.schedule_id.as_str(), + claim.run.run_id.as_str(), + claim.run.lease_id.as_str(), + completed_at, + goal_id.as_deref(), + error.clone(), + ) + .await + { + Ok(true) => {} + Ok(false) => return, + Err(err) => { + warn!( + schedule_id = %claim.schedule.schedule_id, + "failed to persist pre-start scheduled occurrence failure: {err}" + ); + return; + } + } + } match finish_scheduled_run_state( &state_db, &claim.schedule.schedule_id, @@ -734,7 +531,7 @@ impl ThreadScheduleRuntime { &claim.run.lease_id, goal_id.as_deref(), Some(error), - Utc::now(), + completed_at, ) .await { @@ -820,6 +617,7 @@ impl ThreadScheduleRuntime { state_db, claim, ScheduleRunDeferral { + kind: ScheduleRunDeferralKind::Capacity, retry_at: wait.retry_at, error: wait.to_string(), }, @@ -833,6 +631,21 @@ impl ThreadScheduleRuntime { claim: codex_state::ThreadScheduleClaim, deferral: ScheduleRunDeferral, ) { + if deferral.kind == ScheduleRunDeferralKind::IdleAdmission { + match wait_scheduled_run_for_idle_state(&state_db, &claim, &deferral, Utc::now()).await + { + Ok(Some(schedule)) => { + self.emit_schedule_updated(claim.schedule.thread_id, schedule) + .await; + } + Ok(None) => {} + Err(err) => warn!( + schedule_id = %claim.schedule.schedule_id, + "failed to preserve scheduled occurrence while waiting for idle: {err}" + ), + } + return; + } match defer_scheduled_run_state(&state_db, &claim, &deferral, Utc::now()).await { Ok(Some((schedule, run))) => { self.emit_schedule_updated(claim.schedule.thread_id, schedule) @@ -879,6 +692,7 @@ impl ThreadScheduleRuntime { } } +#[cfg(test)] async fn submit_scheduled_turn_if_owned( state_db: &StateDbHandle, lease: codex_state::ThreadScheduleRunLeaseParams<'_>, @@ -960,14 +774,17 @@ fn scheduled_goal_objective(prompt: &str) -> Option<&str> { } fn scheduled_goal_is_held(goal: &codex_state::ThreadGoal, objective: &str) -> bool { - goal.objective == objective - && matches!( - goal.status, - codex_state::ThreadGoalStatus::Paused - | codex_state::ThreadGoalStatus::Blocked - | codex_state::ThreadGoalStatus::UsageLimited - | codex_state::ThreadGoalStatus::BudgetLimited - ) + goal.objective == objective && scheduled_goal_status_is_held(goal.status) +} + +fn scheduled_goal_status_is_held(status: codex_state::ThreadGoalStatus) -> bool { + matches!( + status, + codex_state::ThreadGoalStatus::Paused + | codex_state::ThreadGoalStatus::Blocked + | codex_state::ThreadGoalStatus::UsageLimited + | codex_state::ThreadGoalStatus::BudgetLimited + ) } fn scheduled_goal_thread_prompt( @@ -1011,195 +828,6 @@ fn scheduled_loop_nesting_guidance(schedule: &codex_state::ThreadSchedule) -> St } } -#[cfg_attr(test, derive(Debug, PartialEq, Eq))] -enum ScheduledTurnFinish { - Complete, - Failed(String), -} - -fn scheduled_turn_finish(event: &EventMsg) -> Option { - match event { - EventMsg::TurnComplete(completed) - if completed - .last_agent_message - .as_deref() - .is_some_and(|message| !message.trim().is_empty()) => - { - Some(ScheduledTurnFinish::Complete) - } - EventMsg::TurnComplete(_) => Some(ScheduledTurnFinish::Failed(schedule_run_error( - "scheduled turn completed without a final assistant message", - ))), - EventMsg::TurnAborted(aborted) => Some(ScheduledTurnFinish::Failed(schedule_run_error( - format!("scheduled turn aborted: {:?}", aborted.reason), - ))), - EventMsg::Error(error) => Some(ScheduledTurnFinish::Failed(schedule_turn_event_error( - error, - ))), - _ => None, - } -} - -pub(super) fn default_thread_schedule_expires_at(now: DateTime) -> Option> { - now.checked_add_signed(ChronoDuration::days(DEFAULT_SCHEDULE_EXPIRATION_DAYS)) -} - -pub(super) fn next_thread_schedule_run_at( - schedule: &codex_state::ThreadScheduleSpec, - timezone: &str, - after: DateTime, -) -> anyhow::Result>> { - let next = match schedule { - codex_state::ThreadScheduleSpec::Once => None, - codex_state::ThreadScheduleSpec::Dynamic => { - after.checked_add_signed(ChronoDuration::minutes(DEFAULT_DYNAMIC_INTERVAL_MINUTES)) - } - codex_state::ThreadScheduleSpec::Interval(interval) => { - let amount = interval.amount; - let duration = match interval.unit { - codex_state::ThreadScheduleIntervalUnit::Minutes => ChronoDuration::minutes(amount), - codex_state::ThreadScheduleIntervalUnit::Hours => ChronoDuration::hours(amount), - codex_state::ThreadScheduleIntervalUnit::Days => ChronoDuration::days(amount), - }; - after.checked_add_signed(duration) - } - codex_state::ThreadScheduleSpec::Cron { expression } => { - let timezone = parse_schedule_timezone(timezone)?; - let cron = Cron::from_str(expression) - .map_err(|err| anyhow::anyhow!("invalid cron expression `{expression}`: {err}"))?; - let local_after = after.with_timezone(&timezone); - let next = cron.find_next_occurrence(&local_after, /*inclusive*/ false)?; - Some(next.with_timezone(&Utc)) - } - }; - Ok(next) -} - -fn next_thread_schedule_run_after_completion( - schedule: &codex_state::ThreadScheduleSpec, - timezone: &str, - scheduled_for: Option>, - completed_at: DateTime, -) -> anyhow::Result>> { - let interval_duration = match schedule { - codex_state::ThreadScheduleSpec::Dynamic => { - Some(ChronoDuration::minutes(DEFAULT_DYNAMIC_INTERVAL_MINUTES)) - } - codex_state::ThreadScheduleSpec::Interval(interval) => { - let amount = interval.amount; - Some(match interval.unit { - codex_state::ThreadScheduleIntervalUnit::Minutes => ChronoDuration::minutes(amount), - codex_state::ThreadScheduleIntervalUnit::Hours => ChronoDuration::hours(amount), - codex_state::ThreadScheduleIntervalUnit::Days => ChronoDuration::days(amount), - }) - } - codex_state::ThreadScheduleSpec::Once | codex_state::ThreadScheduleSpec::Cron { .. } => { - None - } - }; - - if let (Some(interval_duration), Some(scheduled_for)) = (interval_duration, scheduled_for) { - let Some(next_run_at) = scheduled_for.checked_add_signed(interval_duration) else { - return Ok(None); - }; - if next_run_at > completed_at { - return Ok(Some(next_run_at)); - } - let duration_ms = interval_duration.num_milliseconds(); - if duration_ms <= 0 { - return Ok(None); - } - let elapsed_ms = completed_at - .signed_duration_since(scheduled_for) - .num_milliseconds(); - let periods_elapsed = elapsed_ms.div_euclid(duration_ms).saturating_add(1); - return Ok(duration_ms - .checked_mul(periods_elapsed) - .and_then(|advance_ms| { - scheduled_for.checked_add_signed(ChronoDuration::milliseconds(advance_ms)) - })); - } - - next_thread_schedule_run_at(schedule, timezone, completed_at) -} - -pub(super) fn normalize_schedule_timezone(timezone: &str) -> anyhow::Result { - parse_schedule_timezone(timezone).map(|timezone| timezone.name().to_string()) -} - -pub(super) async fn finish_scheduled_run_after_turn( - thread_id: ThreadId, - scheduled_run: crate::thread_state::ScheduledThreadScheduleRun, - event: &EventMsg, - turn_error: Option, - outgoing: &Arc, -) { - let completed_at = Utc::now(); - let error = match (scheduled_turn_finish(event), turn_error) { - (Some(_), Some(error)) => Some(schedule_turn_error(&error)), - (Some(ScheduledTurnFinish::Complete), None) => None, - (Some(ScheduledTurnFinish::Failed(error)), None) => Some(error), - (None, _) => return, - }; - match finish_scheduled_run_state( - &scheduled_run.state_db, - scheduled_run.schedule_id.as_str(), - scheduled_run.run_id.as_str(), - scheduled_run.lease_id.as_str(), - scheduled_run.goal_id.as_deref(), - error, - completed_at, - ) - .await - { - Ok(Some((schedule, run))) => { - outgoing - .send_server_notification(ServerNotification::ThreadScheduleUpdated( - ThreadScheduleUpdatedNotification { - thread_id: thread_id.to_string(), - schedule: api_thread_schedule_from_state(schedule), - }, - )) - .await; - outgoing - .send_server_notification(ServerNotification::ThreadScheduleRunUpdated( - ThreadScheduleRunUpdatedNotification { - thread_id: thread_id.to_string(), - run: api_thread_schedule_run_from_state(run), - }, - )) - .await; - } - Ok(None) => {} - Err(err) => warn!( - schedule_id = %scheduled_run.schedule_id, - thread_id = %thread_id, - "failed to finish scheduled thread run: {err}" - ), - } -} - -pub(super) async fn recover_scheduled_run_for_terminal_turn( - state_db: &StateDbHandle, - thread_id: ThreadId, - turn_id: &str, -) -> anyhow::Result> { - let Some(run) = state_db - .thread_schedules() - .get_running_thread_schedule_run_for_turn(thread_id, turn_id) - .await? - else { - return Ok(None); - }; - Ok(Some(crate::thread_state::ScheduledThreadScheduleRun { - schedule_id: run.schedule_id, - run_id: run.run_id, - lease_id: run.lease_id, - goal_id: run.goal_id, - state_db: state_db.clone(), - })) -} - /// Maximum consecutive failed runs before a recurring schedule stops /// re-arming itself (circuit breaker). The streak resets after a success. const MAX_CONSECUTIVE_SCHEDULE_FAILURES: i64 = 10; @@ -1223,11 +851,6 @@ struct ScheduleUsageProfileWait { retry_at: DateTime, } -struct ScheduleSubmitError { - error: anyhow::Error, - goal_id: Option, -} - impl std::fmt::Display for ScheduleUsageProfileWait { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( @@ -1277,10 +900,17 @@ impl std::error::Error for ScheduledGoalHeld {} #[derive(Debug, Clone, PartialEq, Eq)] struct ScheduleRunDeferral { + kind: ScheduleRunDeferralKind, retry_at: DateTime, error: String, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ScheduleRunDeferralKind { + IdleAdmission, + Capacity, +} + impl std::fmt::Display for ScheduleRunDeferral { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( @@ -1309,6 +939,7 @@ fn schedule_deferral_for_idle_rejection( codex_core::TryStartTurnIfIdleRejectionReason::PlanMode => return None, }; Some(ScheduleRunDeferral { + kind: ScheduleRunDeferralKind::IdleAdmission, retry_at: now + ChronoDuration::seconds(SCHEDULE_IDLE_RETRY_DELAY_SECONDS), error, }) @@ -1497,6 +1128,7 @@ async fn defer_scheduled_run_for_usage_profile_wait_state( state_db, claim, &ScheduleRunDeferral { + kind: ScheduleRunDeferralKind::Capacity, retry_at: wait.retry_at, error: wait.to_string(), }, @@ -1505,6 +1137,31 @@ async fn defer_scheduled_run_for_usage_profile_wait_state( .await } +async fn wait_scheduled_run_for_idle_state( + state_db: &StateDbHandle, + claim: &codex_state::ThreadScheduleClaim, + deferral: &ScheduleRunDeferral, + now: DateTime, +) -> anyhow::Result> { + let updated = state_db + .thread_schedules() + .wait_thread_schedule_run_for_idle( + claim.schedule.schedule_id.as_str(), + claim.run.run_id.as_str(), + claim.run.lease_id.as_str(), + deferral.retry_at, + now, + ) + .await?; + if !updated { + return Ok(None); + } + state_db + .thread_schedules() + .get_thread_schedule(&claim.schedule.schedule_id) + .await +} + async fn defer_scheduled_run_state( state_db: &StateDbHandle, claim: &codex_state::ThreadScheduleClaim, @@ -2053,8 +1710,11 @@ mod tests { use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; + use codex_protocol::protocol::TurnAbortReason; + use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; use codex_protocol::protocol::TurnContextItem; + use codex_protocol::protocol::TurnStartedEvent; use codex_state::ThreadMetadataBuilder; use pretty_assertions::assert_eq; use std::sync::atomic::AtomicBool; @@ -2064,6 +1724,49 @@ mod tests { DateTime::::from_timestamp(seconds, 0).expect("valid timestamp") } + async fn enqueue_and_start_claim( + state_db: &codex_state::StateRuntime, + claim: &codex_state::ThreadScheduleClaim, + goal_id: Option<&str>, + now: DateTime, + lease_duration: Duration, + ) -> codex_state::ThreadScheduleRun { + let turn_id = claim + .run + .turn_id + .as_deref() + .expect("claimed occurrence should have a stable turn id"); + state_db + .thread_schedules() + .enqueue_thread_schedule_run(codex_state::ThreadScheduleRunEnqueueParams { + schedule_id: claim.schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + goal_id, + auth_profile_recorded: false, + auth_profile: None, + turn_input: "app-server schedule test input", + now, + }) + .await + .expect("claimed occurrence should enqueue") + .expect("claimed occurrence should retain its lease"); + state_db + .thread_schedules() + .mark_thread_schedule_run_started(codex_state::ThreadScheduleRunStartParams { + schedule_id: claim.schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + turn_id, + goal_id, + now, + lease_duration, + }) + .await + .expect("enqueued occurrence should start") + .expect("enqueued occurrence should retain its lease") + } + fn prompt_test_schedule(nesting_depth: i64) -> codex_state::ThreadSchedule { codex_state::ThreadSchedule { thread_id: ThreadId::new(), @@ -2162,6 +1865,27 @@ mod tests { }) } + fn resumed_history_with_turn_events( + thread_id: ThreadId, + events: impl IntoIterator, + ) -> InitialHistory { + InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history: events.into_iter().map(RolloutItem::EventMsg).collect(), + rollout_path: None, + }) + } + + fn turn_started(turn_id: &str, started_at: i64) -> EventMsg { + EventMsg::TurnStarted(TurnStartedEvent { + turn_id: turn_id.to_string(), + trace_id: None, + started_at: Some(started_at), + model_context_window: None, + collaboration_mode_kind: Default::default(), + }) + } + #[test] fn schedule_failure_backoff_grows_and_caps() { let completed = at(/*seconds*/ 1_000_000); @@ -2709,6 +2433,7 @@ mod tests { assert_eq!( Some(ScheduleRunDeferral { + kind: ScheduleRunDeferralKind::IdleAdmission, retry_at: at(/*seconds*/ 1_700_000_030), error: "scheduled thread is busy".to_string(), }), @@ -2721,6 +2446,7 @@ mod tests { ); assert_eq!( Some(ScheduleRunDeferral { + kind: ScheduleRunDeferralKind::IdleAdmission, retry_at: at(/*seconds*/ 1_700_000_030), error: "scheduled thread has pending mailbox trigger-turn work".to_string(), }), @@ -2742,88 +2468,8 @@ mod tests { ); } - #[tokio::test] - async fn idle_rejection_deferral_rearms_without_incrementing_failure_count() { - let temp_dir = tempfile::tempdir().expect("temp dir should be created"); - let state_db = codex_state::StateRuntime::init( - temp_dir.path().to_path_buf(), - "fallback-provider".to_string(), - ) - .await - .expect("state db should initialize"); - let thread_id = ThreadId::new(); - let mut builder = ThreadMetadataBuilder::new( - thread_id, - temp_dir.path().join("thread.jsonl"), - at(/*seconds*/ 1_700_000_000), - SessionSource::Cli, - ); - builder.cwd = temp_dir.path().join("workspace"); - state_db - .upsert_thread(&builder.build("fallback-provider")) - .await - .expect("thread metadata should persist"); - let now = at(/*seconds*/ 1_700_000_000); - let schedule = state_db - .thread_schedules() - .create_thread_schedule(codex_state::ThreadScheduleCreateParams { - thread_id, - prompt: "wait until the thread is idle".to_string(), - prompt_source: codex_state::ThreadSchedulePromptSource::Inline, - schedule: codex_state::ThreadScheduleSpec::Interval( - codex_state::ThreadScheduleInterval { - amount: 5, - unit: codex_state::ThreadScheduleIntervalUnit::Minutes, - }, - ), - timezone: "UTC".to_string(), - status: codex_state::ThreadScheduleStatus::Active, - next_run_at: Some(now), - expires_at: None, - }) - .await - .expect("schedule should create"); - let claim = state_db - .thread_schedules() - .claim_due_thread_schedule(now, "lease-busy", Duration::from_secs(300)) - .await - .expect("claim should succeed") - .expect("schedule should claim"); - let completed_at = now + chrono::Duration::seconds(5); - let deferral = ScheduleRunDeferral { - retry_at: now + chrono::Duration::seconds(SCHEDULE_IDLE_RETRY_DELAY_SECONDS), - error: "scheduled thread is busy".to_string(), - }; - - let (deferred_schedule, deferred_run) = - defer_scheduled_run_state(&state_db, &claim, &deferral, completed_at) - .await - .expect("idle rejection should defer") - .expect("deferred rows should load"); - - assert_eq!( - codex_state::ThreadSchedule { - next_run_at: Some(deferral.retry_at), - last_run_at: Some(completed_at), - failure_count: 0, - lease_id: None, - lease_expires_at: None, - updated_at: deferred_schedule.updated_at, - ..schedule - }, - deferred_schedule - ); - assert_eq!( - codex_state::ThreadScheduleRun { - status: codex_state::ThreadScheduleRunStatus::Deferred, - turn_id: None, - error: Some(deferral.error), - completed_at: Some(completed_at), - ..claim.run - }, - deferred_run - ); - } + #[path = "occurrence_tests.rs"] + mod occurrence_tests; #[tokio::test] async fn usage_profile_wait_defers_claim_without_incrementing_failure_count() { @@ -3255,20 +2901,14 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); - state_db - .thread_schedules() - .mark_thread_schedule_run_started(codex_state::ThreadScheduleRunStartParams { - schedule_id: schedule.schedule_id.as_str(), - run_id: claim.run.run_id.as_str(), - lease_id: claim.run.lease_id.as_str(), - turn_id: "turn-after-restart", - goal_id: Some(goal.goal_id.as_str()), - now: scheduled_for, - lease_duration: Duration::from_secs(300), - }) - .await - .expect("run start should persist") - .expect("run should still exist"); + enqueue_and_start_claim( + &state_db, + &claim, + Some(goal.goal_id.as_str()), + scheduled_for, + Duration::from_secs(300), + ) + .await; drop(state_db); let reopened = codex_state::StateRuntime::init( @@ -3335,156 +2975,6 @@ mod tests { ); } - #[tokio::test] - async fn stale_started_run_cannot_submit_after_reaper_replacement() { - let temp_dir = tempfile::tempdir().expect("temp dir should be created"); - let state_db = codex_state::StateRuntime::init( - temp_dir.path().to_path_buf(), - "fallback-provider".to_string(), - ) - .await - .expect("state db should initialize"); - let thread_id = ThreadId::new(); - let mut builder = ThreadMetadataBuilder::new( - thread_id, - temp_dir.path().join("thread.jsonl"), - at(/*seconds*/ 1_700_000_000), - SessionSource::Cli, - ); - builder.cwd = temp_dir.path().join("workspace"); - state_db - .upsert_thread(&builder.build("fallback-provider")) - .await - .expect("thread metadata should persist"); - let now = at(/*seconds*/ 1_700_000_000); - let schedule = state_db - .thread_schedules() - .create_thread_schedule(codex_state::ThreadScheduleCreateParams { - thread_id, - prompt: "never submit stale work".to_string(), - prompt_source: codex_state::ThreadSchedulePromptSource::Inline, - schedule: codex_state::ThreadScheduleSpec::Interval( - codex_state::ThreadScheduleInterval { - amount: 1, - unit: codex_state::ThreadScheduleIntervalUnit::Minutes, - }, - ), - timezone: "UTC".to_string(), - status: codex_state::ThreadScheduleStatus::Active, - next_run_at: Some(now), - expires_at: None, - }) - .await - .expect("schedule should create"); - let claim = state_db - .thread_schedules() - .claim_due_thread_schedule(now, "lease-suspended", Duration::from_secs(30)) - .await - .expect("claim should succeed") - .expect("schedule should claim"); - state_db - .thread_schedules() - .mark_thread_schedule_run_started(codex_state::ThreadScheduleRunStartParams { - schedule_id: schedule.schedule_id.as_str(), - run_id: claim.run.run_id.as_str(), - lease_id: claim.run.lease_id.as_str(), - turn_id: "turn-suspended", - goal_id: None, - now, - lease_duration: Duration::from_secs(30), - }) - .await - .expect("run start should persist") - .expect("run should still own the lease before suspension"); - - let contender = codex_state::StateRuntime::init( - temp_dir.path().to_path_buf(), - "fallback-provider".to_string(), - ) - .await - .expect("contending state db should initialize"); - let resumed_at = now + chrono::Duration::seconds(31); - let replacement = contender - .thread_schedules() - .claim_due_thread_schedule(resumed_at, "lease-replacement", Duration::from_secs(30)) - .await - .expect("expired lease reaper should not error") - .expect("expired started run should be replaced"); - let submitted = Arc::new(AtomicBool::new(false)); - let submission_observer = Arc::clone(&submitted); - let ownership_lost = CancellationToken::new(); - - let submission = submit_scheduled_turn_if_owned( - &state_db, - codex_state::ThreadScheduleRunLeaseParams { - schedule_id: schedule.schedule_id.as_str(), - run_id: claim.run.run_id.as_str(), - lease_id: claim.run.lease_id.as_str(), - now: resumed_at, - lease_duration: Duration::from_secs(30), - }, - &ownership_lost, - async move { - submission_observer.store(true, Ordering::SeqCst); - }, - ) - .await - .expect("stale dispatch validation should not error"); - - assert_eq!(None, submission); - assert!(!submitted.load(Ordering::SeqCst)); - assert_eq!( - codex_state::ThreadScheduleRunStatus::Failed, - state_db - .thread_schedules() - .get_thread_schedule_run(claim.run.run_id.as_str()) - .await - .expect("old run should load") - .expect("old run should exist") - .status - ); - assert_eq!( - codex_state::ThreadScheduleRunStatus::Leased, - replacement.run.status - ); - - let replacement_started_at = resumed_at + chrono::Duration::seconds(1); - state_db - .thread_schedules() - .mark_thread_schedule_run_started(codex_state::ThreadScheduleRunStartParams { - schedule_id: schedule.schedule_id.as_str(), - run_id: replacement.run.run_id.as_str(), - lease_id: replacement.run.lease_id.as_str(), - turn_id: "turn-replacement", - goal_id: None, - now: replacement_started_at, - lease_duration: Duration::from_secs(30), - }) - .await - .expect("replacement start should persist") - .expect("replacement should start"); - ownership_lost.cancel(); - let cancelled_submission_observer = Arc::clone(&submitted); - let cancelled_submission = submit_scheduled_turn_if_owned( - &state_db, - codex_state::ThreadScheduleRunLeaseParams { - schedule_id: schedule.schedule_id.as_str(), - run_id: replacement.run.run_id.as_str(), - lease_id: replacement.run.lease_id.as_str(), - now: replacement_started_at + chrono::Duration::seconds(1), - lease_duration: Duration::from_secs(30), - }, - &ownership_lost, - async move { - cancelled_submission_observer.store(true, Ordering::SeqCst); - }, - ) - .await - .expect("cancelled dispatch validation should not error"); - assert_eq!(None, cancelled_submission); - assert!(!submitted.load(Ordering::SeqCst)); - } - #[tokio::test] async fn resolved_default_goal_prompt_submit_failure_pauses_schedule_when_goal_remains_held() { let temp_dir = tempfile::tempdir().expect("temp dir should be created"); @@ -4005,37 +3495,6 @@ mod tests { assert!(!once_prompt.contains("parent_schedule_id set to schedule-123")); } - #[test] - fn scheduled_turn_without_agent_message_fails() { - let finish = scheduled_turn_finish(&EventMsg::TurnComplete(TurnCompleteEvent { - turn_id: "turn-1".to_string(), - last_agent_message: None, - completed_at: None, - duration_ms: None, - time_to_first_token_ms: None, - })); - - assert_eq!( - Some(ScheduledTurnFinish::Failed( - "scheduled turn completed without a final assistant message".to_string() - )), - finish - ); - } - - #[test] - fn scheduled_turn_with_agent_message_completes() { - let finish = scheduled_turn_finish(&EventMsg::TurnComplete(TurnCompleteEvent { - turn_id: "turn-1".to_string(), - last_agent_message: Some("done".to_string()), - completed_at: None, - duration_ms: None, - time_to_first_token_ms: None, - })); - - assert_eq!(Some(ScheduledTurnFinish::Complete), finish); - } - #[test] fn scheduled_turn_error_fails() { let finish = scheduled_turn_finish(&EventMsg::Error(ErrorEvent { diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence.rs new file mode 100644 index 000000000..3e862505e --- /dev/null +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence.rs @@ -0,0 +1,32 @@ +//! App-server execution and terminal replay for one scheduled occurrence. + +use super::*; + +mod execution; +mod terminal; + +pub(super) use terminal::PersistedScheduledTurnTerminal; +pub(super) use terminal::ScheduledTurnFinish; +pub(super) use terminal::default_thread_schedule_expires_at; +pub(super) use terminal::finish_scheduled_run_after_turn; +pub(super) use terminal::next_thread_schedule_run_after_completion; +pub(super) use terminal::next_thread_schedule_run_at; +pub(super) use terminal::normalize_schedule_timezone; +pub(super) use terminal::persisted_scheduled_turn_terminal; +pub(super) use terminal::recover_scheduled_run_for_terminal_turn; +pub(super) use terminal::scheduled_turn_finish; + +impl ThreadScheduleRuntime { + pub(super) async fn execute_claim( + &self, + state_db: StateDbHandle, + claim: codex_state::ThreadScheduleClaim, + ) { + self.execute_occurrence_claim(state_db, claim).await; + } +} + +struct ScheduleSubmitError { + error: anyhow::Error, + goal_id: Option, +} diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/execution.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/execution.rs new file mode 100644 index 000000000..5a816208d --- /dev/null +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/execution.rs @@ -0,0 +1,511 @@ +//! Admission, stable-turn submission, and crash recovery execution. + +use super::*; + +impl ThreadScheduleRuntime { + pub(super) async fn execute_occurrence_claim( + &self, + state_db: StateDbHandle, + claim: codex_state::ThreadScheduleClaim, + ) { + let thread_id = claim.schedule.thread_id; + if claim.occurrence_state == codex_state::ThreadScheduleOccurrenceState::Terminal { + self.replay_terminal_claim(state_db, claim).await; + return; + } + if claim.occurrence_state == codex_state::ThreadScheduleOccurrenceState::Started { + let turn_id = match claim.run.turn_id.as_deref() { + Some(turn_id) => turn_id, + None => { + let goal_id = claim.run.goal_id.clone(); + self.fail_claimed_run_after_submit_error( + state_db, + claim, + goal_id, + "accepted scheduled turn cannot be resumed because its stable turn identifier is unavailable" + .to_string(), + ) + .await; + return; + } + }; + match self + .persisted_terminal_scheduled_turn(&state_db, thread_id, turn_id) + .await + { + Ok(Some(terminal)) => { + self.finish_claim_from_persisted_terminal(state_db, claim, terminal) + .await; + return; + } + Ok(None) => match self + .held_started_scheduled_goal(&state_db, thread_id, claim.run.goal_id.as_deref()) + .await + { + Ok(Some(held)) => { + let goal_id = held.goal_id.clone(); + self.fail_claimed_run_after_submit_error( + state_db, + claim, + Some(goal_id), + held.to_string(), + ) + .await; + return; + } + Ok(None) => {} + Err(error) => { + let goal_id = claim.run.goal_id.clone(); + self.fail_claimed_run_after_submit_error( + state_db, + claim, + goal_id, + format!( + "accepted scheduled turn cannot be safely resumed because its goal state could not be inspected: {error}" + ), + ) + .await; + return; + } + }, + Err(error) => { + let goal_id = claim.run.goal_id.clone(); + self.fail_claimed_run_after_submit_error( + state_db, + claim, + goal_id, + format!( + "accepted scheduled turn cannot be safely resumed because its durable rollout could not be inspected: {error}" + ), + ) + .await; + return; + } + } + } + if claim.occurrence_state == codex_state::ThreadScheduleOccurrenceState::Started + && claim.turn_input.is_none() + { + let goal_id = claim.run.goal_id.clone(); + self.fail_claimed_run_after_submit_error( + state_db, + claim, + goal_id, + "accepted scheduled turn cannot be resumed because its persisted input is unavailable" + .to_string(), + ) + .await; + return; + } + + let resolved_prompt = if claim.turn_input.is_some() { + Ok(String::new()) + } else { + self.resolve_claim_prompt(&state_db, thread_id, &claim.schedule) + .await + }; + let result = match resolved_prompt { + Ok(prompt) => { + let scheduled_goal_objective = claim + .turn_input + .is_none() + .then(|| scheduled_goal_objective(&prompt).map(str::to_string)) + .flatten(); + self.submit_claimed_schedule( + thread_id, + state_db.clone(), + &claim, + prompt, + scheduled_goal_objective, + ) + .await + } + Err(error) => Err(ScheduleSubmitError { + error, + goal_id: None, + }), + }; + if let Err(ScheduleSubmitError { error, goal_id }) = result { + warn!( + schedule_id = %claim.schedule.schedule_id, + thread_id = %thread_id, + "failed to submit scheduled thread run: {error}" + ); + if let Some(wait) = error.downcast_ref::() { + self.defer_claimed_run_for_usage_profile_wait(state_db, claim, wait.clone()) + .await; + return; + } + if let Some(deferral) = error.downcast_ref::() { + self.defer_claimed_run(state_db, claim, deferral.clone()) + .await; + return; + } + self.fail_claimed_run_after_submit_error( + state_db, + claim, + goal_id, + schedule_submit_error(&error), + ) + .await; + } + } + + async fn submit_claimed_schedule( + &self, + thread_id: ThreadId, + state_db: StateDbHandle, + claim: &codex_state::ThreadScheduleClaim, + prompt: String, + scheduled_goal_objective: Option, + ) -> Result<(), ScheduleSubmitError> { + let claim_auth_profile = match claim.occurrence_auth_profile.clone() { + Some(auth_profile) => Some(auth_profile), + None => { + self.claim_auth_profile(&state_db, thread_id, &claim.schedule) + .await + } + }; + let claim_auth_profile = if matches!( + claim.occurrence_state, + codex_state::ThreadScheduleOccurrenceState::Enqueued + | codex_state::ThreadScheduleOccurrenceState::Started + ) { + claim_auth_profile + } else { + let broker_decision = super::usage_profile_broker::resolve_dispatch_auth_profile( + &self.auth_manager, + &self.config, + claim_auth_profile.clone(), + ) + .await; + match schedule_auth_profile_after_broker_decision( + claim_auth_profile, + broker_decision, + self.config.usage_self_heal.reset_retry_buffer_secs, + Utc::now(), + ) { + Ok(resolved) => resolved, + Err(wait) => { + return Err(ScheduleSubmitError { + error: anyhow::Error::new(wait), + goal_id: None, + }); + } + } + }; + let thread = self + .load_or_resume_thread(thread_id, claim_auth_profile.clone()) + .await + .map_err(|error| ScheduleSubmitError { + error, + goal_id: None, + })?; + self.ensure_schedule_listener(thread_id, thread.clone()) + .await + .map_err(|error| ScheduleSubmitError { + error, + goal_id: None, + })?; + let thread_state = self.thread_state_manager.thread_state(thread_id).await; + let listener_command_tx = { + let thread_state = thread_state.lock().await; + thread_state.listener_command_tx() + }; + let (turn_prompt, scheduled_goal_id) = if let Some(turn_input) = claim.turn_input.clone() { + (turn_input, claim.run.goal_id.clone()) + } else if let Some(objective) = scheduled_goal_objective.as_deref() { + let scheduled_goal_id = self + .prepare_scheduled_goal( + thread_id, + &state_db, + objective, + listener_command_tx.clone(), + ) + .await + .map_err(|error| { + let goal_id = error + .downcast_ref::() + .map(|held| held.goal_id.clone()); + ScheduleSubmitError { error, goal_id } + })?; + ( + scheduled_goal_thread_prompt( + objective, + claim.run.run_id.as_str(), + claim.run.scheduled_for, + &claim.schedule, + ), + Some(scheduled_goal_id), + ) + } else { + ( + scheduled_thread_prompt( + &prompt, + &claim.schedule, + claim.run.run_id.as_str(), + claim.run.scheduled_for, + ), + None, + ) + }; + let thread_settings = scheduled_thread_settings_from_snapshot( + thread.config_snapshot().await, + claim_auth_profile.clone(), + ); + let turn_id = claim + .run + .turn_id + .clone() + .ok_or_else(|| ScheduleSubmitError { + error: anyhow::anyhow!( + "claimed schedule occurrence {} has no stable turn identifier", + claim.run.run_id + ), + goal_id: scheduled_goal_id.clone(), + })?; + let run = if claim.occurrence_state == codex_state::ThreadScheduleOccurrenceState::Started { + claim.run.clone() + } else { + state_db + .thread_schedules() + .enqueue_thread_schedule_run(codex_state::ThreadScheduleRunEnqueueParams { + schedule_id: claim.schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + goal_id: scheduled_goal_id.as_deref(), + auth_profile_recorded: claim_auth_profile.is_some(), + auth_profile: claim_auth_profile.as_ref().and_then(Option::as_deref), + turn_input: turn_prompt.as_str(), + now: Utc::now(), + }) + .await + .map_err(|error| ScheduleSubmitError { + error, + goal_id: scheduled_goal_id.clone(), + })? + .ok_or_else(|| ScheduleSubmitError { + error: anyhow::anyhow!( + "claimed schedule occurrence {} lost ownership before enqueue", + claim.run.run_id + ), + goal_id: scheduled_goal_id.clone(), + })? + }; + let ownership_lost = match self.start_lease_heartbeat(state_db.clone(), &run).await { + Ok(Some(ownership_lost)) => ownership_lost, + Ok(None) => { + return Err(ScheduleSubmitError { + error: anyhow::anyhow!( + "claimed schedule run {} lost lease ownership before dispatch readiness", + claim.run.run_id + ), + goal_id: scheduled_goal_id, + }); + } + Err(error) => { + return Err(ScheduleSubmitError { + error, + goal_id: scheduled_goal_id, + }); + } + }; + { + let mut thread_state = thread_state.lock().await; + thread_state.track_scheduled_run( + turn_id.clone(), + crate::thread_state::ScheduledThreadScheduleRun { + schedule_id: run.schedule_id.clone(), + run_id: run.run_id.clone(), + lease_id: run.lease_id.clone(), + goal_id: run.goal_id.clone(), + state_db: state_db.clone(), + }, + ); + } + let lease_is_owned = state_db + .thread_schedules() + .extend_thread_schedule_lease(codex_state::ThreadScheduleRunLeaseParams { + schedule_id: run.schedule_id.as_str(), + run_id: run.run_id.as_str(), + lease_id: run.lease_id.as_str(), + now: Utc::now(), + lease_duration: SCHEDULE_LEASE_DURATION, + }) + .await + .map_err(|error| ScheduleSubmitError { + error, + goal_id: scheduled_goal_id.clone(), + })?; + if !lease_is_owned || ownership_lost.is_cancelled() { + thread_state + .lock() + .await + .take_scheduled_run(turn_id.as_str()); + return Err(ScheduleSubmitError { + error: anyhow::anyhow!( + "claimed schedule run {} lost lease ownership before turn submission", + claim.run.run_id + ), + goal_id: scheduled_goal_id, + }); + } + // Once core reserves the idle turn, await it to completion. Dropping + // this future on a concurrent heartbeat cancellation could strand both + // the idle reservation and a newly persisted Started occurrence. + let start_result = thread + .try_start_scheduled_user_input_turn_if_idle( + turn_id.clone(), + vec![CoreInputItem::Text { + text: turn_prompt, + text_elements: Vec::new(), + }], + Default::default(), + thread_settings, + codex_core::ScheduledTurnStart { + schedule_id: run.schedule_id.clone(), + run_id: run.run_id.clone(), + lease_id: run.lease_id.clone(), + goal_id: run.goal_id.clone(), + lease_duration: SCHEDULE_LEASE_DURATION, + }, + ) + .await; + let run = match start_result { + Ok(run) => run, + Err(err) => { + thread_state + .lock() + .await + .take_scheduled_run(turn_id.as_str()); + if let Some(deferral) = schedule_deferral_for_idle_rejection(&err, Utc::now()) { + return Err(ScheduleSubmitError { + error: anyhow::Error::new(deferral), + goal_id: scheduled_goal_id, + }); + } + return Err(ScheduleSubmitError { + error: anyhow::anyhow!("failed to start scheduled prompt: {err}"), + goal_id: scheduled_goal_id, + }); + } + }; + self.emit_schedule_run_updated(thread_id, run).await; + Ok(()) + } + + async fn replay_terminal_claim( + &self, + state_db: StateDbHandle, + claim: codex_state::ThreadScheduleClaim, + ) { + let completed_at = claim.run.completed_at.unwrap_or_else(Utc::now); + let error = (claim.run.status == codex_state::ThreadScheduleRunStatus::Failed).then(|| { + claim + .run + .error + .clone() + .unwrap_or_else(|| "scheduled turn failed without a recorded error".to_string()) + }); + match finish_scheduled_run_state( + &state_db, + claim.schedule.schedule_id.as_str(), + claim.run.run_id.as_str(), + claim.run.lease_id.as_str(), + claim.run.goal_id.as_deref(), + error, + completed_at, + ) + .await + { + Ok(Some((schedule, run))) => { + self.emit_schedule_updated(claim.schedule.thread_id, schedule) + .await; + self.emit_schedule_run_updated(claim.schedule.thread_id, run) + .await; + } + Ok(None) => {} + Err(err) => warn!( + schedule_id = %claim.schedule.schedule_id, + "failed to replay terminal scheduled occurrence finalization: {err}" + ), + } + } + + async fn persisted_terminal_scheduled_turn( + &self, + state_db: &StateDbHandle, + thread_id: ThreadId, + turn_id: &str, + ) -> anyhow::Result> { + let rollout_path = codex_rollout::find_thread_path_by_id_str( + &self.config.codex_home, + &thread_id.to_string(), + Some(state_db), + ) + .await? + .ok_or_else(|| anyhow::anyhow!("thread rollout not found for {thread_id}"))?; + let history = codex_rollout::RolloutRecorder::get_rollout_history(&rollout_path) + .await + .with_context(|| { + format!( + "failed to load rollout {} for scheduled turn recovery", + rollout_path.display() + ) + })?; + Ok(persisted_scheduled_turn_terminal( + &history, + turn_id, + Utc::now(), + )) + } + + async fn held_started_scheduled_goal( + &self, + state_db: &StateDbHandle, + thread_id: ThreadId, + goal_id: Option<&str>, + ) -> anyhow::Result> { + let Some(goal_id) = goal_id else { + return Ok(None); + }; + let goal = state_db.thread_goals().get_thread_goal(thread_id).await?; + Ok(goal + .filter(|goal| goal.goal_id == goal_id && scheduled_goal_status_is_held(goal.status)) + .map(|goal| ScheduledGoalHeld { + goal_id: goal.goal_id, + status: goal.status, + })) + } + + async fn finish_claim_from_persisted_terminal( + &self, + state_db: StateDbHandle, + claim: codex_state::ThreadScheduleClaim, + terminal: PersistedScheduledTurnTerminal, + ) { + match finish_scheduled_run_state( + &state_db, + claim.schedule.schedule_id.as_str(), + claim.run.run_id.as_str(), + claim.run.lease_id.as_str(), + claim.run.goal_id.as_deref(), + terminal.error, + terminal.completed_at, + ) + .await + { + Ok(Some((schedule, run))) => { + self.emit_schedule_updated(claim.schedule.thread_id, schedule) + .await; + self.emit_schedule_run_updated(claim.schedule.thread_id, run) + .await; + } + Ok(None) => {} + Err(err) => warn!( + schedule_id = %claim.schedule.schedule_id, + "failed to finalize scheduled occurrence from durable terminal turn: {err}" + ), + } + } +} diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs new file mode 100644 index 000000000..309aa40e4 --- /dev/null +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs @@ -0,0 +1,269 @@ +//! Durable rollout terminal detection and idempotent schedule finalization. + +use super::*; + +#[cfg_attr(test, derive(Debug, PartialEq, Eq))] +pub(crate) enum ScheduledTurnFinish { + Complete, + Failed(String), +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct PersistedScheduledTurnTerminal { + pub(crate) completed_at: DateTime, + pub(crate) error: Option, +} + +pub(super) fn persisted_scheduled_turn_terminal( + history: &InitialHistory, + turn_id: &str, + fallback_completed_at: DateTime, +) -> Option { + let rollout_items = history.get_rollout_items(); + let turn = build_api_turns_from_rollout_items(&rollout_items) + .into_iter() + .find(|turn| turn.id == turn_id); + let Some(turn) = turn else { + let aborted = rollout_items + .iter() + .filter_map(|item| match item { + RolloutItem::EventMsg(EventMsg::TurnAborted(aborted)) + if aborted.turn_id.as_deref() == Some(turn_id) => + { + Some(aborted) + } + _ => None, + }) + .last()?; + return Some(PersistedScheduledTurnTerminal { + completed_at: aborted + .completed_at + .and_then(|timestamp| DateTime::::from_timestamp(timestamp, 0)) + .unwrap_or(fallback_completed_at), + error: Some(schedule_run_error(format!( + "scheduled turn aborted: {:?}", + aborted.reason + ))), + }); + }; + let completed_at = turn + .completed_at + .and_then(|timestamp| DateTime::::from_timestamp(timestamp, 0)) + .unwrap_or(fallback_completed_at); + let error = match turn.status { + codex_app_server_protocol::TurnStatus::Completed => { + let finish = rollout_items + .iter() + .filter_map(|item| match item { + RolloutItem::EventMsg(event @ EventMsg::TurnComplete(completed)) + if completed.turn_id == turn_id => + { + scheduled_turn_finish(event) + } + _ => None, + }) + .last()?; + match finish { + ScheduledTurnFinish::Complete => None, + ScheduledTurnFinish::Failed(error) => Some(error), + } + } + codex_app_server_protocol::TurnStatus::Failed => Some( + turn.error + .as_ref() + .map(schedule_turn_error) + .unwrap_or_else(|| schedule_run_error("scheduled turn failed")), + ), + codex_app_server_protocol::TurnStatus::Interrupted => { + Some(schedule_run_error("scheduled turn was interrupted")) + } + codex_app_server_protocol::TurnStatus::InProgress => return None, + }; + Some(PersistedScheduledTurnTerminal { + completed_at, + error, + }) +} + +pub(super) fn scheduled_turn_finish(event: &EventMsg) -> Option { + match event { + EventMsg::TurnComplete(completed) + if completed + .last_agent_message + .as_deref() + .is_some_and(|message| !message.trim().is_empty()) => + { + Some(ScheduledTurnFinish::Complete) + } + EventMsg::TurnComplete(_) => Some(ScheduledTurnFinish::Failed(schedule_run_error( + "scheduled turn completed without a final assistant message", + ))), + EventMsg::TurnAborted(aborted) => Some(ScheduledTurnFinish::Failed(schedule_run_error( + format!("scheduled turn aborted: {:?}", aborted.reason), + ))), + EventMsg::Error(error) if error.affects_turn_status() => Some(ScheduledTurnFinish::Failed( + schedule_turn_event_error(error), + )), + _ => None, + } +} + +pub(super) fn default_thread_schedule_expires_at(now: DateTime) -> Option> { + now.checked_add_signed(ChronoDuration::days(DEFAULT_SCHEDULE_EXPIRATION_DAYS)) +} + +pub(super) fn next_thread_schedule_run_at( + schedule: &codex_state::ThreadScheduleSpec, + timezone: &str, + after: DateTime, +) -> anyhow::Result>> { + let next = match schedule { + codex_state::ThreadScheduleSpec::Once => None, + codex_state::ThreadScheduleSpec::Dynamic => { + after.checked_add_signed(ChronoDuration::minutes(DEFAULT_DYNAMIC_INTERVAL_MINUTES)) + } + codex_state::ThreadScheduleSpec::Interval(interval) => { + let amount = interval.amount; + let duration = match interval.unit { + codex_state::ThreadScheduleIntervalUnit::Minutes => ChronoDuration::minutes(amount), + codex_state::ThreadScheduleIntervalUnit::Hours => ChronoDuration::hours(amount), + codex_state::ThreadScheduleIntervalUnit::Days => ChronoDuration::days(amount), + }; + after.checked_add_signed(duration) + } + codex_state::ThreadScheduleSpec::Cron { expression } => { + let timezone = parse_schedule_timezone(timezone)?; + let cron = Cron::from_str(expression) + .map_err(|err| anyhow::anyhow!("invalid cron expression `{expression}`: {err}"))?; + let local_after = after.with_timezone(&timezone); + let next = cron.find_next_occurrence(&local_after, /*inclusive*/ false)?; + Some(next.with_timezone(&Utc)) + } + }; + Ok(next) +} + +pub(super) fn next_thread_schedule_run_after_completion( + schedule: &codex_state::ThreadScheduleSpec, + timezone: &str, + scheduled_for: Option>, + completed_at: DateTime, +) -> anyhow::Result>> { + let interval_duration = match schedule { + codex_state::ThreadScheduleSpec::Dynamic => { + Some(ChronoDuration::minutes(DEFAULT_DYNAMIC_INTERVAL_MINUTES)) + } + codex_state::ThreadScheduleSpec::Interval(interval) => { + let amount = interval.amount; + Some(match interval.unit { + codex_state::ThreadScheduleIntervalUnit::Minutes => ChronoDuration::minutes(amount), + codex_state::ThreadScheduleIntervalUnit::Hours => ChronoDuration::hours(amount), + codex_state::ThreadScheduleIntervalUnit::Days => ChronoDuration::days(amount), + }) + } + codex_state::ThreadScheduleSpec::Once | codex_state::ThreadScheduleSpec::Cron { .. } => { + None + } + }; + + if let (Some(interval_duration), Some(scheduled_for)) = (interval_duration, scheduled_for) { + let Some(next_run_at) = scheduled_for.checked_add_signed(interval_duration) else { + return Ok(None); + }; + if next_run_at > completed_at { + return Ok(Some(next_run_at)); + } + let duration_ms = interval_duration.num_milliseconds(); + if duration_ms <= 0 { + return Ok(None); + } + let elapsed_ms = completed_at + .signed_duration_since(scheduled_for) + .num_milliseconds(); + let periods_elapsed = elapsed_ms.div_euclid(duration_ms).saturating_add(1); + return Ok(duration_ms + .checked_mul(periods_elapsed) + .and_then(|advance_ms| { + scheduled_for.checked_add_signed(ChronoDuration::milliseconds(advance_ms)) + })); + } + + next_thread_schedule_run_at(schedule, timezone, completed_at) +} + +pub(super) fn normalize_schedule_timezone(timezone: &str) -> anyhow::Result { + parse_schedule_timezone(timezone).map(|timezone| timezone.name().to_string()) +} + +pub(super) async fn finish_scheduled_run_after_turn( + thread_id: ThreadId, + scheduled_run: crate::thread_state::ScheduledThreadScheduleRun, + event: &EventMsg, + turn_error: Option, + outgoing: &Arc, +) { + let completed_at = Utc::now(); + let error = match (scheduled_turn_finish(event), turn_error) { + (Some(_), Some(error)) => Some(schedule_turn_error(&error)), + (Some(ScheduledTurnFinish::Complete), None) => None, + (Some(ScheduledTurnFinish::Failed(error)), None) => Some(error), + (None, _) => return, + }; + match finish_scheduled_run_state( + &scheduled_run.state_db, + scheduled_run.schedule_id.as_str(), + scheduled_run.run_id.as_str(), + scheduled_run.lease_id.as_str(), + scheduled_run.goal_id.as_deref(), + error, + completed_at, + ) + .await + { + Ok(Some((schedule, run))) => { + outgoing + .send_server_notification(ServerNotification::ThreadScheduleUpdated( + ThreadScheduleUpdatedNotification { + thread_id: thread_id.to_string(), + schedule: api_thread_schedule_from_state(schedule), + }, + )) + .await; + outgoing + .send_server_notification(ServerNotification::ThreadScheduleRunUpdated( + ThreadScheduleRunUpdatedNotification { + thread_id: thread_id.to_string(), + run: api_thread_schedule_run_from_state(run), + }, + )) + .await; + } + Ok(None) => {} + Err(err) => warn!( + schedule_id = %scheduled_run.schedule_id, + thread_id = %thread_id, + "failed to finish scheduled thread run: {err}" + ), + } +} + +pub(super) async fn recover_scheduled_run_for_terminal_turn( + state_db: &StateDbHandle, + thread_id: ThreadId, + turn_id: &str, +) -> anyhow::Result> { + let Some(run) = state_db + .thread_schedules() + .get_running_thread_schedule_run_for_turn(thread_id, turn_id) + .await? + else { + return Ok(None); + }; + Ok(Some(crate::thread_state::ScheduledThreadScheduleRun { + schedule_id: run.schedule_id, + run_id: run.run_id, + lease_id: run.lease_id, + goal_id: run.goal_id, + state_db: state_db.clone(), + })) +} diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence_tests.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence_tests.rs new file mode 100644 index 000000000..6da851f50 --- /dev/null +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence_tests.rs @@ -0,0 +1,604 @@ +use super::*; + +#[tokio::test] +async fn idle_rejection_reuses_one_pending_occurrence_without_counting_a_run() { + let temp_dir = tempfile::tempdir().expect("temp dir should be created"); + let state_db = codex_state::StateRuntime::init( + temp_dir.path().to_path_buf(), + "fallback-provider".to_string(), + ) + .await + .expect("state db should initialize"); + let thread_id = ThreadId::new(); + let mut builder = ThreadMetadataBuilder::new( + thread_id, + temp_dir.path().join("thread.jsonl"), + at(/*seconds*/ 1_700_000_000), + SessionSource::Cli, + ); + builder.cwd = temp_dir.path().join("workspace"); + state_db + .upsert_thread(&builder.build("fallback-provider")) + .await + .expect("thread metadata should persist"); + let now = at(/*seconds*/ 1_700_000_000); + let schedule = state_db + .thread_schedules() + .create_thread_schedule(codex_state::ThreadScheduleCreateParams { + thread_id, + prompt: "wait until the thread is idle".to_string(), + prompt_source: codex_state::ThreadSchedulePromptSource::Inline, + schedule: codex_state::ThreadScheduleSpec::Interval( + codex_state::ThreadScheduleInterval { + amount: 5, + unit: codex_state::ThreadScheduleIntervalUnit::Minutes, + }, + ), + timezone: "UTC".to_string(), + status: codex_state::ThreadScheduleStatus::Active, + next_run_at: Some(now), + expires_at: None, + }) + .await + .expect("schedule should create"); + assert!( + state_db + .thread_schedules() + .claim_due_thread_schedule( + now - chrono::Duration::seconds(1), + "lease-before-due", + Duration::from_secs(300), + ) + .await + .expect("before-due claim should not fail") + .is_none() + ); + assert_eq!( + codex_state::ThreadScheduleStats::default(), + state_db + .thread_schedules() + .get_thread_schedule_stats(schedule.schedule_id.as_str()) + .await + .expect("before-due stats should load") + ); + let claim = state_db + .thread_schedules() + .claim_due_thread_schedule(now, "lease-busy", Duration::from_secs(300)) + .await + .expect("claim should succeed") + .expect("schedule should claim"); + let completed_at = now + chrono::Duration::seconds(5); + let deferral = ScheduleRunDeferral { + kind: ScheduleRunDeferralKind::IdleAdmission, + retry_at: now + chrono::Duration::seconds(SCHEDULE_IDLE_RETRY_DELAY_SECONDS), + error: "scheduled thread is busy".to_string(), + }; + + let deferred_schedule = + wait_scheduled_run_for_idle_state(&state_db, &claim, &deferral, completed_at) + .await + .expect("idle rejection should defer") + .expect("waiting schedule should load"); + + assert_eq!( + codex_state::ThreadSchedule { + next_run_at: Some(now), + last_run_at: None, + failure_count: 0, + lease_id: None, + lease_expires_at: None, + updated_at: deferred_schedule.updated_at, + ..schedule + }, + deferred_schedule + ); + let stats = state_db + .thread_schedules() + .get_thread_schedule_stats(schedule.schedule_id.as_str()) + .await + .expect("schedule stats should load while waiting for idle"); + assert_eq!( + codex_state::ThreadScheduleStats::default(), + stats, + "idle waiting is admission state, not a durable run" + ); + + let retry_claim = state_db + .thread_schedules() + .claim_due_thread_schedule( + deferral.retry_at, + "lease-busy-retry", + Duration::from_secs(300), + ) + .await + .expect("idle retry claim should succeed") + .expect("pending occurrence should be reclaimed"); + assert_eq!(claim.run.run_id, retry_claim.run.run_id); + assert_eq!(claim.run.turn_id, retry_claim.run.turn_id); + let retry_stats = state_db + .thread_schedules() + .get_thread_schedule_stats(schedule.schedule_id.as_str()) + .await + .expect("schedule stats should load after idle retry claim"); + assert_eq!(codex_state::ThreadScheduleStats::default(), retry_stats); +} + +#[tokio::test] +async fn restart_finalizes_a_started_terminal_turn_before_honoring_its_held_goal() { + let temp_dir = tempfile::tempdir().expect("temp dir should be created"); + let state_db = codex_state::StateRuntime::init( + temp_dir.path().to_path_buf(), + "fallback-provider".to_string(), + ) + .await + .expect("state db should initialize"); + let thread_id = ThreadId::new(); + let mut builder = ThreadMetadataBuilder::new( + thread_id, + temp_dir.path().join("thread.jsonl"), + at(/*seconds*/ 1_700_000_000), + SessionSource::Cli, + ); + builder.cwd = temp_dir.path().join("workspace"); + state_db + .upsert_thread(&builder.build("fallback-provider")) + .await + .expect("thread metadata should persist"); + let goal = state_db + .thread_goals() + .replace_thread_goal( + thread_id, + "recover terminal rollout", + codex_state::ThreadGoalStatus::Blocked, + /*token_budget*/ None, + ) + .await + .expect("blocked goal should persist"); + let scheduled_for = at(/*seconds*/ 1_700_000_000); + let schedule = state_db + .thread_schedules() + .create_thread_schedule(codex_state::ThreadScheduleCreateParams { + thread_id, + prompt: "/goal recover terminal rollout".to_string(), + prompt_source: codex_state::ThreadSchedulePromptSource::Inline, + schedule: codex_state::ThreadScheduleSpec::Interval( + codex_state::ThreadScheduleInterval { + amount: 1, + unit: codex_state::ThreadScheduleIntervalUnit::Minutes, + }, + ), + timezone: "UTC".to_string(), + status: codex_state::ThreadScheduleStatus::Active, + next_run_at: Some(scheduled_for), + expires_at: None, + }) + .await + .expect("schedule should create"); + let claim = state_db + .thread_schedules() + .claim_due_thread_schedule(scheduled_for, "lease-first", Duration::from_secs(1)) + .await + .expect("schedule claim should succeed") + .expect("schedule should claim"); + let started = enqueue_and_start_claim( + &state_db, + &claim, + Some(goal.goal_id.as_str()), + scheduled_for, + Duration::from_secs(1), + ) + .await; + let turn_id = started + .turn_id + .clone() + .expect("started occurrence should retain its stable turn id"); + drop(state_db); + + let reopened = codex_state::StateRuntime::init( + temp_dir.path().to_path_buf(), + "fallback-provider".to_string(), + ) + .await + .expect("state db should reopen"); + let recovered = reopened + .thread_schedules() + .claim_due_thread_schedule( + scheduled_for + chrono::Duration::seconds(2), + "lease-recovered", + Duration::from_secs(300), + ) + .await + .expect("started occurrence recovery should succeed") + .expect("expired started occurrence should be reclaimed"); + assert_eq!(claim.run.run_id, recovered.run.run_id); + assert_eq!(Some(turn_id.clone()), recovered.run.turn_id); + assert_eq!(Some(goal.goal_id.clone()), recovered.run.goal_id); + assert_eq!( + codex_state::ThreadScheduleOccurrenceState::Started, + recovered.occurrence_state + ); + + let completed_at = scheduled_for + chrono::Duration::seconds(1); + let history = resumed_history_with_turn_events( + thread_id, + [ + turn_started(turn_id.as_str(), scheduled_for.timestamp()), + EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: turn_id.clone(), + last_agent_message: Some("already finished".to_string()), + completed_at: Some(completed_at.timestamp()), + duration_ms: Some(1_000), + time_to_first_token_ms: Some(100), + }), + ], + ); + let terminal = persisted_scheduled_turn_terminal( + &history, + turn_id.as_str(), + scheduled_for + chrono::Duration::seconds(99), + ) + .expect("the persisted matching turn should already be terminal"); + + let finished = finish_scheduled_run_state( + &reopened, + schedule.schedule_id.as_str(), + recovered.run.run_id.as_str(), + recovered.run.lease_id.as_str(), + recovered.run.goal_id.as_deref(), + terminal.error, + terminal.completed_at, + ) + .await + .expect("persisted terminal turn should finalize") + .expect("one schedule run should finalize"); + assert_eq!( + codex_state::ThreadScheduleRunStatus::Completed, + finished.1.status + ); + assert_eq!(Some(completed_at), finished.1.completed_at); + assert_eq!(codex_state::ThreadScheduleStatus::Paused, finished.0.status); + assert_eq!(None, finished.0.next_run_at); + assert!( + finish_scheduled_run_state( + &reopened, + schedule.schedule_id.as_str(), + recovered.run.run_id.as_str(), + recovered.run.lease_id.as_str(), + recovered.run.goal_id.as_deref(), + /*error*/ None, + completed_at, + ) + .await + .expect("replayed finalization should not fail") + .is_none(), + "terminal replay must not create or finish another run" + ); + assert_eq!( + codex_state::ThreadScheduleStats { + total_runs: 1, + completed_runs: 1, + last_started_at: Some(scheduled_for), + last_completed_at: Some(completed_at), + ..codex_state::ThreadScheduleStats::default() + }, + reopened + .thread_schedules() + .get_thread_schedule_stats(schedule.schedule_id.as_str()) + .await + .expect("schedule stats should load") + ); +} + +#[tokio::test] +async fn stale_started_owner_cannot_submit_after_same_occurrence_is_reclaimed() { + let temp_dir = tempfile::tempdir().expect("temp dir should be created"); + let state_db = codex_state::StateRuntime::init( + temp_dir.path().to_path_buf(), + "fallback-provider".to_string(), + ) + .await + .expect("state db should initialize"); + let thread_id = ThreadId::new(); + let mut builder = ThreadMetadataBuilder::new( + thread_id, + temp_dir.path().join("thread.jsonl"), + at(/*seconds*/ 1_700_000_000), + SessionSource::Cli, + ); + builder.cwd = temp_dir.path().join("workspace"); + state_db + .upsert_thread(&builder.build("fallback-provider")) + .await + .expect("thread metadata should persist"); + let now = at(/*seconds*/ 1_700_000_000); + let schedule = state_db + .thread_schedules() + .create_thread_schedule(codex_state::ThreadScheduleCreateParams { + thread_id, + prompt: "never submit stale work".to_string(), + prompt_source: codex_state::ThreadSchedulePromptSource::Inline, + schedule: codex_state::ThreadScheduleSpec::Interval( + codex_state::ThreadScheduleInterval { + amount: 1, + unit: codex_state::ThreadScheduleIntervalUnit::Minutes, + }, + ), + timezone: "UTC".to_string(), + status: codex_state::ThreadScheduleStatus::Active, + next_run_at: Some(now), + expires_at: None, + }) + .await + .expect("schedule should create"); + let claim = state_db + .thread_schedules() + .claim_due_thread_schedule(now, "lease-suspended", Duration::from_secs(30)) + .await + .expect("claim should succeed") + .expect("schedule should claim"); + enqueue_and_start_claim(&state_db, &claim, None, now, Duration::from_secs(30)).await; + + let contender = codex_state::StateRuntime::init( + temp_dir.path().to_path_buf(), + "fallback-provider".to_string(), + ) + .await + .expect("contending state db should initialize"); + let resumed_at = now + chrono::Duration::seconds(31); + let recovered = contender + .thread_schedules() + .claim_due_thread_schedule(resumed_at, "lease-replacement", Duration::from_secs(30)) + .await + .expect("expired lease reaper should not error") + .expect("expired started run should be recovered"); + assert_eq!(claim.run.run_id, recovered.run.run_id); + assert_eq!(claim.run.turn_id, recovered.run.turn_id); + let submitted = Arc::new(AtomicBool::new(false)); + let submission_observer = Arc::clone(&submitted); + let ownership_lost = CancellationToken::new(); + + let submission = submit_scheduled_turn_if_owned( + &state_db, + codex_state::ThreadScheduleRunLeaseParams { + schedule_id: schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + now: resumed_at, + lease_duration: Duration::from_secs(30), + }, + &ownership_lost, + async move { + submission_observer.store(true, Ordering::SeqCst); + }, + ) + .await + .expect("stale dispatch validation should not error"); + + assert_eq!(None, submission); + assert!(!submitted.load(Ordering::SeqCst)); + assert_eq!( + codex_state::ThreadScheduleRunStatus::Running, + state_db + .thread_schedules() + .get_thread_schedule_run(claim.run.run_id.as_str()) + .await + .expect("recovered run should load") + .expect("recovered run should exist") + .status + ); + assert_eq!( + codex_state::ThreadScheduleRunStatus::Running, + recovered.run.status + ); + + ownership_lost.cancel(); + let cancelled_submission_observer = Arc::clone(&submitted); + let cancelled_submission = submit_scheduled_turn_if_owned( + &state_db, + codex_state::ThreadScheduleRunLeaseParams { + schedule_id: schedule.schedule_id.as_str(), + run_id: recovered.run.run_id.as_str(), + lease_id: recovered.run.lease_id.as_str(), + now: resumed_at + chrono::Duration::seconds(1), + lease_duration: Duration::from_secs(30), + }, + &ownership_lost, + async move { + cancelled_submission_observer.store(true, Ordering::SeqCst); + }, + ) + .await + .expect("cancelled dispatch validation should not error"); + assert_eq!(None, cancelled_submission); + assert!(!submitted.load(Ordering::SeqCst)); +} + +#[test] +fn scheduled_turn_without_agent_message_fails() { + let finish = scheduled_turn_finish(&EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: "turn-1".to_string(), + last_agent_message: None, + completed_at: None, + duration_ms: None, + time_to_first_token_ms: None, + })); + + assert_eq!( + Some(ScheduledTurnFinish::Failed( + "scheduled turn completed without a final assistant message".to_string() + )), + finish + ); +} + +#[test] +fn scheduled_turn_with_agent_message_completes() { + let finish = scheduled_turn_finish(&EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: "turn-1".to_string(), + last_agent_message: Some("done".to_string()), + completed_at: None, + duration_ms: None, + time_to_first_token_ms: None, + })); + + assert_eq!(Some(ScheduledTurnFinish::Complete), finish); +} + +#[test] +fn persisted_completed_scheduled_turn_is_terminal_for_the_matching_turn_only() { + let thread_id = ThreadId::new(); + let completed_at = 1_700_000_005; + let history = resumed_history_with_turn_events( + thread_id, + [ + turn_started("turn-scheduled", 1_700_000_000), + EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: "turn-scheduled".to_string(), + last_agent_message: Some("done".to_string()), + completed_at: Some(completed_at), + duration_ms: Some(5_000), + time_to_first_token_ms: Some(100), + }), + ], + ); + + assert_eq!( + Some(PersistedScheduledTurnTerminal { + completed_at: at(completed_at), + error: None, + }), + persisted_scheduled_turn_terminal( + &history, + "turn-scheduled", + at(/*seconds*/ 1_700_000_999), + ) + ); + assert_eq!( + None, + persisted_scheduled_turn_terminal(&history, "turn-other", at(/*seconds*/ 1_700_000_999),) + ); +} + +#[test] +fn persisted_failed_scheduled_turn_keeps_the_replayed_failure() { + let thread_id = ThreadId::new(); + let history = resumed_history_with_turn_events( + thread_id, + [ + turn_started("turn-scheduled", 1_700_000_000), + EventMsg::Error(ErrorEvent { + message: "model failed".to_string(), + codex_error_info: None, + }), + EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: "turn-scheduled".to_string(), + last_agent_message: None, + completed_at: Some(1_700_000_006), + duration_ms: Some(6_000), + time_to_first_token_ms: None, + }), + ], + ); + + assert_eq!( + Some(PersistedScheduledTurnTerminal { + completed_at: at(/*seconds*/ 1_700_000_006), + error: Some("scheduled turn failed: model failed".to_string()), + }), + persisted_scheduled_turn_terminal( + &history, + "turn-scheduled", + at(/*seconds*/ 1_700_000_999), + ) + ); +} + +#[test] +fn persisted_aborted_scheduled_turn_is_an_explicit_failure() { + let thread_id = ThreadId::new(); + let history = resumed_history_with_turn_events( + thread_id, + [ + turn_started("turn-scheduled", 1_700_000_000), + EventMsg::TurnAborted(TurnAbortedEvent { + turn_id: Some("turn-scheduled".to_string()), + reason: TurnAbortReason::Interrupted, + completed_at: Some(1_700_000_007), + duration_ms: Some(7_000), + }), + ], + ); + + assert_eq!( + Some(PersistedScheduledTurnTerminal { + completed_at: at(/*seconds*/ 1_700_000_007), + error: Some("scheduled turn was interrupted".to_string()), + }), + persisted_scheduled_turn_terminal( + &history, + "turn-scheduled", + at(/*seconds*/ 1_700_000_999), + ) + ); +} + +#[test] +fn persisted_aborted_scheduled_turn_without_a_start_is_an_explicit_failure() { + let thread_id = ThreadId::new(); + let history = resumed_history_with_turn_events( + thread_id, + [EventMsg::TurnAborted(TurnAbortedEvent { + turn_id: Some("turn-scheduled".to_string()), + reason: TurnAbortReason::Interrupted, + completed_at: Some(1_700_000_007), + duration_ms: Some(7_000), + })], + ); + + assert_eq!( + Some(PersistedScheduledTurnTerminal { + completed_at: at(/*seconds*/ 1_700_000_007), + error: Some("scheduled turn aborted: Interrupted".to_string()), + }), + persisted_scheduled_turn_terminal( + &history, + "turn-scheduled", + at(/*seconds*/ 1_700_000_999), + ) + ); +} + +#[test] +fn persisted_in_progress_scheduled_turn_is_not_terminal() { + let thread_id = ThreadId::new(); + let history = resumed_history_with_turn_events( + thread_id, + [ + turn_started("turn-scheduled", 1_700_000_000), + EventMsg::Error(ErrorEvent { + message: "rollback request failed".to_string(), + codex_error_info: Some(CoreCodexErrorInfo::ThreadRollbackFailed), + }), + ], + ); + + assert_eq!( + None, + persisted_scheduled_turn_terminal( + &history, + "turn-scheduled", + at(/*seconds*/ 1_700_000_999), + ) + ); +} + +#[test] +fn scheduled_turn_non_affecting_error_is_not_terminal() { + assert_eq!( + None, + scheduled_turn_finish(&EventMsg::Error(ErrorEvent { + message: "rollback request failed".to_string(), + codex_error_info: Some(CoreCodexErrorInfo::ThreadRollbackFailed), + })) + ); +} diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index e12006ff9..d03a68cff 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -131,6 +131,8 @@ pub enum TryStartUserInputTurnIfIdleError { Rejected(TryStartTurnIfIdleRejectionReason), /// The requested turn settings were rejected before the turn started. InvalidRequest(CodexErr), + /// The durable scheduled occurrence could not be materialized as started. + ScheduleState(anyhow::Error), } impl TryStartUserInputTurnIfIdleError { @@ -139,6 +141,7 @@ impl TryStartUserInputTurnIfIdleError { Self::EmptyInput => None, Self::Rejected(reason) => Some(*reason), Self::InvalidRequest(_) => None, + Self::ScheduleState(_) => None, } } } @@ -149,6 +152,7 @@ impl std::fmt::Display for TryStartUserInputTurnIfIdleError { Self::EmptyInput => write!(f, "turn input must not be empty"), Self::Rejected(reason) => write!(f, "thread is not idle: {reason:?}"), Self::InvalidRequest(err) => write!(f, "{err}"), + Self::ScheduleState(err) => write!(f, "{err}"), } } } @@ -188,6 +192,17 @@ pub struct CodexThreadSettingsOverrides { pub personality: Option, } +/// Durable schedule state that must become `Started` inside the idle-turn +/// reservation, immediately before model work begins. +#[derive(Clone, Debug)] +pub struct ScheduledTurnStart { + pub schedule_id: String, + pub run_id: String, + pub lease_id: String, + pub goal_id: Option, + pub lease_duration: std::time::Duration, +} + pub struct CodexThread { pub(crate) codex: Codex, pub(crate) session_source: SessionSource, @@ -477,6 +492,29 @@ impl CodexThread { .await } + /// Starts a scheduled user-input turn only if idle, materializing the + /// durable run as `Started` under the same reservation before task launch. + pub async fn try_start_scheduled_user_input_turn_if_idle( + &self, + sub_id: String, + items: Vec, + additional_context: BTreeMap, + overrides: CodexThreadSettingsOverrides, + scheduled_start: ScheduledTurnStart, + ) -> Result { + let updates = self.thread_settings_update(overrides).await; + self.codex + .session + .try_start_scheduled_user_input_turn_if_idle( + sub_id, + items, + additional_context, + updates, + scheduled_start, + ) + .await + } + /// Starts a regular turn when trigger-turn mailbox work is pending and the /// thread is idle. pub async fn maybe_start_turn_for_pending_work(&self) -> bool { diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index af5f26a9b..62ffa088b 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -23,6 +23,7 @@ mod config_lock; mod remote_compaction_budget; pub use codex_thread::CodexThread; pub use codex_thread::CodexThreadSettingsOverrides; +pub use codex_thread::ScheduledTurnStart; pub use codex_thread::ThreadConfigSnapshot; pub use codex_thread::TryStartTurnIfIdleError; pub use codex_thread::TryStartTurnIfIdleRejectionReason; diff --git a/codex-rs/core/src/session/inject.rs b/codex-rs/core/src/session/inject.rs index 22dec4626..aeb14c1cd 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/codex-rs/core/src/session/inject.rs @@ -209,10 +209,14 @@ impl Session { input.into_iter().map(TurnInput::ResponseItem).collect(), ) .await; - if self - .start_task(turn_context, Vec::new(), RegularTask::new()) + if !self + .start_task_after_policy_preflight( + turn_context, + Vec::new(), + RegularTask::new(), + Some(Arc::clone(&turn_state)), + ) .await - .is_err() { self.clear_reserved_idle_turn(&turn_state).await; self.maybe_start_turn_for_pending_work().await; @@ -235,6 +239,50 @@ impl Session { additional_context: BTreeMap, updates: SessionSettingsUpdate, ) -> Result { + self.try_start_user_input_turn_if_idle_inner( + sub_id, + input, + additional_context, + updates, + /*scheduled_start*/ None, + ) + .await + .map(|(sub_id, _)| sub_id) + } + + pub(crate) async fn try_start_scheduled_user_input_turn_if_idle( + self: &Arc, + sub_id: String, + input: Vec, + additional_context: BTreeMap, + updates: SessionSettingsUpdate, + scheduled_start: crate::codex_thread::ScheduledTurnStart, + ) -> Result { + self.try_start_user_input_turn_if_idle_inner( + sub_id, + input, + additional_context, + updates, + Some(scheduled_start), + ) + .await? + .1 + .ok_or_else(|| { + TryStartUserInputTurnIfIdleError::ScheduleState(anyhow::anyhow!( + "scheduled turn started without a durable schedule run" + )) + }) + } + + async fn try_start_user_input_turn_if_idle_inner( + self: &Arc, + sub_id: String, + input: Vec, + additional_context: BTreeMap, + updates: SessionSettingsUpdate, + scheduled_start: Option, + ) -> Result<(String, Option), TryStartUserInputTurnIfIdleError> + { if input.is_empty() { return Err(TryStartUserInputTurnIfIdleError::EmptyInput); } @@ -328,15 +376,72 @@ impl Session { content: input, client_id: None, }); - if let Err(error) = self - .start_task(turn_context, task_input, RegularTask::new()) + if !self.still_holds_reserved_idle_turn(&turn_state).await { + self.clear_reserved_idle_turn(&turn_state).await; + return Err(TryStartUserInputTurnIfIdleError::Rejected( + TryStartTurnIfIdleRejectionReason::Busy, + )); + } + let scheduled_run_result = if let Some(scheduled_start) = scheduled_start { + match self.state_db() { + Some(state_db) => state_db + .thread_schedules() + .mark_thread_schedule_run_started(codex_state::ThreadScheduleRunStartParams { + schedule_id: scheduled_start.schedule_id.as_str(), + run_id: scheduled_start.run_id.as_str(), + lease_id: scheduled_start.lease_id.as_str(), + turn_id: sub_id.as_str(), + goal_id: scheduled_start.goal_id.as_deref(), + now: chrono::Utc::now(), + lease_duration: scheduled_start.lease_duration, + }) + .await + .map_err(TryStartUserInputTurnIfIdleError::ScheduleState) + .and_then(|run| { + run.ok_or_else(|| { + TryStartUserInputTurnIfIdleError::ScheduleState(anyhow::anyhow!( + "scheduled occurrence {} no longer owns its active lease", + scheduled_start.run_id + )) + }) + }) + .map(Some), + None => Err(TryStartUserInputTurnIfIdleError::ScheduleState( + anyhow::anyhow!( + "state persistence is unavailable for scheduled turn {}", + scheduled_start.run_id + ), + )), + } + } else { + Ok(None) + }; + let scheduled_run = match scheduled_run_result { + Ok(scheduled_run) => scheduled_run, + Err(error) => { + self.clear_reserved_idle_turn(&turn_state).await; + self.maybe_start_turn_for_pending_work().await; + return Err(error); + } + }; + if !self + .start_task_after_policy_preflight( + turn_context, + task_input, + RegularTask::new(), + Some(Arc::clone(&turn_state)), + ) .await { self.clear_reserved_idle_turn(&turn_state).await; self.maybe_start_turn_for_pending_work().await; - return Err(TryStartUserInputTurnIfIdleError::InvalidRequest(error)); + return Err(TryStartUserInputTurnIfIdleError::InvalidRequest( + CodexErr::InvalidRequest( + "thread became active before scheduled task start".to_string(), + ), + )); } - Ok(sub_id) + Ok((sub_id, scheduled_run)) } /// Returns true when `turn_state` is still the bare (task-less) idle-turn reservation diff --git a/codex-rs/core/src/session/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index 2b5814230..a1a114732 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/codex-rs/core/src/session/input_queue.rs @@ -341,6 +341,34 @@ impl InputQueue { turn_state.lock().await.pending_input.items.split_off(0) } + pub(crate) async fn get_pending_input_for_turn_state( + &self, + turn_state: &Mutex, + ) -> Vec { + let (pending_input, accepts_mailbox_delivery) = { + let mut turn_state = turn_state.lock().await; + ( + turn_state.pending_input.items.split_off(0), + turn_state.accepts_mailbox_delivery_for_current_turn(), + ) + }; + if !accepts_mailbox_delivery { + return pending_input; + } + let mailbox_items = self + .drain_mailbox_input_items() + .await + .into_iter() + .map(TurnInput::ResponseItem); + if pending_input.is_empty() { + mailbox_items.collect() + } else { + let mut pending_input = pending_input; + pending_input.extend(mailbox_items); + pending_input + } + } + #[expect( clippy::await_holding_invalid_type, reason = "active turn checks and turn state updates must remain atomic" diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index ff713f5a6..626303528 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -10928,6 +10928,52 @@ async fn try_start_turn_if_idle_rejects_active_turn_without_injecting() { sess.abort_all_tasks(TurnAbortReason::Interrupted).await; } +#[tokio::test] +async fn reserved_task_start_cannot_overwrite_a_replacement_active_turn() { + let (sess, tc, _rx) = make_session_and_context_with_rx().await; + let reservation = ActiveTurn::default(); + let reserved_turn_state = Arc::clone(&reservation.turn_state); + *sess.active_turn.lock().await = Some(reservation); + + let replacement = ActiveTurn::default(); + let replacement_turn_state = Arc::clone(&replacement.turn_state); + *sess.active_turn.lock().await = Some(replacement); + sess.input_queue + .enqueue_mailbox_communication(InterAgentCommunication::new( + AgentPath::root(), + AgentPath::root(), + Vec::new(), + "replacement trigger".to_string(), + /*trigger_turn*/ true, + )) + .await + .expect("mailbox queue has room"); + + let installed = sess + .start_task_after_policy_preflight( + Arc::clone(&tc), + Vec::new(), + NeverEndingTask { + kind: TaskKind::Regular, + listen_to_cancellation_token: true, + }, + Some(reserved_turn_state), + ) + .await; + + assert!(!installed); + { + let active_turn = sess.active_turn.lock().await; + let active_turn = active_turn.as_ref().expect("replacement remains active"); + assert!(active_turn.task.is_none()); + assert!(Arc::ptr_eq( + &active_turn.turn_state, + &replacement_turn_state + )); + } + assert!(sess.input_queue.has_trigger_turn_mailbox_items().await); +} + #[tokio::test] async fn session_continuation_rejects_active_destination_without_injecting() { let (sess, tc, _rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 2ecd5b8d4..163097045 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -11,7 +11,9 @@ use std::time::Instant; use codex_extension_api::ExtensionData; use futures::future::BoxFuture; use tokio::select; +use tokio::sync::Mutex; use tokio::sync::Notify; +use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; use tokio_util::task::AbortOnDropHandle; use tracing::Instrument; @@ -32,6 +34,7 @@ use crate::session::turn_context::TurnContext; use crate::state::ActiveTurn; use crate::state::RunningTask; use crate::state::TaskKind; +use crate::state::TurnState; use codex_analytics::TurnProfileFact; use codex_analytics::TurnTokenUsageFact; use codex_login::AuthManager; @@ -362,8 +365,17 @@ impl Session { } self.abort_all_tasks(TurnAbortReason::Replaced).await; self.clear_connector_selection().await; - self.start_task_after_policy_preflight(turn_context, input, task) - .await; + if !self + .start_task_after_policy_preflight( + turn_context, + input, + task, + /*expected_turn_state*/ None, + ) + .await + { + warn!("task start lost ownership of its active-turn reservation"); + } } // The explicit return type keeps the `Send` contract visible for lifecycle @@ -387,8 +399,19 @@ impl Session { .await; return Err(error); } - self.start_task_after_policy_preflight(turn_context, input, task) - .await; + if !self + .start_task_after_policy_preflight( + turn_context, + input, + task, + /*expected_turn_state*/ None, + ) + .await + { + return Err(codex_protocol::error::CodexErr::InvalidRequest( + "thread became active before task start".to_string(), + )); + } Ok(()) } } @@ -396,53 +419,33 @@ impl Session { // Keep the explicit `Send` bound visible: this helper participates in a // recursive task-start path that can be awaited from a spawned turn. #[allow(clippy::manual_async_fn)] - fn start_task_after_policy_preflight( + pub(crate) fn start_task_after_policy_preflight( self: &Arc, turn_context: Arc, input: Vec, task: T, - ) -> impl std::future::Future + Send + '_ { + expected_turn_state: Option>>, + ) -> impl std::future::Future + Send + '_ { async move { + let turn_state = match expected_turn_state { + Some(turn_state) => turn_state, + None => { + let mut active = self.active_turn.lock().await; + let turn = active.get_or_insert_with(ActiveTurn::default); + if turn.task.is_some() { + return false; + } + Arc::clone(&turn.turn_state) + } + }; let task: Arc = Arc::new(task); let task_kind = task.kind(); let span_name = task.span_name(); - let started_at = Instant::now(); - let turn_started_at_unix_ms = turn_context - .turn_timing_state - .mark_turn_started(started_at) - .await; - turn_context - .turn_metadata_state - .set_turn_started_at_unix_ms(turn_started_at_unix_ms); - let token_usage_at_turn_start = self.total_token_usage().await.unwrap_or_default(); let cancellation_token = CancellationToken::new(); let done = Arc::new(Notify::new()); - self.services - .guardian_rejection_circuit_breaker - .lock() - .await - .clear_turn(&turn_context.sub_id); - - let pending_items = self.input_queue.get_pending_input(&self.active_turn).await; - let turn_state = { - let mut active = self.active_turn.lock().await; - let turn = active.get_or_insert_with(ActiveTurn::default); - debug_assert!(turn.task.is_none()); - Arc::clone(&turn.turn_state) - }; - turn_state.lock().await.token_usage_at_turn_start = token_usage_at_turn_start.clone(); - self.input_queue - .extend_pending_input_for_turn_state(turn_state.as_ref(), pending_items) - .await; - self.emit_turn_start_lifecycle(turn_context.as_ref(), &token_usage_at_turn_start) - .await; - let turn_extension_data = Arc::clone(&turn_context.extension_data); - let mut active = self.active_turn.lock().await; - let turn = active.get_or_insert_with(ActiveTurn::default); - debug_assert!(turn.task.is_none()); let done_clone = Arc::clone(&done); let session_ctx = Arc::new(SessionTaskContext::new( Arc::clone(self), @@ -451,6 +454,8 @@ impl Session { let ctx = Arc::clone(&turn_context); let task_for_run = Arc::clone(&task); let task_input = input; + let task_turn_state = Arc::clone(&turn_state); + let (start_tx, start_rx) = oneshot::channel(); let task_cancellation_token = cancellation_token.child_token(); // Task-owned turn spans keep a core-owned span open for the // full task lifecycle after the submission dispatch span ends. @@ -472,6 +477,47 @@ impl Session { ); let handle = tokio::spawn( async move { + let may_start = select! { + result = start_rx => result.is_ok(), + _ = task_cancellation_token.cancelled() => false, + }; + if !may_start { + done_clone.notify_one(); + return; + } + let sess = session_ctx.clone_session(); + let started_at = Instant::now(); + let turn_started_at_unix_ms = ctx + .turn_timing_state + .mark_turn_started(started_at) + .await; + ctx.turn_metadata_state + .set_turn_started_at_unix_ms(turn_started_at_unix_ms); + let token_usage_at_turn_start = + sess.total_token_usage().await.unwrap_or_default(); + sess.services + .guardian_rejection_circuit_breaker + .lock() + .await + .clear_turn(&ctx.sub_id); + let pending_items = sess + .input_queue + .get_pending_input_for_turn_state(task_turn_state.as_ref()) + .await; + task_turn_state.lock().await.token_usage_at_turn_start = + token_usage_at_turn_start.clone(); + sess.input_queue + .extend_pending_input_for_turn_state( + task_turn_state.as_ref(), + pending_items, + ) + .await; + sess.emit_turn_start_lifecycle(ctx.as_ref(), &token_usage_at_turn_start) + .await; + if task_cancellation_token.is_cancelled() { + done_clone.notify_one(); + return; + } let ctx_for_finish = Arc::clone(&ctx); let last_agent_message = task_for_run .run( @@ -481,7 +527,6 @@ impl Session { task_cancellation_token.child_token(), ) .await; - let sess = session_ctx.clone_session(); if let Err(err) = sess.flush_rollout().await { warn!("failed to flush rollout before completing turn: {err}"); sess.send_event( @@ -517,7 +562,22 @@ impl Session { turn_extension_data, _timer: timer, }; - turn.task = Some(running_task); + let installed = { + let mut active = self.active_turn.lock().await; + match active.as_mut() { + Some(turn) + if turn.task.is_none() && Arc::ptr_eq(&turn.turn_state, &turn_state) => + { + turn.task = Some(running_task); + true + } + _ => false, + } + }; + if installed { + let _ = start_tx.send(()); + } + installed } } @@ -580,7 +640,7 @@ impl Session { return false; } - { + let turn_state = { let mut active_turn = self.active_turn.lock().await; if active_turn .as_ref() @@ -593,13 +653,23 @@ impl Session { // mail is user/client-directed, so it takes over stale or lower-priority // reservations and lets their owners observe that the reservation was lost. *active_turn = Some(ActiveTurn::default()); - } + Arc::clone( + &active_turn + .as_ref() + .expect("pending-work reservation should be present") + .turn_state, + ) + }; self.maybe_emit_unknown_model_warning_for_turn(turn_context.as_ref()) .await; - self.start_task_after_policy_preflight(turn_context, Vec::new(), RegularTask::new()) - .await; - true + self.start_task_after_policy_preflight( + turn_context, + Vec::new(), + RegularTask::new(), + Some(turn_state), + ) + .await } pub async fn abort_all_tasks(self: &Arc, reason: TurnAbortReason) { diff --git a/codex-rs/core/src/tools/handlers/loop_control.rs b/codex-rs/core/src/tools/handlers/loop_control.rs index 4b69de053..d3b3464ce 100644 --- a/codex-rs/core/src/tools/handlers/loop_control.rs +++ b/codex-rs/core/src/tools/handlers/loop_control.rs @@ -985,6 +985,47 @@ mod tests { .expect("schedule should be created") } + async fn enqueue_and_start_claim( + runtime: &codex_state::StateRuntime, + claim: &codex_state::ThreadScheduleClaim, + now: DateTime, + ) -> codex_state::ThreadScheduleRun { + let turn_id = claim + .run + .turn_id + .as_deref() + .expect("claimed occurrence should have a stable turn id"); + runtime + .thread_schedules() + .enqueue_thread_schedule_run(codex_state::ThreadScheduleRunEnqueueParams { + schedule_id: claim.schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + goal_id: None, + auth_profile_recorded: false, + auth_profile: None, + turn_input: "loop-control test input", + now, + }) + .await + .expect("claimed occurrence should enqueue") + .expect("claimed occurrence should retain its lease"); + runtime + .thread_schedules() + .mark_thread_schedule_run_started(codex_state::ThreadScheduleRunStartParams { + schedule_id: claim.schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + turn_id, + goal_id: None, + now, + lease_duration: Duration::from_secs(300), + }) + .await + .expect("enqueued occurrence should start") + .expect("enqueued occurrence should retain its lease") + } + #[test] fn create_action_deserializes() { let args: ManageLoopArgs = serde_json::from_str( @@ -1301,20 +1342,7 @@ mod tests { .await .expect("first run should claim") .expect("first run should be due"); - runtime - .thread_schedules() - .mark_thread_schedule_run_started(codex_state::ThreadScheduleRunStartParams { - schedule_id: &schedule.schedule_id, - run_id: &first_claim.run.run_id, - lease_id: "lease-complete", - turn_id: "turn-complete", - goal_id: None, - now: first_run_at, - lease_duration: Duration::from_secs(300), - }) - .await - .expect("first run should start") - .expect("first run should exist"); + enqueue_and_start_claim(&runtime, &first_claim, first_run_at).await; let second_run_at = at(/*seconds*/ 1_700_000_600); runtime @@ -1335,20 +1363,7 @@ mod tests { .await .expect("second run should claim") .expect("second run should be due"); - runtime - .thread_schedules() - .mark_thread_schedule_run_started(codex_state::ThreadScheduleRunStartParams { - schedule_id: &schedule.schedule_id, - run_id: &second_claim.run.run_id, - lease_id: "lease-fail", - turn_id: "turn-fail", - goal_id: None, - now: second_run_at, - lease_duration: Duration::from_secs(300), - }) - .await - .expect("second run should start") - .expect("second run should exist"); + enqueue_and_start_claim(&runtime, &second_claim, second_run_at).await; runtime .thread_schedules() .fail_thread_schedule_run( @@ -1574,20 +1589,7 @@ mod tests { .await .expect("run should claim") .expect("run should be due"); - runtime - .thread_schedules() - .mark_thread_schedule_run_started(codex_state::ThreadScheduleRunStartParams { - schedule_id: &first.schedule_id, - run_id: &claim.run.run_id, - lease_id: "lease-complete", - turn_id: "turn-complete", - goal_id: None, - now: first_run_at, - lease_duration: Duration::from_secs(300), - }) - .await - .expect("run should start") - .expect("run should exist"); + enqueue_and_start_claim(&runtime, &claim, first_run_at).await; runtime .thread_schedules() .complete_thread_schedule_run( @@ -1682,6 +1684,7 @@ mod tests { .await .expect("run should claim") .expect("run should be due"); + enqueue_and_start_claim(&runtime, &claim, at(/*seconds*/ 1_700_000_300)).await; runtime .thread_schedules() .fail_thread_schedule_run( diff --git a/codex-rs/core/src/tools/handlers/schedule_control.rs b/codex-rs/core/src/tools/handlers/schedule_control.rs index ebe936317..0c90f7b4a 100644 --- a/codex-rs/core/src/tools/handlers/schedule_control.rs +++ b/codex-rs/core/src/tools/handlers/schedule_control.rs @@ -918,6 +918,47 @@ mod tests { .expect("schedule should be created") } + async fn enqueue_and_start_claim( + runtime: &codex_state::StateRuntime, + claim: &codex_state::ThreadScheduleClaim, + now: DateTime, + ) { + let turn_id = claim + .run + .turn_id + .as_deref() + .expect("claimed occurrence should have a stable turn id"); + runtime + .thread_schedules() + .enqueue_thread_schedule_run(codex_state::ThreadScheduleRunEnqueueParams { + schedule_id: claim.schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + goal_id: None, + auth_profile_recorded: false, + auth_profile: None, + turn_input: "schedule-control test input", + now, + }) + .await + .expect("claimed occurrence should enqueue") + .expect("claimed occurrence should retain its lease"); + runtime + .thread_schedules() + .mark_thread_schedule_run_started(codex_state::ThreadScheduleRunStartParams { + schedule_id: claim.schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + turn_id, + goal_id: None, + now, + lease_duration: std::time::Duration::from_secs(60), + }) + .await + .expect("enqueued occurrence should start") + .expect("enqueued occurrence should retain its lease"); + } + #[tokio::test] async fn create_adds_one_time_schedule() { let (_temp_dir, runtime) = test_runtime().await; @@ -1122,6 +1163,7 @@ mod tests { .await .expect("schedule claim should succeed") .expect("schedule should be due"); + enqueue_and_start_claim(&runtime, &claim, at(/*seconds*/ 1_700_000_301)).await; runtime .thread_schedules() .fail_thread_schedule_run( diff --git a/codex-rs/state/migrations/0063_thread_schedule_occurrence_state.sql b/codex-rs/state/migrations/0063_thread_schedule_occurrence_state.sql new file mode 100644 index 000000000..6b5212e29 --- /dev/null +++ b/codex-rs/state/migrations/0063_thread_schedule_occurrence_state.sql @@ -0,0 +1,198 @@ +ALTER TABLE thread_schedule_runs +ADD COLUMN deferral_kind TEXT CHECK(deferral_kind IS NULL OR deferral_kind IN ('idle', 'capacity')); + +UPDATE thread_schedule_runs +SET deferral_kind = 'idle' +WHERE status = 'deferred' + AND error IN ( + 'scheduled thread is busy', + 'scheduled thread has pending mailbox trigger-turn work' + ); + +UPDATE thread_schedule_runs +SET deferral_kind = 'capacity' +WHERE status = 'deferred' + AND deferral_kind IS NULL; + +UPDATE thread_schedule_runs +SET status = 'failed', + error = 'superseded by later active schedule run during occurrence migration', + completed_at_ms = COALESCE(completed_at_ms, started_at_ms) +WHERE status IN ('leased', 'running') + AND EXISTS ( + SELECT 1 + FROM thread_schedule_runs AS later + WHERE later.schedule_id = thread_schedule_runs.schedule_id + AND later.status IN ('leased', 'running') + AND ( + later.started_at_ms > thread_schedule_runs.started_at_ms + OR ( + later.started_at_ms = thread_schedule_runs.started_at_ms + AND later.rowid > thread_schedule_runs.rowid + ) + ) + ); + +CREATE TABLE thread_schedule_occurrences ( + occurrence_id TEXT PRIMARY KEY NOT NULL, + schedule_id TEXT NOT NULL UNIQUE, + thread_id TEXT NOT NULL, + state TEXT NOT NULL CHECK(state IN ('waiting_idle', 'enqueued', 'started', 'terminal')), + turn_id TEXT NOT NULL, + goal_id TEXT, + auth_profile_recorded INTEGER NOT NULL DEFAULT 0 CHECK(auth_profile_recorded IN (0, 1)), + auth_profile TEXT, + scheduled_for_ms INTEGER, + retry_at_ms INTEGER, + turn_input TEXT, + created_at_ms INTEGER NOT NULL, + updated_at_ms INTEGER NOT NULL, + FOREIGN KEY(schedule_id) REFERENCES thread_schedules(schedule_id) ON DELETE CASCADE +); + +INSERT INTO thread_schedule_occurrences ( + occurrence_id, + schedule_id, + thread_id, + state, + turn_id, + goal_id, + scheduled_for_ms, + retry_at_ms, + turn_input, + created_at_ms, + updated_at_ms +) +SELECT + run_id, + schedule_id, + thread_id, + CASE status WHEN 'running' THEN 'started' ELSE 'waiting_idle' END, + COALESCE(turn_id, run_id), + goal_id, + scheduled_for_ms, + NULL, + NULL, + started_at_ms, + started_at_ms +FROM thread_schedule_runs +WHERE status IN ('leased', 'running'); + +DELETE FROM thread_schedule_runs +WHERE status = 'leased' + AND EXISTS ( + SELECT 1 + FROM thread_schedule_occurrences + WHERE thread_schedule_occurrences.occurrence_id = thread_schedule_runs.run_id + AND thread_schedule_occurrences.state = 'waiting_idle' + ); + +CREATE INDEX idx_thread_schedule_occurrences_ready + ON thread_schedule_occurrences(state, retry_at_ms, updated_at_ms); + +CREATE TRIGGER thread_schedule_runs_reject_legacy_duplicate_occurrence_insert +BEFORE INSERT ON thread_schedule_runs +WHEN ( + NEW.status IN ('leased', 'running') + AND NOT EXISTS ( + SELECT 1 + FROM thread_schedule_occurrences + WHERE thread_schedule_occurrences.schedule_id = NEW.schedule_id + AND thread_schedule_occurrences.occurrence_id = NEW.run_id + AND thread_schedule_occurrences.state != 'terminal' + ) +) OR EXISTS ( + SELECT 1 + FROM thread_schedule_occurrences + WHERE thread_schedule_occurrences.schedule_id = NEW.schedule_id + AND ( + thread_schedule_occurrences.occurrence_id != NEW.run_id + OR thread_schedule_occurrences.state = 'terminal' + ) + ) +BEGIN + SELECT RAISE(ABORT, 'active schedule occurrence must be reused'); +END; + +CREATE TRIGGER thread_schedule_occurrences_follow_legacy_terminal_update +AFTER UPDATE OF status ON thread_schedule_runs +WHEN NEW.status IN ('deferred', 'completed', 'failed') + AND EXISTS ( + SELECT 1 + FROM thread_schedules + WHERE thread_schedules.schedule_id = NEW.schedule_id + AND thread_schedules.lease_id IS NULL + ) +BEGIN + UPDATE thread_schedule_occurrences + SET state = 'terminal', updated_at_ms = COALESCE(NEW.completed_at_ms, updated_at_ms) + WHERE occurrence_id = NEW.run_id; + + DELETE FROM thread_schedule_occurrences + WHERE occurrence_id = NEW.run_id + AND EXISTS ( + SELECT 1 + FROM thread_schedules + WHERE thread_schedules.schedule_id = thread_schedule_occurrences.schedule_id + AND thread_schedules.lease_id IS NULL + ); +END; + +CREATE TRIGGER thread_schedule_occurrences_follow_legacy_schedule_hold +AFTER UPDATE OF status, lease_id ON thread_schedules +WHEN NEW.status IN ('paused', 'expired') + AND NEW.lease_id IS NULL + AND EXISTS ( + SELECT 1 + FROM thread_schedule_occurrences + WHERE thread_schedule_occurrences.schedule_id = NEW.schedule_id + ) +BEGIN + UPDATE thread_schedule_runs + SET status = 'failed', + error = CASE NEW.status + WHEN 'paused' THEN 'scheduled run cancelled because schedule was paused' + ELSE 'scheduled run cancelled because schedule expired' + END, + completed_at_ms = COALESCE(completed_at_ms, NEW.updated_at_ms) + WHERE run_id IN ( + SELECT occurrence_id + FROM thread_schedule_occurrences + WHERE schedule_id = NEW.schedule_id AND state = 'started' + ) + AND status = 'running'; + + INSERT INTO thread_schedule_runs ( + run_id, + schedule_id, + thread_id, + status, + lease_id, + turn_id, + goal_id, + error, + scheduled_for_ms, + started_at_ms, + completed_at_ms + ) + SELECT + occurrence_id, + schedule_id, + thread_id, + 'failed', + COALESCE(OLD.lease_id, 'legacy-schedule-hold'), + turn_id, + goal_id, + CASE NEW.status + WHEN 'paused' THEN 'scheduled run cancelled because schedule was paused' + ELSE 'scheduled run cancelled because schedule expired' + END, + scheduled_for_ms, + created_at_ms, + NEW.updated_at_ms + FROM thread_schedule_occurrences + WHERE schedule_id = NEW.schedule_id AND state = 'enqueued'; + + DELETE FROM thread_schedule_occurrences + WHERE schedule_id = NEW.schedule_id; +END; diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index 6c407144c..706f46c79 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -249,6 +249,8 @@ pub use runtime::ThreadScheduleClaim; pub use runtime::ThreadScheduleCreateParams; pub use runtime::ThreadScheduleDueClaimParams; pub use runtime::ThreadScheduleNowClaimParams; +pub use runtime::ThreadScheduleOccurrenceState; +pub use runtime::ThreadScheduleRunEnqueueParams; pub use runtime::ThreadScheduleRunForGoalFinishParams; pub use runtime::ThreadScheduleRunLeaseParams; pub use runtime::ThreadScheduleRunStartParams; diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index caea0489a..8c76a3074 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -252,6 +252,8 @@ pub use schedules::ThreadScheduleClaim; pub use schedules::ThreadScheduleCreateParams; pub use schedules::ThreadScheduleDueClaimParams; pub use schedules::ThreadScheduleNowClaimParams; +pub use schedules::ThreadScheduleOccurrenceState; +pub use schedules::ThreadScheduleRunEnqueueParams; pub use schedules::ThreadScheduleRunForGoalFinishParams; pub use schedules::ThreadScheduleRunLeaseParams; pub use schedules::ThreadScheduleRunStartParams; diff --git a/codex-rs/state/src/runtime/schedules.rs b/codex-rs/state/src/runtime/schedules.rs index ccab8e046..2e662e95e 100644 --- a/codex-rs/state/src/runtime/schedules.rs +++ b/codex-rs/state/src/runtime/schedules.rs @@ -42,81 +42,21 @@ pub struct ThreadScheduleUpdate { pub expires_at: Option>>, } -#[derive(Clone)] -pub struct ThreadScheduleClaim { - pub schedule: crate::ThreadSchedule, - pub run: crate::ThreadScheduleRun, -} - -#[derive(Clone)] -pub struct ThreadScheduleDueClaimParams<'a> { - pub now: DateTime, - pub lease_id: &'a str, - pub lease_duration: Duration, - pub local_active_owner_id: Option<&'a str>, - pub local_active_fresh_after: Option>, -} - -#[derive(Clone)] -pub struct ThreadScheduleNowClaimParams<'a> { - pub schedule_id: &'a str, - pub now: DateTime, - pub lease_id: &'a str, - pub lease_duration: Duration, - pub local_active_owner_id: Option<&'a str>, - pub local_active_fresh_after: Option>, -} - -pub struct ThreadScheduleRunForGoalFinishParams<'a> { - pub schedule_id: &'a str, - pub run_id: &'a str, - pub lease_id: &'a str, - pub completed_at: DateTime, - pub next_run_at: Option>, - pub expected_goal_id: &'a str, -} - -#[derive(Clone)] -pub struct ThreadScheduleRunStartParams<'a> { - pub schedule_id: &'a str, - pub run_id: &'a str, - pub lease_id: &'a str, - pub turn_id: &'a str, - pub goal_id: Option<&'a str>, - pub now: DateTime, - pub lease_duration: Duration, -} - -#[derive(Clone)] -pub struct ThreadScheduleRunLeaseParams<'a> { - pub schedule_id: &'a str, - pub run_id: &'a str, - pub lease_id: &'a str, - pub now: DateTime, - pub lease_duration: Duration, -} +mod occurrence; +pub use occurrence::ThreadScheduleClaim; +pub use occurrence::ThreadScheduleDueClaimParams; +pub use occurrence::ThreadScheduleNowClaimParams; +pub use occurrence::ThreadScheduleOccurrenceState; +pub use occurrence::ThreadScheduleRunEnqueueParams; +pub use occurrence::ThreadScheduleRunForGoalFinishParams; +pub use occurrence::ThreadScheduleRunLeaseParams; +pub use occurrence::ThreadScheduleRunStartParams; struct ScheduleNesting { parent_schedule_id: Option, nesting_depth: i64, } -#[derive(Clone, Copy)] -enum ThreadScheduleClaimTarget<'a> { - Due, - Now { schedule_id: &'a str }, -} - -#[derive(Clone)] -struct ClaimThreadScheduleParams<'a> { - target: ThreadScheduleClaimTarget<'a>, - now: DateTime, - lease_id: &'a str, - lease_duration: Duration, - local_active_owner_id: Option<&'a str>, - local_active_fresh_after: Option>, -} - impl ScheduleStore { pub async fn create_thread_schedule( &self, @@ -391,6 +331,10 @@ WHERE schedule_id = ? AND status IN ('leased', 'running') .bind(schedule_id) .execute(&mut *tx) .await?; + sqlx::query("DELETE FROM thread_schedule_occurrences WHERE schedule_id = ?") + .bind(schedule_id) + .execute(&mut *tx) + .await?; } } let schedule = row.map(|row| thread_schedule_from_row(&row)).transpose()?; @@ -582,1089 +526,84 @@ ORDER BY depth, created_at_ms, schedule_id } let parent = self .get_thread_schedule(parent_schedule_id) - .await? - .ok_or_else(|| { - anyhow::anyhow!( - "invalid nested loop: parent schedule not found: {parent_schedule_id}" - ) - })?; - self.validate_parent_schedule(&parent, params.thread_id, ¶ms.schedule)?; - Ok(ScheduleNesting { - parent_schedule_id: Some(parent.schedule_id), - nesting_depth: parent.nesting_depth + 1, - }) - } - - async fn validate_schedule_update_nesting( - &self, - existing: &crate::ThreadSchedule, - schedule: &crate::ThreadScheduleSpec, - ) -> anyhow::Result<()> { - if self - .has_child_thread_schedules(existing.schedule_id.as_str()) - .await? - { - anyhow::bail!( - "invalid nested loop: cannot update loop cadence while it has nested child loops; update or clear child loops first" - ); - } - let Some(parent_schedule_id) = existing.parent_schedule_id.as_deref() else { - return Ok(()); - }; - if matches!(schedule, crate::ThreadScheduleSpec::Once) { - anyhow::bail!("invalid nested loop: one-time schedules cannot be nested"); - } - let parent = self - .get_thread_schedule(parent_schedule_id) - .await? - .ok_or_else(|| { - anyhow::anyhow!( - "invalid nested loop: parent schedule not found: {parent_schedule_id}" - ) - })?; - self.validate_parent_schedule(&parent, existing.thread_id, schedule) - } - - fn validate_parent_schedule( - &self, - parent: &crate::ThreadSchedule, - thread_id: ThreadId, - child_schedule: &crate::ThreadScheduleSpec, - ) -> anyhow::Result<()> { - if parent.thread_id != thread_id { - anyhow::bail!("invalid nested loop: parent schedule must belong to the same thread"); - } - if matches!(parent.schedule, crate::ThreadScheduleSpec::Once) { - anyhow::bail!("invalid nested loop: parent schedule must be recurring"); - } - if parent.nesting_depth >= MAX_THREAD_SCHEDULE_NESTING_DEPTH { - anyhow::bail!( - "invalid nested loop: maximum nesting depth is {MAX_THREAD_SCHEDULE_NESTING_DEPTH}" - ); - } - validate_nested_loop_cadence(&parent.schedule, child_schedule) - } - - async fn has_child_thread_schedules(&self, schedule_id: &str) -> anyhow::Result { - let count: i64 = sqlx::query_scalar( - r#" -SELECT COUNT(*) -FROM thread_schedules -WHERE parent_schedule_id = ? - "#, - ) - .bind(schedule_id) - .fetch_one(self.pool.as_ref()) - .await?; - Ok(count > 0) - } - - pub async fn get_thread_schedule_run( - &self, - run_id: &str, - ) -> anyhow::Result> { - let sql = run_returning( - r#" -SELECT - "#, - ); - let row = sqlx::query(sqlx::AssertSqlSafe(format!( - "{sql}FROM thread_schedule_runs WHERE run_id = ?" - ))) - .bind(run_id) - .fetch_optional(self.pool.as_ref()) - .await?; - row.map(|row| thread_schedule_run_from_row(&row)) - .transpose() - } - - pub async fn get_running_thread_schedule_run_for_turn( - &self, - thread_id: ThreadId, - turn_id: &str, - ) -> anyhow::Result> { - let sql = run_returning( - r#" -SELECT -"#, - ); - let row = sqlx::query(sqlx::AssertSqlSafe(format!( - r#"{sql} -FROM thread_schedule_runs -WHERE thread_id = ? - AND turn_id = ? - AND status = 'running' -ORDER BY started_at_ms DESC -LIMIT 1 -"# - ))) - .bind(thread_id.to_string()) - .bind(turn_id) - .fetch_optional(self.pool.as_ref()) - .await?; - row.map(|row| thread_schedule_run_from_row(&row)) - .transpose() - } - - pub async fn get_thread_schedule_stats( - &self, - schedule_id: &str, - ) -> anyhow::Result { - let row = sqlx::query( - r#" -SELECT - COUNT(*) AS total_runs, - COALESCE(SUM(CASE WHEN status = 'leased' THEN 1 ELSE 0 END), 0) AS leased_runs, - COALESCE(SUM(CASE WHEN status = 'running' THEN 1 ELSE 0 END), 0) AS running_runs, - COALESCE(SUM(CASE WHEN status = 'deferred' THEN 1 ELSE 0 END), 0) AS deferred_runs, - COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0) AS completed_runs, - COALESCE(SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END), 0) AS failed_runs, - MAX(started_at_ms) AS last_started_at_ms, - -- Only successfully completed runs contribute to last_completed_at. The - -- completed_at_ms column is also written for deferred and failed runs (it is - -- really a "finished at" timestamp), so deriving last_completed_at from the - -- raw MAX would populate it even when completed_runs is 0. Keeping this - -- filtered ensures last_completed_at is non-null iff completed_runs > 0. - MAX(CASE WHEN status = 'completed' THEN completed_at_ms END) AS last_completed_at_ms -FROM thread_schedule_runs -WHERE schedule_id = ? - "#, - ) - .bind(schedule_id) - .fetch_one(self.pool.as_ref()) - .await?; - let last_error = sqlx::query_scalar( - r#" -SELECT error -FROM thread_schedule_runs -WHERE schedule_id = ? - AND status = 'failed' - AND error IS NOT NULL - AND TRIM(error) != '' -ORDER BY completed_at_ms DESC, started_at_ms DESC -LIMIT 1 - "#, - ) - .bind(schedule_id) - .fetch_optional(self.pool.as_ref()) - .await?; - Ok(crate::ThreadScheduleStats { - total_runs: row.try_get("total_runs")?, - leased_runs: row.try_get("leased_runs")?, - running_runs: row.try_get("running_runs")?, - deferred_runs: row.try_get("deferred_runs")?, - completed_runs: row.try_get("completed_runs")?, - failed_runs: row.try_get("failed_runs")?, - last_started_at: row - .try_get::, _>("last_started_at_ms")? - .map(epoch_millis_to_datetime) - .transpose()?, - last_completed_at: row - .try_get::, _>("last_completed_at_ms")? - .map(epoch_millis_to_datetime) - .transpose()?, - last_error, - }) - } - - pub async fn claim_due_thread_schedule( - &self, - now: DateTime, - lease_id: &str, - lease_duration: Duration, - ) -> anyhow::Result> { - self.claim_due_thread_schedule_with_params(ThreadScheduleDueClaimParams { - now, - lease_id, - lease_duration, - local_active_owner_id: None, - local_active_fresh_after: None, - }) - .await - } - - pub async fn claim_due_thread_schedule_with_params( - &self, - params: ThreadScheduleDueClaimParams<'_>, - ) -> anyhow::Result> { - let ThreadScheduleDueClaimParams { - now, - lease_id, - lease_duration, - local_active_owner_id, - local_active_fresh_after, - } = params; - let params = ClaimThreadScheduleParams { - target: ThreadScheduleClaimTarget::Due, - now, - lease_id, - lease_duration, - local_active_owner_id, - local_active_fresh_after, - }; - crate::busy_retry::retry_on_busy("claim due thread schedule", || { - self.claim_thread_schedule_once(params.clone()) - }) - .await - } - - pub async fn claim_thread_schedule_now( - &self, - schedule_id: &str, - now: DateTime, - lease_id: &str, - lease_duration: Duration, - ) -> anyhow::Result> { - self.claim_thread_schedule_now_with_params(ThreadScheduleNowClaimParams { - schedule_id, - now, - lease_id, - lease_duration, - local_active_owner_id: None, - local_active_fresh_after: None, - }) - .await - } - - pub async fn claim_thread_schedule_now_with_params( - &self, - params: ThreadScheduleNowClaimParams<'_>, - ) -> anyhow::Result> { - let ThreadScheduleNowClaimParams { - schedule_id, - now, - lease_id, - lease_duration, - local_active_owner_id, - local_active_fresh_after, - } = params; - let params = ClaimThreadScheduleParams { - target: ThreadScheduleClaimTarget::Now { schedule_id }, - now, - lease_id, - lease_duration, - local_active_owner_id, - local_active_fresh_after, - }; - crate::busy_retry::retry_on_busy("claim thread schedule now", || { - self.claim_thread_schedule_once(params.clone()) - }) - .await - } - - async fn claim_thread_schedule_once( - &self, - params: ClaimThreadScheduleParams<'_>, - ) -> anyhow::Result> { - let ClaimThreadScheduleParams { - target, - now, - lease_id, - lease_duration, - local_active_owner_id, - local_active_fresh_after, - } = params; - let now_ms = datetime_to_epoch_millis(now); - let lease_expires_at = now + chrono::Duration::from_std(lease_duration)?; - let lease_expires_at_ms = datetime_to_epoch_millis(lease_expires_at); - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; - let owner_filter = match (local_active_owner_id, local_active_fresh_after) { - (Some(owner_id), Some(fresh_after)) => { - Some((owner_id, datetime_to_epoch_millis(fresh_after))) - } - _ => None, - }; - let owner_scoped_lease_id = owner_filter.as_ref().map(|_| format!("owner:{lease_id}")); - let lease_id = owner_scoped_lease_id.as_deref().unwrap_or(lease_id); - let active_owner_filter = if owner_filter.is_some() { - r#" - AND NOT EXISTS ( - SELECT 1 - FROM local_active_sessions - WHERE local_active_sessions.thread_id = thread_schedules.thread_id - AND local_active_sessions.last_seen_at_ms >= ? - AND local_active_sessions.owner_id != ? - ) -"# - } else { - "" - }; - let sql = match target { - ThreadScheduleClaimTarget::Due => format!( - r#" -SELECT {SCHEDULE_COLUMNS} -FROM thread_schedules -WHERE status = 'active' - AND next_run_at_ms IS NOT NULL - AND next_run_at_ms <= ? - AND (expires_at_ms IS NULL OR expires_at_ms > ?) - AND (lease_id IS NULL OR lease_expires_at_ms <= ?) -{active_owner_filter} -ORDER BY next_run_at_ms, created_at_ms -LIMIT 1 -"# - ), - ThreadScheduleClaimTarget::Now { .. } => format!( - r#" -SELECT {SCHEDULE_COLUMNS} -FROM thread_schedules -WHERE schedule_id = ? - AND status = 'active' - AND (expires_at_ms IS NULL OR expires_at_ms > ?) - AND (lease_id IS NULL OR lease_expires_at_ms <= ?) -{active_owner_filter} -"# - ), - }; - let mut query = sqlx::query(sqlx::AssertSqlSafe(sql)); - query = match target { - ThreadScheduleClaimTarget::Due => query.bind(now_ms).bind(now_ms).bind(now_ms), - ThreadScheduleClaimTarget::Now { schedule_id } => { - query.bind(schedule_id).bind(now_ms).bind(now_ms) - } - }; - if let Some((owner_id, fresh_after_ms)) = owner_filter { - query = query.bind(fresh_after_ms).bind(owner_id); - } - let schedule_row = query.fetch_optional(&mut *tx).await?; - let Some(schedule_row) = schedule_row else { - tx.commit().await?; - return Ok(None); - }; - let selected_schedule = thread_schedule_from_row(&schedule_row)?; - let active_goal_ids: Vec> = sqlx::query_scalar( - r#" -SELECT goal_id -FROM thread_schedule_runs -WHERE schedule_id = ? AND status IN ('leased', 'running') -ORDER BY started_at_ms DESC - "#, - ) - .bind(selected_schedule.schedule_id.as_str()) - .fetch_all(&mut *tx) - .await?; - let mut goal_ids = active_goal_ids - .iter() - .filter_map(Clone::clone) - .collect::>(); - goal_ids.sort(); - goal_ids.dedup(); - let goal_hold_can_pause = - !goal_ids.is_empty() && selected_schedule.schedule != crate::ThreadScheduleSpec::Once; - // Read-only probe: this transaction only runs `SELECT EXISTS` against - // goals.db and is always rolled back, so a deferred `BEGIN` is enough. A - // `BEGIN IMMEDIATE` would take a goals.db write lock and hold it across the - // state.db commit for no benefit. Lock order is consistently state -> goals - // at every site that touches both, so there is no inversion to guard against. - let mut goal_tx = if goal_hold_can_pause { - Some(self.goals_pool.begin().await?) - } else { - None - }; - let mut pause_for_goal_hold = false; - if let Some(goal_tx) = goal_tx.as_mut() { - for goal_id in goal_ids { - pause_for_goal_hold = sqlx::query_scalar::<_, bool>( - r#" -SELECT EXISTS( - SELECT 1 - FROM thread_goals - WHERE thread_id = ? - AND goal_id = ? - AND status IN ('paused', 'blocked', 'usage_limited', 'budget_limited') -) - "#, - ) - .bind(selected_schedule.thread_id.to_string()) - .bind(goal_id) - .fetch_one(&mut **goal_tx) - .await?; - if pause_for_goal_hold { - break; - } - } - } - if !active_goal_ids.is_empty() { - sqlx::query( - r#" -UPDATE thread_schedule_runs -SET status = 'failed', - error = ?, - completed_at_ms = ? -WHERE schedule_id = ? AND status IN ('leased', 'running') - "#, - ) - .bind(redact_state_string( - "scheduled run lease expired before terminal completion", - )) - .bind(now_ms) - .bind(selected_schedule.schedule_id.as_str()) - .execute(&mut *tx) - .await?; - } - if pause_for_goal_hold { - sqlx::query( - r#" -UPDATE thread_schedules -SET status = 'paused', - next_run_at_ms = NULL, - last_run_at_ms = ?, - failure_count = failure_count + 1, - lease_id = NULL, - lease_expires_at_ms = NULL, - updated_at_ms = ? -WHERE schedule_id = ? AND status = 'active' - "#, - ) - .bind(now_ms) - .bind(now_ms) - .bind(selected_schedule.schedule_id.as_str()) - .execute(&mut *tx) - .await?; - tx.commit().await?; - if let Some(goal_tx) = goal_tx { - let _ = goal_tx.rollback().await; - } - return Ok(None); - } - let sql = schedule_returning( - r#" -UPDATE thread_schedules -SET lease_id = ?, - lease_expires_at_ms = ?, - last_run_at_ms = CASE WHEN ? THEN ? ELSE last_run_at_ms END, - failure_count = CASE WHEN ? THEN failure_count + 1 ELSE failure_count END, - updated_at_ms = ? -WHERE schedule_id = ? AND status = 'active' -RETURNING -"#, - ); - let reaped_expired_run = !active_goal_ids.is_empty(); - let schedule_row = sqlx::query(sqlx::AssertSqlSafe(sql)) - .bind(lease_id) - .bind(lease_expires_at_ms) - .bind(reaped_expired_run) - .bind(now_ms) - .bind(reaped_expired_run) - .bind(now_ms) - .bind(selected_schedule.schedule_id.as_str()) - .fetch_optional(&mut *tx) - .await?; - let Some(schedule_row) = schedule_row else { - // `thread_schedules_ignore_legacy_live_owner_claim` silently drops - // the lease update (RAISE(IGNORE)) when a legacy, non-owner-scoped - // lease is claimed while a local session is live. Treat that as an - // unclaimed schedule and discard the speculative reap above so the - // live owner keeps ownership of its runs. - tx.rollback().await?; - if let Some(goal_tx) = goal_tx { - let _ = goal_tx.rollback().await; - } - return Ok(None); - }; - let schedule = thread_schedule_from_row(&schedule_row)?; - let scheduled_for_ms = match target { - ThreadScheduleClaimTarget::Due => schedule.next_run_at.map(datetime_to_epoch_millis), - ThreadScheduleClaimTarget::Now { .. } => Some(now_ms), - }; - let run = - Self::insert_leased_run(&mut tx, &schedule, lease_id, scheduled_for_ms, now_ms).await?; - tx.commit().await?; - if let Some(goal_tx) = goal_tx { - let _ = goal_tx.rollback().await; - } - Ok(Some(ThreadScheduleClaim { schedule, run })) - } - - async fn insert_leased_run( - tx: &mut sqlx::Transaction<'_, Sqlite>, - schedule: &crate::ThreadSchedule, - lease_id: &str, - scheduled_for_ms: Option, - started_at_ms: i64, - ) -> anyhow::Result { - let run_id = Uuid::new_v4().to_string(); - let run_row = sqlx::query( - r#" -INSERT INTO thread_schedule_runs ( - run_id, - schedule_id, - thread_id, - status, - lease_id, - scheduled_for_ms, - started_at_ms -) VALUES (?, ?, ?, ?, ?, ?, ?) -RETURNING - run_id, - schedule_id, - thread_id, - status, - lease_id, - turn_id, - goal_id, - error, - scheduled_for_ms, - started_at_ms, - completed_at_ms - "#, - ) - .bind(run_id) - .bind(schedule.schedule_id.as_str()) - .bind(schedule.thread_id.to_string()) - .bind(crate::ThreadScheduleRunStatus::Leased.as_str()) - .bind(lease_id) - .bind(scheduled_for_ms) - .bind(started_at_ms) - .fetch_one(&mut **tx) - .await?; - thread_schedule_run_from_row(&run_row) - } - - pub async fn mark_thread_schedule_run_started( - &self, - params: ThreadScheduleRunStartParams<'_>, - ) -> anyhow::Result> { - crate::busy_retry::retry_on_busy("mark thread schedule run started", || { - self.mark_thread_schedule_run_started_once(params.clone()) - }) - .await - } - - async fn mark_thread_schedule_run_started_once( - &self, - params: ThreadScheduleRunStartParams<'_>, - ) -> anyhow::Result> { - let ThreadScheduleRunStartParams { - schedule_id, - run_id, - lease_id, - turn_id, - goal_id, - now, - lease_duration, - } = params; - let now_ms = datetime_to_epoch_millis(now); - let lease_expires_at_ms = - datetime_to_epoch_millis(now + chrono::Duration::from_std(lease_duration)?); - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; - let schedule_result = sqlx::query( - r#" -UPDATE thread_schedules -SET lease_expires_at_ms = MAX(lease_expires_at_ms, ?), - updated_at_ms = ? -WHERE schedule_id = ? - AND lease_id = ? - AND lease_expires_at_ms > ? - AND EXISTS ( - SELECT 1 - FROM thread_schedule_runs - WHERE thread_schedule_runs.schedule_id = thread_schedules.schedule_id - AND thread_schedule_runs.run_id = ? - AND thread_schedule_runs.lease_id = ? - AND thread_schedule_runs.status = 'leased' - ) - "#, - ) - .bind(lease_expires_at_ms) - .bind(now_ms) - .bind(schedule_id) - .bind(lease_id) - .bind(now_ms) - .bind(run_id) - .bind(lease_id) - .execute(&mut *tx) - .await?; - if schedule_result.rows_affected() == 0 { - tx.commit().await?; - return Ok(None); - } - let sql = run_returning( - r#" -UPDATE thread_schedule_runs -SET status = ?, turn_id = ?, goal_id = ? -WHERE schedule_id = ? AND run_id = ? AND lease_id = ? AND status = 'leased' -RETURNING -"#, - ); - let row = sqlx::query(sqlx::AssertSqlSafe(sql)) - .bind(crate::ThreadScheduleRunStatus::Running.as_str()) - .bind(turn_id) - .bind(goal_id) - .bind(schedule_id) - .bind(run_id) - .bind(lease_id) - .fetch_optional(&mut *tx) - .await?; - let Some(row) = row else { - tx.rollback().await?; - return Ok(None); - }; - let run = thread_schedule_run_from_row(&row)?; - tx.commit().await?; - Ok(Some(run)) - } - - pub async fn extend_thread_schedule_lease( - &self, - params: ThreadScheduleRunLeaseParams<'_>, - ) -> anyhow::Result { - let ThreadScheduleRunLeaseParams { - schedule_id, - run_id, - lease_id, - now, - lease_duration, - } = params; - let now_ms = datetime_to_epoch_millis(now); - let lease_expires_at = now + chrono::Duration::from_std(lease_duration)?; - let result = sqlx::query( - r#" -UPDATE thread_schedules -SET lease_expires_at_ms = ?, updated_at_ms = ? -WHERE schedule_id = ? - AND status = 'active' - AND lease_id = ? - AND lease_expires_at_ms > ? - AND (expires_at_ms IS NULL OR expires_at_ms > ?) - AND EXISTS ( - SELECT 1 - FROM thread_schedule_runs - WHERE thread_schedule_runs.schedule_id = thread_schedules.schedule_id - AND thread_schedule_runs.run_id = ? - AND thread_schedule_runs.lease_id = ? - AND thread_schedule_runs.status IN ('leased', 'running') - ) - "#, - ) - .bind(datetime_to_epoch_millis(lease_expires_at)) - .bind(now_ms) - .bind(schedule_id) - .bind(lease_id) - .bind(now_ms) - .bind(now_ms) - .bind(run_id) - .bind(lease_id) - .execute(self.pool.as_ref()) - .await?; - Ok(result.rows_affected() > 0) - } - - pub async fn complete_thread_schedule_run( - &self, - schedule_id: &str, - run_id: &str, - lease_id: &str, - completed_at: DateTime, - next_run_at: Option>, - ) -> anyhow::Result { - self.finish_thread_schedule_run(FinishThreadScheduleRunParams { - schedule_id, - run_id, - lease_id, - completed_at, - next_run_at, - expected_goal_id: None, - finish: FinishScheduleRun::Completed, - }) - .await - } - - pub async fn complete_thread_schedule_run_for_goal( - &self, - params: ThreadScheduleRunForGoalFinishParams<'_>, - ) -> anyhow::Result { - let ThreadScheduleRunForGoalFinishParams { - schedule_id, - run_id, - lease_id, - completed_at, - next_run_at, - expected_goal_id, - } = params; - self.finish_thread_schedule_run(FinishThreadScheduleRunParams { - schedule_id, - run_id, - lease_id, - completed_at, - next_run_at, - expected_goal_id: Some(expected_goal_id), - finish: FinishScheduleRun::Completed, - }) - .await - } - - pub async fn fail_thread_schedule_run( - &self, - schedule_id: &str, - run_id: &str, - lease_id: &str, - completed_at: DateTime, - next_run_at: Option>, - error: String, - ) -> anyhow::Result { - self.finish_thread_schedule_run(FinishThreadScheduleRunParams { - schedule_id, - run_id, - lease_id, - completed_at, - next_run_at, - expected_goal_id: None, - finish: FinishScheduleRun::Failed { error }, - }) - .await - } - - pub async fn fail_thread_schedule_run_for_goal( - &self, - params: ThreadScheduleRunForGoalFinishParams<'_>, - error: String, - ) -> anyhow::Result { - let ThreadScheduleRunForGoalFinishParams { - schedule_id, - run_id, - lease_id, - completed_at, - next_run_at, - expected_goal_id, - } = params; - self.finish_thread_schedule_run(FinishThreadScheduleRunParams { - schedule_id, - run_id, - lease_id, - completed_at, - next_run_at, - expected_goal_id: Some(expected_goal_id), - finish: FinishScheduleRun::Failed { error }, - }) - .await - } - - pub async fn defer_thread_schedule_run( - &self, - schedule_id: &str, - run_id: &str, - lease_id: &str, - completed_at: DateTime, - next_run_at: DateTime, - error: String, - ) -> anyhow::Result { - let completed_at_ms = datetime_to_epoch_millis(completed_at); - let requested_next_run_at_ms = datetime_to_epoch_millis(next_run_at); - let mut tx = self.pool.begin().await?; - let schedule_result = sqlx::query( - r#" -UPDATE thread_schedules -SET - status = CASE - WHEN status = 'expired' THEN 'expired' - WHEN expires_at_ms IS NOT NULL AND ? >= expires_at_ms THEN 'expired' - WHEN status = 'paused' THEN 'paused' - ELSE status - END, - lease_id = NULL, - lease_expires_at_ms = NULL, - last_run_at_ms = ?, - next_run_at_ms = CASE - WHEN status IN ('expired', 'paused') THEN NULL - WHEN expires_at_ms IS NOT NULL AND ? >= expires_at_ms THEN NULL - ELSE ? - END, - updated_at_ms = ? -WHERE schedule_id = ? AND lease_id = ? - "#, - ) - .bind(requested_next_run_at_ms) - .bind(completed_at_ms) - .bind(requested_next_run_at_ms) - .bind(requested_next_run_at_ms) - .bind(completed_at_ms) - .bind(schedule_id) - .bind(lease_id) - .execute(&mut *tx) - .await?; - if schedule_result.rows_affected() == 0 { - tx.commit().await?; - return Ok(false); - } - let run_result = sqlx::query( - r#" -UPDATE thread_schedule_runs -SET status = ?, turn_id = NULL, error = ?, completed_at_ms = ? -WHERE schedule_id = ? AND run_id = ? AND lease_id = ? - "#, - ) - .bind(crate::ThreadScheduleRunStatus::Deferred.as_str()) - .bind(redact_state_string(error)) - .bind(completed_at_ms) - .bind(schedule_id) - .bind(run_id) - .bind(lease_id) - .execute(&mut *tx) - .await?; - if run_result.rows_affected() == 0 { - tx.rollback().await?; - return Ok(false); - } - tx.commit().await?; - Ok(true) - } - - pub async fn expire_thread_schedules(&self, now: DateTime) -> anyhow::Result { - let now_ms = datetime_to_epoch_millis(now); - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; - sqlx::query( - r#" -UPDATE thread_schedule_runs -SET status = 'failed', - error = ?, - completed_at_ms = COALESCE(completed_at_ms, ?) -WHERE status IN ('leased', 'running') - AND EXISTS ( - SELECT 1 - FROM thread_schedules - WHERE thread_schedules.schedule_id = thread_schedule_runs.schedule_id - AND thread_schedules.status = 'active' - AND thread_schedules.expires_at_ms IS NOT NULL - AND thread_schedules.expires_at_ms <= ? - AND ( - thread_schedules.lease_id IS NULL - OR thread_schedules.lease_expires_at_ms <= ? - ) - ) - "#, - ) - .bind(redact_state_string(EXPIRED_SCHEDULE_RUN_ERROR)) - .bind(now_ms) - .bind(now_ms) - .bind(now_ms) - .execute(&mut *tx) - .await?; - let result = sqlx::query( - r#" -UPDATE thread_schedules -SET - status = 'expired', - next_run_at_ms = NULL, - lease_id = NULL, - lease_expires_at_ms = NULL, - updated_at_ms = ? -WHERE status = 'active' - AND expires_at_ms IS NOT NULL - AND expires_at_ms <= ? - AND (lease_id IS NULL OR lease_expires_at_ms <= ?) - "#, - ) - .bind(now_ms) - .bind(now_ms) - .bind(now_ms) - .execute(&mut *tx) - .await?; - tx.commit().await?; - Ok(result.rows_affected()) - } - - async fn finish_thread_schedule_run( - &self, - params: FinishThreadScheduleRunParams<'_>, - ) -> anyhow::Result { - crate::busy_retry::retry_on_busy("finish thread schedule run", || { - self.finish_thread_schedule_run_once(params.clone()) + .await? + .ok_or_else(|| { + anyhow::anyhow!( + "invalid nested loop: parent schedule not found: {parent_schedule_id}" + ) + })?; + self.validate_parent_schedule(&parent, params.thread_id, ¶ms.schedule)?; + Ok(ScheduleNesting { + parent_schedule_id: Some(parent.schedule_id), + nesting_depth: parent.nesting_depth + 1, }) - .await } - async fn finish_thread_schedule_run_once( + async fn validate_schedule_update_nesting( &self, - params: FinishThreadScheduleRunParams<'_>, - ) -> anyhow::Result { - let FinishThreadScheduleRunParams { - schedule_id, - run_id, - lease_id, - completed_at, - next_run_at, - expected_goal_id, - finish, - } = params; - let completed_at_ms = datetime_to_epoch_millis(completed_at); - let next_run_at_ms = next_run_at.map(datetime_to_epoch_millis); - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; - let schedule_context: Option<(String, String)> = sqlx::query_as( - r#" -SELECT thread_id, schedule_kind -FROM thread_schedules -WHERE schedule_id = ? AND lease_id = ? - AND EXISTS ( - SELECT 1 - FROM thread_schedule_runs - WHERE thread_schedule_runs.schedule_id = thread_schedules.schedule_id - AND thread_schedule_runs.run_id = ? - AND thread_schedule_runs.lease_id = ? - AND thread_schedule_runs.status IN ('leased', 'running') - AND (? IS NULL OR thread_schedule_runs.goal_id IS NULL OR thread_schedule_runs.goal_id = ?) - ) - "#, - ) - .bind(schedule_id) - .bind(lease_id) - .bind(run_id) - .bind(lease_id) - .bind(expected_goal_id) - .bind(expected_goal_id) - .fetch_optional(&mut *tx) - .await?; - let Some((thread_id, schedule_kind)) = schedule_context else { - tx.commit().await?; - return Ok(false); - }; - let goal_hold_can_pause = expected_goal_id.is_some() && schedule_kind != ONCE_SCHEDULE_KIND; - // Read-only probe: this transaction only runs `SELECT EXISTS` against - // goals.db and is always rolled back, so a deferred `BEGIN` is enough. A - // `BEGIN IMMEDIATE` would take a goals.db write lock and hold it across the - // state.db commit for no benefit. Lock order is consistently state -> goals - // at every site that touches both, so there is no inversion to guard against. - let mut goal_tx = if goal_hold_can_pause { - Some(self.goals_pool.begin().await?) - } else { - None + existing: &crate::ThreadSchedule, + schedule: &crate::ThreadScheduleSpec, + ) -> anyhow::Result<()> { + if self + .has_child_thread_schedules(existing.schedule_id.as_str()) + .await? + { + anyhow::bail!( + "invalid nested loop: cannot update loop cadence while it has nested child loops; update or clear child loops first" + ); + } + let Some(parent_schedule_id) = existing.parent_schedule_id.as_deref() else { + return Ok(()); }; - let pause_for_goal_hold = match (expected_goal_id, goal_hold_can_pause, goal_tx.as_mut()) { - (Some(expected_goal_id), true, Some(goal_tx)) => { - sqlx::query_scalar::<_, bool>( - r#" -SELECT EXISTS( - SELECT 1 - FROM thread_goals - WHERE thread_id = ? - AND goal_id = ? - AND status IN ('paused', 'blocked', 'usage_limited', 'budget_limited') -) - "#, + if matches!(schedule, crate::ThreadScheduleSpec::Once) { + anyhow::bail!("invalid nested loop: one-time schedules cannot be nested"); + } + let parent = self + .get_thread_schedule(parent_schedule_id) + .await? + .ok_or_else(|| { + anyhow::anyhow!( + "invalid nested loop: parent schedule not found: {parent_schedule_id}" ) - .bind(thread_id) - .bind(expected_goal_id) - .fetch_one(&mut **goal_tx) - .await? - } - (Some(_), false, None) | (None, false, None) => false, - // `goal_hold_can_pause` is what decides whether `goal_tx` was opened, so - // the arms above are exhaustive in practice. Fail the write instead of - // panicking out of a state-store transaction if that ever drifts. - (expected_goal_id, goal_hold_can_pause, goal_tx) => { - anyhow::bail!( - "goal transaction presence does not match the recurring goal schedule invariant (expected_goal_id={}, goal_hold_can_pause={goal_hold_can_pause}, goal_tx={})", - expected_goal_id.is_some(), - goal_tx.is_some(), - ); - } - }; - let failed = matches!(finish, FinishScheduleRun::Failed { .. }); - // The only thing that pauses a schedule at finish time is a goal hold; there - // is deliberately no caller-supplied pause flag. - let pause_schedule = pause_for_goal_hold; - let schedule_result = sqlx::query( - r#" -UPDATE thread_schedules -SET - status = CASE - WHEN status = 'expired' THEN 'expired' - WHEN expires_at_ms IS NOT NULL AND ? >= expires_at_ms THEN 'expired' - WHEN status = 'paused' THEN 'paused' - WHEN ? THEN 'paused' - WHEN ? IS NULL THEN 'expired' - ELSE status - END, - lease_id = NULL, - lease_expires_at_ms = NULL, - last_run_at_ms = ?, - next_run_at_ms = CASE - WHEN status IN ('expired', 'paused') THEN NULL - WHEN expires_at_ms IS NOT NULL AND ? >= expires_at_ms THEN NULL - WHEN ? THEN NULL - WHEN ? IS NULL THEN NULL - ELSE ? - END, - failure_count = CASE WHEN ? THEN failure_count + 1 ELSE 0 END, - updated_at_ms = ? -WHERE schedule_id = ? AND lease_id = ? - "#, - ) - .bind(completed_at_ms) - .bind(pause_schedule) - .bind(next_run_at_ms) - .bind(completed_at_ms) - .bind(completed_at_ms) - .bind(pause_schedule) - .bind(next_run_at_ms) - .bind(next_run_at_ms) - .bind(failed) - .bind(completed_at_ms) - .bind(schedule_id) - .bind(lease_id) - .execute(&mut *tx) - .await?; - if schedule_result.rows_affected() == 0 { - tx.commit().await?; - if let Some(goal_tx) = goal_tx { - let _ = goal_tx.rollback().await; - } - return Ok(false); + })?; + self.validate_parent_schedule(&parent, existing.thread_id, schedule) + } + + fn validate_parent_schedule( + &self, + parent: &crate::ThreadSchedule, + thread_id: ThreadId, + child_schedule: &crate::ThreadScheduleSpec, + ) -> anyhow::Result<()> { + if parent.thread_id != thread_id { + anyhow::bail!("invalid nested loop: parent schedule must belong to the same thread"); } - let (status, error) = match &finish { - FinishScheduleRun::Completed => (crate::ThreadScheduleRunStatus::Completed, None), - FinishScheduleRun::Failed { error, .. } => { - (crate::ThreadScheduleRunStatus::Failed, Some(error.as_str())) - } - }; - let error = error.map(redact_state_string); - let run_result = sqlx::query( + if matches!(parent.schedule, crate::ThreadScheduleSpec::Once) { + anyhow::bail!("invalid nested loop: parent schedule must be recurring"); + } + if parent.nesting_depth >= MAX_THREAD_SCHEDULE_NESTING_DEPTH { + anyhow::bail!( + "invalid nested loop: maximum nesting depth is {MAX_THREAD_SCHEDULE_NESTING_DEPTH}" + ); + } + validate_nested_loop_cadence(&parent.schedule, child_schedule) + } + + async fn has_child_thread_schedules(&self, schedule_id: &str) -> anyhow::Result { + let count: i64 = sqlx::query_scalar( r#" -UPDATE thread_schedule_runs -SET status = ?, error = ?, completed_at_ms = ? -WHERE schedule_id = ? AND run_id = ? AND lease_id = ? +SELECT COUNT(*) +FROM thread_schedules +WHERE parent_schedule_id = ? "#, ) - .bind(status.as_str()) - .bind(error) - .bind(completed_at_ms) .bind(schedule_id) - .bind(run_id) - .bind(lease_id) - .execute(&mut *tx) + .fetch_one(self.pool.as_ref()) .await?; - if run_result.rows_affected() == 0 { - tx.rollback().await?; - if let Some(goal_tx) = goal_tx { - let _ = goal_tx.rollback().await; - } - return Ok(false); - } - tx.commit().await?; - if let Some(goal_tx) = goal_tx { - let _ = goal_tx.rollback().await; - } - Ok(true) + Ok(count > 0) } } -#[derive(Clone)] -struct FinishThreadScheduleRunParams<'a> { - schedule_id: &'a str, - run_id: &'a str, - lease_id: &'a str, - completed_at: DateTime, - next_run_at: Option>, - expected_goal_id: Option<&'a str>, - finish: FinishScheduleRun, -} - -#[derive(Clone)] -enum FinishScheduleRun { - Completed, - Failed { error: String }, -} - struct ScheduleBindings<'a> { kind: &'static str, interval_amount: Option, @@ -1890,6 +829,50 @@ mod tests { .expect("schedule should be created") } + async fn enqueue_and_start_claim( + runtime: &StateRuntime, + claim: &ThreadScheduleClaim, + goal_id: Option<&str>, + turn_input: &str, + now: DateTime, + lease_duration: Duration, + ) -> crate::ThreadScheduleRun { + let turn_id = claim + .run + .turn_id + .as_deref() + .expect("claimed occurrence should have a stable turn id"); + runtime + .thread_schedules() + .enqueue_thread_schedule_run(ThreadScheduleRunEnqueueParams { + schedule_id: claim.schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + goal_id, + auth_profile_recorded: false, + auth_profile: None, + turn_input, + now, + }) + .await + .expect("claimed occurrence should enqueue") + .expect("claimed occurrence should retain its lease"); + runtime + .thread_schedules() + .mark_thread_schedule_run_started(ThreadScheduleRunStartParams { + schedule_id: claim.schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + turn_id, + goal_id, + now, + lease_duration, + }) + .await + .expect("enqueued occurrence should start") + .expect("enqueued occurrence should retain its lease") + } + #[tokio::test] async fn create_update_list_and_delete_thread_schedule() { let runtime = test_runtime().await; @@ -2618,20 +1601,15 @@ mod tests { .await .expect("claim should succeed") .expect("one-time schedule should claim"); - runtime - .thread_schedules() - .mark_thread_schedule_run_started(ThreadScheduleRunStartParams { - schedule_id: &schedule.schedule_id, - run_id: &claim.run.run_id, - lease_id: "lease-once", - turn_id: "turn-once", - goal_id: None, - now, - lease_duration: Duration::from_secs(300), - }) - .await - .expect("run should update") - .expect("run should exist"); + enqueue_and_start_claim( + &runtime, + &claim, + None, + "once input", + now, + Duration::from_secs(300), + ) + .await; assert!( runtime @@ -2692,229 +1670,52 @@ mod tests { ) .await; create_interval_schedule( - &runtime, - thread_id, - "future task", - Some(now + chrono::Duration::minutes(1)), - ) - .await; - - let first_claim = runtime - .thread_schedules() - .claim_due_thread_schedule(now, "lease-a", Duration::from_secs(300)) - .await - .expect("claim should succeed") - .expect("first due schedule should claim"); - assert_eq!(first.schedule_id, first_claim.schedule.schedule_id); - assert_eq!(Some("lease-a".to_string()), first_claim.schedule.lease_id); - assert_eq!( - crate::ThreadScheduleRunStatus::Leased, - first_claim.run.status - ); - assert_eq!("lease-a", first_claim.run.lease_id); - assert_eq!( - Some(now - chrono::Duration::minutes(2)), - first_claim.run.scheduled_for - ); - - let second_claim = runtime - .thread_schedules() - .claim_due_thread_schedule(now, "lease-b", Duration::from_secs(300)) - .await - .expect("claim should succeed") - .expect("second due schedule should claim"); - assert_eq!(second.schedule_id, second_claim.schedule.schedule_id); - - assert!( - runtime - .thread_schedules() - .claim_due_thread_schedule(now, "lease-c", Duration::from_secs(300)) - .await - .expect("no more schedules should be claimable") - .is_none() - ); - } - - #[tokio::test] - async fn claim_due_thread_schedule_reaps_expired_run_before_retrying_once() { - let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); - let thread_id = test_thread_id(/*id*/ 44); - upsert_test_thread(runtime.as_ref(), thread_id).await; - let now = at(/*seconds*/ 1_700_000_000); - let schedule = - create_interval_schedule(runtime.as_ref(), thread_id, "restart retry", Some(now)).await; - let original_claim = runtime - .thread_schedules() - .claim_due_thread_schedule(now, "lease-before-restart", Duration::from_secs(30)) - .await - .expect("initial claim should succeed") - .expect("schedule should claim"); - runtime - .thread_schedules() - .mark_thread_schedule_run_started(ThreadScheduleRunStartParams { - schedule_id: &schedule.schedule_id, - run_id: &original_claim.run.run_id, - lease_id: "lease-before-restart", - turn_id: "turn-before-restart", - goal_id: None, - now, - lease_duration: Duration::from_secs(30), - }) - .await - .expect("run should start") - .expect("run should still exist"); - drop(runtime); - - let reopened = StateRuntime::init(codex_home, "test-provider".to_string()) - .await - .expect("state db should reopen after process restart"); - let retry_at = now + chrono::Duration::seconds(31); - let retry_claim = reopened - .thread_schedules() - .claim_due_thread_schedule(retry_at, "lease-after-restart", Duration::from_secs(30)) - .await - .expect("expired run recovery should succeed") - .expect("expired non-goal run should retry exactly once"); - - let original_run = reopened - .thread_schedules() - .get_thread_schedule_run(&original_claim.run.run_id) - .await - .expect("original run should load") - .expect("original run should exist"); - assert_eq!(crate::ThreadScheduleRunStatus::Failed, original_run.status); - assert_eq!(Some(retry_at), original_run.completed_at); - assert_eq!( - Some("scheduled run lease expired before terminal completion".to_string()), - original_run.error - ); - assert_eq!( - crate::ThreadScheduleRunStatus::Leased, - retry_claim.run.status - ); - assert_eq!( - original_claim.run.scheduled_for, - retry_claim.run.scheduled_for - ); - assert_ne!(original_claim.run.run_id, retry_claim.run.run_id); - let stats = reopened - .thread_schedules() - .get_thread_schedule_stats(&schedule.schedule_id) - .await - .expect("schedule stats should load"); - assert_eq!(2, stats.total_runs); - assert_eq!(1, stats.leased_runs); - assert_eq!(0, stats.running_runs); - assert_eq!(1, stats.failed_runs); - assert!( - reopened - .thread_schedules() - .claim_due_thread_schedule( - retry_at, - "lease-duplicate-retry", - Duration::from_secs(30), - ) - .await - .expect("duplicate claim check should succeed") - .is_none(), - "one expired lease may create at most one replacement claim" - ); - } - - #[tokio::test] - async fn claim_due_thread_schedule_pauses_expired_run_for_held_goal_after_restart() { - let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) - .await - .expect("state db should initialize"); - let thread_id = test_thread_id(/*id*/ 45); - upsert_test_thread(runtime.as_ref(), thread_id).await; - let goal = runtime - .thread_goals() - .replace_thread_goal( - thread_id, - "hold after restart", - crate::ThreadGoalStatus::Blocked, - /*token_budget*/ None, - ) - .await - .expect("blocked goal should persist"); - let now = at(/*seconds*/ 1_700_000_000); - let schedule = - create_interval_schedule(runtime.as_ref(), thread_id, "hold after restart", Some(now)) - .await; - let original_claim = runtime + &runtime, + thread_id, + "future task", + Some(now + chrono::Duration::minutes(1)), + ) + .await; + + let first_claim = runtime .thread_schedules() - .claim_due_thread_schedule(now, "lease-goal-restart", Duration::from_secs(30)) + .claim_due_thread_schedule(now, "lease-a", Duration::from_secs(300)) .await - .expect("initial claim should succeed") - .expect("schedule should claim"); - runtime + .expect("claim should succeed") + .expect("first due schedule should claim"); + assert_eq!(first.schedule_id, first_claim.schedule.schedule_id); + assert_eq!(Some("lease-a".to_string()), first_claim.schedule.lease_id); + assert_eq!( + crate::ThreadScheduleRunStatus::Leased, + first_claim.run.status + ); + assert_eq!("lease-a", first_claim.run.lease_id); + assert_eq!( + Some(now - chrono::Duration::minutes(2)), + first_claim.run.scheduled_for + ); + + let second_claim = runtime .thread_schedules() - .mark_thread_schedule_run_started(ThreadScheduleRunStartParams { - schedule_id: &schedule.schedule_id, - run_id: &original_claim.run.run_id, - lease_id: "lease-goal-restart", - turn_id: "turn-goal-restart", - goal_id: Some(&goal.goal_id), - now, - lease_duration: Duration::from_secs(30), - }) + .claim_due_thread_schedule(now, "lease-b", Duration::from_secs(300)) .await - .expect("goal run should start") - .expect("goal run should still exist"); - drop(runtime); + .expect("claim should succeed") + .expect("second due schedule should claim"); + assert_eq!(second.schedule_id, second_claim.schedule.schedule_id); - let reopened = StateRuntime::init(codex_home, "test-provider".to_string()) - .await - .expect("state db should reopen after process restart"); - let retry_at = now + chrono::Duration::seconds(31); assert!( - reopened + runtime .thread_schedules() - .claim_due_thread_schedule( - retry_at, - "lease-held-replacement", - Duration::from_secs(30), - ) + .claim_due_thread_schedule(now, "lease-c", Duration::from_secs(300)) .await - .expect("expired goal run recovery should succeed") - .is_none(), - "a persisted held goal must pause instead of creating a replacement run" + .expect("no more schedules should be claimable") + .is_none() ); - - let held_schedule = reopened - .thread_schedules() - .get_thread_schedule(&schedule.schedule_id) - .await - .expect("schedule should load") - .expect("schedule should exist"); - assert_eq!(crate::ThreadScheduleStatus::Paused, held_schedule.status); - assert_eq!(None, held_schedule.next_run_at); - assert_eq!(None, held_schedule.lease_id); - let original_run = reopened - .thread_schedules() - .get_thread_schedule_run(&original_claim.run.run_id) - .await - .expect("original run should load") - .expect("original run should exist"); - assert_eq!(crate::ThreadScheduleRunStatus::Failed, original_run.status); - assert_eq!(Some(goal.goal_id), original_run.goal_id); - assert_eq!(Some(retry_at), original_run.completed_at); - let stats = reopened - .thread_schedules() - .get_thread_schedule_stats(&schedule.schedule_id) - .await - .expect("schedule stats should load"); - assert_eq!(1, stats.total_runs); - assert_eq!(0, stats.leased_runs); - assert_eq!(0, stats.running_runs); - assert_eq!(1, stats.failed_runs); } + #[path = "occurrence_tests.rs"] + mod occurrence_tests; + #[tokio::test] async fn claim_due_thread_schedule_skips_fresh_foreign_active_owner() { let runtime = test_runtime().await; @@ -3221,6 +2022,15 @@ mod tests { .expect("claim should succeed") .expect("schedule should claim"); assert_eq!(schedule.schedule_id, claim.schedule.schedule_id); + enqueue_and_start_claim( + &runtime, + &claim, + None, + "long running input", + now, + Duration::from_secs(300), + ) + .await; assert!( runtime @@ -3274,20 +2084,15 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); - runtime - .thread_schedules() - .mark_thread_schedule_run_started(ThreadScheduleRunStartParams { - schedule_id: &schedule.schedule_id, - run_id: &claim.run.run_id, - lease_id: "lease-stale", - turn_id: "turn-stale", - goal_id: None, - now, - lease_duration: Duration::from_secs(30), - }) - .await - .expect("run start should persist") - .expect("run should start"); + enqueue_and_start_claim( + &runtime, + &claim, + None, + "stale input", + now, + Duration::from_secs(30), + ) + .await; let expired_at = now + chrono::Duration::seconds(31); assert!( @@ -3313,13 +2118,25 @@ mod tests { .expect("schedule should exist") .lease_expires_at ); - let replacement = runtime + let recovered = runtime .thread_schedules() .claim_due_thread_schedule(expired_at, "lease-new", Duration::from_secs(30)) .await .expect("reaper should not error") - .expect("expired run should be replaceable"); - assert_ne!(claim.run.run_id, replacement.run.run_id); + .expect("expired run should be recoverable"); + assert_eq!(claim.run.run_id, recovered.run.run_id); + assert_eq!(claim.run.turn_id, recovered.run.turn_id); + assert_eq!( + crate::ThreadScheduleRunStatus::Running, + recovered.run.status + ); + let stats = runtime + .thread_schedules() + .get_thread_schedule_stats(&schedule.schedule_id) + .await + .expect("schedule stats should load"); + assert_eq!(1, stats.total_runs); + assert_eq!(1, stats.running_runs); } #[tokio::test] @@ -3345,6 +2162,15 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); + enqueue_and_start_claim( + &runtime, + &claim, + Some(goal.goal_id.as_str()), + "held task input", + now, + Duration::from_secs(300), + ) + .await; // The caller still asks for a rearm; the goal hold must win. assert!( @@ -3419,6 +2245,15 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); + enqueue_and_start_claim( + &runtime, + &claim, + Some(original_goal.goal_id.as_str()), + "original goal input", + now, + Duration::from_secs(300), + ) + .await; let replacement_goal = runtime .thread_goals() .replace_thread_goal( @@ -3493,6 +2328,15 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); + enqueue_and_start_claim( + &runtime, + &claim, + Some(goal.goal_id.as_str()), + "once goal input", + now, + Duration::from_secs(300), + ) + .await; assert!( runtime @@ -3552,6 +2396,15 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); + enqueue_and_start_claim( + &runtime, + &claim, + Some(goal.goal_id.as_str()), + "contended goal input", + now, + Duration::from_secs(300), + ) + .await; let contender_state_pool = SqlitePoolOptions::new() .max_connections(1) @@ -3648,20 +2501,15 @@ mod tests { .await .expect("initial claim should succeed") .expect("schedule should claim"); - runtime - .thread_schedules() - .mark_thread_schedule_run_started(ThreadScheduleRunStartParams { - schedule_id: &schedule.schedule_id, - run_id: &claim.run.run_id, - lease_id: "lease-terminal-race", - turn_id: "turn-terminal-race", - goal_id: None, - now, - lease_duration: Duration::from_secs(30), - }) - .await - .expect("run should start") - .expect("run should still exist"); + enqueue_and_start_claim( + &runtime, + &claim, + None, + "terminal race input", + now, + Duration::from_secs(30), + ) + .await; let contender = StateRuntime::init(codex_home, "test-provider".to_string()) .await @@ -3694,21 +2542,27 @@ mod tests { .await .expect("original run should load") .expect("original run should exist"); - assert!( - matches!( - original_run.status, - crate::ThreadScheduleRunStatus::Completed | crate::ThreadScheduleRunStatus::Failed - ), - "the old running row must be terminal after the race" + assert_eq!( + if completion { + crate::ThreadScheduleRunStatus::Completed + } else { + crate::ThreadScheduleRunStatus::Running + }, + original_run.status, + "the same run is either terminalized or reclaimed without replacement" ); + if let Some(replacement) = replacement.as_ref() { + assert_eq!(claim.run.run_id, replacement.run.run_id); + assert_eq!(claim.run.turn_id, replacement.run.turn_id); + } let stats = runtime .thread_schedules() .get_thread_schedule_stats(&schedule.schedule_id) .await .expect("schedule stats should load"); - assert_eq!(0, stats.running_runs); - assert_eq!(i64::from(replacement.is_some()), stats.leased_runs); - assert_eq!(if replacement.is_some() { 2 } else { 1 }, stats.total_runs); + assert_eq!(i64::from(replacement.is_some()), stats.running_runs); + assert_eq!(0, stats.leased_runs); + assert_eq!(1, stats.total_runs); } #[tokio::test] @@ -3756,37 +2610,30 @@ mod tests { .is_none(), "a reaped expired run must never become dispatchable" ); - let replacement = replacement + let recovered = replacement .expect("expired lease reaper should not error") - .expect("expired run should produce one replacement claim"); - - let original_run = runtime - .thread_schedules() - .get_thread_schedule_run(&original_claim.run.run_id) - .await - .expect("original run should load") - .expect("original run should exist"); - assert_eq!(crate::ThreadScheduleRunStatus::Failed, original_run.status); - assert_eq!(Some(retry_at), original_run.completed_at); - assert_eq!( - Some("scheduled run lease expired before terminal completion".to_string()), - original_run.error + .expect("expired occurrence should be reclaimed"); + assert_eq!(original_claim.run.run_id, recovered.run.run_id); + assert_eq!(original_claim.run.turn_id, recovered.run.turn_id); + assert!( + runtime + .thread_schedules() + .get_thread_schedule_run(&original_claim.run.run_id) + .await + .expect("run lookup should succeed") + .is_none(), + "waiting occurrence must not create a run row" ); let replacement_started_at = retry_at + chrono::Duration::seconds(1); - let replacement_run = runtime - .thread_schedules() - .mark_thread_schedule_run_started(ThreadScheduleRunStartParams { - schedule_id: &schedule.schedule_id, - run_id: &replacement.run.run_id, - lease_id: "lease-replacement", - turn_id: "turn-replacement", - goal_id: None, - now: replacement_started_at, - lease_duration: Duration::from_secs(30), - }) - .await - .expect("replacement start should not error") - .expect("replacement should remain the sole dispatchable run"); + let replacement_run = enqueue_and_start_claim( + &runtime, + &recovered, + None, + "recovered waiting input", + replacement_started_at, + Duration::from_secs(30), + ) + .await; assert_eq!( crate::ThreadScheduleRunStatus::Running, replacement_run.status @@ -3796,7 +2643,7 @@ mod tests { .thread_schedules() .complete_thread_schedule_run( &schedule.schedule_id, - &replacement.run.run_id, + &recovered.run.run_id, "lease-replacement", replacement_started_at + chrono::Duration::seconds(1), Some(now + chrono::Duration::hours(1)), @@ -3810,11 +2657,11 @@ mod tests { .get_thread_schedule_stats(&schedule.schedule_id) .await .expect("schedule stats should load"); - assert_eq!(2, stats.total_runs); + assert_eq!(1, stats.total_runs); assert_eq!(0, stats.leased_runs); assert_eq!(0, stats.running_runs); assert_eq!(1, stats.completed_runs); - assert_eq!(1, stats.failed_runs); + assert_eq!(0, stats.failed_runs); } #[tokio::test] @@ -3841,6 +2688,15 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); + enqueue_and_start_claim( + &runtime, + &claim, + Some(goal.goal_id.as_str()), + "expired active input", + now, + Duration::from_secs(300), + ) + .await; let expired = runtime .thread_schedules() @@ -3924,6 +2780,15 @@ mod tests { .await .expect("complete schedule should claim") .expect("complete schedule should be due"); + enqueue_and_start_claim( + &runtime, + &complete_claim, + None, + "late complete input", + now, + Duration::from_secs(300), + ) + .await; runtime .thread_schedules() .set_thread_schedule_status(&complete_schedule.schedule_id, held_status) @@ -3977,6 +2842,15 @@ mod tests { .await .expect("deferred schedule should claim") .expect("deferred schedule should be due"); + enqueue_and_start_claim( + &runtime, + &defer_claim, + None, + "late defer input", + now, + Duration::from_secs(300), + ) + .await; runtime .thread_schedules() .set_thread_schedule_status(&defer_schedule.schedule_id, held_status) @@ -4085,22 +2959,17 @@ mod tests { .expect("claim should succeed") .expect("schedule should claim"); - let running = runtime - .thread_schedules() - .mark_thread_schedule_run_started(ThreadScheduleRunStartParams { - schedule_id: &completed_schedule.schedule_id, - run_id: &completed_claim.run.run_id, - lease_id: "lease-complete", - turn_id: "turn-1", - goal_id: None, - now, - lease_duration: Duration::from_secs(300), - }) - .await - .expect("run should update") - .expect("run should exist"); + let running = enqueue_and_start_claim( + &runtime, + &completed_claim, + None, + "completed input", + now, + Duration::from_secs(300), + ) + .await; assert_eq!(crate::ThreadScheduleRunStatus::Running, running.status); - assert_eq!(Some("turn-1".to_string()), running.turn_id); + assert_eq!(completed_claim.run.turn_id, running.turn_id); let next_run_at = now + chrono::Duration::minutes(5); assert!( @@ -4175,6 +3044,15 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); + enqueue_and_start_claim( + &runtime, + &failed_claim, + None, + "failed input", + now, + Duration::from_secs(300), + ) + .await; assert!( runtime .thread_schedules() @@ -4253,6 +3131,15 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); + enqueue_and_start_claim( + &runtime, + &claim, + None, + "resume failure input", + now, + Duration::from_secs(300), + ) + .await; runtime .thread_schedules() .fail_thread_schedule_run( @@ -4308,6 +3195,15 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); + enqueue_and_start_claim( + &runtime, + &claim, + None, + "update failure input", + now, + Duration::from_secs(300), + ) + .await; runtime .thread_schedules() .fail_thread_schedule_run( @@ -4523,6 +3419,15 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); + enqueue_and_start_claim( + &runtime, + &claim_one, + None, + "completed stats input", + now, + Duration::from_secs(300), + ) + .await; assert!( runtime .thread_schedules() @@ -4572,6 +3477,15 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); + enqueue_and_start_claim( + &runtime, + &claim_three, + None, + "failed stats input", + third_run_at, + Duration::from_secs(300), + ) + .await; assert!( runtime .thread_schedules() @@ -4716,20 +3630,15 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); - runtime - .thread_schedules() - .mark_thread_schedule_run_started(ThreadScheduleRunStartParams { - schedule_id: &schedule.schedule_id, - run_id: &claim.run.run_id, - lease_id: "lease-live", - turn_id: "turn-live", - goal_id: None, - now, - lease_duration: Duration::from_secs(300), - }) - .await - .expect("run should update") - .expect("run should exist"); + enqueue_and_start_claim( + &runtime, + &claim, + None, + "live input", + now, + Duration::from_secs(300), + ) + .await; let after_expiry = now + chrono::Duration::seconds(20); assert_eq!( @@ -4832,20 +3741,15 @@ mod tests { .await .expect("paused schedule claim should succeed") .expect("paused schedule should claim"); - runtime - .thread_schedules() - .mark_thread_schedule_run_started(ThreadScheduleRunStartParams { - schedule_id: &paused_schedule.schedule_id, - run_id: &paused_claim.run.run_id, - lease_id: "lease-paused", - turn_id: "turn-paused", - goal_id: None, - now, - lease_duration: Duration::from_secs(300), - }) - .await - .expect("paused run start should persist") - .expect("paused run should start"); + enqueue_and_start_claim( + &runtime, + &paused_claim, + None, + "paused input", + now, + Duration::from_secs(300), + ) + .await; let paused = runtime .thread_schedules() .set_thread_schedule_status( @@ -4886,20 +3790,15 @@ mod tests { .await .expect("expiring schedule claim should succeed") .expect("expiring schedule should claim"); - runtime - .thread_schedules() - .mark_thread_schedule_run_started(ThreadScheduleRunStartParams { - schedule_id: &expiring_schedule.schedule_id, - run_id: &expiring_claim.run.run_id, - lease_id: "lease-expiring", - turn_id: "turn-expiring", - goal_id: None, - now, - lease_duration: Duration::from_secs(30), - }) - .await - .expect("expiring run start should persist") - .expect("expiring run should start"); + enqueue_and_start_claim( + &runtime, + &expiring_claim, + None, + "expiring input", + now, + Duration::from_secs(30), + ) + .await; let expired_at = now + chrono::Duration::seconds(31); assert_eq!( 1, diff --git a/codex-rs/state/src/runtime/schedules/occurrence.rs b/codex-rs/state/src/runtime/schedules/occurrence.rs new file mode 100644 index 000000000..ccf51978f --- /dev/null +++ b/codex-rs/state/src/runtime/schedules/occurrence.rs @@ -0,0 +1,159 @@ +//! Durable state machine for one scheduled occurrence. + +use super::*; + +mod claim; +mod finish; +mod start; + +#[derive(Clone)] +pub struct ThreadScheduleClaim { + pub schedule: crate::ThreadSchedule, + pub run: crate::ThreadScheduleRun, + pub occurrence_state: ThreadScheduleOccurrenceState, + pub turn_input: Option, + pub occurrence_auth_profile: Option>, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ThreadScheduleOccurrenceState { + WaitingIdle, + Enqueued, + Started, + Terminal, +} + +#[derive(Clone)] +pub struct ThreadScheduleDueClaimParams<'a> { + pub now: DateTime, + pub lease_id: &'a str, + pub lease_duration: Duration, + pub local_active_owner_id: Option<&'a str>, + pub local_active_fresh_after: Option>, +} + +#[derive(Clone)] +pub struct ThreadScheduleNowClaimParams<'a> { + pub schedule_id: &'a str, + pub now: DateTime, + pub lease_id: &'a str, + pub lease_duration: Duration, + pub local_active_owner_id: Option<&'a str>, + pub local_active_fresh_after: Option>, +} + +pub struct ThreadScheduleRunForGoalFinishParams<'a> { + pub schedule_id: &'a str, + pub run_id: &'a str, + pub lease_id: &'a str, + pub completed_at: DateTime, + pub next_run_at: Option>, + pub expected_goal_id: &'a str, +} + +#[derive(Clone)] +pub struct ThreadScheduleRunStartParams<'a> { + pub schedule_id: &'a str, + pub run_id: &'a str, + pub lease_id: &'a str, + pub turn_id: &'a str, + pub goal_id: Option<&'a str>, + pub now: DateTime, + pub lease_duration: Duration, +} + +#[derive(Clone)] +pub struct ThreadScheduleRunEnqueueParams<'a> { + pub schedule_id: &'a str, + pub run_id: &'a str, + pub lease_id: &'a str, + pub goal_id: Option<&'a str>, + pub auth_profile_recorded: bool, + pub auth_profile: Option<&'a str>, + pub turn_input: &'a str, + pub now: DateTime, +} + +#[derive(Clone)] +pub struct ThreadScheduleRunLeaseParams<'a> { + pub schedule_id: &'a str, + pub run_id: &'a str, + pub lease_id: &'a str, + pub now: DateTime, + pub lease_duration: Duration, +} + +#[derive(Clone, sqlx::FromRow)] +struct ThreadScheduleOccurrenceRow { + occurrence_id: String, + schedule_id: String, + thread_id: String, + state: String, + turn_id: String, + goal_id: Option, + auth_profile_recorded: bool, + auth_profile: Option, + scheduled_for_ms: Option, + turn_input: Option, + created_at_ms: i64, +} + +const OCCURRENCE_WAITING_IDLE: &str = "waiting_idle"; +const OCCURRENCE_ENQUEUED: &str = "enqueued"; +const OCCURRENCE_STARTED: &str = "started"; +const OCCURRENCE_TERMINAL: &str = "terminal"; + +impl ThreadScheduleOccurrenceState { + fn from_str(state: &str) -> anyhow::Result { + match state { + OCCURRENCE_WAITING_IDLE => Ok(Self::WaitingIdle), + OCCURRENCE_ENQUEUED => Ok(Self::Enqueued), + OCCURRENCE_STARTED => Ok(Self::Started), + OCCURRENCE_TERMINAL => Ok(Self::Terminal), + state => anyhow::bail!("unsupported thread schedule occurrence state {state}"), + } + } +} + +#[derive(Clone, Copy)] +enum ThreadScheduleClaimTarget<'a> { + Due, + Now { schedule_id: &'a str }, +} + +#[derive(Clone)] +struct ClaimThreadScheduleParams<'a> { + target: ThreadScheduleClaimTarget<'a>, + now: DateTime, + lease_id: &'a str, + lease_duration: Duration, + local_active_owner_id: Option<&'a str>, + local_active_fresh_after: Option>, +} + +#[derive(Clone)] +struct FinishThreadScheduleRunParams<'a> { + schedule_id: &'a str, + run_id: &'a str, + lease_id: &'a str, + completed_at: DateTime, + next_run_at: Option>, + expected_goal_id: Option<&'a str>, + finish: FinishScheduleRun, +} + +#[derive(Clone)] +struct FinalizeThreadScheduleRunParams<'a> { + schedule_id: &'a str, + run_id: &'a str, + lease_id: &'a str, + completed_at: DateTime, + next_run_at: Option>, + expected_goal_id: Option<&'a str>, +} + +#[derive(Clone)] +enum FinishScheduleRun { + Completed, + Failed { error: String }, +} diff --git a/codex-rs/state/src/runtime/schedules/occurrence/claim.rs b/codex-rs/state/src/runtime/schedules/occurrence/claim.rs new file mode 100644 index 000000000..94ca955e4 --- /dev/null +++ b/codex-rs/state/src/runtime/schedules/occurrence/claim.rs @@ -0,0 +1,564 @@ +//! Claim and restart recovery for pending scheduled occurrences. + +use super::*; + +impl ScheduleStore { + pub async fn get_thread_schedule_run( + &self, + run_id: &str, + ) -> anyhow::Result> { + let sql = run_returning( + r#" +SELECT + "#, + ); + let row = sqlx::query(sqlx::AssertSqlSafe(format!( + "{sql}FROM thread_schedule_runs WHERE run_id = ?" + ))) + .bind(run_id) + .fetch_optional(self.pool.as_ref()) + .await?; + row.map(|row| thread_schedule_run_from_row(&row)) + .transpose() + } + + pub async fn get_running_thread_schedule_run_for_turn( + &self, + thread_id: ThreadId, + turn_id: &str, + ) -> anyhow::Result> { + let sql = run_returning( + r#" +SELECT +"#, + ); + let row = sqlx::query(sqlx::AssertSqlSafe(format!( + r#"{sql} +FROM thread_schedule_runs +WHERE thread_id = ? + AND turn_id = ? + AND status = 'running' +ORDER BY started_at_ms DESC +LIMIT 1 +"# + ))) + .bind(thread_id.to_string()) + .bind(turn_id) + .fetch_optional(self.pool.as_ref()) + .await?; + row.map(|row| thread_schedule_run_from_row(&row)) + .transpose() + } + + pub async fn get_thread_schedule_stats( + &self, + schedule_id: &str, + ) -> anyhow::Result { + let row = sqlx::query( + r#" +SELECT + COALESCE(SUM(CASE WHEN deferral_kind IS NULL OR deferral_kind != 'idle' THEN 1 ELSE 0 END), 0) AS total_runs, + COALESCE(SUM(CASE WHEN status = 'leased' THEN 1 ELSE 0 END), 0) AS leased_runs, + COALESCE(SUM(CASE WHEN status = 'running' THEN 1 ELSE 0 END), 0) AS running_runs, + COALESCE(SUM(CASE WHEN status = 'deferred' AND deferral_kind = 'capacity' THEN 1 ELSE 0 END), 0) AS deferred_runs, + COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0) AS completed_runs, + COALESCE(SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END), 0) AS failed_runs, + MAX(CASE WHEN deferral_kind IS NULL OR deferral_kind != 'idle' THEN started_at_ms END) AS last_started_at_ms, + -- Only successfully completed runs contribute to last_completed_at. The + -- completed_at_ms column is also written for deferred and failed runs (it is + -- really a "finished at" timestamp), so deriving last_completed_at from the + -- raw MAX would populate it even when completed_runs is 0. Keeping this + -- filtered ensures last_completed_at is non-null iff completed_runs > 0. + MAX(CASE WHEN status = 'completed' THEN completed_at_ms END) AS last_completed_at_ms +FROM thread_schedule_runs +WHERE schedule_id = ? + "#, + ) + .bind(schedule_id) + .fetch_one(self.pool.as_ref()) + .await?; + let last_error = sqlx::query_scalar( + r#" +SELECT error +FROM thread_schedule_runs +WHERE schedule_id = ? + AND status = 'failed' + AND error IS NOT NULL + AND TRIM(error) != '' +ORDER BY completed_at_ms DESC, started_at_ms DESC +LIMIT 1 + "#, + ) + .bind(schedule_id) + .fetch_optional(self.pool.as_ref()) + .await?; + Ok(crate::ThreadScheduleStats { + total_runs: row.try_get("total_runs")?, + leased_runs: row.try_get("leased_runs")?, + running_runs: row.try_get("running_runs")?, + deferred_runs: row.try_get("deferred_runs")?, + completed_runs: row.try_get("completed_runs")?, + failed_runs: row.try_get("failed_runs")?, + last_started_at: row + .try_get::, _>("last_started_at_ms")? + .map(epoch_millis_to_datetime) + .transpose()?, + last_completed_at: row + .try_get::, _>("last_completed_at_ms")? + .map(epoch_millis_to_datetime) + .transpose()?, + last_error, + }) + } + + pub async fn claim_due_thread_schedule( + &self, + now: DateTime, + lease_id: &str, + lease_duration: Duration, + ) -> anyhow::Result> { + self.claim_due_thread_schedule_with_params(ThreadScheduleDueClaimParams { + now, + lease_id, + lease_duration, + local_active_owner_id: None, + local_active_fresh_after: None, + }) + .await + } + + pub async fn claim_due_thread_schedule_with_params( + &self, + params: ThreadScheduleDueClaimParams<'_>, + ) -> anyhow::Result> { + let ThreadScheduleDueClaimParams { + now, + lease_id, + lease_duration, + local_active_owner_id, + local_active_fresh_after, + } = params; + let params = ClaimThreadScheduleParams { + target: ThreadScheduleClaimTarget::Due, + now, + lease_id, + lease_duration, + local_active_owner_id, + local_active_fresh_after, + }; + crate::busy_retry::retry_on_busy("claim due thread schedule", || { + self.claim_thread_schedule_once(params.clone()) + }) + .await + } + + pub async fn claim_thread_schedule_now( + &self, + schedule_id: &str, + now: DateTime, + lease_id: &str, + lease_duration: Duration, + ) -> anyhow::Result> { + self.claim_thread_schedule_now_with_params(ThreadScheduleNowClaimParams { + schedule_id, + now, + lease_id, + lease_duration, + local_active_owner_id: None, + local_active_fresh_after: None, + }) + .await + } + + pub async fn claim_thread_schedule_now_with_params( + &self, + params: ThreadScheduleNowClaimParams<'_>, + ) -> anyhow::Result> { + let ThreadScheduleNowClaimParams { + schedule_id, + now, + lease_id, + lease_duration, + local_active_owner_id, + local_active_fresh_after, + } = params; + let params = ClaimThreadScheduleParams { + target: ThreadScheduleClaimTarget::Now { schedule_id }, + now, + lease_id, + lease_duration, + local_active_owner_id, + local_active_fresh_after, + }; + crate::busy_retry::retry_on_busy("claim thread schedule now", || { + self.claim_thread_schedule_once(params.clone()) + }) + .await + } + + async fn claim_thread_schedule_once( + &self, + params: ClaimThreadScheduleParams<'_>, + ) -> anyhow::Result> { + let ClaimThreadScheduleParams { + target, + now, + lease_id, + lease_duration, + local_active_owner_id, + local_active_fresh_after, + } = params; + let now_ms = datetime_to_epoch_millis(now); + let lease_expires_at = now + chrono::Duration::from_std(lease_duration)?; + let lease_expires_at_ms = datetime_to_epoch_millis(lease_expires_at); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let owner_filter = match (local_active_owner_id, local_active_fresh_after) { + (Some(owner_id), Some(fresh_after)) => { + Some((owner_id, datetime_to_epoch_millis(fresh_after))) + } + _ => None, + }; + let owner_scoped_lease_id = owner_filter.as_ref().map(|_| format!("owner:{lease_id}")); + let lease_id = owner_scoped_lease_id.as_deref().unwrap_or(lease_id); + let active_owner_filter = if owner_filter.is_some() { + r#" + AND NOT EXISTS ( + SELECT 1 + FROM local_active_sessions + WHERE local_active_sessions.thread_id = thread_schedules.thread_id + AND local_active_sessions.last_seen_at_ms >= ? + AND local_active_sessions.owner_id != ? + ) +"# + } else { + "" + }; + let sql = match target { + ThreadScheduleClaimTarget::Due => format!( + r#" +SELECT {SCHEDULE_COLUMNS} +FROM thread_schedules +WHERE status = 'active' + AND next_run_at_ms IS NOT NULL + AND next_run_at_ms <= ? + AND (expires_at_ms IS NULL OR expires_at_ms > ?) + AND (lease_id IS NULL OR lease_expires_at_ms <= ?) + AND NOT EXISTS ( + SELECT 1 + FROM thread_schedule_occurrences + WHERE thread_schedule_occurrences.schedule_id = thread_schedules.schedule_id + AND thread_schedule_occurrences.state IN ('waiting_idle', 'enqueued', 'started') + AND thread_schedule_occurrences.retry_at_ms > ? + ) +{active_owner_filter} +ORDER BY next_run_at_ms, created_at_ms +LIMIT 1 +"# + ), + ThreadScheduleClaimTarget::Now { .. } => format!( + r#" +SELECT {SCHEDULE_COLUMNS} +FROM thread_schedules +WHERE schedule_id = ? + AND status = 'active' + AND (expires_at_ms IS NULL OR expires_at_ms > ?) + AND (lease_id IS NULL OR lease_expires_at_ms <= ?) +{active_owner_filter} +"# + ), + }; + let mut query = sqlx::query(sqlx::AssertSqlSafe(sql)); + query = match target { + ThreadScheduleClaimTarget::Due => { + query.bind(now_ms).bind(now_ms).bind(now_ms).bind(now_ms) + } + ThreadScheduleClaimTarget::Now { schedule_id } => { + query.bind(schedule_id).bind(now_ms).bind(now_ms) + } + }; + if let Some((owner_id, fresh_after_ms)) = owner_filter { + query = query.bind(fresh_after_ms).bind(owner_id); + } + let schedule_row = query.fetch_optional(&mut *tx).await?; + let Some(schedule_row) = schedule_row else { + tx.commit().await?; + return Ok(None); + }; + let selected_schedule = thread_schedule_from_row(&schedule_row)?; + let existing_occurrence = sqlx::query_as::<_, ThreadScheduleOccurrenceRow>( + r#" +SELECT + occurrence_id, + schedule_id, + thread_id, + state, + turn_id, + goal_id, + auth_profile_recorded, + auth_profile, + scheduled_for_ms, + turn_input, + created_at_ms +FROM thread_schedule_occurrences +WHERE schedule_id = ? + "#, + ) + .bind(selected_schedule.schedule_id.as_str()) + .fetch_optional(&mut *tx) + .await?; + let goal_id = existing_occurrence + .as_ref() + .and_then(|occurrence| occurrence.goal_id.as_deref()); + let goal_hold_can_pause = + goal_id.is_some() && selected_schedule.schedule != crate::ThreadScheduleSpec::Once; + // Read-only probe: this transaction only runs `SELECT EXISTS` against + // goals.db and is always rolled back, so a deferred `BEGIN` is enough. A + // `BEGIN IMMEDIATE` would take a goals.db write lock and hold it across the + // state.db commit for no benefit. Lock order is consistently state -> goals + // at every site that touches both, so there is no inversion to guard against. + let mut goal_tx = if goal_hold_can_pause { + Some(self.goals_pool.begin().await?) + } else { + None + }; + let mut pause_for_goal_hold = false; + if let (Some(goal_tx), Some(goal_id)) = (goal_tx.as_mut(), goal_id) { + pause_for_goal_hold = sqlx::query_scalar::<_, bool>( + r#" +SELECT EXISTS( + SELECT 1 + FROM thread_goals + WHERE thread_id = ? + AND goal_id = ? + AND status IN ('paused', 'blocked', 'usage_limited', 'budget_limited') +) + "#, + ) + .bind(selected_schedule.thread_id.to_string()) + .bind(goal_id) + .fetch_one(&mut **goal_tx) + .await?; + } + if pause_for_goal_hold + && existing_occurrence.as_ref().is_none_or(|occurrence| { + occurrence.state != OCCURRENCE_STARTED && occurrence.state != OCCURRENCE_TERMINAL + }) + { + if let Some(occurrence) = existing_occurrence.as_ref() { + let error = redact_state_string("scheduled run stopped because its goal is held"); + if occurrence.state == OCCURRENCE_ENQUEUED { + sqlx::query( + r#" +INSERT INTO thread_schedule_runs ( + run_id, + schedule_id, + thread_id, + status, + lease_id, + turn_id, + goal_id, + error, + scheduled_for_ms, + started_at_ms, + completed_at_ms +) +SELECT + occurrence_id, + schedule_id, + thread_id, + 'failed', + ?, + turn_id, + goal_id, + ?, + scheduled_for_ms, + created_at_ms, + ? +FROM thread_schedule_occurrences +WHERE occurrence_id = ? + "#, + ) + .bind(lease_id) + .bind(error) + .bind(now_ms) + .bind(occurrence.occurrence_id.as_str()) + .execute(&mut *tx) + .await?; + } + } + sqlx::query("DELETE FROM thread_schedule_occurrences WHERE schedule_id = ?") + .bind(selected_schedule.schedule_id.as_str()) + .execute(&mut *tx) + .await?; + sqlx::query( + r#" +UPDATE thread_schedules +SET status = 'paused', + next_run_at_ms = NULL, + last_run_at_ms = ?, + failure_count = failure_count + 1, + lease_id = NULL, + lease_expires_at_ms = NULL, + updated_at_ms = ? +WHERE schedule_id = ? AND status = 'active' + "#, + ) + .bind(now_ms) + .bind(now_ms) + .bind(selected_schedule.schedule_id.as_str()) + .execute(&mut *tx) + .await?; + tx.commit().await?; + if let Some(goal_tx) = goal_tx { + let _ = goal_tx.rollback().await; + } + return Ok(None); + } + let sql = schedule_returning( + r#" + UPDATE thread_schedules +SET lease_id = ?, + lease_expires_at_ms = ?, + updated_at_ms = ? +WHERE schedule_id = ? AND status = 'active' +RETURNING +"#, + ); + let schedule_row = sqlx::query(sqlx::AssertSqlSafe(sql)) + .bind(lease_id) + .bind(lease_expires_at_ms) + .bind(now_ms) + .bind(selected_schedule.schedule_id.as_str()) + .fetch_optional(&mut *tx) + .await?; + let Some(schedule_row) = schedule_row else { + // `thread_schedules_ignore_legacy_live_owner_claim` silently drops + // the lease update (RAISE(IGNORE)) when a legacy, non-owner-scoped + // lease is claimed while a local session is live. Treat that as an + // unclaimed schedule and discard the speculative reap above so the + // live owner keeps ownership of its runs. + tx.rollback().await?; + if let Some(goal_tx) = goal_tx { + let _ = goal_tx.rollback().await; + } + return Ok(None); + }; + let schedule = thread_schedule_from_row(&schedule_row)?; + let occurrence = match existing_occurrence { + Some(occurrence) => { + sqlx::query( + r#" +UPDATE thread_schedule_occurrences +SET state = ?, updated_at_ms = ? +WHERE occurrence_id = ? + "#, + ) + .bind(occurrence.state.as_str()) + .bind(now_ms) + .bind(occurrence.occurrence_id.as_str()) + .execute(&mut *tx) + .await?; + occurrence + } + None => { + let occurrence_id = Uuid::new_v4().to_string(); + let turn_id = Uuid::now_v7().to_string(); + let scheduled_for_ms = match target { + ThreadScheduleClaimTarget::Due => { + selected_schedule.next_run_at.map(datetime_to_epoch_millis) + } + ThreadScheduleClaimTarget::Now { .. } => Some(now_ms), + }; + sqlx::query( + r#" +INSERT INTO thread_schedule_occurrences ( + occurrence_id, + schedule_id, + thread_id, + state, + turn_id, + scheduled_for_ms, + created_at_ms, + updated_at_ms +) VALUES (?, ?, ?, 'waiting_idle', ?, ?, ?, ?) + "#, + ) + .bind(occurrence_id.as_str()) + .bind(schedule.schedule_id.as_str()) + .bind(schedule.thread_id.to_string()) + .bind(turn_id.as_str()) + .bind(scheduled_for_ms) + .bind(now_ms) + .bind(now_ms) + .execute(&mut *tx) + .await?; + ThreadScheduleOccurrenceRow { + occurrence_id, + schedule_id: schedule.schedule_id.clone(), + thread_id: schedule.thread_id.to_string(), + state: OCCURRENCE_WAITING_IDLE.to_string(), + turn_id, + goal_id: None, + auth_profile_recorded: false, + auth_profile: None, + scheduled_for_ms, + turn_input: None, + created_at_ms: now_ms, + } + } + }; + if occurrence.state == OCCURRENCE_STARTED || occurrence.state == OCCURRENCE_TERMINAL { + sqlx::query("UPDATE thread_schedule_runs SET lease_id = ? WHERE run_id = ?") + .bind(lease_id) + .bind(occurrence.occurrence_id.as_str()) + .execute(&mut *tx) + .await?; + } + let run = Self::occurrence_run(&mut tx, &occurrence, lease_id).await?; + tx.commit().await?; + if let Some(goal_tx) = goal_tx { + let _ = goal_tx.rollback().await; + } + let occurrence_state = ThreadScheduleOccurrenceState::from_str(&occurrence.state)?; + let occurrence_auth_profile = occurrence + .auth_profile_recorded + .then(|| occurrence.auth_profile.clone()); + Ok(Some(ThreadScheduleClaim { + schedule, + run, + occurrence_state, + turn_input: occurrence.turn_input, + occurrence_auth_profile, + })) + } + + async fn occurrence_run( + tx: &mut sqlx::Transaction<'_, Sqlite>, + occurrence: &ThreadScheduleOccurrenceRow, + lease_id: &str, + ) -> anyhow::Result { + if occurrence.state == OCCURRENCE_STARTED || occurrence.state == OCCURRENCE_TERMINAL { + let sql = run_returning("SELECT"); + let row = sqlx::query(sqlx::AssertSqlSafe(format!( + "{sql} FROM thread_schedule_runs WHERE run_id = ?" + ))) + .bind(occurrence.occurrence_id.as_str()) + .fetch_one(&mut **tx) + .await?; + return thread_schedule_run_from_row(&row); + } + Ok(crate::ThreadScheduleRun { + thread_id: ThreadId::try_from(occurrence.thread_id.clone())?, + schedule_id: occurrence.schedule_id.clone(), + run_id: occurrence.occurrence_id.clone(), + status: crate::ThreadScheduleRunStatus::Leased, + lease_id: lease_id.to_string(), + turn_id: Some(occurrence.turn_id.clone()), + goal_id: occurrence.goal_id.clone(), + error: None, + scheduled_for: optional_epoch_millis_to_datetime(occurrence.scheduled_for_ms)?, + started_at: epoch_millis_to_datetime(occurrence.created_at_ms)?, + completed_at: None, + }) + } +} diff --git a/codex-rs/state/src/runtime/schedules/occurrence/finish.rs b/codex-rs/state/src/runtime/schedules/occurrence/finish.rs new file mode 100644 index 000000000..8ce5aa6d8 --- /dev/null +++ b/codex-rs/state/src/runtime/schedules/occurrence/finish.rs @@ -0,0 +1,730 @@ +//! Terminal recording, cadence finalization, and explicit deferrals. + +use super::*; + +impl ScheduleStore { + pub async fn record_thread_schedule_run_terminal( + &self, + schedule_id: &str, + run_id: &str, + lease_id: &str, + completed_at: DateTime, + expected_goal_id: Option<&str>, + error: Option, + ) -> anyhow::Result { + let completed_at_ms = datetime_to_epoch_millis(completed_at); + let status = if error.is_some() { + crate::ThreadScheduleRunStatus::Failed + } else { + crate::ThreadScheduleRunStatus::Completed + }; + let error = error.map(redact_state_string); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let run_result = sqlx::query( + r#" +UPDATE thread_schedule_runs +SET status = ?, error = ?, completed_at_ms = ? +WHERE schedule_id = ? + AND run_id = ? + AND lease_id = ? + AND status = 'running' + AND (? IS NULL OR goal_id IS NULL OR goal_id = ?) + AND EXISTS ( + SELECT 1 + FROM thread_schedules + WHERE thread_schedules.schedule_id = thread_schedule_runs.schedule_id + AND thread_schedules.lease_id = ? + ) + AND EXISTS ( + SELECT 1 + FROM thread_schedule_occurrences + WHERE thread_schedule_occurrences.occurrence_id = thread_schedule_runs.run_id + AND thread_schedule_occurrences.state = 'started' + ) + "#, + ) + .bind(status.as_str()) + .bind(error) + .bind(completed_at_ms) + .bind(schedule_id) + .bind(run_id) + .bind(lease_id) + .bind(expected_goal_id) + .bind(expected_goal_id) + .bind(lease_id) + .execute(&mut *tx) + .await?; + if run_result.rows_affected() == 0 { + let already_terminal: bool = sqlx::query_scalar( + r#" +SELECT EXISTS( + SELECT 1 + FROM thread_schedule_runs + JOIN thread_schedule_occurrences + ON thread_schedule_occurrences.occurrence_id = thread_schedule_runs.run_id + JOIN thread_schedules + ON thread_schedules.schedule_id = thread_schedule_runs.schedule_id + WHERE thread_schedule_runs.schedule_id = ? + AND thread_schedule_runs.run_id = ? + AND thread_schedule_runs.lease_id = ? + AND thread_schedule_runs.status IN ('completed', 'failed') + AND thread_schedule_occurrences.state = 'terminal' + AND thread_schedules.lease_id = ? + AND (? IS NULL OR thread_schedule_runs.goal_id IS NULL OR thread_schedule_runs.goal_id = ?) +) + "#, + ) + .bind(schedule_id) + .bind(run_id) + .bind(lease_id) + .bind(lease_id) + .bind(expected_goal_id) + .bind(expected_goal_id) + .fetch_one(&mut *tx) + .await?; + tx.commit().await?; + return Ok(already_terminal); + } + let occurrence_result = sqlx::query( + r#" +UPDATE thread_schedule_occurrences +SET state = 'terminal', updated_at_ms = ? +WHERE occurrence_id = ? AND schedule_id = ? AND state = 'started' + "#, + ) + .bind(completed_at_ms) + .bind(run_id) + .bind(schedule_id) + .execute(&mut *tx) + .await?; + if occurrence_result.rows_affected() == 0 { + tx.rollback().await?; + return Ok(false); + } + tx.commit().await?; + Ok(true) + } + + pub async fn fail_thread_schedule_occurrence_before_start( + &self, + schedule_id: &str, + run_id: &str, + lease_id: &str, + completed_at: DateTime, + goal_id: Option<&str>, + error: String, + ) -> anyhow::Result { + let completed_at_ms = datetime_to_epoch_millis(completed_at); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let sql = run_returning( + r#" +INSERT INTO thread_schedule_runs ( + run_id, + schedule_id, + thread_id, + status, + lease_id, + turn_id, + goal_id, + error, + scheduled_for_ms, + started_at_ms, + completed_at_ms +) +SELECT + occurrence_id, + schedule_id, + thread_id, + 'failed', + ?, + turn_id, + COALESCE(goal_id, ?), + ?, + scheduled_for_ms, + created_at_ms, + ? +FROM thread_schedule_occurrences +WHERE occurrence_id = ? + AND schedule_id = ? + AND state IN ('waiting_idle', 'enqueued') + AND EXISTS ( + SELECT 1 + FROM thread_schedules + WHERE thread_schedules.schedule_id = thread_schedule_occurrences.schedule_id + AND thread_schedules.lease_id = ? + ) +RETURNING +"#, + ); + let row = sqlx::query(sqlx::AssertSqlSafe(sql)) + .bind(lease_id) + .bind(goal_id) + .bind(redact_state_string(error)) + .bind(completed_at_ms) + .bind(run_id) + .bind(schedule_id) + .bind(lease_id) + .fetch_optional(&mut *tx) + .await?; + if row.is_none() { + tx.commit().await?; + return Ok(false); + } + let occurrence_result = sqlx::query( + r#" +UPDATE thread_schedule_occurrences +SET state = 'terminal', + goal_id = COALESCE(goal_id, ?), + updated_at_ms = ? +WHERE occurrence_id = ? AND state IN ('waiting_idle', 'enqueued') + "#, + ) + .bind(goal_id) + .bind(completed_at_ms) + .bind(run_id) + .execute(&mut *tx) + .await?; + if occurrence_result.rows_affected() == 0 { + tx.rollback().await?; + return Ok(false); + } + tx.commit().await?; + Ok(true) + } + + pub async fn finalize_terminal_thread_schedule_run( + &self, + schedule_id: &str, + run_id: &str, + lease_id: &str, + completed_at: DateTime, + next_run_at: Option>, + expected_goal_id: Option<&str>, + ) -> anyhow::Result { + self.finalize_terminal_thread_schedule_run_once(FinalizeThreadScheduleRunParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + expected_goal_id, + }) + .await + } + + pub async fn complete_thread_schedule_run( + &self, + schedule_id: &str, + run_id: &str, + lease_id: &str, + completed_at: DateTime, + next_run_at: Option>, + ) -> anyhow::Result { + self.finish_thread_schedule_run(FinishThreadScheduleRunParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + expected_goal_id: None, + finish: FinishScheduleRun::Completed, + }) + .await + } + + pub async fn complete_thread_schedule_run_for_goal( + &self, + params: ThreadScheduleRunForGoalFinishParams<'_>, + ) -> anyhow::Result { + let ThreadScheduleRunForGoalFinishParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + expected_goal_id, + } = params; + self.finish_thread_schedule_run(FinishThreadScheduleRunParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + expected_goal_id: Some(expected_goal_id), + finish: FinishScheduleRun::Completed, + }) + .await + } + + pub async fn fail_thread_schedule_run( + &self, + schedule_id: &str, + run_id: &str, + lease_id: &str, + completed_at: DateTime, + next_run_at: Option>, + error: String, + ) -> anyhow::Result { + self.finish_thread_schedule_run(FinishThreadScheduleRunParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + expected_goal_id: None, + finish: FinishScheduleRun::Failed { error }, + }) + .await + } + + pub async fn fail_thread_schedule_run_for_goal( + &self, + params: ThreadScheduleRunForGoalFinishParams<'_>, + error: String, + ) -> anyhow::Result { + let ThreadScheduleRunForGoalFinishParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + expected_goal_id, + } = params; + self.finish_thread_schedule_run(FinishThreadScheduleRunParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + expected_goal_id: Some(expected_goal_id), + finish: FinishScheduleRun::Failed { error }, + }) + .await + } + + pub async fn defer_thread_schedule_run( + &self, + schedule_id: &str, + run_id: &str, + lease_id: &str, + completed_at: DateTime, + next_run_at: DateTime, + error: String, + ) -> anyhow::Result { + let completed_at_ms = datetime_to_epoch_millis(completed_at); + let requested_next_run_at_ms = datetime_to_epoch_millis(next_run_at); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let run_result = sqlx::query( + r#" +INSERT INTO thread_schedule_runs ( + run_id, + schedule_id, + thread_id, + status, + lease_id, + turn_id, + goal_id, + error, + scheduled_for_ms, + started_at_ms, + completed_at_ms, + deferral_kind +) +SELECT + occurrence_id, + schedule_id, + thread_id, + 'deferred', + ?, + NULL, + goal_id, + ?, + scheduled_for_ms, + created_at_ms, + ?, + 'capacity' +FROM thread_schedule_occurrences +WHERE schedule_id = ? + AND occurrence_id = ? + AND EXISTS ( + SELECT 1 + FROM thread_schedules + WHERE thread_schedules.schedule_id = thread_schedule_occurrences.schedule_id + AND thread_schedules.lease_id = ? + ) +ON CONFLICT(run_id) DO UPDATE SET + status = 'deferred', + error = excluded.error, + completed_at_ms = excluded.completed_at_ms, + deferral_kind = 'capacity' + "#, + ) + .bind(lease_id) + .bind(redact_state_string(error)) + .bind(completed_at_ms) + .bind(schedule_id) + .bind(run_id) + .bind(lease_id) + .execute(&mut *tx) + .await?; + if run_result.rows_affected() == 0 { + tx.rollback().await?; + return Ok(false); + } + sqlx::query("DELETE FROM thread_schedule_occurrences WHERE occurrence_id = ?") + .bind(run_id) + .execute(&mut *tx) + .await?; + let schedule_result = sqlx::query( + r#" +UPDATE thread_schedules +SET + status = CASE + WHEN status = 'expired' THEN 'expired' + WHEN expires_at_ms IS NOT NULL AND ? >= expires_at_ms THEN 'expired' + WHEN status = 'paused' THEN 'paused' + ELSE status + END, + lease_id = NULL, + lease_expires_at_ms = NULL, + last_run_at_ms = ?, + next_run_at_ms = CASE + WHEN status IN ('expired', 'paused') THEN NULL + WHEN expires_at_ms IS NOT NULL AND ? >= expires_at_ms THEN NULL + ELSE ? + END, + updated_at_ms = ? +WHERE schedule_id = ? AND lease_id = ? + "#, + ) + .bind(requested_next_run_at_ms) + .bind(completed_at_ms) + .bind(requested_next_run_at_ms) + .bind(requested_next_run_at_ms) + .bind(completed_at_ms) + .bind(schedule_id) + .bind(lease_id) + .execute(&mut *tx) + .await?; + if schedule_result.rows_affected() == 0 { + tx.rollback().await?; + return Ok(false); + } + tx.commit().await?; + Ok(true) + } + + pub async fn wait_thread_schedule_run_for_idle( + &self, + schedule_id: &str, + run_id: &str, + lease_id: &str, + retry_at: DateTime, + now: DateTime, + ) -> anyhow::Result { + let now_ms = datetime_to_epoch_millis(now); + let retry_at_ms = datetime_to_epoch_millis(retry_at); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let occurrence_result = sqlx::query( + r#" +UPDATE thread_schedule_occurrences +SET retry_at_ms = ?, + updated_at_ms = ? +WHERE schedule_id = ? + AND occurrence_id = ? + AND state IN ('waiting_idle', 'enqueued', 'started') + "#, + ) + .bind(retry_at_ms) + .bind(now_ms) + .bind(schedule_id) + .bind(run_id) + .execute(&mut *tx) + .await?; + if occurrence_result.rows_affected() == 0 { + tx.commit().await?; + return Ok(false); + } + let schedule_result = sqlx::query( + r#" +UPDATE thread_schedules +SET lease_id = NULL, lease_expires_at_ms = NULL, updated_at_ms = ? +WHERE schedule_id = ? AND lease_id = ? + "#, + ) + .bind(now_ms) + .bind(schedule_id) + .bind(lease_id) + .execute(&mut *tx) + .await?; + if schedule_result.rows_affected() == 0 { + tx.rollback().await?; + return Ok(false); + } + tx.commit().await?; + Ok(true) + } + + pub async fn expire_thread_schedules(&self, now: DateTime) -> anyhow::Result { + let now_ms = datetime_to_epoch_millis(now); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + sqlx::query( + r#" +UPDATE thread_schedule_runs +SET status = 'failed', + error = ?, + completed_at_ms = COALESCE(completed_at_ms, ?) +WHERE status IN ('leased', 'running') + AND EXISTS ( + SELECT 1 + FROM thread_schedules + WHERE thread_schedules.schedule_id = thread_schedule_runs.schedule_id + AND thread_schedules.status = 'active' + AND thread_schedules.expires_at_ms IS NOT NULL + AND thread_schedules.expires_at_ms <= ? + AND ( + thread_schedules.lease_id IS NULL + OR thread_schedules.lease_expires_at_ms <= ? + ) + ) + "#, + ) + .bind(redact_state_string(EXPIRED_SCHEDULE_RUN_ERROR)) + .bind(now_ms) + .bind(now_ms) + .bind(now_ms) + .execute(&mut *tx) + .await?; + let result = sqlx::query( + r#" +UPDATE thread_schedules +SET + status = 'expired', + next_run_at_ms = NULL, + lease_id = NULL, + lease_expires_at_ms = NULL, + updated_at_ms = ? +WHERE status = 'active' + AND expires_at_ms IS NOT NULL + AND expires_at_ms <= ? + AND (lease_id IS NULL OR lease_expires_at_ms <= ?) + "#, + ) + .bind(now_ms) + .bind(now_ms) + .bind(now_ms) + .execute(&mut *tx) + .await?; + sqlx::query( + r#" +DELETE FROM thread_schedule_occurrences +WHERE EXISTS ( + SELECT 1 + FROM thread_schedules + WHERE thread_schedules.schedule_id = thread_schedule_occurrences.schedule_id + AND thread_schedules.status = 'expired' +) + "#, + ) + .execute(&mut *tx) + .await?; + tx.commit().await?; + Ok(result.rows_affected()) + } + + async fn finish_thread_schedule_run( + &self, + params: FinishThreadScheduleRunParams<'_>, + ) -> anyhow::Result { + crate::busy_retry::retry_on_busy("finish thread schedule run", || { + self.finish_thread_schedule_run_once(params.clone()) + }) + .await + } + + async fn finish_thread_schedule_run_once( + &self, + params: FinishThreadScheduleRunParams<'_>, + ) -> anyhow::Result { + let FinishThreadScheduleRunParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + expected_goal_id, + finish, + } = params; + let error = match finish { + FinishScheduleRun::Completed => None, + FinishScheduleRun::Failed { error } => Some(error), + }; + if !self + .record_thread_schedule_run_terminal( + schedule_id, + run_id, + lease_id, + completed_at, + expected_goal_id, + error, + ) + .await? + { + return Ok(false); + } + self.finalize_terminal_thread_schedule_run_once(FinalizeThreadScheduleRunParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + expected_goal_id, + }) + .await + } + + async fn finalize_terminal_thread_schedule_run_once( + &self, + params: FinalizeThreadScheduleRunParams<'_>, + ) -> anyhow::Result { + let FinalizeThreadScheduleRunParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + expected_goal_id, + } = params; + let completed_at_ms = datetime_to_epoch_millis(completed_at); + let next_run_at_ms = next_run_at.map(datetime_to_epoch_millis); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let schedule_context: Option<(String, String, String)> = sqlx::query_as( + r#" +SELECT thread_schedules.thread_id, thread_schedules.schedule_kind, thread_schedule_runs.status +FROM thread_schedules +JOIN thread_schedule_runs ON thread_schedule_runs.schedule_id = thread_schedules.schedule_id +JOIN thread_schedule_occurrences ON thread_schedule_occurrences.occurrence_id = thread_schedule_runs.run_id +WHERE thread_schedules.schedule_id = ? AND thread_schedules.lease_id = ? + AND thread_schedule_runs.run_id = ? + AND thread_schedule_runs.lease_id = ? + AND thread_schedule_runs.status IN ('completed', 'failed') + AND thread_schedule_occurrences.state = 'terminal' + AND (? IS NULL OR thread_schedule_runs.goal_id IS NULL OR thread_schedule_runs.goal_id = ?) + "#, + ) + .bind(schedule_id) + .bind(lease_id) + .bind(run_id) + .bind(lease_id) + .bind(expected_goal_id) + .bind(expected_goal_id) + .fetch_optional(&mut *tx) + .await?; + let Some((thread_id, schedule_kind, run_status)) = schedule_context else { + tx.commit().await?; + return Ok(false); + }; + let goal_hold_can_pause = expected_goal_id.is_some() && schedule_kind != ONCE_SCHEDULE_KIND; + // Read-only probe: this transaction only runs `SELECT EXISTS` against + // goals.db and is always rolled back, so a deferred `BEGIN` is enough. A + // `BEGIN IMMEDIATE` would take a goals.db write lock and hold it across the + // state.db commit for no benefit. Lock order is consistently state -> goals + // at every site that touches both, so there is no inversion to guard against. + let mut goal_tx = if goal_hold_can_pause { + Some(self.goals_pool.begin().await?) + } else { + None + }; + let pause_for_goal_hold = match (expected_goal_id, goal_hold_can_pause, goal_tx.as_mut()) { + (Some(expected_goal_id), true, Some(goal_tx)) => { + sqlx::query_scalar::<_, bool>( + r#" +SELECT EXISTS( + SELECT 1 + FROM thread_goals + WHERE thread_id = ? + AND goal_id = ? + AND status IN ('paused', 'blocked', 'usage_limited', 'budget_limited') +) + "#, + ) + .bind(thread_id) + .bind(expected_goal_id) + .fetch_one(&mut **goal_tx) + .await? + } + (Some(_), false, None) | (None, false, None) => false, + // `goal_hold_can_pause` is what decides whether `goal_tx` was opened, so + // the arms above are exhaustive in practice. Fail the write instead of + // panicking out of a state-store transaction if that ever drifts. + (expected_goal_id, goal_hold_can_pause, goal_tx) => { + anyhow::bail!( + "goal transaction presence does not match the recurring goal schedule invariant (expected_goal_id={}, goal_hold_can_pause={goal_hold_can_pause}, goal_tx={})", + expected_goal_id.is_some(), + goal_tx.is_some(), + ); + } + }; + let failed = run_status == crate::ThreadScheduleRunStatus::Failed.as_str(); + // The only thing that pauses a schedule at finish time is a goal hold; there + // is deliberately no caller-supplied pause flag. + let pause_schedule = pause_for_goal_hold; + let schedule_result = sqlx::query( + r#" +UPDATE thread_schedules +SET + status = CASE + WHEN status = 'expired' THEN 'expired' + WHEN expires_at_ms IS NOT NULL AND ? >= expires_at_ms THEN 'expired' + WHEN status = 'paused' THEN 'paused' + WHEN ? THEN 'paused' + WHEN ? IS NULL THEN 'expired' + ELSE status + END, + lease_id = NULL, + lease_expires_at_ms = NULL, + last_run_at_ms = ?, + next_run_at_ms = CASE + WHEN status IN ('expired', 'paused') THEN NULL + WHEN expires_at_ms IS NOT NULL AND ? >= expires_at_ms THEN NULL + WHEN ? THEN NULL + WHEN ? IS NULL THEN NULL + ELSE ? + END, + failure_count = CASE WHEN ? THEN failure_count + 1 ELSE 0 END, + updated_at_ms = ? +WHERE schedule_id = ? AND lease_id = ? + "#, + ) + .bind(completed_at_ms) + .bind(pause_schedule) + .bind(next_run_at_ms) + .bind(completed_at_ms) + .bind(completed_at_ms) + .bind(pause_schedule) + .bind(next_run_at_ms) + .bind(next_run_at_ms) + .bind(failed) + .bind(completed_at_ms) + .bind(schedule_id) + .bind(lease_id) + .execute(&mut *tx) + .await?; + if schedule_result.rows_affected() == 0 { + tx.commit().await?; + if let Some(goal_tx) = goal_tx { + let _ = goal_tx.rollback().await; + } + return Ok(false); + } + sqlx::query("DELETE FROM thread_schedule_occurrences WHERE occurrence_id = ?") + .bind(run_id) + .execute(&mut *tx) + .await?; + tx.commit().await?; + if let Some(goal_tx) = goal_tx { + let _ = goal_tx.rollback().await; + } + Ok(true) + } +} diff --git a/codex-rs/state/src/runtime/schedules/occurrence/start.rs b/codex-rs/state/src/runtime/schedules/occurrence/start.rs new file mode 100644 index 000000000..10698910a --- /dev/null +++ b/codex-rs/state/src/runtime/schedules/occurrence/start.rs @@ -0,0 +1,251 @@ +//! Enqueue and durable start transitions for scheduled occurrences. + +use super::*; + +impl ScheduleStore { + pub async fn enqueue_thread_schedule_run( + &self, + params: ThreadScheduleRunEnqueueParams<'_>, + ) -> anyhow::Result> { + let ThreadScheduleRunEnqueueParams { + schedule_id, + run_id, + lease_id, + goal_id, + auth_profile_recorded, + auth_profile, + turn_input, + now, + } = params; + let now_ms = datetime_to_epoch_millis(now); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let updated = sqlx::query( + r#" +UPDATE thread_schedule_occurrences +SET state = 'enqueued', + goal_id = COALESCE(goal_id, ?), + auth_profile_recorded = CASE WHEN ? THEN 1 ELSE auth_profile_recorded END, + auth_profile = CASE WHEN ? THEN ? ELSE auth_profile END, + turn_input = ?, + retry_at_ms = NULL, + updated_at_ms = ? +WHERE occurrence_id = ? + AND schedule_id = ? + AND state IN ('waiting_idle', 'enqueued') + AND EXISTS ( + SELECT 1 + FROM thread_schedules + WHERE thread_schedules.schedule_id = thread_schedule_occurrences.schedule_id + AND thread_schedules.lease_id = ? + AND thread_schedules.lease_expires_at_ms > ? + ) + "#, + ) + .bind(goal_id) + .bind(auth_profile_recorded) + .bind(auth_profile_recorded) + .bind(auth_profile) + .bind(redact_state_string(turn_input)) + .bind(now_ms) + .bind(run_id) + .bind(schedule_id) + .bind(lease_id) + .bind(now_ms) + .execute(&mut *tx) + .await?; + if updated.rows_affected() == 0 { + tx.commit().await?; + return Ok(None); + } + let occurrence = sqlx::query_as::<_, ThreadScheduleOccurrenceRow>( + r#" +SELECT occurrence_id, schedule_id, thread_id, state, turn_id, goal_id, + auth_profile_recorded, auth_profile, scheduled_for_ms, turn_input, created_at_ms +FROM thread_schedule_occurrences +WHERE occurrence_id = ? + "#, + ) + .bind(run_id) + .fetch_one(&mut *tx) + .await?; + let run = Self::occurrence_run(&mut tx, &occurrence, lease_id).await?; + tx.commit().await?; + Ok(Some(run)) + } + + pub async fn mark_thread_schedule_run_started( + &self, + params: ThreadScheduleRunStartParams<'_>, + ) -> anyhow::Result> { + crate::busy_retry::retry_on_busy("mark thread schedule run started", || { + self.mark_thread_schedule_run_started_once(params.clone()) + }) + .await + } + + async fn mark_thread_schedule_run_started_once( + &self, + params: ThreadScheduleRunStartParams<'_>, + ) -> anyhow::Result> { + let ThreadScheduleRunStartParams { + schedule_id, + run_id, + lease_id, + turn_id, + goal_id, + now, + lease_duration, + } = params; + let now_ms = datetime_to_epoch_millis(now); + let lease_expires_at_ms = + datetime_to_epoch_millis(now + chrono::Duration::from_std(lease_duration)?); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let schedule_result = sqlx::query( + r#" +UPDATE thread_schedules +SET lease_expires_at_ms = MAX(lease_expires_at_ms, ?), + updated_at_ms = ? +WHERE schedule_id = ? + AND lease_id = ? + AND lease_expires_at_ms > ? + AND EXISTS ( + SELECT 1 + FROM thread_schedule_occurrences + WHERE thread_schedule_occurrences.schedule_id = thread_schedules.schedule_id + AND thread_schedule_occurrences.occurrence_id = ? + AND thread_schedule_occurrences.turn_id = ? + AND thread_schedule_occurrences.state IN ('enqueued', 'started') + ) + "#, + ) + .bind(lease_expires_at_ms) + .bind(now_ms) + .bind(schedule_id) + .bind(lease_id) + .bind(now_ms) + .bind(run_id) + .bind(turn_id) + .execute(&mut *tx) + .await?; + if schedule_result.rows_affected() == 0 { + tx.commit().await?; + return Ok(None); + } + let occurrence_result = sqlx::query( + r#" +UPDATE thread_schedule_occurrences +SET state = 'started', goal_id = COALESCE(goal_id, ?), updated_at_ms = ? +WHERE schedule_id = ? + AND occurrence_id = ? + AND turn_id = ? + AND state IN ('enqueued', 'started') +"#, + ) + .bind(goal_id) + .bind(now_ms) + .bind(schedule_id) + .bind(run_id) + .bind(turn_id) + .execute(&mut *tx) + .await?; + if occurrence_result.rows_affected() == 0 { + tx.rollback().await?; + return Ok(None); + } + let sql = run_returning( + r#" +INSERT INTO thread_schedule_runs ( + run_id, + schedule_id, + thread_id, + status, + lease_id, + turn_id, + goal_id, + scheduled_for_ms, + started_at_ms +) +SELECT + occurrence_id, + schedule_id, + thread_id, + 'running', + ?, + turn_id, + goal_id, + scheduled_for_ms, + ? +FROM thread_schedule_occurrences +WHERE occurrence_id = ? AND state = 'started' +ON CONFLICT(run_id) DO UPDATE SET + lease_id = excluded.lease_id, + goal_id = COALESCE(thread_schedule_runs.goal_id, excluded.goal_id) +WHERE thread_schedule_runs.status = 'running' +RETURNING +"#, + ); + let row = sqlx::query(sqlx::AssertSqlSafe(sql)) + .bind(lease_id) + .bind(now_ms) + .bind(run_id) + .fetch_one(&mut *tx) + .await?; + let run = thread_schedule_run_from_row(&row)?; + tx.commit().await?; + Ok(Some(run)) + } + + pub async fn extend_thread_schedule_lease( + &self, + params: ThreadScheduleRunLeaseParams<'_>, + ) -> anyhow::Result { + let ThreadScheduleRunLeaseParams { + schedule_id, + run_id, + lease_id, + now, + lease_duration, + } = params; + let now_ms = datetime_to_epoch_millis(now); + let lease_expires_at = now + chrono::Duration::from_std(lease_duration)?; + let result = sqlx::query( + r#" +UPDATE thread_schedules +SET lease_expires_at_ms = ?, updated_at_ms = ? +WHERE schedule_id = ? + AND status = 'active' + AND lease_id = ? + AND lease_expires_at_ms > ? + AND (expires_at_ms IS NULL OR expires_at_ms > ?) + AND EXISTS ( + SELECT 1 + FROM thread_schedule_occurrences + WHERE thread_schedule_occurrences.schedule_id = thread_schedules.schedule_id + AND thread_schedule_occurrences.occurrence_id = ? + AND thread_schedule_occurrences.state IN ('enqueued', 'started') + AND ( + thread_schedule_occurrences.state = 'enqueued' + OR EXISTS ( + SELECT 1 + FROM thread_schedule_runs + WHERE thread_schedule_runs.run_id = thread_schedule_occurrences.occurrence_id + AND thread_schedule_runs.lease_id = ? + AND thread_schedule_runs.status = 'running' + ) + ) + ) + "#, + ) + .bind(datetime_to_epoch_millis(lease_expires_at)) + .bind(now_ms) + .bind(schedule_id) + .bind(lease_id) + .bind(now_ms) + .bind(now_ms) + .bind(run_id) + .bind(lease_id) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } +} diff --git a/codex-rs/state/src/runtime/schedules/occurrence_tests.rs b/codex-rs/state/src/runtime/schedules/occurrence_tests.rs new file mode 100644 index 000000000..717c8f809 --- /dev/null +++ b/codex-rs/state/src/runtime/schedules/occurrence_tests.rs @@ -0,0 +1,719 @@ +use super::*; + +#[tokio::test] +async fn legacy_active_run_insert_fails_closed_without_matching_occurrence() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(/*id*/ 52); + upsert_test_thread(&runtime, thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let schedule = + create_interval_schedule(&runtime, thread_id, "rollback safety", Some(now)).await; + + let error = sqlx::query( + r#" +INSERT INTO thread_schedule_runs ( +run_id, +schedule_id, +thread_id, +status, +lease_id, +scheduled_for_ms, +started_at_ms +) VALUES (?, ?, ?, 'leased', ?, ?, ?) + "#, + ) + .bind("legacy-fresh-run") + .bind(schedule.schedule_id.as_str()) + .bind(thread_id.to_string()) + .bind("legacy-lease") + .bind(datetime_to_epoch_millis(now)) + .bind(datetime_to_epoch_millis(now)) + .execute(runtime.pool.as_ref()) + .await + .expect_err("an older binary must not create active work without an occurrence"); + assert!( + error + .to_string() + .contains("active schedule occurrence must be reused"), + "unexpected rollback guard error: {error}" + ); + + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "new-runtime-lease", Duration::from_secs(30)) + .await + .expect("new runtime claim should succeed") + .expect("schedule should claim"); + let run = enqueue_and_start_claim( + &runtime, + &claim, + None, + "new runtime input", + now, + Duration::from_secs(30), + ) + .await; + assert_eq!(claim.run.run_id, run.run_id); + assert_eq!(crate::ThreadScheduleRunStatus::Running, run.status); +} + +#[tokio::test] +async fn legacy_schedule_hold_cannot_resurrect_a_pending_occurrence_after_roll_forward() { + for (index, phase, held_status) in [ + (0, "waiting", crate::ThreadScheduleStatus::Paused), + (1, "enqueued", crate::ThreadScheduleStatus::Expired), + (2, "started", crate::ThreadScheduleStatus::Expired), + ] { + let runtime = test_runtime().await; + let thread_id = test_thread_id(/*id*/ 53 + index); + upsert_test_thread(&runtime, thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let schedule = + create_interval_schedule(&runtime, thread_id, &format!("rollback {phase}"), Some(now)) + .await; + let lease_id = format!("legacy-{phase}-lease"); + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, lease_id.as_str(), Duration::from_secs(300)) + .await + .expect("initial claim should succeed") + .expect("schedule should claim"); + if phase == "enqueued" { + runtime + .thread_schedules() + .enqueue_thread_schedule_run(ThreadScheduleRunEnqueueParams { + schedule_id: schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: lease_id.as_str(), + goal_id: None, + auth_profile_recorded: false, + auth_profile: None, + turn_input: "accepted before downgrade", + now, + }) + .await + .expect("occurrence should enqueue") + .expect("enqueue should retain ownership"); + } else if phase == "started" { + enqueue_and_start_claim( + &runtime, + &claim, + None, + "started before downgrade", + now, + Duration::from_secs(300), + ) + .await; + // Match the older expiration order: terminalize the legacy run + // first while the schedule still carries its lease, then clear + // the schedule lease in a separate statement. + sqlx::query( + r#" +UPDATE thread_schedule_runs +SET status = 'failed', error = 'legacy expiry', completed_at_ms = ? +WHERE run_id = ? AND status = 'running' + "#, + ) + .bind(datetime_to_epoch_millis(now + chrono::Duration::seconds(1))) + .bind(claim.run.run_id.as_str()) + .execute(runtime.pool.as_ref()) + .await + .expect("legacy run terminal update should succeed"); + } + + let held_at = now + chrono::Duration::seconds(2); + sqlx::query( + r#" +UPDATE thread_schedules +SET status = ?, +next_run_at_ms = NULL, +lease_id = NULL, +lease_expires_at_ms = NULL, +updated_at_ms = ? +WHERE schedule_id = ? + "#, + ) + .bind(held_status.as_str()) + .bind(datetime_to_epoch_millis(held_at)) + .bind(schedule.schedule_id.as_str()) + .execute(runtime.pool.as_ref()) + .await + .expect("legacy schedule hold should succeed"); + + let occurrence_count: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM thread_schedule_occurrences WHERE schedule_id = ?", + ) + .bind(schedule.schedule_id.as_str()) + .fetch_one(runtime.pool.as_ref()) + .await + .expect("occurrence count should load"); + assert_eq!(0, occurrence_count, "stale {phase} occurrence survived"); + let held_run = runtime + .thread_schedules() + .get_thread_schedule_run(claim.run.run_id.as_str()) + .await + .expect("legacy run lookup should succeed"); + if phase == "waiting" { + assert_eq!(None, held_run, "idle waiting is not a durable run"); + } else { + assert_eq!( + Some(crate::ThreadScheduleRunStatus::Failed), + held_run.map(|run| run.status), + "accepted {phase} work should be explicitly terminal" + ); + } + + let resume_at = now + chrono::Duration::minutes(5); + runtime + .thread_schedules() + .resume_thread_schedule_at(schedule.schedule_id.as_str(), resume_at) + .await + .expect("roll-forward resume should succeed") + .expect("held schedule should still exist"); + let resumed_claim = runtime + .thread_schedules() + .claim_due_thread_schedule( + resume_at, + &format!("roll-forward-{phase}"), + Duration::from_secs(300), + ) + .await + .expect("roll-forward claim should succeed") + .expect("resumed schedule should claim"); + assert_ne!( + claim.run.run_id, resumed_claim.run.run_id, + "roll-forward must create a new occurrence after legacy {phase} termination" + ); + assert_eq!( + ThreadScheduleOccurrenceState::WaitingIdle, + resumed_claim.occurrence_state + ); + } +} + +#[tokio::test] +async fn claim_due_thread_schedule_recovers_started_occurrence_without_duplicate_run() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("state db should initialize"); + let thread_id = test_thread_id(/*id*/ 44); + upsert_test_thread(runtime.as_ref(), thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let schedule = + create_interval_schedule(runtime.as_ref(), thread_id, "restart retry", Some(now)).await; + let original_claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-before-restart", Duration::from_secs(30)) + .await + .expect("initial claim should succeed") + .expect("schedule should claim"); + enqueue_and_start_claim( + &runtime, + &original_claim, + None, + "restart input", + now, + Duration::from_secs(30), + ) + .await; + drop(runtime); + + let reopened = StateRuntime::init(codex_home, "test-provider".to_string()) + .await + .expect("state db should reopen after process restart"); + let retry_at = now + chrono::Duration::seconds(31); + let retry_claim = reopened + .thread_schedules() + .claim_due_thread_schedule(retry_at, "lease-after-restart", Duration::from_secs(30)) + .await + .expect("expired run recovery should succeed") + .expect("expired non-goal run should retry exactly once"); + + let recovered_run = reopened + .thread_schedules() + .get_thread_schedule_run(&original_claim.run.run_id) + .await + .expect("recovered run should load") + .expect("recovered run should exist"); + assert_eq!( + crate::ThreadScheduleRunStatus::Running, + recovered_run.status + ); + assert_eq!( + crate::ThreadScheduleRunStatus::Running, + retry_claim.run.status + ); + assert_eq!( + original_claim.run.scheduled_for, + retry_claim.run.scheduled_for + ); + assert_eq!(original_claim.run.run_id, retry_claim.run.run_id); + assert_eq!(original_claim.run.turn_id, retry_claim.run.turn_id); + let stats = reopened + .thread_schedules() + .get_thread_schedule_stats(&schedule.schedule_id) + .await + .expect("schedule stats should load"); + assert_eq!(1, stats.total_runs); + assert_eq!(0, stats.leased_runs); + assert_eq!(1, stats.running_runs); + assert_eq!(0, stats.failed_runs); + assert!( + reopened + .thread_schedules() + .claim_due_thread_schedule(retry_at, "lease-duplicate-retry", Duration::from_secs(30),) + .await + .expect("duplicate claim check should succeed") + .is_none(), + "one expired lease may create at most one replacement claim" + ); +} + +#[tokio::test] +async fn claim_due_thread_schedule_recovers_waiting_idle_occurrence_after_restart() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("state db should initialize"); + let thread_id = test_thread_id(/*id*/ 51); + upsert_test_thread(runtime.as_ref(), thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let retry_at = now + chrono::Duration::seconds(30); + let schedule = + create_interval_schedule(runtime.as_ref(), thread_id, "waiting restart", Some(now)).await; + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-waiting", Duration::from_secs(30)) + .await + .expect("initial claim should succeed") + .expect("schedule should claim"); + assert!( + runtime + .thread_schedules() + .wait_thread_schedule_run_for_idle( + schedule.schedule_id.as_str(), + claim.run.run_id.as_str(), + claim.run.lease_id.as_str(), + retry_at, + now + chrono::Duration::seconds(1), + ) + .await + .expect("waiting occurrence should persist its retry") + ); + drop(runtime); + + let reopened = StateRuntime::init(codex_home, "test-provider".to_string()) + .await + .expect("state db should reopen"); + let recovered = reopened + .thread_schedules() + .claim_due_thread_schedule(retry_at, "lease-waiting-recovery", Duration::from_secs(30)) + .await + .expect("waiting recovery should succeed") + .expect("waiting occurrence should be reclaimed"); + assert_eq!(claim.run.run_id, recovered.run.run_id); + assert_eq!(claim.run.turn_id, recovered.run.turn_id); + assert_eq!( + ThreadScheduleOccurrenceState::WaitingIdle, + recovered.occurrence_state + ); + assert!(recovered.turn_input.is_none()); + assert_eq!( + crate::ThreadScheduleStats::default(), + reopened + .thread_schedules() + .get_thread_schedule_stats(schedule.schedule_id.as_str()) + .await + .expect("stats should load") + ); +} + +#[tokio::test] +async fn claim_due_thread_schedule_recovers_enqueued_input_profile_and_stable_ids() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("state db should initialize"); + let thread_id = test_thread_id(/*id*/ 48); + upsert_test_thread(runtime.as_ref(), thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let schedule = + create_interval_schedule(runtime.as_ref(), thread_id, "enqueued restart", Some(now)).await; + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-enqueued", Duration::from_secs(30)) + .await + .expect("initial claim should succeed") + .expect("schedule should claim"); + runtime + .thread_schedules() + .enqueue_thread_schedule_run(ThreadScheduleRunEnqueueParams { + schedule_id: schedule.schedule_id.as_str(), + run_id: claim.run.run_id.as_str(), + lease_id: claim.run.lease_id.as_str(), + goal_id: None, + auth_profile_recorded: true, + auth_profile: Some("alternate-profile"), + turn_input: "persisted enqueued input", + now, + }) + .await + .expect("occurrence should enqueue") + .expect("occurrence should retain its lease"); + let retry_at = now + chrono::Duration::seconds(30); + assert!( + runtime + .thread_schedules() + .wait_thread_schedule_run_for_idle( + schedule.schedule_id.as_str(), + claim.run.run_id.as_str(), + claim.run.lease_id.as_str(), + retry_at, + now + chrono::Duration::seconds(1), + ) + .await + .expect("enqueued occurrence should wait for idle") + ); + drop(runtime); + + let reopened = StateRuntime::init(codex_home, "test-provider".to_string()) + .await + .expect("state db should reopen"); + assert!( + reopened + .thread_schedules() + .claim_due_thread_schedule( + retry_at - chrono::Duration::seconds(1), + "lease-enqueued-too-early", + Duration::from_secs(30), + ) + .await + .expect("early enqueued recovery should not fail") + .is_none(), + "idle retry must honor the occurrence retry time" + ); + let recovered = reopened + .thread_schedules() + .claim_due_thread_schedule(retry_at, "lease-enqueued-recovery", Duration::from_secs(30)) + .await + .expect("enqueued recovery should succeed") + .expect("enqueued occurrence should be reclaimed"); + assert_eq!(claim.run.run_id, recovered.run.run_id); + assert_eq!(claim.run.turn_id, recovered.run.turn_id); + assert_eq!( + ThreadScheduleOccurrenceState::Enqueued, + recovered.occurrence_state + ); + assert_eq!( + Some("persisted enqueued input".to_string()), + recovered.turn_input + ); + assert_eq!( + Some(Some("alternate-profile".to_string())), + recovered.occurrence_auth_profile + ); + assert_eq!( + crate::ThreadScheduleStats::default(), + reopened + .thread_schedules() + .get_thread_schedule_stats(schedule.schedule_id.as_str()) + .await + .expect("stats should load"), + "Enqueued is not a durable run" + ); +} + +#[tokio::test] +async fn terminal_recovery_finalizes_interval_cron_and_once_cadence_once() { + let now = at(/*seconds*/ 1_700_000_000); + let cases = [ + ( + "interval", + crate::ThreadScheduleSpec::Interval(crate::ThreadScheduleInterval { + amount: 5, + unit: crate::ThreadScheduleIntervalUnit::Minutes, + }), + Some(now + chrono::Duration::minutes(5)), + crate::ThreadScheduleStatus::Active, + ), + ( + "cron", + crate::ThreadScheduleSpec::Cron { + expression: "*/5 * * * *".to_string(), + }, + Some(now + chrono::Duration::minutes(5)), + crate::ThreadScheduleStatus::Active, + ), + ( + "once", + crate::ThreadScheduleSpec::Once, + None, + crate::ThreadScheduleStatus::Expired, + ), + ]; + + for (name, schedule_spec, next_run_at, expected_status) in cases { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("state db should initialize"); + let thread_id = test_thread_id(/*id*/ 49); + upsert_test_thread(runtime.as_ref(), thread_id).await; + let schedule = runtime + .thread_schedules() + .create_thread_schedule(ThreadScheduleCreateParams { + thread_id, + prompt: format!("{name} terminal recovery"), + prompt_source: crate::ThreadSchedulePromptSource::Inline, + schedule: schedule_spec, + timezone: "UTC".to_string(), + status: crate::ThreadScheduleStatus::Active, + next_run_at: Some(now), + expires_at: None, + }) + .await + .expect("schedule should create"); + let lease_id = format!("lease-{name}"); + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, lease_id.as_str(), Duration::from_secs(30)) + .await + .expect("schedule should claim") + .expect("schedule should be due"); + enqueue_and_start_claim(&runtime, &claim, None, name, now, Duration::from_secs(30)).await; + let completed_at = now + chrono::Duration::seconds(1); + assert!( + runtime + .thread_schedules() + .record_thread_schedule_run_terminal( + schedule.schedule_id.as_str(), + claim.run.run_id.as_str(), + claim.run.lease_id.as_str(), + completed_at, + None, + None, + ) + .await + .expect("terminal outcome should persist") + ); + let before_finalization = runtime + .thread_schedules() + .get_thread_schedule(schedule.schedule_id.as_str()) + .await + .expect("schedule should load") + .expect("schedule should exist"); + assert_eq!(Some(now), before_finalization.next_run_at); + assert_eq!(None, before_finalization.last_run_at); + drop(runtime); + + let reopened = StateRuntime::init(codex_home, "test-provider".to_string()) + .await + .expect("state db should reopen after terminal persistence"); + let recovery_lease_id = format!("lease-{name}-recovery"); + let recovered = reopened + .thread_schedules() + .claim_due_thread_schedule( + now + chrono::Duration::seconds(31), + recovery_lease_id.as_str(), + Duration::from_secs(30), + ) + .await + .expect("terminal recovery claim should succeed") + .expect("terminal occurrence should be reclaimed"); + assert_eq!(claim.run.run_id, recovered.run.run_id); + assert_eq!( + ThreadScheduleOccurrenceState::Terminal, + recovered.occurrence_state + ); + assert!( + reopened + .thread_schedules() + .finalize_terminal_thread_schedule_run( + schedule.schedule_id.as_str(), + recovered.run.run_id.as_str(), + recovered.run.lease_id.as_str(), + completed_at, + next_run_at, + None, + ) + .await + .expect("terminal finalization should succeed") + ); + assert!( + !reopened + .thread_schedules() + .finalize_terminal_thread_schedule_run( + schedule.schedule_id.as_str(), + recovered.run.run_id.as_str(), + recovered.run.lease_id.as_str(), + completed_at, + next_run_at, + None, + ) + .await + .expect("replayed finalization should be idempotent") + ); + let finalized = reopened + .thread_schedules() + .get_thread_schedule(schedule.schedule_id.as_str()) + .await + .expect("schedule should load") + .expect("schedule should exist"); + assert_eq!(expected_status, finalized.status); + assert_eq!(next_run_at, finalized.next_run_at); + assert_eq!(Some(completed_at), finalized.last_run_at); + let stats = reopened + .thread_schedules() + .get_thread_schedule_stats(schedule.schedule_id.as_str()) + .await + .expect("stats should load"); + assert_eq!(1, stats.total_runs); + assert_eq!(1, stats.completed_runs); + } +} + +#[tokio::test] +async fn fatal_pre_start_failure_creates_one_failed_terminal_run() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(/*id*/ 50); + upsert_test_thread(&runtime, thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let schedule = + create_interval_schedule(&runtime, thread_id, "fatal preparation", Some(now)).await; + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-fatal", Duration::from_secs(30)) + .await + .expect("schedule should claim") + .expect("schedule should be due"); + let completed_at = now + chrono::Duration::seconds(1); + assert!( + runtime + .thread_schedules() + .fail_thread_schedule_occurrence_before_start( + schedule.schedule_id.as_str(), + claim.run.run_id.as_str(), + claim.run.lease_id.as_str(), + completed_at, + None, + "prompt source unavailable".to_string(), + ) + .await + .expect("pre-start failure should persist") + ); + assert!( + runtime + .thread_schedules() + .finalize_terminal_thread_schedule_run( + schedule.schedule_id.as_str(), + claim.run.run_id.as_str(), + claim.run.lease_id.as_str(), + completed_at, + Some(now + chrono::Duration::minutes(5)), + None, + ) + .await + .expect("failed terminal should finalize") + ); + let run = runtime + .thread_schedules() + .get_thread_schedule_run(claim.run.run_id.as_str()) + .await + .expect("run should load") + .expect("failed run should exist"); + assert_eq!(crate::ThreadScheduleRunStatus::Failed, run.status); + assert_eq!(Some("prompt source unavailable".to_string()), run.error); + let stats = runtime + .thread_schedules() + .get_thread_schedule_stats(schedule.schedule_id.as_str()) + .await + .expect("stats should load"); + assert_eq!(1, stats.total_runs); + assert_eq!(1, stats.failed_runs); + assert_eq!(0, stats.deferred_runs); +} + +#[tokio::test] +async fn claim_due_thread_schedule_returns_started_run_before_honoring_held_goal() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("state db should initialize"); + let thread_id = test_thread_id(/*id*/ 45); + upsert_test_thread(runtime.as_ref(), thread_id).await; + let goal = runtime + .thread_goals() + .replace_thread_goal( + thread_id, + "hold after restart", + crate::ThreadGoalStatus::Blocked, + /*token_budget*/ None, + ) + .await + .expect("blocked goal should persist"); + let now = at(/*seconds*/ 1_700_000_000); + let schedule = + create_interval_schedule(runtime.as_ref(), thread_id, "hold after restart", Some(now)) + .await; + let original_claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-goal-restart", Duration::from_secs(30)) + .await + .expect("initial claim should succeed") + .expect("schedule should claim"); + enqueue_and_start_claim( + &runtime, + &original_claim, + Some(&goal.goal_id), + "goal restart input", + now, + Duration::from_secs(30), + ) + .await; + drop(runtime); + + let reopened = StateRuntime::init(codex_home, "test-provider".to_string()) + .await + .expect("state db should reopen after process restart"); + let retry_at = now + chrono::Duration::seconds(31); + let recovered = reopened + .thread_schedules() + .claim_due_thread_schedule(retry_at, "lease-held-replacement", Duration::from_secs(30)) + .await + .expect("expired goal run recovery should succeed") + .expect("started run must reach durable rollout recovery before goal handling"); + assert_eq!(original_claim.run.run_id, recovered.run.run_id); + assert_eq!( + crate::ThreadScheduleOccurrenceState::Started, + recovered.occurrence_state + ); + + let held_schedule = reopened + .thread_schedules() + .get_thread_schedule(&schedule.schedule_id) + .await + .expect("schedule should load") + .expect("schedule should exist"); + assert_eq!(crate::ThreadScheduleStatus::Active, held_schedule.status); + assert_eq!(Some(now), held_schedule.next_run_at); + assert_eq!(Some(recovered.run.lease_id.clone()), held_schedule.lease_id); + let original_run = reopened + .thread_schedules() + .get_thread_schedule_run(&original_claim.run.run_id) + .await + .expect("original run should load") + .expect("original run should exist"); + assert_eq!(crate::ThreadScheduleRunStatus::Running, original_run.status); + assert_eq!(Some(goal.goal_id), original_run.goal_id); + assert_eq!(None, original_run.completed_at); + let stats = reopened + .thread_schedules() + .get_thread_schedule_stats(&schedule.schedule_id) + .await + .expect("schedule stats should load"); + assert_eq!(1, stats.total_runs); + assert_eq!(0, stats.leased_runs); + assert_eq!(1, stats.running_runs); + assert_eq!(0, stats.failed_runs); +} From caf083f5bf3ff62e0778cce21e75648d261b7cc8 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 8 Aug 2026 17:01:26 +0300 Subject: [PATCH 02/20] test(state): expose occurrence finalization fence Move the occurrence migration after the live schema sequence and repair the three compiler blockers so the focused regression can execute remotely. Agent: cossus --- ...0066_thread_schedule_occurrence_state.sql} | 0 .../src/runtime/schedules/occurrence/claim.rs | 7 +- .../schedules/{ => tests}/occurrence_tests.rs | 86 +++++++++++++++++++ 3 files changed, 91 insertions(+), 2 deletions(-) rename codex-rs/state/migrations/{0063_thread_schedule_occurrence_state.sql => 0066_thread_schedule_occurrence_state.sql} (100%) rename codex-rs/state/src/runtime/schedules/{ => tests}/occurrence_tests.rs (90%) diff --git a/codex-rs/state/migrations/0063_thread_schedule_occurrence_state.sql b/codex-rs/state/migrations/0066_thread_schedule_occurrence_state.sql similarity index 100% rename from codex-rs/state/migrations/0063_thread_schedule_occurrence_state.sql rename to codex-rs/state/migrations/0066_thread_schedule_occurrence_state.sql diff --git a/codex-rs/state/src/runtime/schedules/occurrence/claim.rs b/codex-rs/state/src/runtime/schedules/occurrence/claim.rs index 94ca955e4..192ea919c 100644 --- a/codex-rs/state/src/runtime/schedules/occurrence/claim.rs +++ b/codex-rs/state/src/runtime/schedules/occurrence/claim.rs @@ -532,7 +532,7 @@ INSERT INTO thread_schedule_occurrences ( })) } - async fn occurrence_run( + pub(super) async fn occurrence_run( tx: &mut sqlx::Transaction<'_, Sqlite>, occurrence: &ThreadScheduleOccurrenceRow, lease_id: &str, @@ -556,7 +556,10 @@ INSERT INTO thread_schedule_occurrences ( turn_id: Some(occurrence.turn_id.clone()), goal_id: occurrence.goal_id.clone(), error: None, - scheduled_for: optional_epoch_millis_to_datetime(occurrence.scheduled_for_ms)?, + scheduled_for: occurrence + .scheduled_for_ms + .map(epoch_millis_to_datetime) + .transpose()?, started_at: epoch_millis_to_datetime(occurrence.created_at_ms)?, completed_at: None, }) diff --git a/codex-rs/state/src/runtime/schedules/occurrence_tests.rs b/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs similarity index 90% rename from codex-rs/state/src/runtime/schedules/occurrence_tests.rs rename to codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs index 717c8f809..a6801e5f0 100644 --- a/codex-rs/state/src/runtime/schedules/occurrence_tests.rs +++ b/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs @@ -191,6 +191,92 @@ WHERE schedule_id = ? } } +#[tokio::test] +async fn terminal_once_occurrence_survives_the_legacy_schedule_hold_trigger() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(/*id*/ 57); + upsert_test_thread(&runtime, thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let schedule = runtime + .thread_schedules() + .create_thread_schedule(ThreadScheduleCreateParams { + thread_id, + prompt: "once trigger fencing".to_string(), + prompt_source: crate::ThreadSchedulePromptSource::Inline, + schedule: crate::ThreadScheduleSpec::Once, + timezone: "UTC".to_string(), + status: crate::ThreadScheduleStatus::Active, + next_run_at: Some(now), + expires_at: None, + }) + .await + .expect("once schedule should create"); + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-once-trigger", Duration::from_secs(30)) + .await + .expect("once schedule should claim") + .expect("once schedule should be due"); + enqueue_and_start_claim( + &runtime, + &claim, + None, + "once trigger fencing", + now, + Duration::from_secs(30), + ) + .await; + let completed_at = now + chrono::Duration::seconds(1); + assert!( + runtime + .thread_schedules() + .record_thread_schedule_run_terminal( + schedule.schedule_id.as_str(), + claim.run.run_id.as_str(), + claim.run.lease_id.as_str(), + completed_at, + None, + None, + ) + .await + .expect("terminal outcome should persist") + ); + + let result = sqlx::query( + r#" +UPDATE thread_schedules +SET status = 'expired', + lease_id = NULL, + lease_expires_at_ms = NULL, + last_run_at_ms = ?, + next_run_at_ms = NULL, + updated_at_ms = ? +WHERE schedule_id = ? AND lease_id = ? + "#, + ) + .bind(datetime_to_epoch_millis(completed_at)) + .bind(datetime_to_epoch_millis(completed_at)) + .bind(schedule.schedule_id.as_str()) + .bind(claim.run.lease_id.as_str()) + .execute(runtime.pool.as_ref()) + .await + .expect("new-runtime once finalization step should update the schedule"); + assert_eq!(1, result.rows_affected()); + + let occurrence_state: Option = sqlx::query_scalar( + "SELECT state FROM thread_schedule_occurrences WHERE occurrence_id = ?", + ) + .bind(claim.run.run_id.as_str()) + .fetch_optional(runtime.pool.as_ref()) + .await + .expect("terminal occurrence should load"); + assert_eq!( + Some("terminal".to_string()), + occurrence_state, + "the compatibility trigger must leave terminal work for fenced runtime finalization" + ); +} + #[tokio::test] async fn claim_due_thread_schedule_recovers_started_occurrence_without_duplicate_run() { let codex_home = unique_temp_dir(); From d28191b9b77ba4c392a391f44f6c82eaacc3adc0 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 8 Aug 2026 17:07:18 +0300 Subject: [PATCH 03/20] test(state): disambiguate occurrence assertions Import the repository-standard pretty assertion macro explicitly in the nested occurrence regression module. Agent: cossus --- codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs b/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs index a6801e5f0..4412fce1b 100644 --- a/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs +++ b/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs @@ -1,4 +1,5 @@ use super::*; +use pretty_assertions::assert_eq; #[tokio::test] async fn legacy_active_run_insert_fails_closed_without_matching_occurrence() { From 6497f2384d5f6a075462724654bb79e4b8394da3 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 8 Aug 2026 17:10:50 +0300 Subject: [PATCH 04/20] fix(state): preserve terminal occurrence fencing Restrict the legacy pause/expiry compatibility trigger to non-terminal occurrences so the current runtime remains responsible for its fenced terminal delete. Agent: cossus --- .../0066_thread_schedule_occurrence_state.sql | 1 + .../src/runtime/schedules/tests/occurrence_tests.rs | 13 ++++++------- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/codex-rs/state/migrations/0066_thread_schedule_occurrence_state.sql b/codex-rs/state/migrations/0066_thread_schedule_occurrence_state.sql index 6b5212e29..d3a5110d9 100644 --- a/codex-rs/state/migrations/0066_thread_schedule_occurrence_state.sql +++ b/codex-rs/state/migrations/0066_thread_schedule_occurrence_state.sql @@ -146,6 +146,7 @@ WHEN NEW.status IN ('paused', 'expired') SELECT 1 FROM thread_schedule_occurrences WHERE thread_schedule_occurrences.schedule_id = NEW.schedule_id + AND thread_schedule_occurrences.state != 'terminal' ) BEGIN UPDATE thread_schedule_runs diff --git a/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs b/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs index 4412fce1b..40009715a 100644 --- a/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs +++ b/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs @@ -264,13 +264,12 @@ WHERE schedule_id = ? AND lease_id = ? .expect("new-runtime once finalization step should update the schedule"); assert_eq!(1, result.rows_affected()); - let occurrence_state: Option = sqlx::query_scalar( - "SELECT state FROM thread_schedule_occurrences WHERE occurrence_id = ?", - ) - .bind(claim.run.run_id.as_str()) - .fetch_optional(runtime.pool.as_ref()) - .await - .expect("terminal occurrence should load"); + let occurrence_state: Option = + sqlx::query_scalar("SELECT state FROM thread_schedule_occurrences WHERE occurrence_id = ?") + .bind(claim.run.run_id.as_str()) + .fetch_optional(runtime.pool.as_ref()) + .await + .expect("terminal occurrence should load"); assert_eq!( Some("terminal".to_string()), occurrence_state, From 09d2d8a9c1bf7cce3761f153c6f72424e555ee26 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 8 Aug 2026 17:29:41 +0300 Subject: [PATCH 05/20] fix(app-server): expose occurrence terminal helpers Keep the split terminal implementation private to the schedule runtime while making its re-exports visible across the nested occurrence module boundary. Agent: cossus --- .../occurrence/terminal.rs | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs index 309aa40e4..0833c0720 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs @@ -14,7 +14,7 @@ pub(crate) struct PersistedScheduledTurnTerminal { pub(crate) error: Option, } -pub(super) fn persisted_scheduled_turn_terminal( +pub(in super::super) fn persisted_scheduled_turn_terminal( history: &InitialHistory, turn_id: &str, fallback_completed_at: DateTime, @@ -85,7 +85,7 @@ pub(super) fn persisted_scheduled_turn_terminal( }) } -pub(super) fn scheduled_turn_finish(event: &EventMsg) -> Option { +pub(in super::super) fn scheduled_turn_finish(event: &EventMsg) -> Option { match event { EventMsg::TurnComplete(completed) if completed @@ -108,11 +108,13 @@ pub(super) fn scheduled_turn_finish(event: &EventMsg) -> Option) -> Option> { +pub(in super::super::super) fn default_thread_schedule_expires_at( + now: DateTime, +) -> Option> { now.checked_add_signed(ChronoDuration::days(DEFAULT_SCHEDULE_EXPIRATION_DAYS)) } -pub(super) fn next_thread_schedule_run_at( +pub(in super::super::super) fn next_thread_schedule_run_at( schedule: &codex_state::ThreadScheduleSpec, timezone: &str, after: DateTime, @@ -143,7 +145,7 @@ pub(super) fn next_thread_schedule_run_at( Ok(next) } -pub(super) fn next_thread_schedule_run_after_completion( +pub(in super::super) fn next_thread_schedule_run_after_completion( schedule: &codex_state::ThreadScheduleSpec, timezone: &str, scheduled_for: Option>, @@ -191,11 +193,13 @@ pub(super) fn next_thread_schedule_run_after_completion( next_thread_schedule_run_at(schedule, timezone, completed_at) } -pub(super) fn normalize_schedule_timezone(timezone: &str) -> anyhow::Result { +pub(in super::super::super) fn normalize_schedule_timezone( + timezone: &str, +) -> anyhow::Result { parse_schedule_timezone(timezone).map(|timezone| timezone.name().to_string()) } -pub(super) async fn finish_scheduled_run_after_turn( +pub(in super::super::super) async fn finish_scheduled_run_after_turn( thread_id: ThreadId, scheduled_run: crate::thread_state::ScheduledThreadScheduleRun, event: &EventMsg, @@ -247,7 +251,7 @@ pub(super) async fn finish_scheduled_run_after_turn( } } -pub(super) async fn recover_scheduled_run_for_terminal_turn( +pub(in super::super::super) async fn recover_scheduled_run_for_terminal_turn( state_db: &StateDbHandle, thread_id: ThreadId, turn_id: &str, From df1f662f1ebf8f3a6864ecd109410401d0f27b28 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 8 Aug 2026 17:46:09 +0300 Subject: [PATCH 06/20] fix(app-server): widen occurrence re-export hop Expose only the five terminal helpers consumed by sibling request processors through the intermediate occurrence module. Agent: cossus --- .../thread_schedule_runtime/occurrence.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence.rs index 3e862505e..d8c55ee33 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence.rs @@ -7,13 +7,13 @@ mod terminal; pub(super) use terminal::PersistedScheduledTurnTerminal; pub(super) use terminal::ScheduledTurnFinish; -pub(super) use terminal::default_thread_schedule_expires_at; -pub(super) use terminal::finish_scheduled_run_after_turn; +pub(in super::super) use terminal::default_thread_schedule_expires_at; +pub(in super::super) use terminal::finish_scheduled_run_after_turn; pub(super) use terminal::next_thread_schedule_run_after_completion; -pub(super) use terminal::next_thread_schedule_run_at; -pub(super) use terminal::normalize_schedule_timezone; +pub(in super::super) use terminal::next_thread_schedule_run_at; +pub(in super::super) use terminal::normalize_schedule_timezone; pub(super) use terminal::persisted_scheduled_turn_terminal; -pub(super) use terminal::recover_scheduled_run_for_terminal_turn; +pub(in super::super) use terminal::recover_scheduled_run_for_terminal_turn; pub(super) use terminal::scheduled_turn_finish; impl ThreadScheduleRuntime { From 8935e06fb997b2529a5bf314e91df9405197d064 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 8 Aug 2026 17:56:39 +0300 Subject: [PATCH 07/20] test(loop): link occurrence regressions conventionally Place split occurrence tests on conventional nested module paths so both rustc and cargo-shear follow the same linkage. Agent: cossus --- .../app-server/src/request_processors/thread_schedule_runtime.rs | 1 - .../thread_schedule_runtime/{ => tests}/occurrence_tests.rs | 0 codex-rs/state/src/runtime/schedules.rs | 1 - 3 files changed, 2 deletions(-) rename codex-rs/app-server/src/request_processors/thread_schedule_runtime/{ => tests}/occurrence_tests.rs (100%) diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs index 6bc70c64f..7ae8d9515 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs @@ -2468,7 +2468,6 @@ mod tests { ); } - #[path = "occurrence_tests.rs"] mod occurrence_tests; #[tokio::test] diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence_tests.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs similarity index 100% rename from codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence_tests.rs rename to codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs diff --git a/codex-rs/state/src/runtime/schedules.rs b/codex-rs/state/src/runtime/schedules.rs index 2e662e95e..6ed2e65d8 100644 --- a/codex-rs/state/src/runtime/schedules.rs +++ b/codex-rs/state/src/runtime/schedules.rs @@ -1713,7 +1713,6 @@ mod tests { ); } - #[path = "occurrence_tests.rs"] mod occurrence_tests; #[tokio::test] From 998f86c930a1368a0f9d264cc7387630c67332a1 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 8 Aug 2026 18:11:12 +0300 Subject: [PATCH 08/20] test(loop): compile relinked regressions Disambiguate assertions, retain the expected schedule for later lookup, import the Tokio test channel, and gate test-only re-exports. Agent: cossus --- .../src/request_processors/thread_lifecycle/scheduled_runs.rs | 1 + .../request_processors/thread_schedule_runtime/occurrence.rs | 2 ++ .../thread_schedule_runtime/tests/occurrence_tests.rs | 3 ++- 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle/scheduled_runs.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle/scheduled_runs.rs index e2c2c3acf..73d07168c 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle/scheduled_runs.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle/scheduled_runs.rs @@ -38,6 +38,7 @@ mod tests { use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::TurnCompleteEvent; use pretty_assertions::assert_eq; + use tokio::sync::mpsc; #[tokio::test] async fn non_affecting_error_keeps_scheduled_run_running_until_turn_complete() { diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence.rs index d8c55ee33..4960e0e09 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence.rs @@ -6,6 +6,7 @@ mod execution; mod terminal; pub(super) use terminal::PersistedScheduledTurnTerminal; +#[cfg(test)] pub(super) use terminal::ScheduledTurnFinish; pub(in super::super) use terminal::default_thread_schedule_expires_at; pub(in super::super) use terminal::finish_scheduled_run_after_turn; @@ -14,6 +15,7 @@ pub(in super::super) use terminal::next_thread_schedule_run_at; pub(in super::super) use terminal::normalize_schedule_timezone; pub(super) use terminal::persisted_scheduled_turn_terminal; pub(in super::super) use terminal::recover_scheduled_run_for_terminal_turn; +#[cfg(test)] pub(super) use terminal::scheduled_turn_finish; impl ThreadScheduleRuntime { diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs index 6da851f50..46c581dc7 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs @@ -1,4 +1,5 @@ use super::*; +use pretty_assertions::assert_eq; #[tokio::test] async fn idle_rejection_reuses_one_pending_occurrence_without_counting_a_run() { @@ -88,7 +89,7 @@ async fn idle_rejection_reuses_one_pending_occurrence_without_counting_a_run() { lease_id: None, lease_expires_at: None, updated_at: deferred_schedule.updated_at, - ..schedule + ..schedule.clone() }, deferred_schedule ); From a5a681a5e3ecd6e2360f9239e5b219b478489733 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 8 Aug 2026 18:24:55 +0300 Subject: [PATCH 09/20] test(loop): satisfy argument comment lint Agent: cossus --- .../tests/occurrence_tests.rs | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs index 46c581dc7..4f1606f66 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs @@ -337,7 +337,14 @@ async fn stale_started_owner_cannot_submit_after_same_occurrence_is_reclaimed() .await .expect("claim should succeed") .expect("schedule should claim"); - enqueue_and_start_claim(&state_db, &claim, None, now, Duration::from_secs(30)).await; + enqueue_and_start_claim( + &state_db, + &claim, + /*goal_id*/ None, + now, + Duration::from_secs(30), + ) + .await; let contender = codex_state::StateRuntime::init( temp_dir.path().to_path_buf(), @@ -452,7 +459,7 @@ fn persisted_completed_scheduled_turn_is_terminal_for_the_matching_turn_only() { let history = resumed_history_with_turn_events( thread_id, [ - turn_started("turn-scheduled", 1_700_000_000), + turn_started("turn-scheduled", /*started_at*/ 1_700_000_000), EventMsg::TurnComplete(TurnCompleteEvent { turn_id: "turn-scheduled".to_string(), last_agent_message: Some("done".to_string()), @@ -486,7 +493,7 @@ fn persisted_failed_scheduled_turn_keeps_the_replayed_failure() { let history = resumed_history_with_turn_events( thread_id, [ - turn_started("turn-scheduled", 1_700_000_000), + turn_started("turn-scheduled", /*started_at*/ 1_700_000_000), EventMsg::Error(ErrorEvent { message: "model failed".to_string(), codex_error_info: None, @@ -520,7 +527,7 @@ fn persisted_aborted_scheduled_turn_is_an_explicit_failure() { let history = resumed_history_with_turn_events( thread_id, [ - turn_started("turn-scheduled", 1_700_000_000), + turn_started("turn-scheduled", /*started_at*/ 1_700_000_000), EventMsg::TurnAborted(TurnAbortedEvent { turn_id: Some("turn-scheduled".to_string()), reason: TurnAbortReason::Interrupted, @@ -575,7 +582,7 @@ fn persisted_in_progress_scheduled_turn_is_not_terminal() { let history = resumed_history_with_turn_events( thread_id, [ - turn_started("turn-scheduled", 1_700_000_000), + turn_started("turn-scheduled", /*started_at*/ 1_700_000_000), EventMsg::Error(ErrorEvent { message: "rollback request failed".to_string(), codex_error_info: Some(CoreCodexErrorInfo::ThreadRollbackFailed), From 7a9e227568fd8e9425410c2fd219f6f6f391d709 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 8 Aug 2026 18:34:38 +0300 Subject: [PATCH 10/20] fix(loop): honor occurrence terminal lifecycle Agent: cossus --- .../thread_schedule_runtime.rs | 49 ++++++++++- .../occurrence/terminal.rs | 83 +++++++++++++------ .../tests/occurrence_tests.rs | 18 +++- 3 files changed, 122 insertions(+), 28 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs index 7ae8d9515..6a1452ce7 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs @@ -2631,6 +2631,14 @@ mod tests { .expect("retry should claim deferred schedule"); assert_eq!(Some(wait.retry_at), retry_claim.run.scheduled_for); let completed_at = wait.retry_at + chrono::Duration::seconds(5); + enqueue_and_start_claim( + &state_db, + &retry_claim, + /*goal_id*/ None, + wait.retry_at, + Duration::from_secs(300), + ) + .await; let (finished_schedule, finished_run) = finish_scheduled_run_state( &state_db, @@ -2714,6 +2722,14 @@ mod tests { .expect("claim should succeed") .expect("schedule should claim"); let completed_at = scheduled_for + chrono::Duration::seconds(5); + enqueue_and_start_claim( + &state_db, + &claim, + /*goal_id*/ None, + scheduled_for, + Duration::from_secs(300), + ) + .await; let (finished_schedule, finished_run) = finish_scheduled_run_state( &state_db, @@ -2809,6 +2825,14 @@ mod tests { .expect("claim should succeed") .expect("schedule should claim"); let completed_at = scheduled_for + chrono::Duration::seconds(5); + enqueue_and_start_claim( + &state_db, + &claim, + Some(goal.goal_id.as_str()), + scheduled_for, + Duration::from_secs(300), + ) + .await; let (finished_schedule, finished_run) = finish_scheduled_run_state( &state_db, @@ -2900,6 +2924,11 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); + let turn_id = claim + .run + .turn_id + .clone() + .expect("claimed occurrence should have a stable turn id"); enqueue_and_start_claim( &state_db, &claim, @@ -2917,7 +2946,7 @@ mod tests { .await .expect("state db should reopen"); let recovered = - recover_scheduled_run_for_terminal_turn(&reopened, thread_id, "turn-after-restart") + recover_scheduled_run_for_terminal_turn(&reopened, thread_id, turn_id.as_str()) .await .expect("run recovery should succeed") .expect("running schedule should recover after restart"); @@ -2946,7 +2975,7 @@ mod tests { ); assert_eq!(Some(completed_at), finished.1.completed_at); assert!( - recover_scheduled_run_for_terminal_turn(&reopened, thread_id, "turn-after-restart",) + recover_scheduled_run_for_terminal_turn(&reopened, thread_id, turn_id.as_str()) .await .expect("completed run lookup should succeed") .is_none() @@ -3033,6 +3062,14 @@ mod tests { .await .expect("claim should succeed") .expect("schedule should claim"); + enqueue_and_start_claim( + &state_db, + &claim, + Some(goal.goal_id.as_str()), + scheduled_for, + Duration::from_secs(300), + ) + .await; let (finished_schedule, finished_run) = finish_scheduled_run_state( &state_db, @@ -3121,6 +3158,14 @@ mod tests { .expect("claim should succeed") .expect("schedule should claim"); let completed_at = scheduled_for + chrono::Duration::seconds(5); + enqueue_and_start_claim( + &state_db, + &claim, + /*goal_id*/ None, + scheduled_for, + Duration::from_secs(300), + ) + .await; let (finished_schedule, finished_run) = finish_scheduled_run_state( &state_db, diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs index 0833c0720..f0ed94a20 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs @@ -50,34 +50,67 @@ pub(in super::super) fn persisted_scheduled_turn_terminal( .completed_at .and_then(|timestamp| DateTime::::from_timestamp(timestamp, 0)) .unwrap_or(fallback_completed_at); - let error = match turn.status { - codex_app_server_protocol::TurnStatus::Completed => { - let finish = rollout_items - .iter() - .filter_map(|item| match item { - RolloutItem::EventMsg(event @ EventMsg::TurnComplete(completed)) - if completed.turn_id == turn_id => - { - scheduled_turn_finish(event) - } - _ => None, - }) - .last()?; - match finish { - ScheduledTurnFinish::Complete => None, - ScheduledTurnFinish::Failed(error) => Some(error), + let mut replaying_turn = false; + let mut replayed_error = None; + for item in &rollout_items { + match item { + RolloutItem::EventMsg(EventMsg::TurnStarted(started)) => { + if replaying_turn { + break; + } + replaying_turn = started.turn_id == turn_id; + } + RolloutItem::EventMsg(EventMsg::Error(error)) + if replaying_turn && error.affects_turn_status() => + { + replayed_error = Some(schedule_turn_event_error(error)); + break; + } + RolloutItem::EventMsg(EventMsg::TurnComplete(completed)) + if replaying_turn && completed.turn_id == turn_id => + { + break; } + RolloutItem::EventMsg(EventMsg::TurnAborted(aborted)) + if replaying_turn && aborted.turn_id.as_deref() == Some(turn_id) => + { + break; + } + _ => {} } - codex_app_server_protocol::TurnStatus::Failed => Some( - turn.error - .as_ref() - .map(schedule_turn_error) - .unwrap_or_else(|| schedule_run_error("scheduled turn failed")), - ), - codex_app_server_protocol::TurnStatus::Interrupted => { - Some(schedule_run_error("scheduled turn was interrupted")) + } + let error = if let Some(error) = replayed_error { + Some(error) + } else { + match turn.status { + codex_app_server_protocol::TurnStatus::Completed => { + let finish = rollout_items + .iter() + .filter_map(|item| match item { + RolloutItem::EventMsg(event @ EventMsg::TurnComplete(completed)) + if completed.turn_id == turn_id => + { + scheduled_turn_finish(event) + } + _ => None, + }) + .last()?; + match finish { + ScheduledTurnFinish::Complete => None, + ScheduledTurnFinish::Failed(error) => Some(error), + } + } + codex_app_server_protocol::TurnStatus::Failed => Some( + turn.error + .as_ref() + .map(schedule_turn_error) + .unwrap_or_else(|| schedule_run_error("scheduled turn failed")), + ), + codex_app_server_protocol::TurnStatus::Interrupted => { + Some(schedule_run_error("scheduled turn was interrupted")) + } + codex_app_server_protocol::TurnStatus::InProgress => return None, } - codex_app_server_protocol::TurnStatus::InProgress => return None, }; Some(PersistedScheduledTurnTerminal { completed_at, diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs index 4f1606f66..6cd1dfb64 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/tests/occurrence_tests.rs @@ -459,7 +459,23 @@ fn persisted_completed_scheduled_turn_is_terminal_for_the_matching_turn_only() { let history = resumed_history_with_turn_events( thread_id, [ + turn_started("turn-other", /*started_at*/ 1_699_999_990), + EventMsg::Error(ErrorEvent { + message: "other turn failed".to_string(), + codex_error_info: None, + }), + EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: "turn-other".to_string(), + last_agent_message: None, + completed_at: Some(1_699_999_999), + duration_ms: Some(9_000), + time_to_first_token_ms: None, + }), turn_started("turn-scheduled", /*started_at*/ 1_700_000_000), + EventMsg::Error(ErrorEvent { + message: "non-terminal rollback warning".to_string(), + codex_error_info: Some(CoreCodexErrorInfo::ThreadRollbackFailed), + }), EventMsg::TurnComplete(TurnCompleteEvent { turn_id: "turn-scheduled".to_string(), last_agent_message: Some("done".to_string()), @@ -483,7 +499,7 @@ fn persisted_completed_scheduled_turn_is_terminal_for_the_matching_turn_only() { ); assert_eq!( None, - persisted_scheduled_turn_terminal(&history, "turn-other", at(/*seconds*/ 1_700_000_999),) + persisted_scheduled_turn_terminal(&history, "turn-missing", at(/*seconds*/ 1_700_000_999),) ); } From fc317a03f0d088388d1a49da11a1f1764007e026 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 8 Aug 2026 18:44:46 +0300 Subject: [PATCH 11/20] test(state): satisfy occurrence argument lint Agent: cossus --- codex-rs/state/src/runtime/schedules.rs | 32 ++++++++--------- .../schedules/tests/occurrence_tests.rs | 34 ++++++++++++------- 2 files changed, 37 insertions(+), 29 deletions(-) diff --git a/codex-rs/state/src/runtime/schedules.rs b/codex-rs/state/src/runtime/schedules.rs index 6ed2e65d8..ca8013c51 100644 --- a/codex-rs/state/src/runtime/schedules.rs +++ b/codex-rs/state/src/runtime/schedules.rs @@ -1604,7 +1604,7 @@ mod tests { enqueue_and_start_claim( &runtime, &claim, - None, + /*goal_id*/ None, "once input", now, Duration::from_secs(300), @@ -2024,7 +2024,7 @@ mod tests { enqueue_and_start_claim( &runtime, &claim, - None, + /*goal_id*/ None, "long running input", now, Duration::from_secs(300), @@ -2086,7 +2086,7 @@ mod tests { enqueue_and_start_claim( &runtime, &claim, - None, + /*goal_id*/ None, "stale input", now, Duration::from_secs(30), @@ -2503,7 +2503,7 @@ mod tests { enqueue_and_start_claim( &runtime, &claim, - None, + /*goal_id*/ None, "terminal race input", now, Duration::from_secs(30), @@ -2627,7 +2627,7 @@ mod tests { let replacement_run = enqueue_and_start_claim( &runtime, &recovered, - None, + /*goal_id*/ None, "recovered waiting input", replacement_started_at, Duration::from_secs(30), @@ -2782,7 +2782,7 @@ mod tests { enqueue_and_start_claim( &runtime, &complete_claim, - None, + /*goal_id*/ None, "late complete input", now, Duration::from_secs(300), @@ -2844,7 +2844,7 @@ mod tests { enqueue_and_start_claim( &runtime, &defer_claim, - None, + /*goal_id*/ None, "late defer input", now, Duration::from_secs(300), @@ -2961,7 +2961,7 @@ mod tests { let running = enqueue_and_start_claim( &runtime, &completed_claim, - None, + /*goal_id*/ None, "completed input", now, Duration::from_secs(300), @@ -3046,7 +3046,7 @@ mod tests { enqueue_and_start_claim( &runtime, &failed_claim, - None, + /*goal_id*/ None, "failed input", now, Duration::from_secs(300), @@ -3133,7 +3133,7 @@ mod tests { enqueue_and_start_claim( &runtime, &claim, - None, + /*goal_id*/ None, "resume failure input", now, Duration::from_secs(300), @@ -3197,7 +3197,7 @@ mod tests { enqueue_and_start_claim( &runtime, &claim, - None, + /*goal_id*/ None, "update failure input", now, Duration::from_secs(300), @@ -3421,7 +3421,7 @@ mod tests { enqueue_and_start_claim( &runtime, &claim_one, - None, + /*goal_id*/ None, "completed stats input", now, Duration::from_secs(300), @@ -3479,7 +3479,7 @@ mod tests { enqueue_and_start_claim( &runtime, &claim_three, - None, + /*goal_id*/ None, "failed stats input", third_run_at, Duration::from_secs(300), @@ -3632,7 +3632,7 @@ mod tests { enqueue_and_start_claim( &runtime, &claim, - None, + /*goal_id*/ None, "live input", now, Duration::from_secs(300), @@ -3743,7 +3743,7 @@ mod tests { enqueue_and_start_claim( &runtime, &paused_claim, - None, + /*goal_id*/ None, "paused input", now, Duration::from_secs(300), @@ -3792,7 +3792,7 @@ mod tests { enqueue_and_start_claim( &runtime, &expiring_claim, - None, + /*goal_id*/ None, "expiring input", now, Duration::from_secs(30), diff --git a/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs b/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs index 40009715a..c67878f2b 100644 --- a/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs +++ b/codex-rs/state/src/runtime/schedules/tests/occurrence_tests.rs @@ -48,7 +48,7 @@ started_at_ms let run = enqueue_and_start_claim( &runtime, &claim, - None, + /*goal_id*/ None, "new runtime input", now, Duration::from_secs(30), @@ -99,7 +99,7 @@ async fn legacy_schedule_hold_cannot_resurrect_a_pending_occurrence_after_roll_f enqueue_and_start_claim( &runtime, &claim, - None, + /*goal_id*/ None, "started before downgrade", now, Duration::from_secs(300), @@ -221,7 +221,7 @@ async fn terminal_once_occurrence_survives_the_legacy_schedule_hold_trigger() { enqueue_and_start_claim( &runtime, &claim, - None, + /*goal_id*/ None, "once trigger fencing", now, Duration::from_secs(30), @@ -236,8 +236,8 @@ async fn terminal_once_occurrence_survives_the_legacy_schedule_hold_trigger() { claim.run.run_id.as_str(), claim.run.lease_id.as_str(), completed_at, - None, - None, + /*expected_goal_id*/ None, + /*error*/ None, ) .await .expect("terminal outcome should persist") @@ -297,7 +297,7 @@ async fn claim_due_thread_schedule_recovers_started_occurrence_without_duplicate enqueue_and_start_claim( &runtime, &original_claim, - None, + /*goal_id*/ None, "restart input", now, Duration::from_secs(30), @@ -567,7 +567,15 @@ async fn terminal_recovery_finalizes_interval_cron_and_once_cadence_once() { .await .expect("schedule should claim") .expect("schedule should be due"); - enqueue_and_start_claim(&runtime, &claim, None, name, now, Duration::from_secs(30)).await; + enqueue_and_start_claim( + &runtime, + &claim, + /*goal_id*/ None, + name, + now, + Duration::from_secs(30), + ) + .await; let completed_at = now + chrono::Duration::seconds(1); assert!( runtime @@ -577,8 +585,8 @@ async fn terminal_recovery_finalizes_interval_cron_and_once_cadence_once() { claim.run.run_id.as_str(), claim.run.lease_id.as_str(), completed_at, - None, - None, + /*expected_goal_id*/ None, + /*error*/ None, ) .await .expect("terminal outcome should persist") @@ -621,7 +629,7 @@ async fn terminal_recovery_finalizes_interval_cron_and_once_cadence_once() { recovered.run.lease_id.as_str(), completed_at, next_run_at, - None, + /*expected_goal_id*/ None, ) .await .expect("terminal finalization should succeed") @@ -635,7 +643,7 @@ async fn terminal_recovery_finalizes_interval_cron_and_once_cadence_once() { recovered.run.lease_id.as_str(), completed_at, next_run_at, - None, + /*expected_goal_id*/ None, ) .await .expect("replayed finalization should be idempotent") @@ -682,7 +690,7 @@ async fn fatal_pre_start_failure_creates_one_failed_terminal_run() { claim.run.run_id.as_str(), claim.run.lease_id.as_str(), completed_at, - None, + /*goal_id*/ None, "prompt source unavailable".to_string(), ) .await @@ -697,7 +705,7 @@ async fn fatal_pre_start_failure_creates_one_failed_terminal_run() { claim.run.lease_id.as_str(), completed_at, Some(now + chrono::Duration::minutes(5)), - None, + /*expected_goal_id*/ None, ) .await .expect("failed terminal should finalize") From d1c9b8381c0b5698db62243d92634053c5a1ecdc Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 8 Aug 2026 19:49:59 +0300 Subject: [PATCH 12/20] fix(core): serialize task start mailbox snapshot Complete the installed turn start snapshot before releasing its worker so trigger mail queued after spawn_task belongs to the next turn. Agent: cossus --- codex-rs/core/src/tasks/mod.rs | 57 +++++++++++++++++----------------- 1 file changed, 28 insertions(+), 29 deletions(-) diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 163097045..269e70d9c 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -454,7 +454,6 @@ impl Session { let ctx = Arc::clone(&turn_context); let task_for_run = Arc::clone(&task); let task_input = input; - let task_turn_state = Arc::clone(&turn_state); let (start_tx, start_rx) = oneshot::channel(); let task_cancellation_token = cancellation_token.child_token(); // Task-owned turn spans keep a core-owned span open for the @@ -486,34 +485,6 @@ impl Session { return; } let sess = session_ctx.clone_session(); - let started_at = Instant::now(); - let turn_started_at_unix_ms = ctx - .turn_timing_state - .mark_turn_started(started_at) - .await; - ctx.turn_metadata_state - .set_turn_started_at_unix_ms(turn_started_at_unix_ms); - let token_usage_at_turn_start = - sess.total_token_usage().await.unwrap_or_default(); - sess.services - .guardian_rejection_circuit_breaker - .lock() - .await - .clear_turn(&ctx.sub_id); - let pending_items = sess - .input_queue - .get_pending_input_for_turn_state(task_turn_state.as_ref()) - .await; - task_turn_state.lock().await.token_usage_at_turn_start = - token_usage_at_turn_start.clone(); - sess.input_queue - .extend_pending_input_for_turn_state( - task_turn_state.as_ref(), - pending_items, - ) - .await; - sess.emit_turn_start_lifecycle(ctx.as_ref(), &token_usage_at_turn_start) - .await; if task_cancellation_token.is_cancelled() { done_clone.notify_one(); return; @@ -575,6 +546,34 @@ impl Session { } }; if installed { + // Finish the start snapshot before releasing the worker or returning: callers may + // enqueue trigger mail immediately after `spawn_task`, and that mail belongs to the + // next turn rather than the one that was just installed. + let started_at = Instant::now(); + let turn_started_at_unix_ms = turn_context + .turn_timing_state + .mark_turn_started(started_at) + .await; + turn_context + .turn_metadata_state + .set_turn_started_at_unix_ms(turn_started_at_unix_ms); + let token_usage_at_turn_start = self.total_token_usage().await.unwrap_or_default(); + self.services + .guardian_rejection_circuit_breaker + .lock() + .await + .clear_turn(&turn_context.sub_id); + let pending_items = self + .input_queue + .get_pending_input_for_turn_state(turn_state.as_ref()) + .await; + turn_state.lock().await.token_usage_at_turn_start = + token_usage_at_turn_start.clone(); + self.input_queue + .extend_pending_input_for_turn_state(turn_state.as_ref(), pending_items) + .await; + self.emit_turn_start_lifecycle(turn_context.as_ref(), &token_usage_at_turn_start) + .await; let _ = start_tx.send(()); } installed From f77d901eed18a1abb6d1b7c365088cc8e04d4a69 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 8 Aug 2026 20:55:56 +0300 Subject: [PATCH 13/20] fix(state): fence terminal writes by lease expiry Agent: cossus --- codex-rs/state/src/runtime/schedules.rs | 57 ++++++++----------- .../runtime/schedules/occurrence/finish.rs | 4 ++ 2 files changed, 28 insertions(+), 33 deletions(-) diff --git a/codex-rs/state/src/runtime/schedules.rs b/codex-rs/state/src/runtime/schedules.rs index ca8013c51..780456707 100644 --- a/codex-rs/state/src/runtime/schedules.rs +++ b/codex-rs/state/src/runtime/schedules.rs @@ -2514,26 +2514,27 @@ mod tests { .await .expect("contending state runtime should initialize"); let retry_at = now + chrono::Duration::seconds(31); - let completion = runtime.thread_schedules().complete_thread_schedule_run( - &schedule.schedule_id, - &claim.run.run_id, - "lease-terminal-race", - retry_at, - Some(now + chrono::Duration::hours(1)), - ); - let replacement = contender.thread_schedules().claim_due_thread_schedule( - retry_at, - "lease-reaper-race", - Duration::from_secs(30), - ); - let (completion, replacement) = tokio::join!(completion, replacement); - let completion = completion.expect("late completion should not error"); - let replacement = replacement.expect("expired lease reaper should not error"); - assert_ne!( - completion, - replacement.is_some(), - "either the terminal event or the reaper may own the old lease, never both" + let completion = runtime + .thread_schedules() + .complete_thread_schedule_run( + &schedule.schedule_id, + &claim.run.run_id, + "lease-terminal-race", + retry_at, + Some(now + chrono::Duration::hours(1)), + ) + .await + .expect("late completion should not error"); + assert!( + !completion, + "a terminal event at or after lease expiry must be fenced out" ); + let replacement = contender + .thread_schedules() + .claim_due_thread_schedule(retry_at, "lease-reaper-race", Duration::from_secs(30)) + .await + .expect("expired lease reaper should not error") + .expect("expired lease reaper should reclaim the original occurrence"); let original_run = runtime .thread_schedules() @@ -2541,25 +2542,15 @@ mod tests { .await .expect("original run should load") .expect("original run should exist"); - assert_eq!( - if completion { - crate::ThreadScheduleRunStatus::Completed - } else { - crate::ThreadScheduleRunStatus::Running - }, - original_run.status, - "the same run is either terminalized or reclaimed without replacement" - ); - if let Some(replacement) = replacement.as_ref() { - assert_eq!(claim.run.run_id, replacement.run.run_id); - assert_eq!(claim.run.turn_id, replacement.run.turn_id); - } + assert_eq!(crate::ThreadScheduleRunStatus::Running, original_run.status); + assert_eq!(claim.run.run_id, replacement.run.run_id); + assert_eq!(claim.run.turn_id, replacement.run.turn_id); let stats = runtime .thread_schedules() .get_thread_schedule_stats(&schedule.schedule_id) .await .expect("schedule stats should load"); - assert_eq!(i64::from(replacement.is_some()), stats.running_runs); + assert_eq!(1, stats.running_runs); assert_eq!(0, stats.leased_runs); assert_eq!(1, stats.total_runs); } diff --git a/codex-rs/state/src/runtime/schedules/occurrence/finish.rs b/codex-rs/state/src/runtime/schedules/occurrence/finish.rs index 8ce5aa6d8..10a4bdb1c 100644 --- a/codex-rs/state/src/runtime/schedules/occurrence/finish.rs +++ b/codex-rs/state/src/runtime/schedules/occurrence/finish.rs @@ -34,6 +34,7 @@ WHERE schedule_id = ? FROM thread_schedules WHERE thread_schedules.schedule_id = thread_schedule_runs.schedule_id AND thread_schedules.lease_id = ? + AND thread_schedules.lease_expires_at_ms > ? ) AND EXISTS ( SELECT 1 @@ -52,6 +53,7 @@ WHERE schedule_id = ? .bind(expected_goal_id) .bind(expected_goal_id) .bind(lease_id) + .bind(completed_at_ms) .execute(&mut *tx) .await?; if run_result.rows_affected() == 0 { @@ -70,6 +72,7 @@ SELECT EXISTS( AND thread_schedule_runs.status IN ('completed', 'failed') AND thread_schedule_occurrences.state = 'terminal' AND thread_schedules.lease_id = ? + AND thread_schedules.lease_expires_at_ms > ? AND (? IS NULL OR thread_schedule_runs.goal_id IS NULL OR thread_schedule_runs.goal_id = ?) ) "#, @@ -78,6 +81,7 @@ SELECT EXISTS( .bind(run_id) .bind(lease_id) .bind(lease_id) + .bind(completed_at_ms) .bind(expected_goal_id) .bind(expected_goal_id) .fetch_one(&mut *tx) From 2d28c89b34891ab31040e4290ade24f396de160c Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sat, 8 Aug 2026 21:46:28 +0300 Subject: [PATCH 14/20] test(app-server): keep scheduled-run lease current Agent: cossus --- .../request_processors/thread_lifecycle/scheduled_runs.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle/scheduled_runs.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle/scheduled_runs.rs index 73d07168c..41395afab 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle/scheduled_runs.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle/scheduled_runs.rs @@ -50,8 +50,7 @@ mod tests { .await .expect("state db should initialize"); let thread_id = ThreadId::new(); - let now = chrono::DateTime::::from_timestamp(1_700_000_000, 0) - .expect("test timestamp should be valid"); + let now = Utc::now(); let mut builder = ThreadMetadataBuilder::new( thread_id, temp_dir.path().join("thread.jsonl"), @@ -158,7 +157,7 @@ mod tests { let complete = EventMsg::TurnComplete(TurnCompleteEvent { turn_id: turn_id.clone(), last_agent_message: Some("finished after rollback warning".to_string()), - completed_at: Some(1_700_000_001), + completed_at: Some(now.timestamp() + 1), duration_ms: Some(1_000), time_to_first_token_ms: Some(100), }); From 1f0597f20061ae05d0805846246e894984e357ec Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sun, 9 Aug 2026 00:33:49 +0300 Subject: [PATCH 15/20] test(app-server): assemble credential fixtures at runtime Agent: cossus --- .../src/message_processor_schedule_tests.rs | 9 +++- .../thread_schedule_runtime.rs | 54 ++++++++++++------- 2 files changed, 41 insertions(+), 22 deletions(-) diff --git a/codex-rs/app-server/src/message_processor_schedule_tests.rs b/codex-rs/app-server/src/message_processor_schedule_tests.rs index bc62cb088..69329b30b 100644 --- a/codex-rs/app-server/src/message_processor_schedule_tests.rs +++ b/codex-rs/app-server/src/message_processor_schedule_tests.rs @@ -801,13 +801,18 @@ fn message_for_test_connection(envelope: OutgoingEnvelope) -> Option String { + ["s", "k-test-schedule-secret"].concat() +} + async fn create_mock_responses_server_unauthorized() -> MockServer { let server = MockServer::start().await; + let api_key = schedule_test_api_key(); Mock::given(method("POST")) .and(path_regex(".*/responses$")) .respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({ "error": { - "message": "Incorrect API key provided: sk-test-schedule-secret", + "message": format!("Incorrect API key provided: {api_key}"), "type": "invalid_request_error", "param": null, "code": "invalid_api_key" @@ -2883,7 +2888,7 @@ fn thread_schedule_run_now_records_model_errors_as_failed_runs() -> Result<()> { "schedule run error should discard raw provider auth messages: {error}" ); assert!( - !error.contains("sk-test-schedule-secret"), + !error.contains(schedule_test_api_key().as_str()), "schedule run error should redact API keys: {error}" ); assert!(failed.run.completed_at.is_some()); diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs index 6a1452ce7..c24cc2ad5 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime.rs @@ -1724,6 +1724,10 @@ mod tests { DateTime::::from_timestamp(seconds, 0).expect("valid timestamp") } + fn credential_fixture(parts: &[&str]) -> String { + parts.concat() + } + async fn enqueue_and_start_claim( state_db: &codex_state::StateRuntime, claim: &codex_state::ThreadScheduleClaim, @@ -3556,8 +3560,10 @@ mod tests { #[test] fn scheduled_turn_usage_limit_error_is_classified_and_redacted() { + let api_key_name = credential_fixture(&["OPENAI_API_", "KEY"]); + let api_key = credential_fixture(&["s", "k-test-secret"]); let finish = scheduled_turn_finish(&EventMsg::Error(ErrorEvent { - message: "You've hit your usage limit. OPENAI_API_KEY=sk-test-secret".to_string(), + message: format!("You've hit your usage limit. {api_key_name}={api_key}"), codex_error_info: Some(CoreCodexErrorInfo::UsageLimitExceeded), })); @@ -3615,27 +3621,32 @@ mod tests { #[test] fn redacts_sensitive_schedule_run_error_values() { - let sanitized = schedule_run_error( - "failed with OPENAI_API_KEY=sk-test-secret token: plain-secret Bearer sk-bearer-secret", - ); + let api_key_name = credential_fixture(&["OPENAI_API_", "KEY"]); + let api_key = credential_fixture(&["s", "k-test-secret"]); + let plain_secret = credential_fixture(&["plain", "-secret"]); + let bearer_token = credential_fixture(&["s", "k-bearer-secret"]); + let sanitized = schedule_run_error(format!( + "failed with {api_key_name}={api_key} token: {plain_secret} Bearer {bearer_token}" + )); assert_eq!( "failed with OPENAI_API_KEY=[redacted] token: [redacted] Bearer [redacted]", sanitized ); - assert!(!sanitized.contains("sk-test-secret")); - assert!(!sanitized.contains("plain-secret")); - assert!(!sanitized.contains("sk-bearer-secret")); + assert!(!sanitized.contains(api_key.as_str())); + assert!(!sanitized.contains(plain_secret.as_str())); + assert!(!sanitized.contains(bearer_token.as_str())); } #[test] fn redacts_short_prefixed_api_keys_from_schedule_run_errors() { - let sanitized = schedule_run_error( - "unexpected status 401 Unauthorized: Incorrect API key provided: sk-work.", - ); + let api_key = credential_fixture(&["s", "k-work"]); + let sanitized = schedule_run_error(format!( + "unexpected status 401 Unauthorized: Incorrect API key provided: {api_key}." + )); assert!(sanitized.contains("[redacted]")); - assert!(!sanitized.contains("sk-work")); + assert!(!sanitized.contains(api_key.as_str())); } #[test] @@ -3643,18 +3654,21 @@ mod tests { // Standalone credentials that are NOT `sk-` prefixed must still be // stripped from persisted run errors and notifications. for secret in [ - "ghp_0123456789abcdefghijABCDEFGHIJ01", - "github_pat_11ABCDEFG0abcdefghijKLMNOP", - "AKIAIOSFODNN7EXAMPLE", - "AIzaSyA1234567890abcdefghijklmnopqrs", - "xai-abcdef0123456789abcdef0123", - "glpat-abcdefghij0123456789", - "xoxb-1234567890-abcdefghijkl", - "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0In0.dozjgNryP4J3jVmNHl0w5N_XgL0n3I9P", + credential_fixture(&["g", "hp_0123456789abcdefghijABCDEFGHIJ01"]), + credential_fixture(&["github", "_pat_11ABCDEFG0abcdefghijKLMNOP"]), + credential_fixture(&["A", "KIAIOSFODNN7EXAMPLE"]), + credential_fixture(&["A", "IzaSyA1234567890abcdefghijklmnopqrs"]), + credential_fixture(&["x", "ai-abcdef0123456789abcdef0123"]), + credential_fixture(&["g", "lpat-abcdefghij0123456789"]), + credential_fixture(&["x", "oxb-1234567890-abcdefghijkl"]), + credential_fixture(&[ + "ey", + "JhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0In0.dozjgNryP4J3jVmNHl0w5N_XgL0n3I9P", + ]), ] { let sanitized = schedule_run_error(format!("run failed near {secret} end")); assert!( - !sanitized.contains(secret), + !sanitized.contains(secret.as_str()), "expected `{secret}` to be redacted, got: {sanitized}" ); } From 9c36376bbc816135397b58b65c4b8bf68979cbf5 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sun, 9 Aug 2026 00:55:41 +0300 Subject: [PATCH 16/20] Fix exact-head Bazel clippy failures Agent: cossus --- .../occurrence/terminal.rs | 4 +- codex-rs/core/src/tasks/mod.rs | 49 ++----------------- 2 files changed, 6 insertions(+), 47 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs index f0ed94a20..da12e1bb5 100644 --- a/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs +++ b/codex-rs/app-server/src/request_processors/thread_schedule_runtime/occurrence/terminal.rs @@ -34,7 +34,7 @@ pub(in super::super) fn persisted_scheduled_turn_terminal( } _ => None, }) - .last()?; + .next_back()?; return Some(PersistedScheduledTurnTerminal { completed_at: aborted .completed_at @@ -94,7 +94,7 @@ pub(in super::super) fn persisted_scheduled_turn_terminal( } _ => None, }) - .last()?; + .next_back()?; match finish { ScheduledTurnFinish::Complete => None, ScheduledTurnFinish::Failed(error) => Some(error), diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 269e70d9c..cb51df2ee 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -378,44 +378,6 @@ impl Session { } } - // The explicit return type keeps the `Send` contract visible for lifecycle - // code that may await this helper inside spawned tasks. - #[allow(clippy::manual_async_fn)] - pub(crate) fn start_task( - self: &Arc, - turn_context: Arc, - input: Vec, - task: T, - ) -> impl std::future::Future> + Send + '_ { - async move { - if let Err(error) = self - .task_policy_ready(turn_context.as_ref(), task.kind()) - .await - { - self.send_event_raw(Event { - id: turn_context.sub_id.clone(), - msg: EventMsg::Error(error.to_error_event(/*message_prefix*/ None)), - }) - .await; - return Err(error); - } - if !self - .start_task_after_policy_preflight( - turn_context, - input, - task, - /*expected_turn_state*/ None, - ) - .await - { - return Err(codex_protocol::error::CodexErr::InvalidRequest( - "thread became active before task start".to_string(), - )); - } - Ok(()) - } - } - // Keep the explicit `Send` bound visible: this helper participates in a // recursive task-start path that can be awaited from a spawned turn. #[allow(clippy::manual_async_fn)] @@ -651,13 +613,10 @@ impl Session { // Empty active turns are reservations, not running work. Trigger-turn mailbox // mail is user/client-directed, so it takes over stale or lower-priority // reservations and lets their owners observe that the reservation was lost. - *active_turn = Some(ActiveTurn::default()); - Arc::clone( - &active_turn - .as_ref() - .expect("pending-work reservation should be present") - .turn_state, - ) + let reservation = ActiveTurn::default(); + let turn_state = Arc::clone(&reservation.turn_state); + *active_turn = Some(reservation); + turn_state }; self.maybe_emit_unknown_model_warning_for_turn(turn_context.as_ref()) From 73745071e884bfadd7e08991b6bab25a9890aa65 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sun, 9 Aug 2026 17:18:05 +0300 Subject: [PATCH 17/20] chore(state): renumber occurrence migration Agent: nausicaa --- ...rrence_state.sql => 0068_thread_schedule_occurrence_state.sql} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename codex-rs/state/migrations/{0066_thread_schedule_occurrence_state.sql => 0068_thread_schedule_occurrence_state.sql} (100%) diff --git a/codex-rs/state/migrations/0066_thread_schedule_occurrence_state.sql b/codex-rs/state/migrations/0068_thread_schedule_occurrence_state.sql similarity index 100% rename from codex-rs/state/migrations/0066_thread_schedule_occurrence_state.sql rename to codex-rs/state/migrations/0068_thread_schedule_occurrence_state.sql From 6fa99395f39ac1dce9b9116dda89e169c3f1e078 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sun, 9 Aug 2026 18:18:45 +0300 Subject: [PATCH 18/20] fix(monitor): require current-fence authorization Agent: nausicaa --- .../thread_monitor_processor.rs | 193 ++++++++++++- .../thread_monitor_runtime.rs | 265 +++++++++++------- codex-rs/core/src/codex_thread.rs | 9 + codex-rs/core/src/exec.rs | 30 ++ codex-rs/core/src/session/session.rs | 6 + .../src/tools/handlers/monitor_control.rs | 143 +++++++++- .../0068_thread_schedule_occurrence_state.sql | 3 + codex-rs/state/src/lib.rs | 1 + codex-rs/state/src/model/mod.rs | 1 + codex-rs/state/src/model/thread_monitor.rs | 197 +++++++++++++ codex-rs/state/src/runtime/monitors.rs | 158 +++++++++-- .../state/src/runtime/workflow_automation.rs | 2 + 12 files changed, 875 insertions(+), 133 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_monitor_processor.rs b/codex-rs/app-server/src/request_processors/thread_monitor_processor.rs index 291b439c5..a37296b0d 100644 --- a/codex-rs/app-server/src/request_processors/thread_monitor_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_monitor_processor.rs @@ -1,5 +1,10 @@ use super::thread_monitor_api::*; use super::*; +use codex_app_server_protocol::CommandExecutionApprovalDecision; +use codex_app_server_protocol::CommandExecutionRequestApprovalParams; +use codex_app_server_protocol::CommandExecutionRequestApprovalResponse; +use codex_app_server_protocol::ServerRequestPayload; +use codex_shell_command::parse_command::shlex_join; #[derive(Clone)] pub(crate) struct ThreadMonitorRequestProcessor { @@ -106,6 +111,17 @@ impl ThreadMonitorRequestProcessor { let state_db = self.prepare_monitor_mutation(thread_id).await?; self.ensure_thread_monitor_capacity(&state_db, thread_id) .await?; + let writer = self.current_monitor_writer(thread_id).await?; + let authorization = self + .authorize_monitor_command( + &request_id, + thread_id, + &writer, + /*generation*/ 0, + command.as_str(), + cwd.as_deref(), + ) + .await?; let monitor = state_db .thread_monitors() .create_thread_monitor(codex_state::ThreadMonitorCreateParams { @@ -117,6 +133,7 @@ impl ThreadMonitorRequestProcessor { routing, output_file, status: codex_state::ThreadMonitorStatus::Running, + authorization: Some(authorization), }) .await .map_err(|err| internal_error(format!("failed to create thread monitor: {err}")))?; @@ -257,9 +274,24 @@ impl ThreadMonitorRequestProcessor { let monitor_id = self .resolve_monitor_id_for_thread(&state_db, thread_id, params.monitor_id.as_str()) .await?; + let monitor = self + .load_monitor_for_thread(&state_db, thread_id, monitor_id.as_str()) + .await?; + let cwd = validate_optional_monitor_relative_path("monitor cwd", monitor.cwd.clone())?; + let writer = self.current_monitor_writer(thread_id).await?; + let authorization = self + .authorize_monitor_command( + &request_id, + thread_id, + &writer, + monitor.generation + 1, + monitor.command.as_str(), + cwd.as_deref(), + ) + .await?; let monitor = state_db .thread_monitors() - .restart_thread_monitor(monitor_id.as_str()) + .restart_thread_monitor(monitor_id.as_str(), monitor.generation, authorization) .await .map_err(|err| internal_error(format!("failed to restart thread monitor: {err}")))? .ok_or_else(|| invalid_request(format!("monitor not found: {monitor_id}")))?; @@ -333,6 +365,128 @@ impl ThreadMonitorRequestProcessor { Ok(()) } + async fn current_monitor_writer( + &self, + thread_id: ThreadId, + ) -> Result, JSONRPCErrorError> { + self.thread_manager + .get_thread(thread_id) + .await + .map_err(|_| { + invalid_request(format!( + "thread must be loaded to authorize monitor commands: {thread_id}" + )) + }) + } + + #[allow(clippy::too_many_arguments)] + async fn authorize_monitor_command( + &self, + request_id: &ConnectionRequestId, + thread_id: ThreadId, + writer: &Arc, + generation: i64, + command: &str, + cwd: Option<&str>, + ) -> Result { + let writer_fence = writer.monitor_writer_fence().to_string(); + let before = writer.config_snapshot().await; + let approval_cwd = monitor_approval_cwd(&before.cwd, cwd)?; + self.request_monitor_command_approval(request_id, thread_id, command, approval_cwd) + .await?; + + let current_writer = self.current_monitor_writer(thread_id).await?; + let after = current_writer.config_snapshot().await; + if writer_fence != current_writer.monitor_writer_fence() + || before.cwd != after.cwd + || before.permission_profile != after.permission_profile + { + return Err(invalid_request( + "monitor command authorization became stale before it could be recorded", + )); + } + Ok(codex_state::ThreadMonitorAuthorization::new( + thread_id, + generation, + command, + cwd, + writer_fence, + after.permission_profile, + after.cwd.display().to_string(), + )) + } + + async fn request_monitor_command_approval( + &self, + request_id: &ConnectionRequestId, + thread_id: ThreadId, + command: &str, + cwd: AbsolutePathBuf, + ) -> Result<(), JSONRPCErrorError> { + let command = codex_core::exec::persistent_shell_command_args(command); + let item_id = format!("monitor-authorization-{}", Uuid::new_v4()); + let params = CommandExecutionRequestApprovalParams { + thread_id: thread_id.to_string(), + turn_id: item_id.clone(), + item_id, + started_at_ms: Utc::now().timestamp_millis(), + approval_id: None, + reason: Some( + "Authorize this persistent monitor command. It may run in the background until stopped." + .to_string(), + ), + network_approval_context: None, + command: Some(shlex_join(&command)), + cwd: Some(cwd), + command_actions: None, + additional_permissions: None, + proposed_execpolicy_amendment: None, + proposed_network_policy_amendments: None, + available_decisions: Some(vec![ + CommandExecutionApprovalDecision::Accept, + CommandExecutionApprovalDecision::Decline, + CommandExecutionApprovalDecision::Cancel, + ]), + }; + let connection_ids = [request_id.connection_id]; + let (_, receiver) = self + .outgoing + .send_request_to_connections( + Some(&connection_ids), + ServerRequestPayload::CommandExecutionRequestApproval(params), + Some(thread_id), + ) + .await; + let response = receiver + .await + .map_err(|err| { + warn!( + thread_id = %thread_id, + "monitor command approval callback closed: {err}" + ); + invalid_request("monitor command approval was not granted") + })? + .map_err(|err| { + warn!( + thread_id = %thread_id, + "monitor command approval failed: {err:?}" + ); + invalid_request("monitor command approval was not granted") + })?; + let response = serde_json::from_value::(response) + .map_err(|err| { + warn!( + thread_id = %thread_id, + "invalid monitor command approval response: {err}" + ); + invalid_request("monitor command approval was not granted") + })?; + if !monitor_command_approval_granted(&response.decision) { + return Err(invalid_request("monitor command approval was not granted")); + } + Ok(()) + } + async fn prepare_monitor_mutation( &self, thread_id: ThreadId, @@ -579,3 +733,40 @@ fn parse_thread_id_for_monitor_request(thread_id: &str) -> Result, +) -> Result { + let cwd = match cwd { + Some(cwd) => thread_cwd.join(cwd), + None => thread_cwd.to_path_buf(), + }; + AbsolutePathBuf::try_from(cwd) + .map_err(|err| invalid_request(format!("monitor cwd must be absolute: {err}"))) +} + +fn monitor_command_approval_granted(decision: &CommandExecutionApprovalDecision) -> bool { + matches!(decision, CommandExecutionApprovalDecision::Accept) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn monitor_command_approval_is_one_shot_and_fails_closed() { + assert!(monitor_command_approval_granted( + &CommandExecutionApprovalDecision::Accept + )); + assert!(!monitor_command_approval_granted( + &CommandExecutionApprovalDecision::AcceptForSession + )); + assert!(!monitor_command_approval_granted( + &CommandExecutionApprovalDecision::Decline + )); + assert!(!monitor_command_approval_granted( + &CommandExecutionApprovalDecision::Cancel + )); + } +} diff --git a/codex-rs/app-server/src/request_processors/thread_monitor_runtime.rs b/codex-rs/app-server/src/request_processors/thread_monitor_runtime.rs index 5df05078f..7db39d7b6 100644 --- a/codex-rs/app-server/src/request_processors/thread_monitor_runtime.rs +++ b/codex-rs/app-server/src/request_processors/thread_monitor_runtime.rs @@ -2,21 +2,14 @@ use super::thread_monitor_api::api_thread_monitor_event_from_state; use super::thread_monitor_api::api_thread_monitor_from_state; use super::*; use codex_protocol::AgentPath; +use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::InterAgentCommunication; -#[cfg(not(target_os = "windows"))] -use codex_shell_command::shell_detect::ShellType; -#[cfg(not(target_os = "windows"))] -use codex_shell_command::shell_detect::get_shell; -#[cfg(not(target_os = "windows"))] -use codex_shell_command::shell_detect::ultimate_fallback_shell; use std::path::Component; use std::path::Path; use std::path::PathBuf; -use std::process::Stdio; use tokio::io::AsyncBufReadExt; use tokio::io::AsyncWriteExt; use tokio::io::BufReader; -use tokio::process::Command; const MONITOR_POLL_INTERVAL: Duration = Duration::from_secs(2); const MAX_MONITOR_EVENT_CHARS: usize = 8_000; @@ -231,25 +224,17 @@ impl ThreadMonitorRuntime { .await; return; }; - self.record_monitor_event( - &state_db, - &monitor, - codex_state::ThreadMonitorEventStream::System, - "monitor process starting", - ) - .await; - - let thread_cwd = match monitor_thread_cwd(&state_db, &monitor).await { - Ok(cwd) => cwd, + let (initial_snapshot, _) = match self.authorized_monitor_execution(&monitor).await { + Ok(context) => context, Err(err) => { - let error = monitor_error(format!("invalid monitor thread cwd: {err}")); + let error = monitor_error(err); self.mark_monitor_failed(&state_db, &monitor, error).await; self.remove_active_monitor(&monitor.monitor_id, monitor.generation) .await; return; } }; - let cwd = match resolve_monitor_cwd(&monitor, thread_cwd.as_path()).await { + let cwd = match resolve_monitor_cwd(&monitor, initial_snapshot.cwd.as_path()).await { Ok(cwd) => cwd, Err(err) => { let error = monitor_error(format!("invalid monitor cwd: {err}")); @@ -259,17 +244,59 @@ impl ThreadMonitorRuntime { return; } }; - let mut command = monitor_command(&monitor.command); - command - .current_dir(&cwd) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .stdin(Stdio::null()); - - let mut child = match command.spawn() { + let (snapshot, config) = match self.authorized_monitor_execution(&monitor).await { + Ok(context) => context, + Err(err) => { + let error = monitor_error(err); + self.mark_monitor_failed(&state_db, &monitor, error).await; + self.remove_active_monitor(&monitor.monitor_id, monitor.generation) + .await; + return; + } + }; + if initial_snapshot.cwd != snapshot.cwd { + let error = monitor_error( + "monitor command authorization became stale while resolving its cwd; restart the monitor to reauthorize" + .to_string(), + ); + self.mark_monitor_failed(&state_db, &monitor, error).await; + self.remove_active_monitor(&monitor.monitor_id, monitor.generation) + .await; + return; + } + let cwd = match AbsolutePathBuf::try_from(cwd) { + Ok(cwd) => cwd, + Err(err) => { + let error = monitor_error(format!("invalid absolute monitor cwd: {err}")); + self.mark_monitor_failed(&state_db, &monitor, error).await; + self.remove_active_monitor(&monitor.monitor_id, monitor.generation) + .await; + return; + } + }; + self.record_monitor_event( + &state_db, + &monitor, + codex_state::ThreadMonitorEventStream::System, + "monitor process starting", + ) + .await; + let env = create_env(&config.shell_environment_policy, Some(monitor.thread_id)); + let mut child = match codex_core::exec::spawn_streaming_command_under_sandbox( + codex_core::exec::persistent_shell_command_args(&monitor.command), + cwd.clone(), + env, + &snapshot.permission_profile, + &snapshot.cwd, + &config.codex_linux_sandbox_exe, + config.features.use_legacy_landlock(), + ) + .await + { Ok(child) => child, Err(err) => { - let error = monitor_error(format!("failed to start monitor command: {err}")); + let error = + monitor_error(format!("failed to start sandboxed monitor command: {err}")); self.mark_monitor_failed(&state_db, &monitor, error).await; self.remove_active_monitor(&monitor.monitor_id, monitor.generation) .await; @@ -298,7 +325,7 @@ impl ThreadMonitorRuntime { let runtime = self.clone(); let state_db = state_db.clone(); let monitor = monitor.clone(); - let cwd = cwd.clone(); + let cwd = cwd.to_path_buf(); let cancel_token = cancel_token.clone(); self.tasks.spawn(async move { runtime @@ -318,7 +345,7 @@ impl ThreadMonitorRuntime { let runtime = self.clone(); let state_db = state_db.clone(); let monitor = monitor.clone(); - let cwd = cwd.clone(); + let cwd = cwd.to_path_buf(); let cancel_token = cancel_token.clone(); self.tasks.spawn(async move { runtime @@ -369,6 +396,34 @@ impl ThreadMonitorRuntime { .await; } + async fn authorized_monitor_execution( + &self, + monitor: &codex_state::ThreadMonitor, + ) -> Result<(ThreadConfigSnapshot, Arc), String> { + let thread = self + .thread_manager + .get_thread(monitor.thread_id) + .await + .map_err(|_| { + "monitor command authorization has no current thread writer; restart the monitor to reauthorize" + .to_string() + })?; + let snapshot = thread.config_snapshot().await; + let thread_cwd = snapshot.cwd.display().to_string(); + if !monitor_authorization_is_current( + monitor, + thread.monitor_writer_fence(), + &snapshot.permission_profile, + thread_cwd.as_str(), + ) { + return Err( + "monitor command authorization is missing or stale; restart the monitor to reauthorize" + .to_string(), + ); + } + Ok((snapshot, thread.config().await)) + } + async fn read_monitor_stream( &self, state_db: StateDbHandle, @@ -635,20 +690,6 @@ impl ThreadMonitorRuntime { } } -async fn monitor_thread_cwd( - state_db: &StateDbHandle, - monitor: &codex_state::ThreadMonitor, -) -> anyhow::Result { - let metadata = state_db - .get_thread(monitor.thread_id) - .await? - .ok_or_else(|| anyhow::anyhow!("monitor thread metadata not found"))?; - if metadata.cwd.as_os_str().is_empty() { - anyhow::bail!("monitor thread cwd is empty"); - } - Ok(metadata.cwd) -} - /// Builds the mailbox communication injected into a thread for a monitor line. /// /// `trigger_turn` is `true` so an idle thread is woken to observe the event @@ -679,26 +720,6 @@ Output: ) } -fn monitor_command(command: &str) -> Command { - #[cfg(target_os = "windows")] - { - let mut cmd = Command::new("cmd"); - cmd.arg("/C").arg(command); - cmd - } - - #[cfg(not(target_os = "windows"))] - { - let shell = get_shell(ShellType::Bash, /*path*/ None) - .or_else(|| get_shell(ShellType::Zsh, /*path*/ None)) - .or_else(|| get_shell(ShellType::Sh, /*path*/ None)) - .unwrap_or_else(ultimate_fallback_shell); - let mut cmd = Command::new(shell.shell_path); - cmd.arg("-lc").arg(command); - cmd - } -} - async fn resolve_monitor_cwd( monitor: &codex_state::ThreadMonitor, fallback: &Path, @@ -715,6 +736,17 @@ async fn resolve_monitor_cwd( Ok(canonical_cwd) } +fn monitor_authorization_is_current( + monitor: &codex_state::ThreadMonitor, + writer_fence: &str, + permission_profile: &PermissionProfile, + thread_cwd: &str, +) -> bool { + monitor.authorization.as_ref().is_some_and(|authorization| { + authorization.authorizes(monitor, writer_fence, permission_profile, thread_cwd) + }) +} + fn resolve_monitor_relative_path( field_name: &str, base: &Path, @@ -755,6 +787,7 @@ fn truncate_chars(value: String, max_chars: usize) -> String { mod tests { use super::*; use pretty_assertions::assert_eq; + use tokio::process::Command; #[test] fn monitor_relative_path_resolver_stays_within_base() { @@ -802,44 +835,23 @@ mod tests { ); } - #[tokio::test] - async fn monitor_thread_cwd_reads_persisted_thread_metadata() -> anyhow::Result<()> { - let tempdir = tempfile::TempDir::new()?; - let app_server_cwd = tempdir.path().join("server"); - let thread_cwd = tempdir.path().join("thread"); - tokio::fs::create_dir_all(&app_server_cwd).await?; - tokio::fs::create_dir_all(&thread_cwd).await?; - let state_db = - codex_state::StateRuntime::init(tempdir.path().join("state"), "test-provider".into()) - .await?; - let thread_id = codex_protocol::ThreadId::new(); - let mut builder = codex_state::ThreadMetadataBuilder::new( - thread_id, - tempdir.path().join("rollout.jsonl"), - chrono::Utc::now(), - codex_protocol::protocol::SessionSource::default(), - ); - builder.cwd = thread_cwd.clone(); - state_db - .upsert_thread(&builder.build("test-provider")) - .await?; - let monitor = test_monitor_for_thread(thread_id, /*cwd*/ None); - - assert_eq!(monitor_thread_cwd(&state_db, &monitor).await?, thread_cwd); - assert_ne!( - monitor_thread_cwd(&state_db, &monitor).await?, - app_server_cwd - ); - Ok(()) - } - #[tokio::test] async fn monitor_command_supports_bash_source_when_bash_is_available() { - if get_shell(ShellType::Bash, /*path*/ None).is_none() { + if codex_shell_command::shell_detect::get_shell( + codex_shell_command::shell_detect::ShellType::Bash, + /*path*/ None, + ) + .is_none() + { return; } - - let output = monitor_command("source /dev/null && printf ok") + let command = + codex_core::exec::persistent_shell_command_args("source /dev/null && printf ok"); + let Some((program, args)) = command.split_first() else { + return; + }; + let output = Command::new(program) + .args(args) .output() .await .expect("monitor command should run"); @@ -852,6 +864,45 @@ mod tests { assert_eq!(String::from_utf8_lossy(&output.stdout), "ok"); } + #[test] + fn monitor_authorization_requires_current_fence_profile_and_subject() { + let permission_profile = PermissionProfile::read_only(); + let monitor = test_monitor(/*cwd*/ None); + assert!(monitor_authorization_is_current( + &monitor, + "writer-fence", + &permission_profile, + "/workspace", + )); + + let mut missing = monitor.clone(); + missing.authorization = None; + assert!(!monitor_authorization_is_current( + &missing, + "writer-fence", + &permission_profile, + "/workspace", + )); + assert!(!monitor_authorization_is_current( + &monitor, + "stale-writer", + &permission_profile, + "/workspace", + )); + assert!(!monitor_authorization_is_current( + &monitor, + "writer-fence", + &PermissionProfile::Disabled, + "/workspace", + )); + assert!(!monitor_authorization_is_current( + &monitor, + "writer-fence", + &permission_profile, + "/different-workspace", + )); + } + #[test] fn monitor_output_communication_uses_wake_if_idle_mailbox_shape() { let monitor = test_monitor(/*cwd*/ None); @@ -882,17 +933,29 @@ mod tests { cwd: Option<&str>, ) -> codex_state::ThreadMonitor { let now = chrono::Utc::now(); + let command = "printf ok"; + let cwd = cwd.map(str::to_string); + let authorization = codex_state::ThreadMonitorAuthorization::new( + thread_id, + /*generation*/ 1, + command, + cwd.as_deref(), + "writer-fence".to_string(), + PermissionProfile::read_only(), + "/workspace".to_string(), + ); codex_state::ThreadMonitor { thread_id, monitor_id: "monitor-id".to_string(), name: "monitor".to_string(), prompt: "watch".to_string(), - command: "printf ok".to_string(), - cwd: cwd.map(str::to_string), + command: command.to_string(), + cwd, routing: codex_state::ThreadMonitorRouting::File, output_file: Some("monitor.log".to_string()), status: codex_state::ThreadMonitorStatus::Running, generation: 1, + authorization: Some(authorization), process_id: None, last_event_at: None, last_error: None, diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index d03a68cff..c6b9cefc6 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -741,6 +741,15 @@ impl CodexThread { self.codex.thread_config_snapshot().await } + /// Returns the process-local writer fence for monitor command authorization. + /// + /// The fence changes whenever this thread is materialized as a new live + /// session, so persisted monitor authorizations cannot be replayed by a + /// later app-server process without an explicit create/restart approval. + pub fn monitor_writer_fence(&self) -> &str { + self.codex.session.monitor_writer_fence() + } + /// Returns the files that supplied the thread's loaded model instructions. pub async fn instruction_sources(&self) -> Vec { self.codex.instruction_sources().await diff --git a/codex-rs/core/src/exec.rs b/codex-rs/core/src/exec.rs index 4d9f58454..e55f8f844 100644 --- a/codex-rs/core/src/exec.rs +++ b/codex-rs/core/src/exec.rs @@ -44,6 +44,12 @@ use codex_sandboxing::SandboxTransformRequest; use codex_sandboxing::SandboxType; use codex_sandboxing::SandboxablePreference; use codex_sandboxing::compatibility_sandbox_policy_for_permission_profile; +#[cfg(not(target_os = "windows"))] +use codex_shell_command::shell_detect::ShellType; +#[cfg(not(target_os = "windows"))] +use codex_shell_command::shell_detect::get_shell; +#[cfg(not(target_os = "windows"))] +use codex_shell_command::shell_detect::ultimate_fallback_shell; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_pty::DEFAULT_OUTPUT_BYTES_CAP; use codex_utils_pty::process_group::kill_child_process_group; @@ -81,6 +87,30 @@ pub(crate) const MAX_EXEC_OUTPUT_DELTAS_PER_CALL: usize = 10_000; // indefinitely, effectively hanging the whole agent. pub const IO_DRAIN_TIMEOUT_MS: u64 = 2_000; // 2 s should be plenty for local pipes +/// Builds the exact shell argv used for persistent model-designed commands. +/// +/// Approval and execution must both use this helper so the command the user +/// authorizes is the command the monitor runtime later passes to the sandbox. +pub fn persistent_shell_command_args(command: &str) -> Vec { + #[cfg(target_os = "windows")] + { + vec!["cmd".to_string(), "/C".to_string(), command.to_string()] + } + + #[cfg(not(target_os = "windows"))] + { + let shell = get_shell(ShellType::Bash, /*path*/ None) + .or_else(|| get_shell(ShellType::Zsh, /*path*/ None)) + .or_else(|| get_shell(ShellType::Sh, /*path*/ None)) + .unwrap_or_else(ultimate_fallback_shell); + vec![ + shell.shell_path.to_string_lossy().into_owned(), + "-lc".to_string(), + command.to_string(), + ] + } +} + #[derive(Debug)] pub struct ExecParams { pub command: Vec, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index e77138556..610bee385 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -26,6 +26,7 @@ use tokio::sync::Semaphore; /// A session has at most 1 running task at a time, and can be interrupted by user input. pub(crate) struct Session { pub(crate) thread_id: ThreadId, + pub(crate) monitor_writer_fence: String, pub(crate) installation_id: String, pub(super) tx_event: Sender, pub(super) agent_status: watch::Sender, @@ -662,6 +663,10 @@ impl Session { self.services.agent_control.session_id() } + pub(crate) fn monitor_writer_fence(&self) -> &str { + self.monitor_writer_fence.as_str() + } + #[instrument(name = "session_init", level = "info", skip_all)] #[allow(clippy::too_many_arguments)] #[expect( @@ -1262,6 +1267,7 @@ impl Session { let sess = Arc::new(Session { thread_id, + monitor_writer_fence: uuid::Uuid::new_v4().to_string(), installation_id, tx_event: tx_event.clone(), agent_status, diff --git a/codex-rs/core/src/tools/handlers/monitor_control.rs b/codex-rs/core/src/tools/handlers/monitor_control.rs index 42ca1945a..d98bd79c9 100644 --- a/codex-rs/core/src/tools/handlers/monitor_control.rs +++ b/codex-rs/core/src/tools/handlers/monitor_control.rs @@ -1,6 +1,8 @@ //! Built-in model tool handler for managing thread monitors. use crate::function_tool::FunctionCallError; +use crate::session::session::Session; +use crate::session::turn_context::TurnContext; use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; @@ -11,8 +13,10 @@ use crate::tools::handlers::parse_arguments; use crate::tools::registry::CoreToolRuntime; use crate::tools::registry::ToolExecutor; use codex_protocol::ThreadId; +use codex_protocol::protocol::ReviewDecision; use codex_tools::ToolName; use codex_tools::ToolSpec; +use codex_utils_absolute_path::AbsolutePathBuf; use serde::Deserialize; use serde::Serialize; use serde_json::Value as JsonValue; @@ -129,7 +133,11 @@ impl ToolExecutor for ManageMonitorHandler { invocation: ToolInvocation, ) -> Result, FunctionCallError> { let ToolInvocation { - session, payload, .. + session, + turn, + call_id, + payload, + .. } = invocation; let arguments = match payload { @@ -146,7 +154,15 @@ impl ToolExecutor for ManageMonitorHandler { let state_db = session.state_db().ok_or_else(|| { FunctionCallError::Fatal("sqlite state db is unavailable for this session".to_string()) })?; - let response = manage_monitor(state_db, session.thread_id(), args).await?; + let response = manage_monitor( + state_db, + session.clone(), + turn, + call_id, + session.thread_id(), + args, + ) + .await?; monitor_response(response, verbose).map(boxed_tool_output) } } @@ -155,11 +171,16 @@ impl CoreToolRuntime for ManageMonitorHandler {} async fn manage_monitor( state_db: Arc, + session: Arc, + turn: Arc, + call_id: String, thread_id: ThreadId, args: ManageMonitorArgs, ) -> Result { match args.action { - MonitorAction::Create => create_monitor(state_db, thread_id, args).await, + MonitorAction::Create => { + create_monitor(state_db, &session, &turn, call_id.as_str(), thread_id, args).await + } MonitorAction::List => { let monitors = list_monitor_snapshots(&state_db, thread_id).await?; Ok(ManageMonitorResponse { @@ -175,13 +196,18 @@ async fn manage_monitor( } MonitorAction::Read => read_monitor(state_db, thread_id, args).await, MonitorAction::Stop => set_monitor_stopped(state_db, thread_id, args).await, - MonitorAction::Restart => restart_monitor(state_db, thread_id, args).await, + MonitorAction::Restart => { + restart_monitor(state_db, &session, &turn, call_id.as_str(), thread_id, args).await + } MonitorAction::Delete => delete_monitor(state_db, thread_id, args).await, } } async fn create_monitor( state_db: Arc, + session: &Arc, + turn: &Arc, + call_id: &str, thread_id: ThreadId, args: ManageMonitorArgs, ) -> Result { @@ -215,6 +241,16 @@ async fn create_monitor( let output_file = validate_optional_monitor_relative_path("output_file", args.output_file.as_deref())?; let output_file = validate_output_file_for_routing(routing, output_file)?; + let authorization = authorize_monitor_command( + session, + turn, + call_id, + thread_id, + /*generation*/ 0, + command.as_str(), + cwd.as_deref(), + ) + .await?; let monitor = state_db .thread_monitors() .create_thread_monitor(codex_state::ThreadMonitorCreateParams { @@ -226,6 +262,7 @@ async fn create_monitor( routing, output_file, status: codex_state::ThreadMonitorStatus::Running, + authorization: Some(authorization), }) .await .map_err(|err| { @@ -344,14 +381,28 @@ async fn set_monitor_stopped( async fn restart_monitor( state_db: Arc, + session: &Arc, + turn: &Arc, + call_id: &str, thread_id: ThreadId, args: ManageMonitorArgs, ) -> Result { let monitor_id = resolve_monitor_id(&state_db, thread_id, args.monitor_id.as_deref()).await?; - load_monitor_for_thread(&state_db, thread_id, monitor_id.as_str()).await?; + let monitor = load_monitor_for_thread(&state_db, thread_id, monitor_id.as_str()).await?; + let cwd = validate_optional_monitor_relative_path("cwd", monitor.cwd.as_deref())?; + let authorization = authorize_monitor_command( + session, + turn, + call_id, + thread_id, + monitor.generation + 1, + monitor.command.as_str(), + cwd.as_deref(), + ) + .await?; let monitor = state_db .thread_monitors() - .restart_thread_monitor(monitor_id.as_str()) + .restart_thread_monitor(monitor_id.as_str(), monitor.generation, authorization) .await .map_err(|err| FunctionCallError::Fatal(format!("failed to restart monitor: {err}")))? .ok_or_else(|| model_error(format!("monitor not found: {monitor_id}")))?; @@ -369,6 +420,75 @@ async fn restart_monitor( }) } +async fn authorize_monitor_command( + session: &Arc, + turn: &Arc, + call_id: &str, + thread_id: ThreadId, + generation: i64, + command: &str, + cwd: Option<&str>, +) -> Result { + let before = session.thread_config_snapshot().await; + let approval_cwd = monitor_approval_cwd(&before.cwd, cwd)?; + let decision = session + .request_command_approval( + turn, + call_id.to_string(), + /*approval_id*/ None, + crate::exec::persistent_shell_command_args(command), + approval_cwd, + Some( + "Authorize this persistent monitor command. It may run in the background until stopped." + .to_string(), + ), + /*network_approval_context*/ None, + /*proposed_execpolicy_amendment*/ None, + /*additional_permissions*/ None, + Some(vec![ + ReviewDecision::Approved, + ReviewDecision::Denied, + ReviewDecision::Abort, + ]), + ) + .await; + if !monitor_command_approval_granted(&decision) { + return Err(model_error("monitor command approval was not granted")); + } + + let after = session.thread_config_snapshot().await; + if before.cwd != after.cwd || before.permission_profile != after.permission_profile { + return Err(model_error( + "monitor command authorization became stale before it could be recorded", + )); + } + Ok(codex_state::ThreadMonitorAuthorization::new( + thread_id, + generation, + command, + cwd, + session.monitor_writer_fence().to_string(), + after.permission_profile, + after.cwd.display().to_string(), + )) +} + +fn monitor_approval_cwd( + thread_cwd: &AbsolutePathBuf, + cwd: Option<&str>, +) -> Result { + let cwd = match cwd { + Some(cwd) => thread_cwd.join(cwd), + None => thread_cwd.to_path_buf(), + }; + AbsolutePathBuf::try_from(cwd) + .map_err(|err| model_error(format!("monitor cwd must be absolute: {err}"))) +} + +fn monitor_command_approval_granted(decision: &ReviewDecision) -> bool { + matches!(decision, ReviewDecision::Approved) +} + async fn delete_monitor( state_db: Arc, thread_id: ThreadId, @@ -760,6 +880,7 @@ mod tests { routing: codex_state::ThreadMonitorRouting::Stream, output_file: None, status: codex_state::ThreadMonitorStatus::Running, + authorization: None, }) .await .expect("monitor should be created") @@ -909,4 +1030,14 @@ mod tests { assert!(validate_optional_monitor_relative_path("output_file", Some("../out")).is_err()); assert!(validate_optional_monitor_relative_path("output_file", Some(".")).is_err()); } + + #[test] + fn monitor_command_approval_is_one_shot_and_fails_closed() { + assert!(monitor_command_approval_granted(&ReviewDecision::Approved)); + assert!(!monitor_command_approval_granted( + &ReviewDecision::ApprovedForSession + )); + assert!(!monitor_command_approval_granted(&ReviewDecision::Denied)); + assert!(!monitor_command_approval_granted(&ReviewDecision::Abort)); + } } diff --git a/codex-rs/state/migrations/0068_thread_schedule_occurrence_state.sql b/codex-rs/state/migrations/0068_thread_schedule_occurrence_state.sql index d3a5110d9..3e03869ff 100644 --- a/codex-rs/state/migrations/0068_thread_schedule_occurrence_state.sql +++ b/codex-rs/state/migrations/0068_thread_schedule_occurrence_state.sql @@ -1,3 +1,6 @@ +ALTER TABLE thread_monitors +ADD COLUMN authorization_json TEXT; + ALTER TABLE thread_schedule_runs ADD COLUMN deferral_kind TEXT CHECK(deferral_kind IS NULL OR deferral_kind IN ('idle', 'capacity')); diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index 706f46c79..bd0b98bd7 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -134,6 +134,7 @@ pub use model::ThreadGoalStatus; pub use model::ThreadMetadata; pub use model::ThreadMetadataBuilder; pub use model::ThreadMonitor; +pub use model::ThreadMonitorAuthorization; pub use model::ThreadMonitorEvent; pub use model::ThreadMonitorEventStream; pub use model::ThreadMonitorRouting; diff --git a/codex-rs/state/src/model/mod.rs b/codex-rs/state/src/model/mod.rs index 190acf485..46fc5f344 100644 --- a/codex-rs/state/src/model/mod.rs +++ b/codex-rs/state/src/model/mod.rs @@ -107,6 +107,7 @@ pub use thread_metadata::ThreadMetadata; pub use thread_metadata::ThreadMetadataBuilder; pub use thread_metadata::ThreadsPage; pub use thread_monitor::ThreadMonitor; +pub use thread_monitor::ThreadMonitorAuthorization; pub use thread_monitor::ThreadMonitorEvent; pub use thread_monitor::ThreadMonitorEventStream; pub use thread_monitor::ThreadMonitorRouting; diff --git a/codex-rs/state/src/model/thread_monitor.rs b/codex-rs/state/src/model/thread_monitor.rs index 3c9985a95..6e958acd5 100644 --- a/codex-rs/state/src/model/thread_monitor.rs +++ b/codex-rs/state/src/model/thread_monitor.rs @@ -3,6 +3,11 @@ use anyhow::anyhow; use chrono::DateTime; use chrono::Utc; use codex_protocol::ThreadId; +use codex_protocol::models::PermissionProfile; +use serde::Deserialize; +use serde::Serialize; +use sha2::Digest; +use sha2::Sha256; use sqlx::Row; use sqlx::sqlite::SqliteRow; @@ -118,6 +123,7 @@ pub struct ThreadMonitor { pub output_file: Option, pub status: ThreadMonitorStatus, pub generation: i64, + pub authorization: Option, pub process_id: Option, pub last_event_at: Option>, pub last_error: Option, @@ -125,6 +131,79 @@ pub struct ThreadMonitor { pub updated_at: DateTime, } +#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)] +pub struct ThreadMonitorAuthorization { + pub generation: i64, + pub writer_fence: String, + pub subject_digest: String, + pub permission_profile: PermissionProfile, + pub thread_cwd: String, +} + +impl ThreadMonitorAuthorization { + pub fn new( + thread_id: ThreadId, + generation: i64, + command: &str, + cwd: Option<&str>, + writer_fence: String, + permission_profile: PermissionProfile, + thread_cwd: String, + ) -> Self { + let subject_digest = thread_monitor_authorization_subject_digest( + thread_id, + generation, + command, + cwd, + &thread_cwd, + ); + Self { + generation, + writer_fence, + subject_digest, + permission_profile, + thread_cwd, + } + } + + pub fn authorizes( + &self, + monitor: &ThreadMonitor, + writer_fence: &str, + permission_profile: &PermissionProfile, + thread_cwd: &str, + ) -> bool { + self.generation == monitor.generation + && self.writer_fence == writer_fence + && self.permission_profile == *permission_profile + && self.thread_cwd == thread_cwd + && self.matches_subject( + monitor.thread_id, + monitor.generation, + monitor.command.as_str(), + monitor.cwd.as_deref(), + ) + } + + pub fn matches_subject( + &self, + thread_id: ThreadId, + generation: i64, + command: &str, + cwd: Option<&str>, + ) -> bool { + self.generation == generation + && self.subject_digest + == thread_monitor_authorization_subject_digest( + thread_id, + generation, + command, + cwd, + &self.thread_cwd, + ) + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct ThreadMonitorEvent { pub thread_id: ThreadId, @@ -146,6 +225,7 @@ pub(crate) struct ThreadMonitorRow { pub output_file: Option, pub status: String, pub generation: i64, + pub authorization_json: Option, pub process_id: Option, pub last_event_at_ms: Option, pub last_error: Option, @@ -166,6 +246,7 @@ impl ThreadMonitorRow { output_file: row.try_get("output_file")?, status: row.try_get("status")?, generation: row.try_get("generation")?, + authorization_json: row.try_get("authorization_json")?, process_id: row.try_get("process_id")?, last_event_at_ms: row.try_get("last_event_at_ms")?, last_error: row.try_get("last_error")?, @@ -190,6 +271,10 @@ impl TryFrom for ThreadMonitor { output_file: row.output_file, status: ThreadMonitorStatus::try_from(row.status.as_str())?, generation: row.generation, + authorization: row + .authorization_json + .as_deref() + .and_then(|value| serde_json::from_str(value).ok()), process_id: row.process_id, last_event_at: optional_epoch_millis_to_datetime(row.last_event_at_ms)?, last_error: row.last_error, @@ -236,6 +321,118 @@ impl TryFrom for ThreadMonitorEvent { } } +fn thread_monitor_authorization_subject_digest( + thread_id: ThreadId, + generation: i64, + command: &str, + cwd: Option<&str>, + thread_cwd: &str, +) -> String { + let mut hasher = Sha256::new(); + hash_monitor_authorization_field(&mut hasher, b"version", b"1"); + hash_monitor_authorization_field(&mut hasher, b"thread_id", thread_id.to_string().as_bytes()); + hash_monitor_authorization_field( + &mut hasher, + b"generation", + generation.to_string().as_bytes(), + ); + hash_monitor_authorization_field(&mut hasher, b"command", command.as_bytes()); + match cwd { + Some(cwd) => { + hash_monitor_authorization_field(&mut hasher, b"cwd_present", b"1"); + hash_monitor_authorization_field(&mut hasher, b"cwd", cwd.as_bytes()); + } + None => hash_monitor_authorization_field(&mut hasher, b"cwd_present", b"0"), + } + hash_monitor_authorization_field(&mut hasher, b"thread_cwd", thread_cwd.as_bytes()); + let digest = hasher.finalize(); + digest.iter().map(|byte| format!("{byte:02x}")).collect() +} + +fn hash_monitor_authorization_field(hasher: &mut Sha256, label: &[u8], value: &[u8]) { + hasher.update((label.len() as u64).to_le_bytes()); + hasher.update(label); + hasher.update((value.len() as u64).to_le_bytes()); + hasher.update(value); +} + fn optional_epoch_millis_to_datetime(value: Option) -> Result>> { value.map(epoch_millis_to_datetime).transpose() } + +#[cfg(test)] +mod tests { + use super::*; + + fn test_monitor(authorization: Option) -> ThreadMonitor { + let now = Utc::now(); + ThreadMonitor { + thread_id: ThreadId::new(), + monitor_id: "monitor-id".to_string(), + name: "monitor".to_string(), + prompt: "watch".to_string(), + command: "printf ok".to_string(), + cwd: Some("logs".to_string()), + routing: ThreadMonitorRouting::Stream, + output_file: None, + status: ThreadMonitorStatus::Running, + generation: 2, + authorization, + process_id: None, + last_event_at: None, + last_error: None, + created_at: now, + updated_at: now, + } + } + + #[test] + fn monitor_authorization_accepts_only_the_current_exact_subject_and_fence() { + let permission_profile = PermissionProfile::read_only(); + let thread_cwd = "/workspace".to_string(); + let mut monitor = test_monitor(/*authorization*/ None); + let authorization = ThreadMonitorAuthorization::new( + monitor.thread_id, + monitor.generation, + monitor.command.as_str(), + monitor.cwd.as_deref(), + "writer-fence".to_string(), + permission_profile.clone(), + thread_cwd.clone(), + ); + monitor.authorization = Some(authorization.clone()); + + assert!(authorization.authorizes( + &monitor, + "writer-fence", + &permission_profile, + thread_cwd.as_str(), + )); + assert!(!authorization.authorizes( + &monitor, + "stale-writer", + &permission_profile, + thread_cwd.as_str(), + )); + assert!(!authorization.authorizes( + &monitor, + "writer-fence", + &PermissionProfile::Disabled, + thread_cwd.as_str(), + )); + assert!(!authorization.authorizes( + &monitor, + "writer-fence", + &permission_profile, + "/different-workspace", + )); + + monitor.command = "printf changed".to_string(); + assert!(!authorization.authorizes( + &monitor, + "writer-fence", + &permission_profile, + thread_cwd.as_str(), + )); + } +} diff --git a/codex-rs/state/src/runtime/monitors.rs b/codex-rs/state/src/runtime/monitors.rs index 79bfaba8a..500223bd4 100644 --- a/codex-rs/state/src/runtime/monitors.rs +++ b/codex-rs/state/src/runtime/monitors.rs @@ -23,6 +23,7 @@ pub struct ThreadMonitorCreateParams { pub routing: crate::ThreadMonitorRouting, pub output_file: Option, pub status: crate::ThreadMonitorStatus, + pub authorization: Option, } pub struct ThreadMonitorUpdate { @@ -34,6 +35,7 @@ pub struct ThreadMonitorUpdate { pub output_file: Option>, pub status: Option, pub generation: Option, + pub authorization: Option, pub process_id: Option>, pub last_event_at: Option>>, pub last_error: Option>, @@ -53,6 +55,22 @@ impl MonitorStore { ) -> anyhow::Result { let monitor_id = Uuid::new_v4().to_string(); let now_ms = datetime_to_epoch_millis(Utc::now()); + let name = redact_state_string(params.name); + let prompt = redact_state_string(params.prompt); + let command = redact_state_string(params.command); + let cwd = redact_state_optional_string(params.cwd); + let output_file = redact_state_optional_string(params.output_file); + if let Some(authorization) = params.authorization.as_ref() + && !authorization.matches_subject( + params.thread_id, + /*generation*/ 0, + command.as_str(), + cwd.as_deref(), + ) + { + anyhow::bail!("thread monitor authorization does not match the create subject"); + } + let authorization_json = serialize_monitor_authorization(params.authorization.as_ref())?; let sql = thread_monitor_returning( r#" INSERT INTO thread_monitors ( @@ -65,22 +83,24 @@ INSERT INTO thread_monitors ( routing, output_file, status, + authorization_json, created_at_ms, updated_at_ms -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) RETURNING "#, ); let row = sqlx::query(sqlx::AssertSqlSafe(sql)) .bind(monitor_id) .bind(params.thread_id.to_string()) - .bind(redact_state_string(params.name)) - .bind(redact_state_string(params.prompt)) - .bind(redact_state_string(params.command)) - .bind(redact_state_optional_string(params.cwd)) + .bind(name) + .bind(prompt) + .bind(command) + .bind(cwd) .bind(params.routing.as_str()) - .bind(redact_state_optional_string(params.output_file)) + .bind(output_file) .bind(params.status.as_str()) + .bind(authorization_json) .bind(now_ms) .bind(now_ms) .fetch_one(self.pool.as_ref()) @@ -121,6 +141,7 @@ SELECT output_file, status, generation, + authorization_json, process_id, last_event_at_ms, last_error, @@ -151,6 +172,7 @@ SELECT output_file, status, generation, + authorization_json, process_id, last_event_at_ms, last_error, @@ -182,6 +204,7 @@ ORDER BY updated_at_ms, monitor_id let output_file = update.output_file.unwrap_or(existing.output_file); let status = update.status.unwrap_or(existing.status); let generation = update.generation.unwrap_or(existing.generation); + let authorization = update.authorization.or(existing.authorization); let process_id = update.process_id.unwrap_or(existing.process_id); let last_event_at = update.last_event_at.unwrap_or(existing.last_event_at); let last_error = update.last_error.unwrap_or(existing.last_error); @@ -191,6 +214,7 @@ ORDER BY updated_at_ms, monitor_id let cwd = redact_state_optional_string(cwd); let output_file = redact_state_optional_string(output_file); let last_error = redact_state_optional_string(last_error); + let authorization_json = serialize_monitor_authorization(authorization.as_ref())?; let sql = thread_monitor_returning( r#" UPDATE thread_monitors @@ -203,6 +227,7 @@ SET output_file = ?, status = ?, generation = ?, + authorization_json = ?, process_id = ?, last_event_at_ms = ?, last_error = ?, @@ -220,6 +245,7 @@ RETURNING .bind(output_file) .bind(status.as_str()) .bind(generation) + .bind(authorization_json) .bind(process_id) .bind(last_event_at.map(datetime_to_epoch_millis)) .bind(last_error) @@ -247,6 +273,7 @@ RETURNING output_file: None, status: Some(status), generation: None, + authorization: None, process_id: Some(None), last_event_at: None, last_error: Some(last_error), @@ -258,27 +285,51 @@ RETURNING pub async fn restart_thread_monitor( &self, monitor_id: &str, + expected_generation: i64, + authorization: crate::ThreadMonitorAuthorization, ) -> anyhow::Result> { let Some(existing) = self.get_thread_monitor(monitor_id).await? else { return Ok(None); }; - self.update_thread_monitor( - monitor_id, - ThreadMonitorUpdate { - name: None, - prompt: None, - command: None, - cwd: None, - routing: None, - output_file: None, - status: Some(crate::ThreadMonitorStatus::Running), - generation: Some(existing.generation + 1), - process_id: Some(None), - last_event_at: None, - last_error: Some(None), - }, - ) - .await + if existing.generation != expected_generation { + anyhow::bail!("thread monitor generation changed before restart authorization"); + } + let next_generation = expected_generation + 1; + if !authorization.matches_subject( + existing.thread_id, + next_generation, + existing.command.as_str(), + existing.cwd.as_deref(), + ) { + anyhow::bail!("thread monitor authorization does not match the restart subject"); + } + let authorization_json = serialize_monitor_authorization(Some(&authorization))?; + let sql = thread_monitor_returning( + r#" +UPDATE thread_monitors +SET + status = 'running', + generation = ?, + authorization_json = ?, + process_id = NULL, + last_error = NULL, + updated_at_ms = ? +WHERE monitor_id = ? AND generation = ? +RETURNING +"#, + ); + let row = sqlx::query(sqlx::AssertSqlSafe(sql)) + .bind(next_generation) + .bind(authorization_json) + .bind(datetime_to_epoch_millis(Utc::now())) + .bind(monitor_id) + .bind(expected_generation) + .fetch_optional(self.pool.as_ref()) + .await?; + match row { + Some(row) => Ok(Some(thread_monitor_from_row(&row)?)), + None => anyhow::bail!("thread monitor generation changed before restart commit"), + } } pub async fn mark_thread_monitor_started( @@ -304,6 +355,7 @@ RETURNING output_file: None, status: Some(crate::ThreadMonitorStatus::Running), generation: None, + authorization: None, process_id: Some(process_id), last_event_at: None, last_error: Some(None), @@ -467,6 +519,7 @@ fn thread_monitor_select_columns() -> &'static str { output_file, status, generation, + authorization_json, process_id, last_event_at_ms, last_error, @@ -487,6 +540,14 @@ fn thread_monitor_select_by_id(prefix: &'static str) -> String { ) } +fn serialize_monitor_authorization( + authorization: Option<&crate::ThreadMonitorAuthorization>, +) -> anyhow::Result> { + authorization + .map(crate::redacted_local_state_serialized_json_string) + .transpose() +} + #[cfg(test)] mod tests { use super::*; @@ -525,17 +586,27 @@ mod tests { runtime: &StateRuntime, thread_id: ThreadId, ) -> crate::ThreadMonitor { + let command = "while true; do echo ok; sleep 60; done"; runtime .thread_monitors() .create_thread_monitor(ThreadMonitorCreateParams { thread_id, name: "CI watcher".to_string(), prompt: "watch CI".to_string(), - command: "while true; do echo ok; sleep 60; done".to_string(), + command: command.to_string(), cwd: None, routing: crate::ThreadMonitorRouting::Stream, output_file: None, status: crate::ThreadMonitorStatus::Running, + authorization: Some(crate::ThreadMonitorAuthorization::new( + thread_id, + /*generation*/ 0, + command, + /*cwd*/ None, + "writer-fence".to_string(), + codex_protocol::models::PermissionProfile::read_only(), + "/workspace".to_string(), + )), }) .await .expect("monitor should be created") @@ -552,6 +623,7 @@ mod tests { assert_eq!("CI watcher", created.name); assert_eq!(crate::ThreadMonitorStatus::Running, created.status); assert_eq!(0, created.generation); + assert!(created.authorization.is_some()); let event = runtime .thread_monitors() @@ -576,14 +648,49 @@ mod tests { .expect("events should list"); assert_eq!(vec![event], events); + let restart_authorization = crate::ThreadMonitorAuthorization::new( + thread_id, + /*generation*/ 1, + created.command.as_str(), + created.cwd.as_deref(), + "writer-fence".to_string(), + codex_protocol::models::PermissionProfile::read_only(), + "/workspace".to_string(), + ); let restarted = runtime .thread_monitors() - .restart_thread_monitor(created.monitor_id.as_str()) + .restart_thread_monitor( + created.monitor_id.as_str(), + /*expected_generation*/ 0, + restart_authorization, + ) .await .expect("restart should succeed") .expect("monitor should exist"); assert_eq!(crate::ThreadMonitorStatus::Running, restarted.status); assert_eq!(1, restarted.generation); + assert!(restarted.authorization.is_some()); + let stale_authorization = crate::ThreadMonitorAuthorization::new( + thread_id, + /*generation*/ 1, + created.command.as_str(), + created.cwd.as_deref(), + "writer-fence".to_string(), + codex_protocol::models::PermissionProfile::read_only(), + "/workspace".to_string(), + ); + assert!( + runtime + .thread_monitors() + .restart_thread_monitor( + created.monitor_id.as_str(), + /*expected_generation*/ 0, + stale_authorization, + ) + .await + .is_err(), + "a stale generation must not commit a restart authorization" + ); let deleted = runtime .thread_monitors() @@ -619,6 +726,7 @@ mod tests { routing: crate::ThreadMonitorRouting::Stream, output_file: None, status: crate::ThreadMonitorStatus::Running, + authorization: None, }) .await .expect("monitor should be created"); diff --git a/codex-rs/state/src/runtime/workflow_automation.rs b/codex-rs/state/src/runtime/workflow_automation.rs index aa19ce5e0..c4fd630da 100644 --- a/codex-rs/state/src/runtime/workflow_automation.rs +++ b/codex-rs/state/src/runtime/workflow_automation.rs @@ -1325,6 +1325,7 @@ WHERE timer_id = ? routing: crate::ThreadMonitorRouting::Stream, output_file: None, status: crate::ThreadMonitorStatus::Running, + authorization: None, }) .await .expect("monitor should create"); @@ -1408,6 +1409,7 @@ WHERE timer_id = ? routing: crate::ThreadMonitorRouting::Stream, output_file: None, status: crate::ThreadMonitorStatus::Running, + authorization: None, }) .await .expect("monitor should create"); From 840814b96bfbeb79e2efb50a2212466f725ccb16 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sun, 9 Aug 2026 18:43:34 +0300 Subject: [PATCH 19/20] fix(monitor): expose current session snapshot Agent: nausicaa --- .../src/request_processors/thread_monitor_processor.rs | 2 +- codex-rs/core/src/session/mod.rs | 5 +++++ codex-rs/core/src/tools/handlers/monitor_control.rs | 2 +- 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/thread_monitor_processor.rs b/codex-rs/app-server/src/request_processors/thread_monitor_processor.rs index a37296b0d..0e9f496db 100644 --- a/codex-rs/app-server/src/request_processors/thread_monitor_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_monitor_processor.rs @@ -740,7 +740,7 @@ fn monitor_approval_cwd( ) -> Result { let cwd = match cwd { Some(cwd) => thread_cwd.join(cwd), - None => thread_cwd.to_path_buf(), + None => thread_cwd.clone(), }; AbsolutePathBuf::try_from(cwd) .map_err(|err| invalid_request(format!("monitor cwd must be absolute: {err}"))) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 9071ca8c0..14c471058 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1869,6 +1869,11 @@ impl Session { .clone() } + pub(crate) async fn thread_config_snapshot(&self) -> ThreadConfigSnapshot { + let state = self.state.lock().await; + state.session_configuration.thread_config_snapshot() + } + pub(crate) async fn provider(&self) -> ModelProviderInfo { let state = self.state.lock().await; state.session_configuration.provider.clone() diff --git a/codex-rs/core/src/tools/handlers/monitor_control.rs b/codex-rs/core/src/tools/handlers/monitor_control.rs index d98bd79c9..485c6f0d0 100644 --- a/codex-rs/core/src/tools/handlers/monitor_control.rs +++ b/codex-rs/core/src/tools/handlers/monitor_control.rs @@ -479,7 +479,7 @@ fn monitor_approval_cwd( ) -> Result { let cwd = match cwd { Some(cwd) => thread_cwd.join(cwd), - None => thread_cwd.to_path_buf(), + None => thread_cwd.clone(), }; AbsolutePathBuf::try_from(cwd) .map_err(|err| model_error(format!("monitor cwd must be absolute: {err}"))) From 0af635e7c1b22750362b49195951bb030d97f55a Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Sun, 9 Aug 2026 18:59:57 +0300 Subject: [PATCH 20/20] fix(app-server): use configured monitor environment Agent: nausicaa --- .../src/request_processors/thread_monitor_runtime.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/codex-rs/app-server/src/request_processors/thread_monitor_runtime.rs b/codex-rs/app-server/src/request_processors/thread_monitor_runtime.rs index 7db39d7b6..0e358e13d 100644 --- a/codex-rs/app-server/src/request_processors/thread_monitor_runtime.rs +++ b/codex-rs/app-server/src/request_processors/thread_monitor_runtime.rs @@ -281,7 +281,10 @@ impl ThreadMonitorRuntime { "monitor process starting", ) .await; - let env = create_env(&config.shell_environment_policy, Some(monitor.thread_id)); + let env = create_env( + &config.permissions.shell_environment_policy, + Some(monitor.thread_id), + ); let mut child = match codex_core::exec::spawn_streaming_command_under_sandbox( codex_core::exec::persistent_shell_command_args(&monitor.command), cwd.clone(),