ai-dispatch 8.91.0

Multi-AI CLI team orchestrator
// Task-focused Store tests.
// Exports: task query/mutation tests.
// Deps: Store, rusqlite.

use super::*;
use crate::store::TaskCompletionUpdate;

#[test]
fn insert_and_get_task() {
    let store = Store::open_memory().unwrap();
    let task = make_task("t-0001", AgentKind::Codex, TaskStatus::Running);
    store.insert_task(&task).unwrap();

    let loaded = store.get_task("t-0001").unwrap().unwrap();
    assert_eq!(loaded.id, task.id);
    assert_eq!(loaded.agent, AgentKind::Codex);
    assert_eq!(loaded.status, TaskStatus::Running);
    assert_eq!(loaded.pending_reason, None);
}

#[test]
fn insert_and_get_task_persists_dispatch_flags() {
    let store = Store::open_memory().unwrap();
    let mut task = make_task("t-0004", AgentKind::Codex, TaskStatus::Pending);
    task.verify = Some("cargo test".to_string());
    task.read_only = true;
    task.budget = true;
    store.insert_task(&task).unwrap();

    let loaded = store.get_task("t-0004").unwrap().unwrap();
    assert_eq!(loaded.verify.as_deref(), Some("cargo test"));
    assert!(loaded.read_only);
    assert!(loaded.budget);
}

#[test]
fn insert_and_update_task_persists_delivery_assessment() {
    let store = Store::open_memory().unwrap();
    let mut task = make_task("t-0004d", AgentKind::Codex, TaskStatus::Done);
    task.delivery_assessment = Some(DeliveryAssessment::HollowOutput);
    store.insert_task(&task).unwrap();

    let loaded = store.get_task("t-0004d").unwrap().unwrap();
    assert_eq!(loaded.delivery_assessment, Some(DeliveryAssessment::HollowOutput));

    store
        .update_delivery_assessment("t-0004d", Some(DeliveryAssessment::EmptyDiff))
        .unwrap();
    let updated = store.get_task("t-0004d").unwrap().unwrap();
    assert_eq!(updated.delivery_assessment, Some(DeliveryAssessment::EmptyDiff));

    store.update_delivery_assessment("t-0004d", None).unwrap();
    let cleared = store.get_task("t-0004d").unwrap().unwrap();
    assert_eq!(cleared.delivery_assessment, None);
}

#[test]
fn insert_waiting_task_persists_retry_fields() {
    let store = Store::open_memory().unwrap();
    let prompt = "x".repeat(160);
    store
        .insert_waiting_task(
            "t-0004w",
            "codex",
            &prompt,
            Some("resolved prompt"),
            Some("wg-batch"),
            Some("/tmp/repo"),
            Some("feat/retry"),
            Some("o3"),
            Some("cargo check"),
            true,
            true,
        )
        .unwrap();

    let loaded = store.get_task("t-0004w").unwrap().unwrap();
    assert_eq!(loaded.status, TaskStatus::Waiting);
    assert_eq!(loaded.prompt, prompt);
    assert_eq!(loaded.resolved_prompt.as_deref(), Some("resolved prompt"));
    assert_eq!(loaded.workgroup_id.as_deref(), Some("wg-batch"));
    assert_eq!(loaded.repo_path.as_deref(), Some("/tmp/repo"));
    assert_eq!(loaded.worktree_branch.as_deref(), Some("feat/retry"));
    assert_eq!(loaded.model.as_deref(), Some("o3"));
    assert_eq!(loaded.verify.as_deref(), Some("cargo check"));
    assert!(loaded.read_only);
    assert!(loaded.budget);
}

#[test]
fn migrate_adds_repo_path_column() {
    let conn = rusqlite::Connection::open_in_memory().unwrap();
    conn.execute_batch(
        "CREATE TABLE tasks (
            id TEXT PRIMARY KEY,
            agent TEXT NOT NULL,
            prompt TEXT NOT NULL,
            status TEXT NOT NULL DEFAULT 'pending',
            worktree_path TEXT,
            worktree_branch TEXT,
            log_path TEXT,
            output_path TEXT,
            tokens INTEGER,
            duration_ms INTEGER,
            created_at TEXT NOT NULL,
            completed_at TEXT
        );
        CREATE TABLE events (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            task_id TEXT NOT NULL,
            timestamp TEXT NOT NULL,
            event_type TEXT NOT NULL,
            detail TEXT NOT NULL
        );",
    )
    .unwrap();
    let store = Store {
        conn: std::sync::Mutex::new(conn),
    };

    store.migrate().unwrap();

    let conn = store.db();
    let mut stmt = conn.prepare("PRAGMA table_info(tasks)").unwrap();
    let columns = stmt
        .query_map([], |row| row.get::<_, String>(1))
        .unwrap()
        .map(|row| row.unwrap())
        .collect::<Vec<_>>();
    assert!(columns.contains(&"repo_path".to_string()));
    assert!(columns.contains(&"resolved_prompt".to_string()));
    assert!(columns.contains(&"verify".to_string()));
    assert!(columns.contains(&"read_only".to_string()));
    assert!(columns.contains(&"budget".to_string()));
    assert!(columns.contains(&"category".to_string()));
    assert!(columns.contains(&"pending_reason".to_string()));
    assert!(columns.contains(&"delivery_assessment".to_string()));
}

