selfware 0.6.1

//! Run supervisor — process model that owns N concurrent agent runs.
//!
//! This is **Increment 1** of the `RunSupervisor`. It provides the core
//! concurrent-run lifecycle (`spawn`, `abort`, `status`, `list`) plus an
//! agent-launching convenience (`start`), per-run broadcast event streams
//! (`attach`, `wait_for_terminal`), and documented stubs for `pause` and
//! `resume` which land in later increments.
//!
//! The design deliberately keeps the testable primitive — [`RunSupervisor::spawn`]
//! — independent of any live LLM. Tests spawn trivial futures (`async { Ok(()) }`,
//! `async { Err(...) }`, long sleeps) and assert state transitions without
//! touching [`crate::agent::Agent`], which requires a configured API client.
//!
//! Cancellation uses [`std::sync::atomic::AtomicBool`] (no `tokio_util`
//! dependency). An `abort()` flips the flag **and** calls `JoinHandle::abort`,
//! so cooperatively-polling futures can observe the flag while also receiving an
//! immediate hard-cancel.

use std::collections::HashMap;
use std::future::Future;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

use tokio::sync::RwLock;
use tokio::task::JoinHandle;

use crate::agent::tui_events::{AgentEvent, BroadcastEmitter, EventEmitter};

/// Lifecycle status of a single run.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RunStatus {
    /// The run is currently executing.
    Running,
    /// The run is paused (not yet implemented — see [`RunSupervisor::pause`]).
    Paused,
    /// The run completed successfully (`Ok(())`).
    Completed,
    /// The run failed (`Err(_)`).
    Failed,
    /// The run was aborted via [`RunSupervisor::abort`].
    Aborted,
}

impl RunStatus {
    /// Whether the run has settled and will not progress further.
    pub fn is_terminal(self) -> bool {
        matches!(
            self,
            RunStatus::Completed | RunStatus::Failed | RunStatus::Aborted
        )
    }

    /// The event broadcast to a run's subscribers when it settles in this
    /// status. Emitted exactly once per run, *after* the status write, so
    /// any waiter woken by the event already observes the final status
    /// (P0-3). Only meaningful for terminal statuses.
    fn terminal_event(self) -> AgentEvent {
        match self {
            RunStatus::Completed => AgentEvent::Completed {
                message: "run completed".to_string(),
            },
            RunStatus::Failed => AgentEvent::Error {
                message: "run failed".to_string(),
            },
            other => AgentEvent::Status {
                message: format!("run status: {other:?}"),
            },
        }
    }
}

/// Opaque, monotonic identifier for a run.
///
/// Generated by an [`AtomicU64`] counter on the supervisor — no dates or
/// randomness involved, so IDs are strictly monotonic within a supervisor
/// instance and trivially comparable/orderable.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct RunId(pub u64);

impl std::fmt::Display for RunId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "run-{}", self.0)
    }
}

/// Capacity of each run's broadcast event channel.
const EVENT_CHANNEL_CAPACITY: usize = 256;

/// How often [`RunSupervisor::wait_for_terminal`] re-polls a run's status
/// while waiting on its event stream. The poll guarantees the wait ends
/// even when no terminal event is observed (receiver attached late, lag,
/// or a run whose sender is otherwise silent).
const STATUS_POLL_INTERVAL: Duration = Duration::from_millis(100);

/// Internal handle for a single run.
struct RunHandle {
    status: Arc<RwLock<RunStatus>>,
    cancel: Arc<AtomicBool>,
    join: JoinHandle<()>,
    /// Human-readable task description for diagnostics/listing.
    #[allow(dead_code)]
    task: String,
    /// Per-run broadcast sender for `AgentEvent`s.
    ///
    /// Every run — whether created via [`RunSupervisor::spawn`] (pure future)
    /// or [`RunSupervisor::start`] (real agent) — gets its own broadcast
    /// channel so that [`attach`](RunSupervisor::attach) always works.
    /// Exactly one **terminal event** is broadcast when the run settles
    /// (completed/failed/aborted); the sender itself lives here for the
    /// supervisor's lifetime, so receivers must not rely on the channel
    /// closing — use [`RunSupervisor::wait_for_terminal`] to wait for the
    /// end of a run.
    events: tokio::sync::broadcast::Sender<AgentEvent>,
}

