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::{
    AgentRestartScheduled, AgentSelfRestartScheduler, AppError, AttachLiveAgent,
    AttachLiveAgentInput, CellAgentPresentation, CloseTerminal, LiveRuntimeState, LiveSessionKind,
    LiveSessions, RestartCellAgentRuntime, RestartCellAgentRuntimeInput, StopAgentExecutionCascade,
    StopLiveAgent, StopLiveAgentInput, StructuredSessions, TerminalSessions,
};
use domain::ports::{
    AgentSession, AgentSessionError, ExitStatus, OutputStream, PtyError, PtyHandle, PtyPort,
    ReplyStream, SpawnSpec,
};
use domain::Project;
use domain::{
    AgentId, LayoutId, 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 wait(&self, _handle: &PtyHandle) -> Result<ExitStatus, PtyError> {
        Ok(ExitStatus { code: Some(0) })
    }
    fn try_wait(&self, _handle: &PtyHandle) -> Result<Option<ExitStatus>, PtyError> {
        Ok(Some(ExitStatus { code: Some(0) }))
    }
    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 stop/cancel calls.
struct FakeSession {
    id: SessionId,
    shutdown_called: Arc<AtomicBool>,
    cancel_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 cancel_current_turn(&self) -> Result<(), AgentSessionError> {
        self.cancel_called.store(true, Ordering::SeqCst);
        Ok(())
    }
    async fn shutdown(&self) -> Result<(), AgentSessionError> {
        self.shutdown_called.store(true, Ordering::SeqCst);
        Ok(())
    }
}

#[derive(Default)]
struct FakeCascade {
    calls: Mutex<Vec<(ProjectId, AgentId)>>,
}

#[derive(Default)]
struct FakeRestartScheduler {
    calls: Mutex<Vec<(ProjectId, AgentId, NodeId)>>,
}

impl AgentSelfRestartScheduler for FakeRestartScheduler {
    fn schedule(
        &self,
        _project: &Project,
        _agent: AgentId,
        _reason: Option<String>,
    ) -> Result<AgentRestartScheduled, AppError> {
        unreachable!("cell restart must use the cell-aware scheduler seam")
    }

    fn schedule_cell(
        &self,
        project: &Project,
        agent: AgentId,
        node: NodeId,
        _reason: Option<String>,
    ) -> Result<AgentRestartScheduled, AppError> {
        self.calls.lock().unwrap().push((project.id, agent, node));
        Ok(AgentRestartScheduled {
            restart_id: "restart-cell-test".to_owned(),
        })
    }
}

#[async_trait]
impl StopAgentExecutionCascade for FakeCascade {
    async fn stop_agent_execution_children(&self, project: &Project, agent_id: AgentId) {
        self.calls.lock().unwrap().push((project.id, agent_id));
    }
}

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

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

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),
    ));
    let restart_scheduler = Arc::new(FakeRestartScheduler::default());
    Fixture {
        attach: AttachLiveAgent::new(Arc::clone(&live)),
        stop: StopLiveAgent::new(Arc::clone(&live), close),
        pty,
        structured,
        pty_port,
        restart: RestartCellAgentRuntime::new(
            Arc::clone(&live),
            Arc::clone(&restart_scheduler) as Arc<dyn AgentSelfRestartScheduler>,
        ),
        restart_scheduler,
    }
}

fn restart_input(
    agent_id: AgentId,
    node_id: NodeId,
    presentation: CellAgentPresentation,
) -> RestartCellAgentRuntimeInput {
    RestartCellAgentRuntimeInput {
        project: project(),
        node_id,
        agent_id,
        presentation,
        layout_id: Some(LayoutId::from_uuid(Uuid::from_u128(500))),
    }
}

fn insert_pty(
    sessions: &TerminalSessions,
    session_id: SessionId,
    agent_id: AgentId,
    node_id: NodeId,
) {
    sessions.insert_in_project(
        project().id,
        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>, Arc<AtomicBool>) {
    let shutdown_called = Arc::new(AtomicBool::new(false));
    let cancel_called = Arc::new(AtomicBool::new(false));
    sessions.insert_in_project(
        project().id,
        Arc::new(FakeSession {
            id: session_id,
            shutdown_called: Arc::clone(&shutdown_called),
            cancel_called: Arc::clone(&cancel_called),
        }),
        agent_id,
        node_id,
    );
    (shutdown_called, cancel_called)
}

// ---------------------------------------------------------------------------
// 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_in_project(project().id, &a),
        Some(nid(200))
    );
    assert_eq!(
        f.pty.session_for_agent_in_project(project().id, &a),
        Some(sid(1))
    );
}

#[test]
fn attach_structured_rebinds_node() {
    let f = fixture();
    let a = aid(10);
    let _flags = 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_in_project(project().id, &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");
}

// ---------------------------------------------------------------------------
// RestartCellAgentRuntime
// ---------------------------------------------------------------------------

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

    let out = f
        .restart
        .execute(restart_input(agent, node, CellAgentPresentation::Tui))
        .unwrap();

    assert_eq!(out.restart_id, "restart-cell-test");
    assert_eq!(out.agent_id, agent);
    assert_eq!(out.node_id, node);
    assert_eq!(out.kind, LiveSessionKind::Pty);
    assert_eq!(out.presentation, CellAgentPresentation::Tui);
    assert_eq!(
        *f.restart_scheduler.calls.lock().unwrap(),
        vec![(project().id, agent, node)]
    );
}

