From a332c196c43025078b82d0139d2bd5bd984f6e38 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Thu, 9 Apr 2026 18:08:17 -0700 Subject: [PATCH] Move thread timers to sqlite --- codex-rs/core/src/codex.rs | 456 +++++++++++++----- codex-rs/core/src/codex_tests.rs | 100 +--- codex-rs/core/src/timers.rs | 286 ++--------- codex-rs/core/tests/suite/timers.rs | 131 +++++ codex-rs/rollout/src/lib.rs | 1 - codex-rs/rollout/src/list.rs | 10 +- codex-rs/rollout/src/timer_sidecar.rs | 20 - .../state/migrations/0025_thread_timers.sql | 22 + codex-rs/state/src/lib.rs | 6 +- codex-rs/state/src/model/mod.rs | 5 + codex-rs/state/src/model/thread_timer.rs | 73 +++ codex-rs/state/src/runtime.rs | 5 + codex-rs/state/src/runtime/timers.rs | 371 ++++++++++++++ 13 files changed, 1013 insertions(+), 473 deletions(-) delete mode 100644 codex-rs/rollout/src/timer_sidecar.rs create mode 100644 codex-rs/state/migrations/0025_thread_timers.sql create mode 100644 codex-rs/state/src/model/thread_timer.rs create mode 100644 codex-rs/state/src/runtime/timers.rs diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index c8601b1d5b..300cc11564 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1,11 +1,13 @@ use std::collections::HashMap; use std::collections::HashSet; use std::fmt::Debug; -use std::future::Future; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; @@ -307,6 +309,7 @@ use crate::tasks::SessionTaskContext; use crate::timers::ClaimedTimer; use crate::timers::CreateTimer; use crate::timers::PersistedTimer; +use crate::timers::RestoredTimerTask; use crate::timers::TIMER_FIRED_BACKGROUND_EVENT_PREFIX; use crate::timers::TIMER_UPDATED_BACKGROUND_EVENT_PREFIX; use crate::timers::ThreadTimer; @@ -314,9 +317,7 @@ use crate::timers::ThreadTimerTrigger; use crate::timers::TimerDelivery; use crate::timers::TimerTaskSpec; use crate::timers::TimersState; -use crate::timers::load_timer_sidecar; use crate::timers::timer_prompt_input_item; -use crate::timers::write_timer_sidecar; use crate::tools::ToolRouter; use crate::tools::context::SharedTurnDiffTracker; use crate::tools::js_repl::JsReplHandle; @@ -415,6 +416,44 @@ pub struct Codex { pub(crate) type SessionLoopTermination = Shared>; +const TIMER_SOURCE_AGENT: &str = "agent"; +const TIMER_CLIENT_ID_FALLBACK: &str = "codex-cli"; +const TIMER_DB_SYNC_INTERVAL: Duration = Duration::from_secs(15); +const TIMER_DB_MAX_REFRESH_INTERVAL: Duration = Duration::from_secs(60); + +fn db_timer_to_persisted_timer(row: codex_state::ThreadTimer) -> Option { + let trigger = match serde_json::from_str(&row.trigger_json) { + Ok(trigger) => trigger, + Err(err) => { + warn!("skipping invalid persisted timer {} trigger: {err}", row.id); + return None; + } + }; + let delivery = + match serde_json::from_value::(serde_json::Value::String(row.delivery)) { + Ok(delivery) => delivery, + Err(err) => { + warn!( + "skipping invalid persisted timer {} delivery: {err}", + row.id + ); + return None; + } + }; + Some(PersistedTimer { + timer: ThreadTimer { + id: row.id, + trigger, + prompt: row.prompt, + delivery, + created_at: row.created_at, + next_run_at: row.next_run_at, + last_run_at: row.last_run_at, + }, + pending_run: row.pending_run, + }) +} + /// Wrapper returned by [`Codex::spawn`] containing the spawned [`Codex`], /// the submission id for the initial `ConfigureSession` request and the /// unique session id. @@ -838,9 +877,7 @@ pub(crate) struct Session { /// Prevents concurrent timers from claiming multiple timers before a /// newly started turn becomes the active turn. timer_start_in_progress: Mutex, - /// Serializes timer sidecar snapshots and file writes so stale timer state - /// cannot overwrite newer persisted state. - timer_sidecar_write_lock: Mutex<()>, + timer_db_sync_started: AtomicBool, mailbox: Mailbox, mailbox_rx: Mutex, idle_pending_input: Mutex>, // TODO (jif) merge with mailbox! @@ -2014,7 +2051,7 @@ impl Session { conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), timer_start_in_progress: Mutex::new(false), - timer_sidecar_write_lock: Mutex::new(()), + timer_db_sync_started: AtomicBool::new(false), mailbox, mailbox_rx: Mutex::new(mailbox_rx), idle_pending_input: Mutex::new(Vec::new()), @@ -2057,7 +2094,7 @@ impl Session { for event in events { sess.send_event_raw(event).await; } - sess.restore_timers_from_sidecar().await; + sess.restore_timers_from_db().await; // Start the watcher after SessionConfigured so it cannot emit earlier events. sess.start_skills_watcher_listener(); @@ -2170,6 +2207,36 @@ impl Session { self.services.state_db.clone() } + async fn timer_state_db(&self) -> Result { + if let Some(state_db) = self.state_db() { + return Ok(state_db); + } + + let config = { + let state = self.state.lock().await; + state + .session_configuration + .original_config_do_not_use + .clone() + }; + match codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.model_provider_id.clone(), + ) + .await + { + Ok(runtime) => Ok(runtime), + Err(err) => { + let message = format!( + "failed to initialize SQLite state db for timers at {}: {err}", + config.sqlite_home.display() + ); + warn!("{message}"); + Err(message) + } + } + } + /// Ensure rollout file writes are durably flushed. pub(crate) async fn flush_rollout(&self) { let recorder = { @@ -2700,7 +2767,12 @@ impl Session { .await } - pub(crate) async fn list_timers(&self) -> Vec { + pub(crate) async fn list_timers(self: &Arc) -> Vec { + self.sync_timers_from_db(/*emit_update*/ false).await; + self.list_timers_from_memory().await + } + + async fn list_timers_from_memory(&self) -> Vec { self.timers.lock().await.list_timers() } @@ -2711,15 +2783,14 @@ impl Session { delivery: TimerDelivery, ) -> Result { self.ensure_rollout_materialized().await; - let Some(rollout_path) = self.current_rollout_path().await else { - return Err("timers require a persisted rollout path for restoration".to_string()); - }; + let state_db = self.timer_state_db().await?; + self.start_timer_db_sync_task(state_db.clone()); let timer_cancel = CancellationToken::new(); let id = uuid::Uuid::new_v4().to_string(); - let (timer, timer_spec) = { + let (timer, persisted_timer, timer_spec) = { let mut timers = self.timers.lock().await; - timers.create_timer( + let (timer, timer_spec) = timers.create_timer( CreateTimer { id: id.clone(), trigger, @@ -2728,13 +2799,20 @@ impl Session { now: Utc::now(), }, Some(timer_cancel.clone()), - )? + )?; + let persisted_timer = timers + .persisted_timer(&id) + .ok_or_else(|| format!("created timer {id} was not stored in memory"))?; + (timer, persisted_timer, timer_spec) }; - if let Err(err) = self.persist_timers_to_rollout_sidecar(&rollout_path).await { + let params = self + .thread_timer_create_params(&persisted_timer, TIMER_SOURCE_AGENT) + .await?; + if let Err(err) = state_db.create_thread_timer(¶ms).await { if let Some(runtime) = self.timers.lock().await.remove_timer(&id) { TimersState::cancel_runtime(&runtime); } - return Err(err); + return Err(format!("failed to persist timer to sqlite: {err}")); } if let Some(timer_spec) = timer_spec { @@ -2745,17 +2823,30 @@ impl Session { Ok(timer) } - pub(crate) async fn delete_timer(&self, id: &str) -> Result { - let Some(runtime) = self.timers.lock().await.remove_timer(id) else { + pub(crate) async fn delete_timer(self: &Arc, id: &str) -> Result { + self.sync_timers_from_db(/*emit_update*/ false).await; + let state_db = self.timer_state_db().await?; + self.start_timer_db_sync_task(state_db.clone()); + + let runtime = self.timers.lock().await.remove_timer(id); + let deleted = match state_db + .delete_thread_timer(&self.thread_id_string(), id) + .await + { + Ok(deleted) => deleted, + Err(err) => { + if let Some(runtime) = runtime { + self.timers.lock().await.restore_runtime(runtime); + } + return Err(format!("failed to delete timer from sqlite: {err}")); + } + }; + let Some(runtime) = runtime else { + return Ok(deleted); + }; + if !deleted { + self.timers.lock().await.restore_runtime(runtime); return Ok(false); - }; - let Some(rollout_path) = self.current_rollout_path().await else { - self.timers.lock().await.restore_runtime(runtime); - return Err("timers require a persisted rollout path for restoration".to_string()); - }; - if let Err(err) = self.persist_timers_to_rollout_sidecar(&rollout_path).await { - self.timers.lock().await.restore_runtime(runtime); - return Err(err); } TimersState::cancel_runtime(&runtime); self.emit_timer_updated_notification().await; @@ -2776,8 +2867,6 @@ impl Session { if deleted_one_shot_timer { self.emit_timer_updated_notification().await; } - self.persist_timers_sidecar_best_effort().await; - let input_item = timer_prompt_input_item(&context); match context.delivery { TimerDelivery::SteerCurrentTurn => { @@ -2822,10 +2911,17 @@ impl Session { can_after_turn, active_turn_is_regular, ); - if claimed.is_none() { + let Some(claimed) = claimed else { *self.timer_start_in_progress.lock().await = false; + return None; + }; + + if !self.try_claim_timer_in_db(&claimed).await { + self.sync_timers_from_db(/*emit_update*/ true).await; + *self.timer_start_in_progress.lock().await = false; + return None; } - claimed + Some(claimed) } async fn inject_timer_into_active_turn(&self, item: ResponseInputItem) -> bool { @@ -2868,8 +2964,9 @@ impl Session { break; }; let changed = session.timers.lock().await.mark_timer_due(&id, Utc::now()); - if changed { - session.persist_timers_sidecar_best_effort().await; + if changed && !session.persist_timer_due_best_effort(&id).await { + session.sync_timers_from_db(/*emit_update*/ true).await; + continue; } session.maybe_start_pending_timer().await; let next_timer_spec = session @@ -2885,98 +2982,239 @@ impl Session { }); } - async fn persist_timers_to_rollout_sidecar(&self, rollout_path: &Path) -> Result<(), String> { - let rollout_path = rollout_path.to_path_buf(); - self.persist_timers_sidecar_with_writer(|timers| async move { - write_timer_sidecar(&rollout_path, &timers).await - }) - .await - } - - async fn persist_timers_sidecar_with_writer( + async fn thread_timer_create_params( &self, - write_sidecar: W, - ) -> Result<(), String> - where - W: FnOnce(Vec) -> Fut, - Fut: Future>, - { - let _write_lock = self.timer_sidecar_write_lock.lock().await; - let timers = self.timers.lock().await.persisted_timers(); - write_sidecar(timers).await + persisted_timer: &PersistedTimer, + source: &str, + ) -> Result { + let timer = &persisted_timer.timer; + Ok(codex_state::ThreadTimerCreateParams { + id: timer.id.clone(), + thread_id: self.thread_id_string(), + source: source.to_string(), + client_id: self.timer_client_id().await, + trigger_json: serde_json::to_string(&timer.trigger) + .map_err(|err| format!("failed to serialize timer trigger: {err}"))?, + prompt: timer.prompt.clone(), + delivery: timer.delivery.as_str().to_string(), + created_at: timer.created_at, + next_run_at: timer.next_run_at, + last_run_at: timer.last_run_at, + pending_run: persisted_timer.pending_run, + }) } - async fn persist_timers_sidecar_best_effort(&self) { - let Some(rollout_path) = self.current_rollout_path().await else { - return; + fn thread_timer_update_params( + &self, + persisted_timer: &PersistedTimer, + ) -> Result { + let timer = &persisted_timer.timer; + Ok(codex_state::ThreadTimerUpdateParams { + trigger_json: serde_json::to_string(&timer.trigger) + .map_err(|err| format!("failed to serialize timer trigger: {err}"))?, + delivery: timer.delivery.as_str().to_string(), + next_run_at: timer.next_run_at, + last_run_at: timer.last_run_at, + pending_run: persisted_timer.pending_run, + }) + } + + async fn persist_timer_due_best_effort(&self, id: &str) -> bool { + let state_db = match self.timer_state_db().await { + Ok(state_db) => state_db, + Err(err) => { + warn!("failed to persist due timer {id}: {err}"); + return false; + } }; - if let Err(err) = self.persist_timers_to_rollout_sidecar(&rollout_path).await { - warn!("failed to persist timers sidecar: {err}"); + let Some(persisted_timer) = self.timers.lock().await.persisted_timer(id) else { + return false; + }; + match state_db + .update_thread_timer_due( + &self.thread_id_string(), + id, + persisted_timer.timer.next_run_at, + ) + .await + { + Ok(updated) => updated, + Err(err) => { + warn!("failed to persist due timer {id}: {err}"); + false + } } } - async fn restore_timers_from_sidecar(self: &Arc) { + async fn try_claim_timer_in_db(&self, claimed: &ClaimedTimer) -> bool { + let state_db = match self.timer_state_db().await { + Ok(state_db) => state_db, + Err(err) => { + warn!( + "failed to claim timer {} in sqlite: {err}", + claimed.timer.id + ); + return true; + } + }; + let thread_id = self.thread_id_string(); + let result = if claimed.deleted_one_shot_timer { + state_db + .claim_one_shot_thread_timer(&thread_id, &claimed.timer.id) + .await + } else { + let persisted_timer = PersistedTimer { + timer: claimed.timer.clone(), + pending_run: claimed.timer.trigger.is_idle_recurring(), + }; + let Ok(params) = self.thread_timer_update_params(&persisted_timer) else { + return false; + }; + state_db + .claim_recurring_thread_timer(&thread_id, &claimed.timer.id, ¶ms) + .await + }; + match result { + Ok(claimed) => claimed, + Err(err) => { + warn!( + "failed to claim timer {} in sqlite: {err}", + claimed.timer.id + ); + false + } + } + } + + async fn restore_timers_from_db(self: &Arc) { if !self.features.enabled(Feature::TimerScheduler) { return; } - let Some(rollout_path) = self.current_rollout_path().await else { + let Ok(state_db) = self.timer_state_db().await else { return; }; - let persisted = match load_timer_sidecar(&rollout_path).await { - Ok(persisted) => persisted, - Err(err) => { - warn!("{err}"); - return; - } - }; - if persisted.is_empty() { - return; - } - - { - let _write_lock = self.timer_sidecar_write_lock.lock().await; - let mut normalized = Vec::::new(); - for persisted_timer in persisted { - let timer_cancel = CancellationToken::new(); - let timer_spec = { - let mut timers = self.timers.lock().await; - match timers.restore_timer( - persisted_timer.clone(), - Utc::now(), - Some(timer_cancel.clone()), - ) { - Ok(timer_spec) => { - normalized = timers.persisted_timers(); - timer_spec - } - Err(err) => { - warn!( - "skipping invalid persisted timer {}: {err}", - persisted_timer.timer.id - ); - continue; - } - } - }; - if let Some(timer_spec) = timer_spec { - self.spawn_timer_task( - persisted_timer.timer.id.clone(), - timer_spec, - timer_cancel, - ); - } - } - - if let Err(err) = write_timer_sidecar(&rollout_path, &normalized).await { - warn!("failed to rewrite normalized timer sidecar: {err}"); - } - } - self.emit_timer_updated_notification().await; + self.start_timer_db_sync_task(state_db); + self.sync_timers_from_db(/*emit_update*/ true).await; self.maybe_start_pending_timer().await; } + fn start_timer_db_sync_task(self: &Arc, state_db: state_db::StateDbHandle) { + if !self.features.enabled(Feature::TimerScheduler) { + return; + } + if self + .timer_db_sync_started + .swap(/*val*/ true, Ordering::SeqCst) + { + return; + } + let weak = Arc::downgrade(self); + let session_cancel = self.timer_tasks_cancellation_token.clone(); + tokio::spawn(async move { + let checker = match state_db.timer_data_version_checker().await { + Ok(checker) => checker, + Err(err) => { + warn!("failed to start timer db sync: {err}"); + if let Some(session) = weak.upgrade() { + session.timer_db_sync_started.store(false, Ordering::SeqCst); + } + return; + } + }; + let mut last_data_version = checker.data_version().await.ok(); + let mut last_full_refresh = tokio::time::Instant::now(); + let mut interval = tokio::time::interval(TIMER_DB_SYNC_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + tokio::select! { + _ = session_cancel.cancelled() => break, + _ = interval.tick() => {} + } + + let current_data_version = match checker.data_version().await { + Ok(version) => Some(version), + Err(err) => { + warn!("failed to poll timer db data_version: {err}"); + None + } + }; + let version_changed = + current_data_version.is_some() && current_data_version != last_data_version; + let max_refresh_elapsed = + last_full_refresh.elapsed() >= TIMER_DB_MAX_REFRESH_INTERVAL; + if !version_changed && !max_refresh_elapsed { + continue; + } + last_data_version = current_data_version.or(last_data_version); + let Some(session) = weak.upgrade() else { + break; + }; + session.sync_timers_from_db(/*emit_update*/ true).await; + session.maybe_start_pending_timer().await; + last_full_refresh = tokio::time::Instant::now(); + } + }); + } + + async fn sync_timers_from_db(self: &Arc, emit_update: bool) -> bool { + if !self.features.enabled(Feature::TimerScheduler) { + return false; + } + let Ok(state_db) = self.timer_state_db().await else { + return false; + }; + self.start_timer_db_sync_task(state_db.clone()); + let thread_id = self.thread_id_string(); + let db_timers = match state_db.list_thread_timers(&thread_id).await { + Ok(timers) => timers, + Err(err) => { + warn!("failed to load timers from sqlite for thread {thread_id}: {err}"); + return false; + } + }; + let persisted = db_timers + .into_iter() + .filter_map(db_timer_to_persisted_timer) + .collect::>(); + let (changed, restored_tasks) = self + .timers + .lock() + .await + .replace_timers_if_changed(persisted, Utc::now()); + self.spawn_restored_timer_tasks(restored_tasks); + if changed && emit_update { + self.emit_timer_updated_notification().await; + } + changed + } + + fn spawn_restored_timer_tasks(self: &Arc, restored_tasks: Vec) { + for RestoredTimerTask { + id, + timer_spec, + timer_cancel, + } in restored_tasks + { + self.spawn_timer_task(id, timer_spec, timer_cancel); + } + } + + fn thread_id_string(&self) -> String { + self.conversation_id.to_string() + } + + async fn timer_client_id(&self) -> String { + let state = self.state.lock().await; + state + .session_configuration + .app_server_client_name + .clone() + .unwrap_or_else(|| TIMER_CLIENT_ID_FALLBACK.to_string()) + } + async fn emit_timer_updated_notification(&self) { - let timers = self.list_timers().await; + let timers = self.list_timers_from_memory().await; let Ok(payload) = serde_json::to_string(&timers) else { warn!("failed to serialize timer update payload"); return; diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 7988a5a09a..b4ecc545b2 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -14,8 +14,6 @@ use crate::shell::default_user_shell; use crate::timers::ThreadTimerTrigger; use crate::timers::TimerDelivery; use crate::timers::TimersState; -use crate::timers::load_timer_sidecar; -use crate::timers::write_timer_sidecar; use crate::tools::format_exec_output_str; use codex_features::Features; @@ -120,8 +118,6 @@ use serde_json::json; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration as StdDuration; -use tempfile::TempDir; -use tokio::sync::Notify; #[path = "codex_tests_guardian.rs"] mod guardian_tests; @@ -2823,7 +2819,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), timer_start_in_progress: Mutex::new(false), - timer_sidecar_write_lock: Mutex::new(()), + timer_db_sync_started: AtomicBool::new(false), mailbox, mailbox_rx: Mutex::new(mailbox_rx), idle_pending_input: Mutex::new(Vec::new()), @@ -3668,7 +3664,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), timer_start_in_progress: Mutex::new(false), - timer_sidecar_write_lock: Mutex::new(()), + timer_db_sync_started: AtomicBool::new(false), mailbox, mailbox_rx: Mutex::new(mailbox_rx), idle_pending_input: Mutex::new(Vec::new()), @@ -3760,98 +3756,6 @@ async fn maybe_start_pending_timer_claims_only_one_timer_while_start_is_in_progr ); } -#[tokio::test] -async fn timer_sidecar_persistence_serializes_stale_and_newer_snapshots() { - let (session, _, _) = make_session_and_context_with_rx().await; - let tempdir = TempDir::new().expect("tempdir"); - let rollout_path = tempdir.path().join("rollout.jsonl"); - let now = chrono::Utc::now(); - { - let mut timers = session.timers.lock().await; - timers - .create_timer( - crate::timers::CreateTimer { - id: "timer-1".to_string(), - trigger: ThreadTimerTrigger::Delay { - seconds: 10, - repeat: Some(true), - }, - prompt: "first".to_string(), - delivery: TimerDelivery::AfterTurn, - now, - }, - /*timer_cancel*/ None, - ) - .expect("first timer should be created"); - } - - let first_snapshot_taken = Arc::new(Notify::new()); - let allow_first_write = Arc::new(Notify::new()); - let first_session = Arc::clone(&session); - let first_rollout_path = rollout_path.clone(); - let first_snapshot_taken_for_task = Arc::clone(&first_snapshot_taken); - let allow_first_write_for_task = Arc::clone(&allow_first_write); - let first_write = tokio::spawn(async move { - first_session - .persist_timers_sidecar_with_writer(|timers| async move { - first_snapshot_taken_for_task.notify_one(); - allow_first_write_for_task.notified().await; - write_timer_sidecar(&first_rollout_path, &timers).await - }) - .await - .expect("first persist should succeed"); - }); - first_snapshot_taken.notified().await; - - { - let mut timers = session.timers.lock().await; - let removed = timers - .remove_timer("timer-1") - .expect("first timer should be removable"); - TimersState::cancel_runtime(&removed); - timers - .create_timer( - crate::timers::CreateTimer { - id: "timer-2".to_string(), - trigger: ThreadTimerTrigger::Delay { - seconds: 10, - repeat: Some(true), - }, - prompt: "second".to_string(), - delivery: TimerDelivery::AfterTurn, - now, - }, - /*timer_cancel*/ None, - ) - .expect("second timer should be created"); - } - - let second_session = Arc::clone(&session); - let second_rollout_path = rollout_path.clone(); - let second_write = tokio::spawn(async move { - second_session - .persist_timers_sidecar_with_writer(|timers| async move { - write_timer_sidecar(&second_rollout_path, &timers).await - }) - .await - .expect("second persist should succeed"); - }); - - allow_first_write.notify_one(); - first_write - .await - .expect("first persist task should complete"); - second_write - .await - .expect("second persist task should complete"); - - let loaded = load_timer_sidecar(&rollout_path) - .await - .expect("timer sidecar should load"); - let persisted = session.timers.lock().await.persisted_timers(); - assert_eq!(loaded, persisted); -} - #[tokio::test] async fn refresh_mcp_servers_is_deferred_until_next_turn() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/timers.rs b/codex-rs/core/src/timers.rs index 3e77ddd9a3..d5185908a7 100644 --- a/codex-rs/core/src/timers.rs +++ b/codex-rs/core/src/timers.rs @@ -1,7 +1,7 @@ //! Persistent thread-local timer scheduling for follow-on turns and same-turn steer delivery. //! //! This module owns the in-memory timer registry, trigger evaluation, the user -//! message injected when a timer fires, and the JSON sidecar format used to +//! message injected when a timer fires, and the persistent state shape used to //! restore timers after a harness restart. use crate::timer_trigger::TimerTrigger; @@ -15,11 +15,7 @@ use codex_protocol::models::ResponseInputItem; use serde::Deserialize; use serde::Serialize; use std::collections::HashMap; -use std::io::ErrorKind; -use std::path::Path; -use std::path::PathBuf; use tokio_util::sync::CancellationToken; -use uuid::Uuid; pub const TIMER_UPDATED_BACKGROUND_EVENT_PREFIX: &str = "timer_updated:"; pub const TIMER_FIRED_BACKGROUND_EVENT_PREFIX: &str = "timer_fired:"; @@ -104,6 +100,13 @@ pub(crate) struct TimerTaskSpec { pub(crate) delay: std::time::Duration, } +#[derive(Debug)] +pub(crate) struct RestoredTimerTask { + pub(crate) id: String, + pub(crate) timer_spec: TimerTaskSpec, + pub(crate) timer_cancel: CancellationToken, +} + impl TimersState { pub(crate) fn list_timers(&self) -> Vec { let mut timers = self @@ -137,6 +140,48 @@ impl TimersState { timers } + pub(crate) fn persisted_timer(&self, id: &str) -> Option { + self.timers.get(id).map(|runtime| PersistedTimer { + timer: runtime.timer.clone(), + pending_run: runtime.pending_run, + }) + } + + pub(crate) fn replace_timers_if_changed( + &mut self, + persisted: Vec, + now: chrono::DateTime, + ) -> (bool, Vec) { + if self.persisted_timers() == persisted { + return (false, Vec::new()); + } + + for runtime in self.timers.values() { + Self::cancel_runtime(runtime); + } + self.timers.clear(); + + let mut restored_tasks = Vec::new(); + for persisted_timer in persisted { + let timer_cancel = CancellationToken::new(); + let timer_id = persisted_timer.timer.id.clone(); + match self.restore_timer(persisted_timer, now, Some(timer_cancel.clone())) { + Ok(Some(timer_spec)) => { + restored_tasks.push(RestoredTimerTask { + id: timer_id, + timer_spec, + timer_cancel, + }); + } + Ok(None) => {} + Err(err) => { + tracing::warn!("skipping invalid persisted timer {timer_id}: {err}"); + } + } + } + (true, restored_tasks) + } + pub(crate) fn create_timer( &mut self, create_timer: CreateTimer, @@ -379,113 +424,6 @@ fn render_timer_prompt_template(template: &str, timer: &TimerInvocationContext) .to_string() } -pub fn timer_sidecar_path_for_rollout(rollout_path: &Path) -> PathBuf { - PathBuf::from(format!("{}.timers.json", rollout_path.display())) -} - -pub(crate) async fn load_timer_sidecar(rollout_path: &Path) -> Result, String> { - let sidecar_path = timer_sidecar_path_for_rollout(rollout_path); - let bytes = match tokio::fs::read(&sidecar_path).await { - Ok(bytes) => bytes, - Err(err) if err.kind() == ErrorKind::NotFound => return Ok(Vec::new()), - Err(err) => { - return Err(format!( - "failed to read timer sidecar `{}`: {err}", - sidecar_path.display() - )); - } - }; - let raw_timers: Vec = serde_json::from_slice(&bytes).map_err(|err| { - format!( - "failed to parse timer sidecar `{}`: {err}", - sidecar_path.display() - ) - })?; - let mut timers = Vec::new(); - for raw_timer in raw_timers { - match serde_json::from_value::(raw_timer) { - Ok(timer) => timers.push(timer), - Err(err) => { - tracing::warn!( - "skipping invalid persisted timer from `{}`: {err}", - sidecar_path.display() - ); - } - } - } - Ok(timers) -} - -pub(crate) async fn write_timer_sidecar( - rollout_path: &Path, - timers: &[PersistedTimer], -) -> Result<(), String> { - let sidecar_path = timer_sidecar_path_for_rollout(rollout_path); - if timers.is_empty() { - match tokio::fs::remove_file(&sidecar_path).await { - Ok(()) => return Ok(()), - Err(err) if err.kind() == ErrorKind::NotFound => return Ok(()), - Err(err) => { - return Err(format!( - "failed to remove empty timer sidecar `{}`: {err}", - sidecar_path.display() - )); - } - } - } - - let bytes = serde_json::to_vec_pretty(timers).map_err(|err| { - format!( - "failed to serialize timer sidecar `{}`: {err}", - sidecar_path.display() - ) - })?; - let tmp_path = PathBuf::from(format!("{}.tmp-{}", sidecar_path.display(), Uuid::new_v4())); - tokio::fs::write(&tmp_path, bytes).await.map_err(|err| { - format!( - "failed to write temporary timer sidecar `{}`: {err}", - tmp_path.display() - ) - })?; - match tokio::fs::rename(&tmp_path, &sidecar_path).await { - Ok(()) => Ok(()), - Err(initial_error) => { - #[cfg(target_os = "windows")] - { - match tokio::fs::remove_file(&sidecar_path).await { - Ok(()) => { - tokio::fs::rename(&tmp_path, &sidecar_path) - .await - .map_err(|err| { - format!( - "failed to replace timer sidecar `{}` with `{}`: {err}", - sidecar_path.display(), - tmp_path.display() - ) - })?; - return Ok(()); - } - Err(err) if err.kind() == ErrorKind::NotFound => {} - Err(err) => { - let _ = tokio::fs::remove_file(&tmp_path).await; - return Err(format!( - "failed to remove existing timer sidecar `{}` before replace: {err}", - sidecar_path.display() - )); - } - } - } - - let _ = tokio::fs::remove_file(&tmp_path).await; - Err(format!( - "failed to atomically replace timer sidecar `{}` with `{}`: {initial_error}", - sidecar_path.display(), - tmp_path.display() - )) - } - } -} - #[cfg(test)] mod tests { use super::CreateTimer; @@ -495,17 +433,13 @@ mod tests { use super::TimerDelivery; use super::TimerInvocationContext; use super::TimersState; - use super::load_timer_sidecar; use super::timer_prompt_input_item; - use super::timer_sidecar_path_for_rollout; - use super::write_timer_sidecar; use crate::timer_trigger::TimerTrigger; use chrono::TimeZone; use chrono::Utc; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseInputItem; use pretty_assertions::assert_eq; - use tempfile::TempDir; const ZERO_SECONDS: u64 = 0; const TEN_SECONDS: u64 = 10; @@ -741,124 +675,4 @@ mod tests { } ); } - - #[tokio::test] - async fn timer_sidecar_round_trips_persisted_timers() { - let tempdir = TempDir::new().expect("tempdir"); - let rollout_path = tempdir.path().join("rollout.jsonl"); - let persisted = vec![PersistedTimer { - timer: super::ThreadTimer { - id: "timer-1".to_string(), - trigger: delay(ZERO_SECONDS, /*repeat*/ None), - prompt: "run tests".to_string(), - delivery: TimerDelivery::AfterTurn, - created_at: 1, - next_run_at: None, - last_run_at: None, - }, - pending_run: true, - }]; - - write_timer_sidecar(&rollout_path, &persisted) - .await - .expect("write sidecar"); - let loaded = load_timer_sidecar(&rollout_path) - .await - .expect("load sidecar"); - - assert_eq!(loaded, persisted); - assert_eq!( - timer_sidecar_path_for_rollout(&rollout_path), - tempdir.path().join("rollout.jsonl.timers.json") - ); - } - - #[tokio::test] - async fn timer_sidecar_overwrites_existing_file() { - let tempdir = TempDir::new().expect("tempdir"); - let rollout_path = tempdir.path().join("rollout.jsonl"); - let original = vec![PersistedTimer { - timer: super::ThreadTimer { - id: "timer-1".to_string(), - trigger: delay(ZERO_SECONDS, /*repeat*/ None), - prompt: "run tests".to_string(), - delivery: TimerDelivery::AfterTurn, - created_at: 1, - next_run_at: None, - last_run_at: None, - }, - pending_run: true, - }]; - let replacement = vec![PersistedTimer { - timer: super::ThreadTimer { - id: "timer-2".to_string(), - trigger: delay(SIXTY_SECONDS, Some(true)), - prompt: "run different tests".to_string(), - delivery: TimerDelivery::SteerCurrentTurn, - created_at: 2, - next_run_at: None, - last_run_at: Some(3), - }, - pending_run: false, - }]; - - write_timer_sidecar(&rollout_path, &original) - .await - .expect("write original sidecar"); - write_timer_sidecar(&rollout_path, &replacement) - .await - .expect("overwrite sidecar"); - - let loaded = load_timer_sidecar(&rollout_path) - .await - .expect("load overwritten sidecar"); - assert_eq!(loaded, replacement); - } - - #[tokio::test] - async fn timer_sidecar_skips_invalid_entries() { - let tempdir = TempDir::new().expect("tempdir"); - let rollout_path = tempdir.path().join("rollout.jsonl"); - let sidecar_path = timer_sidecar_path_for_rollout(&rollout_path); - tokio::fs::write( - &sidecar_path, - r#"[ - { "timer": { "id": "old", "cronExpression": "@after-turn" }, "pending_run": true }, - { - "timer": { - "id": "timer-1", - "trigger": { "kind": "delay", "seconds": 0, "repeat": null }, - "prompt": "run tests", - "delivery": "after-turn", - "created_at": 1, - "next_run_at": null, - "last_run_at": null - }, - "pending_run": true - } - ]"#, - ) - .await - .expect("write sidecar"); - - let loaded = load_timer_sidecar(&rollout_path) - .await - .expect("load sidecar"); - - assert_eq!( - loaded, - vec![PersistedTimer { - timer: super::ThreadTimer { - id: "timer-1".to_string(), - trigger: delay(ZERO_SECONDS, /*repeat*/ None), - prompt: "run tests".to_string(), - delivery: TimerDelivery::AfterTurn, - created_at: 1, - next_run_at: None, - last_run_at: None, - }, - pending_run: true, - }] - ); - } } diff --git a/codex-rs/core/tests/suite/timers.rs b/codex-rs/core/tests/suite/timers.rs index f105319a28..c96f2e9996 100644 --- a/codex-rs/core/tests/suite/timers.rs +++ b/codex-rs/core/tests/suite/timers.rs @@ -1,5 +1,6 @@ use anyhow::Result; use anyhow::anyhow; +use chrono::Utc; use codex_core::timers::TIMER_FIRED_BACKGROUND_EVENT_PREFIX; use codex_core::timers::ThreadTimer; use codex_core::timers::ThreadTimerTrigger; @@ -44,6 +45,10 @@ async fn assert_after_turn_timer_starts_and_emits_fired_event() -> Result<()> { .features .enable(Feature::TimerScheduler) .unwrap_or_else(|err| panic!("test config should allow feature update: {err}")); + config + .features + .enable(Feature::Sqlite) + .unwrap_or_else(|err| panic!("test config should allow feature update: {err}")); }); let test = builder.build(&server).await?; @@ -104,3 +109,129 @@ Do not expose scheduler internals unless they matter to the user. Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn create_timer_persists_source_and_client_metadata() -> Result<()> { + let server = start_mock_server().await; + let mut builder = test_codex().with_config(|config| { + config + .features + .enable(Feature::TimerScheduler) + .unwrap_or_else(|err| panic!("test config should allow feature update: {err}")); + config + .features + .enable(Feature::Sqlite) + .unwrap_or_else(|err| panic!("test config should allow feature update: {err}")); + }); + let test = builder.build(&server).await?; + + let created = test + .codex + .create_timer( + ThreadTimerTrigger::Delay { + seconds: 60, + repeat: Some(true), + }, + "run timer".to_string(), + TimerDelivery::AfterTurn, + ) + .await + .map_err(|err| anyhow!("{err}"))?; + + let db = test.codex.state_db().expect("state db enabled"); + let timers = db + .list_thread_timers(&test.session_configured.session_id.to_string()) + .await?; + + assert_eq!(timers.len(), 1); + assert_eq!(timers[0].id, created.id); + assert_eq!(timers[0].source, "agent"); + assert_eq!(timers[0].client_id, "codex-cli"); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn create_timer_lazily_opens_sqlite_for_ephemeral_thread() -> Result<()> { + let server = start_mock_server().await; + let mut builder = test_codex().with_config(|config| { + config.ephemeral = true; + config + .features + .enable(Feature::TimerScheduler) + .unwrap_or_else(|err| panic!("test config should allow feature update: {err}")); + }); + let test = builder.build(&server).await?; + + let created = test + .codex + .create_timer( + ThreadTimerTrigger::Delay { + seconds: 60, + repeat: Some(true), + }, + "run timer".to_string(), + TimerDelivery::AfterTurn, + ) + .await + .map_err(|err| anyhow!("{err}"))?; + + assert_eq!( + test.codex.list_timers().await, + vec![created], + "ephemeral threads should still open sqlite timer storage lazily" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn list_timers_discovers_externally_inserted_timer() -> Result<()> { + let server = start_mock_server().await; + let mut builder = test_codex().with_config(|config| { + config + .features + .enable(Feature::TimerScheduler) + .unwrap_or_else(|err| panic!("test config should allow feature update: {err}")); + config + .features + .enable(Feature::Sqlite) + .unwrap_or_else(|err| panic!("test config should allow feature update: {err}")); + }); + let test = builder.build(&server).await?; + let db = test.codex.state_db().expect("state db enabled"); + let created_at = Utc::now().timestamp(); + + db.create_thread_timer(&codex_state::ThreadTimerCreateParams { + id: "external-timer".to_string(), + thread_id: test.session_configured.session_id.to_string(), + source: "client".to_string(), + client_id: "external-client".to_string(), + trigger_json: r#"{"kind":"delay","seconds":60,"repeat":true}"#.to_string(), + prompt: "external timer".to_string(), + delivery: "after-turn".to_string(), + created_at, + next_run_at: Some(created_at + 60), + last_run_at: None, + pending_run: false, + }) + .await?; + + let timers = test.codex.list_timers().await; + + assert_eq!(timers.len(), 1); + assert_eq!(timers[0].id, "external-timer"); + assert_eq!( + timers[0].trigger, + ThreadTimerTrigger::Delay { + seconds: 60, + repeat: Some(true), + } + ); + assert_eq!(timers[0].prompt, "external timer"); + assert_eq!(timers[0].delivery, TimerDelivery::AfterTurn); + assert_eq!(timers[0].created_at, created_at); + assert_eq!(timers[0].last_run_at, None); + + Ok(()) +} diff --git a/codex-rs/rollout/src/lib.rs b/codex-rs/rollout/src/lib.rs index 3b5f7618a8..6acf2c3be4 100644 --- a/codex-rs/rollout/src/lib.rs +++ b/codex-rs/rollout/src/lib.rs @@ -11,7 +11,6 @@ pub(crate) mod policy; pub(crate) mod recorder; pub(crate) mod session_index; pub mod state_db; -mod timer_sidecar; pub(crate) mod default_client { pub use codex_login::default_client::*; diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index 869792fc0d..f89bcf58c8 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -17,7 +17,6 @@ use uuid::Uuid; use super::ARCHIVED_SESSIONS_SUBDIR; use super::SESSIONS_SUBDIR; -use super::timer_sidecar::thread_preview_from_timer_sidecar; use crate::protocol::EventMsg; use crate::state_db; use codex_file_search as file_search; @@ -714,12 +713,7 @@ async fn build_thread_item( { return None; } - let timer_preview = if summary.saw_user_event { - None - } else { - thread_preview_from_timer_sidecar(&path).await - }; - if summary.saw_session_meta && (summary.saw_user_event || timer_preview.is_some()) { + if summary.saw_session_meta && summary.saw_user_event { let HeadTailSummary { thread_id, first_user_message, @@ -742,7 +736,7 @@ async fn build_thread_item( return Some(ThreadItem { path, thread_id, - first_user_message: first_user_message.or(timer_preview), + first_user_message, cwd, git_branch, git_sha, diff --git a/codex-rs/rollout/src/timer_sidecar.rs b/codex-rs/rollout/src/timer_sidecar.rs deleted file mode 100644 index af1997b85d..0000000000 --- a/codex-rs/rollout/src/timer_sidecar.rs +++ /dev/null @@ -1,20 +0,0 @@ -//! Helpers for locating timer sidecars and surfacing scheduler-only threads in -//! session listings. - -use std::path::Path; -use std::path::PathBuf; - -const TIMER_THREAD_PREVIEW: &str = "(timer configured)"; - -pub(crate) fn timer_sidecar_path_for_rollout(rollout_path: &Path) -> PathBuf { - PathBuf::from(format!("{}.timers.json", rollout_path.display())) -} - -pub(crate) async fn thread_preview_from_timer_sidecar(rollout_path: &Path) -> Option { - let sidecar_path = timer_sidecar_path_for_rollout(rollout_path); - tokio::fs::try_exists(sidecar_path) - .await - .ok() - .filter(|exists| *exists) - .map(|_| TIMER_THREAD_PREVIEW.to_string()) -} diff --git a/codex-rs/state/migrations/0025_thread_timers.sql b/codex-rs/state/migrations/0025_thread_timers.sql new file mode 100644 index 0000000000..e77a3c5f74 --- /dev/null +++ b/codex-rs/state/migrations/0025_thread_timers.sql @@ -0,0 +1,22 @@ +CREATE TABLE thread_timers ( + id TEXT PRIMARY KEY, + thread_id TEXT NOT NULL, + source TEXT NOT NULL, + client_id TEXT NOT NULL, + trigger_json TEXT NOT NULL, + prompt TEXT NOT NULL, + delivery TEXT NOT NULL, + created_at INTEGER NOT NULL, + next_run_at INTEGER, + last_run_at INTEGER, + pending_run INTEGER NOT NULL +); + +CREATE INDEX idx_thread_timers_thread_created + ON thread_timers(thread_id, created_at, id); + +CREATE INDEX idx_thread_timers_thread_pending + ON thread_timers(thread_id, pending_run, created_at, id); + +CREATE INDEX idx_thread_timers_thread_next_run + ON thread_timers(thread_id, next_run_at); diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index efad7651f5..96c1a1ba6d 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -45,8 +45,12 @@ pub use model::Stage1OutputRef; pub use model::Stage1StartupClaimParams; pub use model::ThreadMetadata; pub use model::ThreadMetadataBuilder; +pub use model::ThreadTimer; +pub use model::ThreadTimerCreateParams; +pub use model::ThreadTimerUpdateParams; pub use model::ThreadsPage; pub use runtime::RemoteControlEnrollmentRecord; +pub use runtime::TimerDataVersionChecker; pub use runtime::logs_db_filename; pub use runtime::logs_db_path; pub use runtime::state_db_filename; @@ -58,7 +62,7 @@ pub const SQLITE_HOME_ENV: &str = "CODEX_SQLITE_HOME"; pub const LOGS_DB_FILENAME: &str = "logs"; pub const LOGS_DB_VERSION: u32 = 2; pub const STATE_DB_FILENAME: &str = "state"; -pub const STATE_DB_VERSION: u32 = 5; +pub const STATE_DB_VERSION: u32 = 6; /// Errors encountered during DB operations. Tags: [stage] pub const DB_ERROR_METRIC: &str = "codex.db.error"; diff --git a/codex-rs/state/src/model/mod.rs b/codex-rs/state/src/model/mod.rs index 39f0e800fc..f7fc8d283b 100644 --- a/codex-rs/state/src/model/mod.rs +++ b/codex-rs/state/src/model/mod.rs @@ -4,6 +4,7 @@ mod graph; mod log; mod memories; mod thread_metadata; +mod thread_timer; pub use agent_job::AgentJob; pub use agent_job::AgentJobCreateParams; @@ -32,6 +33,9 @@ pub use thread_metadata::SortKey; pub use thread_metadata::ThreadMetadata; pub use thread_metadata::ThreadMetadataBuilder; pub use thread_metadata::ThreadsPage; +pub use thread_timer::ThreadTimer; +pub use thread_timer::ThreadTimerCreateParams; +pub use thread_timer::ThreadTimerUpdateParams; pub(crate) use agent_job::AgentJobItemRow; pub(crate) use agent_job::AgentJobRow; @@ -40,3 +44,4 @@ pub(crate) use memories::stage1_output_ref_from_parts; pub(crate) use thread_metadata::ThreadRow; pub(crate) use thread_metadata::anchor_from_item; pub(crate) use thread_metadata::datetime_to_epoch_seconds; +pub(crate) use thread_timer::ThreadTimerRow; diff --git a/codex-rs/state/src/model/thread_timer.rs b/codex-rs/state/src/model/thread_timer.rs new file mode 100644 index 0000000000..fc59a256be --- /dev/null +++ b/codex-rs/state/src/model/thread_timer.rs @@ -0,0 +1,73 @@ +use sqlx::FromRow; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ThreadTimerCreateParams { + pub id: String, + pub thread_id: String, + pub source: String, + pub client_id: String, + pub trigger_json: String, + pub prompt: String, + pub delivery: String, + pub created_at: i64, + pub next_run_at: Option, + pub last_run_at: Option, + pub pending_run: bool, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ThreadTimerUpdateParams { + pub trigger_json: String, + pub delivery: String, + pub next_run_at: Option, + pub last_run_at: Option, + pub pending_run: bool, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ThreadTimer { + pub id: String, + pub thread_id: String, + pub source: String, + pub client_id: String, + pub trigger_json: String, + pub prompt: String, + pub delivery: String, + pub created_at: i64, + pub next_run_at: Option, + pub last_run_at: Option, + pub pending_run: bool, +} + +#[derive(Debug, FromRow)] +pub(crate) struct ThreadTimerRow { + pub id: String, + pub thread_id: String, + pub source: String, + pub client_id: String, + pub trigger_json: String, + pub prompt: String, + pub delivery: String, + pub created_at: i64, + pub next_run_at: Option, + pub last_run_at: Option, + pub pending_run: i64, +} + +impl From for ThreadTimer { + fn from(row: ThreadTimerRow) -> Self { + Self { + id: row.id, + thread_id: row.thread_id, + source: row.source, + client_id: row.client_id, + trigger_json: row.trigger_json, + prompt: row.prompt, + delivery: row.delivery, + created_at: row.created_at, + next_run_at: row.next_run_at, + last_run_at: row.last_run_at, + pending_run: row.pending_run != 0, + } + } +} diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index f71b6adf05..21bf493dcc 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -15,6 +15,9 @@ use crate::STATE_DB_VERSION; use crate::SortKey; use crate::ThreadMetadata; use crate::ThreadMetadataBuilder; +use crate::ThreadTimer; +use crate::ThreadTimerCreateParams; +use crate::ThreadTimerUpdateParams; use crate::ThreadsPage; use crate::apply_rollout_item; use crate::migrations::runtime_logs_migrator; @@ -58,8 +61,10 @@ mod remote_control; #[cfg(test)] mod test_support; mod threads; +mod timers; pub use remote_control::RemoteControlEnrollmentRecord; +pub use timers::TimerDataVersionChecker; // "Partition" is the retained-log-content bucket we cap at 10 MiB: // - one bucket per non-null thread_id diff --git a/codex-rs/state/src/runtime/timers.rs b/codex-rs/state/src/runtime/timers.rs new file mode 100644 index 0000000000..9c548403a7 --- /dev/null +++ b/codex-rs/state/src/runtime/timers.rs @@ -0,0 +1,371 @@ +use super::*; +use crate::model::ThreadTimerRow; +use tokio::sync::Mutex; + +pub struct TimerDataVersionChecker { + conn: Mutex, +} + +impl TimerDataVersionChecker { + pub async fn data_version(&self) -> anyhow::Result { + let mut conn = self.conn.lock().await; + let version = sqlx::query_scalar::<_, i64>("PRAGMA data_version") + .fetch_one(&mut *conn) + .await?; + Ok(version) + } +} + +impl StateRuntime { + pub async fn timer_data_version_checker(&self) -> anyhow::Result { + let state_path = state_db_path(self.codex_home()); + let options = base_sqlite_options(state_path.as_path()); + let conn = options.connect().await?; + Ok(TimerDataVersionChecker { + conn: Mutex::new(conn), + }) + } + + pub async fn create_thread_timer( + &self, + params: &ThreadTimerCreateParams, + ) -> anyhow::Result<()> { + sqlx::query( + r#" +INSERT INTO thread_timers ( + id, + thread_id, + source, + client_id, + trigger_json, + prompt, + delivery, + created_at, + next_run_at, + last_run_at, + pending_run +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + "#, + ) + .bind(params.id.as_str()) + .bind(params.thread_id.as_str()) + .bind(params.source.as_str()) + .bind(params.client_id.as_str()) + .bind(params.trigger_json.as_str()) + .bind(params.prompt.as_str()) + .bind(params.delivery.as_str()) + .bind(params.created_at) + .bind(params.next_run_at) + .bind(params.last_run_at) + .bind(i64::from(params.pending_run)) + .execute(self.pool.as_ref()) + .await?; + Ok(()) + } + + pub async fn list_thread_timers(&self, thread_id: &str) -> anyhow::Result> { + let rows = sqlx::query_as::<_, ThreadTimerRow>( + r#" +SELECT + id, + thread_id, + source, + client_id, + trigger_json, + prompt, + delivery, + created_at, + next_run_at, + last_run_at, + pending_run +FROM thread_timers +WHERE thread_id = ? +ORDER BY created_at ASC, id ASC + "#, + ) + .bind(thread_id) + .fetch_all(self.pool.as_ref()) + .await?; + Ok(rows.into_iter().map(ThreadTimer::from).collect()) + } + + pub async fn delete_thread_timer(&self, thread_id: &str, id: &str) -> anyhow::Result { + let result = sqlx::query("DELETE FROM thread_timers WHERE thread_id = ? AND id = ?") + .bind(thread_id) + .bind(id) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn update_thread_timer_due( + &self, + thread_id: &str, + id: &str, + next_run_at: Option, + ) -> anyhow::Result { + let result = sqlx::query( + r#" +UPDATE thread_timers +SET pending_run = 1, + next_run_at = ? +WHERE thread_id = ? + AND id = ? + "#, + ) + .bind(next_run_at) + .bind(thread_id) + .bind(id) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn claim_one_shot_thread_timer( + &self, + thread_id: &str, + id: &str, + ) -> anyhow::Result { + let result = sqlx::query( + r#" +DELETE FROM thread_timers +WHERE thread_id = ? + AND id = ? + AND pending_run = 1 + "#, + ) + .bind(thread_id) + .bind(id) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn claim_recurring_thread_timer( + &self, + thread_id: &str, + id: &str, + params: &ThreadTimerUpdateParams, + ) -> anyhow::Result { + let result = sqlx::query( + r#" +UPDATE thread_timers +SET trigger_json = ?, + delivery = ?, + next_run_at = ?, + last_run_at = ?, + pending_run = ? +WHERE thread_id = ? + AND id = ? + AND pending_run = 1 + "#, + ) + .bind(params.trigger_json.as_str()) + .bind(params.delivery.as_str()) + .bind(params.next_run_at) + .bind(params.last_run_at) + .bind(i64::from(params.pending_run)) + .bind(thread_id) + .bind(id) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } +} + +#[cfg(test)] +mod tests { + use super::StateRuntime; + use super::test_support::unique_temp_dir; + use crate::ThreadTimerCreateParams; + use crate::ThreadTimerUpdateParams; + use pretty_assertions::assert_eq; + + fn timer_params(id: &str, thread_id: &str) -> ThreadTimerCreateParams { + ThreadTimerCreateParams { + id: id.to_string(), + thread_id: thread_id.to_string(), + source: "agent".to_string(), + client_id: "codex-tui".to_string(), + trigger_json: r#"{"kind":"delay","seconds":10,"repeat":true}"#.to_string(), + prompt: "run tests".to_string(), + delivery: "after-turn".to_string(), + created_at: 100, + next_run_at: Some(110), + last_run_at: None, + pending_run: false, + } + } + + async fn test_runtime() -> std::sync::Arc { + StateRuntime::init(unique_temp_dir(), "test-provider".to_string()) + .await + .expect("initialize runtime") + } + + #[tokio::test] + async fn thread_timers_table_and_indexes_exist() { + let runtime = test_runtime().await; + let names = sqlx::query_scalar::<_, String>( + r#" +SELECT name +FROM sqlite_master +WHERE tbl_name = 'thread_timers' + AND name NOT LIKE 'sqlite_autoindex_%' +ORDER BY name + "#, + ) + .fetch_all(runtime.pool.as_ref()) + .await + .expect("query schema objects"); + + assert_eq!( + names, + vec![ + "idx_thread_timers_thread_created", + "idx_thread_timers_thread_next_run", + "idx_thread_timers_thread_pending", + "thread_timers", + ] + ); + } + + #[tokio::test] + async fn thread_timer_rows_round_trip_source_and_client_metadata() { + let runtime = test_runtime().await; + let mut params = timer_params("timer-1", "thread-1"); + params.pending_run = true; + params.last_run_at = Some(105); + + runtime + .create_thread_timer(¶ms) + .await + .expect("create timer"); + let timers = runtime + .list_thread_timers("thread-1") + .await + .expect("list timers"); + + assert_eq!(timers.len(), 1); + let timer = &timers[0]; + assert_eq!(timer.id, params.id); + assert_eq!(timer.thread_id, params.thread_id); + assert_eq!(timer.source, params.source); + assert_eq!(timer.client_id, params.client_id); + assert_eq!(timer.trigger_json, params.trigger_json); + assert_eq!(timer.prompt, params.prompt); + assert_eq!(timer.delivery, params.delivery); + assert_eq!(timer.created_at, params.created_at); + assert_eq!(timer.next_run_at, params.next_run_at); + assert_eq!(timer.last_run_at, params.last_run_at); + assert_eq!(timer.pending_run, params.pending_run); + } + + #[tokio::test] + async fn thread_timer_crud_is_scoped_to_thread_id() { + let runtime = test_runtime().await; + runtime + .create_thread_timer(&timer_params("timer-1", "thread-1")) + .await + .expect("create thread-1 timer"); + runtime + .create_thread_timer(&timer_params("timer-2", "thread-2")) + .await + .expect("create thread-2 timer"); + + assert_eq!( + runtime + .list_thread_timers("thread-1") + .await + .expect("list thread-1 timers") + .into_iter() + .map(|timer| timer.id) + .collect::>(), + vec!["timer-1".to_string()] + ); + assert!( + !runtime + .delete_thread_timer("thread-1", "timer-2") + .await + .expect("delete wrong thread timer") + ); + assert!( + runtime + .delete_thread_timer("thread-2", "timer-2") + .await + .expect("delete correct thread timer") + ); + } + + #[tokio::test] + async fn one_shot_claim_consumes_pending_timer_once() { + let runtime = test_runtime().await; + let mut params = timer_params("timer-1", "thread-1"); + params.pending_run = true; + params.next_run_at = None; + runtime + .create_thread_timer(¶ms) + .await + .expect("create pending timer"); + + assert!( + runtime + .claim_one_shot_thread_timer("thread-1", "timer-1") + .await + .expect("claim timer") + ); + assert!( + !runtime + .claim_one_shot_thread_timer("thread-1", "timer-1") + .await + .expect("claim timer again") + ); + assert!( + runtime + .list_thread_timers("thread-1") + .await + .expect("list timers") + .is_empty() + ); + } + + #[tokio::test] + async fn recurring_claim_updates_pending_timer_once() { + let runtime = test_runtime().await; + let mut params = timer_params("timer-1", "thread-1"); + params.pending_run = true; + runtime + .create_thread_timer(¶ms) + .await + .expect("create pending timer"); + let update = ThreadTimerUpdateParams { + trigger_json: params.trigger_json.clone(), + delivery: "steer-current-turn".to_string(), + next_run_at: Some(120), + last_run_at: Some(110), + pending_run: false, + }; + + assert!( + runtime + .claim_recurring_thread_timer("thread-1", "timer-1", &update) + .await + .expect("claim recurring timer") + ); + assert!( + !runtime + .claim_recurring_thread_timer("thread-1", "timer-1", &update) + .await + .expect("claim recurring timer again") + ); + let timers = runtime + .list_thread_timers("thread-1") + .await + .expect("list timers"); + assert_eq!(timers.len(), 1); + assert_eq!(timers[0].delivery, "steer-current-turn"); + assert_eq!(timers[0].next_run_at, Some(120)); + assert_eq!(timers[0].last_run_at, Some(110)); + assert!(!timers[0].pending_run); + } +}