IdeA/crates/application/tests/snapshot_running_agents.rs

//! T5 tests for [`SnapshotRunningAgents`].
//!
//! At close time, before the global PTY kill, the use case must freeze on each
//! agent-bearing leaf whether that agent's PTY was still live
//! (`agent_was_running`). Liveness is derived purely from the live-session
//! registry (a [`LiveAgentRegistry`]), never from CLI parsing.
//!
//! Every port is faked in-memory. We assert the persisted `layouts.json` flags
//! and the ordering guarantee (snapshot reads liveness BEFORE the kill).

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

use async_trait::async_trait;
use domain::layout::Workspace;
use domain::ports::{DirEntry, FileSystem, FsError, ProjectStore, RemotePath, StoreError};
use domain::{
    AgentId, Direction, LayoutId, LayoutNode, LayoutTree, LeafCell, NodeId, Project, ProjectId,
    ProjectPath, RemoteRef, SplitContainer, WeightedChild,
};
use uuid::Uuid;

use application::{LiveAgentRegistry, SnapshotRunningAgents, SnapshotRunningAgentsInput};

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

#[derive(Default)]
struct FakeFsInner {
    files: HashMap<String, Vec<u8>>,
    dirs: HashSet<String>,
}

#[derive(Default, Clone)]
struct FakeFs(Arc<Mutex<FakeFsInner>>);

impl FakeFs {
    fn read_file(&self, path: &str) -> Option<Vec<u8>> {
        self.0.lock().unwrap().files.get(path).cloned()
    }
    fn put(&self, path: &str, data: &[u8]) {
        self.0
            .lock()
            .unwrap()
            .files
            .insert(path.to_owned(), data.to_vec());
    }
    fn writes(&self) -> usize {
        self.0.lock().unwrap().files.len()
    }
}

#[async_trait]
impl FileSystem for FakeFs {
    async fn read(&self, path: &RemotePath) -> Result<Vec<u8>, FsError> {
        self.0
            .lock()
            .unwrap()
            .files
            .get(path.as_str())
            .cloned()
            .ok_or_else(|| FsError::NotFound(path.as_str().to_owned()))
    }
    async fn write(&self, path: &RemotePath, data: &[u8]) -> Result<(), FsError> {
        self.0
            .lock()
            .unwrap()
            .files
            .insert(path.as_str().to_owned(), data.to_vec());
        Ok(())
    }
    async fn exists(&self, path: &RemotePath) -> Result<bool, FsError> {
        let inner = self.0.lock().unwrap();
        Ok(inner.files.contains_key(path.as_str()) || inner.dirs.contains(path.as_str()))
    }
    async fn create_dir_all(&self, path: &RemotePath) -> Result<(), FsError> {
        self.0.lock().unwrap().dirs.insert(path.as_str().to_owned());
        Ok(())
    }
    async fn list(&self, _path: &RemotePath) -> Result<Vec<DirEntry>, FsError> {
        Ok(Vec::new())
    }
    async fn symlink(&self, _src: &RemotePath, _dst: &RemotePath) -> Result<(), FsError> {
        Ok(())
    }
}

#[derive(Default)]
struct FakeStoreInner {
    projects: Vec<Project>,
}

#[derive(Default, Clone)]
struct FakeStore(Arc<Mutex<FakeStoreInner>>);

#[async_trait]
impl ProjectStore for FakeStore {
    async fn list_projects(&self) -> Result<Vec<Project>, StoreError> {
        Ok(self.0.lock().unwrap().projects.clone())
    }
    async fn load_project(&self, id: ProjectId) -> Result<Project, StoreError> {
        self.0
            .lock()
            .unwrap()
            .projects
            .iter()
            .find(|p| p.id == id)
            .cloned()
            .ok_or(StoreError::NotFound)
    }
    async fn save_project(&self, project: &Project) -> Result<(), StoreError> {
        self.0.lock().unwrap().projects.push(project.clone());
        Ok(())
    }
    async fn save_workspace(&self, _w: &Workspace) -> Result<(), StoreError> {
        Ok(())
    }
    async fn load_workspace(&self) -> Result<Workspace, StoreError> {
        Ok(Workspace::default())
    }
}

