IdeA/crates/infrastructure/tests/background_task_store.rs

//! B2 integration tests for [`FsBackgroundTaskStore`].

use std::path::PathBuf;
use std::sync::Arc;

use domain::{
    AgentId, BackgroundTask, BackgroundTaskKind, BackgroundTaskResult, BackgroundTaskState,
    BackgroundTaskStore, BackgroundTaskWakePolicy, ProjectId, ProjectPath, TaskId,
};
use infrastructure::FsBackgroundTaskStore;
use uuid::Uuid;

struct TempDir(PathBuf);

impl TempDir {
    fn new() -> Self {
        let path = std::env::temp_dir().join(format!("idea-b2-bg-tasks-{}", Uuid::new_v4()));
        std::fs::create_dir_all(&path).unwrap();
        Self(path)
    }

    fn project_path(&self) -> ProjectPath {
        ProjectPath::new(self.0.to_string_lossy().into_owned()).unwrap()
    }

    fn task_dir(&self) -> PathBuf {
        self.0.join(".ideai").join("background-tasks")
    }

    fn task_file(&self, project_id: ProjectId) -> PathBuf {
        self.task_dir().join(format!("{project_id}.json"))
    }

    fn task_tmp_file(&self, project_id: ProjectId) -> PathBuf {
        self.task_dir().join(format!("{project_id}.json.tmp"))
    }

    fn tmp_files(&self) -> Vec<PathBuf> {
        let Ok(entries) = std::fs::read_dir(self.task_dir()) else {
            return Vec::new();
        };
        entries
            .filter_map(Result::ok)
            .map(|entry| entry.path())
            .filter(|path| path.extension().and_then(|ext| ext.to_str()) == Some("tmp"))
            .collect()
    }
}

impl Drop for TempDir {
    fn drop(&mut self) {
        let _ = std::fs::remove_dir_all(&self.0);
    }
}

fn task_id(n: u128) -> TaskId {
    TaskId::from_uuid(Uuid::from_u128(n))
}

fn project_id(n: u128) -> ProjectId {
    ProjectId::from_uuid(Uuid::from_u128(n))
}

fn agent_id(n: u128) -> AgentId {
    AgentId::from_uuid(Uuid::from_u128(n))
}

fn queued_task(id: u128, project: u128, owner: u128, now_ms: u64) -> BackgroundTask {
    BackgroundTask::new(
        task_id(id),
        project_id(project),
        agent_id(owner),
        BackgroundTaskKind::Command {
            label: format!("task-{id}"),
        },
        BackgroundTaskWakePolicy::WakeOwner,
        now_ms,
        None,
    )
    .unwrap()
}

fn running_task(id: u128, project: u128, owner: u128, now_ms: u64) -> BackgroundTask {
    queued_task(id, project, owner, now_ms)
        .transition(BackgroundTaskState::Running, now_ms + 1)
        .unwrap()
}

fn completed_task(id: u128, project: u128, owner: u128, now_ms: u64) -> BackgroundTask {
    running_task(id, project, owner, now_ms)
        .complete(BackgroundTaskResult::Success {
            finished_at_ms: now_ms + 2,
            exit_code: Some(0),
            summary: "ok".into(),
            stdout_tail: Some("done".into()),
            stderr_tail: None,
        })
        .unwrap()
}

#[tokio::test]
async fn create_get_save_round_trips_across_fresh_store() {
    let tmp = TempDir::new();
    let store = FsBackgroundTaskStore::new(&tmp.project_path());
    let task = queued_task(1, 10, 100, 1_000);

    store.create(&task).await.unwrap();
    assert_eq!(store.get(task.id).await.unwrap(), Some(task.clone()));

    let running = task
        .transition(BackgroundTaskState::Running, 1_010)
        .unwrap();
    store.save(&running).await.unwrap();

    let reborn = FsBackgroundTaskStore::new(&tmp.project_path());
    assert_eq!(reborn.get(task.id).await.unwrap(), Some(running));
    assert!(tmp.task_file(project_id(10)).exists());
    assert!(!tmp.task_tmp_file(project_id(10)).exists());

    let raw = std::fs::read_to_string(tmp.task_file(project_id(10))).unwrap();
    assert!(raw.contains("\"projectId\""), "camelCase doc: {raw}");
    assert!(raw.contains("\"ownerAgentId\""), "camelCase task: {raw}");
}

