hyperhive/swarm-controller/src/term_stream.rs

//! `GET /api/agents/{name}/term/stream` — an SSE relay of one agent's
//! terminal rows, sourced straight from the swarm queue.
//!
//! **No `hive` in the path, on purpose.** An agent is not pinned to a
//! hive forever — it can move — so a URL naming one would go stale the
//! moment it did. `crate::agent_status`'s bucket is read only to answer
//! 404 for an agent that has never reported.
//!
//! **Replay, then the live tail, through one consumer.** `hive-agent`
//! publishes each already-classified `TermMsg` row to the core subject
//! `$SWARM.term.{agent}`, and the stream `term-main-{agent}`
//! (`swarm_queue_client::agent_term`) keeps them. The handler creates an
//! ephemeral consumer on that stream starting at the newest
//! `agent_term::REPLAY_ROWS` rows, so the reader gets those and then every
//! new row, in stream order, with no gap and no duplicate between the two.
//! Each event's `id` is the row's stream sequence; a browser that
//! reconnects sends it back as `Last-Event-ID` and resumes right after it.
//!
//! **No stream, no replay.** While the stream does not exist yet (up to a
//! minute after the agent is declared) or cannot be read, the handler
//! falls back to a core subscription: rows published while the reader is
//! attached, and nothing before.
//!
//! **The payload is passed through opaquely**, the same way
//! `crate::status::HiveStatus::snapshot` stores a hive's status
//! snapshot: no reason to depend on `hive-agent`'s `TermMsg` type just
//! to re-serialise a row this daemon never inspects. The bytes on the
//! wire are already what `hive-agent::swarm_term::publish` sent.

use std::convert::Infallible;
use std::time::SystemTime;

use axum::extract::{Path, State};
use axum::http::HeaderMap;
use axum::response::sse::{Event, KeepAlive, Sse};
use futures_util::stream::BoxStream;
use futures_util::{Stream, StreamExt as _};
use swarm_queue_client::agent_term;

use crate::AppState;

