IdeA/crates/application/tests/agent_wake.rs

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

use application::{AgentWakeService, WakeSessionProvider};
use async_trait::async_trait;
use domain::background_task::{
    BackgroundTask, BackgroundTaskKind, BackgroundTaskResult, BackgroundTaskState,
    BackgroundTaskWakePolicy,
};
use domain::ids::{AgentId, ProjectId, SessionId, TaskId};
use domain::inbox::{
    AgentInbox, AgentInboxSnapshot, InboxError, InboxItem, InboxItemKind, InboxReceipt,
    InboxReceiptStatus, InboxSource,
};
use domain::input::{AgentBusyState, InputMediator};
use domain::mailbox::{AgentMailbox, MailboxError, PendingReply, Ticket, TicketId};
use domain::ports::{
    AgentSession, AgentSessionError, AgentWakePort, BackgroundTaskPortError, BackgroundTaskStore,
    ReplyEvent, ReplyStream, WakeError, WakeReason,
};
use domain::project::{Project, ProjectPath};
use domain::remote::RemoteRef;
use uuid::Uuid;

fn id(n: u128) -> Uuid {
    Uuid::from_u128(n)
}

fn agent(n: u128) -> AgentId {
    AgentId::from_uuid(id(n))
}

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

fn ticket(n: u128) -> TicketId {
    TicketId::from_uuid(id(n))
}

fn project() -> Project {
    Project::new(
        ProjectId::from_uuid(id(100)),
        "p",
        ProjectPath::new("/tmp/idea-agent-wake-test").unwrap(),
        RemoteRef::local(),
        0,
    )
    .unwrap()
}

fn completion_item(agent_id: AgentId, task_id: TaskId, ticket_id: TicketId) -> InboxItem {
    InboxItem {
        id: ticket_id,
        agent_id,
        source: InboxSource::BackgroundTask { task_id },
        kind: InboxItemKind::BackgroundCompletion,
        body: format!("Background task {task_id} completed."),
        created_at_ms: 10,
        correlation_id: Some(task_id.to_string()),
    }
}

fn completed_task(
    project: &Project,
    owner: AgentId,
    task_id: TaskId,
    summary: &str,
) -> BackgroundTask {
    BackgroundTask::new(
        task_id,
        project.id,
        owner,
        BackgroundTaskKind::Command {
            label: "build".to_owned(),
        },
        BackgroundTaskWakePolicy::WakeOwner,
        1,
        None,
    )
    .unwrap()
    .transition(BackgroundTaskState::Running, 10)
    .unwrap()
    .complete(BackgroundTaskResult::Success {
        finished_at_ms: 20,
        exit_code: Some(0),
        summary: summary.to_owned(),
        stdout_tail: Some("ok".to_owned()),
        stderr_tail: None,
    })
    .unwrap()
}

#[derive(Default)]
struct FakeInbox {
    queues: Mutex<HashMap<AgentId, VecDeque<InboxItem>>>,
}

impl AgentInbox for FakeInbox {
    fn enqueue_message(
        &self,
        agent_id: AgentId,
        item: InboxItem,
    ) -> Result<InboxReceipt, InboxError> {
        let mut queues = self.queues.lock().unwrap();
        let queue = queues.entry(agent_id).or_default();
        let item_id = item.id;
        queue.push_back(item);
        Ok(InboxReceipt {
            item_id,
            agent_id,
            depth: queue.len(),
            status: InboxReceiptStatus::Queued,
        })
    }

    fn dequeue_next(&self, agent_id: AgentId) -> Option<InboxItem> {
        self.queues
            .lock()
            .unwrap()
            .entry(agent_id)
            .or_default()
            .pop_front()
    }

    fn snapshot(&self, agent_id: AgentId) -> AgentInboxSnapshot {
        let queues = self.queues.lock().unwrap();
        let items = queues
            .get(&agent_id)
            .map(|queue| queue.iter().cloned().collect::<Vec<_>>())
            .unwrap_or_default();
        AgentInboxSnapshot {
            agent_id,
            depth: items.len(),
            items,
        }
    }
}

