klieo_runlog/

store_sqlite.rs

//! SQLite-backed `RunLogStore`. Single-table schema with `steps` stored as JSON.
//!
//! Available behind the `sqlite` feature.

use crate::error::RunLogError;
use crate::store::{RunLogQuery, RunLogStore};
use crate::types::{Cost, RunLog, RunStatus, Usage};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use klieo_core::ids::RunId;
use rusqlite::{params, Connection, OptionalExtension};
use std::path::Path;
use std::sync::Arc;
use tokio::sync::Mutex;

/// SQLite-backed `RunLogStore`.
pub struct SqliteRunLogStore {
    conn: Arc<Mutex<Connection>>,
}

impl SqliteRunLogStore {
    /// Open or create the database file at `path` and ensure the schema exists.
    ///
    /// On open, ensures both `cost_prompt_usd` and `cost_completion_usd`
    /// columns exist via `ALTER TABLE … ADD COLUMN` for forward-compatible
    /// upgrade of pre-existing databases. Both columns are nullable so rows
    /// written before the columns existed continue to read back as `cost_estimate = None`
    /// (the prompt/completion split is genuinely unknown for rows written
    /// before the split columns existed; fabricating `0.0` halves would
    /// violate the audit-evidence invariant `prompt + completion == total`).
    pub async fn new(path: impl AsRef<Path>) -> Result<Self, RunLogError> {
        let p = path.as_ref().to_path_buf();
        let conn = tokio::task::spawn_blocking(move || -> Result<Connection, RunLogError> {
            let c = Connection::open(&p).map_err(|e| RunLogError::Store(e.to_string()))?;
            // WAL allows concurrent readers alongside a single writer and
            // avoids the rollback-journal fsync per commit; `synchronous=NORMAL`
            // is the documented safe pairing with WAL (durable across app
            // crashes; loses at most the last committed transaction on OS
            // crash/power loss, which is acceptable for replayable run logs).
            c.pragma_update(None, "journal_mode", "WAL")
                .map_err(|e| RunLogError::Store(e.to_string()))?;
            c.pragma_update(None, "synchronous", "NORMAL")
                .map_err(|e| RunLogError::Store(e.to_string()))?;
            c.execute_batch(
                "CREATE TABLE IF NOT EXISTS runlog (
                    run_id TEXT PRIMARY KEY,
                    agent TEXT NOT NULL,
                    started_at TEXT NOT NULL,
                    finished_at TEXT,
                    status TEXT NOT NULL,
                    steps TEXT NOT NULL,
                    prompt_tokens INTEGER NOT NULL,
                    completion_tokens INTEGER NOT NULL,
                    cost_total_usd REAL,
                    cost_prompt_usd REAL,
                    cost_completion_usd REAL
                 );
                 CREATE INDEX IF NOT EXISTS runlog_agent_started
                   ON runlog (agent, started_at DESC);",
            )
            .map_err(|e| RunLogError::Store(e.to_string()))?;
            ensure_cost_split_columns(&c)?;
            Ok(c)
        })
        .await
        .map_err(|e| RunLogError::Store(e.to_string()))??;
        Ok(Self {
            conn: Arc::new(Mutex::new(conn)),
        })
    }
}

/// Add the `cost_prompt_usd` / `cost_completion_usd` columns to a pre-existing
/// `runlog` table when they are missing. Idempotent and safe to call on a
/// freshly created table.
fn ensure_cost_split_columns(c: &Connection) -> Result<(), RunLogError> {
    let mut have_prompt = false;
    let mut have_completion = false;
    {
        let mut stmt = c
            .prepare("PRAGMA table_info(runlog)")
            .map_err(|e| RunLogError::Store(e.to_string()))?;
        let rows = stmt
            .query_map([], |row| row.get::<_, String>(1))
            .map_err(|e| RunLogError::Store(e.to_string()))?;
        for r in rows {
            let name = r.map_err(|e| RunLogError::Store(e.to_string()))?;
            if name == "cost_prompt_usd" {
                have_prompt = true;
            } else if name == "cost_completion_usd" {
                have_completion = true;
            }
        }
    }
    if !have_prompt {
        add_column_if_missing(c, "ALTER TABLE runlog ADD COLUMN cost_prompt_usd REAL")?;
    }
    if !have_completion {
        add_column_if_missing(c, "ALTER TABLE runlog ADD COLUMN cost_completion_usd REAL")?;
    }
    Ok(())
}

