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

orch/store/
mod.rs