#[tokio::test]
async fn create_rejects_duplicate_task_id() {
    let tmp = TempDir::new();
    let store = FsBackgroundTaskStore::new(&tmp.project_path());
    let task = queued_task(1, 10, 100, 1_000);

    store.create(&task).await.unwrap();
    let err = store.create(&task).await.unwrap_err();

    assert_eq!(err, domain::BackgroundTaskPortError::AlreadyExists);
}

#[tokio::test]
async fn list_open_for_agent_uses_registry_and_filters_terminal_tasks() {
    let tmp = TempDir::new();
    let store = FsBackgroundTaskStore::new(&tmp.project_path());
    let open_a = queued_task(1, 10, 100, 1_000);
    let open_b = running_task(2, 10, 200, 1_000);
    let done_a = completed_task(3, 10, 100, 1_000);

    store.create(&open_a).await.unwrap();
    store.create(&open_b).await.unwrap();
    store.create(&done_a).await.unwrap();

    assert_eq!(
        store.registry_open_task_ids_for_agent(agent_id(100)).await,
        vec![open_a.id]
    );

    let listed = store.list_open_for_agent(agent_id(100)).await.unwrap();
    assert_eq!(listed, vec![open_a]);
    assert_eq!(
        store.list_open_for_agent(agent_id(200)).await.unwrap(),
        vec![open_b]
    );
}

#[tokio::test]
async fn list_undelivered_completions_and_mark_are_idempotent() {
    let tmp = TempDir::new();
    let store = FsBackgroundTaskStore::new(&tmp.project_path());
    let task = completed_task(1, 10, 100, 1_000);

    store.create(&task).await.unwrap();
    assert_eq!(
        store.list_undelivered_completions().await.unwrap(),
        vec![task.clone()]
    );

    store.mark_completion_delivered(task.id).await.unwrap();
    store.mark_completion_delivered(task.id).await.unwrap();

    let delivered = store.get(task.id).await.unwrap().unwrap();
    assert!(delivered.completion_delivered);
    assert!(store
        .list_undelivered_completions()
        .await
        .unwrap()
        .is_empty());
    assert!(store
        .list_open_for_agent(agent_id(100))
        .await
        .unwrap()
        .is_empty());
}

#[tokio::test]
async fn persistence_is_segmented_by_project_id() {
    let tmp = TempDir::new();
    let store = FsBackgroundTaskStore::new(&tmp.project_path());
    let p1 = queued_task(1, 10, 100, 1_000);
    let p2 = queued_task(2, 20, 100, 1_000);

    store.create(&p1).await.unwrap();
    store.create(&p2).await.unwrap();

    assert!(tmp.task_file(project_id(10)).exists());
    assert!(tmp.task_file(project_id(20)).exists());
    assert_ne!(tmp.task_file(project_id(10)), tmp.task_file(project_id(20)));

    let raw_p1 = std::fs::read_to_string(tmp.task_file(project_id(10))).unwrap();
    let raw_p2 = std::fs::read_to_string(tmp.task_file(project_id(20))).unwrap();
    assert!(raw_p1.contains(&p1.id.to_string()));
    assert!(!raw_p1.contains(&p2.id.to_string()));
    assert!(raw_p2.contains(&p2.id.to_string()));
    assert!(!raw_p2.contains(&p1.id.to_string()));
}