#[test]
fn restart_cell_accepts_structured_custom_cli_and_headless_hints() {
    for presentation in [
        CellAgentPresentation::CustomCli,
        CellAgentPresentation::Headless,
    ] {
        let f = fixture();
        let agent = aid(10);
        let node = nid(100);
        insert_structured(&f.structured, sid(2), agent, node);

        let out = f
            .restart
            .execute(restart_input(agent, node, presentation))
            .unwrap();

        assert_eq!(out.kind, LiveSessionKind::Structured);
        assert_eq!(out.presentation, presentation);
        assert_eq!(f.restart_scheduler.calls.lock().unwrap().len(), 1);
    }
}

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

    let error = f
        .restart
        .execute(restart_input(
            agent,
            nid(999),
            CellAgentPresentation::Unknown,
        ))
        .unwrap_err();

    assert_eq!(error.code(), "AGENT_NOT_BOUND_TO_CELL");
    assert!(f.restart_scheduler.calls.lock().unwrap().is_empty());
}

#[test]
fn restart_cell_rejects_missing_runtime_without_scheduling() {
    let f = fixture();
    let error = f
        .restart
        .execute(restart_input(
            aid(999),
            nid(100),
            CellAgentPresentation::Unknown,
        ))
        .unwrap_err();

    assert_eq!(error.code(), "RUNTIME_NOT_FOUND");
    assert!(f.restart_scheduler.calls.lock().unwrap().is_empty());
}

#[test]
fn restart_cell_rejects_agent_visible_in_multiple_live_cells() {
    let f = fixture();
    let agent = aid(10);
    insert_pty(&f.pty, sid(1), agent, nid(100));
    insert_structured(&f.structured, sid(2), agent, nid(200));

    let error = f
        .restart
        .execute(restart_input(agent, nid(100), CellAgentPresentation::Tui))
        .unwrap_err();

    assert_eq!(error.code(), "CONFLICT_AGENT_VISIBLE_ELSEWHERE");
    assert!(f.restart_scheduler.calls.lock().unwrap().is_empty());
}

// ---------------------------------------------------------------------------
// 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,
            reason: None,
        })
        .await
        .unwrap();

    assert_eq!(out.agent_id, a);
    assert!(!out.already_stopped);
    assert_eq!(out.stopped[0].session_id, sid(1));
    assert_eq!(out.stopped[0].node_id, nid(100));
    assert_eq!(out.stopped[0].kind, LiveSessionKind::Pty);
    assert_eq!(out.stopped[0].state, LiveRuntimeState::Stopped);
    // 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_in_project(project().id, &a), None);
}

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

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

    assert_eq!(out.stopped[0].session_id, sid(2));
    assert_eq!(out.stopped[0].kind, LiveSessionKind::Structured);
    assert_eq!(out.stopped[0].state, LiveRuntimeState::Stopped);
    assert!(!cancel_called.load(Ordering::SeqCst));
    assert!(
        shutdown_called.load(Ordering::SeqCst),
        "session.shutdown() must be called"
    );
    assert!(f.structured.is_empty(), "live session is removed");
    assert_eq!(
        f.structured
            .session_id_for_agent_in_project(project().id, &a),
        None
    );
    // No PTY was touched.
    assert!(f.pty_port.kills().is_empty());
}

#[tokio::test]
async fn stop_structured_runs_execution_cascade_best_effort() {
    let f = fixture();
    let a = aid(10);
    let (shutdown_called, cancel_called) = insert_structured(&f.structured, sid(2), a, nid(100));
    let cascade = Arc::new(FakeCascade::default());
    let stop = f
        .stop
        .with_execution_cascade(Arc::clone(&cascade) as Arc<dyn StopAgentExecutionCascade>);

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

    assert_eq!(out.stopped[0].session_id, sid(2));
    assert_eq!(*cascade.calls.lock().unwrap(), vec![(project().id, a)]);
    assert!(
        !cancel_called.load(Ordering::SeqCst),
        "session.cancel_current_turn() is not the runtime stop primitive"
    );
    assert!(
        shutdown_called.load(Ordering::SeqCst),
        "session.shutdown() must be called"
    );
    assert!(f.structured.is_empty(), "live session is removed");
}

#[tokio::test]
async fn stop_removes_every_runtime_for_the_agent_across_registries() {
    let f = fixture();
    let agent = aid(10);
    insert_pty(&f.pty, sid(1), agent, nid(100));
    let (shutdown_called, _) = insert_structured(&f.structured, sid(2), agent, nid(200));

    let out = f
        .stop
        .execute(StopLiveAgentInput {
            project: project(),
            agent_id: agent,
            reason: Some("user requested".to_owned()),
        })
        .await
        .unwrap();

    assert_eq!(out.stopped.len(), 2);
    assert!(f.pty.is_empty());
    assert!(f.structured.is_empty());
    assert_eq!(f.pty_port.kills(), vec![sid(1)]);
    assert!(shutdown_called.load(Ordering::SeqCst));
}

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

    let first = f
        .stop
        .execute(StopLiveAgentInput {
            project: project(),
            agent_id: a,
            reason: None,
        })
        .await
        .unwrap();
    let second = f
        .stop
        .execute(StopLiveAgentInput {
            project: project(),
            agent_id: a,
            reason: None,
        })
        .await
        .unwrap();

    assert!(!first.already_stopped);
    assert!(second.already_stopped);
    assert!(second.stopped.is_empty());
    assert_eq!(f.pty_port.kills(), vec![sid(1)], "kill is not repeated");
}

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

    let out = f
        .stop
        .execute(StopLiveAgentInput {
            project: project(),
            agent_id: aid(999),
            reason: None,
        })
        .await
        .expect("absent agent is an idempotent success");
    assert!(out.already_stopped);
    assert!(out.stopped.is_empty());
    assert!(f.pty_port.kills().is_empty(), "nothing killed");
}