#[test]
fn migrate_moves_legacy_delivery_status_out_of_verify_status() {
    let conn = rusqlite::Connection::open_in_memory().unwrap();
    conn.execute_batch(
        "CREATE TABLE tasks (
            id TEXT PRIMARY KEY,
            agent TEXT NOT NULL,
            prompt TEXT NOT NULL,
            status TEXT NOT NULL DEFAULT 'pending',
            created_at TEXT NOT NULL,
            verify_status TEXT NOT NULL DEFAULT 'skipped'
        );
        INSERT INTO tasks (id, agent, prompt, status, created_at, verify_status)
        VALUES ('t-empty', 'codex', 'prompt', 'done', '2026-04-16T00:00:00+00:00', 'empty_diff');
        CREATE TABLE events (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            task_id TEXT NOT NULL,
            timestamp TEXT NOT NULL,
            event_type TEXT NOT NULL,
            detail TEXT NOT NULL
        );",
    )
    .unwrap();
    let store = Store {
        conn: std::sync::Mutex::new(conn),
    };

    store.migrate().unwrap();

    let conn = store.db();
    let (verify_status, delivery_assessment): (String, Option<String>) = conn
        .query_row(
            "SELECT verify_status, delivery_assessment FROM tasks WHERE id = 't-empty'",
            [],
            |row| Ok((row.get(0)?, row.get(1)?)),
        )
        .unwrap();
    assert_eq!(verify_status, "skipped");
    assert_eq!(delivery_assessment.as_deref(), Some("empty_diff"));
}

#[test]
fn update_completion() {
    let store = Store::open_memory().unwrap();
    let task = make_task("t-0002", AgentKind::Gemini, TaskStatus::Running);
    store.insert_task(&task).unwrap();
    store
        .update_task_completion(TaskCompletionUpdate {
            id: "t-0002",
            status: TaskStatus::Done,
            tokens: Some(3000),
            duration_ms: 47000,
            model: Some("gemini-2.5-flash"),
            cost_usd: Some(0.0038),
            exit_code: None,
        })
        .unwrap();

    let loaded = store.get_task("t-0002").unwrap().unwrap();
    assert_eq!(loaded.status, TaskStatus::Done);
        assert_eq!(loaded.tokens, Some(3000));
        assert_eq!(loaded.duration_ms, Some(47000));
        assert_eq!(loaded.model.as_deref(), Some("gemini-2.5-flash"));
        assert!((loaded.cost_usd.unwrap() - 0.0038).abs() < 0.0001);
        assert_eq!(loaded.exit_code, None);
        assert!(loaded.completed_at.is_some());
}

#[test]
fn update_completion_does_not_override_stopped_status() {
    let store = Store::open_memory().unwrap();
    let task = make_task("t-0002s", AgentKind::Codex, TaskStatus::Stopped);
    store.insert_task(&task).unwrap();
    store
        .update_task_completion(TaskCompletionUpdate {
            id: "t-0002s",
            status: TaskStatus::Failed,
            tokens: None,
            duration_ms: 1234,
            model: Some("gpt-5.4"),
            cost_usd: None,
            exit_code: Some(1),
        })
        .unwrap();

    let loaded = store.get_task("t-0002s").unwrap().unwrap();
    assert_eq!(loaded.status, TaskStatus::Stopped);
}

#[test]
fn update_resolved_prompt_persists() {
    let store = Store::open_memory().unwrap();
    let task = make_task("t-0003", AgentKind::Codex, TaskStatus::Pending);
    store.insert_task(&task).unwrap();

    store
        .update_resolved_prompt("t-0003", "resolved prompt")
        .unwrap();

    let loaded = store.get_task("t-0003").unwrap().unwrap();
    assert_eq!(loaded.resolved_prompt.as_deref(), Some("resolved prompt"));
}

