hyperhive/hivectl/src/dag_progress.rs

//! `hivectl` rebuild-queue progress rendering.
//!
//! Split out of `hivectl.rs` (which is already large): everything that
//! polls the daemon's node queue (`HostRequest::QueueNodes`) and renders
//! progress lives here. [`crate::dag_progress::wait_for_nodes`] is the entry point the command
//! handlers call; it dispatches to a live `indicatif` animation on a TTY
//! and a plain line-on-change stream otherwise.
//!
//! Consumes `hive-jobq-wire`'s generic `GraphNode`/`NodePayload`: a
//! node's kind (and, for a root, what submitted it) comes straight off
//! `payload.label`; node-specific extras (`agent`) ride in
//! `payload.data`'s opaque kvps — the shape `hive-c0re`'s
//! `impl WireNode for NodeKind` produces (see its doc comment). There's
//! no separate roll-up field on the wire: a node's own
//! `state` already reflects everything below it (see `hive_jobq_wire`'s
//! doc comment), so nothing here re-derives a roll-up or waits for every
//! node to go terminal before reading a result.
//!
//! `ids` is always a **batch**, not a single id: a hive-wide op (e.g.
//! restarting every agent) submits one root per agent, and the whole
//! batch is polled together in a single `QueueNodes` request per tick —
//! `group_by_root` splits the combined response back into per-root
//! groups rather than issuing one round-trip per id.

use std::collections::{BTreeSet, HashMap, HashSet};
use std::path::Path;

use anyhow::{Context as _, Result, bail};
use hive_host_sock::jobs::State;
use hive_jobq_wire::{GraphDep, GraphNode, WireId};

/// Poll the submitted ids (`HostRequest::QueueNodes`, ~1s interval, all
/// ids in one request per tick) and render progress until they all reach
/// a terminal state. Exits non-zero (via the returned `Err`) when any
/// node ends `failed`; a `cancelled` node terminates the wait but is an
/// operator action, not an error.
///
/// # Errors
///
/// Returns an error if the daemon socket can't be reached, or if any
/// polled node finished in the `failed` state.
pub(crate) async fn wait_for_nodes(socket: &Path, ids: Vec<u64>, no_wait: bool) -> Result<()> {
    use std::io::IsTerminal as _;
    if no_wait || ids.is_empty() {
        return Ok(());
    }
    // Animate only on a real terminal. Piped / CI output falls back to the
    // plain line-on-change stream so logs stay free of spinner redraw noise.
    if std::io::stderr().is_terminal() {
        wait_for_nodes_animated(socket, ids).await
    } else {
        wait_for_nodes_plain(socket, ids).await
    }
}

/// Split a combined `QueueNodes` response back into per-root groups, keyed
/// by each root's own id. A node with `parent: None` is a root and starts
/// its own group (keyed by its own id); everything else joins its direct
/// parent's group. Today's job shapes are root + flat children (no deeper
/// nesting), so a direct-parent lookup is enough — matches the flat
/// child-iteration every render helper below already assumed.
fn group_by_root(nodes: Vec<GraphNode>) -> HashMap<u64, Vec<GraphNode>> {
    let mut groups: HashMap<u64, Vec<GraphNode>> = HashMap::new();
    for n in nodes {
        let key = n.parent.unwrap_or(n.id);
        groups.entry(key).or_default().push(n);
    }
    groups
}

/// This id's own root node within its group, found by id rather than
/// "no parent" — `group_by_root` already keyed each group by its root's
/// id, so the root is just the member whose id matches the key.
fn root_in_group(id: u64, nodes: &[GraphNode]) -> Option<&GraphNode> {
    nodes.iter().find(|n| n.id == id)
}