/// Supervisor owning N concurrent agent runs.
///
/// Create with [`RunSupervisor::new`], then drive runs via [`spawn`](Self::spawn)
/// (generic, testable) or [`start`](Self::start) (builds a real
/// [`crate::agent::Agent`] run — requires a live LLM).
pub struct RunSupervisor {
    runs: RwLock<HashMap<RunId, RunHandle>>,
    next_id: AtomicU64,
}

impl Default for RunSupervisor {
    fn default() -> Self {
        Self::new()
    }
}

impl RunSupervisor {
    /// Create an empty supervisor with no runs.
    pub fn new() -> Self {
        Self {
            runs: RwLock::new(HashMap::new()),
            next_id: AtomicU64::new(1),
        }
    }

    /// Spawn a generic future as a supervised run and return its [`RunId`].
    ///
    /// This is the **core testable primitive**: tests pass trivial futures
    /// (`async { Ok(()) }`, `async { Err(...) }`, a long `sleep`) and assert
    /// state transitions without needing a live LLM.
    ///
    /// The future runs on a dedicated tokio task. On completion the run's
    /// status becomes [`RunStatus::Completed`] (for `Ok`) or
    /// [`RunStatus::Failed`] (for `Err`) — **unless** the cancel flag was set
    /// first, in which case the status becomes [`RunStatus::Aborted`]
    /// regardless of the future's own outcome.
    ///
    /// The `RunId` is allocated with a monotonic `fetch_add` on an internal
    /// [`AtomicU64`]; no dates/randomness are used.
    ///
    /// When the run settles, exactly one terminal event is broadcast on the
    /// run's event channel (after the status write) so that
    /// [`wait_for_terminal`](Self::wait_for_terminal) and `attach`
    /// subscribers observe the end of the run.
    pub async fn spawn<F>(&self, task: String, run: F) -> RunId
    where
        F: Future<Output = anyhow::Result<()>> + Send + 'static,
    {
        let (event_tx, _event_rx) =
            tokio::sync::broadcast::channel::<AgentEvent>(EVENT_CHANNEL_CAPACITY);
        self.spawn_with_events(task, event_tx, run).await
    }

    /// Like [`spawn`](Self::spawn) but uses a caller-provided broadcast sender
    /// for the run's event channel instead of creating a new one.
    ///
    /// This lets [`start`](Self::start) pre-create the channel, wire a
    /// [`BroadcastEmitter`] onto the agent, and then pass the *same* sender
    /// here so that [`attach`](Self::attach) subscribers receive events the
    /// agent emits.
    async fn spawn_with_events<F>(
        &self,
        task: String,
        event_tx: tokio::sync::broadcast::Sender<AgentEvent>,
        run: F,
    ) -> RunId
    where
        F: Future<Output = anyhow::Result<()>> + Send + 'static,
    {
        let id = RunId(self.next_id.fetch_add(1, Ordering::Relaxed));
        let status = Arc::new(RwLock::new(RunStatus::Running));
        let cancel = Arc::new(AtomicBool::new(false));

        let status_clone = Arc::clone(&status);
        let cancel_clone = Arc::clone(&cancel);
        let settle_tx = event_tx.clone();

        let join = tokio::spawn(async move {
            let result = run.await;
            // If the cancel flag was flipped (abort), treat the outcome as
            // Aborted regardless of what the future returned.
            let final_status = if cancel_clone.load(Ordering::Relaxed) {
                RunStatus::Aborted
            } else {
                match result {
                    Ok(()) => RunStatus::Completed,
                    Err(_) => RunStatus::Failed,
                }
            };
            // Status first, THEN the terminal event (P0-3): any waiter woken
            // by the event must already observe the final status. The send
            // may fail when no subscribers exist — that is fine, waiters
            // also poll the status (see `wait_for_terminal`).
            *status_clone.write().await = final_status;
            let _ = settle_tx.send(final_status.terminal_event());
        });

        let handle = RunHandle {
            status,
            cancel,
            join,
            task,
            events: event_tx,
        };

        self.runs.write().await.insert(id, handle);
        id
    }

