IdeA/crates/application/tests/background_task_cancel.rs

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

use async_trait::async_trait;

use application::CancelBackgroundTask;
use domain::ports::{
    BackgroundCompletionStream, BackgroundTaskHandle, BackgroundTaskPortError,
    BackgroundTaskRunner, BackgroundTaskSpec, BackgroundTaskStore, Clock, EventBus, EventStream,
};
use domain::{
    AgentId, BackgroundTask, BackgroundTaskKind, BackgroundTaskResult, BackgroundTaskState,
    BackgroundTaskWakePolicy, DomainEvent, ProjectId, TaskId,
};
use uuid::Uuid;

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

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

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

fn running_task(id: TaskId) -> BackgroundTask {
    BackgroundTask::new(
        id,
        pid(1),
        aid(2),
        BackgroundTaskKind::Command {
            label: "build".to_owned(),
        },
        BackgroundTaskWakePolicy::WakeOwner,
        100,
        None,
    )
    .unwrap()
    .transition(BackgroundTaskState::Running, 110)
    .unwrap()
}

#[derive(Default)]
struct FakeStore {
    tasks: Mutex<HashMap<TaskId, BackgroundTask>>,
    saves: Mutex<Vec<BackgroundTaskState>>,
}

impl FakeStore {
    fn insert(&self, task: BackgroundTask) {
        self.tasks.lock().unwrap().insert(task.id, task);
    }

    fn saved_states(&self) -> Vec<BackgroundTaskState> {
        self.saves.lock().unwrap().clone()
    }
}