/// Non-TTY progress: print a fresh line whenever a node's rendered state
/// changes. No cursor tricks, so it's clean in pipes and CI logs.
async fn wait_for_nodes_plain(socket: &Path, ids: Vec<u64>) -> Result<()> {
    let mut pending: BTreeSet<u64> = ids.into_iter().collect();
    let mut last: HashMap<u64, String> = HashMap::new();
    let mut failed: Vec<String> = Vec::new();
    while !pending.is_empty() {
        let batch: Vec<u64> = pending.iter().copied().collect();
        let resp = crate::client::request(
            socket,
            hive_host_sock::HostRequest::QueueNodes { ids: batch },
        )
        .await
        .with_context(|| format!("connect to daemon socket {}", socket.display()))?;
        let groups = group_by_root(resp.nodes.unwrap_or_default());
        for id in pending.clone() {
            let nodes = groups.get(&id);
            let root = nodes.and_then(|n| root_in_group(id, n));
            let (Some(nodes), Some(root)) = (nodes, root) else {
                // Unknown id — nothing prunes the graph yet (see
                // `JobQueue::node_subtrees`'s doc comment), so an id that
                // resolves to no node never named a real job. A
                // *completed* job's nodes keep riding here instead, with a
                // terminal root `state`, which is what the check below
                // watches for.
                println!("job #{id}: gone from queue history");
                pending.remove(&id);
                continue;
            };
            let line = render_node_line(root, nodes);
            if last.get(&id) != Some(&line) {
                println!("{line}");
                last.insert(id, line);
            }
            if root.state.is_terminal() {
                if root.state == State::Failed {
                    failed.push(format!("{} {}", root.payload.label, node_agents(nodes)));
                }
                pending.remove(&id);
            }
        }
        if !pending.is_empty() {
            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
        }
    }
    finish_wait(failed)
}

/// TTY progress: a live `indicatif` render — one braille-spinner line per
/// node (grouped under a per-root header), each with its own elapsed
/// timer, plus an overall-elapsed footer. A node with more than one
/// dependency (fan-in) gets its own row with an `(after …)` marker rather
/// than being crammed onto a chain line.
async fn wait_for_nodes_animated(socket: &Path, ids: Vec<u64>) -> Result<()> {
    use indicatif::{MultiProgress, ProgressBar, ProgressStyle};

    let spinner = ProgressStyle::with_template("  {spinner} {msg}")
        .unwrap_or_else(|_| ProgressStyle::default_spinner())
        .tick_strings(&["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏", "·"]);
    let plain =
        ProgressStyle::with_template("{msg}").unwrap_or_else(|_| ProgressStyle::default_bar());

    let mp = MultiProgress::new();
    let started = std::time::Instant::now();
    // Per-root header bars + per-node bars, keyed so we update in place.
    // The insertion order groups a root's nodes right under its header.
    let mut root_bars: HashMap<u64, ProgressBar> = HashMap::new();
    let mut node_bars: HashMap<(u64, WireId), ProgressBar> = HashMap::new();
    let mut node_done: HashSet<(u64, WireId)> = HashSet::new();

    let mut pending: BTreeSet<u64> = ids.into_iter().collect();
    let mut failed: Vec<String> = Vec::new();

    while !pending.is_empty() {
        let now = now_unix();
        let batch: Vec<u64> = pending.iter().copied().collect();
        let resp = crate::client::request(
            socket,
            hive_host_sock::HostRequest::QueueNodes { ids: batch },
        )
        .await
        .with_context(|| format!("connect to daemon socket {}", socket.display()))?;
        let groups = group_by_root(resp.nodes.unwrap_or_default());
        for id in pending.clone() {
            let nodes = groups.get(&id);
            let root = nodes.and_then(|n| root_in_group(id, n));
            let (Some(nodes), Some(root)) = (nodes, root) else {
                mp.println(format!("job #{id}: gone from queue history"))
                    .ok();
                pending.remove(&id);
                continue;
            };
            let hdr = root_bars.entry(id).or_insert_with(|| {
                let b = mp.add(ProgressBar::new_spinner());
                b.set_style(plain.clone());
                b
            });
            hdr.set_message(format!(
                "{} {} {} · {}",
                state_glyph(root.state),
                root.payload.label,
                node_agents(nodes),
                fmt_dur(node_elapsed(root, now)),
            ));
            for n in nodes.iter().filter(|n| n.id != root.id) {
                let key = (id, n.id);
                if node_done.contains(&key) {
                    continue;
                }
                let bar = node_bars.entry(key).or_insert_with(|| {
                    let b = mp.add(ProgressBar::new_spinner());
                    b.set_style(spinner.clone());
                    b.enable_steady_tick(std::time::Duration::from_millis(120));
                    b
                });
                if n.state.is_terminal() {
                    bar.set_style(plain.clone());
                    bar.finish_with_message(format!(
                        "  {} {}",
                        state_glyph(n.state),
                        node_line(nodes, n, now)
                    ));
                    node_done.insert(key);
                } else {
                    bar.set_message(node_line(nodes, n, now));
                }
            }
            if root.state.is_terminal() {
                if root.state == State::Failed {
                    failed.push(format!("{} {}", root.payload.label, node_agents(nodes)));
                }
                pending.remove(&id);
            }
        }
        if !pending.is_empty() {
            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
        }
    }
    // Leave the final node lines on screen; drop the still-spinning headers
    // into their terminal state and print an overall-elapsed footer.
    for hdr in root_bars.values() {
        hdr.finish();
    }
    mp.println(format!(
        "done · {} elapsed",
        fmt_dur(started.elapsed().as_secs().try_into().unwrap_or(0))
    ))
    .ok();
    finish_wait(failed)
}