#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
async fn concurrent_creates_for_same_project_do_not_lose_tasks() {
    let tmp = TempDir::new();
    let store = Arc::new(FsBackgroundTaskStore::new(&tmp.project_path()));
    let project = project_id(10);
    let owner = agent_id(100);

    let mut handles = Vec::new();
    for n in 1..=64 {
        let store = Arc::clone(&store);
        handles.push(tokio::spawn(async move {
            let task = queued_task(n, 10, 100, 1_000 + n as u64);
            store.create(&task).await.map(|()| task.id)
        }));
    }

    let mut expected = Vec::new();
    for handle in handles {
        expected.push(handle.await.unwrap().unwrap());
    }
    expected.sort();

    let reborn = FsBackgroundTaskStore::new(&tmp.project_path());
    for id in &expected {
        assert!(
            reborn.get(*id).await.unwrap().is_some(),
            "created task {id} must be durable"
        );
    }

    assert_eq!(
        store.registry_open_task_ids_for_agent(owner).await,
        expected,
        "registry must retain every concurrent create"
    );
    assert!(tmp.tmp_files().is_empty(), "tmp files left behind");
    assert!(tmp.task_file(project).exists());
}

#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
async fn concurrent_saves_for_same_project_do_not_overwrite_each_other() {
    let tmp = TempDir::new();
    let store = Arc::new(FsBackgroundTaskStore::new(&tmp.project_path()));

    for n in 1..=64 {
        store
            .create(&running_task(n, 10, 100, 1_000))
            .await
            .unwrap();
    }

    let mut handles = Vec::new();
    for n in 1..=64 {
        let store = Arc::clone(&store);
        handles.push(tokio::spawn(async move {
            let task = completed_task(n, 10, 100, 2_000 + n as u64);
            store.save(&task).await.map(|()| task.id)
        }));
    }

    let mut expected = Vec::new();
    for handle in handles {
        expected.push(handle.await.unwrap().unwrap());
    }
    expected.sort();

    let reborn = FsBackgroundTaskStore::new(&tmp.project_path());
    for id in &expected {
        let task = reborn.get(*id).await.unwrap().unwrap();
        assert_eq!(task.state, BackgroundTaskState::Completed);
    }

    assert!(
        reborn
            .list_open_for_agent(agent_id(100))
            .await
            .unwrap()
            .is_empty(),
        "all concurrently saved tasks must be terminal"
    );
    assert_eq!(
        reborn.list_undelivered_completions().await.unwrap().len(),
        expected.len()
    );
    assert!(tmp.tmp_files().is_empty(), "tmp files left behind");
}

#[tokio::test]
async fn reconcile_marks_running_without_live_handle_failed_and_pending() {
    let tmp = TempDir::new();
    let store = FsBackgroundTaskStore::new(&tmp.project_path());
    let lost = running_task(1, 10, 100, 1_000);
    let still_live = running_task(2, 10, 100, 1_000);
    let already_pending = completed_task(3, 10, 100, 1_000);

    store.create(&lost).await.unwrap();
    store.create(&still_live).await.unwrap();
    store.create(&already_pending).await.unwrap();

    let report = store.reconcile_boot(&[still_live.id], 2_000).await.unwrap();

    assert_eq!(report.failed_task_ids, vec![lost.id]);
    assert_eq!(
        report.delivery_pending_task_ids,
        vec![lost.id, already_pending.id]
    );

    let lost_after = store.get(lost.id).await.unwrap().unwrap();
    assert_eq!(lost_after.state, BackgroundTaskState::Failed);
    assert!(lost_after.has_pending_completion_delivery());
    assert!(matches!(
        lost_after.result,
        Some(BackgroundTaskResult::Failure {
            finished_at_ms: 2_000,
            exit_code: None,
            ..
        })
    ));

    let live_after = store.get(still_live.id).await.unwrap().unwrap();
    assert_eq!(live_after.state, BackgroundTaskState::Running);
    assert_eq!(
        store.registry_open_task_ids_for_agent(agent_id(100)).await,
        vec![still_live.id]
    );
}

#[tokio::test]
async fn reconcile_uses_monotonic_timestamp_when_clock_is_behind() {
    let tmp = TempDir::new();
    let store = FsBackgroundTaskStore::new(&tmp.project_path());
    let lost = running_task(1, 10, 100, 1_000);

    store.create(&lost).await.unwrap();
    store.reconcile_boot(&[], 900).await.unwrap();

    let after = store.get(lost.id).await.unwrap().unwrap();
    assert!(matches!(
        after.result,
        Some(BackgroundTaskResult::Failure {
            finished_at_ms: 1_001,
            ..
        })
    ));
}