ninox-core 0.17.1

use crate::{github::{GitHubClient, GithubApi}, store::Store, types::*};
use std::{collections::HashMap, sync::Arc};
use tokio::sync::{broadcast, Mutex};

#[derive(Debug, Clone)]
pub enum Event {
    OrchestratorSpawned(Orchestrator),
    OrchestratorRemoved(OrchestratorId),
    SessionUpdated(Session),
    SessionSpawned(Session),
    SessionDone(SessionId),
    TerminalOutput { session_id: SessionId, bytes: Vec<u8> },
    /// Rendering stream from an attached tmux client (AttachedClient).
    /// `generation` identifies which `AttachedClient::spawn` call produced
    /// this event — consumers must ignore events whose generation doesn't
    /// match the currently-live client for `session_id` so a stale client
    /// (superseded by a fresh attach) cannot clobber the new one.
    ClientOutput   { session_id: SessionId, generation: u64, bytes: Vec<u8> },
    /// The attached tmux client process exited (detach, kill, server gone).
    /// See `ClientOutput` for why `generation` matters.
    ClientClosed   { session_id: SessionId, generation: u64 },
    CiUpdated      { pr_id: PrId, status: CIStatus },
    PrOpened       { session_id: SessionId, pr: PR },
    ReviewComment  { pr_id: PrId, comment: Comment },
    Notification(Notification),
}

pub struct Engine {
    pub store: Arc<Store>,
    tx: broadcast::Sender<Event>,
    pty_writers:  Mutex<HashMap<SessionId, tokio::sync::mpsc::UnboundedSender<Vec<u8>>>>,
    /// Per-session cancellation senders for active FIFO reader tasks.
    /// Sending () to the stored sender stops the running reader immediately.
    stream_cancel: Mutex<HashMap<SessionId, tokio::sync::oneshot::Sender<()>>>,
    /// Optional GitHub API client. None when no token is configured.
    pub github: Option<Arc<dyn GithubApi>>,
}

impl Engine {
    pub fn new(store: Arc<Store>) -> Arc<Self> {
        let (tx, _) = broadcast::channel(256);
        Arc::new(Self {
            store,
            tx,
            pty_writers:   Mutex::new(HashMap::new()),
            stream_cancel: Mutex::new(HashMap::new()),
            github:        None,
        })
    }

    pub fn new_with_github(store: Arc<Store>, token: String) -> Arc<Self> {
        let (tx, _) = broadcast::channel(256);
        let github = GitHubClient::new(token).ok()
            .map(|c| Arc::new(c) as Arc<dyn GithubApi>);
        Arc::new(Self {
            store,
            tx,
            pty_writers:   Mutex::new(HashMap::new()),
            stream_cancel: Mutex::new(HashMap::new()),
            github,
        })
    }

    /// Construct an `Engine` with a caller-supplied `GithubApi` — the
    /// dependency-injection seam tests use to drive the GitHub-enrichment
    /// poller against a fake instead of the real network.
    pub fn new_with_github_api(store: Arc<Store>, github: Arc<dyn GithubApi>) -> Arc<Self> {
        let (tx, _) = broadcast::channel(256);
        Arc::new(Self {
            store,
            tx,
            pty_writers:   Mutex::new(HashMap::new()),
            stream_cancel: Mutex::new(HashMap::new()),
            github:        Some(github),
        })
    }

    /// Cancel any running FIFO reader for `session_id` and return a fresh
    /// cancellation receiver for the new reader.  Call this at the top of
    /// every `start_streaming` invocation.
    pub async fn register_stream(
        &self,
        session_id: SessionId,
    ) -> tokio::sync::oneshot::Receiver<()> {
        let mut map = self.stream_cancel.lock().await;
        if let Some(old_tx) = map.remove(&session_id) {
            let _ = old_tx.send(());
        }
        let (tx, rx) = tokio::sync::oneshot::channel();
        map.insert(session_id, tx);
        rx
    }

    pub fn emit(&self, event: Event) {
        let _ = self.tx.send(event);
    }

    pub fn subscribe(&self) -> broadcast::Receiver<Event> {
        self.tx.subscribe()
    }

    pub async fn register_pty_writer(
        &self,
        session_id: SessionId,
        writer: tokio::sync::mpsc::UnboundedSender<Vec<u8>>,
    ) {
        self.pty_writers.lock().await.insert(session_id, writer);
    }

    pub async fn get_pty_writer(
        &self,
        session_id: &str,
    ) -> Option<tokio::sync::mpsc::UnboundedSender<Vec<u8>>> {
        self.pty_writers.lock().await.get(session_id).cloned()
    }