#[derive(Default)]
struct SharedTurnState {
    busy: Mutex<HashMap<AgentId, AgentBusyState>>,
    tickets: Mutex<HashMap<AgentId, VecDeque<Ticket>>>,
}

impl SharedTurnState {
    fn force_busy(&self, agent_id: AgentId, ticket_id: TicketId) {
        self.busy.lock().unwrap().insert(
            agent_id,
            AgentBusyState::Busy {
                ticket: ticket_id,
                since_ms: 1,
            },
        );
    }

    fn ticket_depth(&self, agent_id: AgentId) -> usize {
        self.tickets
            .lock()
            .unwrap()
            .get(&agent_id)
            .map(VecDeque::len)
            .unwrap_or_default()
    }
}

impl InputMediator for SharedTurnState {
    fn enqueue(&self, agent_id: AgentId, ticket: Ticket) -> PendingReply {
        self.enqueue_silent(agent_id, ticket)
    }

    fn enqueue_silent(&self, agent_id: AgentId, ticket: Ticket) -> PendingReply {
        self.busy.lock().unwrap().insert(
            agent_id,
            AgentBusyState::Busy {
                ticket: ticket.id,
                since_ms: 1,
            },
        );
        self.tickets
            .lock()
            .unwrap()
            .entry(agent_id)
            .or_default()
            .push_back(ticket);
        PendingReply::new(Box::pin(async { Err(MailboxError::Cancelled) }))
    }

    fn preempt(&self, _agent: AgentId) {}

    fn mark_idle(&self, agent_id: AgentId) {
        self.busy
            .lock()
            .unwrap()
            .insert(agent_id, AgentBusyState::Idle);
    }

    fn busy_state(&self, agent_id: AgentId) -> AgentBusyState {
        self.busy
            .lock()
            .unwrap()
            .get(&agent_id)
            .copied()
            .unwrap_or(AgentBusyState::Idle)
    }
}

impl AgentMailbox for SharedTurnState {
    fn enqueue(&self, agent_id: AgentId, ticket: Ticket) -> PendingReply {
        <Self as InputMediator>::enqueue(self, agent_id, ticket)
    }

    fn resolve(&self, _agent: AgentId, _result: String) -> Result<(), MailboxError> {
        Ok(())
    }

    fn cancel_head(&self, agent_id: AgentId, ticket_id: TicketId) {
        let mut tickets = self.tickets.lock().unwrap();
        if let Some(queue) = tickets.get_mut(&agent_id) {
            if queue.front().is_some_and(|ticket| ticket.id == ticket_id) {
                queue.pop_front();
            }
        }
    }
}

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

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

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

#[async_trait]
impl BackgroundTaskStore for FakeTaskStore {
    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.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(self
            .tasks
            .lock()
            .unwrap()
            .values()
            .filter(|task| task.is_terminal() && !task.completion_delivered)
            .cloned()
            .collect())
    }

    async fn mark_completion_delivered(
        &self,
        task_id: TaskId,
    ) -> Result<(), BackgroundTaskPortError> {
        self.delivered.lock().unwrap().push(task_id);
        if let Some(task) = self.tasks.lock().unwrap().get_mut(&task_id) {
            task.completion_delivered = true;
        }
        Ok(())
    }
}

#[derive(Default)]
struct FakeSession {
    prompts: Mutex<Vec<String>>,
    events: Mutex<Option<Vec<ReplyEvent>>>,
}

impl FakeSession {
    fn with_events(events: Vec<ReplyEvent>) -> Self {
        Self {
            prompts: Mutex::new(Vec::new()),
            events: Mutex::new(Some(events)),
        }
    }
}

#[async_trait]
impl AgentSession for FakeSession {
    fn id(&self) -> SessionId {
        SessionId::from_uuid(id(500))
    }

    fn conversation_id(&self) -> Option<String> {
        Some("conversation".to_owned())
    }

    async fn send(&self, prompt: &str) -> Result<ReplyStream, AgentSessionError> {
        self.prompts.lock().unwrap().push(prompt.to_owned());
        let events = self.events.lock().unwrap().take().unwrap_or_else(|| {
            vec![ReplyEvent::Final {
                content: "ack".to_owned(),
            }]
        });
        Ok(Box::new(events.into_iter()))
    }

