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
41pub use control::{ChatCostSummary, ControlMessage, MemoryEntry, parse_since_duration};
42pub use dependencies::{
43    BlockerSnapshot, DepSource, NewDependency, ObservedPullRequest, TaskDependency,
44    WaitingOnDependency, blocker_failure_kind, blocker_is_satisfied, blocker_snapshot_key,
45    dep_label, worktree_refresh_key,
46};
47pub use helpers::{
48    StoreResult, get_cost_estimate, get_cost_estimate_result, get_recent_memory,
49    get_recent_memory_result, get_task_field, get_task_field_direct, get_token_usage,
50    get_token_usage_result, opt_store_get_task, opt_store_get_task_by_id, resolve_store_id,
51    store_increment, store_increment_by_id, store_log_activity, store_reset_counters,
52    store_reset_failure_counters, store_set, store_set_by_id, store_set_result,
53    store_set_result_by_id, store_touch_updated_at,
54};
55pub use jobs::JobState;
56pub use metrics::{
57    AgentStat, CostByGroup, CostPeriod, CostSummary, ErrorStat, HighReviewCycleTask,
58    InsertTaskMetric, MetricsSummary, SlowTaskInfo,
59};
60pub use pricing::{CostEstimate, ModelPricing, TokenUsage, pricing_for_model};
61pub use tasks::{
62    CliTaskRef, CliTaskSpec, CompleteRun, NewTask, RunTokenUsage, StartRun, StoreRoute, Task,
63    TaskActivity, TaskRun, TaskStatus, UpsertExternal, parse_cli_task_spec,
64};
65
66/// Default database path: `~/.orch/orch.db`
67pub async fn default_db_path() -> anyhow::Result<std::path::PathBuf> {
68    crate::home::db_path().await
69}
70
71/// The unified task store backed by SQLite (via sqlx).
72#[derive(Clone)]
73pub struct TaskStore {
74    pool: SqlitePool,
75    /// Set once on the daemon. Without it, replicated writes apply directly, as in tests and the CLI
76    raft: Arc<OnceLock<crate::cluster::Node>>,
77    /// Set on the CLI when a daemon runs, so its writes go through the daemon's log
78    daemon: Arc<OnceLock<crate::cluster::local::Endpoint>>,
79}
80
81impl std::fmt::Debug for TaskStore {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        f.debug_struct("TaskStore").finish_non_exhaustive()
84    }
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    ///
116    /// Uses `max_connections(1)` to avoid WAL visibility issues where a write
117    /// on one pooled connection isn't visible to a read on another. Used by
118    /// integration tests that need read-your-own-writes consistency.
119    pub async fn open_single(db_path: &Path) -> anyhow::Result<Self> {
120        let options = SqliteConnectOptions::new()
121            .filename(db_path)
122            .journal_mode(SqliteJournalMode::Wal)
123            .create_if_missing(true)
124            .busy_timeout(std::time::Duration::from_secs(5))
125            .pragma("foreign_keys", "ON");
126
127        let pool = SqlitePoolOptions::new()
128            .max_connections(1)
129            .connect_with(options)
130            .await
131            .with_context(|| format!("opening task store (single): {}", db_path.display()))?;
132
133        let store = Self {
134            pool,
135            raft: Arc::default(),
136            daemon: Arc::default(),
137        };
138        store.migrate().await?;
139        Ok(store)
140    }
141
142    /// Open an in-memory store (for testing).
143    #[cfg(test)]
144    pub async fn open_memory() -> anyhow::Result<Self> {
145        let options = SqliteConnectOptions::new()
146            .filename(":memory:")
147            .pragma("foreign_keys", "ON");
148        let pool = SqlitePoolOptions::new()
149            .max_connections(1)
150            .connect_with(options)
151            .await?;
152
153        let store = Self {
154            pool,
155            raft: Arc::default(),
156            daemon: Arc::default(),
157        };
158        store.migrate().await?;
159        Ok(store)
160    }
161
162    /// Run migrations.
163    async fn migrate(&self) -> anyhow::Result<()> {
164        sqlx::migrate!("./migrations")
165            .run(&self.pool)
166            .await
167            .context("running task store migrations")?;
168        Ok(())
169    }
170
171    pub fn attach_node(&self, node: crate::cluster::Node) {
172        if self.raft.set(node).is_err() {
173            tracing::warn!("raft already attached to the task store");
174        }
175    }
176
177    /// True on the node that decides: the Raft leader, or a store with no Raft, as in the CLI and tests.
178    /// A single node always leads, so one machine runs every phase as before
179    pub fn is_leader(&self) -> bool {
180        self.raft.get().is_none_or(super::cluster::Node::is_leader)
181    }
182
183    /// Writes go through Raft or a daemon that runs it, so they reach the other machines
184    pub(crate) fn replicates(&self) -> bool {
185        self.raft.get().is_some() || self.daemon.get().is_some()
186    }
187
188    pub fn node(&self) -> Option<&crate::cluster::Node> {
189        self.raft.get()
190    }
191
192    pub fn attach_daemon(&self, endpoint: crate::cluster::local::Endpoint) {
193        let _ = self.daemon.set(endpoint);
194    }
195
196    /// Propose a replicated write through Raft, or apply it directly when no Raft is attached
197    pub(crate) async fn write(
198        &self,
199        cmd: crate::cluster::Command,
200    ) -> anyhow::Result<crate::cluster::Reply> {
201        let Some(node) = self.raft.get() else {
202            if let Some(endpoint) = self.daemon.get() {
203                match crate::cluster::local::send(endpoint, &cmd).await {
204                    Ok(reply) => return reply_result(reply),
205                    // Nothing was sent, so applying here cannot apply the command twice
206                    Err(e) if e.is_connect() => {
207                        tracing::debug!(err = %e, "daemon not listening, writing directly");
208                    }
209                    Err(e) => return Err(e.into()),
210                }
211            }
212
213            return self.apply_direct(cmd).await;
214        };
215        let req = crate::cluster::Request {
216            now: crate::cluster::now_stamp(),
217            cmd,
218        };
219        reply_result(node.write(req).await?)
220    }
221
222    pub(crate) async fn write_int(&self, cmd: crate::cluster::Command) -> anyhow::Result<i64> {
223        match self.write(cmd).await? {
224            crate::cluster::Reply::Int(n) => Ok(n),
225            other => anyhow::bail!("expected an integer reply, got {other:?}"),
226        }
227    }
228
229    pub(crate) async fn write_bool(&self, cmd: crate::cluster::Command) -> anyhow::Result<bool> {
230        match self.write(cmd).await? {
231            crate::cluster::Reply::Bool(b) => Ok(b),
232            other => anyhow::bail!("expected a boolean reply, got {other:?}"),
233        }
234    }
235
236    /// The same SQL apply runs, in one transaction on this node only
237    pub(crate) async fn apply_direct(
238        &self,
239        cmd: crate::cluster::Command,
240    ) -> anyhow::Result<crate::cluster::Reply> {
241        let mut attempt = 0;
242        loop {
243            match self.apply_once(cmd.clone()).await {
244                // A dropped transaction rolls back, so running the command again is safe
245                Err(e) if attempt < WRITE_LOCK_RETRIES && is_write_lock_error(&e) => {
246                    attempt += 1;
247                    tracing::warn!(attempt, err = %e, "database is locked, retrying the write");
248                    tokio::time::sleep(std::time::Duration::from_millis(100 << (attempt - 1)))
249                        .await;
250                }
251                result => return result,
252            }
253        }
254    }
255
256    async fn apply_once(
257        &self,
258        cmd: crate::cluster::Command,
259    ) -> anyhow::Result<crate::cluster::Reply> {
260        let mut tx = self.begin_write().await?;
261        let reply = cmd
262            .apply(&mut tx, &crate::cluster::now_stamp(), false)
263            .await?;
264        tx.commit().await?;
265        Ok(reply)
266    }
267
268    /// Start a transaction that holds the write lock from its first statement.
269    ///
270    /// A deferred `BEGIN` that reads and then writes fails at once with `database is locked` when
271    /// another connection commits in between, and `busy_timeout` cannot wait that out. `BEGIN
272    /// IMMEDIATE` waits for the lock up front, where `busy_timeout` applies
273    pub(crate) async fn begin_write(
274        &self,
275    ) -> sqlx::Result<sqlx::Transaction<'static, sqlx::Sqlite>> {
276        self.pool.begin_with("BEGIN IMMEDIATE").await
277    }
278
279    /// Get a reference to the underlying pool (for advanced queries).
280    pub fn pool(&self) -> &SqlitePool {
281        &self.pool
282    }
283}
284
285/// Retries for a write that still finds the database locked after `busy_timeout`
286const WRITE_LOCK_RETRIES: u32 = 4;
287
288/// `SQLITE_BUSY` (5) or `SQLITE_BUSY_SNAPSHOT` (517) anywhere in the error chain
289fn is_write_lock_error(err: &anyhow::Error) -> bool {
290    err.chain().any(|cause| {
291        cause
292            .downcast_ref::<sqlx::Error>()
293            .and_then(|e| e.as_database_error())
294            .and_then(sqlx::error::DatabaseError::code)
295            .is_some_and(|code| code == "5" || code == "517")
296    })
297}
298
299fn reply_result(reply: crate::cluster::Reply) -> anyhow::Result<crate::cluster::Reply> {
300    match reply {
301        crate::cluster::Reply::Err(e) => Err(anyhow::anyhow!(e)),
302        other => Ok(other),
303    }
304}
305
306#[cfg(test)]
307mod tests;
orch/store/mod.rs

orch/store/
mod.rs