    /// Kill all worker sessions belonging to an orchestrator, delete from DB, emit events.
    pub async fn remove_orchestrator(&self, orchestrator_id: &str) -> anyhow::Result<()> {
        let workers = self.store.sessions_by_orchestrator(orchestrator_id)?;
        let sessions_dir = crate::config::AppConfig::sessions_dir();
        for session in &workers {
            let _ = crate::tmux::kill_session(&session.id).await;
            remove_worktree_and_artifacts(&session.id, session.workspace_path.as_deref(), &sessions_dir).await;
        }
        // Also kill the orchestrator's own tmux session (same id as orchestrator).
        let _ = crate::tmux::kill_session(orchestrator_id).await;
        self.store.delete_orchestrator(orchestrator_id)?;
        for session in workers {
            self.emit(Event::SessionDone(session.id));
        }
        self.emit(Event::OrchestratorRemoved(orchestrator_id.to_string()));
        Ok(())
    }

    /// Kill the tmux session and delete it from the DB entirely.
    pub async fn remove_session(&self, session_id: &str) -> anyhow::Result<()> {
        let _ = crate::tmux::kill_session(session_id).await;
        let workspace_path = self.store.get_session(session_id).ok().flatten()
            .and_then(|s| s.workspace_path);
        remove_worktree_and_artifacts(
            session_id, workspace_path.as_deref(), &crate::config::AppConfig::sessions_dir(),
        ).await;
        self.store.delete_session(session_id)?;
        self.emit(Event::SessionDone(session_id.to_string()));
        Ok(())
    }

    /// Send a text message to the agent running in a session's tmux window.
    /// The message is injected as keyboard input — the agent sees it as typed text.
    /// Returns Ok(()) if the tmux send succeeded; returns an error if the session
    /// has no active tmux window or tmux is unavailable.
    pub async fn send_to_session(&self, session_id: &str, message: &str) -> anyhow::Result<()> {
        crate::tmux::send_keys(session_id, message).await
    }

    /// Kill the tmux session, mark it Terminated in the DB, and emit SessionUpdated.
    pub async fn terminate_session(&self, session_id: &str) -> anyhow::Result<()> {
        // Best-effort tmux kill (session may already be dead).
        let _ = crate::tmux::kill_session(session_id).await;

        if let Some(mut session) = self.store.get_session(session_id)? {
            // Never clobber a terminal status — most importantly, never flip
            // a `Done` session back to `Terminated`. `Done` is only ever
            // reached via `handle_merge_detection`, which already notified
            // the orchestrator; `sweep_retired_sessions`'s notification
            // dedup relies on `Done` meaning "already told" (see
            // `lifecycle::poller::sweep_retired_sessions`), so re-marking it
            // `Terminated` here would make the sweep send a second,
            // contradictory notice for a PR that was in fact merged.
            if matches!(
                session.status,
                SessionStatus::Done | SessionStatus::Terminated | SessionStatus::Interrupted
            ) {
                return Ok(());
            }
            session.status = SessionStatus::Terminated;
            self.store.upsert_session(&session)?;
            self.emit(Event::SessionUpdated(session));
        }
        Ok(())
    }

    /// Kill the tmux session (best-effort), remove its worktree/artifacts,
    /// and mark it Done in the DB. Called automatically when a PR is
    /// merged. Emits SessionUpdated. Mirrors `remove_session`'s worktree
    /// and artifact cleanup so the record isn't deleted here but still
    /// stops leaking a checked-out worktree and per-session hook files.
    pub async fn cleanup_session(&self, session_id: &str) -> anyhow::Result<()> {
        self.cleanup_session_in(session_id, &crate::config::AppConfig::sessions_dir()).await
    }

    /// `cleanup_session`'s implementation, with the sessions dir as a
    /// parameter so tests can point artifact removal at a tempdir instead
    /// of the real `AppConfig::sessions_dir()`.
    async fn cleanup_session_in(
        &self,
        session_id:   &str,
        sessions_dir: &std::path::Path,
    ) -> anyhow::Result<()> {
        // Best-effort tmux kill — session may already be dead.
        let _ = crate::tmux::kill_session(session_id).await;

        if let Some(mut session) = self.store.get_session(session_id)? {
            remove_worktree_and_artifacts(session_id, session.workspace_path.as_deref(), sessions_dir).await;
            session.status = crate::types::SessionStatus::Done;
            session.terminal_at = Some(crate::lifecycle::poller::now_millis());
            self.store.upsert_session(&session)?;
            self.emit(Event::SessionUpdated(session));
        }
        Ok(())
    }
}