/// Shared tail: succeed when nothing failed, else bail listing the failures.
fn finish_wait(mut failed: Vec<String>) -> Result<()> {
    if failed.is_empty() {
        Ok(())
    } else {
        failed.sort();
        failed.dedup();
        bail!("queued job(s) failed: {}", failed.join(", "))
    }
}

/// Distinct agents across a root's nodes, comma-joined for display. Each
/// node's `agent` (when it targets one) rides in `payload.data["agent"]` —
/// an opaque kvp, not a typed field, since `GraphNode` carries nothing
/// domain-specific (see `hive_jobq_wire::WireNode::data`'s doc comment).
fn node_agents(nodes: &[GraphNode]) -> String {
    let mut seen: Vec<&str> = Vec::new();
    for n in nodes {
        if let Some(agent) = n
            .payload
            .data
            .get("agent")
            .and_then(serde_json::Value::as_str)
            && !seen.contains(&agent)
        {
            seen.push(agent);
        }
    }
    seen.join(",")
}

/// Current unix time in seconds (0 on the impossible pre-epoch error).
fn now_unix() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .ok()
        .and_then(|d| i64::try_from(d.as_secs()).ok())
        .unwrap_or(0)
}

/// Elapsed seconds for a node: `started_at` → `finished_at`/`now` once
/// it's run; `created_at` → `now` while it's still queued. `GraphNode`
/// carries no group-level timestamp separate from its own, so a header
/// line just reads this off whichever node it's summarizing.
fn node_elapsed(n: &GraphNode, now: i64) -> i64 {
    let start = n.started_at.unwrap_or(n.created_at);
    (n.finished_at.map_or(now, |t| t.timestamp()) - start.timestamp()).max(0)
}

/// Compact duration: `45s` under a minute, else `1m03s`.
fn fmt_dur(secs: i64) -> String {
    let s = secs.max(0);
    if s < 60 {
        format!("{s}s")
    } else {
        format!("{}m{:02}s", s / 60, s % 60)
    }
}

/// One animated node line: kind, an `(after …)` marker for a fan-in node
/// (>1 node-dependency), its elapsed timer, and a truncated error tail.
/// Resource deps (`GraphDep::Resource`) don't name another node, so they're
/// filtered out of the fan-in count the same way `hive-c0re`'s own
/// `dag_view` projection already does.
fn node_line(nodes: &[GraphNode], n: &GraphNode, now: i64) -> String {
    use std::fmt::Write as _;
    let mut s = n.payload.label.clone();
    let dep_ids: Vec<WireId> = n
        .deps
        .iter()
        .filter_map(|d| match d {
            GraphDep::Node { id, .. } => Some(*id),
            GraphDep::Resource { .. } => None,
        })
        .collect();
    if dep_ids.len() > 1 {
        let after: Vec<&str> = dep_ids
            .iter()
            .filter_map(|id| nodes.iter().find(|m| m.id == *id))
            .map(|m| m.payload.label.as_str())
            .collect();
        if !after.is_empty() {
            let _ = write!(s, " (after {})", after.join(", "));
        }
    }
    let el = node_elapsed(n, now);
    if el > 0 {
        let _ = write!(s, " · {}", fmt_dur(el));
    }
    if let Some(err) = &n.error {
        let short: String = err.chars().take(100).collect();
        let _ = write!(s, " — {short}");
    }
    s
}

fn state_glyph(state: State) -> &'static str {
    match state {
        State::Pending => "⏸",
        // `Finishing` is own-work-done with sub-nodes still going — in flight,
        // so it reads the same as running.
        State::Running | State::Finishing => "▶",
        State::Done => "✔",
        State::Failed => "✖",
        State::Cancelled => "⊘",
        // Distinct from cancelled: nothing went wrong, this branch just
        // wasn't the one the run took.
        State::Skipped => "·",
    }
}

