use anyhow::Context;
use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions};
use std::path::Path;
use std::sync::{Arc, OnceLock};
pub(crate) mod control;
pub(crate) mod dependencies;
mod helpers;
pub(crate) mod jobs;
pub(crate) mod kv;
mod metrics;
mod pricing;
pub(crate) mod replica;
pub(crate) mod tasks;
#[allow(unused_imports)]
pub use control::{parse_since_duration, ChatCostSummary, ControlMessage, MemoryEntry};
#[allow(unused_imports)]
pub use dependencies::{
blocker_failure_kind, blocker_is_satisfied, blocker_snapshot_key, dep_label,
worktree_refresh_key, BlockerSnapshot, DepSource, NewDependency, ObservedPullRequest,
TaskDependency, WaitingOnDependency,
};
#[allow(unused_imports)]
pub use helpers::{
get_cost_estimate, get_cost_estimate_result, get_recent_memory, get_recent_memory_result,
get_task_field, get_task_field_direct, get_token_usage, get_token_usage_result,
opt_store_get_task, opt_store_get_task_by_id, resolve_store_id, review_session_expected,
set_review_session_expected, store_increment, store_increment_by_id, store_log_activity,
store_reset_counters, store_reset_failure_counters, store_set, store_set_by_id,
store_set_result, store_set_result_by_id, store_touch_updated_at, StoreResult,
};
#[allow(unused_imports)]
pub use jobs::JobState;
#[allow(unused_imports)]
pub use metrics::{
AgentStat, CostByGroup, CostPeriod, CostSummary, ErrorStat, HighReviewCycleTask,
InsertTaskMetric, MetricsSummary, SlowTaskInfo,
};
#[allow(unused_imports)]
pub use pricing::{pricing_for_model, CostEstimate, ModelPricing, TokenUsage};
pub use tasks::{
parse_cli_task_spec, CliTaskRef, CliTaskSpec, CompleteRun, NewTask, RunTokenUsage, StartRun,
StoreRoute, Task, TaskActivity, TaskRun, TaskStatus, UpsertExternal,
};
pub async fn default_db_path() -> anyhow::Result<std::path::PathBuf> {
crate::home::db_path().await
}
#[derive(Clone)]
pub struct TaskStore {
pool: SqlitePool,
raft: Arc<OnceLock<crate::cluster::Node>>,
daemon: Arc<OnceLock<crate::cluster::local::Endpoint>>,
}
impl TaskStore {
pub async fn open(db_path: &Path) -> anyhow::Result<Self> {
let options = SqliteConnectOptions::new()
.filename(db_path)
.journal_mode(SqliteJournalMode::Wal)
.create_if_missing(true)
.busy_timeout(std::time::Duration::from_secs(5))
.pragma("foreign_keys", "ON");
let pool = SqlitePoolOptions::new()
.max_connections(5)
.connect_with(options)
.await
.with_context(|| format!("opening task store: {}", db_path.display()))?;
let store = Self {
pool,
raft: Arc::default(),
daemon: Arc::default(),
};
store.migrate().await?;
Ok(store)
}
#[allow(dead_code)]
pub async fn open_single(db_path: &Path) -> anyhow::Result<Self> {
let options = SqliteConnectOptions::new()
.filename(db_path)
.journal_mode(SqliteJournalMode::Wal)
.create_if_missing(true)
.busy_timeout(std::time::Duration::from_secs(5))
.pragma("foreign_keys", "ON");
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect_with(options)
.await
.with_context(|| format!("opening task store (single): {}", db_path.display()))?;
let store = Self {
pool,
raft: Arc::default(),
daemon: Arc::default(),
};
store.migrate().await?;
Ok(store)
}
#[cfg(test)]
pub async fn open_memory() -> anyhow::Result<Self> {
let options = SqliteConnectOptions::new()
.filename(":memory:")
.pragma("foreign_keys", "ON");
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect_with(options)
.await?;
let store = Self {
pool,
raft: Arc::default(),
daemon: Arc::default(),
};
store.migrate().await?;
Ok(store)
}
async fn migrate(&self) -> anyhow::Result<()> {
sqlx::migrate!("./migrations")
.run(&self.pool)
.await
.context("running task store migrations")?;
Ok(())
}
pub fn attach_node(&self, node: crate::cluster::Node) {
if self.raft.set(node).is_err() {
tracing::warn!("raft already attached to the task store");
}
}
pub fn is_leader(&self) -> bool {
self.raft.get().is_none_or(|node| node.is_leader())
}
pub(crate) fn replicates(&self) -> bool {
self.raft.get().is_some() || self.daemon.get().is_some()
}
pub fn node(&self) -> Option<&crate::cluster::Node> {
self.raft.get()
}
pub fn attach_daemon(&self, endpoint: crate::cluster::local::Endpoint) {
let _ = self.daemon.set(endpoint);
}
pub(crate) async fn write(
&self,
cmd: crate::cluster::Command,
) -> anyhow::Result<crate::cluster::Reply> {
let Some(node) = self.raft.get() else {
if let Some(endpoint) = self.daemon.get() {
match crate::cluster::local::send(endpoint, &cmd).await {
Ok(reply) => return reply_result(reply),
Err(e) if e.is_connect() => {
tracing::debug!(err = %e, "daemon not listening, writing directly");
}
Err(e) => return Err(e.into()),
}
}
return self.apply_direct(cmd).await;
};
let req = crate::cluster::Request {
now: crate::cluster::now_stamp(),
cmd,
};
reply_result(node.write(req).await?)
}
pub(crate) async fn write_int(&self, cmd: crate::cluster::Command) -> anyhow::Result<i64> {
match self.write(cmd).await? {
crate::cluster::Reply::Int(n) => Ok(n),
other => anyhow::bail!("expected an integer reply, got {other:?}"),
}
}
pub(crate) async fn write_bool(&self, cmd: crate::cluster::Command) -> anyhow::Result<bool> {
match self.write(cmd).await? {
crate::cluster::Reply::Bool(b) => Ok(b),
other => anyhow::bail!("expected a boolean reply, got {other:?}"),
}
}
pub(crate) async fn apply_direct(
&self,
cmd: crate::cluster::Command,
) -> anyhow::Result<crate::cluster::Reply> {
let mut tx = self.pool.begin().await?;
let reply = cmd
.apply(&mut tx, &crate::cluster::now_stamp(), false)
.await?;
tx.commit().await?;
Ok(reply)
}
#[allow(dead_code)]
pub fn pool(&self) -> &SqlitePool {
&self.pool
}
}
fn reply_result(reply: crate::cluster::Reply) -> anyhow::Result<crate::cluster::Reply> {
match reply {
crate::cluster::Reply::Err(e) => Err(anyhow::anyhow!(e)),
other => Ok(other),
}
}
#[cfg(test)]
mod tests;