IdeA/crates/infrastructure/tests/background_completion_sink.rs

//! B3 tests for the background completion sink.

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{mpsc, Arc, Mutex};

use async_trait::async_trait;
use domain::ports::{
    BackgroundCompletionStream, BackgroundTaskCompletion, BackgroundTaskHandle,
    BackgroundTaskRunner, BackgroundTaskSpec,
};
use domain::{
    AgentId, BackgroundTask, BackgroundTaskKind, BackgroundTaskPortError, BackgroundTaskResult,
    BackgroundTaskState, BackgroundTaskStore, BackgroundTaskWakePolicy, ProjectId, ProjectPath,
    TaskId,
};
use infrastructure::{
    BackgroundCompletionSink, BackgroundCompletionSinkError, BackgroundCompletionSinkOutcome,
    FsBackgroundTaskStore,
};
use tokio::sync::mpsc::unbounded_channel;
use uuid::Uuid;

struct TempDir(PathBuf);

impl TempDir {
    fn new() -> Self {
        let path = std::env::temp_dir().join(format!("idea-b3-completion-sink-{}", 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()
    }
}

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 running_task(id: u128) -> BackgroundTask {
    BackgroundTask::new(
        task_id(id),
        project_id(10),
        agent_id(100),
        BackgroundTaskKind::Command {
            label: format!("task-{id}"),
        },
        BackgroundTaskWakePolicy::WakeOwner,
        1_000,
        None,
    )
    .unwrap()
    .transition(BackgroundTaskState::Running, 1_010)
    .unwrap()
}

fn success_completion(id: u128) -> BackgroundTaskCompletion {
    BackgroundTaskCompletion {
        task_id: task_id(id),
        result: BackgroundTaskResult::Success {
            finished_at_ms: 1_100,
            exit_code: Some(0),
            summary: "ok".into(),
            stdout_tail: Some("done".into()),
            stderr_tail: None,
        },
    }
}

#[tokio::test]
async fn runner_stream_completion_is_persisted_then_signalled() {
    let tmp = TempDir::new();
    let store = Arc::new(FsBackgroundTaskStore::new(&tmp.project_path()));
    let task = running_task(1);
    store.create(&task).await.unwrap();

    let (ready_tx, mut ready_rx) = unbounded_channel();
    let sink = Arc::new(BackgroundCompletionSink::new(store.clone(), ready_tx));
    let (completion_tx, completion_rx) = mpsc::channel();
    let runner = Arc::new(FakeRunner::new(completion_rx));
    let handle = sink.start_from_runner(runner);

    completion_tx.send(success_completion(1)).unwrap();
    drop(completion_tx);

    let ready = tokio::time::timeout(std::time::Duration::from_secs(2), ready_rx.recv())
        .await
        .unwrap()
        .unwrap();
    assert_eq!(ready.task_id, task.id);

    let persisted = store.get(task.id).await.unwrap().unwrap();
    assert_eq!(persisted.state, BackgroundTaskState::Completed);
    assert!(persisted.has_pending_completion_delivery());
    handle.await.unwrap();
}

#[tokio::test]
async fn persistence_happens_before_ready_signal() {
    let task = running_task(1);
    let store = Arc::new(RecordingStore::with_task(task.clone()));
    let order = store.order.clone();
    let (ready_tx, mut ready_rx) = unbounded_channel();
    let sink = BackgroundCompletionSink::new(store, ready_tx);

    let receiver_order = order.clone();
    let receiver = tokio::spawn(async move {
        let _ = ready_rx.recv().await;
        receiver_order.lock().unwrap().push("ready");
    });

    let outcome = sink
        .process_completion(success_completion(1))
        .await
        .unwrap();
    assert!(matches!(
        outcome,
        BackgroundCompletionSinkOutcome::PersistedAndSignalled(_)
    ));
    receiver.await.unwrap();

    assert_eq!(*order.lock().unwrap(), vec!["save", "ready"]);
}

#[tokio::test]
async fn duplicate_completion_same_task_signals_once() {
    let tmp = TempDir::new();
    let store = Arc::new(FsBackgroundTaskStore::new(&tmp.project_path()));
    let task = running_task(1);
    store.create(&task).await.unwrap();
    let (ready_tx, mut ready_rx) = unbounded_channel();
    let sink = BackgroundCompletionSink::new(store.clone(), ready_tx);

    let first = sink
        .process_completion(success_completion(1))
        .await
        .unwrap();
    let second = sink
        .process_completion(success_completion(1))
        .await
        .unwrap();

    assert!(matches!(
        first,
        BackgroundCompletionSinkOutcome::PersistedAndSignalled(_)
    ));
    assert_eq!(
        second,
        BackgroundCompletionSinkOutcome::IgnoredDuplicate { task_id: task.id }
    );
    assert!(ready_rx.recv().await.is_some());
    assert!(ready_rx.try_recv().is_err());

    let persisted = store.get(task.id).await.unwrap().unwrap();
    assert_eq!(persisted.state, BackgroundTaskState::Completed);
    assert!(persisted.has_pending_completion_delivery());
}

#[tokio::test]
async fn completion_for_already_terminal_task_is_ignored() {
    let tmp = TempDir::new();
    let store = Arc::new(FsBackgroundTaskStore::new(&tmp.project_path()));
    let completed = running_task(1)
        .complete(success_completion(1).result)
        .unwrap();
    store.create(&completed).await.unwrap();
    let (ready_tx, mut ready_rx) = unbounded_channel();
    let sink = BackgroundCompletionSink::new(store.clone(), ready_tx);

    let outcome = sink
        .process_completion(success_completion(1))
        .await
        .unwrap();

    assert_eq!(
        outcome,
        BackgroundCompletionSinkOutcome::IgnoredAlreadyTerminal {
            task_id: completed.id
        }
    );
    assert!(ready_rx.try_recv().is_err());
    assert_eq!(store.get(completed.id).await.unwrap(), Some(completed));
}

#[tokio::test]
async fn failed_persistence_does_not_signal_ready() {
    let task = running_task(1);
    let store = Arc::new(RecordingStore::with_task(task.clone()).failing_save());
    let (ready_tx, mut ready_rx) = unbounded_channel();
    let sink = BackgroundCompletionSink::new(store, ready_tx);

    let err = sink
        .process_completion(success_completion(1))
        .await
        .unwrap_err();

    assert_eq!(
        err,
        BackgroundCompletionSinkError::Store(BackgroundTaskPortError::Store("boom".into()))
    );
    assert!(ready_rx.try_recv().is_err());
}

#[tokio::test]
async fn failed_persistence_releases_task_for_later_retry() {
    let task = running_task(1);
    let store = Arc::new(RecordingStore::with_task(task.clone()).fail_first_save());
    let (ready_tx, mut ready_rx) = unbounded_channel();
    let sink = BackgroundCompletionSink::new(store, ready_tx);

    assert!(sink
        .process_completion(success_completion(1))
        .await
        .is_err());
    let second = sink
        .process_completion(success_completion(1))
        .await
        .unwrap();

    assert!(matches!(
        second,
        BackgroundCompletionSinkOutcome::PersistedAndSignalled(_)
    ));
    assert!(ready_rx.recv().await.is_some());
}

struct FakeRunner {
    rx: Mutex<Option<mpsc::Receiver<BackgroundTaskCompletion>>>,
}

impl FakeRunner {
    fn new(rx: mpsc::Receiver<BackgroundTaskCompletion>) -> Self {
        Self {
            rx: Mutex::new(Some(rx)),
        }
    }
}

#[async_trait]
impl BackgroundTaskRunner for FakeRunner {
    async fn spawn(
        &self,
        _spec: BackgroundTaskSpec,
    ) -> Result<BackgroundTaskHandle, BackgroundTaskPortError> {
        Err(BackgroundTaskPortError::Runner("not implemented".into()))
    }

    async fn cancel(&self, _task_id: TaskId) -> Result<(), BackgroundTaskPortError> {
        Err(BackgroundTaskPortError::Runner("not implemented".into()))
    }

    fn subscribe_completions(&self) -> BackgroundCompletionStream {
        let rx = self.rx.lock().unwrap().take().unwrap();
        Box::new(rx.into_iter())
    }
}

struct RecordingStore {
    tasks: Mutex<HashMap<TaskId, BackgroundTask>>,
    order: Arc<Mutex<Vec<&'static str>>>,
    fail_saves_remaining: Mutex<usize>,
}

impl RecordingStore {
    fn with_task(task: BackgroundTask) -> Self {
        Self {
            tasks: Mutex::new(HashMap::from([(task.id, task)])),
            order: Arc::new(Mutex::new(Vec::new())),
            fail_saves_remaining: Mutex::new(0),
        }
    }

    fn failing_save(self) -> Self {
        *self.fail_saves_remaining.lock().unwrap() = usize::MAX;
        self
    }

    fn fail_first_save(self) -> Self {
        *self.fail_saves_remaining.lock().unwrap() = 1;
        self
    }
}

#[async_trait]
impl BackgroundTaskStore for RecordingStore {
    async fn create(&self, task: &BackgroundTask) -> Result<(), BackgroundTaskPortError> {
        self.tasks.lock().unwrap().insert(task.id, task.clone());
        Ok(())
    }

    async fn get(&self, id: TaskId) -> Result<Option<BackgroundTask>, BackgroundTaskPortError> {
        Ok(self.tasks.lock().unwrap().get(&id).cloned())
    }

    async fn save(&self, task: &BackgroundTask) -> Result<(), BackgroundTaskPortError> {
        let mut failures = self.fail_saves_remaining.lock().unwrap();
        if *failures > 0 {
            if *failures != usize::MAX {
                *failures -= 1;
            }
            return Err(BackgroundTaskPortError::Store("boom".into()));
        }
        drop(failures);
        self.tasks.lock().unwrap().insert(task.id, task.clone());
        self.order.lock().unwrap().push("save");
        Ok(())
    }

    async fn list_open_for_agent(
        &self,
        agent_id: AgentId,
    ) -> Result<Vec<BackgroundTask>, BackgroundTaskPortError> {
        Ok(self
            .tasks
            .lock()
            .unwrap()
            .values()
            .filter(|task| task.owner_agent_id == agent_id && !task.is_terminal())
            .cloned()
            .collect())
    }

    async fn list_undelivered_completions(
        &self,
    ) -> Result<Vec<BackgroundTask>, BackgroundTaskPortError> {
        Ok(self
            .tasks
            .lock()
            .unwrap()
            .values()
            .filter(|task| task.has_pending_completion_delivery())
            .cloned()
            .collect())
    }

    async fn mark_completion_delivered(
        &self,
        task_id: TaskId,
    ) -> Result<(), BackgroundTaskPortError> {
        let mut tasks = self.tasks.lock().unwrap();
        let task = tasks
            .get(&task_id)
            .cloned()
            .ok_or(BackgroundTaskPortError::NotFound)?;
        tasks.insert(
            task_id,
            task.mark_completion_delivered()
                .map_err(|e| BackgroundTaskPortError::Invalid(e.to_string()))?,
        );
        Ok(())
    }
}