IdeA/crates/application/tests/workstate_actions.rs

//! Tests for the Lot D agent-level controlled actions (`AttachLiveAgent`,
//! `StopLiveAgent`) over the live-session registries.
//!
//! Every port is faked in-memory so the use cases run without a real PTY or
//! agent-session backend: [`FakePty`] records `kill`s, [`FakeSession`] records
//! `shutdown`s.

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};

use async_trait::async_trait;

use application::{
    AttachLiveAgent, AttachLiveAgentInput, CloseTerminal, LiveSessionKind, LiveSessions,
    StopLiveAgent, StopLiveAgentInput, StructuredSessions, TerminalSessions,
};
use domain::ports::{
    AgentSession, AgentSessionError, ExitStatus, OutputStream, PtyError, PtyHandle, PtyPort,
    ReplyStream, SpawnSpec,
};
use domain::Project;
use domain::{
    AgentId, NodeId, ProjectId, ProjectPath, PtySize, RemoteRef, SessionId, SessionKind,
    TerminalSession,
};
use uuid::Uuid;

// ---------------------------------------------------------------------------
// ids / fixtures
// ---------------------------------------------------------------------------

fn aid(n: u128) -> AgentId {
    AgentId::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 project() -> Project {
    Project::new(
        ProjectId::from_uuid(Uuid::from_u128(1)),
        "demo",
        ProjectPath::new("/tmp/idea-workstate-actions").unwrap(),
        RemoteRef::local(),
        1_700_000_000_000,
    )
    .unwrap()
}

// ---------------------------------------------------------------------------
// Fakes
// ---------------------------------------------------------------------------

/// A recording [`PtyPort`] whose `kill` records the killed session id.
#[derive(Default)]
struct FakePty {
    kills: Mutex<Vec<SessionId>>,
}

impl FakePty {
    fn kills(&self) -> Vec<SessionId> {
        self.kills.lock().unwrap().clone()
    }
}

#[async_trait]
impl PtyPort for FakePty {
    async fn spawn(&self, _spec: SpawnSpec, _size: PtySize) -> Result<PtyHandle, PtyError> {
        unreachable!("Lot D never spawns")
    }
    fn write(&self, _handle: &PtyHandle, _data: &[u8]) -> Result<(), PtyError> {
        Ok(())
    }
    fn resize(&self, _handle: &PtyHandle, _size: PtySize) -> Result<(), PtyError> {
        Ok(())
    }
    fn subscribe_output(&self, _handle: &PtyHandle) -> Result<OutputStream, PtyError> {
        Ok(Box::new(std::iter::empty()))
    }
    fn scrollback(&self, _handle: &PtyHandle) -> Result<Vec<u8>, PtyError> {
        Ok(Vec::new())
    }
    async fn kill(&self, handle: &PtyHandle) -> Result<ExitStatus, PtyError> {
        self.kills.lock().unwrap().push(handle.session_id);
        Ok(ExitStatus { code: Some(0) })
    }
}

/// A fake structured [`AgentSession`] recording whether `shutdown` was called.
struct FakeSession {
    id: SessionId,
    shutdown_called: Arc<AtomicBool>,
}

#[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> {
        self.shutdown_called.store(true, Ordering::SeqCst);
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// wiring helpers
// ---------------------------------------------------------------------------

struct Fixture {
    attach: AttachLiveAgent,
    stop: StopLiveAgent,
    pty: Arc<TerminalSessions>,
    structured: Arc<StructuredSessions>,
    pty_port: Arc<FakePty>,
}

fn fixture() -> 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 pty_port = Arc::new(FakePty::default());
    let close = Arc::new(CloseTerminal::new(
        Arc::clone(&pty_port) as Arc<dyn PtyPort>,
        Arc::clone(&pty),
    ));
    Fixture {
        attach: AttachLiveAgent::new(Arc::clone(&live)),
        stop: StopLiveAgent::new(Arc::clone(&live), close),
        pty,
        structured,
        pty_port,
    }
}

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-actions").unwrap(),
            SessionKind::Agent { agent_id },
            PtySize { rows: 24, cols: 80 },
        ),
    );
}

fn insert_structured(
    sessions: &StructuredSessions,
    session_id: SessionId,
    agent_id: AgentId,
    node_id: NodeId,
) -> Arc<AtomicBool> {
    let flag = Arc::new(AtomicBool::new(false));
    sessions.insert(
        Arc::new(FakeSession {
            id: session_id,
            shutdown_called: Arc::clone(&flag),
        }),
        agent_id,
        node_id,
    );
    flag
}

// ---------------------------------------------------------------------------
// AttachLiveAgent
// ---------------------------------------------------------------------------