#[utoipa::path(
    get,
    path = "/api/agents/{name}/term/stream",
    params(("name" = String, Path, description = "agent whose terminal to stream")),
    responses(
        (status = 200, description = "server-sent event stream; each event's `data` is \
            one already-classified TermMsg row, JSON as hive-agent published it \
            (opaque to this daemon); the newest retained rows first, then live rows. \
            Each event's `id` is the row's stream sequence, and a `Last-Event-ID` \
            request header resumes right after that row",
            body = String, content_type = "text/event-stream"),
        (status = 400, description = "the agent name is not shaped like an identifier \
            (problem+json)", body = String),
        (status = 404, description = "the agent has never reported to the swarm \
            (problem+json)", body = String),
        (status = 503, description = "no swarm queue is configured on this host, no \
            agent-status reader is wired up, or the queue could not be reached \
            (problem+json)", body = String),
        (status = 500, description = "reading the agent's status failed, or the \
            subscribe itself failed (problem+json)", body = String),
    ),
    tag = "agents"
)]
pub(crate) async fn stream_agent_term(
    State(state): State<AppState>,
    Path(agent): Path<String>,
    headers: HeaderMap,
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, problem_details::ProblemDetails> {
    let Some(status) = state.status.as_ref() else {
        return Err(crate::error_problem(
            axum::http::StatusCode::SERVICE_UNAVAILABLE,
            "no swarm queue is configured on this host",
        ));
    };
    let Some(agent_status) = state.agent_status.as_ref() else {
        return Err(crate::error_problem(
            axum::http::StatusCode::SERVICE_UNAVAILABLE,
            "no agent-status reader is configured on this host",
        ));
    };
    let agent = hive_types::Ident::parse(&agent)
        .map_err(|reason| crate::error_problem(axum::http::StatusCode::BAD_REQUEST, reason))?
        .into_string();

    // One-agent "roster": reuses `AgentStatusReader::view`'s already-tested
    // row logic rather than a second, narrower lookup that could disagree
    // with it about what "never reported" means.
    agent_status
        .view(std::slice::from_ref(&agent), SystemTime::now())
        .await
        .map_err(|e| {
            let detail = format!("{e:#}");
            tracing::warn!(%agent, error = %detail, "term stream: reading agent status failed");
            crate::error_problem(axum::http::StatusCode::INTERNAL_SERVER_ERROR, &detail)
        })?
        .into_iter()
        .find_map(|row| row.hive)
        .ok_or_else(|| {
            crate::error_problem(
                axum::http::StatusCode::NOT_FOUND,
                &format!("{agent:?} has never reported to the swarm"),
            )
        })?;

    let client = status.queue_client();
    // Before subscribing: an unconnected client does not fail a
    // subscribe request outright, but there is no point opening a
    // subscription against a queue that is not there — the honest
    // answer is 503, the same shape every other queue-backed route
    // here already uses.
    swarm_queue_client::ensure_connected(&client).map_err(|e| {
        crate::error_problem(
            axum::http::StatusCode::SERVICE_UNAVAILABLE,
            &swarm_queue_client::chain(&e),
        )
    })?;

    let stream = match replay(&client, &agent, last_event_id(&headers)).await {
        Ok(rows) => rows.map(|row| Ok(row_event(&row))).boxed(),
        Err(error) => {
            tracing::warn!(%agent, %error, "term stream: no replay, live rows only");
            let subscriber =
                subscribe(&client, swarm_queue_client::term_subject(&agent), "term").await?;
            subscriber
                .map(|msg| Ok(Event::default().data(String::from_utf8_lossy(&msg.payload))))
                .boxed()
        }
    };
    Ok(Sse::new(stream).keep_alive(KeepAlive::default()))
}

/// One row read from the agent's terminal stream.
struct Row {
    /// Its stream sequence, when the delivery names one.
    seq: Option<u64>,
    payload: bytes::Bytes,
}

/// The agent's rows from its terminal stream: the replay
/// [`agent_term::replay_start`] picks, then every new row.
async fn replay(
    client: &async_nats::Client,
    agent: &str,
    resume_after: Option<u64>,
) -> Result<BoxStream<'static, Row>, String> {
    let js = async_nats::jetstream::new(client.clone());
    let name = agent_term::stream_name(agent);
    let stream = js
        .get_stream(&name)
        .await
        .map_err(|e| format!("opening {name}: {}", swarm_queue_client::chain(&e)))?;
    let start = agent_term::replay_start(stream.cached_info().state.last_sequence, resume_after);
    let deliver_subject = client.new_inbox();
    // The inbox's last token is a fresh nuid, which is also a valid name.
    let consumer_name = format!(
        "term-replay-{}",
        deliver_subject.rsplit('.').next().unwrap_or_default()
    );
    let config = agent_term::replay_consumer_config(agent, consumer_name, deliver_subject, start);
    let messages = stream
        .create_consumer(config)
        .await
        .map_err(|e| {
            format!(
                "creating a consumer on {name}: {}",
                swarm_queue_client::chain(&e)
            )
        })?
        .messages()
        .await
        .map_err(|e| format!("subscribing to {name}: {}", swarm_queue_client::chain(&e)))?;
    tracing::info!(stream = %name, start, "term stream: client attached");
    Ok(rows_until_error(messages, agent.to_owned()))
}

/// The rows of a consumer's deliveries, ending at the first error. A missed
/// heartbeat is such an error, and ending there closes the SSE response, so
/// the browser reconnects with `Last-Event-ID` and gets a fresh consumer.
fn rows_until_error<E: std::fmt::Display>(
    messages: impl Stream<Item = Result<async_nats::jetstream::Message, E>> + Send + 'static,
    agent: String,
) -> BoxStream<'static, Row> {
    messages
        .scan((), move |(), msg| {
            let row = match msg {
                Ok(msg) => Some(Row {
                    seq: msg.info().ok().map(|info| info.stream_sequence),
                    payload: msg.message.payload,
                }),
                Err(e) => {
                    tracing::warn!(%agent, error = %e, "term stream: consumer failed, ending the stream");
                    None
                }
            };
            std::future::ready(row)
        })
        .boxed()
}

/// One row as an SSE event, carrying its stream sequence as the event `id`.
fn row_event(row: &Row) -> Event {
    let event = Event::default().data(String::from_utf8_lossy(&row.payload));
    match row.seq {
        Some(seq) => event.id(seq.to_string()),
        None => event,
    }
}

/// The stream sequence a reconnecting browser last received, from its
/// `Last-Event-ID` header. Absent or unparseable means a fresh reader.
fn last_event_id(headers: &HeaderMap) -> Option<u64> {
    headers
        .get("last-event-id")?
        .to_str()
        .ok()?
        .trim()
        .parse()
        .ok()
}