#[async_trait]
impl BackgroundTaskStore for FakeStore {
    async fn create(&self, task: &BackgroundTask) -> Result<(), BackgroundTaskPortError> {
        self.insert(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> {
        self.saves.lock().unwrap().push(task.state);
        self.insert(task.clone());
        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(Vec::new())
    }

    async fn mark_completion_delivered(
        &self,
        _task_id: TaskId,
    ) -> Result<(), BackgroundTaskPortError> {
        Ok(())
    }
}

struct FakeRunner {
    result: Mutex<Result<(), BackgroundTaskPortError>>,
    cancels: Mutex<Vec<TaskId>>,
}

impl FakeRunner {
    fn ok() -> Self {
        Self {
            result: Mutex::new(Ok(())),
            cancels: Mutex::new(Vec::new()),
        }
    }

    fn failing(message: &str) -> Self {
        Self {
            result: Mutex::new(Err(BackgroundTaskPortError::Runner(message.to_owned()))),
            cancels: Mutex::new(Vec::new()),
        }
    }

    fn cancels(&self) -> Vec<TaskId> {
        self.cancels.lock().unwrap().clone()
    }
}

#[async_trait]
impl BackgroundTaskRunner for FakeRunner {
    async fn spawn(
        &self,
        spec: BackgroundTaskSpec,
    ) -> Result<BackgroundTaskHandle, BackgroundTaskPortError> {
        Ok(BackgroundTaskHandle {
            task_id: spec.task_id,
        })
    }

    async fn cancel(&self, task_id: TaskId) -> Result<(), BackgroundTaskPortError> {
        self.cancels.lock().unwrap().push(task_id);
        self.result.lock().unwrap().clone()
    }

    fn subscribe_completions(&self) -> BackgroundCompletionStream {
        Box::new(std::iter::empty())
    }
}

#[derive(Default)]
struct SpyEvents {
    events: Mutex<Vec<DomainEvent>>,
}

impl SpyEvents {
    fn events(&self) -> Vec<DomainEvent> {
        self.events.lock().unwrap().clone()
    }
}

impl EventBus for SpyEvents {
    fn publish(&self, event: DomainEvent) {
        self.events.lock().unwrap().push(event);
    }

    fn subscribe(&self) -> EventStream {
        Box::new(std::iter::empty())
    }
}

struct FixedClock(i64);

impl Clock for FixedClock {
    fn now_millis(&self) -> i64 {
        self.0
    }
}

fn usecase(
    store: Arc<FakeStore>,
    runner: Arc<FakeRunner>,
    events: Arc<SpyEvents>,
) -> CancelBackgroundTask {
    CancelBackgroundTask::new(store, runner, Arc::new(FixedClock(200))).with_events(events)
}

#[tokio::test]
async fn cancel_running_projects_cancelling_immediately_then_calls_runner() {
    let store = Arc::new(FakeStore::default());
    let runner = Arc::new(FakeRunner::ok());
    let events = Arc::new(SpyEvents::default());
    store.insert(running_task(tid(10)));

    let out = usecase(Arc::clone(&store), Arc::clone(&runner), Arc::clone(&events))
        .execute(tid(10))
        .await
        .unwrap();

    assert_eq!(out.task.state, BackgroundTaskState::Cancelling);
    assert_eq!(store.saved_states(), vec![BackgroundTaskState::Cancelling]);
    assert_eq!(runner.cancels(), vec![tid(10)]);
    assert_eq!(
        events.events(),
        vec![DomainEvent::BackgroundTaskStateChanged {
            project_id: pid(1),
            task_id: tid(10),
            owner_agent_id: aid(2),
            state: BackgroundTaskState::Cancelling,
        }]
    );
}

#[tokio::test]
async fn cancel_terminal_task_is_explicit_error_without_runner_call() {
    let store = Arc::new(FakeStore::default());
    let runner = Arc::new(FakeRunner::ok());
    let events = Arc::new(SpyEvents::default());
    let done = running_task(tid(10))
        .complete(BackgroundTaskResult::Success {
            finished_at_ms: 120,
            exit_code: Some(0),
            summary: "ok".to_owned(),
            stdout_tail: None,
            stderr_tail: None,
        })
        .unwrap();
    store.insert(done);

    let err = usecase(Arc::clone(&store), Arc::clone(&runner), events)
        .execute(tid(10))
        .await
        .unwrap_err();

    assert!(matches!(err, application::AppError::Invalid(_)));
    assert!(runner.cancels().is_empty());
    assert!(store.saved_states().is_empty());
}

#[tokio::test]
async fn cancel_already_cancelling_is_idempotent_and_retries_runner_cancel() {
    let store = Arc::new(FakeStore::default());
    let runner = Arc::new(FakeRunner::ok());
    let events = Arc::new(SpyEvents::default());
    store.insert(
        running_task(tid(10))
            .transition(BackgroundTaskState::Cancelling, 120)
            .unwrap(),
    );

    let out = usecase(Arc::clone(&store), Arc::clone(&runner), events)
        .execute(tid(10))
        .await
        .unwrap();

    assert_eq!(out.task.state, BackgroundTaskState::Cancelling);
    assert_eq!(runner.cancels(), vec![tid(10)]);
    assert!(store.saved_states().is_empty());
}

#[tokio::test]
async fn cancel_kill_failure_persists_cancel_failed_and_returns_error() {
    let store = Arc::new(FakeStore::default());
    let runner = Arc::new(FakeRunner::failing("kill denied"));
    let events = Arc::new(SpyEvents::default());
    store.insert(running_task(tid(10)));

    let err = usecase(Arc::clone(&store), Arc::clone(&runner), Arc::clone(&events))
        .execute(tid(10))
        .await
        .unwrap_err();

    assert!(matches!(err, application::AppError::Process(_)));
    assert_eq!(
        store.saved_states(),
        vec![
            BackgroundTaskState::Cancelling,
            BackgroundTaskState::CancelFailed,
        ]
    );
    let persisted = store.get(tid(10)).await.unwrap().unwrap();
    assert_eq!(persisted.state, BackgroundTaskState::CancelFailed);
    assert!(matches!(
        persisted.result,
        Some(BackgroundTaskResult::CancelFailed { .. })
    ));
    assert_eq!(
        events.events(),
        vec![
            DomainEvent::BackgroundTaskStateChanged {
                project_id: pid(1),
                task_id: tid(10),
                owner_agent_id: aid(2),
                state: BackgroundTaskState::Cancelling,
            },
            DomainEvent::BackgroundTaskStateChanged {
                project_id: pid(1),
                task_id: tid(10),
                owner_agent_id: aid(2),
                state: BackgroundTaskState::CancelFailed,
            },
        ]
    );
}