/// Remove a session's worktree (if any) and hook artifacts — the shared
/// teardown tail of `remove_orchestrator`, `remove_session`, and
/// `cleanup_session`, so none of them leak a checked-out worktree or
/// per-session hook files.
async fn remove_worktree_and_artifacts(
    session_id:     &str,
    workspace_path: Option<&str>,
    sessions_dir:   &std::path::Path,
) {
    if let Some(wp) = workspace_path {
        remove_worker_worktree(wp, session_id).await;
    }
    crate::hooks::remove_session_artifacts(sessions_dir, session_id);
}

/// Remove a Ninox-managed git worktree (best-effort, never propagates errors).
///
/// Only acts on paths that match the `.claude/worktrees/{session_id}` pattern
/// so it never touches unrelated directories.
async fn remove_worker_worktree(workspace_path: &str, session_id: &str) {
    let suffix = format!("/.claude/worktrees/{session_id}");
    if !workspace_path.ends_with(&suffix) {
        return; // Not a Ninox worktree — leave it alone.
    }
    // Derive the repo root by stripping the worktree suffix.
    let repo_root = &workspace_path[..workspace_path.len() - suffix.len()];
    if repo_root.is_empty() {
        return;
    }
    let _ = tokio::process::Command::new("git")
        .args(["-C", repo_root, "worktree", "remove", "--force", workspace_path])
        .output()
        .await;
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::store::Store;
    use tempfile::tempdir;

    #[tokio::test]
    async fn emit_received_by_subscriber() {
        let store = Arc::new(Store::open(tempdir().unwrap().keep().join("t.db")).unwrap());
        let engine = Engine::new(store);
        let mut rx = engine.subscribe();
        engine.emit(Event::SessionDone("s1".into()));
        let event = rx.recv().await.unwrap();
        assert!(matches!(event, Event::SessionDone(id) if id == "s1"));
    }

    #[tokio::test]
    async fn terminate_emits_session_updated() {
        let store = Arc::new(Store::open(tempdir().unwrap().keep().join("t.db")).unwrap());
        let session = crate::types::Session {
            id: "s1".into(), orchestrator_id: None, name: "w".into(),
            repo: "r".into(), status: crate::types::SessionStatus::Working,
            agent_type: "c".into(), cost_usd: 0.0, started_at: 0,
            pr_number: None, pr_id: None, workspace_path: None, pid: None,
            model: None, context_tokens: None, catalogue_path: None,
            context_used_pct: None, context_total_tokens: None, context_window_size: None,
            claude_session_id: None,
            summary: None,
            terminal_at: None,
        };
        store.upsert_session(&session).unwrap();
        let engine = Engine::new(store);
        let mut rx = engine.subscribe();

        engine.terminate_session("s1").await.unwrap();

        let evt = rx.recv().await.unwrap();
        if let Event::SessionUpdated(s) = evt {
            assert!(matches!(s.status, crate::types::SessionStatus::Terminated));
            assert_eq!(
                s.terminal_at, None,
                "user-initiated terminate_session must not stamp terminal_at — \
                 that's reserved for the automatic lifecycle path, so this \
                 session is purged on sight rather than held for the \
                 retention grace period",
            );
        } else {
            panic!("expected SessionUpdated");
        }
    }

    /// A `Done` session (reached via `cleanup_session`, which already
    /// notified the orchestrator about the merge) must never be flipped back
    /// to `Terminated` by a later `terminate_session` call (e.g. the `DELETE
    /// /sessions/:id` route firing on a stale client, or a race with the
    /// poller). `sweep_retired_sessions`'s notification dedup relies on
    /// `Done` meaning "already told" — reopening it as `Terminated` would
    /// make the sweep send a second, contradictory notification claiming the
    /// PR was never detected merged.
    #[tokio::test]
    async fn terminate_session_does_not_reopen_a_done_session() {
        let store = Arc::new(Store::open(tempdir().unwrap().keep().join("t.db")).unwrap());
        let session = crate::types::Session {
            id: "s1".into(), orchestrator_id: None, name: "w".into(),
            repo: "r".into(), status: crate::types::SessionStatus::Done,
            agent_type: "c".into(), cost_usd: 0.0, started_at: 0,
            pr_number: Some(1), pr_id: Some(1), workspace_path: None, pid: None,
            model: None, context_tokens: None, catalogue_path: None,
            context_used_pct: None, context_total_tokens: None, context_window_size: None,
            claude_session_id: None,
            summary: None,
            terminal_at: Some(1_000),
        };
        store.upsert_session(&session).unwrap();
        let engine = Engine::new(Arc::clone(&store));
        let mut rx = engine.subscribe();

        engine.terminate_session("s1").await.unwrap();

        assert!(
            rx.try_recv().is_err(),
            "a no-op terminate_session on a Done session must not emit SessionUpdated"
        );
        let after = store.get_session("s1").unwrap().unwrap();
        assert!(matches!(after.status, crate::types::SessionStatus::Done), "status must stay Done");
        assert_eq!(after.terminal_at, Some(1_000), "terminal_at must be untouched");
    }

    #[tokio::test]
    async fn cleanup_session_sets_done_status() {
        let store = Arc::new(
            Store::open(tempdir().unwrap().keep().join("t.db")).unwrap()
        );
        let session = crate::types::Session {
            id: "s1".into(), orchestrator_id: None, name: "w".into(),
            repo: "r".into(), status: crate::types::SessionStatus::PrOpen,
            agent_type: "c".into(), cost_usd: 0.0, started_at: 0,
            pr_number: Some(1), pr_id: Some(1),
            workspace_path: None, pid: None,
            model: None, context_tokens: None, catalogue_path: None,
            context_used_pct: None, context_total_tokens: None, context_window_size: None,
            claude_session_id: None,
            summary: None,
            terminal_at: None,
        };
        store.upsert_session(&session).unwrap();
        let engine = Engine::new(Arc::clone(&store));
        let mut rx = engine.subscribe();

        engine.cleanup_session("s1").await.unwrap();

        let evt = rx.recv().await.unwrap();
        if let Event::SessionUpdated(s) = evt {
            assert!(matches!(s.status, crate::types::SessionStatus::Done));
            assert!(
                s.terminal_at.is_some(),
                "cleanup_session must stamp terminal_at so the retention sweep can hold this \
                 record for the fleet board's grace window instead of purging it on sight",
            );
        } else {
            panic!("expected SessionUpdated");
        }
    }

    #[tokio::test]
    async fn cleanup_session_removes_worktree_and_artifacts() {
        // Real git repo + real worktree so we exercise the same
        // `remove_worker_worktree` path `remove_session` uses, not a mock.
        let repo_dir = tempdir().unwrap();
        let run_git = |args: &[&str]| {
            let status = std::process::Command::new("git")
                .args(args)
                .current_dir(repo_dir.path())
                .status()
                .unwrap();
            assert!(status.success(), "git {args:?} failed");
        };
        run_git(&["init", "-q"]);
        run_git(&["-c", "user.email=t@t.com", "-c", "user.name=t", "commit", "--allow-empty", "-q", "-m", "init"]);

        let session_id = "s1";
        let worktree_path = repo_dir.path().join(".claude/worktrees").join(session_id);
        run_git(&["worktree", "add", "-q", worktree_path.to_str().unwrap()]);
        assert!(worktree_path.exists(), "test setup: worktree must exist before cleanup");

        let sessions_dir = tempdir().unwrap();
        let artifact_path = sessions_dir.path().join(format!("{session_id}.json"));
        std::fs::write(&artifact_path, "{}").unwrap();

        let store = Arc::new(Store::open(tempdir().unwrap().keep().join("t.db")).unwrap());
        let session = crate::types::Session {
            id: session_id.into(), orchestrator_id: None, name: "w".into(),
            repo: "r".into(), status: crate::types::SessionStatus::PrOpen,
            agent_type: "c".into(), cost_usd: 0.0, started_at: 0,
            pr_number: Some(1), pr_id: Some(1),
            workspace_path: Some(worktree_path.to_string_lossy().to_string()),
            pid: None,
            model: None, context_tokens: None, catalogue_path: None,
            context_used_pct: None, context_total_tokens: None, context_window_size: None,
            claude_session_id: None,
            summary: None,
            terminal_at: None,
        };
        store.upsert_session(&session).unwrap();
        let engine = Engine::new(Arc::clone(&store));

        engine.cleanup_session_in(session_id, sessions_dir.path()).await.unwrap();

        assert!(!worktree_path.exists(), "cleanup_session must remove the worktree, like remove_session does");
        assert!(!artifact_path.exists(), "cleanup_session must remove session artifacts, like remove_session does");
        let after = store.get_session(session_id).unwrap().unwrap();
        assert!(matches!(after.status, crate::types::SessionStatus::Done));
    }

    #[tokio::test]
    async fn remove_worker_worktree_ignores_non_ninox_paths() {
        // Should be a no-op for paths that don't match .claude/worktrees/{id}.
        // We just verify it doesn't panic or error.
        remove_worker_worktree("/some/random/path", "s1").await;
        remove_worker_worktree("/repo/.claude/worktrees/other-id", "s1").await;
        remove_worker_worktree("", "s1").await;
    }
}