/// Run an `ALTER TABLE … ADD COLUMN` that is safe under concurrent first-open.
///
/// SQLite serialises writes via the file lock, so two processes racing the
/// migration both call `PRAGMA table_info` simultaneously, see the column
/// missing, and both attempt the ALTER. The second's ALTER fails with
/// `SQLITE_ERROR (1)` and message containing `duplicate column name` — that
/// outcome is benign (the column is already there) and is swallowed here.
fn add_column_if_missing(c: &Connection, sql: &str) -> Result<(), RunLogError> {
    match c.execute(sql, []) {
        Ok(_) => Ok(()),
        Err(e) => {
            let msg = e.to_string();
            if msg.contains("duplicate column name") {
                Ok(())
            } else {
                Err(RunLogError::Store(msg))
            }
        }
    }
}

fn status_to_str(s: RunStatus) -> &'static str {
    match s {
        RunStatus::Running => "running",
        RunStatus::Completed => "completed",
        RunStatus::Failed => "failed",
        RunStatus::Cancelled => "cancelled",
    }
}

fn status_from_str(s: &str) -> Result<RunStatus, RunLogError> {
    match s {
        "running" => Ok(RunStatus::Running),
        "completed" => Ok(RunStatus::Completed),
        "failed" => Ok(RunStatus::Failed),
        "cancelled" => Ok(RunStatus::Cancelled),
        other => Err(RunLogError::Store(format!("unknown status {other}"))),
    }
}

/// Parse a `RunId` from its `Display` form. `RunId` does not implement
/// `FromStr` (Plan #11 trait freeze), so we parse via the underlying ULID.
fn parse_run_id(s: &str) -> Result<RunId, RunLogError> {
    let u = ulid::Ulid::from_string(s)
        .map_err(|e| RunLogError::Store(format!("bad run_id {s}: {e}")))?;
    Ok(RunId(u))
}

