IdeA/crates/application/tests/workstate.rs

//! Tests for the project work-state read model.

use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};

use async_trait::async_trait;

use application::{
    GetProjectWorkState, GetProjectWorkStateInput, LiveSessionKind, LiveSessions,
    StructuredSessions, TerminalSessions, TicketWorkSource, TicketWorkStatus,
};
use domain::mailbox::{
    AgentQueueSnapshot, MailboxError, PendingReply, QueuedTicketSnapshot, Ticket,
};
use domain::ports::{
    AgentContextStore, AgentSession, AgentSessionError, PtyHandle, ReplyStream, StoreError,
};
use domain::{
    Agent, AgentBusyState, AgentId, AgentManifest, AgentOrigin, ConversationId, InputMediator,
    InputSource, ManifestEntry, MarkdownDoc, NodeId, ProfileId, Project, ProjectId, ProjectPath,
    PtySize, RemoteRef, SessionId, SessionKind, TerminalSession, TicketId,
};
use uuid::Uuid;

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

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

fn sid(n: u128) -> SessionId {
    SessionId::from_uuid(Uuid::from_u128(n))
}

fn nid(n: u128) -> NodeId {
    NodeId::from_uuid(Uuid::from_u128(n))
}

fn ticket_id(n: u128) -> domain::TicketId {
    domain::TicketId::from_uuid(Uuid::from_u128(n))
}

fn project() -> Project {
    Project::new(
        ProjectId::from_uuid(Uuid::from_u128(1)),
        "demo",
        ProjectPath::new("/tmp/idea-workstate-test").unwrap(),
        RemoteRef::local(),
        1_700_000_000_000,
    )
    .unwrap()
}

fn agent(n: u128, name: &str) -> Agent {
    Agent::new(
        aid(n),
        name,
        format!("agents/{name}.md"),
        pid(100 + n),
        AgentOrigin::Scratch,
        false,
    )
    .unwrap()
}

fn manifest(agents: &[Agent]) -> AgentManifest {
    AgentManifest::new(1, agents.iter().map(ManifestEntry::from_agent).collect()).unwrap()
}

#[derive(Clone)]
struct FakeContexts {
    manifest: AgentManifest,
}

#[async_trait]
impl AgentContextStore for FakeContexts {
    async fn read_context(
        &self,
        _project: &Project,
        _agent: &AgentId,
    ) -> Result<MarkdownDoc, StoreError> {
        Err(StoreError::NotFound)
    }

    async fn write_context(
        &self,
        _project: &Project,
        _agent: &AgentId,
        _md: &MarkdownDoc,
    ) -> Result<(), StoreError> {
        Ok(())
    }

    async fn load_manifest(&self, _project: &Project) -> Result<AgentManifest, StoreError> {
        Ok(self.manifest.clone())
    }

    async fn save_manifest(
        &self,
        _project: &Project,
        _manifest: &AgentManifest,
    ) -> Result<(), StoreError> {
        Ok(())
    }
}

#[derive(Default)]
struct FakeInput {
    busy: Mutex<HashMap<AgentId, AgentBusyState>>,
}

impl FakeInput {
    fn set_busy(&self, agent: AgentId, busy: AgentBusyState) {
        self.busy.lock().unwrap().insert(agent, busy);
    }
}

impl InputMediator for FakeInput {
    fn enqueue(&self, _agent: AgentId, _ticket: Ticket) -> PendingReply {
        let fut: Pin<Box<dyn Future<Output = Result<String, MailboxError>> + Send>> =
            Box::pin(async { Err(MailboxError::Cancelled) });
        PendingReply::new(fut)
    }

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

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

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

#[derive(Default)]
struct FakeQueue {
    queues: Mutex<HashMap<AgentId, Vec<QueuedTicketSnapshot>>>,
}

impl FakeQueue {
    fn set(&self, agent: AgentId, tickets: Vec<QueuedTicketSnapshot>) {
        self.queues.lock().unwrap().insert(agent, tickets);
    }
}

impl AgentQueueSnapshot for FakeQueue {
    fn queue_for(&self, agent: AgentId) -> Vec<QueuedTicketSnapshot> {
        self.queues
            .lock()
            .unwrap()
            .get(&agent)
            .cloned()
            .unwrap_or_default()
    }
}

fn snapshot(
    id: u128,
    position: u32,
    source: InputSource,
    requester: &str,
    task: &str,
) -> QueuedTicketSnapshot {
    QueuedTicketSnapshot {
        id: TicketId::from_uuid(Uuid::from_u128(id)),
        source,
        conversation: ConversationId::from_uuid(Uuid::from_u128(id + 1000)),
        requester: requester.to_owned(),
        task: task.to_owned(),
        position,
    }
}

struct FakeSession {
    id: SessionId,
}

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

    fn conversation_id(&self) -> Option<String> {
        None
    }