    /// Build a real agent run from a [`crate::config::Config`] and spawn it.
    ///
    /// This constructs an [`crate::agent::Agent`] and calls
    /// [`Agent::run_task`](crate::agent::Agent::run_task). **Requires a live
    /// LLM** (a configured API client); tests should use [`spawn`](Self::spawn)
    /// instead.
    ///
    /// A per-run `tokio::sync::broadcast` channel is created and a
    /// [`BroadcastEmitter`] is wired onto the agent via
    /// [`Agent::with_event_emitter`](crate::agent::Agent::with_event_emitter)
    /// before `run_task` is called, so that [`attach`](Self::attach) returns a
    /// live subscriber that receives the agent's `AgentEvent`s.
    pub async fn start(&self, task: String, config: crate::config::Config) -> RunId {
        let task_for_run = task.clone();
        // We create the broadcast channel here and clone the sender into the
        // RunHandle (via spawn) while also wrapping a clone in a
        // BroadcastEmitter for the agent.
        let (event_tx, _event_rx) =
            tokio::sync::broadcast::channel::<AgentEvent>(EVENT_CHANNEL_CAPACITY);
        let emitter: Arc<dyn EventEmitter> = Arc::new(BroadcastEmitter::new(event_tx.clone()));

        let task_for_spawn = task.clone();
        // Call spawn but override its internal channel with our pre-built one
        // by using a lower-level path: we can't easily inject into spawn, so
        // instead we use spawn_with_events.
        let id = self
            .spawn_with_events(task_for_spawn, event_tx, async move {
                let mut agent = crate::agent::Agent::new(config).await?;
                agent = agent.with_event_emitter(emitter);
                agent.run_task(&task_for_run).await
            })
            .await;
        id
    }

    /// Abort a run: set the cancel flag, abort the `JoinHandle`, and mark the
    /// status [`RunStatus::Aborted`].
    ///
    /// Returns `true` if the run existed, `false` otherwise. Aborting an
    /// already-completed run is harmless (it just re-sets an already-finished
    /// status).
    pub async fn abort(&self, id: &RunId) -> bool {
        let runs = self.runs.write().await;
        let Some(handle) = runs.get(id) else {
            return false;
        };
        handle.cancel.store(true, Ordering::Relaxed);
        handle.join.abort();
        let was_terminal = {
            let mut st = handle.status.write().await;
            let prev = *st;
            *st = RunStatus::Aborted;
            prev.is_terminal()
        };
        // P0-3: wake event waiters — but only if the run hadn't already
        // settled, so each run emits exactly one terminal event.
        if !was_terminal {
            let _ = handle.events.send(RunStatus::Aborted.terminal_event());
        }
        true
    }

    /// Read the current [`RunStatus`] of a run, or `None` if it doesn't exist.
    pub async fn status(&self, id: &RunId) -> Option<RunStatus> {
        let runs = self.runs.read().await;
        match runs.get(id) {
            Some(h) => {
                let st = *h.status.read().await;
                Some(st)
            }
            None => None,
        }
    }

    /// Snapshot of all runs and their current statuses.
    pub async fn list(&self) -> Vec<(RunId, RunStatus)> {
        let runs = self.runs.read().await;
        let mut out = Vec::with_capacity(runs.len());
        for (id, handle) in runs.iter() {
            let st = *handle.status.read().await;
            out.push((*id, st));
        }
        out
    }