/// One progress line for a root: roll-up glyph, label, agents, then the
/// node chain — the CLI twin of the dashboard's queue card. The glyph and
/// label come off `root`; everything else comes straight off the wire.
/// Used by the plain (non-TTY) path.
///
/// The chain shows every entry `wire_snapshot` sends, `Done` ones
/// included — that filtering was a dashboard-history-bounding concern,
/// not relevant to a single actively-watched job, so a completed step
/// keeps its checkmark instead of vanishing from the line, matching how
/// the animated path already behaves.
fn render_node_line(root: &GraphNode, nodes: &[GraphNode]) -> String {
    use std::fmt::Write as _;
    let mut out = format!(
        "{} {} {:<12}",
        state_glyph(root.state),
        root.payload.label,
        node_agents(nodes)
    );
    let children: Vec<&GraphNode> = nodes.iter().filter(|n| n.id != root.id).collect();
    for (i, n) in children.iter().enumerate() {
        let sep = if i == 0 { "  " } else { " → " };
        let _ = write!(out, "{sep}{} {}", state_glyph(n.state), n.payload.label);
    }
    if let Some(err) = children.iter().find_map(|n| n.error.as_deref()) {
        let short: String = err.chars().take(120).collect();
        let _ = write!(out, "  — {short}");
    }
    out
}

#[cfg(test)]
mod tests {
    use hive_host_sock::jobs::State;
    use hive_jobq_wire::{GraphNode, NodePayload};
    use serde_json::json;

    use super::render_node_line;

    fn root_node(id: u64, label: &str, state: State) -> GraphNode {
        GraphNode {
            id,
            parent: None,
            state,
            deps: Vec::new(),
            created_at: hive_sh4re::wire_time::from_secs(0),
            started_at: None,
            finished_at: None,
            error: None,
            payload: NodePayload {
                label: label.to_owned(),
                data: json!({}),
            },
        }
    }

    fn work_node(id: u64, root: u64, agent: &str, label: &str, state: State) -> GraphNode {
        GraphNode {
            id,
            parent: Some(root),
            state,
            deps: Vec::new(),
            created_at: hive_sh4re::wire_time::from_secs(0),
            started_at: None,
            finished_at: None,
            error: None,
            payload: NodePayload {
                label: label.to_owned(),
                data: json!({ "agent": agent }),
            },
        }
    }

    #[test]
    fn render_node_line_shows_label_and_chain() {
        // The header shows what the backend sends: the glyph off `root`'s
        // own `state` (`Running` here) + `root`'s own `payload.label`
        // ("manual"); the operation is read off the node chain, not any
        // extra metadata. `Done` nodes stay in the chain — nothing here
        // filters them off.
        let root = root_node(7, "manual", State::Running);
        let nodes = vec![
            root.clone(),
            work_node(0, 7, "alice", "prebuild", State::Done),
            work_node(1, 7, "alice", "stop_for_update", State::Done),
            work_node(2, 7, "alice", "swap", State::Running),
            work_node(3, 7, "alice", "reconcile", State::Pending),
        ];
        let line = render_node_line(&root, &nodes);
        assert!(line.starts_with("▶ manual alice"), "{line}");
        assert!(
            line.contains("✔ prebuild → ✔ stop_for_update → ▶ swap → ⏸ reconcile"),
            "{line}"
        );
    }

    #[test]
    fn render_node_line_surfaces_first_node_error() {
        let root = root_node(8, "manual", State::Failed);
        let mut failed = work_node(0, 8, "bob", "prebuild", State::Failed);
        failed.error = Some("nix build exploded".to_owned());
        let nodes = vec![root.clone(), failed];
        let line = render_node_line(&root, &nodes);
        assert!(line.contains("✖ manual"), "{line}");
        assert!(line.contains("— nix build exploded"), "{line}");
    }

    #[test]
    fn group_by_root_splits_a_combined_batch_response() {
        // Two independent roots' subtrees riding the same QueueNodes
        // response (the whole point of batching): each node must land in
        // its own root's group, not get mixed into the other's.
        let nodes = vec![
            root_node(1, "manual", State::Running),
            work_node(10, 1, "alice", "prebuild", State::Running),
            root_node(2, "manual", State::Done),
            work_node(20, 2, "bob", "prebuild", State::Done),
        ];
        let groups = super::group_by_root(nodes);
        assert_eq!(groups.len(), 2);
        assert_eq!(groups[&1].len(), 2, "root 1's group must have its child");
        assert_eq!(groups[&2].len(), 2, "root 2's group must have its child");
        assert!(groups[&1].iter().any(|n| n.id == 10));
        assert!(groups[&2].iter().any(|n| n.id == 20));
    }
}