/// Subscribe to `subject`. `what` names the relay in logs and errors.
pub(crate) async fn subscribe(
    client: &async_nats::Client,
    subject: String,
    what: &str,
) -> Result<async_nats::Subscriber, problem_details::ProblemDetails> {
    let subscriber = client.subscribe(subject.clone()).await.map_err(|e| {
        tracing::warn!(%subject, error = %e, "{what} stream: subscribe failed");
        crate::error_problem(
            axum::http::StatusCode::INTERNAL_SERVER_ERROR,
            &format!("subscribing to {subject} failed: {e}"),
        )
    })?;
    tracing::info!(%subject, "{what} stream: client attached");
    Ok(subscriber)
}

#[cfg(test)]
mod tests {
    use super::*;

    use std::time::Duration;

    use crate::test_nats::TestNats;

    const AGENT: &str = "iris";

    async fn publish_rows(client: &async_nats::Client, rows: std::ops::RangeInclusive<u64>) {
        for i in rows {
            client
                .publish(
                    swarm_queue_client::term_subject(AGENT),
                    format!("row{i}").into(),
                )
                .await
                .expect("publish");
        }
        client.flush().await.expect("flush");
    }

    /// Wait until the stream holds row `last`, so the next step sees it.
    async fn stored(client: &async_nats::Client, last: u64) {
        let js = async_nats::jetstream::new(client.clone());
        for _ in 0..100 {
            let mut stream = js
                .get_stream(agent_term::stream_name(AGENT))
                .await
                .expect("the stream");
            if stream.info().await.expect("info").state.last_sequence >= last {
                return;
            }
            tokio::time::sleep(Duration::from_millis(50)).await;
        }
        panic!("the stream never stored row {last}");
    }

    async fn take(rows: &mut BoxStream<'static, Row>, n: usize) -> Vec<(u64, String)> {
        let mut out = Vec::with_capacity(n);
        for _ in 0..n {
            let row = tokio::time::timeout(Duration::from_secs(5), rows.next())
                .await
                .unwrap_or_else(|_| panic!("row {} of {n} never came; got {out:?}", out.len() + 1))
                .expect("the row stream ended");
            out.push((
                row.seq.expect("a delivered row names its sequence"),
                String::from_utf8_lossy(&row.payload).into_owned(),
            ));
        }
        out
    }

    fn rows(range: std::ops::RangeInclusive<u64>) -> Vec<(u64, String)> {
        range.map(|i| (i, format!("row{i}"))).collect()
    }

    async fn nothing_more(rows: &mut BoxStream<'static, Row>) {
        if let Ok(Some(row)) = tokio::time::timeout(Duration::from_millis(500), rows.next()).await {
            panic!("an extra row: seq {:?}", row.seq);
        }
    }

    /// A consumer the server dropped, as on a restart, ends the row stream
    /// once its heartbeats stop, so the SSE response closes and the browser
    /// reconnects instead of keeping a tail that never delivers.
    #[tokio::test]
    async fn a_dropped_consumer_ends_the_row_stream() {
        use futures_util::TryStreamExt as _;

        let nats = TestNats::start();
        let client = nats.client().await;
        let stream = agent_term::open_or_create(&client, AGENT)
            .await
            .expect("the stream");
        publish_rows(&client, 1..=3).await;
        stored(&client, 3).await;
        let mut replayed = replay(&client, AGENT, None).await.expect("a replay");
        assert_eq!(take(&mut replayed, 3).await, rows(1..=3));
        let consumers: Vec<String> = stream
            .consumer_names()
            .try_collect()
            .await
            .expect("consumer names");
        let [consumer] = consumers.as_slice() else {
            panic!("expected the one replay consumer, got {consumers:?}");
        };
        stream
            .delete_consumer(consumer)
            .await
            .expect("deleting the consumer");
        let end = tokio::time::timeout(agent_term::IDLE_HEARTBEAT * 4, replayed.next())
            .await
            .expect("the row stream ends once heartbeats stop");
        assert!(
            end.is_none(),
            "a row arrived after the consumer was deleted"
        );
    }

