1use anyhow::Context;
27use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions};
28use std::path::Path;
29
30mod control;
31mod dependencies;
32mod helpers;
33mod jobs;
34mod kv;
35mod metrics;
36mod pricing;
37mod tasks;
38
39#[allow(unused_imports)]
40pub use control::{parse_since_duration, ChatCostSummary, ControlMessage, MemoryEntry};
41#[allow(unused_imports)]
42pub use dependencies::{
43 blocker_failure_kind, blocker_is_satisfied, blocker_snapshot_key, dep_label,
44 worktree_refresh_key, BlockerSnapshot, DepSource, NewDependency, ObservedPullRequest,
45 TaskDependency, WaitingOnDependency,
46};
47#[allow(unused_imports)]
48pub use helpers::{
49 get_cost_estimate, get_cost_estimate_result, get_recent_memory, get_recent_memory_result,
50 get_task_field, get_task_field_direct, get_token_usage, get_token_usage_result,
51 opt_store_get_task, opt_store_get_task_by_id, resolve_store_id, review_session_expected,
52 set_review_session_expected, store_increment, store_increment_by_id, store_log_activity,
53 store_reset_counters, store_reset_failure_counters, store_set, store_set_by_id,
54 store_set_result, store_set_result_by_id, store_touch_updated_at, store_touch_updated_at_by_id,
55 StoreResult,
56};
57#[allow(unused_imports)]
58pub use jobs::JobState;
59#[allow(unused_imports)]
60pub use metrics::{
61 AgentStat, CostByGroup, CostPeriod, CostSummary, ErrorStat, HighReviewCycleTask,
62 InsertTaskMetric, MetricsSummary, SlowTaskInfo,
63};
64#[allow(unused_imports)]
65pub use pricing::{pricing_for_model, CostEstimate, ModelPricing, TokenUsage};
66pub use tasks::{
67 parse_cli_task_spec, CliTaskRef, CliTaskSpec, CompleteRun, NewTask, RunTokenUsage, StartRun,
68 StoreRoute, Task, TaskActivity, TaskRun, TaskStatus, UpsertExternal,
69};
70
71pub async fn default_db_path() -> anyhow::Result<std::path::PathBuf> {
73 crate::home::db_path().await
74}
75
76#[derive(Clone)]
78pub struct TaskStore {
79 pool: SqlitePool,
80}
81
82impl TaskStore {
83 pub async fn open(db_path: &Path) -> anyhow::Result<Self> {
87 let options = SqliteConnectOptions::new()
88 .filename(db_path)
89 .journal_mode(SqliteJournalMode::Wal)
90 .create_if_missing(true)
91 .busy_timeout(std::time::Duration::from_secs(5))
92 .pragma("foreign_keys", "ON");
93
94 let pool = SqlitePoolOptions::new()
95 .max_connections(5)
96 .connect_with(options)
97 .await
98 .with_context(|| format!("opening task store: {}", db_path.display()))?;
99
100 let store = Self { pool };
101 store.migrate().await?;
102 Ok(store)
103 }
104
105 #[allow(dead_code)]
107 pub async fn open_single(db_path: &Path) -> anyhow::Result<Self> {
112 let options = SqliteConnectOptions::new()
113 .filename(db_path)
114 .journal_mode(SqliteJournalMode::Wal)
115 .create_if_missing(true)
116 .busy_timeout(std::time::Duration::from_secs(5))
117 .pragma("foreign_keys", "ON");
118
119 let pool = SqlitePoolOptions::new()
120 .max_connections(1)
121 .connect_with(options)
122 .await
123 .with_context(|| format!("opening task store (single): {}", db_path.display()))?;
124
125 let store = Self { pool };
126 store.migrate().await?;
127 Ok(store)
128 }
129
130 #[cfg(test)]
132 pub async fn open_memory() -> anyhow::Result<Self> {
133 let options = SqliteConnectOptions::new()
134 .filename(":memory:")
135 .pragma("foreign_keys", "ON");
136 let pool = SqlitePoolOptions::new()
137 .max_connections(1)
138 .connect_with(options)
139 .await?;
140
141 let store = Self { pool };
142 store.migrate().await?;
143 Ok(store)
144 }
145
146 async fn migrate(&self) -> anyhow::Result<()> {
148 sqlx::migrate!("./migrations")
149 .run(&self.pool)
150 .await
151 .context("running task store migrations")?;
152 Ok(())
153 }
154
155 #[allow(dead_code)]
157 pub fn pool(&self) -> &SqlitePool {
158 &self.pool
159 }
160}
161
162#[cfg(test)]
163mod tests;