#[async_trait]
impl RunLogStore for SqliteRunLogStore {
    async fn put(&self, run_log: &RunLog) -> Result<(), RunLogError> {
        let conn = self.conn.clone();
        let log = run_log.clone();
        tokio::task::spawn_blocking(move || -> Result<(), RunLogError> {
            let steps_json =
                serde_json::to_string(&log.steps).map_err(|e| RunLogError::Store(e.to_string()))?;
            let started = log.started_at.to_rfc3339();
            let finished = log.finished_at.map(|d| d.to_rfc3339());
            let status = status_to_str(log.status);
            let cost_total = log.cost_estimate.map(|c| c.total_usd);
            let cost_prompt = log.cost_estimate.map(|c| c.prompt_usd);
            let cost_completion = log.cost_estimate.map(|c| c.completion_usd);
            let g = conn.blocking_lock();
            let mut stmt = g
                .prepare_cached(
                    "INSERT OR REPLACE INTO runlog
                     (run_id, agent, started_at, finished_at, status, steps,
                      prompt_tokens, completion_tokens, cost_total_usd,
                      cost_prompt_usd, cost_completion_usd)
                     VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)",
                )
                .map_err(|e| RunLogError::Store(e.to_string()))?;
            stmt.execute(params![
                log.run_id.to_string(),
                log.agent,
                started,
                finished,
                status,
                steps_json,
                log.tokens.prompt_tokens as i64,
                log.tokens.completion_tokens as i64,
                cost_total,
                cost_prompt,
                cost_completion,
            ])
            .map_err(|e| RunLogError::Store(e.to_string()))?;
            Ok(())
        })
        .await
        .map_err(|e| RunLogError::Store(e.to_string()))??;
        Ok(())
    }

    async fn get(&self, run_id: RunId) -> Result<Option<RunLog>, RunLogError> {
        let conn = self.conn.clone();
        let id = run_id.to_string();
        let row: Option<RunLog> =
            tokio::task::spawn_blocking(move || -> Result<Option<RunLog>, RunLogError> {
                let g = conn.blocking_lock();
                let mut stmt = g
                    .prepare_cached(
                        "SELECT run_id, agent, started_at, finished_at, status, steps,
                                prompt_tokens, completion_tokens, cost_total_usd,
                                cost_prompt_usd, cost_completion_usd
                         FROM runlog WHERE run_id = ?1",
                    )
                    .map_err(|e| RunLogError::Store(e.to_string()))?;
                let row = stmt
                    .query_row(params![id], row_to_runlog)
                    .optional()
                    .map_err(|e| RunLogError::Store(e.to_string()))?;
                row.transpose()
            })
            .await
            .map_err(|e| RunLogError::Store(e.to_string()))??;
        Ok(row)
    }

    async fn list(&self, query: &RunLogQuery) -> Result<Vec<RunLog>, RunLogError> {
        let conn = self.conn.clone();
        // Bind optional filters as NULL-able params so a single statement covers
        // every agent/status combination; the `?n IS NULL OR col = ?n` guard
        // skips the predicate when the filter is unset. Status is pushed into the
        // query so it composes with LIMIT/OFFSET (a post-filter would drop
        // matching rows that fell outside the page).
        let agent = query.agent.clone();
        let status = query.status.map(|s| status_to_str(s).to_string());
        // SQLite treats a negative LIMIT as "no limit". `None` -> -1; a `Some`
        // larger than i64::MAX clamps to i64::MAX (a real cap), never wrapping
        // negative into an accidental "no limit".
        let limit_i64 = query
            .limit
            .map(|n| i64::try_from(n).unwrap_or(i64::MAX))
            .unwrap_or(-1);
        let offset_i64 = query.offset as i64;
        let rows = tokio::task::spawn_blocking(move || -> Result<Vec<RunLog>, RunLogError> {
            let g = conn.blocking_lock();
            let mut stmt = g
                .prepare_cached(
                    "SELECT run_id, agent, started_at, finished_at, status, steps,
                            prompt_tokens, completion_tokens, cost_total_usd,
                            cost_prompt_usd, cost_completion_usd
                     FROM runlog
                     WHERE (?1 IS NULL OR agent = ?1)
                       AND (?2 IS NULL OR status = ?2)
                     ORDER BY started_at DESC LIMIT ?3 OFFSET ?4",
                )
                .map_err(|e| RunLogError::Store(e.to_string()))?;
            let mapped = stmt
                .query_map(params![agent, status, limit_i64, offset_i64], row_to_runlog)
                .map_err(|e| RunLogError::Store(e.to_string()))?;
            let out: Vec<Result<RunLog, RunLogError>> = mapped
                .map(|r| r.map_err(|e| RunLogError::Store(e.to_string())))
                .collect::<Result<Vec<_>, _>>()?;
            out.into_iter().collect::<Result<Vec<_>, _>>()
        })
        .await
        .map_err(|e| RunLogError::Store(e.to_string()))??;
        Ok(rows)
    }

    async fn delete(&self, run_id: RunId) -> Result<(), RunLogError> {
        let conn = self.conn.clone();
        let id = run_id.to_string();
        tokio::task::spawn_blocking(move || -> Result<(), RunLogError> {
            let g = conn.blocking_lock();
            let mut stmt = g
                .prepare_cached("DELETE FROM runlog WHERE run_id = ?1")
                .map_err(|e| RunLogError::Store(e.to_string()))?;
            stmt.execute(params![id])
                .map_err(|e| RunLogError::Store(e.to_string()))?;
            Ok(())
        })
        .await
        .map_err(|e| RunLogError::Store(e.to_string()))??;
        Ok(())
    }
}