/// A controllable liveness registry. Liveness is keyed on the hosting **node**
/// (the leaf), matching the snapshot's per-cell query: a leaf is "live" iff its
/// node id is in the set. An optional agent set backs the legacy
/// `is_agent_live` query (unused by the snapshot now, kept for the trait).
#[derive(Default, Clone)]
struct FakeLive {
    nodes: Arc<Mutex<HashSet<NodeId>>>,
    agents: Arc<Mutex<HashSet<AgentId>>>,
}

impl FakeLive {
    /// Seeds the set of live **nodes** (the cells whose PTY is alive).
    fn with_nodes(nodes: &[NodeId]) -> Self {
        Self {
            nodes: Arc::new(Mutex::new(nodes.iter().copied().collect())),
            agents: Arc::new(Mutex::new(HashSet::new())),
        }
    }
    /// Simulates a PTY kill: forget every live node/agent.
    fn kill_all(&self) {
        self.nodes.lock().unwrap().clear();
        self.agents.lock().unwrap().clear();
    }
}

impl LiveAgentRegistry for FakeLive {
    fn is_agent_live(&self, _project_id: domain::ProjectId, agent_id: &AgentId) -> bool {
        self.agents.lock().unwrap().contains(agent_id)
    }
    fn is_node_live(&self, node_id: &NodeId) -> bool {
        self.nodes.lock().unwrap().contains(node_id)
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

const ROOT: &str = "/home/me/proj";
const LAYOUTS_PATH: &str = "/home/me/proj/.ideai/layouts.json";

fn pid(n: u128) -> ProjectId {
    ProjectId::from_uuid(Uuid::from_u128(n))
}
fn nid(n: u128) -> NodeId {
    NodeId::from_uuid(Uuid::from_u128(n))
}
fn aid(n: u128) -> AgentId {
    AgentId::from_uuid(Uuid::from_u128(n))
}
fn lid(n: u128) -> LayoutId {
    LayoutId::from_uuid(Uuid::from_u128(n))
}

fn agent_leaf(node: NodeId, agent: Option<AgentId>) -> LeafCell {
    LeafCell {
        id: node,
        session: None,
        agent,
        conversation_id: None,
        engine_session_id: None,
        agent_was_running: false,
        preferred_view: domain::PreferredView::Tui,
    }
}

async fn register_project(store: &FakeStore, id: ProjectId) {
    let project = Project::new(
        id,
        "Demo",
        ProjectPath::new(ROOT).unwrap(),
        RemoteRef::Local,
        0,
    )
    .unwrap();
    store.save_project(&project).await.unwrap();
}

/// Seeds a valid `layouts.json` with one active layout holding `tree`.
fn seed_layouts(fs: &FakeFs, id: LayoutId, tree: &LayoutTree) {
    let doc = serde_json::json!({
        "version": 1,
        "activeId": id.to_string(),
        "layouts": [ { "id": id.to_string(), "name": "Default", "tree": tree } ],
    });
    fs.put(LAYOUTS_PATH, &serde_json::to_vec(&doc).unwrap());
}

fn doc_json(fs: &FakeFs) -> serde_json::Value {
    serde_json::from_slice(&fs.read_file(LAYOUTS_PATH).expect("layouts.json present")).unwrap()
}

/// Walks the active layout tree and returns `agent_was_running` for the leaf
/// `node`, or `None` if absent.
fn was_running(fs: &FakeFs, node: NodeId) -> Option<bool> {
    let doc = doc_json(fs);
    let tree: LayoutTree =
        serde_json::from_value(doc["layouts"][0]["tree"].clone()).expect("tree parseable");
    fn find(node: &LayoutNode, target: NodeId) -> Option<bool> {
        match node {
            LayoutNode::Leaf(l) if l.id == target => Some(l.agent_was_running),
            LayoutNode::Leaf(_) => None,
            LayoutNode::CustomPluginLayout(_) => None,
            LayoutNode::Split(s) => s.children.iter().find_map(|c| find(&c.node, target)),
            LayoutNode::Grid(g) => g.cells.iter().find_map(|c| find(&c.node, target)),
        }
    }
    find(&tree.root, node)
}

fn make_use_case(store: &FakeStore, fs: &FakeFs, live: &FakeLive) -> SnapshotRunningAgents {
    SnapshotRunningAgents::new(
        Arc::new(store.clone()) as Arc<dyn ProjectStore>,
        Arc::new(fs.clone()) as Arc<dyn FileSystem>,
        Arc::new(live.clone()) as Arc<dyn LiveAgentRegistry>,
    )
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

/// A live agent's leaf is persisted with `agent_was_running = true`.
#[tokio::test]
async fn live_agent_is_marked_running() {
    let store = FakeStore::default();
    let fs = FakeFs::default();
    register_project(&store, pid(1)).await;

    let leaf = nid(10);
    let agent = aid(100);
    seed_layouts(
        &fs,
        lid(1),
        &LayoutTree::single(agent_leaf(leaf, Some(agent))),
    );

    let live = FakeLive::with_nodes(&[leaf]);
    let uc = make_use_case(&store, &fs, &live);

    let out = uc
        .execute(SnapshotRunningAgentsInput { project_id: pid(1) })
        .await
        .unwrap();

    assert_eq!(out.running, 1);
    assert_eq!(out.stopped, 0);
    assert_eq!(was_running(&fs, leaf), Some(true));
}

/// An agent whose PTY has already exited is persisted with `false`.
#[tokio::test]
async fn stopped_agent_is_marked_not_running() {
    let store = FakeStore::default();
    let fs = FakeFs::default();
    register_project(&store, pid(1)).await;

    let leaf = nid(10);
    let agent = aid(100);
    seed_layouts(
        &fs,
        lid(1),
        &LayoutTree::single(agent_leaf(leaf, Some(agent))),
    );

    // Registry is empty: the agent has no live session.
    let live = FakeLive::default();
    let uc = make_use_case(&store, &fs, &live);

    let out = uc
        .execute(SnapshotRunningAgentsInput { project_id: pid(1) })
        .await
        .unwrap();

    assert_eq!(out.running, 0);
    assert_eq!(out.stopped, 1);
    assert_eq!(was_running(&fs, leaf), Some(false));
}

/// Mixed tree (split with one live agent, one stopped agent, one plain leaf):
/// each agent leaf gets the right flag; the non-agent leaf is untouched.
#[tokio::test]
async fn mixed_tree_flags_each_agent_independently() {
    let store = FakeStore::default();
    let fs = FakeFs::default();
    register_project(&store, pid(1)).await;

    let live_leaf = nid(10);
    let dead_leaf = nid(11);
    let plain_leaf = nid(12);
    let live_agent = aid(100);
    let dead_agent = aid(101);

    let tree = LayoutTree::new(LayoutNode::Split(SplitContainer {
        id: nid(1),
        direction: Direction::Row,
        children: vec![
            WeightedChild {
                node: LayoutNode::Leaf(agent_leaf(live_leaf, Some(live_agent))),
                weight: 1.0,
            },
            WeightedChild {
                node: LayoutNode::Leaf(agent_leaf(dead_leaf, Some(dead_agent))),
                weight: 1.0,
            },
            WeightedChild {
                node: LayoutNode::Leaf(agent_leaf(plain_leaf, None)),
                weight: 1.0,
            },
        ],
    }));
    seed_layouts(&fs, lid(1), &tree);

    let live = FakeLive::with_nodes(&[live_leaf]);
    let uc = make_use_case(&store, &fs, &live);

    let out = uc
        .execute(SnapshotRunningAgentsInput { project_id: pid(1) })
        .await
        .unwrap();

    assert_eq!(out.running, 1);
    assert_eq!(out.stopped, 1);
    assert_eq!(was_running(&fs, live_leaf), Some(true));
    assert_eq!(was_running(&fs, dead_leaf), Some(false));
    assert_eq!(was_running(&fs, plain_leaf), Some(false)); // default, untouched
}

/// Ordering guarantee: the snapshot reads liveness BEFORE the kill. Running the
/// snapshot on the live registry persists `true`; running it AFTER `kill_all`
/// (simulating a kill-then-snapshot mistake) would persist `false`.
#[tokio::test]
async fn snapshot_reads_liveness_before_kill() {
    let store = FakeStore::default();
    let fs = FakeFs::default();
    register_project(&store, pid(1)).await;

    let leaf = nid(10);
    let agent = aid(100);
    seed_layouts(
        &fs,
        lid(1),
        &LayoutTree::single(agent_leaf(leaf, Some(agent))),
    );

    let live = FakeLive::with_nodes(&[leaf]);
    let uc = make_use_case(&store, &fs, &live);

    // Correct order (composition root contract): snapshot THEN kill.
    uc.execute(SnapshotRunningAgentsInput { project_id: pid(1) })
        .await
        .unwrap();
    live.kill_all();
    assert_eq!(
        was_running(&fs, leaf),
        Some(true),
        "snapshot taken before kill must record running"
    );

    // Re-seed and demonstrate the opposite order yields `false` — proving the
    // flag is sensitive to registry state at call time (hence order matters).
    seed_layouts(
        &fs,
        lid(1),
        &LayoutTree::single(agent_leaf(leaf, Some(agent))),
    );
    uc.execute(SnapshotRunningAgentsInput { project_id: pid(1) })
        .await
        .unwrap();
    assert_eq!(
        was_running(&fs, leaf),
        Some(false),
        "snapshot taken after kill records not-running"
    );
}

/// A layout with no agent leaf is a no-op: nothing is written.
#[tokio::test]
async fn no_agent_leaf_is_noop() {
    let store = FakeStore::default();
    let fs = FakeFs::default();
    register_project(&store, pid(1)).await;

    let leaf = nid(10);
    seed_layouts(&fs, lid(1), &LayoutTree::single(agent_leaf(leaf, None)));
    let writes_before = fs.writes();

    let live = FakeLive::default();
    let uc = make_use_case(&store, &fs, &live);

    let out = uc
        .execute(SnapshotRunningAgentsInput { project_id: pid(1) })
        .await
        .unwrap();

    assert_eq!(out.running, 0);
    assert_eq!(out.stopped, 0);
    // No agent leaf → no persistence triggered (file count unchanged).
    assert_eq!(fs.writes(), writes_before);
    assert_eq!(was_running(&fs, leaf), Some(false));
}

/// Duplicate leaves for the **same** agent (the singleton-invariant case): only
/// the cell whose PTY is actually live is flagged running. The other leaf —
/// pinning the same agent but never launched (or refused by the guard) — stays
/// `false`. This is exactly why liveness is keyed on the node, not the agent:
/// an agent-keyed query would wrongly mark BOTH leaves as running.
#[tokio::test]
async fn duplicate_leaves_same_agent_only_live_node_is_running() {
    let store = FakeStore::default();
    let fs = FakeFs::default();
    register_project(&store, pid(1)).await;

    let leaf_a = nid(10);
    let leaf_b = nid(11);
    let agent = aid(100); // SAME agent on both leaves

    let tree = LayoutTree::new(LayoutNode::Split(SplitContainer {
        id: nid(1),
        direction: Direction::Row,
        children: vec![
            WeightedChild {
                node: LayoutNode::Leaf(agent_leaf(leaf_a, Some(agent))),
                weight: 1.0,
            },
            WeightedChild {
                node: LayoutNode::Leaf(agent_leaf(leaf_b, Some(agent))),
                weight: 1.0,
            },
        ],
    }));
    seed_layouts(&fs, lid(1), &tree);

    // Only node A hosts a live session.
    let live = FakeLive::with_nodes(&[leaf_a]);
    let uc = make_use_case(&store, &fs, &live);

    let out = uc
        .execute(SnapshotRunningAgentsInput { project_id: pid(1) })
        .await
        .unwrap();

    assert_eq!(out.running, 1, "only the live cell counts as running");
    assert_eq!(
        out.stopped, 1,
        "the duplicate (dead) cell counts as stopped"
    );
    assert_eq!(was_running(&fs, leaf_a), Some(true));
    assert_eq!(
        was_running(&fs, leaf_b),
        Some(false),
        "the duplicate leaf for the same agent must NOT be marked running"
    );
}