huskies/server/src/crdt_state.rs

/// CRDT state layer for pipeline state, backed by SQLite.
///
/// Replaces the filesystem as the primary source of truth for pipeline item
/// metadata (stage, name, agent, etc.).  CRDT ops are persisted to SQLite so
/// state survives restarts.  The filesystem `.huskies/work/` directories are
/// still updated as a secondary output for backwards compatibility.
use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};

use bft_json_crdt::json_crdt::*;
use bft_json_crdt::keypair::make_keypair;
use bft_json_crdt::list_crdt::ListCrdt;
use bft_json_crdt::lww_crdt::LwwRegisterCrdt;
use bft_json_crdt::op::ROOT_ID;
use fastcrypto::ed25519::Ed25519KeyPair;
use fastcrypto::traits::ToFromBytes;
use serde_json::json;
use sqlx::sqlite::SqliteConnectOptions;
use sqlx::SqlitePool;
use std::path::Path;
use tokio::sync::mpsc;

use crate::slog;

// ── CRDT document types ──────────────────────────────────────────────

#[add_crdt_fields]
#[derive(Clone, CrdtNode, Debug)]
pub struct PipelineDoc {
    pub items: ListCrdt<PipelineItemCrdt>,
}

#[add_crdt_fields]
#[derive(Clone, CrdtNode, Debug)]
pub struct PipelineItemCrdt {
    pub story_id: LwwRegisterCrdt<String>,
    pub stage: LwwRegisterCrdt<String>,
    pub name: LwwRegisterCrdt<String>,
    pub agent: LwwRegisterCrdt<String>,
    pub retry_count: LwwRegisterCrdt<f64>,
    pub blocked: LwwRegisterCrdt<bool>,
    pub depends_on: LwwRegisterCrdt<String>,
}

// ── Read-side view types ─────────────────────────────────────────────

/// A snapshot of a single pipeline item derived from the CRDT document.
#[derive(Clone, Debug)]
pub struct PipelineItemView {
    pub story_id: String,
    pub stage: String,
    pub name: Option<String>,
    pub agent: Option<String>,
    pub retry_count: Option<i64>,
    pub blocked: Option<bool>,
    pub depends_on: Option<Vec<u32>>,
}

// ── Internal state ───────────────────────────────────────────────────

struct CrdtState {
    crdt: BaseCrdt<PipelineDoc>,
    keypair: Ed25519KeyPair,
    /// Maps story_id → index in the ListCrdt for O(1) lookup.
    index: HashMap<String, usize>,
    /// Channel sender for fire-and-forget op persistence.
    persist_tx: mpsc::UnboundedSender<SignedOp>,
}

static CRDT_STATE: OnceLock<Mutex<CrdtState>> = OnceLock::new();

// ── Initialisation ───────────────────────────────────────────────────

