diff --git a/src/storage/error.rs b/src/storage/error.rs index b82e3a5..3c76d31 100644 --- a/src/storage/error.rs +++ b/src/storage/error.rs @@ -16,6 +16,9 @@ pub enum StorageError { #[error("schema migration error: {0}")] Migration(String), + + #[error("storage conflict: {0}")] + Conflict(String), } impl StorageError { diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 59741f1..811b66e 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -15,6 +15,29 @@ use std::path::Path; use tokio::time::{Duration, sleep}; const SCHEMA_VERSION: i64 = 1; +const INSERT_MESSAGE_SQL: &str = r#" + INSERT INTO messages (id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +"#; + +fn insert_message_query<'a>( + session_id: &'a str, + msg: &'a crate::storage::message::MessageMeta, +) -> sqlx::query::Query<'a, Sqlite, sqlx::sqlite::SqliteArguments<'a>> { + sqlx::query(INSERT_MESSAGE_SQL) + .bind(&msg.id) + .bind(session_id) + .bind(msg.seq) + .bind(&msg.role) + .bind(&msg.content) + .bind(&msg.reasoning_content) + .bind(&msg.media_refs) + .bind(&msg.tool_call_id) + .bind(&msg.tool_name) + .bind(&msg.tool_calls) + .bind(&msg.source) + .bind(msg.created_at) +} pub struct Storage { pub(crate) pool: Pool, @@ -164,6 +187,12 @@ impl Storage { .execute(&self.pool) .await?; + let memory_fts_exists: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'memory_fts')", + ) + .fetch_one(&self.pool) + .await?; + // FTS5 virtual table for full-text search on memories sqlx::query( r#" @@ -213,10 +242,14 @@ impl Storage { .execute(&self.pool) .await?; - // Rebuild FTS5 index for any existing records - sqlx::query("INSERT INTO memory_fts(memory_fts) VALUES ('rebuild')") - .execute(&self.pool) - .await?; + // Only a newly-created index needs a backfill. Triggers keep an + // existing index current, so rebuilding it on every startup is wasted + // work proportional to the total memory corpus. + if !memory_fts_exists { + sqlx::query("INSERT INTO memory_fts(memory_fts) VALUES ('rebuild')") + .execute(&self.pool) + .await?; + } sqlx::query( r#" @@ -632,64 +665,13 @@ impl Storage { session_id: &str, msg: &crate::storage::message::MessageMeta, ) -> Result { - sqlx::query( - r#" - INSERT INTO messages (id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(&msg.id) - .bind(session_id) - .bind(msg.seq) - .bind(&msg.role) - .bind(&msg.content) - .bind(&msg.reasoning_content) - .bind(&msg.media_refs) - .bind(&msg.tool_call_id) - .bind(&msg.tool_name) - .bind(&msg.tool_calls) - .bind(&msg.source) - .bind(msg.created_at) - .execute(self.pool()) - .await?; + insert_message_query(session_id, msg) + .execute(self.pool()) + .await?; Ok(msg.seq) } - pub async fn append_messages( - &self, - session_id: &str, - msgs: &[crate::storage::message::MessageMeta], - ) -> Result, StorageError> { - let mut tx = self.pool.begin().await?; - let mut seqs = Vec::with_capacity(msgs.len()); - for msg in msgs { - sqlx::query( - r#" - INSERT INTO messages (id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(&msg.id) - .bind(session_id) - .bind(msg.seq) - .bind(&msg.role) - .bind(&msg.content) - .bind(&msg.reasoning_content) - .bind(&msg.media_refs) - .bind(&msg.tool_call_id) - .bind(&msg.tool_name) - .bind(&msg.tool_calls) - .bind(&msg.source) - .bind(msg.created_at) - .execute(&mut *tx) - .await?; - seqs.push(msg.seq); - } - tx.commit().await?; - Ok(seqs) - } - /// Atomically persist all messages produced by one logical turn together /// with the resulting session metadata. A turn is either fully visible /// after restart or not visible at all. @@ -702,26 +684,9 @@ impl Storage { let mut tx = self.pool.begin().await?; for msg in msgs { - sqlx::query( - r#" - INSERT INTO messages (id, session_id, seq, role, content, reasoning_content, media_refs, tool_call_id, tool_name, tool_calls, source, created_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(&msg.id) - .bind(session_id) - .bind(msg.seq) - .bind(&msg.role) - .bind(&msg.content) - .bind(&msg.reasoning_content) - .bind(&msg.media_refs) - .bind(&msg.tool_call_id) - .bind(&msg.tool_name) - .bind(&msg.tool_calls) - .bind(&msg.source) - .bind(msg.created_at) - .execute(&mut *tx) - .await?; + insert_message_query(session_id, msg) + .execute(&mut *tx) + .await?; } sqlx::query( @@ -1257,6 +1222,30 @@ mod tests { assert!(orphan.is_err()); } + #[tokio::test] + async fn reopening_database_does_not_rebuild_existing_fts_index() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("fts.db"); + let storage = Storage::new(&db_path).await.unwrap(); + sqlx::query( + "INSERT INTO memory_fts(rowid, key, content) VALUES (999999, 'startup_sentinel', 'startup_sentinel')", + ) + .execute(storage.pool()) + .await + .unwrap(); + drop(storage); + + let reopened = Storage::new(&db_path).await.unwrap(); + let sentinel_count: i64 = sqlx::query_scalar( + "SELECT count(*) FROM memory_fts WHERE memory_fts MATCH 'startup_sentinel'", + ) + .fetch_one(reopened.pool()) + .await + .unwrap(); + + assert_eq!(sentinel_count, 1); + } + #[tokio::test] async fn background_task_completion_persists_execution_metrics() { let (storage, _dir) = create_test_storage().await; diff --git a/src/storage/scheduler.rs b/src/storage/scheduler.rs index 26805b1..464a5d5 100644 --- a/src/storage/scheduler.rs +++ b/src/storage/scheduler.rs @@ -2,6 +2,7 @@ use serde::{Deserialize, Serialize}; use sqlx::Row; use crate::scheduler::Schedule; +use crate::storage::StorageError; /// A scheduled job stored in the database. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -39,8 +40,8 @@ pub struct JobRun { impl crate::storage::Storage { /// Insert a new scheduled job. - pub async fn add_scheduled_job(&self, job: &ScheduledJob) -> anyhow::Result<()> { - let schedule_json = serde_json::to_string(&job.schedule)?; + pub async fn add_scheduled_job(&self, job: &ScheduledJob) -> Result<(), StorageError> { + let schedule_json = serialize_schedule(&job.schedule)?; sqlx::query( r#" INSERT INTO scheduled_jobs @@ -71,17 +72,17 @@ impl crate::storage::Storage { } /// Fetch a single scheduled job by ID. - pub async fn get_scheduled_job(&self, id: &str) -> anyhow::Result { + pub async fn get_scheduled_job(&self, id: &str) -> Result { let row = sqlx::query("SELECT * FROM scheduled_jobs WHERE id = ?") .bind(id) .fetch_optional(self.pool()) .await? - .ok_or_else(|| anyhow::anyhow!("job not found: {id}"))?; + .ok_or_else(|| StorageError::NotFound(format!("scheduled job {id}")))?; row_to_job(&row) } /// List all scheduled jobs, ordered by next_run_at ascending. - pub async fn list_scheduled_jobs(&self) -> anyhow::Result> { + pub async fn list_scheduled_jobs(&self) -> Result, StorageError> { let rows = sqlx::query("SELECT * FROM scheduled_jobs ORDER BY next_run_at ASC") .fetch_all(self.pool()) .await?; @@ -89,7 +90,7 @@ impl crate::storage::Storage { } /// Delete a scheduled job (cascades to job_runs). - pub async fn remove_scheduled_job(&self, id: &str) -> anyhow::Result<()> { + pub async fn remove_scheduled_job(&self, id: &str) -> Result<(), StorageError> { sqlx::query("DELETE FROM scheduled_jobs WHERE id = ?") .bind(id) .execute(self.pool()) @@ -98,7 +99,11 @@ impl crate::storage::Storage { } /// Enable or disable a scheduled job. - pub async fn set_scheduled_job_enabled(&self, id: &str, enabled: bool) -> anyhow::Result<()> { + pub async fn set_scheduled_job_enabled( + &self, + id: &str, + enabled: bool, + ) -> Result<(), StorageError> { sqlx::query("UPDATE scheduled_jobs SET enabled = ?, updated_at = ? WHERE id = ?") .bind(enabled as i32) .bind(now_ms()) @@ -117,7 +122,7 @@ impl crate::storage::Storage { channel: Option, chat_id: Option, model: Option, - ) -> anyhow::Result<()> { + ) -> Result<(), StorageError> { let now = now_ms(); if let Some(p) = prompt { @@ -129,7 +134,7 @@ impl crate::storage::Storage { .await?; } if let Some(s) = schedule { - let json = serde_json::to_string(&s)?; + let json = serialize_schedule(&s)?; sqlx::query("UPDATE scheduled_jobs SET schedule = ?, updated_at = ? WHERE id = ?") .bind(&json) .bind(now) @@ -169,7 +174,7 @@ impl crate::storage::Storage { &self, id: &str, next_run_at: i64, - ) -> anyhow::Result<()> { + ) -> Result<(), StorageError> { let now = now_ms(); sqlx::query( "UPDATE scheduled_jobs SET next_run_at = ?, last_run_at = ?, updated_at = ? WHERE id = ?", @@ -183,53 +188,6 @@ impl crate::storage::Storage { Ok(()) } - /// Set last_run_at for a job (used when starting execution). - pub async fn touch_scheduled_job_last_run(&self, id: &str, at: i64) -> anyhow::Result<()> { - sqlx::query("UPDATE scheduled_jobs SET last_run_at = ?, updated_at = ? WHERE id = ?") - .bind(at) - .bind(at) - .bind(id) - .execute(self.pool()) - .await?; - Ok(()) - } - - /// Set last_status and last_error after job completion. - pub async fn set_scheduled_job_last_status( - &self, - id: &str, - status: &str, - error: Option<&str>, - ) -> anyhow::Result<()> { - let now = now_ms(); - sqlx::query( - "UPDATE scheduled_jobs SET last_status = ?, last_error = ?, updated_at = ? WHERE id = ?", - ) - .bind(status) - .bind(error) - .bind(now) - .bind(id) - .execute(self.pool()) - .await?; - Ok(()) - } - - /// Fetch enabled jobs whose next_run_at <= now, up to `limit`. - pub async fn due_scheduled_jobs( - &self, - now: i64, - limit: usize, - ) -> anyhow::Result> { - let rows = sqlx::query( - "SELECT * FROM scheduled_jobs WHERE enabled = 1 AND next_run_at <= ? ORDER BY next_run_at ASC LIMIT ?", - ) - .bind(now) - .bind(i64::try_from(limit).unwrap_or(i64::MAX)) - .fetch_all(self.pool()) - .await?; - rows.iter().map(row_to_job).collect() - } - /// Atomically claim due jobs for one scheduler instance. A crashed worker's /// claims become eligible again after `lease_until`. pub async fn claim_due_scheduled_jobs( @@ -238,7 +196,7 @@ impl crate::storage::Storage { lease_until: i64, owner: &str, limit: usize, - ) -> anyhow::Result> { + ) -> Result, StorageError> { if limit == 0 { return Ok(Vec::new()); } @@ -282,7 +240,7 @@ impl crate::storage::Storage { next_run_at: Option, disable: bool, delete: bool, - ) -> anyhow::Result<()> { + ) -> Result<(), StorageError> { let mut tx = self.pool().begin().await?; if delete { @@ -292,7 +250,10 @@ impl crate::storage::Storage { .execute(&mut *tx) .await?; if result.rows_affected() != 1 { - anyhow::bail!("scheduled job lease lost before delete: {}", run.job_id); + return Err(StorageError::Conflict(format!( + "scheduled job lease lost before delete: {}", + run.job_id + ))); } } else { sqlx::query( @@ -332,7 +293,10 @@ impl crate::storage::Storage { .execute(&mut *tx) .await?; if result.rows_affected() != 1 { - anyhow::bail!("scheduled job lease lost before completion: {}", run.job_id); + return Err(StorageError::Conflict(format!( + "scheduled job lease lost before completion: {}", + run.job_id + ))); } } @@ -344,7 +308,7 @@ impl crate::storage::Storage { &self, job_id: &str, owner: &str, - ) -> anyhow::Result<()> { + ) -> Result<(), StorageError> { sqlx::query( "UPDATE scheduled_jobs SET locked_at = NULL, lock_owner = NULL, lease_until = NULL WHERE id = ? AND lock_owner = ?", ) @@ -355,32 +319,12 @@ impl crate::storage::Storage { Ok(()) } - /// Record a job execution run. - pub async fn record_scheduled_job_run(&self, run: &JobRun) -> anyhow::Result<()> { - sqlx::query( - r#" - INSERT INTO job_runs (job_id, started_at, finished_at, status, output, error, duration_ms) - VALUES (?, ?, ?, ?, ?, ?, ?) - "#, - ) - .bind(&run.job_id) - .bind(run.started_at) - .bind(run.finished_at) - .bind(&run.status) - .bind(&run.output) - .bind(&run.error) - .bind(run.duration_ms) - .execute(self.pool()) - .await?; - Ok(()) - } - /// List recent runs for a job, newest first. pub async fn list_scheduled_job_runs( &self, job_id: &str, limit: usize, - ) -> anyhow::Result> { + ) -> Result, StorageError> { let rows = sqlx::query( "SELECT * FROM job_runs WHERE job_id = ? ORDER BY finished_at DESC LIMIT ?", ) @@ -405,7 +349,7 @@ impl crate::storage::Storage { } /// Delete disabled jobs whose updated_at is before `before`. - pub async fn cleanup_disabled_scheduled_jobs(&self, before: i64) -> anyhow::Result<()> { + pub async fn cleanup_disabled_scheduled_jobs(&self, before: i64) -> Result<(), StorageError> { sqlx::query("DELETE FROM scheduled_jobs WHERE enabled = 0 AND updated_at < ?") .bind(before) .execute(self.pool()) @@ -421,9 +365,14 @@ fn now_ms() -> i64 { .as_millis() as i64 } -fn row_to_job(row: &sqlx::sqlite::SqliteRow) -> anyhow::Result { +fn serialize_schedule(schedule: &Schedule) -> Result { + serde_json::to_string(schedule).map_err(|error| StorageError::Serialization(error.to_string())) +} + +fn row_to_job(row: &sqlx::sqlite::SqliteRow) -> Result { let schedule_json: String = row.try_get("schedule")?; - let schedule: Schedule = serde_json::from_str(&schedule_json)?; + let schedule: Schedule = serde_json::from_str(&schedule_json) + .map_err(|error| StorageError::Serialization(error.to_string()))?; Ok(ScheduledJob { id: row.try_get("id")?, name: row.try_get("name")?, @@ -589,114 +538,6 @@ mod tests { assert!(!got.enabled); } - #[tokio::test] - async fn test_due_jobs_only_returns_enabled_and_overdue() { - let storage = setup_storage().await; - let t = now(); - let jobs = vec![ - ScheduledJob { - id: "due".into(), - name: "due".into(), - schedule: Schedule::At { at: t }, - prompt: "1".into(), - channel: "cli_chat".into(), - chat_id: "c".into(), - model: None, - enabled: true, - delete_after_run: false, - next_run_at: t - 1000, - last_run_at: None, - last_status: None, - last_error: None, - created_at: t, - updated_at: t, - }, - ScheduledJob { - id: "future".into(), - name: "future".into(), - schedule: Schedule::At { at: t + 99999999 }, - prompt: "2".into(), - channel: "cli_chat".into(), - chat_id: "c".into(), - model: None, - enabled: true, - delete_after_run: false, - next_run_at: t + 99999999, - last_run_at: None, - last_status: None, - last_error: None, - created_at: t, - updated_at: t, - }, - ScheduledJob { - id: "disabled-due".into(), - name: "disabled due".into(), - schedule: Schedule::At { at: t }, - prompt: "3".into(), - channel: "cli_chat".into(), - chat_id: "c".into(), - model: None, - enabled: false, - delete_after_run: false, - next_run_at: t - 1000, - last_run_at: None, - last_status: None, - last_error: None, - created_at: t, - updated_at: t, - }, - ]; - for j in &jobs { - storage.add_scheduled_job(j).await.unwrap(); - } - let due = storage.due_scheduled_jobs(t, 10).await.unwrap(); - assert_eq!(due.len(), 1); - assert_eq!(due[0].id, "due"); - } - - #[tokio::test] - async fn test_record_run_and_list_runs() { - let storage = setup_storage().await; - let t = now(); - let job = ScheduledJob { - id: "job-run".into(), - name: "run test".into(), - schedule: Schedule::Every { every_ms: 1000 }, - prompt: "hi".into(), - channel: "cli_chat".into(), - chat_id: "c".into(), - model: None, - enabled: true, - delete_after_run: false, - next_run_at: t, - last_run_at: None, - last_status: None, - last_error: None, - created_at: t, - updated_at: t, - }; - storage.add_scheduled_job(&job).await.unwrap(); - - let run = super::JobRun { - id: 0, - job_id: "job-run".into(), - started_at: t, - finished_at: t + 500, - status: "ok".into(), - output: Some("hello".into()), - error: None, - duration_ms: 500, - }; - storage.record_scheduled_job_run(&run).await.unwrap(); - let runs = storage - .list_scheduled_job_runs("job-run", 10) - .await - .unwrap(); - assert_eq!(runs.len(), 1); - assert_eq!(runs[0].status, "ok"); - assert_eq!(runs[0].output.as_deref(), Some("hello")); - } - #[tokio::test] async fn test_update_job() { let storage = setup_storage().await;