    #[tokio::test]
    async fn ensuring_the_terminal_stream_again_leaves_it_as_it_is() {
        let nats = TestNats::start();
        let client = nats.client().await;
        let created = agent_term::open_or_create(&client, AGENT)
            .await
            .expect("the first ensure creates the stream");
        publish_rows(&client, 1..=3).await;
        stored(&client, 3).await;
        agent_term::open_or_create(&client, AGENT)
            .await
            .expect("a second ensure succeeds");
        // Read back from the server, not from what the ensure returned.
        let after = async_nats::jetstream::new(client.clone())
            .get_stream(agent_term::stream_name(AGENT))
            .await
            .expect("the stream")
            .info()
            .await
            .expect("its info")
            .clone();
        assert_eq!(after.config, created.cached_info().config);
        assert_eq!(after.state.messages, 3);
    }

    /// An existing stream is opened, not created again: one made when the
    /// queue announced a different `max_payload` keeps its own config and
    /// rows, where a second create would be refused for the mismatch.
    #[tokio::test]
    async fn an_existing_terminal_stream_with_other_limits_is_opened_as_it_is() {
        let nats = TestNats::start();
        let client = nats.client().await;
        let js = async_nats::jetstream::new(client.clone());
        js.create_stream(async_nats::jetstream::stream::Config {
            name: agent_term::stream_name(AGENT),
            subjects: vec![agent_term::stream_subject(AGENT)],
            max_message_size: 4096,
            ..Default::default()
        })
        .await
        .expect("an earlier stream");
        publish_rows(&client, 1..=2).await;
        stored(&client, 2).await;
        let opened = agent_term::open_or_create(&client, AGENT)
            .await
            .expect("an existing stream opens");
        assert_eq!(opened.cached_info().config.max_message_size, 4096);
        assert_eq!(opened.cached_info().state.messages, 2);
    }

    /// Live rows published after the reader attached, before it read any
    /// replayed row, still arrive after the whole replay.
    #[tokio::test]
    async fn the_replay_reaches_the_reader_before_the_live_tail() {
        let nats = TestNats::start();
        let client = nats.client().await;
        agent_term::open_or_create(&client, AGENT)
            .await
            .expect("the stream");
        publish_rows(&client, 1..=150).await;
        stored(&client, 150).await;
        let mut replayed = replay(&client, AGENT, None).await.expect("a replay");
        publish_rows(&client, 151..=153).await;
        assert_eq!(take(&mut replayed, 103).await, rows(51..=153));
        nothing_more(&mut replayed).await;
    }

    /// A reader that reconnects with the last id it was sent gets every row
    /// after it exactly once, including those published while it was away.
    #[tokio::test]
    async fn a_reconnect_gets_every_later_row_once() {
        let nats = TestNats::start();
        let client = nats.client().await;
        agent_term::open_or_create(&client, AGENT)
            .await
            .expect("the stream");
        publish_rows(&client, 1..=150).await;
        stored(&client, 150).await;
        let mut first = replay(&client, AGENT, None).await.expect("a replay");
        let seen = take(&mut first, 100).await;
        drop(first);
        publish_rows(&client, 151..=160).await;
        stored(&client, 160).await;
        let last_id = seen.last().expect("rows were seen").0.to_string();
        let resume_after = last_event_id(&with_last_event_id(&last_id));
        let mut resumed = replay(&client, AGENT, resume_after)
            .await
            .expect("a resumed replay");
        publish_rows(&client, 161..=163).await;
        assert_eq!(take(&mut resumed, 13).await, rows(151..=163));
        nothing_more(&mut resumed).await;
    }

    fn with_last_event_id(value: &str) -> HeaderMap {
        let mut headers = HeaderMap::new();
        headers.insert(
            "last-event-id",
            value.parse().expect("a valid header value"),
        );
        headers
    }

    #[test]
    fn a_reconnect_resumes_after_the_sequence_it_sends() {
        assert_eq!(last_event_id(&with_last_event_id("917")), Some(917));
        assert_eq!(
            agent_term::replay_start(1000, last_event_id(&with_last_event_id("917"))),
            918
        );
    }

    #[test]
    fn no_or_a_foreign_last_event_id_is_a_fresh_reader() {
        assert_eq!(last_event_id(&HeaderMap::new()), None);
        assert_eq!(last_event_id(&with_last_event_id("r12")), None);
        assert_eq!(last_event_id(&with_last_event_id("")), None);
    }
}