/// Initialise the CRDT state layer.
///
/// Opens the SQLite database, loads or creates a node keypair, replays any
/// persisted ops to reconstruct state, and spawns a background persistence
/// task.  Safe to call only once; subsequent calls are no-ops.
pub async fn init(db_path: &Path) -> Result<(), sqlx::Error> {
    if CRDT_STATE.get().is_some() {
        return Ok(());
    }

    let options = SqliteConnectOptions::new()
        .filename(db_path)
        .create_if_missing(true);
    let pool = SqlitePool::connect_with(options).await?;
    sqlx::migrate!("./migrations").run(&pool).await?;

    // Load or create the node keypair.
    let keypair = load_or_create_keypair(&pool).await?;
    let mut crdt = BaseCrdt::<PipelineDoc>::new(&keypair);

    // Replay persisted ops to reconstruct state.
    let rows: Vec<(String,)> =
        sqlx::query_as("SELECT op_json FROM crdt_ops ORDER BY seq ASC")
            .fetch_all(&pool)
            .await?;

    for (op_json,) in &rows {
        if let Ok(signed_op) = serde_json::from_str::<SignedOp>(op_json) {
            crdt.apply(signed_op);
        } else {
            slog!("[crdt] Warning: failed to deserialize stored op");
        }
    }

    // Build the index from the reconstructed state.
    let index = rebuild_index(&crdt);

    slog!(
        "[crdt] Initialised: {} ops replayed, {} items indexed",
        rows.len(),
        index.len()
    );

    // Spawn background persistence task.
    let (persist_tx, mut persist_rx) = mpsc::unbounded_channel::<SignedOp>();

    tokio::spawn(async move {
        while let Some(op) = persist_rx.recv().await {
            let op_json = match serde_json::to_string(&op) {
                Ok(j) => j,
                Err(e) => {
                    slog!("[crdt] Failed to serialize op: {e}");
                    continue;
                }
            };
            let op_id = hex::encode(&op.id());
            let seq = op.inner.seq as i64;
            let now = chrono::Utc::now().to_rfc3339();

            let result = sqlx::query(
                "INSERT INTO crdt_ops (op_id, seq, op_json, created_at) \
                 VALUES (?1, ?2, ?3, ?4) \
                 ON CONFLICT(op_id) DO NOTHING",
            )
            .bind(&op_id)
            .bind(seq)
            .bind(&op_json)
            .bind(&now)
            .execute(&pool)
            .await;

            if let Err(e) = result {
                slog!("[crdt] Failed to persist op {}: {e}", &op_id[..12]);
            }
        }
    });

    let state = CrdtState {
        crdt,
        keypair,
        index,
        persist_tx,
    };

    let _ = CRDT_STATE.set(Mutex::new(state));
    Ok(())
}

/// Load or create the Ed25519 keypair used by this node.
async fn load_or_create_keypair(pool: &SqlitePool) -> Result<Ed25519KeyPair, sqlx::Error> {
    let row: Option<(Vec<u8>,)> =
        sqlx::query_as("SELECT seed FROM crdt_node_identity WHERE id = 1")
            .fetch_optional(pool)
            .await?;

    if let Some((seed,)) = row {
        // Reconstruct from stored seed.  The seed is the 32-byte private key.
        if let Ok(kp) = Ed25519KeyPair::from_bytes(&seed) {
            return Ok(kp);
        }
        slog!("[crdt] Stored keypair invalid, regenerating");
    }

    let kp = make_keypair();
    let seed = kp.as_bytes().to_vec();
    sqlx::query("INSERT INTO crdt_node_identity (id, seed) VALUES (1, ?1) ON CONFLICT(id) DO UPDATE SET seed = excluded.seed")
        .bind(&seed)
        .execute(pool)
        .await?;

    Ok(kp)
}

/// Rebuild the story_id → list index mapping from the current CRDT state.
fn rebuild_index(crdt: &BaseCrdt<PipelineDoc>) -> HashMap<String, usize> {
    let mut map = HashMap::new();
    for (i, item) in crdt.doc.items.iter().enumerate() {
        if let JsonValue::String(ref sid) = item.story_id.view() {
            map.insert(sid.clone(), i);
        }
    }
    map
}

// ── Write path ───────────────────────────────────────────────────────

/// Create a CRDT op via `op_fn`, sign it, apply it, and send it to the
/// persistence channel.  The closure receives `&mut CrdtState` so it can
/// mutably access the CRDT document, while `sign` only needs `&keypair`.
fn apply_and_persist<F>(state: &mut CrdtState, op_fn: F)
where
    F: FnOnce(&mut CrdtState) -> bft_json_crdt::op::Op<JsonValue>,
{
    let raw_op = op_fn(state);
    let signed = raw_op.sign(&state.keypair);
    state.crdt.apply(signed.clone());
    let _ = state.persist_tx.send(signed);
}

