orch/store/
mod.rs

1//! Unified task store — SQLite as single source of truth.
2//!
3//! `TaskStore` replaces the legacy `internal_tasks` table and GitHub labels as the
4//! authoritative task state. External backends (GitHub, Linear) become sync adapters
5//! that mirror state changes outward.
6//!
7//! Uses sqlx for async SQLite access with file-based migrations.
8//!
9//! All task state, metrics, KV, and rate limits flow through this module.
10//!
11//! # Safety Warning: Never Use `SELECT *`
12//!
13//! Using `SELECT *` with `sqlx-sqlite`'s prepared statement cache can cause **OOB
14//! (out-of-bounds) panics** when the schema changes (e.g., new migration adds a column).
15//!
16//! The issue occurs because:
17//! 1. SQLite caches prepared statements with column metadata at prepare time
18//! 2. When a migration adds a new column, `sqlite3_column_count()` returns the updated count
19//! 3. The cached statement's column metadata vector was built with the old count
20//! 4. Accessing column index N when only N columns exist (0..N-1) causes a panic
21//!
22//! **Solution:** Always use explicit column lists via constants like `TASK_COLS`,
23//! `CONTROL_MESSAGE_COLS`, `JOB_STATE_COLS`, etc. See `tasks.rs` and `control.rs`
24//! for examples. CI will fail if `SELECT *` is detected in store files.
25
26use anyhow::Context;
27use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions};
28use std::path::Path;
29use std::sync::{Arc, OnceLock};
30
31pub(crate) mod control;
32pub(crate) mod dependencies;
33mod helpers;
34pub(crate) mod jobs;
35pub(crate) mod kv;
36mod metrics;
37mod pricing;
38pub(crate) mod replica;
39pub(crate) mod tasks;
40
41#[allow(unused_imports)]
42pub use control::{parse_since_duration, ChatCostSummary, ControlMessage, MemoryEntry};
43#[allow(unused_imports)]
44pub use dependencies::{
45    blocker_failure_kind, blocker_is_satisfied, blocker_snapshot_key, dep_label,
46    worktree_refresh_key, BlockerSnapshot, DepSource, NewDependency, ObservedPullRequest,
47    TaskDependency, WaitingOnDependency,
48};
49#[allow(unused_imports)]
50pub use helpers::{
51    get_cost_estimate, get_cost_estimate_result, get_recent_memory, get_recent_memory_result,
52    get_task_field, get_task_field_direct, get_token_usage, get_token_usage_result,
53    opt_store_get_task, opt_store_get_task_by_id, resolve_store_id, store_increment,
54    store_increment_by_id, store_log_activity, store_reset_counters, store_reset_failure_counters,
55    store_set, store_set_by_id, store_set_result, store_set_result_by_id, store_touch_updated_at,
56    StoreResult,
57};
58#[allow(unused_imports)]
59pub use jobs::JobState;
60#[allow(unused_imports)]
61pub use metrics::{
62    AgentStat, CostByGroup, CostPeriod, CostSummary, ErrorStat, HighReviewCycleTask,
63    InsertTaskMetric, MetricsSummary, SlowTaskInfo,
64};
65#[allow(unused_imports)]
66pub use pricing::{pricing_for_model, CostEstimate, ModelPricing, TokenUsage};
67pub use tasks::{
68    parse_cli_task_spec, CliTaskRef, CliTaskSpec, CompleteRun, NewTask, RunTokenUsage, StartRun,
69    StoreRoute, Task, TaskActivity, TaskRun, TaskStatus, UpsertExternal,
70};
71
72/// Default database path: `~/.orch/orch.db`
73pub async fn default_db_path() -> anyhow::Result<std::path::PathBuf> {
74    crate::home::db_path().await
75}
76
77/// The unified task store backed by SQLite (via sqlx).
78#[derive(Clone)]
79pub struct TaskStore {
80    pool: SqlitePool,
81    /// Set once on the daemon. Without it, replicated writes apply directly, as in tests and the CLI
82    raft: Arc<OnceLock<crate::cluster::Node>>,
83    /// Set on the CLI when a daemon runs, so its writes go through the daemon's log
84    daemon: Arc<OnceLock<crate::cluster::local::Endpoint>>,
85}
86
87impl TaskStore {
88    /// Open (or create) the store at the given database path.
89    ///
90    /// Runs file-based migrations from the `migrations/` directory.
91    pub async fn open(db_path: &Path) -> anyhow::Result<Self> {
92        let options = SqliteConnectOptions::new()
93            .filename(db_path)
94            .journal_mode(SqliteJournalMode::Wal)
95            .create_if_missing(true)
96            .busy_timeout(std::time::Duration::from_secs(5))
97            .pragma("foreign_keys", "ON");
98
99        let pool = SqlitePoolOptions::new()
100            .max_connections(5)
101            .connect_with(options)
102            .await
103            .with_context(|| format!("opening task store: {}", db_path.display()))?;
104
105        let store = Self {
106            pool,
107            raft: Arc::default(),
108            daemon: Arc::default(),
109        };
110        store.migrate().await?;
111        Ok(store)
112    }
113
114    /// Open a store with a single connection.
115    #[allow(dead_code)]
116    ///
117    /// Uses `max_connections(1)` to avoid WAL visibility issues where a write
118    /// on one pooled connection isn't visible to a read on another. Used by
119    /// integration tests that need read-your-own-writes consistency.
120    pub async fn open_single(db_path: &Path) -> anyhow::Result<Self> {
121        let options = SqliteConnectOptions::new()
122            .filename(db_path)
123            .journal_mode(SqliteJournalMode::Wal)
124            .create_if_missing(true)
125            .busy_timeout(std::time::Duration::from_secs(5))
126            .pragma("foreign_keys", "ON");
127
128        let pool = SqlitePoolOptions::new()
129            .max_connections(1)
130            .connect_with(options)
131            .await
132            .with_context(|| format!("opening task store (single): {}", db_path.display()))?;
133
134        let store = Self {
135            pool,
136            raft: Arc::default(),
137            daemon: Arc::default(),
138        };
139        store.migrate().await?;
140        Ok(store)
141    }
142
143    /// Open an in-memory store (for testing).
144    #[cfg(test)]
145    pub async fn open_memory() -> anyhow::Result<Self> {
146        let options = SqliteConnectOptions::new()
147            .filename(":memory:")
148            .pragma("foreign_keys", "ON");
149        let pool = SqlitePoolOptions::new()
150            .max_connections(1)
151            .connect_with(options)
152            .await?;
153
154        let store = Self {
155            pool,
156            raft: Arc::default(),
157            daemon: Arc::default(),
158        };
159        store.migrate().await?;
160        Ok(store)
161    }
162
163    /// Run migrations.
164    async fn migrate(&self) -> anyhow::Result<()> {
165        sqlx::migrate!("./migrations")
166            .run(&self.pool)
167            .await
168            .context("running task store migrations")?;
169        Ok(())
170    }
171
172    pub fn attach_node(&self, node: crate::cluster::Node) {
173        if self.raft.set(node).is_err() {
174            tracing::warn!("raft already attached to the task store");
175        }
176    }
177
178    /// True on the node that decides: the Raft leader, or a store with no Raft, as in the CLI and tests.
179    /// A single node always leads, so one machine runs every phase as before
180    pub fn is_leader(&self) -> bool {
181        self.raft.get().is_none_or(|node| node.is_leader())
182    }
183
184    /// Writes go through Raft or a daemon that runs it, so they reach the other machines
185    pub(crate) fn replicates(&self) -> bool {
186        self.raft.get().is_some() || self.daemon.get().is_some()
187    }
188
189    pub fn node(&self) -> Option<&crate::cluster::Node> {
190        self.raft.get()
191    }
192
193    pub fn attach_daemon(&self, endpoint: crate::cluster::local::Endpoint) {
194        let _ = self.daemon.set(endpoint);
195    }
196
197    /// Propose a replicated write through Raft, or apply it directly when no Raft is attached
198    pub(crate) async fn write(
199        &self,
200        cmd: crate::cluster::Command,
201    ) -> anyhow::Result<crate::cluster::Reply> {
202        let Some(node) = self.raft.get() else {
203            if let Some(endpoint) = self.daemon.get() {
204                match crate::cluster::local::send(endpoint, &cmd).await {
205                    Ok(reply) => return reply_result(reply),
206                    // Nothing was sent, so applying here cannot apply the command twice
207                    Err(e) if e.is_connect() => {
208                        tracing::debug!(err = %e, "daemon not listening, writing directly");
209                    }
210                    Err(e) => return Err(e.into()),
211                }
212            }
213
214            return self.apply_direct(cmd).await;
215        };
216        let req = crate::cluster::Request {
217            now: crate::cluster::now_stamp(),
218            cmd,
219        };
220        reply_result(node.write(req).await?)
221    }
222
223    pub(crate) async fn write_int(&self, cmd: crate::cluster::Command) -> anyhow::Result<i64> {
224        match self.write(cmd).await? {
225            crate::cluster::Reply::Int(n) => Ok(n),
226            other => anyhow::bail!("expected an integer reply, got {other:?}"),
227        }
228    }
229
230    pub(crate) async fn write_bool(&self, cmd: crate::cluster::Command) -> anyhow::Result<bool> {
231        match self.write(cmd).await? {
232            crate::cluster::Reply::Bool(b) => Ok(b),
233            other => anyhow::bail!("expected a boolean reply, got {other:?}"),
234        }
235    }
236
237    /// The same SQL apply runs, in one transaction on this node only
238    pub(crate) async fn apply_direct(
239        &self,
240        cmd: crate::cluster::Command,
241    ) -> anyhow::Result<crate::cluster::Reply> {
242        let mut attempt = 0;
243        loop {
244            match self.apply_once(cmd.clone()).await {
245                // A dropped transaction rolls back, so running the command again is safe
246                Err(e) if attempt < WRITE_LOCK_RETRIES && is_write_lock_error(&e) => {
247                    attempt += 1;
248                    tracing::warn!(attempt, err = %e, "database is locked, retrying the write");
249                    tokio::time::sleep(std::time::Duration::from_millis(100 << (attempt - 1)))
250                        .await;
251                }
252                result => return result,
253            }
254        }
255    }
256
257    async fn apply_once(
258        &self,
259        cmd: crate::cluster::Command,
260    ) -> anyhow::Result<crate::cluster::Reply> {
261        let mut tx = self.begin_write().await?;
262        let reply = cmd
263            .apply(&mut tx, &crate::cluster::now_stamp(), false)
264            .await?;
265        tx.commit().await?;
266        Ok(reply)
267    }
268
269    /// Start a transaction that holds the write lock from its first statement.
270    ///
271    /// A deferred `BEGIN` that reads and then writes fails at once with `database is locked` when
272    /// another connection commits in between, and `busy_timeout` cannot wait that out. `BEGIN
273    /// IMMEDIATE` waits for the lock up front, where `busy_timeout` applies
274    pub(crate) async fn begin_write(
275        &self,
276    ) -> sqlx::Result<sqlx::Transaction<'static, sqlx::Sqlite>> {
277        self.pool.begin_with("BEGIN IMMEDIATE").await
278    }
279
280    /// Get a reference to the underlying pool (for advanced queries).
281    #[allow(dead_code)]
282    pub fn pool(&self) -> &SqlitePool {
283        &self.pool
284    }
285}
286
287/// Retries for a write that still finds the database locked after `busy_timeout`
288const WRITE_LOCK_RETRIES: u32 = 4;
289
290/// `SQLITE_BUSY` (5) or `SQLITE_BUSY_SNAPSHOT` (517) anywhere in the error chain
291fn is_write_lock_error(err: &anyhow::Error) -> bool {
292    err.chain().any(|cause| {
293        cause
294            .downcast_ref::<sqlx::Error>()
295            .and_then(|e| e.as_database_error())
296            .and_then(|e| e.code())
297            .is_some_and(|code| code == "5" || code == "517")
298    })
299}
300
301fn reply_result(reply: crate::cluster::Reply) -> anyhow::Result<crate::cluster::Reply> {
302    match reply {
303        crate::cluster::Reply::Err(e) => Err(anyhow::anyhow!(e)),
304        other => Ok(other),
305    }
306}
307
308#[cfg(test)]
309mod tests;
orch/store/mod.rs

orch/store/
mod.rs
