hyperhive/swarmctl/src/watch.rs

//! `swarmctl agent watch` — follow one agent's terminal rows from the CLI.
//!
//! Reads `GET /api/agents/{name}/term/stream`, the SSE relay of the rows the
//! agent publishes to the swarm queue: the newest retained rows first, then
//! live ones. This is the swarm-level counterpart of `hivectl agent watch`
//! and prints one line per row, `<icon> <summary>`, as that verb prints a
//! tool call.
//!
//! The ordinary request path in [`crate::agent`] buffers a whole response, so
//! the stream has its own: a unix-socket `hyper` request, or a bearer-
//! authenticated gateway request whose body is read chunk by chunk.

use anyhow::{Context as _, Result, bail};
use http_body_util::{BodyExt as _, Empty};
use hyper_util::rt::TokioIo;
use tokio::net::UnixStream;

use crate::agent::{connect_hint, describe_error, parse_ident};
use crate::remote::Target;

/// Run `swarmctl agent watch`. Returns when the stream ends or on Ctrl-C.
pub(crate) fn watch(target: &Target, agent: &str) -> Result<()> {
    let agent = parse_ident(agent, "agent name")?;
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .context("starting a tokio runtime for the controller request")?;
    rt.block_on(async {
        tokio::select! {
            result = follow(target, &agent) => result,
            _ = tokio::signal::ctrl_c() => Ok(()),
        }
    })
}

async fn follow(target: &Target, agent: &str) -> Result<()> {
    let uri = format!("/api/agents/{agent}/term/stream");
    let mut body = open(target, &uri).await?;
    let mut out = std::io::stdout().lock();
    if !emit(&mut out, &format!("watching {agent} — Ctrl-C to stop"))? {
        return Ok(());
    }
    let mut frames = Frames::default();
    while let Some(chunk) = body.next().await? {
        for frame in frames.push(&chunk) {
            if let Some(line) = render_frame(&frame)
                && !emit(&mut out, &line)?
            {
                return Ok(());
            }
        }
    }
    emit(&mut out, "(stream closed by swarm-controller)")?;
    Ok(())
}

/// Write one line and flush. `Ok(false)` means the reader went away (a closed
/// pipe, as in `watch x | head`), which ends the watch quietly; any other
/// write error is an error.
fn emit(out: &mut impl std::io::Write, line: &str) -> Result<bool> {
    match writeln!(out, "{line}").and_then(|()| out.flush()) {
        Ok(()) => Ok(true),
        Err(e) if e.kind() == std::io::ErrorKind::BrokenPipe => Ok(false),
        Err(e) => Err(e).context("writing to stdout"),
    }
}

/// Replace every control character except tab with U+FFFD. Rows are
/// agent-authored, and an escape sequence printed as-is would act on the
/// operator's terminal.
fn sanitize(text: &str) -> String {
    text.chars()
        .map(|c| {
            if c.is_control() && c != '\t' {
                '\u{FFFD}'
            } else {
                c
            }
        })
        .collect()
}

/// A response body read in pieces, whichever transport carried it.
enum Body {
    Socket(hyper::body::Incoming),
    Gateway(reqwest::Response),
}

impl Body {
    /// The next piece of data, or `None` at the end of the stream.
    async fn next(&mut self) -> Result<Option<bytes::Bytes>> {
        match self {
            Self::Socket(body) => {
                while let Some(frame) = body.frame().await {
                    let frame = frame.context("reading swarm-controller's event stream")?;
                    if let Ok(data) = frame.into_data() {
                        return Ok(Some(data));
                    }
                }
                Ok(None)
            }
            Self::Gateway(resp) => resp
                .chunk()
                .await
                .context("reading swarm-controller's event stream"),
        }
    }
}