/// Write a pipeline item state through CRDT operations.
///
/// If the item exists, updates its registers.  If not, inserts a new item
/// into the list.  All ops are signed and persisted to SQLite.
pub fn write_item(
    story_id: &str,
    stage: &str,
    name: Option<&str>,
    agent: Option<&str>,
    retry_count: Option<i64>,
    blocked: Option<bool>,
    depends_on: Option<&str>,
) {
    let Some(state_mutex) = CRDT_STATE.get() else {
        return;
    };
    let Ok(mut state) = state_mutex.lock() else {
        return;
    };

    if let Some(&idx) = state.index.get(story_id) {
        // Update existing item registers.
        // Each op is created, signed, applied, and persisted in a block so
        // borrows do not overlap between &mut crdt (set) and &keypair (sign).
        apply_and_persist(&mut state, |s| {
            s.crdt.doc.items[idx].stage.set(stage.to_string())
        });

        if let Some(n) = name {
            apply_and_persist(&mut state, |s| {
                s.crdt.doc.items[idx].name.set(n.to_string())
            });
        }
        if let Some(a) = agent {
            apply_and_persist(&mut state, |s| {
                s.crdt.doc.items[idx].agent.set(a.to_string())
            });
        }
        if let Some(rc) = retry_count {
            apply_and_persist(&mut state, |s| {
                s.crdt.doc.items[idx].retry_count.set(rc as f64)
            });
        }
        if let Some(b) = blocked {
            apply_and_persist(&mut state, |s| {
                s.crdt.doc.items[idx].blocked.set(b)
            });
        }
        if let Some(d) = depends_on {
            apply_and_persist(&mut state, |s| {
                s.crdt.doc.items[idx].depends_on.set(d.to_string())
            });
        }
    } else {
        // Insert new item.
        let item_json: JsonValue = json!({
            "story_id": story_id,
            "stage": stage,
            "name": name.unwrap_or(""),
            "agent": agent.unwrap_or(""),
            "retry_count": retry_count.unwrap_or(0) as f64,
            "blocked": blocked.unwrap_or(false),
            "depends_on": depends_on.unwrap_or(""),
        })
        .into();

        apply_and_persist(&mut state, |s| {
            s.crdt.doc.items.insert(ROOT_ID, item_json)
        });

        // Rebuild index after insertion (indices may shift).
        state.index = rebuild_index(&state.crdt);
    }
}

// ── Read path ────────────────────────────────────────────────────────

/// Read the full pipeline state from the CRDT document.
///
/// Returns items grouped by stage, or `None` if the CRDT layer is not
/// initialised.
pub fn read_all_items() -> Option<Vec<PipelineItemView>> {
    let state_mutex = CRDT_STATE.get()?;
    let state = state_mutex.lock().ok()?;

    let mut items = Vec::new();
    for item_crdt in state.crdt.doc.items.iter() {
        if let Some(view) = extract_item_view(item_crdt) {
            items.push(view);
        }
    }
    Some(items)
}

/// Read a single pipeline item by story_id.
pub fn read_item(story_id: &str) -> Option<PipelineItemView> {
    let state_mutex = CRDT_STATE.get()?;
    let state = state_mutex.lock().ok()?;
    let &idx = state.index.get(story_id)?;
    extract_item_view(&state.crdt.doc.items[idx])
}

/// Extract a `PipelineItemView` from a `PipelineItemCrdt`.
fn extract_item_view(item: &PipelineItemCrdt) -> Option<PipelineItemView> {
    let story_id = match item.story_id.view() {
        JsonValue::String(s) if !s.is_empty() => s,
        _ => return None,
    };
    let stage = match item.stage.view() {
        JsonValue::String(s) if !s.is_empty() => s,
        _ => return None,
    };
    let name = match item.name.view() {
        JsonValue::String(s) if !s.is_empty() => Some(s),
        _ => None,
    };
    let agent = match item.agent.view() {
        JsonValue::String(s) if !s.is_empty() => Some(s),
        _ => None,
    };
    let retry_count = match item.retry_count.view() {
        JsonValue::Number(n) if n > 0.0 => Some(n as i64),
        _ => None,
    };
    let blocked = match item.blocked.view() {
        JsonValue::Bool(b) => Some(b),
        _ => None,
    };
    let depends_on = match item.depends_on.view() {
        JsonValue::String(s) if !s.is_empty() => {
            serde_json::from_str::<Vec<u32>>(&s).ok()
        }
        _ => None,
    };

    Some(PipelineItemView {
        story_id,
        stage,
        name,
        agent,
        retry_count,
        blocked,
        depends_on,
    })
}

/// Hex-encode a byte slice (no external dep needed).
mod hex {
    pub fn encode(bytes: &[u8]) -> String {
        bytes.iter().map(|b| format!("{b:02x}")).collect()
    }
}