    async fn shutdown(&self) -> Result<(), AgentSessionError> {
        Ok(())
    }
}

#[derive(Default)]
struct FakeSessionProvider {
    session: Mutex<Option<Arc<FakeSession>>>,
    launches: Mutex<usize>,
}

impl FakeSessionProvider {
    fn with_session(session: Arc<FakeSession>) -> Self {
        Self {
            session: Mutex::new(Some(session)),
            launches: Mutex::new(0),
        }
    }

    fn launches(&self) -> usize {
        *self.launches.lock().unwrap()
    }

    fn session(&self) -> Arc<FakeSession> {
        self.session.lock().unwrap().as_ref().unwrap().clone()
    }
}

#[async_trait]
impl WakeSessionProvider for FakeSessionProvider {
    async fn session_for_wake(
        &self,
        _project: &Project,
        _agent: AgentId,
    ) -> Result<Arc<dyn AgentSession>, WakeError> {
        let mut session = self.session.lock().unwrap();
        if session.is_none() {
            *self.launches.lock().unwrap() += 1;
            *session = Some(Arc::new(FakeSession::default()));
        }
        Ok(session.as_ref().unwrap().clone())
    }
}

fn service(
    inbox: Arc<FakeInbox>,
    turns: Arc<SharedTurnState>,
    tasks: Arc<FakeTaskStore>,
    sessions: Arc<FakeSessionProvider>,
) -> AgentWakeService {
    AgentWakeService::new(inbox, turns.clone(), turns, tasks, sessions, None)
}

#[tokio::test]
async fn wake_if_idle_starts_turn_with_background_completion_prompt() {
    let project = project();
    let owner = agent(1);
    let task_id = task_id(10);
    let inbox = Arc::new(FakeInbox::default());
    let turns = Arc::new(SharedTurnState::default());
    let tasks = Arc::new(FakeTaskStore::default());
    let session = Arc::new(FakeSession::default());
    let sessions = Arc::new(FakeSessionProvider::with_session(session.clone()));
    tasks.insert(completed_task(&project, owner, task_id, "build finished"));
    inbox
        .enqueue_message(owner, completion_item(owner, task_id, ticket(20)))
        .unwrap();

    service(inbox, turns, tasks, sessions)
        .wake_agent(
            &project,
            owner,
            WakeReason::BackgroundCompletion { task_id },
        )
        .await
        .unwrap();

    let prompts = session.prompts.lock().unwrap();
    assert_eq!(prompts.len(), 1);
    assert!(prompts[0].contains(&task_id.to_string()));
    assert!(prompts[0].contains("exit 0"));
    assert!(prompts[0].contains("build finished"));
    assert!(prompts[0].contains("stdout"));
}

#[tokio::test]
async fn owner_busy_does_not_start_concurrent_wake_and_keeps_item_queued() {
    let project = project();
    let owner = agent(1);
    let task_id = task_id(10);
    let inbox = Arc::new(FakeInbox::default());
    let turns = Arc::new(SharedTurnState::default());
    let tasks = Arc::new(FakeTaskStore::default());
    let session = Arc::new(FakeSession::default());
    let sessions = Arc::new(FakeSessionProvider::with_session(session.clone()));
    tasks.insert(completed_task(&project, owner, task_id, "done"));
    inbox
        .enqueue_message(owner, completion_item(owner, task_id, ticket(20)))
        .unwrap();
    turns.force_busy(owner, ticket(99));

    let err = service(inbox.clone(), turns, tasks, sessions)
        .wake_agent(
            &project,
            owner,
            WakeReason::BackgroundCompletion { task_id },
        )
        .await
        .unwrap_err();

    assert_eq!(err, WakeError::AgentBusy { agent_id: owner });
    assert!(session.prompts.lock().unwrap().is_empty());
    assert_eq!(inbox.snapshot(owner).depth, 1);
}