fn row_to_runlog(row: &rusqlite::Row<'_>) -> rusqlite::Result<Result<RunLog, RunLogError>> {
    let run_id_s: String = row.get(0)?;
    let agent: String = row.get(1)?;
    let started_s: String = row.get(2)?;
    let finished_s: Option<String> = row.get(3)?;
    let status_s: String = row.get(4)?;
    let steps_s: String = row.get(5)?;
    let prompt_tokens: i64 = row.get(6)?;
    let completion_tokens: i64 = row.get(7)?;
    let cost_total: Option<f64> = row.get(8)?;
    // Legacy rows written before the split columns existed have cost_total
    // set but the split columns NULL. Audit evidence requires the invariant
    // `prompt + completion == total`, so we refuse to fabricate a split:
    // such rows surface as `cost_estimate = None` rather than producing an
    // internally inconsistent `Cost`. Current writers always populate all
    // three columns together.
    let cost_prompt: Option<f64> = row.get(9)?;
    let cost_completion: Option<f64> = row.get(10)?;

    Ok((|| -> Result<RunLog, RunLogError> {
        let run_id = parse_run_id(&run_id_s)?;
        let started_at: DateTime<Utc> = DateTime::parse_from_rfc3339(&started_s)
            .map_err(|e| RunLogError::Store(e.to_string()))?
            .with_timezone(&Utc);
        let finished_at = finished_s
            .map(|s| {
                DateTime::parse_from_rfc3339(&s)
                    .map(|d| d.with_timezone(&Utc))
                    .map_err(|e| RunLogError::Store(e.to_string()))
            })
            .transpose()?;
        let status = status_from_str(&status_s)?;
        let steps =
            serde_json::from_str(&steps_s).map_err(|e| RunLogError::Store(e.to_string()))?;
        Ok(RunLog {
            run_id,
            agent,
            started_at,
            finished_at,
            status,
            steps,
            tokens: Usage {
                prompt_tokens: prompt_tokens.max(0) as u32,
                completion_tokens: completion_tokens.max(0) as u32,
            },
            cost_estimate: match (cost_total, cost_prompt, cost_completion) {
                (Some(total), Some(prompt), Some(completion)) => Some(Cost {
                    prompt_usd: prompt,
                    completion_usd: completion,
                    // No cache column is persisted; the stored total remains
                    // authoritative (it already included any cache cost when
                    // computed), so cache_usd reconstructs as 0 rather than
                    // re-deriving a split that wasn't recorded.
                    cache_usd: 0.0,
                    total_usd: total,
                }),
                // Either nothing recorded, or a legacy row with only
                // cost_total set — refuse to fabricate a split.
                _ => None,
            },
        })
    })())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::projector::project;
    use tempfile::NamedTempFile;

    #[tokio::test]
    async fn legacy_row_with_only_total_returns_no_cost_estimate() {
        // Simulate a legacy row written before the split columns existed:
        // cost_total_usd set, both split columns NULL. The reader must
        // refuse to fabricate a split and surface cost_estimate = None.
        let f = NamedTempFile::new().unwrap();
        let store = SqliteRunLogStore::new(f.path()).await.unwrap();
        let log = project(RunId::new(), "writer", &[]);
        let rid = log.run_id;
        store.put(&log).await.unwrap();
        // Mutate the row to look like a legacy record: total set, split NULL.
        {
            let g = store.conn.lock().await;
            g.execute(
                "UPDATE runlog SET cost_total_usd = 0.05,
                                    cost_prompt_usd = NULL,
                                    cost_completion_usd = NULL
                 WHERE run_id = ?1",
                params![rid.to_string()],
            )
            .unwrap();
        }
        let back = store.get(rid).await.unwrap().expect("row exists");
        assert!(
            back.cost_estimate.is_none(),
            "legacy row with NULL split must surface as None, got {:?}",
            back.cost_estimate
        );
    }

    #[tokio::test]
    async fn ensure_cost_split_columns_is_idempotent_when_invoked_twice() {
        // Open the store twice on the same DB file. The second open re-runs
        // ensure_cost_split_columns; the columns already exist, so it must
        // be a silent no-op (and tolerate any duplicate-column race that
        // would otherwise trip ALTER).
        let f = NamedTempFile::new().unwrap();
        let _store1 = SqliteRunLogStore::new(f.path()).await.unwrap();
        // Second open must succeed without error.
        let store2 = SqliteRunLogStore::new(f.path()).await.unwrap();
        // And operations on it must round-trip cleanly.
        let log = project(RunId::new(), "writer", &[]);
        let rid = log.run_id;
        store2.put(&log).await.unwrap();
        assert!(store2.get(rid).await.unwrap().is_some());
    }
}