// ── Tests ────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use bft_json_crdt::json_crdt::OpState;

    #[test]
    fn crdt_doc_insert_and_view() {
        let kp = make_keypair();
        let mut crdt = BaseCrdt::<PipelineDoc>::new(&kp);

        let item_json: JsonValue = json!({
            "story_id": "10_story_test",
            "stage": "2_current",
            "name": "Test Story",
            "agent": "coder-opus",
            "retry_count": 0.0,
            "blocked": false,
            "depends_on": "",
        })
        .into();

        let op = crdt.doc.items.insert(ROOT_ID, item_json).sign(&kp);
        assert_eq!(crdt.apply(op), OpState::Ok);

        let view = crdt.doc.items.view();
        assert_eq!(view.len(), 1);

        let item = &crdt.doc.items[0];
        assert_eq!(item.story_id.view(), JsonValue::String("10_story_test".to_string()));
        assert_eq!(item.stage.view(), JsonValue::String("2_current".to_string()));
    }

    #[test]
    fn crdt_doc_update_stage() {
        let kp = make_keypair();
        let mut crdt = BaseCrdt::<PipelineDoc>::new(&kp);

        let item_json: JsonValue = json!({
            "story_id": "20_story_move",
            "stage": "1_backlog",
            "name": "Move Me",
            "agent": "",
            "retry_count": 0.0,
            "blocked": false,
            "depends_on": "",
        })
        .into();

        let insert_op = crdt.doc.items.insert(ROOT_ID, item_json).sign(&kp);
        crdt.apply(insert_op);

        // Update stage
        let stage_op = crdt.doc.items[0].stage.set("2_current".to_string()).sign(&kp);
        crdt.apply(stage_op);

        assert_eq!(
            crdt.doc.items[0].stage.view(),
            JsonValue::String("2_current".to_string())
        );
    }

    #[test]
    fn crdt_ops_replay_reconstructs_state() {
        let kp = make_keypair();
        let mut crdt1 = BaseCrdt::<PipelineDoc>::new(&kp);

        // Build state with a series of ops.
        let item_json: JsonValue = json!({
            "story_id": "30_story_replay",
            "stage": "1_backlog",
            "name": "Replay Test",
            "agent": "",
            "retry_count": 0.0,
            "blocked": false,
            "depends_on": "",
        })
        .into();

        let op1 = crdt1.doc.items.insert(ROOT_ID, item_json).sign(&kp);
        crdt1.apply(op1.clone());

        let op2 = crdt1.doc.items[0].stage.set("2_current".to_string()).sign(&kp);
        crdt1.apply(op2.clone());

        let op3 = crdt1.doc.items[0].name.set("Updated Name".to_string()).sign(&kp);
        crdt1.apply(op3.clone());

        // Replay ops on a fresh CRDT.
        let mut crdt2 = BaseCrdt::<PipelineDoc>::new(&kp);
        crdt2.apply(op1);
        crdt2.apply(op2);
        crdt2.apply(op3);

        assert_eq!(
            crdt1.doc.items[0].stage.view(),
            crdt2.doc.items[0].stage.view()
        );
        assert_eq!(
            crdt1.doc.items[0].name.view(),
            crdt2.doc.items[0].name.view()
        );
    }

    #[test]
    fn extract_item_view_parses_crdt_item() {
        let kp = make_keypair();
        let mut crdt = BaseCrdt::<PipelineDoc>::new(&kp);

        let item_json: JsonValue = json!({
            "story_id": "40_story_view",
            "stage": "3_qa",
            "name": "View Test",
            "agent": "coder-1",
            "retry_count": 2.0,
            "blocked": true,
            "depends_on": "[10,20]",
        })
        .into();

        let op = crdt.doc.items.insert(ROOT_ID, item_json).sign(&kp);
        crdt.apply(op);

        let view = extract_item_view(&crdt.doc.items[0]).unwrap();
        assert_eq!(view.story_id, "40_story_view");
        assert_eq!(view.stage, "3_qa");
        assert_eq!(view.name.as_deref(), Some("View Test"));
        assert_eq!(view.agent.as_deref(), Some("coder-1"));
        assert_eq!(view.retry_count, Some(2));
        assert_eq!(view.blocked, Some(true));
        assert_eq!(view.depends_on, Some(vec![10, 20]));
    }

    #[test]
    fn rebuild_index_maps_story_ids() {
        let kp = make_keypair();
        let mut crdt = BaseCrdt::<PipelineDoc>::new(&kp);

        for (sid, stage) in &[("10_story_a", "1_backlog"), ("20_story_b", "2_current")] {
            let item: JsonValue = json!({
                "story_id": sid,
                "stage": stage,
                "name": "",
                "agent": "",
                "retry_count": 0.0,
                "blocked": false,
                "depends_on": "",
            })
            .into();
            let op = crdt.doc.items.insert(ROOT_ID, item).sign(&kp);
            crdt.apply(op);
        }

        let index = rebuild_index(&crdt);
        assert_eq!(index.len(), 2);
        assert!(index.contains_key("10_story_a"));
        assert!(index.contains_key("20_story_b"));
    }

    #[tokio::test]
    async fn init_and_write_read_roundtrip() {
        let tmp = tempfile::tempdir().unwrap();
        let db_path = tmp.path().join("crdt_test.db");

        // Init directly (not via the global singleton, for test isolation).
        let options = SqliteConnectOptions::new()
            .filename(&db_path)
            .create_if_missing(true);
        let pool = SqlitePool::connect_with(options).await.unwrap();
        sqlx::migrate!("./migrations").run(&pool).await.unwrap();

        let keypair = make_keypair();
        let mut crdt = BaseCrdt::<PipelineDoc>::new(&keypair);

        // Insert and update like write_item does.
        let item_json: JsonValue = json!({
            "story_id": "50_story_roundtrip",
            "stage": "1_backlog",
            "name": "Roundtrip",
            "agent": "",
            "retry_count": 0.0,
            "blocked": false,
            "depends_on": "",
        })
        .into();

        let insert_op = crdt.doc.items.insert(ROOT_ID, item_json).sign(&keypair);
        crdt.apply(insert_op.clone());

        // Persist the op.
        let op_json = serde_json::to_string(&insert_op).unwrap();
        let op_id = hex::encode(&insert_op.id());
        let now = chrono::Utc::now().to_rfc3339();
        sqlx::query(
            "INSERT INTO crdt_ops (op_id, seq, op_json, created_at) VALUES (?1, ?2, ?3, ?4)",
        )
        .bind(&op_id)
        .bind(insert_op.inner.seq as i64)
        .bind(&op_json)
        .bind(&now)
        .execute(&pool)
        .await
        .unwrap();

        // Reconstruct from DB.
        let rows: Vec<(String,)> =
            sqlx::query_as("SELECT op_json FROM crdt_ops ORDER BY seq ASC")
                .fetch_all(&pool)
                .await
                .unwrap();

        let mut crdt2 = BaseCrdt::<PipelineDoc>::new(&keypair);
        for (json_str,) in &rows {
            let op: SignedOp = serde_json::from_str(json_str).unwrap();
            crdt2.apply(op);
        }

        let view = extract_item_view(&crdt2.doc.items[0]).unwrap();
        assert_eq!(view.story_id, "50_story_roundtrip");
        assert_eq!(view.stage, "1_backlog");
        assert_eq!(view.name.as_deref(), Some("Roundtrip"));
    }

    #[test]
    fn signed_op_serialization_roundtrip() {
        let kp = make_keypair();
        let mut crdt = BaseCrdt::<PipelineDoc>::new(&kp);

        let item: JsonValue = json!({
            "story_id": "60_story_serde",
            "stage": "1_backlog",
            "name": "Serde Test",
            "agent": "",
            "retry_count": 0.0,
            "blocked": false,
            "depends_on": "",
        })
        .into();

        let op = crdt.doc.items.insert(ROOT_ID, item).sign(&kp);
        let json_str = serde_json::to_string(&op).unwrap();
        let deserialized: SignedOp = serde_json::from_str(&json_str).unwrap();

        assert_eq!(op.id(), deserialized.id());
        assert_eq!(op.inner.seq, deserialized.inner.seq);
    }
}