diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index e76dc7ff33..b6276d63ad 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2983,6 +2983,7 @@ dependencies = [ "pretty_assertions", "serde", "serde_json", + "sqlx", "tempfile", "tokio", "tracing", diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 62c9a3f513..6b20d2545a 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -2844,6 +2844,7 @@ mod tests { schedule_id: "schedule-1".to_string(), run_id: "run-1".to_string(), lease_id: "lease-1".to_string(), + goal_id: None, state_db, }, ); 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 57be3c8c1c..e10970dc90 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -280,6 +280,7 @@ pub(super) async fn ensure_listener_task_running( thread_list_state_permit, fallback_model_provider, codex_home, + state_db, .. } = listener_task_context; let outgoing_for_task = Arc::clone(&outgoing); @@ -320,24 +321,53 @@ 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 (raw_events_enabled, terminal_scheduled_run) = { + let terminal_event = matches!( + event.msg, + EventMsg::TurnComplete(_) | EventMsg::TurnAborted(_) | EventMsg::Error(_) + ); + let (raw_events_enabled, tracked_scheduled_run, turn_error) = { let mut thread_state = thread_state.lock().await; thread_state.track_current_turn_event(&event.id, &event.msg); - let terminal_scheduled_run = if matches!( - event.msg, - EventMsg::TurnComplete(_) - | EventMsg::TurnAborted(_) - | EventMsg::Error(_) - ) { - thread_state - .take_scheduled_run(&event.id) - .map(|scheduled_run| { - (scheduled_run, thread_state.turn_summary.last_error.clone()) - }) + let tracked_scheduled_run = if terminal_event { + thread_state.take_scheduled_run(&event.id) } else { None }; - (thread_state.experimental_raw_events, terminal_scheduled_run) + let turn_error = terminal_event + .then(|| thread_state.turn_summary.last_error.clone()) + .flatten(); + ( + thread_state.experimental_raw_events, + tracked_scheduled_run, + turn_error, + ) + }; + let terminal_scheduled_run = match ( + terminal_event, + tracked_scheduled_run, + state_db.as_ref(), + ) { + (true, Some(scheduled_run), _) => Some(scheduled_run), + (true, None, Some(state_db)) => { + match thread_schedule_runtime::recover_scheduled_run_for_terminal_turn( + state_db, + conversation_id, + &event.id, + ) + .await + { + Ok(scheduled_run) => scheduled_run, + Err(err) => { + tracing::warn!( + thread_id = %conversation_id, + turn_id = %event.id, + "failed to recover scheduled run for terminal turn: {err}" + ); + None + } + } + } + (false, _, _) | (true, None, None) => None, }; let subscribed_connection_ids = thread_state_manager .subscribed_connection_ids(conversation_id) @@ -378,7 +408,7 @@ pub(super) async fn ensure_listener_task_running( fallback_model_provider.clone(), ) .await; - if let Some((scheduled_run, turn_error)) = terminal_scheduled_run { + if let Some(scheduled_run) = terminal_scheduled_run { thread_schedule_runtime::finish_scheduled_run_after_turn( conversation_id, scheduled_run, 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 59d17f6034..b62fc3f7d7 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 @@ -182,27 +182,50 @@ impl ThreadScheduleRuntime { self.emit_schedule_run_updated(thread_id, claim.run.clone()) .await; - let result = self - .submit_claimed_schedule(thread_id, state_db.clone(), &claim) - .await; - if let Err(err) = result { + 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: {err}" + "failed to submit scheduled thread run: {error}" ); - if let Some(wait) = err.downcast_ref::() { + 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) = err.downcast_ref::() { + 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, schedule_submit_error(&err)) - .await; + self.fail_claimed_run_after_submit_error( + state_db, + claim, + goal_id, + schedule_submit_error(&error), + ) + .await; } } @@ -211,11 +234,9 @@ impl ThreadScheduleRuntime { thread_id: ThreadId, state_db: StateDbHandle, claim: &codex_state::ThreadScheduleClaim, - ) -> anyhow::Result<()> { - let prompt = self - .resolve_claim_prompt(&state_db, thread_id, &claim.schedule) - .await?; - let scheduled_goal_objective = scheduled_goal_objective(&prompt).map(str::to_string); + prompt: String, + scheduled_goal_objective: Option, + ) -> Result<(), ScheduleSubmitError> { let claim_auth_profile = self .claim_auth_profile(&state_db, thread_id, &claim.schedule) .await; @@ -232,79 +253,145 @@ impl ThreadScheduleRuntime { Utc::now(), ) { Ok(resolved) => resolved, - Err(wait) => return Err(anyhow::Error::new(wait)), + 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?; + .await + .map_err(|error| ScheduleSubmitError { + error, + goal_id: None, + })?; self.ensure_schedule_listener(thread_id, thread.clone()) - .await?; + .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 = if let Some(objective) = scheduled_goal_objective.as_deref() { - self.prepare_scheduled_goal( - thread_id, - &state_db, - objective, - listener_command_tx.clone(), - ) - .await?; - scheduled_goal_thread_prompt( - objective, - claim.run.run_id.as_str(), - claim.run.scheduled_for, - &claim.schedule, - ) - } else { - scheduled_thread_prompt( - &prompt, - &claim.schedule, - claim.run.run_id.as_str(), - claim.run.scheduled_for, - ) - }; + 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 = match state_db + let run_start = state_db .thread_schedules() - .mark_thread_schedule_run_started( - claim.schedule.schedule_id.as_str(), - claim.run.run_id.as_str(), - claim.run.lease_id.as_str(), - turn_id.as_str(), - ) - .await - { + .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(anyhow::anyhow!( - "claimed schedule run {} disappeared before it could start", - claim.run.run_id - )); + 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, + }); } - Err(err) => return Err(err), }; { let mut thread_state = thread_state.lock().await; thread_state.track_scheduled_run( turn_id.clone(), crate::thread_state::ScheduledThreadScheduleRun { - schedule_id: claim.schedule.schedule_id.clone(), - run_id: claim.run.run_id.clone(), - lease_id: claim.run.lease_id.clone(), + 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 = thread - .try_start_user_input_turn_if_idle( + 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, @@ -312,23 +399,51 @@ impl ThreadScheduleRuntime { }], Default::default(), thread_settings, - ) - .await; + ), + ) + .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(anyhow::Error::new(deferral)); + return Err(ScheduleSubmitError { + error: anyhow::Error::new(deferral), + goal_id: scheduled_goal_id, + }); } - return Err(anyhow::anyhow!("failed to start scheduled prompt: {err}")); + return Err(ScheduleSubmitError { + error: anyhow::anyhow!("failed to start scheduled prompt: {err}"), + goal_id: scheduled_goal_id, + }); } - self.spawn_lease_heartbeat( - state_db, - claim.schedule.schedule_id.clone(), - claim.run.lease_id.clone(), - ); self.emit_schedule_run_updated(thread_id, run).await; Ok(()) } @@ -339,11 +454,24 @@ impl ThreadScheduleRuntime { state_db: &StateDbHandle, objective: &str, listener_command_tx: Option>, - ) -> anyhow::Result<()> { + ) -> anyhow::Result { if !self.config.features.enabled(Feature::Goals) { anyhow::bail!("goals feature is disabled"); } + if let Some(goal) = state_db + .thread_goals() + .get_thread_goal(thread_id) + .await + .map_err(|err| anyhow::anyhow!("failed to read scheduled goal: {err}"))? + && scheduled_goal_is_held(&goal, objective) + { + return Err(anyhow::Error::new(ScheduledGoalHeld { + goal_id: goal.goal_id, + status: goal.status, + })); + } + let outcome = self .goal_service .set_thread_goal( @@ -376,7 +504,7 @@ impl ThreadScheduleRuntime { outcome.apply_runtime_effects(&self.goal_service).await; self.goal_service .suppress_next_idle_continuation(thread_id, goal_id.as_str()); - Ok(()) + Ok(goal_id) } async fn resolve_claim_prompt( @@ -448,38 +576,69 @@ impl ThreadScheduleRuntime { schedule_resume_auth_profile(/*schedule_auth_profile*/ None, &initial_history) } - fn spawn_lease_heartbeat( + async fn start_lease_heartbeat( &self, state_db: StateDbHandle, - schedule_id: String, - lease_id: String, - ) { + run: &codex_state::ThreadScheduleRun, + ) -> anyhow::Result> { + let schedule_id = run.schedule_id.clone(); + let run_id = run.run_id.clone(); + let lease_id = run.lease_id.clone(); + let ownership_lost = CancellationToken::new(); + let owns_lease = state_db + .thread_schedules() + .extend_thread_schedule_lease(codex_state::ThreadScheduleRunLeaseParams { + schedule_id: schedule_id.as_str(), + run_id: run_id.as_str(), + lease_id: lease_id.as_str(), + now: Utc::now(), + lease_duration: SCHEDULE_LEASE_DURATION, + }) + .await?; + if !owns_lease { + ownership_lost.cancel(); + return Ok(None); + } + let cancel_token = self.cancel_token.clone(); + let heartbeat_ownership_lost = ownership_lost.clone(); self.tasks.spawn(async move { loop { tokio::select! { - _ = cancel_token.cancelled() => break, + _ = cancel_token.cancelled() => { + heartbeat_ownership_lost.cancel(); + break; + } _ = tokio::time::sleep(SCHEDULE_LEASE_HEARTBEAT_INTERVAL) => {} } match state_db .thread_schedules() - .extend_thread_schedule_lease( - schedule_id.as_str(), - lease_id.as_str(), - Utc::now(), - SCHEDULE_LEASE_DURATION, - ) + .extend_thread_schedule_lease(codex_state::ThreadScheduleRunLeaseParams { + schedule_id: schedule_id.as_str(), + run_id: run_id.as_str(), + lease_id: lease_id.as_str(), + now: Utc::now(), + lease_duration: SCHEDULE_LEASE_DURATION, + }) .await { Ok(true) => {} - Ok(false) => break, - Err(err) => warn!( - schedule_id = %schedule_id, - "failed to refresh scheduled thread lease: {err}" - ), + Ok(false) => { + heartbeat_ownership_lost.cancel(); + break; + } + Err(err) => { + warn!( + schedule_id = %schedule_id, + "failed to refresh scheduled thread lease: {err}" + ); + heartbeat_ownership_lost.cancel(); + break; + } } } }); + Ok(Some(ownership_lost)) } async fn load_or_resume_thread( @@ -565,6 +724,7 @@ impl ThreadScheduleRuntime { &self, state_db: StateDbHandle, claim: codex_state::ThreadScheduleClaim, + goal_id: Option, error: String, ) { match finish_scheduled_run_state( @@ -572,6 +732,7 @@ impl ThreadScheduleRuntime { &claim.schedule.schedule_id, &claim.run.run_id, &claim.run.lease_id, + goal_id.as_deref(), Some(error), Utc::now(), ) @@ -718,6 +879,30 @@ impl ThreadScheduleRuntime { } } +async fn submit_scheduled_turn_if_owned( + state_db: &StateDbHandle, + lease: codex_state::ThreadScheduleRunLeaseParams<'_>, + ownership_lost: &CancellationToken, + submit: impl std::future::Future, +) -> anyhow::Result> { + if ownership_lost.is_cancelled() { + return Ok(None); + } + if !state_db + .thread_schedules() + .extend_thread_schedule_lease(lease) + .await? + { + return Ok(None); + } + tokio::pin!(submit); + tokio::select! { + biased; + _ = ownership_lost.cancelled() => Ok(None), + result = &mut submit => Ok(Some(result)), + } +} + fn scheduled_thread_prompt( prompt: &str, schedule: &codex_state::ThreadSchedule, @@ -774,6 +959,17 @@ fn scheduled_goal_objective(prompt: &str) -> Option<&str> { Some(rest.trim()) } +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 + ) +} + fn scheduled_goal_thread_prompt( objective: &str, run_id: &str, @@ -950,6 +1146,7 @@ pub(super) async fn finish_scheduled_run_after_turn( 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, ) @@ -982,6 +1179,27 @@ pub(super) async fn finish_scheduled_run_after_turn( } } +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; @@ -1005,6 +1223,11 @@ 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!( @@ -1017,6 +1240,41 @@ impl std::fmt::Display for ScheduleUsageProfileWait { impl std::error::Error for ScheduleUsageProfileWait {} +#[derive(Debug, Clone, PartialEq, Eq)] +struct ScheduledGoalHeld { + goal_id: String, + status: codex_state::ThreadGoalStatus, +} + +impl std::fmt::Display for ScheduledGoalHeld { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let recovery = match self.status { + codex_state::ThreadGoalStatus::Paused => { + "scheduled goal is paused; resume the goal and its schedule explicitly before another run" + } + codex_state::ThreadGoalStatus::Blocked => { + "scheduled goal is blocked; resume the goal and its schedule explicitly before another run" + } + codex_state::ThreadGoalStatus::UsageLimited => { + "scheduled goal is usage limited; wait for or resolve the usage limit, then resume the goal and its schedule explicitly before another run" + } + codex_state::ThreadGoalStatus::BudgetLimited => { + "scheduled goal is budget limited; change the goal token budget before resuming the goal, then resume its schedule explicitly before another run" + } + status => { + return write!( + f, + "scheduled goal is {}; resume the goal and its schedule explicitly before another run", + status.as_str() + ); + } + }; + f.write_str(recovery) + } +} + +impl std::error::Error for ScheduledGoalHeld {} + #[derive(Debug, Clone, PartialEq, Eq)] struct ScheduleRunDeferral { retry_at: DateTime, @@ -1100,6 +1358,7 @@ async fn finish_scheduled_run_state( schedule_id: &str, run_id: &str, lease_id: &str, + goal_id: Option<&str>, error: Option, completed_at: DateTime, ) -> anyhow::Result> { @@ -1151,23 +1410,63 @@ async fn finish_scheduled_run_state( (Some(next_run_at), Some(expires_at)) if next_run_at >= expires_at => None, (next_run_at, _) => next_run_at, }; - let updated = if let Some(error) = error { - state_db - .thread_schedules() - .fail_thread_schedule_run( - schedule_id, - run_id, - lease_id, - completed_at, - next_run_at, - error, - ) - .await? - } else { - state_db - .thread_schedules() - .complete_thread_schedule_run(schedule_id, run_id, lease_id, completed_at, next_run_at) - .await? + let updated = match (error, goal_id) { + (Some(error), Some(goal_id)) => { + state_db + .thread_schedules() + .fail_thread_schedule_run_for_goal( + codex_state::ThreadScheduleRunForGoalFinishParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + expected_goal_id: goal_id, + }, + error, + ) + .await? + } + (Some(error), None) => { + state_db + .thread_schedules() + .fail_thread_schedule_run( + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + error, + ) + .await? + } + (None, Some(goal_id)) => { + state_db + .thread_schedules() + .complete_thread_schedule_run_for_goal( + codex_state::ThreadScheduleRunForGoalFinishParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + expected_goal_id: goal_id, + }, + ) + .await? + } + (None, None) => { + state_db + .thread_schedules() + .complete_thread_schedule_run( + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at, + ) + .await? + } }; if !updated { return Ok(None); @@ -1758,6 +2057,8 @@ mod tests { use codex_protocol::protocol::TurnContextItem; use codex_state::ThreadMetadataBuilder; use pretty_assertions::assert_eq; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::Ordering; fn at(seconds: i64) -> DateTime { DateTime::::from_timestamp(seconds, 0).expect("valid timestamp") @@ -2684,6 +2985,7 @@ mod tests { &retry_claim.run.run_id, "lease-retry", None, + None, completed_at, ) .await @@ -2766,6 +3068,7 @@ mod tests { &claim.run.run_id, "lease-run", None, + None, completed_at, ) .await @@ -2793,49 +3096,605 @@ mod tests { ); } - #[test] - fn computes_interval_next_run() { - assert_eq!( - Some(at(/*seconds*/ 1_700_000_300)), - next_thread_schedule_run_at( - &codex_state::ThreadScheduleSpec::Interval(codex_state::ThreadScheduleInterval { - amount: 5, - unit: codex_state::ThreadScheduleIntervalUnit::Minutes, - }), - "UTC", - at(/*seconds*/ 1_700_000_000), - ) - .expect("next interval should compute") + #[tokio::test] + async fn successful_blocked_goal_run_pauses_schedule_without_deferring_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 objective = "watch the external release gate"; + let goal = state_db + .thread_goals() + .replace_thread_goal( + thread_id, + objective, + 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: format!("/goal {objective}"), + 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-run", Duration::from_secs(300)) + .await + .expect("claim should succeed") + .expect("schedule should claim"); + let completed_at = scheduled_for + chrono::Duration::seconds(5); + + let (finished_schedule, finished_run) = finish_scheduled_run_state( + &state_db, + &schedule.schedule_id, + &claim.run.run_id, + "lease-run", + Some(goal.goal_id.as_str()), + None, + completed_at, + ) + .await + .expect("successful run should finish") + .expect("finished rows should load"); - #[test] - fn computes_recurring_next_run_from_scheduled_time_without_drift() { assert_eq!( - Some(at(/*seconds*/ 1_700_000_060)), - next_thread_schedule_run_after_completion( - &codex_state::ThreadScheduleSpec::Interval(codex_state::ThreadScheduleInterval { - amount: 1, - unit: codex_state::ThreadScheduleIntervalUnit::Minutes, - }), - "UTC", - Some(at(/*seconds*/ 1_700_000_000)), - at(/*seconds*/ 1_700_000_005), - ) - .expect("next interval should compute") + codex_state::ThreadScheduleStatus::Paused, + finished_schedule.status, + "a successful still-blocked result must not rearm its schedule" + ); + assert_eq!(None, finished_schedule.next_run_at); + assert_eq!( + codex_state::ThreadScheduleRunStatus::Completed, + finished_run.status ); + assert_eq!(None, finished_run.error); + let goal = state_db + .thread_goals() + .get_thread_goal(thread_id) + .await + .expect("goal read should succeed") + .expect("goal should remain persisted"); + assert_eq!(codex_state::ThreadGoalStatus::Blocked, goal.status); } - #[test] - fn computes_recurring_next_run_skipping_missed_intervals() { - assert_eq!( - Some(at(/*seconds*/ 1_700_000_180)), - next_thread_schedule_run_after_completion( - &codex_state::ThreadScheduleSpec::Interval(codex_state::ThreadScheduleInterval { - amount: 1, - unit: codex_state::ThreadScheduleIntervalUnit::Minutes, - }), - "UTC", + #[tokio::test] + async fn restart_recovers_goal_schedule_run_and_finishes_exactly_once() { + 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, + "watch the release gate after restart", + 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 watch the release gate after restart".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-restart", Duration::from_secs(300)) + .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"); + 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 = + recover_scheduled_run_for_terminal_turn(&reopened, thread_id, "turn-after-restart") + .await + .expect("run recovery should succeed") + .expect("running schedule should recover after restart"); + assert_eq!(schedule.schedule_id, recovered.schedule_id); + assert_eq!(claim.run.run_id, recovered.run_id); + assert_eq!(claim.run.lease_id, recovered.lease_id); + assert_eq!(Some(goal.goal_id), recovered.goal_id); + + let completed_at = scheduled_for + chrono::Duration::seconds(5); + let finished = finish_scheduled_run_state( + &reopened, + recovered.schedule_id.as_str(), + recovered.run_id.as_str(), + recovered.lease_id.as_str(), + recovered.goal_id.as_deref(), + None, + completed_at, + ) + .await + .expect("recovered run should finish") + .expect("recovered run should update exactly once"); + assert_eq!(codex_state::ThreadScheduleStatus::Paused, finished.0.status); + assert_eq!( + codex_state::ThreadScheduleRunStatus::Completed, + finished.1.status + ); + assert_eq!(Some(completed_at), finished.1.completed_at); + assert!( + recover_scheduled_run_for_terminal_turn(&reopened, thread_id, "turn-after-restart",) + .await + .expect("completed run lookup should succeed") + .is_none() + ); + let stats = reopened + .thread_schedules() + .get_thread_schedule_stats(schedule.schedule_id.as_str()) + .await + .expect("schedule stats should load"); + assert_eq!(1, stats.total_runs); + assert_eq!(0, stats.running_runs); + assert_eq!(1, stats.completed_runs); + assert!( + reopened + .thread_schedules() + .claim_due_thread_schedule( + scheduled_for + chrono::Duration::minutes(10), + "lease-duplicate", + Duration::from_secs(300), + ) + .await + .expect("post-restart claim should succeed") + .is_none(), + "terminal recovery must not leave a lease-generated duplicate run" + ); + } + + #[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"); + 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 objective = "watch the default release gate"; + let goal = state_db + .thread_goals() + .replace_thread_goal( + thread_id, + objective, + codex_state::ThreadGoalStatus::Paused, + /*token_budget*/ None, + ) + .await + .expect("paused 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: "default prompt is resolved when the run starts".to_string(), + prompt_source: codex_state::ThreadSchedulePromptSource::Default, + 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-run", Duration::from_secs(300)) + .await + .expect("claim should succeed") + .expect("schedule should claim"); + + let (finished_schedule, finished_run) = finish_scheduled_run_state( + &state_db, + &schedule.schedule_id, + &claim.run.run_id, + "lease-run", + Some(goal.goal_id.as_str()), + Some("failed after resolving the default goal prompt".to_string()), + scheduled_for + chrono::Duration::seconds(5), + ) + .await + .expect("failed run should finish") + .expect("finished rows should load"); + + assert_eq!( + codex_state::ThreadScheduleStatus::Paused, + finished_schedule.status + ); + assert_eq!(None, finished_schedule.next_run_at); + assert_eq!( + codex_state::ThreadScheduleRunStatus::Failed, + finished_run.status + ); + assert_eq!( + Some("failed after resolving the default goal prompt".to_string()), + finished_run.error + ); + } + + #[tokio::test] + async fn stopped_unrelated_goal_does_not_pause_scheduled_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"); + state_db + .thread_goals() + .replace_thread_goal( + thread_id, + "unrelated stopped objective", + 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 run the scheduled objective".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-run", Duration::from_secs(300)) + .await + .expect("claim should succeed") + .expect("schedule should claim"); + let completed_at = scheduled_for + chrono::Duration::seconds(5); + + let (finished_schedule, finished_run) = finish_scheduled_run_state( + &state_db, + &schedule.schedule_id, + &claim.run.run_id, + "lease-run", + None, + None, + completed_at, + ) + .await + .expect("successful run should finish") + .expect("finished rows should load"); + + assert_eq!( + codex_state::ThreadScheduleStatus::Active, + finished_schedule.status + ); + assert_eq!( + Some(scheduled_for + chrono::Duration::minutes(1)), + finished_schedule.next_run_at + ); + assert_eq!( + codex_state::ThreadScheduleRunStatus::Completed, + finished_run.status + ); + } + + #[test] + fn computes_interval_next_run() { + assert_eq!( + Some(at(/*seconds*/ 1_700_000_300)), + next_thread_schedule_run_at( + &codex_state::ThreadScheduleSpec::Interval(codex_state::ThreadScheduleInterval { + amount: 5, + unit: codex_state::ThreadScheduleIntervalUnit::Minutes, + }), + "UTC", + at(/*seconds*/ 1_700_000_000), + ) + .expect("next interval should compute") + ); + } + + #[test] + fn computes_recurring_next_run_from_scheduled_time_without_drift() { + assert_eq!( + Some(at(/*seconds*/ 1_700_000_060)), + next_thread_schedule_run_after_completion( + &codex_state::ThreadScheduleSpec::Interval(codex_state::ThreadScheduleInterval { + amount: 1, + unit: codex_state::ThreadScheduleIntervalUnit::Minutes, + }), + "UTC", + Some(at(/*seconds*/ 1_700_000_000)), + at(/*seconds*/ 1_700_000_005), + ) + .expect("next interval should compute") + ); + } + + #[test] + fn computes_recurring_next_run_skipping_missed_intervals() { + assert_eq!( + Some(at(/*seconds*/ 1_700_000_180)), + next_thread_schedule_run_after_completion( + &codex_state::ThreadScheduleSpec::Interval(codex_state::ThreadScheduleInterval { + amount: 1, + unit: codex_state::ThreadScheduleIntervalUnit::Minutes, + }), + "UTC", Some(at(/*seconds*/ 1_700_000_000)), at(/*seconds*/ 1_700_000_125), ) @@ -3025,6 +3884,69 @@ mod tests { assert_eq!(Some(""), scheduled_goal_objective("/goal")); } + #[test] + fn scheduled_goal_hold_is_status_and_objective_scoped() { + let goal = |status| codex_state::ThreadGoal { + thread_id: ThreadId::new(), + goal_id: "goal-1".to_string(), + objective: "watch the release gate".to_string(), + title: None, + status, + token_budget: None, + tokens_used: 0, + time_used_seconds: 0, + created_at: at(/*seconds*/ 1_700_000_000), + updated_at: at(/*seconds*/ 1_700_000_000), + }; + for status in [ + codex_state::ThreadGoalStatus::Paused, + codex_state::ThreadGoalStatus::Blocked, + codex_state::ThreadGoalStatus::UsageLimited, + codex_state::ThreadGoalStatus::BudgetLimited, + ] { + assert!(scheduled_goal_is_held( + &goal(status), + "watch the release gate" + )); + } + assert!(!scheduled_goal_is_held( + &goal(codex_state::ThreadGoalStatus::Active), + "watch the release gate" + )); + assert!(!scheduled_goal_is_held( + &goal(codex_state::ThreadGoalStatus::Paused), + "ship a different objective" + )); + } + + #[test] + fn scheduled_goal_hold_recovery_is_status_specific() { + let error = |status| { + ScheduledGoalHeld { + goal_id: "goal-release-gate".to_string(), + status, + } + .to_string() + }; + + assert_eq!( + "scheduled goal is paused; resume the goal and its schedule explicitly before another run", + error(codex_state::ThreadGoalStatus::Paused) + ); + assert_eq!( + "scheduled goal is blocked; resume the goal and its schedule explicitly before another run", + error(codex_state::ThreadGoalStatus::Blocked) + ); + assert_eq!( + "scheduled goal is usage limited; wait for or resolve the usage limit, then resume the goal and its schedule explicitly before another run", + error(codex_state::ThreadGoalStatus::UsageLimited) + ); + assert_eq!( + "scheduled goal is budget limited; change the goal token budget before resuming the goal, then resume its schedule explicitly before another run", + error(codex_state::ThreadGoalStatus::BudgetLimited) + ); + } + #[test] fn scheduled_goal_thread_prompt_tells_model_not_to_spawn_followups() { let prompt = scheduled_goal_thread_prompt( diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index 04e76fe0d7..45351ea9fa 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -54,6 +54,7 @@ pub(crate) struct ScheduledThreadScheduleRun { pub(crate) schedule_id: String, pub(crate) run_id: String, pub(crate) lease_id: String, + pub(crate) goal_id: Option, pub(crate) state_db: StateDbHandle, } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index c70a54637c..7537daed1f 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -2741,6 +2741,7 @@ async fn turn_error_lifecycle_exposes_error_and_stores() { turn_level_id: String, turn_id: String, error: CodexErrorInfo, + error_fingerprint: String, saw_session_store: bool, saw_thread_store: bool, } @@ -2749,9 +2750,17 @@ async fn turn_error_lifecycle_exposes_error_and_stores() { records: Arc>>, } + struct LegacyTurnErrorRecorder { + errors: Arc>>, + } + #[async_trait::async_trait] impl codex_extension_api::TurnLifecycleContributor for TurnErrorRecorder { - async fn on_turn_error(&self, input: codex_extension_api::TurnErrorInput<'_>) { + async fn on_turn_error_with_fingerprint( + &self, + input: codex_extension_api::TurnErrorInput<'_>, + error_fingerprint: &str, + ) { self.records .lock() .expect("turn error records lock") @@ -2761,6 +2770,7 @@ async fn turn_error_lifecycle_exposes_error_and_stores() { turn_level_id: input.turn_store.level_id().to_string(), turn_id: input.turn_id.to_string(), error: input.error, + error_fingerprint: error_fingerprint.to_string(), saw_session_store: input .session_store .get::() @@ -2770,12 +2780,26 @@ async fn turn_error_lifecycle_exposes_error_and_stores() { } } + #[async_trait::async_trait] + impl codex_extension_api::TurnLifecycleContributor for LegacyTurnErrorRecorder { + async fn on_turn_error(&self, input: codex_extension_api::TurnErrorInput<'_>) { + self.errors + .lock() + .expect("legacy turn error records lock") + .push(input.error); + } + } + let (mut session, turn_context) = make_session_and_context().await; let records = Arc::new(std::sync::Mutex::new(Vec::new())); + let legacy_errors = Arc::new(std::sync::Mutex::new(Vec::new())); let mut builder = codex_extension_api::ExtensionRegistryBuilder::::new(); builder.turn_lifecycle_contributor(Arc::new(TurnErrorRecorder { records: Arc::clone(&records), })); + builder.turn_lifecycle_contributor(Arc::new(LegacyTurnErrorRecorder { + errors: Arc::clone(&legacy_errors), + })); session.services.extensions = Arc::new(builder.build()); session .services @@ -2786,18 +2810,38 @@ async fn turn_error_lifecycle_exposes_error_and_stores() { .thread_extension_data .insert(ThreadTurnErrorMarker); - let expected = RecordedTurnError { + let expected = |error, error_fingerprint: &str| RecordedTurnError { session_level_id: session.session_id().to_string(), thread_level_id: session.thread_id.to_string(), turn_level_id: turn_context.sub_id.clone(), turn_id: turn_context.sub_id.clone(), - error: CodexErrorInfo::UsageLimitExceeded, + error, + error_fingerprint: error_fingerprint.to_string(), saw_session_store: true, saw_thread_store: true, }; session - .emit_turn_error_lifecycle(&turn_context, CodexErrorInfo::UsageLimitExceeded) + .emit_turn_error_lifecycle( + &turn_context, + &codex_protocol::error::CodexErr::UsageNotIncluded, + ) + .await; + session + .emit_turn_error_lifecycle(&turn_context, &codex_protocol::error::CodexErr::TurnAborted) + .await; + session + .emit_turn_error_lifecycle( + &turn_context, + &codex_protocol::error::CodexErr::RequestTimeout, + ) + .await; + session + .emit_turn_error_lifecycle_with_protocol_error( + &turn_context, + &codex_protocol::error::CodexErr::InvalidImageRequest(), + CodexErrorInfo::BadRequest, + ) .await; let actual = records @@ -2805,7 +2849,36 @@ async fn turn_error_lifecycle_exposes_error_and_stores() { .expect("turn error records lock") .drain(..) .collect::>(); - assert_eq!(vec![expected], actual); + assert_eq!( + vec![ + expected( + CodexErrorInfo::UsageLimitExceeded, + "codex_err:usage_not_included", + ), + expected(CodexErrorInfo::Other, "codex_err:turn_aborted"), + expected(CodexErrorInfo::Other, "codex_err:request_timeout"), + expected( + CodexErrorInfo::BadRequest, + "codex_err:invalid_image_request" + ), + ], + actual, + "broad client variants must retain distinct host classifications without changing the client error" + ); + assert_eq!( + vec![ + CodexErrorInfo::UsageLimitExceeded, + CodexErrorInfo::Other, + CodexErrorInfo::Other, + CodexErrorInfo::BadRequest, + ], + legacy_errors + .lock() + .expect("legacy turn error records lock") + .drain(..) + .collect::>(), + "contributors implementing only the original callback must still receive errors" + ); } #[tokio::test] diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index d9f9291da7..43a9422dcc 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -172,15 +172,13 @@ pub(crate) async fn run_turn( continue; } let err = CodexErr::UsageLimitReached(err); - let error = err.to_codex_protocol_error(); - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) + sess.emit_turn_error_lifecycle(turn_context.as_ref(), &err) .await; error!("Failed to run pre-sampling compact"); return None; } Err(err) => { - let error = err.to_codex_protocol_error(); - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) + sess.emit_turn_error_lifecycle(turn_context.as_ref(), &err) .await; error!("Failed to run pre-sampling compact"); return None; @@ -333,13 +331,11 @@ pub(crate) async fn run_turn( continue; } let err = CodexErr::UsageLimitReached(usage_err); - let error = err.to_codex_protocol_error(); - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) + sess.emit_turn_error_lifecycle(turn_context.as_ref(), &err) .await; return None; } else { - let error = err.to_codex_protocol_error(); - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) + sess.emit_turn_error_lifecycle(turn_context.as_ref(), &err) .await; return None; } @@ -432,13 +428,11 @@ pub(crate) async fn run_turn( continue; } let err = CodexErr::UsageLimitReached(usage_err); - let error = err.to_codex_protocol_error(); - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error) + sess.emit_turn_error_lifecycle(turn_context.as_ref(), &err) .await; return None; } else { - let error = err.to_codex_protocol_error(); - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error) + sess.emit_turn_error_lifecycle(turn_context.as_ref(), &err) .await; return None; } @@ -462,8 +456,7 @@ pub(crate) async fn run_turn( } let e = CodexErr::UsageLimitReached(err); info!("Turn error: {e:#}"); - let error = e.to_codex_protocol_error(); - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) + sess.emit_turn_error_lifecycle(turn_context.as_ref(), &e) .await; sess.track_turn_codex_error(turn_context.as_ref(), &e); let event = EventMsg::Error(e.to_error_event(/*message_prefix*/ None)); @@ -483,8 +476,12 @@ pub(crate) async fn run_turn( sess.track_turn_codex_error(turn_context.as_ref(), &codex_error); let error = CodexErrorInfo::BadRequest; - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) - .await; + sess.emit_turn_error_lifecycle_with_protocol_error( + turn_context.as_ref(), + &codex_error, + error.clone(), + ) + .await; let event = EventMsg::Error(ErrorEvent { message: "Invalid image in your last message. Please remove it and try again." .to_string(), @@ -495,8 +492,7 @@ pub(crate) async fn run_turn( } Err(e) => { info!("Turn error: {e:#}"); - let error = e.to_codex_protocol_error(); - sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone()) + sess.emit_turn_error_lifecycle(turn_context.as_ref(), &e) .await; sess.track_turn_codex_error(turn_context.as_ref(), &e); let event = EventMsg::Error(e.to_error_event(/*message_prefix*/ None)); diff --git a/codex-rs/core/src/tasks/lifecycle.rs b/codex-rs/core/src/tasks/lifecycle.rs index 782cc8c7de..e7eee99495 100644 --- a/codex-rs/core/src/tasks/lifecycle.rs +++ b/codex-rs/core/src/tasks/lifecycle.rs @@ -1,4 +1,6 @@ use codex_extension_api::ExtensionData; +use codex_protocol::error::CodexErr; +use codex_protocol::error::SandboxErr; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TurnAbortReason; @@ -75,18 +77,194 @@ impl Session { pub(crate) async fn emit_turn_error_lifecycle( &self, turn_context: &TurnContext, - error: CodexErrorInfo, + error: &CodexErr, ) { + self.emit_turn_error_lifecycle_with_protocol_error( + turn_context, + error, + error.to_codex_protocol_error(), + ) + .await; + } + + pub(crate) async fn emit_turn_error_lifecycle_with_protocol_error( + &self, + turn_context: &TurnContext, + error: &CodexErr, + protocol_error: CodexErrorInfo, + ) { + let error_fingerprint = stable_turn_error_fingerprint(error); for contributor in self.services.extensions.turn_lifecycle_contributors() { contributor - .on_turn_error(codex_extension_api::TurnErrorInput { - turn_id: turn_context.sub_id.as_str(), - error: error.clone(), - session_store: &self.services.session_extension_data, - thread_store: &self.services.thread_extension_data, - turn_store: turn_context.extension_data.as_ref(), - }) + .on_turn_error_with_fingerprint( + codex_extension_api::TurnErrorInput { + turn_id: turn_context.sub_id.as_str(), + error: protocol_error.clone(), + session_store: &self.services.session_extension_data, + thread_store: &self.services.thread_extension_data, + turn_store: turn_context.extension_data.as_ref(), + }, + error_fingerprint.as_str(), + ) .await; } } } + +fn stable_turn_error_fingerprint(error: &CodexErr) -> String { + let fingerprint = match error { + CodexErr::TurnAborted => "codex_err:turn_aborted", + CodexErr::Stream(..) => "codex_err:stream", + CodexErr::ContextWindowExceeded => "codex_err:context_window_exceeded", + CodexErr::ThreadNotFound(_) => "codex_err:thread_not_found", + CodexErr::AgentLimitReached { .. } => "codex_err:agent_limit_reached", + CodexErr::SessionConfiguredNotFirstEvent => "codex_err:session_configured_not_first_event", + CodexErr::Timeout => "codex_err:timeout", + CodexErr::RequestTimeout => "codex_err:request_timeout", + CodexErr::Spawn => "codex_err:spawn", + CodexErr::Interrupted => "codex_err:interrupted", + CodexErr::InvalidRequest(_) => "codex_err:invalid_request", + CodexErr::InvalidImageRequest() => "codex_err:invalid_image_request", + CodexErr::UsageLimitReached(_) => "codex_err:usage_limit_reached", + CodexErr::ServerOverloaded => "codex_err:server_overloaded", + CodexErr::CyberPolicy { .. } => "codex_err:cyber_policy", + CodexErr::QuotaExceeded => "codex_err:quota_exceeded", + CodexErr::UsageNotIncluded => "codex_err:usage_not_included", + CodexErr::InternalServerError => "codex_err:internal_server_error", + CodexErr::InternalAgentDied => "codex_err:internal_agent_died", + CodexErr::LandlockSandboxExecutableNotProvided => { + "codex_err:landlock_sandbox_executable_not_provided" + } + CodexErr::UnsupportedOperation(_) => "codex_err:unsupported_operation", + CodexErr::RefreshTokenFailed(_) => "codex_err:refresh_token_failed", + CodexErr::Fatal(_) => "codex_err:fatal", + CodexErr::Io(error) => return stable_io_error_fingerprint(error.kind()), + CodexErr::EnvVar(_) => "codex_err:env_var", + CodexErr::Json(error) => match error.classify() { + serde_json::error::Category::Io => "codex_err:json:io", + serde_json::error::Category::Syntax => "codex_err:json:syntax", + serde_json::error::Category::Data => "codex_err:json:data", + serde_json::error::Category::Eof => "codex_err:json:eof", + }, + CodexErr::TokioJoin(error) if error.is_cancelled() => "codex_err:tokio_join:cancelled", + CodexErr::TokioJoin(error) if error.is_panic() => "codex_err:tokio_join:panic", + CodexErr::TokioJoin(_) => "codex_err:tokio_join:other", + #[cfg(target_os = "linux")] + CodexErr::LandlockRuleset(_) => "codex_err:landlock_ruleset", + #[cfg(target_os = "linux")] + CodexErr::LandlockPathFd(_) => "codex_err:landlock_path_fd", + CodexErr::UnexpectedStatus(error) => { + return fingerprint_with_http_status( + "codex_err:unexpected_status", + Some(error.status.as_u16()), + ); + } + CodexErr::RetryLimit(error) => { + return fingerprint_with_http_status( + "codex_err:retry_limit", + Some(error.status.as_u16()), + ); + } + CodexErr::ConnectionFailed(error) => { + return fingerprint_with_http_status( + "codex_err:connection_failed", + error.source.status().map(|status| status.as_u16()), + ); + } + CodexErr::ResponseStreamFailed(error) => { + return fingerprint_with_http_status( + "codex_err:response_stream_failed", + error.source.status().map(|status| status.as_u16()), + ); + } + CodexErr::Sandbox(error) => return stable_sandbox_error_fingerprint(error), + }; + fingerprint.to_string() +} + +fn fingerprint_with_http_status(kind: &str, status: Option) -> String { + match status { + Some(status) => format!("{kind}:http_{status}"), + None => format!("{kind}:http_unknown"), + } +} + +fn stable_io_error_fingerprint(kind: std::io::ErrorKind) -> String { + let kind = match kind { + std::io::ErrorKind::NotFound => "not_found", + std::io::ErrorKind::PermissionDenied => "permission_denied", + std::io::ErrorKind::ConnectionRefused => "connection_refused", + std::io::ErrorKind::ConnectionReset => "connection_reset", + std::io::ErrorKind::ConnectionAborted => "connection_aborted", + std::io::ErrorKind::NotConnected => "not_connected", + std::io::ErrorKind::AddrInUse => "address_in_use", + std::io::ErrorKind::AddrNotAvailable => "address_not_available", + std::io::ErrorKind::BrokenPipe => "broken_pipe", + std::io::ErrorKind::AlreadyExists => "already_exists", + std::io::ErrorKind::WouldBlock => "would_block", + std::io::ErrorKind::InvalidInput => "invalid_input", + std::io::ErrorKind::InvalidData => "invalid_data", + std::io::ErrorKind::TimedOut => "timed_out", + std::io::ErrorKind::WriteZero => "write_zero", + std::io::ErrorKind::Interrupted => "interrupted", + std::io::ErrorKind::Unsupported => "unsupported", + std::io::ErrorKind::UnexpectedEof => "unexpected_eof", + std::io::ErrorKind::OutOfMemory => "out_of_memory", + _ => "other", + }; + format!("codex_err:io:{kind}") +} + +fn stable_sandbox_error_fingerprint(error: &SandboxErr) -> String { + let fingerprint = match error { + SandboxErr::Denied { .. } => "codex_err:sandbox:denied", + #[cfg(target_os = "linux")] + SandboxErr::SeccompInstall(_) => "codex_err:sandbox:seccomp_install", + #[cfg(target_os = "linux")] + SandboxErr::SeccompBackend(_) => "codex_err:sandbox:seccomp_backend", + SandboxErr::Timeout { .. } => "codex_err:sandbox:timeout", + SandboxErr::LandlockRestrict => "codex_err:sandbox:landlock_restrict", + SandboxErr::Signal(_) => "codex_err:sandbox:signal", + }; + fingerprint.to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + use codex_protocol::error::UnexpectedResponseError; + use pretty_assertions::assert_eq; + + #[test] + fn broad_protocol_errors_keep_distinct_source_fingerprints() { + let turn_aborted = CodexErr::TurnAborted; + let request_timeout = CodexErr::RequestTimeout; + + assert_eq!( + turn_aborted.to_codex_protocol_error(), + request_timeout.to_codex_protocol_error() + ); + assert_ne!( + stable_turn_error_fingerprint(&turn_aborted), + stable_turn_error_fingerprint(&request_timeout) + ); + } + + #[test] + fn fingerprint_omits_error_payloads_and_request_ids() { + let error = CodexErr::UnexpectedStatus(UnexpectedResponseError { + status: reqwest::StatusCode::BAD_GATEWAY, + body: "secret response body".to_string(), + url: Some("https://example.invalid/private".to_string()), + cf_ray: Some("volatile-ray".to_string()), + request_id: Some("volatile-request-id".to_string()), + identity_authorization_error: Some("secret authorization detail".to_string()), + identity_error_code: Some("volatile-code".to_string()), + }); + + assert_eq!( + "codex_err:unexpected_status:http_502", + stable_turn_error_fingerprint(&error) + ); + } +} diff --git a/codex-rs/core/src/tools/handlers/loop_control.rs b/codex-rs/core/src/tools/handlers/loop_control.rs index e7d898e92d..64bb39f847 100644 --- a/codex-rs/core/src/tools/handlers/loop_control.rs +++ b/codex-rs/core/src/tools/handlers/loop_control.rs @@ -1298,12 +1298,15 @@ mod tests { .expect("first run should be due"); runtime .thread_schedules() - .mark_thread_schedule_run_started( - &schedule.schedule_id, - &first_claim.run.run_id, - "lease-complete", - "turn-complete", - ) + .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"); @@ -1329,12 +1332,15 @@ mod tests { .expect("second run should be due"); runtime .thread_schedules() - .mark_thread_schedule_run_started( - &schedule.schedule_id, - &second_claim.run.run_id, - "lease-fail", - "turn-fail", - ) + .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"); @@ -1565,12 +1571,15 @@ mod tests { .expect("run should be due"); runtime .thread_schedules() - .mark_thread_schedule_run_started( - &first.schedule_id, - &claim.run.run_id, - "lease-complete", - "turn-complete", - ) + .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"); diff --git a/codex-rs/core/tests/invalid_image_lifecycle.rs b/codex-rs/core/tests/invalid_image_lifecycle.rs new file mode 100644 index 0000000000..e30ff3b674 --- /dev/null +++ b/codex-rs/core/tests/invalid_image_lifecycle.rs @@ -0,0 +1,123 @@ +#![cfg(not(debug_assertions))] + +use std::sync::Arc; +use std::sync::Mutex; + +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::TurnErrorInput; +use codex_extension_api::TurnLifecycleContributor; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::CodexErrorInfo; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::user_input::UserInput; +use core_test_support::responses; +use core_test_support::responses::start_mock_server; +use core_test_support::test_codex::TestCodex; +use core_test_support::test_codex::test_codex; +use core_test_support::test_codex::turn_permission_fields; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use wiremock::ResponseTemplate; + +#[derive(Default)] +struct RecordingTurnErrors { + errors: Mutex>, +} + +impl RecordingTurnErrors { + fn lock_errors(&self) -> std::sync::MutexGuard<'_, Vec<(CodexErrorInfo, String)>> { + match self.errors.lock() { + Ok(errors) => errors, + Err(poisoned) => poisoned.into_inner(), + } + } +} + +#[async_trait::async_trait] +impl TurnLifecycleContributor for RecordingTurnErrors { + async fn on_turn_error_with_fingerprint( + &self, + input: TurnErrorInput<'_>, + error_fingerprint: &str, + ) { + self.lock_errors() + .push((input.error, error_fingerprint.to_string())); + } +} + +fn disabled_user_turn(test: &TestCodex, items: Vec, model: String) -> Op { + let (sandbox_policy, permission_profile) = + turn_permission_fields(PermissionProfile::Disabled, test.config.cwd.as_path()); + Op::UserInput { + items, + environments: None, + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: codex_protocol::protocol::ThreadSettingsOverrides { + cwd: Some(test.config.cwd.clone()), + approval_policy: Some(AskForApproval::Never), + sandbox_policy: Some(sandbox_policy), + permission_profile, + collaboration_mode: Some(codex_protocol::config_types::CollaborationMode { + mode: codex_protocol::config_types::ModeKind::Default, + settings: codex_protocol::config_types::Settings { + model, + reasoning_effort: None, + developer_instructions: None, + }, + }), + ..Default::default() + }, + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn invalid_image_run_turn_preserves_bad_request_and_stable_fingerprint() -> anyhow::Result<()> +{ + let server = start_mock_server().await; + const INVALID_IMAGE_ERROR: &str = + "The image data you provided does not represent a valid image"; + responses::mount_response_once( + &server, + ResponseTemplate::new(400) + .insert_header("content-type", "text/plain") + .set_body_string(INVALID_IMAGE_ERROR), + ) + .await; + + let recorder = Arc::new(RecordingTurnErrors::default()); + let mut extensions = ExtensionRegistryBuilder::::new(); + extensions.turn_lifecycle_contributor(recorder.clone()); + let mut builder = test_codex().with_extensions(Arc::new(extensions.build())); + let test = builder.build(&server).await?; + let session_model = test.session_configured.model.clone(); + + test.codex + .submit(disabled_user_turn( + &test, + vec![UserInput::Text { + text: "trigger invalid image response mapping".into(), + text_elements: Vec::new(), + }], + session_model, + )) + .await?; + + let event = wait_for_event(&test.codex, |event| matches!(event, EventMsg::Error(_))).await; + let EventMsg::Error(error) = event else { + unreachable!("wait predicate only accepts error events"); + }; + assert_eq!(Some(CodexErrorInfo::BadRequest), error.codex_error_info); + assert_eq!( + vec![( + CodexErrorInfo::BadRequest, + "codex_err:invalid_image_request".to_string(), + )], + *recorder.lock_errors() + ); + + Ok(()) +} diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 8706e8ee7a..3d30c5befd 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -87,6 +87,19 @@ pub trait TurnLifecycleContributor: Send + Sync { /// Called when the host observes an error for a running turn. async fn on_turn_error(&self, _input: TurnErrorInput<'_>) {} + + /// Called with a stable, non-sensitive host classification of the error. + /// + /// The default preserves compatibility for contributors that only need the + /// client-facing error. Hosts must exclude messages, request identifiers, + /// and other volatile or secret-bearing payloads from `error_fingerprint`. + async fn on_turn_error_with_fingerprint( + &self, + input: TurnErrorInput<'_>, + _error_fingerprint: &str, + ) { + self.on_turn_error(input).await; + } } /// Extension contribution that can add turn-local model input. diff --git a/codex-rs/ext/goal/Cargo.toml b/codex-rs/ext/goal/Cargo.toml index 6116383243..e3f234b981 100644 --- a/codex-rs/ext/goal/Cargo.toml +++ b/codex-rs/ext/goal/Cargo.toml @@ -31,5 +31,6 @@ tracing = { workspace = true } anyhow = { workspace = true } chrono = { workspace = true } pretty_assertions = { workspace = true } +sqlx = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/ext/goal/src/accounting.rs b/codex-rs/ext/goal/src/accounting.rs index 6b5d483c1e..e62a8ce67e 100644 --- a/codex-rs/ext/goal/src/accounting.rs +++ b/codex-rs/ext/goal/src/accounting.rs @@ -29,6 +29,7 @@ struct GoalTurnAccounting { last_accounted_token_usage: TokenUsage, active_goal_id: Option, account_tokens: bool, + error_observed: bool, } #[derive(Debug)] @@ -117,6 +118,19 @@ impl GoalAccountingState { turn.account_tokens.then(|| turn.active_goal_id()).flatten() } + pub(crate) fn mark_turn_error_observed(&self, turn_id: &str) { + if let Some(turn) = self.inner().turns.get_mut(turn_id) { + turn.error_observed = true; + } + } + + pub(crate) fn turn_error_observed(&self, turn_id: &str) -> bool { + self.inner() + .turns + .get(turn_id) + .is_some_and(|turn| turn.error_observed) + } + pub(crate) fn record_token_usage( &self, turn_id: impl Into, @@ -373,6 +387,7 @@ impl GoalTurnAccounting { current_token_usage, active_goal_id: None, account_tokens, + error_observed: false, } } diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index 40787d3d2e..136d3f26ca 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -92,6 +92,40 @@ impl GoalExtension { goals_config: Arc::new(goals_config), } } + + async fn handle_turn_error(&self, input: TurnErrorInput<'_>, error_fingerprint: &str) { + let Some(runtime) = goal_runtime_handle(input.thread_store) else { + return; + }; + // Core emits turn-stop after terminal turn errors. Record the error + // before any fallible persistence so that later success-only cleanup + // cannot erase this or an earlier durable blocker observation. + runtime + .accounting_state() + .mark_turn_error_observed(input.turn_id); + + let reason = match &input.error { + CodexErrorInfo::UsageLimitExceeded => ActiveGoalStopReason::UsageLimit, + // The turn has ended because the error was non-retryable or its + // retries were exhausted. Hold the goal to prevent automatic + // continuation from looping and consuming tokens. The runtime + // promotes the same blocker to Blocked only after the required + // number of consecutive goal turns. + _ => ActiveGoalStopReason::TurnError { + error: input.error.clone(), + fingerprint: error_fingerprint.to_string(), + }, + }; + if let Err(err) = runtime + .stop_active_goal_for_turn(input.turn_id, reason) + .await + { + tracing::warn!( + error = ?input.error, + "failed to stop active goal after turn error: {err}" + ); + } + } } #[async_trait] @@ -276,21 +310,40 @@ where } let turn_id = input.turn_store.level_id(); - if let Err(err) = runtime + let accounting_state = runtime.accounting_state(); + let goal_id = (!accounting_state.turn_error_observed(turn_id)) + .then(|| accounting_state.active_goal_id_for_current_turn(turn_id)) + .flatten(); + if let Some(goal_id) = goal_id.as_deref() { + let clear_result = codex_state::busy_retry::retry_on_busy( + "clear blocker audit after successful turn", + || { + self.state_dbs + .thread_goals() + .clear_thread_goal_blocker_audit_for_goal(runtime.thread_id(), goal_id) + }, + ) + .await; + if let Err(err) = clear_result { + tracing::warn!( + "failed to clear blocker audit after successful turn {turn_id}: {err}" + ); + } + } + let accounting_result = runtime .account_active_goal_progress( turn_id, &format!("{turn_id}:turn-stop"), codex_state::GoalAccountingMode::ActiveOnly, BudgetLimitedGoalDisposition::ClearActive, ) - .await - { + .await; + runtime.accounting_state().finish_turn(turn_id); + if let Err(err) = accounting_result { tracing::warn!( "failed to account active goal progress at turn stop for {turn_id}: {err}" ); - return; } - runtime.accounting_state().finish_turn(turn_id); } async fn on_turn_abort(&self, input: TurnAbortInput<'_>) { @@ -302,45 +355,52 @@ where } let turn_id = input.turn_store.level_id(); - if let Err(err) = runtime + let accounting_state = runtime.accounting_state(); + let goal_id = (!accounting_state.turn_error_observed(turn_id)) + .then(|| accounting_state.active_goal_id_for_current_turn(turn_id)) + .flatten(); + if let Some(goal_id) = goal_id.as_deref() { + let clear_result = codex_state::busy_retry::retry_on_busy( + "clear blocker audit after turn abort", + || { + self.state_dbs + .thread_goals() + .clear_thread_goal_blocker_audit_for_goal(runtime.thread_id(), goal_id) + }, + ) + .await; + if let Err(err) = clear_result { + tracing::warn!("failed to clear blocker audit after turn abort {turn_id}: {err}"); + } + } + let accounting_result = runtime .account_active_goal_progress( turn_id, &format!("{turn_id}:turn-abort"), codex_state::GoalAccountingMode::ActiveOnly, BudgetLimitedGoalDisposition::ClearActive, ) - .await - { + .await; + runtime.accounting_state().finish_turn(turn_id); + if let Err(err) = accounting_result { tracing::warn!( "failed to account active goal progress after turn abort for {turn_id}: {err}" ); - return; } - runtime.accounting_state().finish_turn(turn_id); } async fn on_turn_error(&self, input: TurnErrorInput<'_>) { - let Some(runtime) = goal_runtime_handle(input.thread_store) else { - return; - }; + let error_fingerprint = protocol_error_fingerprint(&input.error); + self.handle_turn_error(input, error_fingerprint.as_str()) + .await; + } - let reason = match &input.error { - CodexErrorInfo::UsageLimitExceeded => ActiveGoalStopReason::UsageLimit, - // The turn has ended because the error was non-retryable or its - // retries were exhausted. Block the goal to prevent automatic - // continuation from looping and consuming tokens, as can happen - // with compaction errors. - _ => ActiveGoalStopReason::TurnError(input.error.clone()), - }; - if let Err(err) = runtime - .stop_active_goal_for_turn(input.turn_id, reason) - .await - { - tracing::warn!( - error = ?input.error, - "failed to stop active goal after turn error: {err}" - ); - } + async fn on_turn_error_with_fingerprint( + &self, + input: TurnErrorInput<'_>, + error_fingerprint: &str, + ) { + self.handle_turn_error(input, error_fingerprint).await; } } @@ -534,6 +594,58 @@ fn goal_runtime_handle(thread_store: &ExtensionData) -> Option() } +fn protocol_error_fingerprint(error: &CodexErrorInfo) -> String { + let fingerprint = match error { + CodexErrorInfo::ContextWindowExceeded => "codex_err:protocol:context_window_exceeded", + CodexErrorInfo::UsageLimitExceeded => "codex_err:protocol:usage_limit_exceeded", + CodexErrorInfo::ServerOverloaded => "codex_err:protocol:server_overloaded", + CodexErrorInfo::CyberPolicy => "codex_err:protocol:cyber_policy", + CodexErrorInfo::HttpConnectionFailed { http_status_code } => { + return protocol_http_error_fingerprint("http_connection_failed", *http_status_code); + } + CodexErrorInfo::ResponseStreamConnectionFailed { http_status_code } => { + return protocol_http_error_fingerprint( + "response_stream_connection_failed", + *http_status_code, + ); + } + CodexErrorInfo::InternalServerError => "codex_err:protocol:internal_server_error", + CodexErrorInfo::Unauthorized => "codex_err:protocol:unauthorized", + CodexErrorInfo::BadRequest => "codex_err:protocol:bad_request", + CodexErrorInfo::SandboxError => "codex_err:protocol:sandbox_error", + CodexErrorInfo::ResponseStreamDisconnected { http_status_code } => { + return protocol_http_error_fingerprint( + "response_stream_disconnected", + *http_status_code, + ); + } + CodexErrorInfo::ResponseTooManyFailedAttempts { http_status_code } => { + return protocol_http_error_fingerprint( + "response_too_many_failed_attempts", + *http_status_code, + ); + } + CodexErrorInfo::ActiveTurnNotSteerable { turn_kind } => match turn_kind { + codex_protocol::protocol::NonSteerableTurnKind::Review => { + "codex_err:protocol:active_turn_not_steerable:review" + } + codex_protocol::protocol::NonSteerableTurnKind::Compact => { + "codex_err:protocol:active_turn_not_steerable:compact" + } + }, + CodexErrorInfo::ThreadRollbackFailed => "codex_err:protocol:thread_rollback_failed", + CodexErrorInfo::Other => "codex_err:protocol:other", + }; + fingerprint.to_string() +} + +fn protocol_http_error_fingerprint(kind: &str, status: Option) -> String { + match status { + Some(status) => format!("codex_err:protocol:{kind}:http_{status}"), + None => format!("codex_err:protocol:{kind}:http_unknown"), + } +} + fn tool_attempt_counts_for_goal_progress(outcome: ToolCallOutcome) -> bool { match outcome { ToolCallOutcome::Completed { .. } => true, diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index 59f8db18e9..4342b35631 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/codex-rs/ext/goal/src/runtime.rs @@ -40,10 +40,15 @@ pub(crate) struct GoalRuntimeConfig { } pub(crate) enum ActiveGoalStopReason { - TurnError(CodexErrorInfo), + TurnError { + error: CodexErrorInfo, + fingerprint: String, + }, UsageLimit, } +const REQUIRED_CONSECUTIVE_BLOCKER_TURNS: u8 = 3; + struct GoalRuntimeInner { thread_id: ThreadId, state_dbs: Arc, @@ -450,6 +455,12 @@ impl GoalRuntimeHandle { ); } self.inner.accounting_state.clear_active_goal(); + self.inner + .state_dbs + .thread_goals() + .clear_thread_goal_blocker_audit(self.thread_id()) + .await + .map_err(|err| err.to_string())?; Ok(()) } @@ -479,13 +490,9 @@ impl GoalRuntimeHandle { return Ok(()); } - let (event_name, status) = match &reason { - ActiveGoalStopReason::TurnError(_) => { - ("turn-error", codex_state::ThreadGoalStatus::Blocked) - } - ActiveGoalStopReason::UsageLimit => { - ("usage-limit", codex_state::ThreadGoalStatus::UsageLimited) - } + let event_name = match &reason { + ActiveGoalStopReason::TurnError { .. } => "turn-error", + ActiveGoalStopReason::UsageLimit => "usage-limit", }; self.account_active_goal_progress( turn_id, @@ -506,35 +513,89 @@ impl GoalRuntimeHandle { self.inner.accounting_state.clear_active_goal(); return Ok(()); }; - let can_stop = active_goal.status == codex_state::ThreadGoalStatus::Active - || (active_goal.status == codex_state::ThreadGoalStatus::BudgetLimited - && status == codex_state::ThreadGoalStatus::UsageLimited); - if !can_stop { - self.inner.accounting_state.clear_active_goal(); - return Ok(()); - } let previous_status = Some(active_goal.status); - let Some(goal) = self - .inner - .state_dbs - .thread_goals() - .update_thread_goal( - self.thread_id(), - codex_state::GoalUpdate { - objective: None, - title: None, - status: Some(status), - token_budget: None, - expected_goal_id: Some(active_goal.goal_id), - }, - ) - .await - .map_err(|err| err.to_string())? - else { - return Ok(()); + let goal = match &reason { + ActiveGoalStopReason::TurnError { fingerprint, .. } => { + match self + .inner + .state_dbs + .thread_goals() + .observe_active_thread_goal_blocker( + self.thread_id(), + active_goal.goal_id.as_str(), + turn_id, + fingerprint, + REQUIRED_CONSECUTIVE_BLOCKER_TURNS, + ) + .await + .map_err(|err| err.to_string())? + { + codex_state::GoalBlockerAuditOutcome::Updated { goal, .. } => goal, + codex_state::GoalBlockerAuditOutcome::Unchanged(goal) => { + self.inner.accounting_state.clear_active_goal(); + if let Some(goal) = goal + && !matches!( + goal.status, + codex_state::ThreadGoalStatus::Active + | codex_state::ThreadGoalStatus::Paused + ) + { + self.inner + .state_dbs + .thread_goals() + .clear_thread_goal_blocker_audit_for_goal( + self.thread_id(), + goal.goal_id.as_str(), + ) + .await + .map_err(|err| err.to_string())?; + } + return Ok(()); + } + } + } + ActiveGoalStopReason::UsageLimit => { + let can_stop = active_goal.status == codex_state::ThreadGoalStatus::Active + || active_goal.status == codex_state::ThreadGoalStatus::BudgetLimited; + if !can_stop { + self.inner.accounting_state.clear_active_goal(); + return Ok(()); + } + let Some(goal) = self + .inner + .state_dbs + .thread_goals() + .update_thread_goal( + self.thread_id(), + codex_state::GoalUpdate { + objective: None, + title: None, + status: Some(codex_state::ThreadGoalStatus::UsageLimited), + token_budget: None, + expected_goal_id: Some(active_goal.goal_id), + }, + ) + .await + .map_err(|err| err.to_string())? + else { + return Ok(()); + }; + self.inner + .state_dbs + .thread_goals() + .clear_thread_goal_blocker_audit_for_goal( + self.thread_id(), + goal.goal_id.as_str(), + ) + .await + .map_err(|err| err.to_string())?; + goal + } }; match &reason { - ActiveGoalStopReason::TurnError(error) => { + ActiveGoalStopReason::TurnError { error, .. } + if goal.status == codex_state::ThreadGoalStatus::Blocked => + { crate::pending_interaction::record_goal_turn_error_status_wait( self.inner.state_dbs.as_ref(), self.thread_id(), @@ -544,6 +605,7 @@ impl GoalRuntimeHandle { ) .await?; } + ActiveGoalStopReason::TurnError { .. } => {} ActiveGoalStopReason::UsageLimit => { crate::pending_interaction::record_goal_status_wait( self.inner.state_dbs.as_ref(), diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 2fc045c515..1ef56b3ecc 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -19,6 +19,7 @@ use codex_extension_api::ToolExecutor; use codex_extension_api::ToolFinishInput; use codex_extension_api::ToolPayload; use codex_extension_api::ToolSpec; +use codex_extension_api::TurnAbortInput; use codex_extension_api::TurnErrorInput; use codex_extension_api::TurnStartInput; use codex_extension_api::TurnStopInput; @@ -43,8 +44,11 @@ use codex_protocol::protocol::ThreadGoalStatus; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::TruncationPolicy; +use codex_protocol::protocol::TurnAbortReason; use pretty_assertions::assert_eq; use serde_json::json; +use sqlx::sqlite::SqliteConnectOptions; +use sqlx::sqlite::SqlitePoolOptions; use tempfile::TempDir; #[tokio::test] @@ -1895,7 +1899,7 @@ async fn turn_error_usage_limit_accounts_progress_and_clears_accounting() -> any } #[tokio::test] -async fn turn_error_blocks_goal() -> anyhow::Result<()> { +async fn one_turn_error_keeps_goal_active_without_blocking_it() -> anyhow::Result<()> { let runtime = test_runtime().await?; let thread_id = test_thread_id()?; seed_thread_metadata(runtime.as_ref(), thread_id).await?; @@ -1920,12 +1924,22 @@ async fn turn_error_blocks_goal() -> anyhow::Result<()> { .get_thread_goal(thread_id) .await? .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; - assert_eq!(codex_state::ThreadGoalStatus::Blocked, goal.status); + assert_eq!(codex_state::ThreadGoalStatus::Active, goal.status); + assert!( + pending_interactions_for_kind( + runtime.as_ref(), + thread_id, + codex_state::PendingInteractionKind::Blocked, + ) + .await? + .is_empty(), + "one exhausted turn must not create a blocked wait" + ); Ok(()) } #[tokio::test] -async fn turn_error_blocks_goal_plan_node_with_actionable_wait_payload() -> anyhow::Result<()> { +async fn third_consecutive_matching_turn_error_blocks_goal_plan_node() -> anyhow::Result<()> { let runtime = test_runtime().await?; let thread_id = test_thread_id()?; seed_thread_metadata(runtime.as_ref(), thread_id).await?; @@ -1956,6 +1970,62 @@ async fn turn_error_blocks_goal_plan_node_with_actionable_wait_payload() -> anyh harness .notify_turn_error("turn-1", CodexErrorInfo::ContextWindowExceeded) .await; + harness.stop_turn("turn-1").await; + + let active_goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + assert_eq!(codex_state::ThreadGoalStatus::Active, active_goal.status); + + let active_plan = runtime + .thread_goals() + .list_thread_goal_plans(thread_id) + .await? + .into_iter() + .next() + .ok_or_else(|| anyhow::anyhow!("goal plan should exist"))?; + assert_eq!( + codex_state::ThreadGoalPlanStatus::Active, + active_plan.plan.status + ); + assert_eq!( + codex_state::ThreadGoalPlanNodeStatus::Active, + active_plan.nodes[0].status + ); + assert_eq!( + codex_state::ThreadGoalPlanNodeStatus::Pending, + active_plan.nodes[1].status + ); + assert!( + pending_interactions_for_kind( + runtime.as_ref(), + thread_id, + codex_state::PendingInteractionKind::Blocked, + ) + .await? + .is_empty(), + "a pre-threshold goal plan must remain active without being reported as blocked" + ); + + harness.start_turn("turn-2", &TokenUsage::default()).await; + harness + .notify_turn_error("turn-2", CodexErrorInfo::ContextWindowExceeded) + .await; + harness.stop_turn("turn-2").await; + + let active_goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + assert_eq!(codex_state::ThreadGoalStatus::Active, active_goal.status); + + harness.start_turn("turn-3", &TokenUsage::default()).await; + harness + .notify_turn_error("turn-3", CodexErrorInfo::ContextWindowExceeded) + .await; let goal = runtime .thread_goals() @@ -1989,7 +2059,7 @@ async fn turn_error_blocks_goal_plan_node_with_actionable_wait_payload() -> anyh .await?; assert_eq!(1, pending.len()); assert_eq!(Some(goal.goal_id.as_str()), pending[0].source_id.as_deref()); - assert_eq!(Some("turn-1"), pending[0].turn_id.as_deref()); + assert_eq!(Some("turn-3"), pending[0].turn_id.as_deref()); assert_eq!(pending[0].request_payload_json["reason"], "turn-error"); assert_eq!( pending[0].request_payload_json["terminalError"], @@ -1999,6 +2069,587 @@ async fn turn_error_blocks_goal_plan_node_with_actionable_wait_payload() -> anyh "action": "The turn exceeded the model context window. Reduce prompt/history size or run a cleanup/compaction before retrying the loop.", }) ); + + harness.stop_turn("turn-3").await; + resume_goal_in_turn(&harness, "turn-4").await?; + harness + .notify_turn_error("turn-4", CodexErrorInfo::ContextWindowExceeded) + .await; + + let resumed_goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("resumed goal should exist"))?; + assert_eq!( + codex_state::ThreadGoalStatus::Active, + resumed_goal.status, + "resuming a blocked goal must start a fresh blocker audit" + ); + Ok(()) +} + +#[tokio::test] +async fn blocker_audit_survives_restarts_and_deduplicates_replayed_turns() -> anyhow::Result<()> { + let tempdir = TempDir::new()?; + let codex_home = tempdir.path().to_path_buf(); + let thread_id = test_thread_id()?; + + { + let runtime = + codex_state::StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let harness = GoalExtensionHarness::new(runtime, thread_id).await?; + harness.start_turn("turn-1", &TokenUsage::default()).await; + { + let tools = harness.tools(); + tool_by_name(&tools, "create_goal") + .handle(tool_call( + "create_goal", + "call-create-restart-goal", + json!({ "objective": "preserve blocker audit across restarts" }), + )) + .await?; + } + harness + .notify_turn_error("turn-1", CodexErrorInfo::ContextWindowExceeded) + .await; + harness.stop_turn("turn-1").await; + } + + { + let runtime = + codex_state::StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + harness.resume_thread().await; + harness.start_turn("turn-1", &TokenUsage::default()).await; + harness + .notify_turn_error("turn-1", CodexErrorInfo::ContextWindowExceeded) + .await; + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist after replay"))?; + assert_eq!( + codex_state::ThreadGoalStatus::Active, + goal.status, + "redelivering one host turn after restart must not advance the audit" + ); + harness.stop_turn("turn-1").await; + } + + { + let runtime = + codex_state::StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await?; + let harness = GoalExtensionHarness::new(runtime, thread_id).await?; + harness.resume_thread().await; + harness.start_turn("turn-2", &TokenUsage::default()).await; + harness + .notify_turn_error("turn-2", CodexErrorInfo::ContextWindowExceeded) + .await; + harness.stop_turn("turn-2").await; + } + + { + let runtime = + codex_state::StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + harness.resume_thread().await; + harness.start_turn("turn-1", &TokenUsage::default()).await; + harness + .notify_turn_error("turn-1", CodexErrorInfo::ContextWindowExceeded) + .await; + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist after delayed replay"))?; + assert_eq!( + codex_state::ThreadGoalStatus::Active, + goal.status, + "a delayed replay of any already-counted turn must not advance the audit" + ); + harness.stop_turn("turn-1").await; + } + + { + let runtime = + codex_state::StateRuntime::init(codex_home, "test-provider".to_string()).await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + harness.resume_thread().await; + harness.start_turn("turn-3", &TokenUsage::default()).await; + harness + .notify_turn_error("turn-3", CodexErrorInfo::ContextWindowExceeded) + .await; + + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist after restart"))?; + assert_eq!( + codex_state::ThreadGoalStatus::Blocked, + goal.status, + "three distinct matching blocker turns must survive fresh goal and SQLite runtimes" + ); + } + Ok(()) +} + +#[tokio::test] +async fn turn_stop_after_error_does_not_clear_active_goal_audit() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + harness.start_turn("turn-1", &TokenUsage::default()).await; + + let tools = harness.tools(); + tool_by_name(&tools, "create_goal") + .handle(tool_call( + "create_goal", + "call-create-error-cleanup-goal", + json!({ "objective": "preserve committed error audit through turn stop" }), + )) + .await?; + harness + .notify_turn_error("turn-1", CodexErrorInfo::ContextWindowExceeded) + .await; + + harness.stop_turn("turn-1").await; + + harness.start_turn("turn-2", &TokenUsage::default()).await; + harness + .notify_turn_error("turn-2", CodexErrorInfo::ContextWindowExceeded) + .await; + harness.stop_turn("turn-2").await; + + harness.start_turn("turn-3", &TokenUsage::default()).await; + harness + .notify_turn_error("turn-3", CodexErrorInfo::ContextWindowExceeded) + .await; + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + assert_eq!( + codex_state::ThreadGoalStatus::Blocked, + goal.status, + "turn stop after an error must not erase the committed blocker observation" + ); + Ok(()) +} + +#[derive(Clone, Copy, Debug)] +enum TurnFinalizer { + Stop, + Abort, +} + +#[derive(Clone, Copy, Debug)] +enum InjectedFinalizationFailure { + AuditClear, + Accounting, +} + +#[tokio::test] +async fn turn_finalization_failures_discard_stale_accounting_and_allow_the_next_turn() +-> anyhow::Result<()> { + for (case_index, (finalizer, failure)) in [ + (TurnFinalizer::Stop, InjectedFinalizationFailure::AuditClear), + (TurnFinalizer::Stop, InjectedFinalizationFailure::Accounting), + ( + TurnFinalizer::Abort, + InjectedFinalizationFailure::AuditClear, + ), + ( + TurnFinalizer::Abort, + InjectedFinalizationFailure::Accounting, + ), + ] + .into_iter() + .enumerate() + { + let tempdir = TempDir::new()?; + let runtime = codex_state::StateRuntime::init( + tempdir.path().to_path_buf(), + "test-provider".to_string(), + ) + .await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + let turn_id = format!("turn-finalization-failure-{case_index}"); + harness + .start_turn(turn_id.as_str(), &TokenUsage::default()) + .await; + + let tools = harness.tools(); + tool_by_name(&tools, "create_goal") + .handle(tool_call_for_turn( + turn_id.as_str(), + "create_goal", + format!("call-create-finalization-failure-{case_index}").as_str(), + json!({ "objective": format!("recover after {finalizer:?} {failure:?}") }), + )) + .await?; + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist for failure injection"))?; + + let fault_pool = SqlitePoolOptions::new() + .max_connections(1) + .connect_with( + SqliteConnectOptions::new() + .filename(codex_state::goals_db_path(runtime.codex_home())), + ) + .await?; + if matches!(failure, InjectedFinalizationFailure::AuditClear) { + sqlx::query( + r#" +INSERT INTO thread_goal_blocker_audits ( + thread_id, + goal_id, + fingerprint, + first_turn_id, + last_turn_id, + consecutive_turns, + created_at_ms, + updated_at_ms +) VALUES (?, ?, 'codex_err:injected_audit', ?, ?, 1, 1, 1) + "#, + ) + .bind(thread_id.to_string()) + .bind(goal.goal_id.as_str()) + .bind(turn_id.as_str()) + .bind(turn_id.as_str()) + .execute(&fault_pool) + .await?; + } + harness + .record_token_usage( + turn_id.as_str(), + &token_usage( + /*input_tokens*/ 8, /*cached_input_tokens*/ 0, + /*output_tokens*/ 3, /*reasoning_output_tokens*/ 0, + /*total_tokens*/ 11, + ), + ) + .await; + + let (trigger_sql, drop_trigger_sql) = match failure { + InjectedFinalizationFailure::AuditClear => ( + r#" +CREATE TRIGGER inject_goal_audit_clear_failure +BEFORE DELETE ON thread_goal_blocker_audits +BEGIN + SELECT RAISE(FAIL, 'injected goal audit clear failure'); +END + "#, + "DROP TRIGGER inject_goal_audit_clear_failure", + ), + InjectedFinalizationFailure::Accounting => ( + r#" +CREATE TRIGGER inject_goal_accounting_failure +BEFORE UPDATE OF time_used_seconds, tokens_used ON thread_goals +BEGIN + SELECT RAISE(FAIL, 'injected goal accounting failure'); +END + "#, + "DROP TRIGGER inject_goal_accounting_failure", + ), + }; + sqlx::query(trigger_sql).execute(&fault_pool).await?; + + match finalizer { + TurnFinalizer::Stop => harness.stop_turn(turn_id.as_str()).await, + TurnFinalizer::Abort => harness.abort_turn(turn_id.as_str()).await, + } + sqlx::query(drop_trigger_sql).execute(&fault_pool).await?; + + harness + .runtime_handle() + .usage_limit_active_goal_for_turn(turn_id.as_str()) + .await + .map_err(anyhow::Error::msg)?; + let after_stale_turn_probe = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should survive the stale turn probe"))?; + assert_eq!( + codex_state::ThreadGoalStatus::Active, + after_stale_turn_probe.status, + "{finalizer:?} must finish stale turn state even when {failure:?} fails" + ); + + let next_turn_id = format!("turn-after-finalization-failure-{case_index}"); + harness + .start_turn(next_turn_id.as_str(), &TokenUsage::default()) + .await; + harness + .record_token_usage( + next_turn_id.as_str(), + &token_usage( + /*input_tokens*/ 4, /*cached_input_tokens*/ 0, + /*output_tokens*/ 3, /*reasoning_output_tokens*/ 0, + /*total_tokens*/ 7, + ), + ) + .await; + match finalizer { + TurnFinalizer::Stop => harness.stop_turn(next_turn_id.as_str()).await, + TurnFinalizer::Abort => harness.abort_turn(next_turn_id.as_str()).await, + } + + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should accept a subsequent turn"))?; + let expected_tokens = match failure { + InjectedFinalizationFailure::AuditClear => 18, + InjectedFinalizationFailure::Accounting => 7, + }; + assert_eq!(codex_state::ThreadGoalStatus::Active, goal.status); + assert_eq!( + expected_tokens, goal.tokens_used, + "the next turn must account once without reviving failed-turn usage" + ); + let remaining_audits: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM thread_goal_blocker_audits") + .fetch_one(&fault_pool) + .await?; + assert_eq!( + 0, remaining_audits, + "the next successful finalization must clear any audit left by the injected failure" + ); + fault_pool.close().await; + } + Ok(()) +} + +#[tokio::test] +async fn different_turn_error_resets_consecutive_blocker_count() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + harness.start_turn("turn-1", &TokenUsage::default()).await; + + let tools = harness.tools(); + tool_by_name(&tools, "create_goal") + .handle(tool_call( + "create_goal", + "call-create-goal", + json!({ "objective": "ship goal extension backend" }), + )) + .await?; + + harness + .notify_turn_error("turn-1", CodexErrorInfo::ContextWindowExceeded) + .await; + harness.stop_turn("turn-1").await; + + for turn_id in ["turn-2", "turn-3"] { + harness.start_turn(turn_id, &TokenUsage::default()).await; + harness + .notify_turn_error(turn_id, CodexErrorInfo::Other) + .await; + harness.stop_turn(turn_id).await; + + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + assert_eq!( + codex_state::ThreadGoalStatus::Active, + goal.status, + "a changed blocker must restart the three-turn audit" + ); + } + + harness.start_turn("turn-4", &TokenUsage::default()).await; + harness + .notify_turn_error("turn-4", CodexErrorInfo::Other) + .await; + + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + assert_eq!(codex_state::ThreadGoalStatus::Blocked, goal.status); + Ok(()) +} + +#[tokio::test] +async fn broad_turn_error_variants_use_distinct_stable_fingerprints() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + harness.start_turn("turn-1", &TokenUsage::default()).await; + + let tools = harness.tools(); + tool_by_name(&tools, "create_goal") + .handle(tool_call( + "create_goal", + "call-create-broad-error-goal", + json!({ "objective": "keep broad turn failures separate" }), + )) + .await?; + + harness + .notify_turn_error_with_fingerprint( + "turn-1", + CodexErrorInfo::Other, + "codex_err:turn_aborted", + ) + .await; + harness.stop_turn("turn-1").await; + + for turn_id in ["turn-2", "turn-3"] { + harness.start_turn(turn_id, &TokenUsage::default()).await; + harness + .notify_turn_error_with_fingerprint( + turn_id, + CodexErrorInfo::Other, + "codex_err:request_timeout", + ) + .await; + harness.stop_turn(turn_id).await; + + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + assert_eq!( + codex_state::ThreadGoalStatus::Active, + goal.status, + "unrelated broad error variants must not share blocker turns" + ); + } + + harness.start_turn("turn-4", &TokenUsage::default()).await; + harness + .notify_turn_error_with_fingerprint( + "turn-4", + CodexErrorInfo::Other, + "codex_err:request_timeout", + ) + .await; + + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + assert_eq!(codex_state::ThreadGoalStatus::Blocked, goal.status); + Ok(()) +} + +#[tokio::test] +async fn successful_goal_turn_resets_consecutive_blocker_count() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + harness.start_turn("turn-1", &TokenUsage::default()).await; + + let tools = harness.tools(); + tool_by_name(&tools, "create_goal") + .handle(tool_call( + "create_goal", + "call-create-goal", + json!({ "objective": "ship goal extension backend" }), + )) + .await?; + + harness + .notify_turn_error("turn-1", CodexErrorInfo::Other) + .await; + harness.stop_turn("turn-1").await; + + harness.start_turn("turn-2", &TokenUsage::default()).await; + harness.stop_turn("turn-2").await; + + harness.start_turn("turn-3", &TokenUsage::default()).await; + harness + .notify_turn_error("turn-3", CodexErrorInfo::Other) + .await; + harness.stop_turn("turn-3").await; + + harness.start_turn("turn-4", &TokenUsage::default()).await; + harness + .notify_turn_error("turn-4", CodexErrorInfo::Other) + .await; + + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + assert_eq!( + codex_state::ThreadGoalStatus::Active, + goal.status, + "a successful intervening goal turn must reset the blocker audit" + ); + Ok(()) +} + +#[tokio::test] +async fn non_usage_turn_error_preserves_budget_limited_goal() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + harness.start_turn("turn-1", &TokenUsage::default()).await; + + let tools = harness.tools(); + tool_by_name(&tools, "create_goal") + .handle(tool_call( + "create_goal", + "call-create-goal", + json!({ + "objective": "ship goal extension backend", + "token_budget": 1, + }), + )) + .await?; + harness + .record_token_usage( + "turn-1", + &token_usage( + /*input_tokens*/ 2, /*cached_input_tokens*/ 0, /*output_tokens*/ 0, + /*reasoning_output_tokens*/ 0, /*total_tokens*/ 2, + ), + ) + .await; + harness + .notify_tool_finish("turn-1", "call-shell", "shell") + .await; + + harness + .notify_turn_error("turn-1", CodexErrorInfo::Other) + .await; + + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + assert_eq!(codex_state::ThreadGoalStatus::BudgetLimited, goal.status); Ok(()) } @@ -3503,6 +4154,20 @@ impl GoalExtensionHarness { } } + async fn abort_turn(&self, turn_id: &str) { + let turn_store = ExtensionData::new(turn_id); + for contributor in self.registry.turn_lifecycle_contributors() { + contributor + .on_turn_abort(TurnAbortInput { + reason: TurnAbortReason::Interrupted, + session_store: &self.session_store, + thread_store: &self.thread_store, + turn_store: &turn_store, + }) + .await; + } + } + async fn record_token_usage(&self, turn_id: &str, usage: &TokenUsage) { let turn_store = ExtensionData::new(turn_id); let token_usage = TokenUsageInfo { @@ -3564,16 +4229,30 @@ impl GoalExtensionHarness { } async fn notify_turn_error(&self, turn_id: &str, error: CodexErrorInfo) { + let error_fingerprint = test_error_fingerprint(&error); + self.notify_turn_error_with_fingerprint(turn_id, error, error_fingerprint.as_str()) + .await; + } + + async fn notify_turn_error_with_fingerprint( + &self, + turn_id: &str, + error: CodexErrorInfo, + error_fingerprint: &str, + ) { let turn_store = ExtensionData::new(turn_id); for contributor in self.registry.turn_lifecycle_contributors() { contributor - .on_turn_error(TurnErrorInput { - turn_id, - error: error.clone(), - session_store: &self.session_store, - thread_store: &self.thread_store, - turn_store: &turn_store, - }) + .on_turn_error_with_fingerprint( + TurnErrorInput { + turn_id, + error: error.clone(), + session_store: &self.session_store, + thread_store: &self.thread_store, + turn_store: &turn_store, + }, + error_fingerprint, + ) .await; } } @@ -3585,6 +4264,29 @@ impl GoalExtensionHarness { } } +fn test_error_fingerprint(error: &CodexErrorInfo) -> String { + let kind = match error { + CodexErrorInfo::ContextWindowExceeded => "context_window_exceeded", + CodexErrorInfo::UsageLimitExceeded => "usage_limit_exceeded", + CodexErrorInfo::ServerOverloaded => "server_overloaded", + CodexErrorInfo::CyberPolicy => "cyber_policy", + CodexErrorInfo::HttpConnectionFailed { .. } => "http_connection_failed", + CodexErrorInfo::ResponseStreamConnectionFailed { .. } => { + "response_stream_connection_failed" + } + CodexErrorInfo::InternalServerError => "internal_server_error", + CodexErrorInfo::Unauthorized => "unauthorized", + CodexErrorInfo::BadRequest => "bad_request", + CodexErrorInfo::SandboxError => "sandbox_error", + CodexErrorInfo::ResponseStreamDisconnected { .. } => "response_stream_disconnected", + CodexErrorInfo::ResponseTooManyFailedAttempts { .. } => "response_too_many_failed_attempts", + CodexErrorInfo::ActiveTurnNotSteerable { .. } => "active_turn_not_steerable", + CodexErrorInfo::ThreadRollbackFailed => "thread_rollback_failed", + CodexErrorInfo::Other => "other", + }; + format!("codex_err:test_{kind}") +} + fn tool_by_name<'a>( tools: &'a [Arc>], name: &str, @@ -3596,8 +4298,17 @@ fn tool_by_name<'a>( } fn tool_call(tool_name: &str, call_id: &str, arguments: serde_json::Value) -> ToolCall { + tool_call_for_turn("turn-1", tool_name, call_id, arguments) +} + +fn tool_call_for_turn( + turn_id: &str, + tool_name: &str, + call_id: &str, + arguments: serde_json::Value, +) -> ToolCall { ToolCall { - turn_id: "turn-1".to_string(), + turn_id: turn_id.to_string(), call_id: call_id.to_string(), tool_name: codex_extension_api::ToolName::plain(tool_name), model: "gpt-test".to_string(), @@ -3610,6 +4321,20 @@ fn tool_call(tool_name: &str, call_id: &str, arguments: serde_json::Value) -> To } } +async fn resume_goal_in_turn(harness: &GoalExtensionHarness, turn_id: &str) -> anyhow::Result<()> { + harness.start_turn(turn_id, &TokenUsage::default()).await; + let tools = harness.tools(); + tool_by_name(&tools, "resume_goal") + .handle(tool_call_for_turn( + turn_id, + "resume_goal", + &format!("call-resume-{turn_id}"), + json!({}), + )) + .await?; + Ok(()) +} + async fn test_runtime() -> anyhow::Result> { let tempdir = TempDir::new()?; codex_state::StateRuntime::init(tempdir.keep(), "test-provider".to_string()).await diff --git a/codex-rs/state/goals_migrations/0009_thread_goal_blocker_audits.sql b/codex-rs/state/goals_migrations/0009_thread_goal_blocker_audits.sql new file mode 100644 index 0000000000..d98376ac68 --- /dev/null +++ b/codex-rs/state/goals_migrations/0009_thread_goal_blocker_audits.sql @@ -0,0 +1,55 @@ +CREATE TABLE thread_goal_blocker_audits ( + thread_id TEXT PRIMARY KEY NOT NULL, + goal_id TEXT NOT NULL, + fingerprint TEXT NOT NULL CHECK( + length(fingerprint) > 0 + AND length(fingerprint) <= 256 + ), + -- Host-owned goal turn identifiers for the current fingerprint streak. + -- These are not upstream transport request identifiers. The child table + -- below durably holds every distinct id counted in this bounded streak. + first_turn_id TEXT NOT NULL CHECK( + length(first_turn_id) > 0 + AND length(first_turn_id) <= 256 + ), + last_turn_id TEXT NOT NULL CHECK( + length(last_turn_id) > 0 + AND length(last_turn_id) <= 256 + ), + consecutive_turns INTEGER NOT NULL CHECK( + consecutive_turns >= 1 + AND consecutive_turns <= 3 + ), + created_at_ms INTEGER NOT NULL, + updated_at_ms INTEGER NOT NULL, + FOREIGN KEY(thread_id) REFERENCES thread_goals(thread_id) ON DELETE CASCADE +); + +CREATE TABLE thread_goal_blocker_audit_turns ( + thread_id TEXT NOT NULL, + goal_id TEXT NOT NULL, + turn_id TEXT NOT NULL CHECK( + length(turn_id) > 0 + AND length(turn_id) <= 256 + ), + created_at_ms INTEGER NOT NULL, + PRIMARY KEY(thread_id, goal_id, turn_id), + FOREIGN KEY(thread_id) REFERENCES thread_goal_blocker_audits(thread_id) ON DELETE CASCADE +); + +CREATE TRIGGER clear_thread_goal_blocker_audit_after_lifecycle_change +AFTER UPDATE OF goal_id, objective, title, status, token_budget ON thread_goals +WHEN OLD.goal_id IS NOT NEW.goal_id + OR OLD.objective IS NOT NEW.objective + OR OLD.title IS NOT NEW.title + OR OLD.token_budget IS NOT NEW.token_budget + OR ( + OLD.status IS NOT NEW.status + AND NOT ( + OLD.status = 'paused' AND NEW.status = 'active' + ) + ) +BEGIN + DELETE FROM thread_goal_blocker_audits + WHERE thread_id = NEW.thread_id; +END; diff --git a/codex-rs/state/migrations/0061_thread_schedule_run_goal_id.sql b/codex-rs/state/migrations/0061_thread_schedule_run_goal_id.sql new file mode 100644 index 0000000000..68753e0ab5 --- /dev/null +++ b/codex-rs/state/migrations/0061_thread_schedule_run_goal_id.sql @@ -0,0 +1,6 @@ +ALTER TABLE thread_schedule_runs +ADD COLUMN goal_id TEXT; + +CREATE INDEX idx_thread_schedule_runs_running_turn + ON thread_schedule_runs(thread_id, turn_id) + WHERE status = 'running' AND turn_id IS NOT NULL; diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index d5bd4d0de9..33aaa3cfc6 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -168,6 +168,7 @@ pub use runtime::DEFAULT_THREAD_WORKFLOW_RUN_LIST_LIMIT; pub use runtime::DEFAULT_WEBHOOK_EVENT_LIST_LIMIT; pub use runtime::GoalAccountingMode; pub use runtime::GoalAccountingOutcome; +pub use runtime::GoalBlockerAuditOutcome; pub use runtime::GoalDeleteOutcome; pub use runtime::GoalStore; pub use runtime::GoalUpdate; @@ -233,6 +234,9 @@ pub use runtime::ThreadScheduleClaim; pub use runtime::ThreadScheduleCreateParams; pub use runtime::ThreadScheduleDueClaimParams; pub use runtime::ThreadScheduleNowClaimParams; +pub use runtime::ThreadScheduleRunForGoalFinishParams; +pub use runtime::ThreadScheduleRunLeaseParams; +pub use runtime::ThreadScheduleRunStartParams; pub use runtime::ThreadScheduleUpdate; pub use runtime::WEBHOOK_EVENT_DEDUPE_CONFLICT_MESSAGE; pub use runtime::WORKFLOW_STEP_APPROVAL_APPROVED; diff --git a/codex-rs/state/src/model/thread_schedule.rs b/codex-rs/state/src/model/thread_schedule.rs index 514038bb32..422c83dded 100644 --- a/codex-rs/state/src/model/thread_schedule.rs +++ b/codex-rs/state/src/model/thread_schedule.rs @@ -177,6 +177,7 @@ pub struct ThreadScheduleRun { pub status: ThreadScheduleRunStatus, pub lease_id: String, pub turn_id: Option, + pub goal_id: Option, pub error: Option, pub scheduled_for: Option>, pub started_at: DateTime, @@ -293,6 +294,7 @@ pub(crate) struct ThreadScheduleRunRow { pub status: String, pub lease_id: String, pub turn_id: Option, + pub goal_id: Option, pub error: Option, pub scheduled_for_ms: Option, pub started_at_ms: i64, @@ -308,6 +310,7 @@ impl ThreadScheduleRunRow { status: row.try_get("status")?, lease_id: row.try_get("lease_id")?, turn_id: row.try_get("turn_id")?, + goal_id: row.try_get("goal_id")?, error: row.try_get("error")?, scheduled_for_ms: row.try_get("scheduled_for_ms")?, started_at_ms: row.try_get("started_at_ms")?, @@ -327,6 +330,7 @@ impl TryFrom for ThreadScheduleRun { status: ThreadScheduleRunStatus::try_from(row.status.as_str())?, lease_id: row.lease_id, turn_id: row.turn_id, + goal_id: row.goal_id, error: row.error, scheduled_for: optional_epoch_millis_to_datetime(row.scheduled_for_ms)?, started_at: epoch_millis_to_datetime(row.started_at_ms)?, diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index 2979e6a565..91dd9605ff 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -181,6 +181,7 @@ pub use goal_plans::ThreadGoalPlanListPage; pub use goal_plans::ThreadGoalPlanNodeCreateParams; pub use goals::GoalAccountingMode; pub use goals::GoalAccountingOutcome; +pub use goals::GoalBlockerAuditOutcome; pub use goals::GoalDeleteOutcome; pub use goals::GoalStore; pub use goals::GoalUpdate; @@ -237,6 +238,9 @@ pub use schedules::ThreadScheduleClaim; pub use schedules::ThreadScheduleCreateParams; pub use schedules::ThreadScheduleDueClaimParams; pub use schedules::ThreadScheduleNowClaimParams; +pub use schedules::ThreadScheduleRunForGoalFinishParams; +pub use schedules::ThreadScheduleRunLeaseParams; +pub use schedules::ThreadScheduleRunStartParams; pub use schedules::ThreadScheduleUpdate; pub use threads::ThreadFilterOptions; pub use webhooks::DEFAULT_WEBHOOK_EVENT_LIST_LIMIT; @@ -597,7 +601,7 @@ impl StateRuntime { let thread_updated_at_millis = thread_updated_at_millis.unwrap_or(0); let runtime = Arc::new(Self { thread_goals: GoalStore::new(Arc::clone(&goals_pool)), - thread_schedules: ScheduleStore::new(Arc::clone(&pool)), + thread_schedules: ScheduleStore::new(Arc::clone(&pool), Arc::clone(&goals_pool)), thread_monitors: MonitorStore::new(Arc::clone(&pool)), local_active_sessions: LocalActiveSessionStore::new(Arc::clone(&pool)), webhook_events: WebhookEventStore::new(Arc::clone(&pool)), @@ -1296,6 +1300,119 @@ mod tests { } } + #[tokio::test] + async fn thread_schedule_run_goal_migration_preserves_legacy_running_runs() { + let codex_home = unique_temp_dir(); + tokio::fs::create_dir_all(&codex_home) + .await + .expect("create codex home"); + let state_path = state_db_path(codex_home.as_path()); + let pool = SqlitePool::connect_with( + SqliteConnectOptions::new() + .filename(&state_path) + .create_if_missing(true), + ) + .await + .expect("open state db"); + migrator_through(&STATE_MIGRATOR, /*version*/ 60) + .run(&pool) + .await + .expect("apply pre-goal-correlation state schema"); + sqlx::query( + r#" +INSERT INTO threads ( + id, rollout_path, created_at, updated_at, source, model_provider, cwd, title, sandbox_policy, approval_mode +) VALUES ('00000000-0000-0000-0000-000000000001', '', 0, 0, 'cli', 'test-provider', '/', 'fixture', 'workspace-write', 'on-request') + "#, + ) + .execute(&pool) + .await + .expect("insert legacy schedule thread"); + sqlx::query( + r#" +INSERT INTO thread_schedules ( + schedule_id, + thread_id, + prompt_source, + prompt, + schedule_kind, + interval_amount, + interval_unit, + timezone, + status, + next_run_at_ms, + failure_count, + lease_id, + lease_expires_at_ms, + created_at_ms, + updated_at_ms +) VALUES ( + 'schedule-1', + '00000000-0000-0000-0000-000000000001', + 'inline', + '/goal fixture', + 'interval', + 1, + 'minutes', + 'UTC', + 'active', + 60000, + 0, + 'lease-1', + 120000, + 0, + 0 +) + "#, + ) + .execute(&pool) + .await + .expect("insert legacy schedule"); + sqlx::query( + r#" +INSERT INTO thread_schedule_runs ( + run_id, + schedule_id, + thread_id, + status, + lease_id, + turn_id, + scheduled_for_ms, + started_at_ms +) VALUES ( + 'run-1', + 'schedule-1', + '00000000-0000-0000-0000-000000000001', + 'running', + 'lease-1', + 'turn-1', + 60000, + 60000 +) + "#, + ) + .execute(&pool) + .await + .expect("insert legacy running schedule run"); + pool.close().await; + + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("current state schema should migrate legacy schedule runs"); + let run = runtime + .thread_schedules() + .get_thread_schedule_run("run-1") + .await + .expect("read migrated schedule run") + .expect("legacy schedule run should remain"); + assert_eq!(run.status, crate::ThreadScheduleRunStatus::Running); + assert_eq!(run.turn_id.as_deref(), Some("turn-1")); + assert_eq!(run.goal_id, None); + + drop(runtime); + let _ = tokio::fs::remove_dir_all(codex_home).await; + } + #[tokio::test] async fn pending_interaction_event_sequence_migration_survives_vacuum() { let codex_home = unique_temp_dir(); @@ -2198,8 +2315,8 @@ INSERT INTO thread_goal_plan_nodes ( assert!(matches!(unrepaired_err, MigrateError::VersionMismatch(5))); strict_pool.close().await; - // Full runtime init repairs the stamp and applies the missing - // versions 5-7 on top of the legacy schema. + // Full runtime init repairs the stamp and applies the remaining + // migrations on top of the legacy schema. let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await .expect("state runtime should initialize on a 0.1.48-stamped goals db"); @@ -2214,7 +2331,7 @@ INSERT INTO thread_goal_plan_nodes ( .await .expect("repaired stamps should query"); assert_eq!( - (1..=8).collect::>(), + (1..=9).collect::>(), stamped .iter() .map(|(version, _)| *version) @@ -2228,7 +2345,11 @@ INSERT INTO thread_goal_plan_nodes ( .to_vec(); assert_eq!( embedded_deferred_checksum, - stamped.last().expect("version 8 stamp").1, + stamped + .iter() + .find(|(version, _)| *version == 8) + .expect("version 8 stamp") + .1, "version 8 must carry the embedded deferred checksum after repair" ); // The repaired database converges on the fresh schema: assignment @@ -2254,6 +2375,18 @@ JOIN thread_goal_plan_nodes node ON node.thread_id = goal.thread_id .await .expect("goal context lifecycle table should exist"); assert_eq!(0, lifecycle_count); + let blocker_audit_count: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM thread_goal_blocker_audits") + .fetch_one(&query_pool) + .await + .expect("goal blocker audit table should exist"); + assert_eq!(0, blocker_audit_count); + let blocker_audit_turn_count: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM thread_goal_blocker_audit_turns") + .fetch_one(&query_pool) + .await + .expect("goal blocker audit turn table should exist"); + assert_eq!(0, blocker_audit_turn_count); let indexes: Vec = sqlx::query_scalar( "SELECT name FROM sqlite_master WHERE type = 'index' AND tbl_name = 'thread_goal_plan_nodes' AND name LIKE 'idx_%' ORDER BY name", ) diff --git a/codex-rs/state/src/runtime/goal_plans.rs b/codex-rs/state/src/runtime/goal_plans.rs index 0066363ae1..9d079b37e3 100644 --- a/codex-rs/state/src/runtime/goal_plans.rs +++ b/codex-rs/state/src/runtime/goal_plans.rs @@ -1485,6 +1485,15 @@ WHERE node_id = ? /*tokens_used*/ 0, effective_token_budget, ); + sqlx::query( + r#" +DELETE FROM thread_goal_blocker_audits +WHERE thread_id = ? + "#, + ) + .bind(thread_id.to_string()) + .execute(&mut **tx) + .await?; let goal_row = sqlx::query( r#" INSERT INTO thread_goals ( diff --git a/codex-rs/state/src/runtime/goals.rs b/codex-rs/state/src/runtime/goals.rs index fe78eea295..d26e5c8220 100644 --- a/codex-rs/state/src/runtime/goals.rs +++ b/codex-rs/state/src/runtime/goals.rs @@ -28,6 +28,15 @@ pub enum GoalAccountingOutcome { Updated(crate::ThreadGoal), } +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum GoalBlockerAuditOutcome { + Unchanged(Option), + Updated { + goal: crate::ThreadGoal, + consecutive_turns: u8, + }, +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct GoalDeleteOutcome { pub deleted: bool, @@ -125,6 +134,15 @@ WHERE thread_id = ? block_projected_goal_plan_nodes_in_tx(&mut tx, thread_id, previous_goal_id, now_ms) .await?; } + sqlx::query( + r#" +DELETE FROM thread_goal_blocker_audits +WHERE thread_id = ? + "#, + ) + .bind(thread_id.to_string()) + .execute(&mut *tx) + .await?; let row = sqlx::query( r#" INSERT INTO thread_goals ( @@ -595,6 +613,362 @@ WHERE thread_id = ? .await } + pub async fn observe_active_thread_goal_blocker( + &self, + thread_id: ThreadId, + expected_goal_id: &str, + turn_id: &str, + fingerprint: &str, + required_consecutive_turns: u8, + ) -> anyhow::Result { + validate_goal_blocker_fingerprint(fingerprint)?; + validate_goal_blocker_turn_id(turn_id)?; + if required_consecutive_turns == 0 || required_consecutive_turns > 3 { + anyhow::bail!("required consecutive blocker turns must be between 1 and 3"); + } + + crate::busy_retry::retry_on_busy("observe active thread goal blocker", || { + self.observe_active_thread_goal_blocker_once( + thread_id, + expected_goal_id, + turn_id, + fingerprint, + required_consecutive_turns, + ) + }) + .await + } + + async fn observe_active_thread_goal_blocker_once( + &self, + thread_id: ThreadId, + expected_goal_id: &str, + turn_id: &str, + fingerprint: &str, + required_consecutive_turns: u8, + ) -> anyhow::Result { + let thread_id_string = thread_id.to_string(); + let now_ms = datetime_to_epoch_millis(Utc::now()); + // This is a read-then-write operation shared by app-server and + // reconstructed thread runtimes. Take the writer lock before the read + // so WAL advancement cannot turn the upgrade into BUSY_SNAPSHOT. + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let active_goal_row = sqlx::query( + r#" +SELECT + thread_id, + goal_id, + objective, + title, + status, + token_budget, + tokens_used, + time_used_seconds, + created_at_ms, + updated_at_ms +FROM thread_goals +WHERE thread_id = ? + AND goal_id = ? + "#, + ) + .bind(thread_id_string.as_str()) + .bind(expected_goal_id) + .fetch_optional(&mut *tx) + .await?; + let Some(active_goal_row) = active_goal_row else { + tx.rollback().await?; + return Ok(GoalBlockerAuditOutcome::Unchanged( + self.get_thread_goal(thread_id).await?, + )); + }; + let active_goal = thread_goal_from_row(&active_goal_row)?; + if active_goal.status != crate::ThreadGoalStatus::Active { + tx.rollback().await?; + return Ok(GoalBlockerAuditOutcome::Unchanged(Some(active_goal))); + } + + let mut previous_audit: Option<(String, String, String, String, i64, i64)> = + sqlx::query_as( + r#" +SELECT + goal_id, + fingerprint, + first_turn_id, + last_turn_id, + consecutive_turns, + created_at_ms +FROM thread_goal_blocker_audits +WHERE thread_id = ? + "#, + ) + .bind(thread_id_string.as_str()) + .fetch_optional(&mut *tx) + .await?; + + if previous_audit + .as_ref() + .is_some_and(|(goal_id, ..)| goal_id != expected_goal_id) + { + sqlx::query( + r#" +DELETE FROM thread_goal_blocker_audits +WHERE thread_id = ? + "#, + ) + .bind(thread_id_string.as_str()) + .execute(&mut *tx) + .await?; + previous_audit = None; + } + let observed_turns: Vec<(String, i64)> = sqlx::query_as( + r#" +SELECT turn_id, created_at_ms +FROM thread_goal_blocker_audit_turns +WHERE thread_id = ? + AND goal_id = ? + "#, + ) + .bind(thread_id_string.as_str()) + .bind(expected_goal_id) + .fetch_all(&mut *tx) + .await?; + let preserves_current_streak = + previous_audit + .as_ref() + .is_some_and(|(goal_id, previous_fingerprint, ..)| { + goal_id == expected_goal_id && previous_fingerprint == fingerprint + }); + let turn_already_observed = preserves_current_streak + && observed_turns + .iter() + .any(|(observed_turn_id, _)| observed_turn_id == turn_id); + if !preserves_current_streak { + sqlx::query( + r#" +DELETE FROM thread_goal_blocker_audit_turns +WHERE thread_id = ? + AND goal_id = ? + "#, + ) + .bind(thread_id_string.as_str()) + .bind(expected_goal_id) + .execute(&mut *tx) + .await?; + } + let ( + audit_fingerprint, + first_turn_id, + last_turn_id, + consecutive_turns, + audit_created_at_ms, + ) = match previous_audit { + Some(( + goal_id, + previous_fingerprint, + first_turn_id, + last_turn_id, + consecutive_turns, + created_at_ms, + )) if goal_id == expected_goal_id && turn_already_observed => ( + previous_fingerprint, + first_turn_id, + last_turn_id, + consecutive_turns, + created_at_ms, + ), + Some(( + goal_id, + previous_fingerprint, + first_turn_id, + _last_turn_id, + consecutive_turns, + created_at_ms, + )) if goal_id == expected_goal_id && previous_fingerprint == fingerprint => ( + previous_fingerprint, + first_turn_id, + turn_id.to_string(), + consecutive_turns + .saturating_add(1) + .min(i64::from(required_consecutive_turns)), + created_at_ms, + ), + Some(_) | None => ( + fingerprint.to_string(), + turn_id.to_string(), + turn_id.to_string(), + 1, + now_ms, + ), + }; + let consecutive_turns = u8::try_from(consecutive_turns)?; + let status = if consecutive_turns >= required_consecutive_turns { + crate::ThreadGoalStatus::Blocked + } else { + crate::ThreadGoalStatus::Active + }; + let goal_row = sqlx::query( + r#" +UPDATE thread_goals +SET + status = ?, + updated_at_ms = ? +WHERE thread_id = ? + AND goal_id = ? + AND status = 'active' +RETURNING + thread_id, + goal_id, + objective, + title, + status, + token_budget, + tokens_used, + time_used_seconds, + created_at_ms, + updated_at_ms + "#, + ) + .bind(status.as_str()) + .bind(now_ms) + .bind(thread_id_string.as_str()) + .bind(expected_goal_id) + .fetch_optional(&mut *tx) + .await?; + let Some(goal_row) = goal_row else { + tx.rollback().await?; + return Ok(GoalBlockerAuditOutcome::Unchanged( + self.get_thread_goal(thread_id).await?, + )); + }; + let goal = thread_goal_from_row(&goal_row)?; + if goal.status == crate::ThreadGoalStatus::Active { + sqlx::query( + r#" +INSERT INTO thread_goal_blocker_audits ( + thread_id, + goal_id, + fingerprint, + first_turn_id, + last_turn_id, + consecutive_turns, + created_at_ms, + updated_at_ms +) VALUES (?, ?, ?, ?, ?, ?, ?, ?) +ON CONFLICT(thread_id) DO UPDATE SET + goal_id = excluded.goal_id, + fingerprint = excluded.fingerprint, + first_turn_id = excluded.first_turn_id, + last_turn_id = excluded.last_turn_id, + consecutive_turns = excluded.consecutive_turns, + created_at_ms = excluded.created_at_ms, + updated_at_ms = excluded.updated_at_ms + "#, + ) + .bind(thread_id_string.as_str()) + .bind(expected_goal_id) + .bind(audit_fingerprint) + .bind(first_turn_id) + .bind(last_turn_id) + .bind(i64::from(consecutive_turns)) + .bind(audit_created_at_ms) + .bind(now_ms) + .execute(&mut *tx) + .await?; + if preserves_current_streak { + for (observed_turn_id, observed_at_ms) in &observed_turns { + sqlx::query( + r#" +INSERT INTO thread_goal_blocker_audit_turns ( + thread_id, + goal_id, + turn_id, + created_at_ms +) VALUES (?, ?, ?, ?) +ON CONFLICT(thread_id, goal_id, turn_id) DO NOTHING + "#, + ) + .bind(thread_id_string.as_str()) + .bind(expected_goal_id) + .bind(observed_turn_id) + .bind(observed_at_ms) + .execute(&mut *tx) + .await?; + } + } + if !turn_already_observed { + sqlx::query( + r#" +INSERT INTO thread_goal_blocker_audit_turns ( + thread_id, + goal_id, + turn_id, + created_at_ms +) VALUES (?, ?, ?, ?) + "#, + ) + .bind(thread_id_string.as_str()) + .bind(expected_goal_id) + .bind(turn_id) + .bind(now_ms) + .execute(&mut *tx) + .await?; + } + } else { + sqlx::query( + r#" +DELETE FROM thread_goal_blocker_audits +WHERE thread_id = ? + AND goal_id = ? + "#, + ) + .bind(thread_id_string.as_str()) + .bind(expected_goal_id) + .execute(&mut *tx) + .await?; + } + tx.commit().await?; + + Ok(GoalBlockerAuditOutcome::Updated { + goal, + consecutive_turns, + }) + } + + pub async fn clear_thread_goal_blocker_audit( + &self, + thread_id: ThreadId, + ) -> anyhow::Result { + let result = sqlx::query( + r#" +DELETE FROM thread_goal_blocker_audits +WHERE thread_id = ? + "#, + ) + .bind(thread_id.to_string()) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn clear_thread_goal_blocker_audit_for_goal( + &self, + thread_id: ThreadId, + goal_id: &str, + ) -> anyhow::Result { + let result = sqlx::query( + r#" +DELETE FROM thread_goal_blocker_audits +WHERE thread_id = ? + AND goal_id = ? + "#, + ) + .bind(thread_id.to_string()) + .bind(goal_id) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + pub async fn usage_limit_active_thread_goal( &self, thread_id: ThreadId, @@ -799,6 +1173,25 @@ RETURNING } } +fn validate_goal_blocker_fingerprint(fingerprint: &str) -> anyhow::Result<()> { + let valid = fingerprint.starts_with("codex_err:") + && fingerprint.len() <= 256 + && fingerprint.bytes().all(|byte| { + byte.is_ascii_lowercase() || byte.is_ascii_digit() || b"_:".contains(&byte) + }); + if !valid { + anyhow::bail!("goal blocker fingerprint must be a canonical host error classification"); + } + Ok(()) +} + +fn validate_goal_blocker_turn_id(turn_id: &str) -> anyhow::Result<()> { + if turn_id.is_empty() || turn_id.len() > 256 { + anyhow::bail!("goal blocker turn id must be between 1 and 256 bytes"); + } + Ok(()) +} + fn thread_goal_from_row(row: &sqlx::sqlite::SqliteRow) -> anyhow::Result { ThreadGoalRow::try_from_row(row).and_then(crate::ThreadGoal::try_from) } @@ -884,6 +1277,536 @@ mod tests { .expect("test thread should be upserted"); } + async fn blocker_audit( + runtime: &StateRuntime, + thread_id: ThreadId, + ) -> Option<(String, String, String, String, i64)> { + sqlx::query_as( + r#" +SELECT goal_id, fingerprint, first_turn_id, last_turn_id, consecutive_turns +FROM thread_goal_blocker_audits +WHERE thread_id = ? + "#, + ) + .bind(thread_id.to_string()) + .fetch_optional(runtime.thread_goals().pool.as_ref()) + .await + .expect("blocker audit should load") + } + + async fn blocker_audit_turn_ids(runtime: &StateRuntime, thread_id: ThreadId) -> Vec { + sqlx::query_scalar( + r#" +SELECT turn_id +FROM thread_goal_blocker_audit_turns +WHERE thread_id = ? +ORDER BY turn_id + "#, + ) + .bind(thread_id.to_string()) + .fetch_all(runtime.thread_goals().pool.as_ref()) + .await + .expect("blocker audit turn ids should load") + } + + #[tokio::test] + async fn blocker_audit_keeps_goal_active_but_manual_pause_resets_it() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(); + upsert_test_thread(&runtime, thread_id).await; + let goal = runtime + .thread_goals() + .replace_thread_goal( + thread_id, + "audit restart lifecycle", + crate::ThreadGoalStatus::Active, + /*token_budget*/ None, + ) + .await + .expect("goal should be created"); + + let first = runtime + .thread_goals() + .observe_active_thread_goal_blocker( + thread_id, + goal.goal_id.as_str(), + "turn-1", + "codex_err:test_blocker", + 3, + ) + .await + .expect("first blocker should persist"); + assert!(matches!( + first, + GoalBlockerAuditOutcome::Updated { + consecutive_turns: 1, + .. + } + )); + + assert_eq!( + crate::ThreadGoalStatus::Active, + runtime + .thread_goals() + .get_thread_goal(thread_id) + .await + .expect("goal should load") + .expect("goal should exist") + .status, + "a pre-threshold blocker observation must not create a resume gate" + ); + assert_eq!( + Some(( + goal.goal_id.clone(), + "codex_err:test_blocker".to_string(), + "turn-1".to_string(), + "turn-1".to_string(), + 1, + )), + blocker_audit(runtime.as_ref(), thread_id).await, + "the active pre-threshold lifecycle preserves the audit" + ); + + runtime + .thread_goals() + .update_thread_goal( + thread_id, + GoalUpdate { + objective: None, + title: None, + status: Some(crate::ThreadGoalStatus::Paused), + token_budget: None, + expected_goal_id: Some(goal.goal_id.clone()), + }, + ) + .await + .expect("goal should pause manually"); + assert_eq!( + None, + blocker_audit(runtime.as_ref(), thread_id).await, + "an independent manual pause starts a fresh blocker audit" + ); + + runtime + .thread_goals() + .update_thread_goal( + thread_id, + GoalUpdate { + objective: None, + title: None, + status: Some(crate::ThreadGoalStatus::Active), + token_budget: None, + expected_goal_id: Some(goal.goal_id.clone()), + }, + ) + .await + .expect("goal should resume again"); + let after_manual_pause = runtime + .thread_goals() + .observe_active_thread_goal_blocker( + thread_id, + goal.goal_id.as_str(), + "turn-2", + "codex_err:test_blocker", + 3, + ) + .await + .expect("blocker should restart after manual pause"); + assert!(matches!( + after_manual_pause, + GoalBlockerAuditOutcome::Updated { + consecutive_turns: 1, + .. + } + )); + } + + #[tokio::test] + async fn blocker_audit_retries_writer_contention_without_double_counting() { + 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(); + upsert_test_thread(runtime.as_ref(), thread_id).await; + let goal = runtime + .thread_goals() + .replace_thread_goal( + thread_id, + "retry contended blocker observation", + crate::ThreadGoalStatus::Active, + /*token_budget*/ None, + ) + .await + .expect("goal should be created"); + + let contender_pool = SqlitePoolOptions::new() + .max_connections(1) + .connect_with( + SqliteConnectOptions::new() + .filename(codex_home.join(crate::GOALS_DB_FILENAME)) + .journal_mode(SqliteJournalMode::Wal) + .busy_timeout(Duration::from_millis(1)), + ) + .await + .expect("contending goals pool should open"); + let contender = GoalStore::new(Arc::new(contender_pool)); + + let mut writer = runtime + .thread_goals() + .pool + .acquire() + .await + .expect("writer connection should acquire"); + sqlx::query("BEGIN IMMEDIATE") + .execute(&mut *writer) + .await + .expect("writer lock should acquire"); + + let goal_id = goal.goal_id.clone(); + let mut observation = Box::pin(contender.observe_active_thread_goal_blocker( + thread_id, + goal_id.as_str(), + "turn-contended", + "codex_err:test_blocker", + 3, + )); + assert!( + tokio::time::timeout(Duration::from_millis(75), observation.as_mut()) + .await + .is_err(), + "the path-level observation must keep retrying while the second connection holds the writer lock" + ); + sqlx::query("ROLLBACK") + .execute(&mut *writer) + .await + .expect("writer lock should release"); + drop(writer); + + let outcome = tokio::time::timeout(Duration::from_secs(2), observation) + .await + .expect("observation should complete after lock release") + .expect("observation should retry successfully"); + assert!(matches!( + outcome, + GoalBlockerAuditOutcome::Updated { + consecutive_turns: 1, + .. + } + )); + assert_eq!( + Some(( + goal.goal_id, + "codex_err:test_blocker".to_string(), + "turn-contended".to_string(), + "turn-contended".to_string(), + 1, + )), + blocker_audit(runtime.as_ref(), thread_id).await, + "a retried transaction must commit the observation exactly once" + ); + } + + #[tokio::test] + async fn blocker_observations_keep_goal_active_until_the_third_distinct_turn() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(); + upsert_test_thread(&runtime, thread_id).await; + let goal = runtime + .thread_goals() + .replace_thread_goal( + thread_id, + "stay active during the blocker audit", + crate::ThreadGoalStatus::Active, + /*token_budget*/ None, + ) + .await + .expect("goal should be created"); + + for (turn_id, expected_consecutive_turns) in [("turn-1", 1), ("turn-2", 2)] { + let outcome = runtime + .thread_goals() + .observe_active_thread_goal_blocker( + thread_id, + goal.goal_id.as_str(), + turn_id, + "codex_err:test_blocker", + 3, + ) + .await + .expect("pre-threshold blocker observation should persist"); + assert!(matches!( + outcome, + GoalBlockerAuditOutcome::Updated { + goal: crate::ThreadGoal { + status: crate::ThreadGoalStatus::Active, + .. + }, + consecutive_turns, + } if consecutive_turns == expected_consecutive_turns + )); + } + + let third = runtime + .thread_goals() + .observe_active_thread_goal_blocker( + thread_id, + goal.goal_id.as_str(), + "turn-3", + "codex_err:test_blocker", + 3, + ) + .await + .expect("third blocker observation should persist"); + assert!(matches!( + third, + GoalBlockerAuditOutcome::Updated { + goal: crate::ThreadGoal { + status: crate::ThreadGoalStatus::Blocked, + .. + }, + consecutive_turns: 3, + } + )); + } + + #[tokio::test] + async fn blocker_audit_persists_only_the_bounded_current_streak_across_restarts() { + let codex_home = unique_temp_dir(); + let thread_id = test_thread_id(); + let goal_id = { + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("state db should initialize"); + upsert_test_thread(runtime.as_ref(), thread_id).await; + let goal = runtime + .thread_goals() + .replace_thread_goal( + thread_id, + "deduplicate the bounded current blocker streak", + crate::ThreadGoalStatus::Active, + /*token_budget*/ None, + ) + .await + .expect("goal should be created"); + + runtime + .thread_goals() + .observe_active_thread_goal_blocker( + thread_id, + goal.goal_id.as_str(), + "turn-1", + "codex_err:blocker_a", + 3, + ) + .await + .expect("first blocker observation should persist"); + let second = runtime + .thread_goals() + .observe_active_thread_goal_blocker( + thread_id, + goal.goal_id.as_str(), + "turn-2", + "codex_err:blocker_a", + 3, + ) + .await + .expect("second blocker observation should persist"); + assert!(matches!( + second, + GoalBlockerAuditOutcome::Updated { + consecutive_turns: 2, + .. + } + )); + assert_eq!( + vec!["turn-1".to_string(), "turn-2".to_string()], + blocker_audit_turn_ids(runtime.as_ref(), thread_id).await + ); + goal.goal_id + }; + + { + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("state db should reopen for delayed replay"); + let delayed_replay = runtime + .thread_goals() + .observe_active_thread_goal_blocker( + thread_id, + goal_id.as_str(), + "turn-1", + "codex_err:blocker_a", + 3, + ) + .await + .expect("delayed replay should be idempotent"); + assert!(matches!( + delayed_replay, + GoalBlockerAuditOutcome::Updated { + consecutive_turns: 2, + .. + } + )); + assert_eq!( + vec!["turn-1".to_string(), "turn-2".to_string()], + blocker_audit_turn_ids(runtime.as_ref(), thread_id).await, + "a non-adjacent replay from the durable current streak must not advance it" + ); + + let new_fingerprint = runtime + .thread_goals() + .observe_active_thread_goal_blocker( + thread_id, + goal_id.as_str(), + "turn-3", + "codex_err:blocker_b", + 3, + ) + .await + .expect("new blocker fingerprint should start a fresh streak"); + assert!(matches!( + new_fingerprint, + GoalBlockerAuditOutcome::Updated { + consecutive_turns: 1, + .. + } + )); + assert_eq!( + vec!["turn-3".to_string()], + blocker_audit_turn_ids(runtime.as_ref(), thread_id).await, + "turn ids from the previous fingerprint epoch must not outlive the bounded current streak" + ); + } + + let runtime = StateRuntime::init(codex_home, "test-provider".to_string()) + .await + .expect("state db should reopen with the replacement streak"); + let reused_outside_current_streak = runtime + .thread_goals() + .observe_active_thread_goal_blocker( + thread_id, + goal_id.as_str(), + "turn-1", + "codex_err:blocker_b", + 3, + ) + .await + .expect("an id outside the current streak should count normally"); + assert!(matches!( + reused_outside_current_streak, + GoalBlockerAuditOutcome::Updated { + consecutive_turns: 2, + .. + } + )); + assert_eq!( + vec!["turn-1".to_string(), "turn-3".to_string()], + blocker_audit_turn_ids(runtime.as_ref(), thread_id).await + ); + } + + #[tokio::test] + async fn blocker_audit_clears_for_every_goal_scope_change() { + #[derive(Clone, Copy)] + enum GoalScopeChange { + GoalId, + Objective, + Title, + Status, + TokenBudget, + } + + let runtime = test_runtime().await; + let thread_id = test_thread_id(); + upsert_test_thread(&runtime, thread_id).await; + for (index, change) in [ + GoalScopeChange::GoalId, + GoalScopeChange::Objective, + GoalScopeChange::Title, + GoalScopeChange::Status, + GoalScopeChange::TokenBudget, + ] + .into_iter() + .enumerate() + { + let goal = runtime + .thread_goals() + .replace_thread_goal( + thread_id, + format!("scope lifecycle {index}").as_str(), + crate::ThreadGoalStatus::Active, + /*token_budget*/ None, + ) + .await + .expect("goal should be created"); + runtime + .thread_goals() + .observe_active_thread_goal_blocker( + thread_id, + goal.goal_id.as_str(), + format!("turn-{index}").as_str(), + "codex_err:test_blocker", + 3, + ) + .await + .expect("blocker audit should persist"); + assert!(blocker_audit(runtime.as_ref(), thread_id).await.is_some()); + + match change { + GoalScopeChange::GoalId => { + sqlx::query("UPDATE thread_goals SET goal_id = ? WHERE thread_id = ?") + .bind(Uuid::new_v4().to_string()) + .bind(thread_id.to_string()) + .execute(runtime.thread_goals().pool.as_ref()) + .await + .expect("goal id should change through the trigger boundary"); + } + GoalScopeChange::Objective => { + sqlx::query("UPDATE thread_goals SET objective = ? WHERE thread_id = ?") + .bind("changed objective") + .bind(thread_id.to_string()) + .execute(runtime.thread_goals().pool.as_ref()) + .await + .expect("objective should change through the trigger boundary"); + } + GoalScopeChange::Title => { + sqlx::query("UPDATE thread_goals SET title = ? WHERE thread_id = ?") + .bind("Changed scope title") + .bind(thread_id.to_string()) + .execute(runtime.thread_goals().pool.as_ref()) + .await + .expect("title should change through the trigger boundary"); + } + GoalScopeChange::Status => { + sqlx::query("UPDATE thread_goals SET status = 'blocked' WHERE thread_id = ?") + .bind(thread_id.to_string()) + .execute(runtime.thread_goals().pool.as_ref()) + .await + .expect("status should change through the trigger boundary"); + } + GoalScopeChange::TokenBudget => { + sqlx::query("UPDATE thread_goals SET token_budget = ? WHERE thread_id = ?") + .bind(10_000_i64) + .bind(thread_id.to_string()) + .execute(runtime.thread_goals().pool.as_ref()) + .await + .expect("token budget should change through the trigger boundary"); + } + } + assert_eq!( + None, + blocker_audit(runtime.as_ref(), thread_id).await, + "every goal scope change must start a fresh blocker audit" + ); + assert_eq!( + Vec::::new(), + blocker_audit_turn_ids(runtime.as_ref(), thread_id).await, + "trigger cleanup must cascade to every persisted turn id" + ); + } + } + #[tokio::test] async fn goal_context_lifecycle_policy_round_trips() { let runtime = test_runtime().await; diff --git a/codex-rs/state/src/runtime/schedules.rs b/codex-rs/state/src/runtime/schedules.rs index 75dd46b015..db8a8ecc4f 100644 --- a/codex-rs/state/src/runtime/schedules.rs +++ b/codex-rs/state/src/runtime/schedules.rs @@ -5,15 +5,19 @@ use uuid::Uuid; pub const MAX_THREAD_SCHEDULE_NESTING_DEPTH: i64 = 5; const DYNAMIC_LOOP_CADENCE_SECONDS: i64 = 60; +const ONCE_SCHEDULE_KIND: &str = "once"; +const PAUSED_SCHEDULE_RUN_ERROR: &str = "scheduled run cancelled because schedule was paused"; +const EXPIRED_SCHEDULE_RUN_ERROR: &str = "scheduled run cancelled because schedule expired"; #[derive(Clone)] pub struct ScheduleStore { pool: Arc, + goals_pool: Arc, } impl ScheduleStore { - pub(crate) fn new(pool: Arc) -> Self { - Self { pool } + pub(crate) fn new(pool: Arc, goals_pool: Arc) -> Self { + Self { pool, goals_pool } } } @@ -44,6 +48,7 @@ pub struct ThreadScheduleClaim { pub run: crate::ThreadScheduleRun, } +#[derive(Clone)] pub struct ThreadScheduleDueClaimParams<'a> { pub now: DateTime, pub lease_id: &'a str, @@ -52,6 +57,7 @@ pub struct ThreadScheduleDueClaimParams<'a> { pub local_active_fresh_after: Option>, } +#[derive(Clone)] pub struct ThreadScheduleNowClaimParams<'a> { pub schedule_id: &'a str, pub now: DateTime, @@ -61,11 +67,56 @@ pub struct ThreadScheduleNowClaimParams<'a> { 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, +} + 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, @@ -268,6 +319,9 @@ ORDER BY status, next_run_at_ms IS NULL, next_run_at_ms, created_at_ms let prompt = redact_state_string(prompt); let timezone = redact_state_string(timezone); let cron_expression = spec.cron_expression.map(redact_state_string); + let now = Utc::now(); + let now_ms = datetime_to_epoch_millis(now); + let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; let sql = schedule_returning( r#" UPDATE thread_schedules @@ -280,9 +334,11 @@ SET cron_expression = ?, timezone = ?, status = ?, - next_run_at_ms = ?, + next_run_at_ms = CASE WHEN ? = 'expired' THEN NULL ELSE ? END, expires_at_ms = ?, failure_count = CASE WHEN ? THEN 0 ELSE failure_count END, + lease_id = CASE WHEN ? = 'active' THEN lease_id ELSE NULL END, + lease_expires_at_ms = CASE WHEN ? = 'active' THEN lease_expires_at_ms ELSE NULL END, updated_at_ms = ? WHERE schedule_id = ? RETURNING @@ -297,14 +353,42 @@ RETURNING .bind(cron_expression) .bind(timezone) .bind(status.as_str()) + .bind(status.as_str()) .bind(next_run_at.map(datetime_to_epoch_millis)) .bind(expires_at.map(datetime_to_epoch_millis)) .bind(reset_failure_count) - .bind(datetime_to_epoch_millis(Utc::now())) + .bind(status.as_str()) + .bind(status.as_str()) + .bind(now_ms) .bind(schedule_id) - .fetch_optional(self.pool.as_ref()) + .fetch_optional(&mut *tx) .await?; - row.map(|row| thread_schedule_from_row(&row)).transpose() + if row.is_some() { + let terminal_error = match status { + crate::ThreadScheduleStatus::Active => None, + crate::ThreadScheduleStatus::Paused => Some(PAUSED_SCHEDULE_RUN_ERROR), + crate::ThreadScheduleStatus::Expired => Some(EXPIRED_SCHEDULE_RUN_ERROR), + }; + if let Some(terminal_error) = terminal_error { + sqlx::query( + r#" +UPDATE thread_schedule_runs +SET status = 'failed', + error = ?, + completed_at_ms = COALESCE(completed_at_ms, ?) +WHERE schedule_id = ? AND status IN ('leased', 'running') + "#, + ) + .bind(redact_state_string(terminal_error)) + .bind(now_ms) + .bind(schedule_id) + .execute(&mut *tx) + .await?; + } + } + let schedule = row.map(|row| thread_schedule_from_row(&row)).transpose()?; + tx.commit().await?; + Ok(schedule) } pub async fn set_thread_schedule_status( @@ -572,23 +656,14 @@ WHERE parent_schedule_id = ? &self, run_id: &str, ) -> anyhow::Result> { - let row = sqlx::query( + let sql = run_returning( r#" SELECT - run_id, - schedule_id, - thread_id, - status, - lease_id, - turn_id, - error, - scheduled_for_ms, - started_at_ms, - completed_at_ms -FROM thread_schedule_runs -WHERE run_id = ? "#, - ) + ); + 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?; @@ -596,6 +671,34 @@ WHERE run_id = ? .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, @@ -684,71 +787,18 @@ LIMIT 1 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().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 = schedule_returning(&format!( - r#" -UPDATE thread_schedules -SET lease_id = ?, lease_expires_at_ms = ?, updated_at_ms = ? -WHERE schedule_id = ( - SELECT schedule_id - 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 -) -RETURNING -"#, - )); - let mut query = sqlx::query(sqlx::AssertSqlSafe(sql)) - .bind(lease_id) - .bind(lease_expires_at_ms) - .bind(now_ms) - .bind(now_ms) - .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 params = ClaimThreadScheduleParams { + target: ThreadScheduleClaimTarget::Due, + now, + lease_id, + lease_duration, + local_active_owner_id, + local_active_fresh_after, }; - let schedule = thread_schedule_from_row(&schedule_row)?; - let scheduled_for_ms = schedule.next_run_at.map(datetime_to_epoch_millis); - let run = - Self::insert_leased_run(&mut tx, &schedule, lease_id, scheduled_for_ms, now_ms).await?; - tx.commit().await?; - Ok(Some(ThreadScheduleClaim { schedule, run })) + 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( @@ -781,10 +831,36 @@ RETURNING 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().await?; + 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))) @@ -796,35 +872,50 @@ RETURNING 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 != ? + 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 = schedule_returning(&format!( - r#" -UPDATE thread_schedules -SET lease_id = ?, lease_expires_at_ms = ?, updated_at_ms = ? + 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} -RETURNING -"#, - )); - let mut query = sqlx::query(sqlx::AssertSqlSafe(sql)) - .bind(lease_id) - .bind(lease_expires_at_ms) - .bind(now_ms) - .bind(schedule_id) - .bind(now_ms) - .bind(now_ms); +"# + ), + }; + 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); } @@ -833,10 +924,131 @@ RETURNING 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; + let mut goal_tx = if goal_hold_can_pause { + Some(self.goals_pool.begin_with("BEGIN IMMEDIATE").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_one(&mut *tx) + .await?; 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, Some(now_ms), now_ms).await?; + 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 })) } @@ -866,6 +1078,7 @@ RETURNING status, lease_id, turn_id, + goal_id, error, scheduled_for_ms, started_at_ms, @@ -886,50 +1099,128 @@ RETURNING pub async fn mark_thread_schedule_run_started( &self, - schedule_id: &str, - run_id: &str, - lease_id: &str, - turn_id: &str, + 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 = ? -WHERE schedule_id = ? AND run_id = ? AND lease_id = ? +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(self.pool.as_ref()) + .fetch_optional(&mut *tx) .await?; - row.map(|row| thread_schedule_run_from_row(&row)) - .transpose() + 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, - schedule_id: &str, - lease_id: &str, - now: DateTime, - lease_duration: Duration, + 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 lease_id = ? +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(datetime_to_epoch_millis(now)) + .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) @@ -943,14 +1234,64 @@ WHERE schedule_id = ? AND lease_id = ? completed_at: DateTime, next_run_at: Option>, ) -> anyhow::Result { - self.finish_thread_schedule_run( + self.finish_thread_schedule_run(FinishThreadScheduleRunParams { schedule_id, run_id, lease_id, completed_at, next_run_at, - FinishScheduleRun::Completed, - ) + expected_goal_id: None, + finish: FinishScheduleRun::Completed { + pause_schedule: false, + }, + }) + .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 { + pause_schedule: false, + }, + }) + .await + } + + pub async fn complete_thread_schedule_run_and_pause( + &self, + schedule_id: &str, + run_id: &str, + lease_id: &str, + completed_at: DateTime, + ) -> anyhow::Result { + self.finish_thread_schedule_run(FinishThreadScheduleRunParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at: None, + expected_goal_id: None, + finish: FinishScheduleRun::Completed { + pause_schedule: true, + }, + }) .await } @@ -963,14 +1304,69 @@ WHERE schedule_id = ? AND lease_id = ? next_run_at: Option>, error: String, ) -> anyhow::Result { - self.finish_thread_schedule_run( + self.finish_thread_schedule_run(FinishThreadScheduleRunParams { schedule_id, run_id, lease_id, completed_at, next_run_at, - FinishScheduleRun::Failed { error }, - ) + expected_goal_id: None, + finish: FinishScheduleRun::Failed { + error, + pause_schedule: false, + }, + }) + .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, + pause_schedule: false, + }, + }) + .await + } + + pub async fn fail_thread_schedule_run_and_pause( + &self, + schedule_id: &str, + run_id: &str, + lease_id: &str, + completed_at: DateTime, + error: String, + ) -> anyhow::Result { + self.finish_thread_schedule_run(FinishThreadScheduleRunParams { + schedule_id, + run_id, + lease_id, + completed_at, + next_run_at: None, + expected_goal_id: None, + finish: FinishScheduleRun::Failed { + error, + pause_schedule: true, + }, + }) .await } @@ -991,13 +1387,16 @@ WHERE schedule_id = ? AND lease_id = ? 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, @@ -1043,11 +1442,40 @@ WHERE schedule_id = ? AND run_id = ? AND lease_id = ? 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 = ? @@ -1060,40 +1488,134 @@ WHERE status = 'active' .bind(now_ms) .bind(now_ms) .bind(now_ms) - .execute(self.pool.as_ref()) + .execute(&mut *tx) .await?; + tx.commit().await?; Ok(result.rows_affected()) } async fn finish_thread_schedule_run( &self, - schedule_id: &str, - run_id: &str, - lease_id: &str, - completed_at: DateTime, - next_run_at: Option>, - finish: FinishScheduleRun, + 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 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().await?; + 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; + let mut goal_tx = if goal_hold_can_pause { + Some(self.goals_pool.begin_with("BEGIN IMMEDIATE").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, + _ => unreachable!("goal transaction presence follows recurring goal schedule"), + }; let failed = matches!(finish, FinishScheduleRun::Failed { .. }); + let pause_schedule = match &finish { + FinishScheduleRun::Completed { pause_schedule } + | FinishScheduleRun::Failed { pause_schedule, .. } => { + *pause_schedule || pause_for_goal_hold + } + }; let schedule_result = sqlx::query( r#" UPDATE thread_schedules SET - status = CASE WHEN ? IS NULL THEN 'expired' ELSE status END, + 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 = ?, + 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) @@ -1103,11 +1625,16 @@ WHERE schedule_id = ? AND lease_id = ? .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); } let (status, error) = match &finish { - FinishScheduleRun::Completed => (crate::ThreadScheduleRunStatus::Completed, None), - FinishScheduleRun::Failed { error } => { + FinishScheduleRun::Completed { .. } => { + (crate::ThreadScheduleRunStatus::Completed, None) + } + FinishScheduleRun::Failed { error, .. } => { (crate::ThreadScheduleRunStatus::Failed, Some(error.as_str())) } }; @@ -1129,16 +1656,34 @@ WHERE schedule_id = ? AND run_id = ? AND lease_id = ? .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) } } +#[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 }, + Completed { pause_schedule: bool }, + Failed { error: String, pause_schedule: bool }, } struct ScheduleBindings<'a> { @@ -1151,7 +1696,7 @@ struct ScheduleBindings<'a> { fn schedule_bindings(schedule: &crate::ThreadScheduleSpec) -> ScheduleBindings<'_> { match schedule { crate::ThreadScheduleSpec::Once => ScheduleBindings { - kind: "once", + kind: ONCE_SCHEDULE_KIND, interval_amount: None, interval_unit: None, cron_expression: None, @@ -1267,6 +1812,7 @@ const RUN_COLUMNS: &str = r#" status, lease_id, turn_id, + goal_id, error, scheduled_for_ms, started_at_ms, @@ -2062,12 +2608,15 @@ mod tests { .expect("one-time schedule should claim"); runtime .thread_schedules() - .mark_thread_schedule_run_started( - &schedule.schedule_id, - &claim.run.run_id, - "lease-once", - "turn-once", - ) + .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"); @@ -2175,57 +2724,237 @@ mod tests { } #[tokio::test] - async fn claim_due_thread_schedule_skips_fresh_foreign_active_owner() { - let runtime = test_runtime().await; - let thread_id = test_thread_id(/*id*/ 3); - upsert_test_thread(&runtime, thread_id).await; + 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, thread_id, "live owner task", Some(now)).await; + 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 - .local_active_sessions() - .heartbeat_session(LocalActiveSessionHeartbeatParams { - thread_id, - owner_id: "owner-a".to_string(), - session_id: "session-a".to_string(), - pid: Some(100), + .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("active session should heartbeat"); + .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!( - runtime + reopened .thread_schedules() - .claim_due_thread_schedule_with_params(ThreadScheduleDueClaimParams { - now, - lease_id: "lease-owner-b", - lease_duration: Duration::from_secs(300), - local_active_owner_id: Some("owner-b"), - local_active_fresh_after: Some(now - chrono::Duration::seconds(15)), - }) + .claim_due_thread_schedule( + retry_at, + "lease-duplicate-retry", + Duration::from_secs(30), + ) .await - .expect("claim should not fail") + .expect("duplicate claim check should succeed") .is_none(), - "foreign processes should not claim loops owned by a fresh live session" + "one expired lease may create at most one replacement claim" ); + } - let owner_claim = runtime + #[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 .thread_schedules() - .claim_due_thread_schedule_with_params(ThreadScheduleDueClaimParams { + .claim_due_thread_schedule(now, "lease-goal-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-goal-restart", + turn_id: "turn-goal-restart", + goal_id: Some(&goal.goal_id), now, - lease_id: "lease-owner-a", - lease_duration: Duration::from_secs(300), - local_active_owner_id: Some("owner-a"), - local_active_fresh_after: Some(now - chrono::Duration::seconds(15)), + lease_duration: Duration::from_secs(30), }) .await - .expect("owner claim should succeed") - .expect("live owner should claim its due schedule"); + .expect("goal run should start") + .expect("goal run should still exist"); + drop(runtime); - assert_eq!(schedule.schedule_id, owner_claim.schedule.schedule_id); - assert_eq!( - Some("owner:lease-owner-a".to_string()), + 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 + .thread_schedules() + .claim_due_thread_schedule( + retry_at, + "lease-held-replacement", + Duration::from_secs(30), + ) + .await + .expect("expired goal run recovery should succeed") + .is_none(), + "a persisted held goal must pause instead of creating a replacement run" + ); + + 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); + } + + #[tokio::test] + async fn claim_due_thread_schedule_skips_fresh_foreign_active_owner() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(/*id*/ 3); + upsert_test_thread(&runtime, thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let schedule = + create_interval_schedule(&runtime, thread_id, "live owner task", Some(now)).await; + runtime + .local_active_sessions() + .heartbeat_session(LocalActiveSessionHeartbeatParams { + thread_id, + owner_id: "owner-a".to_string(), + session_id: "session-a".to_string(), + pid: Some(100), + now, + }) + .await + .expect("active session should heartbeat"); + + assert!( + runtime + .thread_schedules() + .claim_due_thread_schedule_with_params(ThreadScheduleDueClaimParams { + now, + lease_id: "lease-owner-b", + lease_duration: Duration::from_secs(300), + local_active_owner_id: Some("owner-b"), + local_active_fresh_after: Some(now - chrono::Duration::seconds(15)), + }) + .await + .expect("claim should not fail") + .is_none(), + "foreign processes should not claim loops owned by a fresh live session" + ); + + let owner_claim = runtime + .thread_schedules() + .claim_due_thread_schedule_with_params(ThreadScheduleDueClaimParams { + now, + lease_id: "lease-owner-a", + lease_duration: Duration::from_secs(300), + local_active_owner_id: Some("owner-a"), + local_active_fresh_after: Some(now - chrono::Duration::seconds(15)), + }) + .await + .expect("owner claim should succeed") + .expect("live owner should claim its due schedule"); + + assert_eq!(schedule.schedule_id, owner_claim.schedule.schedule_id); + assert_eq!( + Some("owner:lease-owner-a".to_string()), owner_claim.schedule.lease_id ); assert_eq!("owner:lease-owner-a", owner_claim.run.lease_id); @@ -2386,135 +3115,908 @@ mod tests { } #[tokio::test] - async fn claim_thread_schedule_now_ignores_legacy_claim_and_allows_live_owner() { - let runtime = test_runtime().await; - let thread_id = test_thread_id(/*id*/ 6); - upsert_test_thread(&runtime, thread_id).await; - let now = at(Utc::now().timestamp()); - let schedule = create_interval_schedule( - &runtime, - thread_id, - "manual live owner task", - Some(now + chrono::Duration::hours(1)), - ) - .await; + async fn claim_thread_schedule_now_ignores_legacy_claim_and_allows_live_owner() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(/*id*/ 6); + upsert_test_thread(&runtime, thread_id).await; + let now = at(Utc::now().timestamp()); + let schedule = create_interval_schedule( + &runtime, + thread_id, + "manual live owner task", + Some(now + chrono::Duration::hours(1)), + ) + .await; + runtime + .local_active_sessions() + .heartbeat_session(LocalActiveSessionHeartbeatParams { + thread_id, + owner_id: "owner-a".to_string(), + session_id: "session-a".to_string(), + pid: Some(100), + now, + }) + .await + .expect("active session should heartbeat"); + + assert!( + runtime + .thread_schedules() + .claim_thread_schedule_now( + &schedule.schedule_id, + now, + "legacy-manual-lease", + Duration::from_secs(300), + ) + .await + .expect("legacy manual claim should be ignored without failing") + .is_none(), + "legacy manual run-now should not steal loops from fresh live sessions" + ); + + assert!( + runtime + .thread_schedules() + .claim_thread_schedule_now_with_params(ThreadScheduleNowClaimParams { + schedule_id: &schedule.schedule_id, + now, + lease_id: "manual-foreign-lease", + lease_duration: Duration::from_secs(300), + local_active_owner_id: Some("owner-b"), + local_active_fresh_after: Some(now - chrono::Duration::seconds(15)), + }) + .await + .expect("foreign manual claim should not fail") + .is_none(), + "new foreign manual run-now should not steal loops from fresh live sessions" + ); + + let owner_claim = runtime + .thread_schedules() + .claim_thread_schedule_now_with_params(ThreadScheduleNowClaimParams { + schedule_id: &schedule.schedule_id, + now, + lease_id: "manual-owner-lease", + lease_duration: Duration::from_secs(300), + local_active_owner_id: Some("owner-a"), + local_active_fresh_after: Some(now - chrono::Duration::seconds(15)), + }) + .await + .expect("owner manual claim should succeed") + .expect("live owner should claim manual run-now"); + + assert_eq!(schedule.schedule_id, owner_claim.schedule.schedule_id); + assert_eq!( + Some("owner:manual-owner-lease".to_string()), + owner_claim.schedule.lease_id + ); + assert_eq!("owner:manual-owner-lease", owner_claim.run.lease_id); + assert_eq!(Some(now), owner_claim.run.scheduled_for); + } + + #[tokio::test] + async fn extend_thread_schedule_lease_refreshes_live_claim() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(/*id*/ 6); + upsert_test_thread(&runtime, thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let schedule = + create_interval_schedule(&runtime, thread_id, "long running task", Some(now)).await; + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-long", Duration::from_secs(300)) + .await + .expect("claim should succeed") + .expect("schedule should claim"); + assert_eq!(schedule.schedule_id, claim.schedule.schedule_id); + + assert!( + runtime + .thread_schedules() + .extend_thread_schedule_lease(ThreadScheduleRunLeaseParams { + schedule_id: &schedule.schedule_id, + run_id: &claim.run.run_id, + lease_id: "lease-long", + now: now + chrono::Duration::seconds(120), + lease_duration: Duration::from_secs(300), + }) + .await + .expect("lease should extend") + ); + let refreshed = runtime + .thread_schedules() + .get_thread_schedule(&schedule.schedule_id) + .await + .expect("schedule should load") + .expect("schedule should exist"); + assert_eq!( + Some(now + chrono::Duration::seconds(420)), + refreshed.lease_expires_at + ); + assert!( + !runtime + .thread_schedules() + .extend_thread_schedule_lease(ThreadScheduleRunLeaseParams { + schedule_id: &schedule.schedule_id, + run_id: &claim.run.run_id, + lease_id: "wrong-lease", + now: now + chrono::Duration::seconds(180), + lease_duration: Duration::from_secs(300), + }) + .await + .expect("wrong lease should not fail") + ); + } + + #[tokio::test] + async fn expired_heartbeat_cannot_revive_schedule_lease() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(/*id*/ 49); + upsert_test_thread(&runtime, thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let schedule = + create_interval_schedule(&runtime, thread_id, "stale heartbeat", Some(now)).await; + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-stale", Duration::from_secs(30)) + .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"); + let expired_at = now + chrono::Duration::seconds(31); + + assert!( + !runtime + .thread_schedules() + .extend_thread_schedule_lease(ThreadScheduleRunLeaseParams { + schedule_id: &schedule.schedule_id, + run_id: &claim.run.run_id, + lease_id: "lease-stale", + now: expired_at, + lease_duration: Duration::from_secs(30), + }) + .await + .expect("expired heartbeat should fail closed") + ); + assert_eq!( + Some(now + chrono::Duration::seconds(30)), + runtime + .thread_schedules() + .get_thread_schedule(&schedule.schedule_id) + .await + .expect("schedule should load") + .expect("schedule should exist") + .lease_expires_at + ); + let replacement = 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); + } + + #[tokio::test] + async fn completed_schedule_run_can_atomically_pause_without_rearming() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(/*id*/ 31); + upsert_test_thread(&runtime, thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let schedule = create_interval_schedule(&runtime, thread_id, "held task", Some(now)).await; + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-held", Duration::from_secs(300)) + .await + .expect("claim should succeed") + .expect("schedule should claim"); + + assert!( + runtime + .thread_schedules() + .complete_thread_schedule_run_and_pause( + &schedule.schedule_id, + &claim.run.run_id, + "lease-held", + now + chrono::Duration::seconds(5), + ) + .await + .expect("run should complete while pausing the schedule") + ); + + let held_schedule = runtime + .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); + assert_eq!(0, held_schedule.failure_count); + let run = runtime + .thread_schedules() + .get_thread_schedule_run(&claim.run.run_id) + .await + .expect("run should load") + .expect("run should exist"); + assert_eq!(crate::ThreadScheduleRunStatus::Completed, run.status); + assert!( + runtime + .thread_schedules() + .claim_due_thread_schedule( + now + chrono::Duration::days(1), + "lease-rearm", + Duration::from_secs(300), + ) + .await + .expect("paused schedule claim should not fail") + .is_none(), + "a held schedule must not become claimable again" + ); + } + + #[tokio::test] + async fn goal_correlated_completion_ignores_replacement_with_same_objective() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(/*id*/ 41); + upsert_test_thread(&runtime, thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let next_run_at = now + chrono::Duration::minutes(5); + let objective = "repeat the same objective"; + let original_goal = runtime + .thread_goals() + .replace_thread_goal( + thread_id, + objective, + crate::ThreadGoalStatus::Active, + /*token_budget*/ None, + ) + .await + .expect("original goal should be created"); + let schedule = create_interval_schedule(&runtime, thread_id, objective, Some(now)).await; + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-original-goal", Duration::from_secs(300)) + .await + .expect("claim should succeed") + .expect("schedule should claim"); + let replacement_goal = runtime + .thread_goals() + .replace_thread_goal( + thread_id, + objective, + crate::ThreadGoalStatus::Blocked, + /*token_budget*/ None, + ) + .await + .expect("replacement goal should be created"); + assert_ne!(original_goal.goal_id, replacement_goal.goal_id); + + assert!( + runtime + .thread_schedules() + .complete_thread_schedule_run_for_goal(ThreadScheduleRunForGoalFinishParams { + schedule_id: &schedule.schedule_id, + run_id: &claim.run.run_id, + lease_id: "lease-original-goal", + completed_at: now + chrono::Duration::seconds(5), + next_run_at: Some(next_run_at), + expected_goal_id: original_goal.goal_id.as_str(), + },) + .await + .expect("run should complete") + ); + + let completed = runtime + .thread_schedules() + .get_thread_schedule(&schedule.schedule_id) + .await + .expect("schedule should load") + .expect("schedule should exist"); + assert_eq!(crate::ThreadScheduleStatus::Active, completed.status); + assert_eq!(Some(next_run_at), completed.next_run_at); + assert_eq!(None, completed.lease_id); + } + + #[tokio::test] + async fn goal_correlated_once_completion_expires_instead_of_pausing() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(/*id*/ 43); + upsert_test_thread(&runtime, thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let goal = runtime + .thread_goals() + .replace_thread_goal( + thread_id, + "finish once while held", + crate::ThreadGoalStatus::Blocked, + /*token_budget*/ None, + ) + .await + .expect("blocked goal should be created"); + let schedule = runtime + .thread_schedules() + .create_thread_schedule(ThreadScheduleCreateParams { + thread_id, + prompt: "finish once while held".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("one-time schedule should be created"); + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-once", Duration::from_secs(300)) + .await + .expect("claim should succeed") + .expect("schedule should claim"); + + assert!( + runtime + .thread_schedules() + .complete_thread_schedule_run_for_goal(ThreadScheduleRunForGoalFinishParams { + schedule_id: &schedule.schedule_id, + run_id: &claim.run.run_id, + lease_id: "lease-once", + completed_at: now + chrono::Duration::seconds(5), + next_run_at: None, + expected_goal_id: goal.goal_id.as_str(), + },) + .await + .expect("run should complete") + ); + + let completed = runtime + .thread_schedules() + .get_thread_schedule(&schedule.schedule_id) + .await + .expect("schedule should load") + .expect("schedule should exist"); + assert_eq!(crate::ThreadScheduleStatus::Expired, completed.status); + assert_eq!(None, completed.next_run_at); + assert_eq!(None, completed.lease_id); + } + + #[tokio::test] + async fn concurrent_goal_correlated_finalizers_complete_exactly_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*/ 42); + upsert_test_thread(runtime.as_ref(), thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + let goal = runtime + .thread_goals() + .replace_thread_goal( + thread_id, + "pause once under contention", + crate::ThreadGoalStatus::Blocked, + /*token_budget*/ None, + ) + .await + .expect("blocked goal should be created"); + let schedule = create_interval_schedule( + runtime.as_ref(), + thread_id, + "pause once under contention", + Some(now), + ) + .await; + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-contended", Duration::from_secs(300)) + .await + .expect("claim should succeed") + .expect("schedule should claim"); + + let contender_state_pool = SqlitePoolOptions::new() + .max_connections(1) + .connect_with( + SqliteConnectOptions::new() + .filename(codex_home.join(crate::STATE_DB_FILENAME)) + .journal_mode(SqliteJournalMode::Wal) + .busy_timeout(Duration::from_millis(1)), + ) + .await + .expect("contending state pool should open"); + let contender_goals_pool = SqlitePoolOptions::new() + .max_connections(1) + .connect_with( + SqliteConnectOptions::new() + .filename(codex_home.join(crate::GOALS_DB_FILENAME)) + .journal_mode(SqliteJournalMode::Wal) + .busy_timeout(Duration::from_millis(1)), + ) + .await + .expect("contending goals pool should open"); + let contender = ScheduleStore::new( + Arc::new(contender_state_pool), + Arc::new(contender_goals_pool), + ); + let completed_at = now + chrono::Duration::seconds(5); + let next_run_at = Some(now + chrono::Duration::minutes(5)); + let primary_completion = runtime + .thread_schedules() + .complete_thread_schedule_run_for_goal(ThreadScheduleRunForGoalFinishParams { + schedule_id: &schedule.schedule_id, + run_id: &claim.run.run_id, + lease_id: "lease-contended", + completed_at, + next_run_at, + expected_goal_id: goal.goal_id.as_str(), + }); + let contender_completion = + contender.complete_thread_schedule_run_for_goal(ThreadScheduleRunForGoalFinishParams { + schedule_id: &schedule.schedule_id, + run_id: &claim.run.run_id, + lease_id: "lease-contended", + completed_at, + next_run_at, + expected_goal_id: goal.goal_id.as_str(), + }); + let (primary_result, contender_result) = + tokio::join!(primary_completion, contender_completion); + let completions = [ + primary_result.expect("primary finalizer should succeed"), + contender_result.expect("contending finalizer should succeed"), + ]; + assert_eq!( + 1, + completions + .into_iter() + .filter(|completed| *completed) + .count(), + "the schedule lease must let exactly one finalizer commit" + ); + + let held_schedule = runtime + .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 run = runtime + .thread_schedules() + .get_thread_schedule_run(&claim.run.run_id) + .await + .expect("run should load") + .expect("run should exist"); + assert_eq!(crate::ThreadScheduleRunStatus::Completed, run.status); + } + + #[tokio::test] + async fn late_terminal_and_expired_lease_reaper_settle_one_run_owner() { + 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*/ 46); + 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, "terminal race", Some(now)).await; + let claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-terminal-race", Duration::from_secs(30)) + .await + .expect("initial claim should succeed") + .expect("schedule should claim"); runtime - .local_active_sessions() - .heartbeat_session(LocalActiveSessionHeartbeatParams { - thread_id, - owner_id: "owner-a".to_string(), - session_id: "session-a".to_string(), - pid: Some(100), + .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("active session should heartbeat"); + .expect("run should start") + .expect("run should still exist"); + + let contender = StateRuntime::init(codex_home, "test-provider".to_string()) + .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 original_run = runtime + .thread_schedules() + .get_thread_schedule_run(&claim.run.run_id) + .await + .expect("original run should load") + .expect("original run should exist"); assert!( - runtime - .thread_schedules() - .claim_thread_schedule_now( - &schedule.schedule_id, - now, - "legacy-manual-lease", - Duration::from_secs(300), - ) - .await - .expect("legacy manual claim should be ignored without failing") - .is_none(), - "legacy manual run-now should not steal loops from fresh live sessions" + matches!( + original_run.status, + crate::ThreadScheduleRunStatus::Completed | crate::ThreadScheduleRunStatus::Failed + ), + "the old running row must be terminal after the race" ); + 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); + } + + #[tokio::test] + async fn expired_lease_reaper_prevents_delayed_start_resurrection() { + 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*/ 47); + 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, "delayed start", Some(now)).await; + let original_claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-delayed-start", Duration::from_secs(30)) + .await + .expect("initial claim should succeed") + .expect("schedule should claim"); + let contender = StateRuntime::init(codex_home, "test-provider".to_string()) + .await + .expect("contending state runtime should initialize"); + let retry_at = now + chrono::Duration::seconds(31); + let delayed_start = runtime.thread_schedules().mark_thread_schedule_run_started( + ThreadScheduleRunStartParams { + schedule_id: &schedule.schedule_id, + run_id: &original_claim.run.run_id, + lease_id: "lease-delayed-start", + turn_id: "turn-delayed-start", + goal_id: None, + now: retry_at, + lease_duration: Duration::from_secs(30), + }, + ); + let replacement = contender.thread_schedules().claim_due_thread_schedule( + retry_at, + "lease-replacement", + Duration::from_secs(30), + ); + let (delayed_start, replacement) = tokio::join!(delayed_start, replacement); assert!( - runtime - .thread_schedules() - .claim_thread_schedule_now_with_params(ThreadScheduleNowClaimParams { - schedule_id: &schedule.schedule_id, - now, - lease_id: "manual-foreign-lease", - lease_duration: Duration::from_secs(300), - local_active_owner_id: Some("owner-b"), - local_active_fresh_after: Some(now - chrono::Duration::seconds(15)), - }) - .await - .expect("foreign manual claim should not fail") + delayed_start + .expect("delayed start should not error") .is_none(), - "new foreign manual run-now should not steal loops from fresh live sessions" + "a reaped expired run must never become dispatchable" ); + let replacement = replacement + .expect("expired lease reaper should not error") + .expect("expired run should produce one replacement claim"); - let owner_claim = runtime + let original_run = runtime .thread_schedules() - .claim_thread_schedule_now_with_params(ThreadScheduleNowClaimParams { + .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 + ); + 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, - now, - lease_id: "manual-owner-lease", - lease_duration: Duration::from_secs(300), - local_active_owner_id: Some("owner-a"), - local_active_fresh_after: Some(now - chrono::Duration::seconds(15)), + 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("owner manual claim should succeed") - .expect("live owner should claim manual run-now"); - - assert_eq!(schedule.schedule_id, owner_claim.schedule.schedule_id); + .expect("replacement start should not error") + .expect("replacement should remain the sole dispatchable run"); assert_eq!( - Some("owner:manual-owner-lease".to_string()), - owner_claim.schedule.lease_id + crate::ThreadScheduleRunStatus::Running, + replacement_run.status ); - assert_eq!("owner:manual-owner-lease", owner_claim.run.lease_id); - assert_eq!(Some(now), owner_claim.run.scheduled_for); + assert!( + runtime + .thread_schedules() + .complete_thread_schedule_run( + &schedule.schedule_id, + &replacement.run.run_id, + "lease-replacement", + replacement_started_at + chrono::Duration::seconds(1), + Some(now + chrono::Duration::hours(1)), + ) + .await + .expect("replacement should remain finalizable") + ); + + let stats = runtime + .thread_schedules() + .get_thread_schedule_stats(&schedule.schedule_id) + .await + .expect("schedule stats should load"); + assert_eq!(2, 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); } #[tokio::test] - async fn extend_thread_schedule_lease_refreshes_live_claim() { + async fn explicit_expiry_terminalizes_active_run_before_late_completion() { let runtime = test_runtime().await; - let thread_id = test_thread_id(/*id*/ 6); + let thread_id = test_thread_id(/*id*/ 32); upsert_test_thread(&runtime, thread_id).await; let now = at(/*seconds*/ 1_700_000_000); let schedule = - create_interval_schedule(&runtime, thread_id, "long running task", Some(now)).await; + create_interval_schedule(&runtime, thread_id, "expired while leased", Some(now)).await; let claim = runtime .thread_schedules() - .claim_due_thread_schedule(now, "lease-long", Duration::from_secs(300)) + .claim_due_thread_schedule(now, "lease-expired", Duration::from_secs(300)) .await .expect("claim should succeed") .expect("schedule should claim"); - assert_eq!(schedule.schedule_id, claim.schedule.schedule_id); + + let expired = runtime + .thread_schedules() + .set_thread_schedule_status(&schedule.schedule_id, crate::ThreadScheduleStatus::Expired) + .await + .expect("schedule should expire") + .expect("schedule should still exist"); + assert_eq!(crate::ThreadScheduleStatus::Expired, expired.status); + assert_eq!(None, expired.lease_id); assert!( - runtime + !runtime .thread_schedules() - .extend_thread_schedule_lease( + .complete_thread_schedule_run_and_pause( &schedule.schedule_id, - "lease-long", - now + chrono::Duration::seconds(120), - Duration::from_secs(300), + &claim.run.run_id, + "lease-expired", + now + chrono::Duration::seconds(5), ) .await - .expect("lease should extend") + .expect("late completion should fail closed") ); - let refreshed = runtime + + let completed = runtime .thread_schedules() .get_thread_schedule(&schedule.schedule_id) .await .expect("schedule should load") .expect("schedule should exist"); assert_eq!( - Some(now + chrono::Duration::seconds(420)), - refreshed.lease_expires_at + crate::ThreadScheduleStatus::Expired, + completed.status, + "a late held completion must not downgrade expired to paused" ); + assert_eq!(None, completed.next_run_at); + assert_eq!(None, completed.lease_id); + let run = runtime + .thread_schedules() + .get_thread_schedule_run(&claim.run.run_id) + .await + .expect("run should load") + .expect("run should exist"); + assert_eq!(crate::ThreadScheduleRunStatus::Failed, run.status); + assert_eq!(Some(expired.updated_at), run.completed_at); + } + + #[tokio::test] + async fn terminal_schedule_status_rejects_late_finalizers() { + let runtime = test_runtime().await; + let thread_id = test_thread_id(/*id*/ 33); + upsert_test_thread(&runtime, thread_id).await; + let now = at(/*seconds*/ 1_700_000_000); + + for (index, held_status) in [ + crate::ThreadScheduleStatus::Paused, + crate::ThreadScheduleStatus::Expired, + ] + .into_iter() + .enumerate() + { + let expected_next_run_at = match held_status { + crate::ThreadScheduleStatus::Active => { + unreachable!("only terminal statuses tested") + } + crate::ThreadScheduleStatus::Paused => Some(now), + crate::ThreadScheduleStatus::Expired => None, + }; + let complete_lease = format!("lease-complete-{index}"); + let complete_schedule = create_interval_schedule( + &runtime, + thread_id, + &format!("late complete {index}"), + Some(now), + ) + .await; + let complete_claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, complete_lease.as_str(), Duration::from_secs(300)) + .await + .expect("complete schedule should claim") + .expect("complete schedule should be due"); + runtime + .thread_schedules() + .set_thread_schedule_status(&complete_schedule.schedule_id, held_status) + .await + .expect("complete schedule status should update") + .expect("complete schedule should exist"); + assert!( + !runtime + .thread_schedules() + .complete_thread_schedule_run( + &complete_schedule.schedule_id, + &complete_claim.run.run_id, + complete_lease.as_str(), + now + chrono::Duration::seconds(5), + Some(now + chrono::Duration::hours(1)), + ) + .await + .expect("late completion should fail closed") + ); + let after_complete = runtime + .thread_schedules() + .get_thread_schedule(&complete_schedule.schedule_id) + .await + .expect("complete schedule should load") + .expect("complete schedule should exist"); + assert_eq!(held_status, after_complete.status); + assert_eq!(expected_next_run_at, after_complete.next_run_at); + assert_eq!(None, after_complete.lease_id); + assert_eq!( + crate::ThreadScheduleRunStatus::Failed, + runtime + .thread_schedules() + .get_thread_schedule_run(&complete_claim.run.run_id) + .await + .expect("completed run should load") + .expect("completed run should exist") + .status + ); + + let defer_lease = format!("lease-defer-{index}"); + let defer_schedule = create_interval_schedule( + &runtime, + thread_id, + &format!("late defer {index}"), + Some(now), + ) + .await; + let defer_claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, defer_lease.as_str(), Duration::from_secs(300)) + .await + .expect("deferred schedule should claim") + .expect("deferred schedule should be due"); + runtime + .thread_schedules() + .set_thread_schedule_status(&defer_schedule.schedule_id, held_status) + .await + .expect("deferred schedule status should update") + .expect("deferred schedule should exist"); + assert!( + !runtime + .thread_schedules() + .defer_thread_schedule_run( + &defer_schedule.schedule_id, + &defer_claim.run.run_id, + defer_lease.as_str(), + now + chrono::Duration::seconds(5), + now + chrono::Duration::hours(1), + "held by goal status".to_string(), + ) + .await + .expect("late deferral should fail closed") + ); + let after_defer = runtime + .thread_schedules() + .get_thread_schedule(&defer_schedule.schedule_id) + .await + .expect("deferred schedule should load") + .expect("deferred schedule should exist"); + assert_eq!(held_status, after_defer.status); + assert_eq!(expected_next_run_at, after_defer.next_run_at); + assert_eq!(None, after_defer.lease_id); + let deferred_run = runtime + .thread_schedules() + .get_thread_schedule_run(&defer_claim.run.run_id) + .await + .expect("deferred run should load") + .expect("deferred run should exist"); + assert_eq!(crate::ThreadScheduleRunStatus::Failed, deferred_run.status); + } + + let failed_schedule = + create_interval_schedule(&runtime, thread_id, "late failed hold", Some(now)).await; + let failed_claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-failed", Duration::from_secs(300)) + .await + .expect("failed schedule should claim") + .expect("failed schedule should be due"); + runtime + .thread_schedules() + .set_thread_schedule_status( + &failed_schedule.schedule_id, + crate::ThreadScheduleStatus::Expired, + ) + .await + .expect("failed schedule should expire") + .expect("failed schedule should exist"); assert!( !runtime .thread_schedules() - .extend_thread_schedule_lease( - &schedule.schedule_id, - "wrong-lease", - now + chrono::Duration::seconds(180), - Duration::from_secs(300), + .fail_thread_schedule_run_and_pause( + &failed_schedule.schedule_id, + &failed_claim.run.run_id, + "lease-failed", + now + chrono::Duration::seconds(5), + "goal held".to_string(), ) .await - .expect("wrong lease should not fail") + .expect("late failure should fail closed") ); + let after_failure = runtime + .thread_schedules() + .get_thread_schedule(&failed_schedule.schedule_id) + .await + .expect("failed schedule should load") + .expect("failed schedule should exist"); + assert_eq!(crate::ThreadScheduleStatus::Expired, after_failure.status); + assert_eq!(None, after_failure.next_run_at); + assert_eq!(None, after_failure.lease_id); } #[tokio::test] @@ -2534,12 +4036,15 @@ mod tests { let running = runtime .thread_schedules() - .mark_thread_schedule_run_started( - &completed_schedule.schedule_id, - &completed_claim.run.run_id, - "lease-complete", - "turn-1", - ) + .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"); @@ -3110,12 +4615,15 @@ mod tests { .expect("schedule should claim"); runtime .thread_schedules() - .mark_thread_schedule_run_started( - &schedule.schedule_id, - &claim.run.run_id, - "lease-live", - "turn-live", - ) + .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"); @@ -3205,4 +4713,119 @@ mod tests { assert_eq!(crate::ThreadScheduleStatus::Expired, expired.status); assert_eq!(None, expired.lease_id); } + + #[tokio::test] + async fn pause_and_expiry_terminalize_active_schedule_runs() { + 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 paused_schedule = + create_interval_schedule(&runtime, thread_id, "pause active run", Some(now)).await; + let paused_claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-paused", Duration::from_secs(300)) + .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"); + let paused = runtime + .thread_schedules() + .set_thread_schedule_status( + &paused_schedule.schedule_id, + crate::ThreadScheduleStatus::Paused, + ) + .await + .expect("pause should succeed") + .expect("paused schedule should exist"); + assert_eq!(crate::ThreadScheduleStatus::Paused, paused.status); + assert_eq!(None, paused.lease_id); + let paused_run = runtime + .thread_schedules() + .get_thread_schedule_run(&paused_claim.run.run_id) + .await + .expect("paused run should load") + .expect("paused run should exist"); + assert_eq!(crate::ThreadScheduleRunStatus::Failed, paused_run.status); + assert_eq!(Some(paused.updated_at), paused_run.completed_at); + + let expiring_schedule = runtime + .thread_schedules() + .create_thread_schedule(ThreadScheduleCreateParams { + thread_id, + prompt: "expire active run".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: Some(now + chrono::Duration::seconds(10)), + }) + .await + .expect("expiring schedule should create"); + let expiring_claim = runtime + .thread_schedules() + .claim_due_thread_schedule(now, "lease-expiring", Duration::from_secs(30)) + .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"); + let expired_at = now + chrono::Duration::seconds(31); + assert_eq!( + 1, + runtime + .thread_schedules() + .expire_thread_schedules(expired_at) + .await + .expect("expiry cleanup should succeed") + ); + let expired_run = runtime + .thread_schedules() + .get_thread_schedule_run(&expiring_claim.run.run_id) + .await + .expect("expired run should load") + .expect("expired run should exist"); + assert_eq!(crate::ThreadScheduleRunStatus::Failed, expired_run.status); + assert_eq!(Some(expired_at), expired_run.completed_at); + + for schedule_id in [ + paused_schedule.schedule_id.as_str(), + expiring_schedule.schedule_id.as_str(), + ] { + let stats = runtime + .thread_schedules() + .get_thread_schedule_stats(schedule_id) + .await + .expect("schedule stats should load"); + assert_eq!(0, stats.leased_runs); + assert_eq!(0, stats.running_runs); + } + } }