#[tokio::test]
async fn absent_session_is_launched_or_reattached_by_provider() {
    let project = project();
    let owner = agent(1);
    let task_id = task_id(10);
    let inbox = Arc::new(FakeInbox::default());
    let turns = Arc::new(SharedTurnState::default());
    let tasks = Arc::new(FakeTaskStore::default());
    let sessions = Arc::new(FakeSessionProvider::default());
    tasks.insert(completed_task(&project, owner, task_id, "done"));
    inbox
        .enqueue_message(owner, completion_item(owner, task_id, ticket(20)))
        .unwrap();

    service(inbox, turns, tasks, sessions.clone())
        .wake_agent(
            &project,
            owner,
            WakeReason::BackgroundCompletion { task_id },
        )
        .await
        .unwrap();

    assert_eq!(sessions.launches(), 1);
    assert_eq!(sessions.session().prompts.lock().unwrap().len(), 1);
}

#[tokio::test]
async fn completion_is_marked_delivered_after_successful_wake() {
    let project = project();
    let owner = agent(1);
    let task_id = task_id(10);
    let inbox = Arc::new(FakeInbox::default());
    let turns = Arc::new(SharedTurnState::default());
    let tasks = Arc::new(FakeTaskStore::default());
    let sessions = Arc::new(FakeSessionProvider::with_session(Arc::new(
        FakeSession::default(),
    )));
    tasks.insert(completed_task(&project, owner, task_id, "done"));
    inbox
        .enqueue_message(owner, completion_item(owner, task_id, ticket(20)))
        .unwrap();

    service(inbox, turns, tasks.clone(), sessions)
        .wake_agent(
            &project,
            owner,
            WakeReason::BackgroundCompletion { task_id },
        )
        .await
        .unwrap();

    assert_eq!(tasks.delivered(), vec![task_id]);
}

#[tokio::test]
async fn completion_is_marked_delivered_once_send_is_accepted_even_if_drain_fails() {
    let project = project();
    let owner = agent(1);
    let task_id = task_id(10);
    let inbox = Arc::new(FakeInbox::default());
    let turns = Arc::new(SharedTurnState::default());
    let tasks = Arc::new(FakeTaskStore::default());
    let session = Arc::new(FakeSession::with_events(Vec::new()));
    let sessions = Arc::new(FakeSessionProvider::with_session(session.clone()));
    tasks.insert(completed_task(&project, owner, task_id, "done"));
    inbox
        .enqueue_message(owner, completion_item(owner, task_id, ticket(20)))
        .unwrap();

    let err = service(inbox, turns, tasks.clone(), sessions)
        .wake_agent(
            &project,
            owner,
            WakeReason::BackgroundCompletion { task_id },
        )
        .await
        .unwrap_err();

    assert!(matches!(err, WakeError::Session(_)));
    assert_eq!(session.prompts.lock().unwrap().len(), 1);
    assert_eq!(tasks.delivered(), vec![task_id]);
}

#[tokio::test]
async fn wake_drains_exactly_one_item_per_turn() {
    let project = project();
    let owner = agent(1);
    let first = task_id(10);
    let second = task_id(11);
    let inbox = Arc::new(FakeInbox::default());
    let turns = Arc::new(SharedTurnState::default());
    let tasks = Arc::new(FakeTaskStore::default());
    let session = Arc::new(FakeSession::default());
    let sessions = Arc::new(FakeSessionProvider::with_session(session.clone()));
    tasks.insert(completed_task(&project, owner, first, "first"));
    tasks.insert(completed_task(&project, owner, second, "second"));
    inbox
        .enqueue_message(owner, completion_item(owner, first, ticket(20)))
        .unwrap();
    inbox
        .enqueue_message(owner, completion_item(owner, second, ticket(21)))
        .unwrap();

    service(inbox.clone(), turns.clone(), tasks, sessions)
        .wake_agent(
            &project,
            owner,
            WakeReason::BackgroundCompletion { task_id: first },
        )
        .await
        .unwrap();

    assert_eq!(session.prompts.lock().unwrap().len(), 1);
    assert_eq!(inbox.snapshot(owner).depth, 1);
    assert_eq!(turns.ticket_depth(owner), 0);
}