#[test]
fn update_output_path_sets_field() {
    let store = Store::open_memory().unwrap();
    let task = make_task("t-0005", AgentKind::Codex, TaskStatus::Pending);
    store.insert_task(&task).unwrap();
    store.update_output_path("t-0005", "/tmp/output.md").unwrap();
    let loaded = store.get_task("t-0005").unwrap().unwrap();
    assert_eq!(loaded.output_path.as_deref(), Some("/tmp/output.md"));
}

#[test]
fn fail_pending_with_reason_persists_reason() {
    let store = Store::open_memory().unwrap();
    let task = make_task("t-0006", AgentKind::Codex, TaskStatus::Pending);
    store.insert_task(&task).unwrap();

    let changed = store
        .fail_pending_with_reason("t-0006", PendingReason::RateLimited)
        .unwrap();

    assert!(changed);
    let loaded = store.get_task("t-0006").unwrap().unwrap();
    assert_eq!(loaded.status, TaskStatus::Failed);
    assert_eq!(loaded.pending_reason.as_deref(), Some("rate_limited"));
}

#[test]
fn list_running_filter() {
    let store = Store::open_memory().unwrap();
    store
        .insert_task(&make_task("t-0010", AgentKind::Codex, TaskStatus::Running))
        .unwrap();
    store
        .insert_task(&make_task(
            "t-0012",
            AgentKind::Cursor,
            TaskStatus::AwaitingInput,
        ))
        .unwrap();
    store
        .insert_task(&make_task("t-0011", AgentKind::Gemini, TaskStatus::Done))
        .unwrap();

    let running = store.list_tasks(TaskFilter::Running).unwrap();
    assert_eq!(running.len(), 2);
    let ids = running
        .into_iter()
        .map(|task| task.id.0)
        .collect::<Vec<_>>();
    assert!(ids.contains(&"t-0010".to_string()));
    assert!(ids.contains(&"t-0012".to_string()));
}

#[test]
fn gets_retry_chain_from_root_to_current() {
    let store = Store::open_memory().unwrap();
    let root = make_task("t-1001", AgentKind::Codex, TaskStatus::Done);
    let mut retry_1 = make_task("t-1002", AgentKind::Codex, TaskStatus::Failed);
    retry_1.parent_task_id = Some("t-1001".to_string());
    let mut retry_2 = make_task("t-1003", AgentKind::Codex, TaskStatus::Done);
    retry_2.parent_task_id = Some("t-1002".to_string());

    store.insert_task(&root).unwrap();
    store.insert_task(&retry_1).unwrap();
    store.insert_task(&retry_2).unwrap();

    let chain = store.get_retry_chain("t-1003").unwrap();
    let ids = chain
        .iter()
        .map(|task| task.id.as_str())
        .collect::<Vec<_>>();
    assert_eq!(ids, vec!["t-1001", "t-1002", "t-1003"]);
}

#[test]
fn recent_tasks_for_agent_filters_to_recent_done_tasks() {
    let store = Store::open_memory().unwrap();
    let now = chrono::Local::now();
    let mut recent = make_task("t-2001", AgentKind::Codex, TaskStatus::Done);
    recent.created_at = now - chrono::Duration::days(1);
    recent.duration_ms = Some(120_000);
    let mut older = make_task("t-2002", AgentKind::Codex, TaskStatus::Done);
    older.created_at = now - chrono::Duration::days(8);
    older.duration_ms = Some(240_000);
    let mut failed = make_task("t-2003", AgentKind::Codex, TaskStatus::Failed);
    failed.created_at = now - chrono::Duration::days(1);
    failed.duration_ms = Some(180_000);
    let mut other_agent = make_task("t-2004", AgentKind::Gemini, TaskStatus::Done);
    other_agent.created_at = now - chrono::Duration::days(1);
    other_agent.duration_ms = Some(90_000);

    store.insert_task(&recent).unwrap();
    store.insert_task(&older).unwrap();
    store.insert_task(&failed).unwrap();
    store.insert_task(&other_agent).unwrap();

    let recent_tasks = store.recent_tasks_for_agent(AgentKind::Codex, 10).unwrap();
    let ids = recent_tasks
        .into_iter()
        .map(|task| task.id.0)
        .collect::<Vec<_>>();
    assert_eq!(ids, vec!["t-2001".to_string()]);
}