/// Open the stream, failing with the controller's own error on a non-2xx.
async fn open(target: &Target, uri: &str) -> Result<Body> {
    match target {
        Target::Socket(socket) => {
            let stream = UnixStream::connect(socket)
                .await
                .with_context(|| connect_hint(socket))?;
            let (mut sender, conn) = hyper::client::conn::http1::handshake(TokioIo::new(stream))
                .await
                .with_context(|| format!("HTTP handshake with {}", socket.display()))?;
            tokio::spawn(conn);
            let req = hyper::Request::builder()
                .method(hyper::Method::GET)
                .uri(uri)
                .header(hyper::header::HOST, "localhost")
                .header(hyper::header::ACCEPT, "text/event-stream")
                .body(Empty::<bytes::Bytes>::new())
                .context("building the watch request")?;
            let resp = sender
                .send_request(req)
                .await
                .context("sending the watch request to swarm-controller")?;
            let status = resp.status();
            if !status.is_success() {
                let body = resp
                    .into_body()
                    .collect()
                    .await
                    .context("reading swarm-controller's response")?
                    .to_bytes();
                bail!("{}", describe_error(status, &body));
            }
            Ok(Body::Socket(resp.into_body()))
        }
        Target::Swarm(swarm) => {
            let resp = swarm.open_stream(uri, "watch").await?;
            let status = resp.status();
            if !status.is_success() {
                let body = resp
                    .bytes()
                    .await
                    .context("reading the gateway's response")?;
                bail!("{}", describe_error(status, &body));
            }
            Ok(Body::Gateway(resp))
        }
    }
}

/// Splits a byte stream into SSE frames. A chunk boundary has no relation to
/// a frame boundary, so bytes are held until the blank line that ends one.
#[derive(Default)]
struct Frames {
    pending: Vec<u8>,
}

impl Frames {
    /// Add `chunk`; returns every frame it completed, without the separator.
    fn push(&mut self, chunk: &[u8]) -> Vec<String> {
        self.pending.extend_from_slice(chunk);
        let mut out = Vec::new();
        while let Some(pos) = self.pending.windows(2).position(|w| w == b"\n\n") {
            let frame: Vec<u8> = self.pending.drain(..pos + 2).take(pos).collect();
            out.push(String::from_utf8_lossy(&frame).into_owned());
        }
        out
    }
}

/// One frame as a terminal line: the row's `<icon> <summary>`. A frame with
/// no `data:` line (a keep-alive comment) prints nothing; a payload that is
/// not a row prints raw, so a shape this doesn't know stays visible.
fn render_frame(frame: &str) -> Option<String> {
    render_raw(frame).map(|line| sanitize(&line))
}

fn render_raw(frame: &str) -> Option<String> {
    let data = frame
        .lines()
        .filter_map(|l| l.strip_prefix("data:"))
        .map(str::trim_start)
        .collect::<Vec<_>>()
        .join("\n");
    if data.is_empty() {
        return None;
    }
    let Ok(row) = serde_json::from_str::<serde_json::Value>(&data) else {
        return Some(format!("! {data}"));
    };
    let field = |k: &str| row.get(k).and_then(serde_json::Value::as_str);
    Some(match (field("icon"), field("summary")) {
        (Some(icon), Some(summary)) => format!("{icon} {summary}"),
        (None, Some(summary)) => summary.to_owned(),
        _ => format!("! {data}"),
    })
}

#[cfg(test)]
mod tests {
    use super::{Frames, render_frame};

    #[test]
    fn a_frame_split_across_chunks_is_emitted_once_whole() {
        let mut frames = Frames::default();
        assert!(frames.push(b"id: 1\ndata: {\"summ").is_empty());
        let done = frames.push(b"ary\":\"x\"}\n\nid: 2\nda");
        assert_eq!(done, vec!["id: 1\ndata: {\"summary\":\"x\"}".to_owned()]);
        let rest = frames.push(b"ta: y\n\n");
        assert_eq!(rest, vec!["id: 2\ndata: y".to_owned()]);
    }

    #[test]
    fn a_multibyte_char_split_across_chunks_survives() {
        let bytes = "data: ✅ ok\n\n".as_bytes();
        let mut frames = Frames::default();
        assert!(frames.push(&bytes[..8]).is_empty());
        assert_eq!(frames.push(&bytes[8..]), vec!["data: ✅ ok".to_owned()]);
    }

    #[test]
    fn a_row_prints_icon_and_summary() {
        let line = render_frame(
            "id: 7\ndata: {\"ts\":\"t\",\"icon\":\"🔧\",\"level\":\"info\",\"summary\":\"Bash ls\"}",
        );
        assert_eq!(line.as_deref(), Some("🔧 Bash ls"));
        let bare = render_frame("data: {\"summary\":\"plain\"}");
        assert_eq!(bare.as_deref(), Some("plain"));
    }

    #[test]
    fn keep_alives_print_nothing_and_unknown_payloads_print_raw() {
        assert_eq!(render_frame(": keep-alive"), None);
        assert_eq!(
            render_frame("data: not json").as_deref(),
            Some("! not json")
        );
        assert_eq!(
            render_frame("data: {\"x\":1}").as_deref(),
            Some("! {\"x\":1}")
        );
    }