    async fn send(&self, _prompt: &str) -> Result<ReplyStream, AgentSessionError> {
        Ok(Box::new(std::iter::empty()))
    }

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

fn fake_session(id: SessionId) -> Arc<dyn AgentSession> {
    Arc::new(FakeSession { id })
}

fn insert_pty(
    sessions: &TerminalSessions,
    session_id: SessionId,
    agent_id: AgentId,
    node_id: NodeId,
) {
    sessions.insert(
        PtyHandle { session_id },
        TerminalSession::starting(
            session_id,
            node_id,
            ProjectPath::new("/tmp/idea-workstate-test").unwrap(),
            SessionKind::Agent { agent_id },
            PtySize { rows: 24, cols: 80 },
        ),
    );
}

struct Fixture {
    usecase: GetProjectWorkState,
    pty: Arc<TerminalSessions>,
    structured: Arc<StructuredSessions>,
    input: Arc<FakeInput>,
    queue: Arc<FakeQueue>,
    project: Project,
}

fn fixture(agents: &[Agent]) -> Fixture {
    let pty = Arc::new(TerminalSessions::new());
    let structured = Arc::new(StructuredSessions::new());
    let live = Arc::new(LiveSessions::new(Arc::clone(&pty), Arc::clone(&structured)));
    let input = Arc::new(FakeInput::default());
    let input_port = Arc::clone(&input) as Arc<dyn InputMediator>;
    let queue = Arc::new(FakeQueue::default());
    let queue_port = Arc::clone(&queue) as Arc<dyn AgentQueueSnapshot>;
    let usecase = GetProjectWorkState::new(
        Arc::new(FakeContexts {
            manifest: manifest(agents),
        }),
        live,
        input_port,
        queue_port,
    );
    Fixture {
        usecase,
        pty,
        structured,
        input,
        queue,
        project: project(),
    }
}

#[tokio::test]
async fn workstate_lists_manifest_agents_idle_without_live_sessions() {
    let a = agent(10, "alpha");
    let b = agent(20, "beta");
    let f = fixture(&[a.clone(), b.clone()]);

    let out = f
        .usecase
        .execute(GetProjectWorkStateInput { project: f.project })
        .await
        .unwrap();

    assert_eq!(out.agents.len(), 2);
    assert_eq!(out.agents[0].agent_id, a.id);
    assert_eq!(out.agents[0].name, "alpha");
    assert_eq!(out.agents[0].profile_id, a.profile_id);
    assert_eq!(out.agents[0].live, None);
    assert_eq!(out.agents[0].busy, AgentBusyState::Idle);
    assert_eq!(out.agents[1].agent_id, b.id);
}

#[tokio::test]
async fn workstate_attaches_live_pty_session_to_manifest_agent() {
    let a = agent(10, "alpha");
    let f = fixture(std::slice::from_ref(&a));
    insert_pty(&f.pty, sid(1), a.id, nid(100));

    let out = f
        .usecase
        .execute(GetProjectWorkStateInput { project: f.project })
        .await
        .unwrap();

    let live = out.agents[0].live.unwrap();
    assert_eq!(live.session_id, sid(1));
    assert_eq!(live.node_id, nid(100));
    assert_eq!(live.kind, LiveSessionKind::Pty);
}

#[tokio::test]
async fn workstate_attaches_live_structured_session_to_manifest_agent() {
    let a = agent(10, "alpha");
    let f = fixture(std::slice::from_ref(&a));
    f.structured.insert(fake_session(sid(2)), a.id, nid(200));

    let out = f
        .usecase
        .execute(GetProjectWorkStateInput { project: f.project })
        .await
        .unwrap();

    let live = out.agents[0].live.unwrap();
    assert_eq!(live.session_id, sid(2));
    assert_eq!(live.node_id, nid(200));
    assert_eq!(live.kind, LiveSessionKind::Structured);
}

#[tokio::test]
async fn workstate_includes_busy_state_from_input_mediator() {
    let a = agent(10, "alpha");
    let f = fixture(std::slice::from_ref(&a));
    let busy = AgentBusyState::Busy {
        ticket: ticket_id(7),
        since_ms: 1_234,
    };
    f.input.set_busy(a.id, busy);

    let out = f
        .usecase
        .execute(GetProjectWorkStateInput { project: f.project })
        .await
        .unwrap();

    assert_eq!(out.agents[0].busy, busy);
}

#[tokio::test]
async fn workstate_ignores_live_agents_absent_from_manifest() {
    let a = agent(10, "alpha");
    let f = fixture(std::slice::from_ref(&a));
    insert_pty(&f.pty, sid(99), aid(999), nid(999));

    let out = f
        .usecase
        .execute(GetProjectWorkStateInput { project: f.project })
        .await
        .unwrap();

    assert_eq!(out.agents.len(), 1);
    assert_eq!(out.agents[0].agent_id, a.id);
    assert_eq!(out.agents[0].live, None);
}

#[tokio::test]
async fn workstate_agent_without_queue_has_no_tickets() {
    let a = agent(10, "alpha");
    let f = fixture(std::slice::from_ref(&a));

    let out = f
        .usecase
        .execute(GetProjectWorkStateInput { project: f.project })
        .await
        .unwrap();

    assert!(out.agents[0].tickets.is_empty());
}

#[tokio::test]
async fn workstate_lists_two_tickets_in_fifo_order() {
    let a = agent(10, "alpha");
    let f = fixture(std::slice::from_ref(&a));
    let from = aid(20);
    f.queue.set(
        a.id,
        vec![
            snapshot(1, 0, InputSource::agent(from), "Main", "first"),
            snapshot(2, 1, InputSource::agent(from), "Main", "second"),
        ],
    );

    let out = f
        .usecase
        .execute(GetProjectWorkStateInput { project: f.project })
        .await
        .unwrap();

    let tickets = &out.agents[0].tickets;
    assert_eq!(tickets.len(), 2);
    assert_eq!(tickets[0].ticket_id, ticket_id(1));
    assert_eq!(tickets[0].position, 0);
    assert_eq!(tickets[1].ticket_id, ticket_id(2));
    assert_eq!(tickets[1].position, 1);
}

#[tokio::test]
async fn workstate_marks_busy_head_in_progress_and_rest_queued() {
    let a = agent(10, "alpha");
    let f = fixture(std::slice::from_ref(&a));
    let from = aid(20);
    f.queue.set(
        a.id,
        vec![
            snapshot(1, 0, InputSource::agent(from), "Main", "first"),
            snapshot(2, 1, InputSource::agent(from), "Main", "second"),
        ],
    );
    f.input.set_busy(
        a.id,
        AgentBusyState::Busy {
            ticket: ticket_id(1),
            since_ms: 5,
        },
    );

    let out = f
        .usecase
        .execute(GetProjectWorkStateInput { project: f.project })
        .await
        .unwrap();

    let tickets = &out.agents[0].tickets;
    assert_eq!(tickets[0].status, TicketWorkStatus::InProgress);
    assert_eq!(tickets[1].status, TicketWorkStatus::Queued);
}

#[tokio::test]
async fn workstate_marks_all_queued_when_agent_idle() {
    let a = agent(10, "alpha");
    let f = fixture(std::slice::from_ref(&a));
    let from = aid(20);
    f.queue.set(
        a.id,
        vec![snapshot(1, 0, InputSource::agent(from), "Main", "first")],
    );
    // Agent left Idle (default): even the head is only queued, not in-progress.

    let out = f
        .usecase
        .execute(GetProjectWorkStateInput { project: f.project })
        .await
        .unwrap();

    assert_eq!(out.agents[0].tickets[0].status, TicketWorkStatus::Queued);
}

#[tokio::test]
async fn workstate_ignores_queue_for_agent_absent_from_manifest() {
    let a = agent(10, "alpha");
    let f = fixture(std::slice::from_ref(&a));
    // A queue exists for an agent that is not in the manifest: it must not surface.
    f.queue.set(
        aid(999),
        vec![snapshot(1, 0, InputSource::Human, "User", "ghost")],
    );

    let out = f
        .usecase
        .execute(GetProjectWorkStateInput { project: f.project })
        .await
        .unwrap();

    assert_eq!(out.agents.len(), 1);
    assert!(out.agents[0].tickets.is_empty());
}

#[tokio::test]
async fn workstate_truncates_task_preview_and_keeps_original_len() {
    let a = agent(10, "alpha");
    let f = fixture(std::slice::from_ref(&a));
    // 400 chars with runs of whitespace to normalise; well over the 160 cap.
    let long_task = format!("start{}end", "  x  ".repeat(80));
    let original_len = long_task.chars().count();
    f.queue.set(
        a.id,
        vec![snapshot(1, 0, InputSource::Human, "User", &long_task)],
    );

    let out = f
        .usecase
        .execute(GetProjectWorkStateInput { project: f.project })
        .await
        .unwrap();

    let ticket = &out.agents[0].tickets[0];
    assert_eq!(ticket.task_preview.chars().count(), 160);
    assert!(!ticket.task_preview.contains("  "), "whitespace normalised");
    assert_eq!(ticket.task_len, original_len);
    assert!(ticket.task_len > 160);
}

#[tokio::test]
async fn workstate_maps_human_and_agent_ticket_sources() {
    let a = agent(10, "alpha");
    let f = fixture(std::slice::from_ref(&a));
    let from = aid(20);
    f.queue.set(
        a.id,
        vec![
            snapshot(1, 0, InputSource::Human, "User", "from human"),
            snapshot(2, 1, InputSource::agent(from), "Main", "from agent"),
        ],
    );

    let out = f
        .usecase
        .execute(GetProjectWorkStateInput { project: f.project })
        .await
        .unwrap();

    let tickets = &out.agents[0].tickets;
    assert_eq!(tickets[0].source, TicketWorkSource::Human);
    assert_eq!(tickets[0].requester_label, "User");
    assert_eq!(
        tickets[1].source,
        TicketWorkSource::Agent { agent_id: from }
    );
    assert_eq!(tickets[1].requester_label, "Main");
}