#[test]
fn attach_pty_rebinds_node_without_changing_session() {
    let f = fixture();
    let a = aid(10);
    insert_pty(&f.pty, sid(1), a, nid(100));

    let out = f
        .attach
        .execute(AttachLiveAgentInput {
            project: project(),
            agent_id: a,
            node_id: nid(200),
        })
        .unwrap();

    assert_eq!(out.agent_id, a);
    assert_eq!(out.session_id, sid(1), "session id is preserved on rebind");
    assert_eq!(out.node_id, nid(200), "view rebound to the new node");
    assert_eq!(out.kind, LiveSessionKind::Pty);
    // The registry reflects the new host node, same session.
    assert_eq!(f.pty.node_for_agent(&a), Some(nid(200)));
    assert_eq!(f.pty.session_for_agent(&a), Some(sid(1)));
}

#[test]
fn attach_structured_rebinds_node() {
    let f = fixture();
    let a = aid(10);
    let _flag = insert_structured(&f.structured, sid(2), a, nid(100));

    let out = f
        .attach
        .execute(AttachLiveAgentInput {
            project: project(),
            agent_id: a,
            node_id: nid(300),
        })
        .unwrap();

    assert_eq!(out.session_id, sid(2));
    assert_eq!(out.node_id, nid(300));
    assert_eq!(out.kind, LiveSessionKind::Structured);
    assert_eq!(f.structured.node_for_agent(&a), Some(nid(300)));
}

#[test]
fn attach_unknown_agent_is_not_found() {
    let f = fixture();

    let err = f
        .attach
        .execute(AttachLiveAgentInput {
            project: project(),
            agent_id: aid(999),
            node_id: nid(1),
        })
        .expect_err("absent agent ⇒ NOT_FOUND");

    assert!(
        matches!(err, application::AppError::NotFound(_)),
        "got {err:?}"
    );
}

#[test]
fn attach_same_node_is_idempotent() {
    let f = fixture();
    let a = aid(10);
    insert_pty(&f.pty, sid(1), a, nid(100));

    let first = f
        .attach
        .execute(AttachLiveAgentInput {
            project: project(),
            agent_id: a,
            node_id: nid(100),
        })
        .unwrap();
    let second = f
        .attach
        .execute(AttachLiveAgentInput {
            project: project(),
            agent_id: a,
            node_id: nid(100),
        })
        .unwrap();

    // Re-attaching to the same node is a no-op: same session, same node, no spawn.
    assert_eq!(first.session_id, second.session_id);
    assert_eq!(first.node_id, nid(100));
    assert_eq!(second.node_id, nid(100));
    assert_eq!(f.pty.len(), 1, "no extra session created");
}

// ---------------------------------------------------------------------------
// StopLiveAgent
// ---------------------------------------------------------------------------

#[tokio::test]
async fn stop_pty_kills_and_removes_session() {
    let f = fixture();
    let a = aid(10);
    insert_pty(&f.pty, sid(1), a, nid(100));

    let out = f
        .stop
        .execute(StopLiveAgentInput {
            project: project(),
            agent_id: a,
        })
        .await
        .unwrap();

    assert_eq!(out.agent_id, a);
    assert_eq!(out.session_id, sid(1));
    assert_eq!(out.kind, LiveSessionKind::Pty);
    // Delegated to the close primitive: process killed and registry emptied.
    assert_eq!(f.pty_port.kills(), vec![sid(1)]);
    assert!(f.pty.is_empty(), "live session removed from the registry");
    assert_eq!(f.pty.session_for_agent(&a), None);
}

#[tokio::test]
async fn stop_structured_shuts_down_and_removes_session() {
    let f = fixture();
    let a = aid(10);
    let flag = insert_structured(&f.structured, sid(2), a, nid(100));

    let out = f
        .stop
        .execute(StopLiveAgentInput {
            project: project(),
            agent_id: a,
        })
        .await
        .unwrap();

    assert_eq!(out.session_id, sid(2));
    assert_eq!(out.kind, LiveSessionKind::Structured);
    assert!(flag.load(Ordering::SeqCst), "session.shutdown() was called");
    assert!(f.structured.is_empty(), "live session removed");
    assert_eq!(f.structured.session_id_for_agent(&a), None);
    // No PTY was touched.
    assert!(f.pty_port.kills().is_empty());
}

#[tokio::test]
async fn stop_unknown_agent_is_not_found() {
    let f = fixture();

    let err = f
        .stop
        .execute(StopLiveAgentInput {
            project: project(),
            agent_id: aid(999),
        })
        .await
        .expect_err("absent agent ⇒ NOT_FOUND");

    assert!(
        matches!(err, application::AppError::NotFound(_)),
        "got {err:?}"
    );
    assert!(f.pty_port.kills().is_empty(), "nothing killed");
}