    #[test]
    fn control_characters_in_a_row_never_reach_the_terminal() {
        let line = render_frame("data: {\"icon\":\"\\u001b]0;x\\u0007\",\"summary\":\"a\\u001b[2Jb\\tc\\r\\u009bd\\u007f\"}")
            .expect("row");
        assert_eq!(
            line,
            "\u{FFFD}]0;x\u{FFFD} a\u{FFFD}[2Jb\tc\u{FFFD}\u{FFFD}d\u{FFFD}"
        );
        let raw = render_frame("data: \u{1b}[2J").expect("raw");
        assert_eq!(raw, "! \u{FFFD}[2J");
    }

    struct Failing(std::io::ErrorKind);

    impl std::io::Write for Failing {
        fn write(&mut self, _: &[u8]) -> std::io::Result<usize> {
            Err(self.0.into())
        }

        fn flush(&mut self) -> std::io::Result<()> {
            Ok(())
        }
    }

    #[test]
    fn a_closed_pipe_ends_the_watch_quietly_and_other_write_errors_fail() {
        let mut ok = Vec::new();
        assert!(super::emit(&mut ok, "x").expect("writes"));
        assert_eq!(ok, b"x\n");
        let closed = super::emit(&mut Failing(std::io::ErrorKind::BrokenPipe), "x");
        assert!(!closed.expect("broken pipe is not an error"));
        let other = super::emit(&mut Failing(std::io::ErrorKind::PermissionDenied), "x");
        assert!(other.is_err());
    }

    /// Serves one canned response on a fresh unix socket and returns its path.
    fn serve_once(name: &str, response: Vec<&'static [u8]>) -> std::path::PathBuf {
        use std::io::{Read as _, Write as _};
        let dir = std::env::temp_dir().join(format!("swarmctl-{name}-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).expect("dir");
        let path = dir.join("controller.sock");
        let listener = std::os::unix::net::UnixListener::bind(&path).expect("bind");
        std::thread::spawn(move || {
            let (mut conn, _) = listener.accept().expect("accept");
            let mut seen = Vec::new();
            let mut buf = [0u8; 512];
            while !seen.windows(4).any(|w| w == b"\r\n\r\n") {
                let n = conn.read(&mut buf).expect("read request");
                seen.extend_from_slice(&buf[..n]);
            }
            for part in response {
                conn.write_all(part).expect("write");
                conn.flush().expect("flush");
                std::thread::sleep(std::time::Duration::from_millis(50));
            }
        });
        path
    }

    fn rt() -> tokio::runtime::Runtime {
        tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("runtime")
    }

    /// Rows arrive as the server sends them, a frame split across two chunks
    /// is rendered once, and the end of the stream ends the loop.
    #[test]
    fn rows_stream_over_the_socket_until_it_closes() {
        let socket = serve_once(
            "watch-ok",
            vec![
                b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\n\r\n",
                b"1b\r\nid: 1\ndata: {\"icon\":\"A\",\"su\r\n",
                b"2e\r\nmmary\":\"one\"}\n\nid: 2\ndata: {\"summary\":\"two\"}\n\n\r\n",
                b"0\r\n\r\n",
            ],
        );
        let lines = rt().block_on(async {
            let mut body = super::open(&crate::remote::Target::Socket(socket), "/x")
                .await
                .expect("opens");
            let mut frames = Frames::default();
            let mut lines = Vec::new();
            while let Some(chunk) = body.next().await.expect("reads") {
                lines.extend(frames.push(&chunk).iter().filter_map(|f| render_frame(f)));
            }
            lines
        });
        assert_eq!(lines, vec!["A one".to_owned(), "two".to_owned()]);
    }

    /// A non-2xx is an error carrying the controller's own sentence.
    #[test]
    fn a_refusal_surfaces_the_controllers_detail() {
        let body: &'static [u8] =
            b"{\"status\":404,\"title\":\"Not Found\",\"detail\":\"\\\"ghost\\\" has never reported\"}";
        let socket = serve_once(
            "watch-404",
            vec![
                b"HTTP/1.1 404 Not Found\r\ncontent-type: application/problem+json\r\nconnection: close\r\n\r\n",
                body,
            ],
        );
        let err = rt()
            .block_on(super::open(&crate::remote::Target::Socket(socket), "/x"))
            .err()
            .expect("refused");
        let msg = err.to_string();
        assert!(
            msg.contains("404") && msg.contains("has never reported"),
            "{msg}"
        );
    }
}