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,
};
use domain::mailbox::{MailboxError, PendingReply, Ticket};
use domain::ports::{
    AgentContextStore, AgentSession, AgentSessionError, PtyHandle, ReplyStream, StoreError,
};
use domain::{
    Agent, AgentBusyState, AgentId, AgentManifest, AgentOrigin, InputMediator, ManifestEntry,
    MarkdownDoc, NodeId, ProfileId, Project, ProjectId, ProjectPath, PtySize, RemoteRef, SessionId,
    SessionKind, TerminalSession,
};
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)
    }
}

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>,
    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 usecase = GetProjectWorkState::new(
        Arc::new(FakeContexts {
            manifest: manifest(agents),
        }),
        live,
        input_port,
    );
    Fixture {
        usecase,
        pty,
        structured,
        input,
        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);
}