    /// Attach to a run's event stream.
    ///
    /// Returns a live `tokio::sync::broadcast::Receiver<AgentEvent>` for the
    /// run's event channel, or `None` if the run id is unknown.
    ///
    /// Multiple callers can attach to the same run simultaneously — the
    /// underlying `tokio::sync::broadcast` channel fans out each emitted
    /// `AgentEvent` to every active subscriber.
    ///
    /// For runs created via [`start`](Self::start), the agent's events flow
    /// through this channel in real time. Every run — `spawn` or `start` —
    /// also broadcasts exactly one terminal event when it settles. The
    /// channel itself is **never closed** (the sender lives in the
    /// supervisor), so to wait for the end of a run use
    /// [`wait_for_terminal`](Self::wait_for_terminal) rather than blocking
    /// on `recv()`.
    pub async fn attach(&self, id: &RunId) -> Option<tokio::sync::broadcast::Receiver<AgentEvent>> {
        let runs = self.runs.read().await;
        runs.get(id).map(|h| h.events.subscribe())
    }

    /// Wait until the run `id` reaches a terminal status
    /// ([`RunStatus::Completed`], [`RunStatus::Failed`], or
    /// [`RunStatus::Aborted`]), draining the run's event stream into
    /// `on_event` as events arrive.
    ///
    /// This is the P0-3 fix: the previous CLI loop checked the status only
    /// *before* each blocking `recv().await`, and the broadcast sender
    /// lives in the run's `RunHandle` for the supervisor's lifetime, so
    /// once the agent stopped emitting events the loop blocked in `recv()`
    /// forever and the final status was never observed (or persisted to
    /// the run registry). Here the status is re-checked after every event
    /// **and** on a fixed interval via `tokio::select!`, so a settled run
    /// is always detected within `STATUS_POLL_INTERVAL` — even if its
    /// terminal event was broadcast before the receiver attached, or was
    /// lost to lag.
    ///
    /// Returns the terminal status, or `None` if the run id is unknown.
    pub async fn wait_for_terminal(
        &self,
        id: &RunId,
        mut rx: tokio::sync::broadcast::Receiver<AgentEvent>,
        mut on_event: impl FnMut(AgentEvent),
    ) -> Option<RunStatus> {
        let mut poll = tokio::time::interval(STATUS_POLL_INTERVAL);
        poll.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
        loop {
            match self.status(id).await {
                Some(st) if st.is_terminal() => return Some(st),
                Some(_) => {}
                None => return None,
            }
            tokio::select! {
                ev = rx.recv() => {
                    match ev {
                        Ok(ev) => on_event(ev),
                        // All senders dropped: no further events can arrive.
                        // Do one final status read and return it as-is.
                        Err(tokio::sync::broadcast::error::RecvError::Closed) => {
                            return self.status(id).await;
                        }
                        Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
                    }
                }
                _ = poll.tick() => {}
            }
        }
    }

    /// Pause a running run.
    ///
    /// **Not yet implemented (Increment 2).** Currently returns an error.
    pub async fn pause(&self, _id: &RunId) -> anyhow::Result<()> {
        anyhow::bail!("pause not yet implemented (inc2)")
    }

    /// Resume a paused run.
    ///
    /// **Not yet implemented (Increment 2).** Currently returns an error.
    pub async fn resume(&self, _id: &RunId) -> anyhow::Result<()> {
        anyhow::bail!("resume not yet implemented (inc2)")
    }

    /// Emit an `AgentEvent` directly to a run's broadcast channel.
    ///
    /// This is primarily intended for testing fan-out without a live LLM:
    /// a test can [`spawn`](Self::spawn) a trivial future, [`attach`](Self::attach)
    /// multiple subscribers, then call this to push an event and verify every
    /// subscriber receives it.
    ///
    /// Returns `true` if the run existed, `false` otherwise.
    #[allow(dead_code)]
    pub(crate) async fn emit_event(&self, id: &RunId, event: AgentEvent) -> bool {
        let runs = self.runs.read().await;
        match runs.get(id) {
            Some(h) => {
                let _ = h.events.send(event);
                true
            }
            None => false,
        }
    }
}

#[cfg(test)]
#[path = "../../tests/unit/supervision/run_supervisor/run_supervisor_test.rs"]
mod tests;