hyperhive/swarm-controller/src/subagent_term.rs

//! An agent's subagents: the stream that keeps their rows, which ones the
//! swarm has rows for, and a live SSE relay of one subagent's terminal.
//!
//! The agent's subagent daemon publishes each subagent's classified rows on
//! `$SWARM.term.{agent}.sub.{subagent}` under the agent's own queue credential,
//! which may publish there and nothing else of these streams. [`spawn`]
//! creates the stream `term-sub-{agent}` for every live agent, with the config
//! `swarm_queue_client::subagent_term` spells, and in the same pass the
//! agent's own terminal stream (`swarm_queue_client::agent_term`). Subagents are spawned on demand
//! under caller-chosen names, so the list is the set of subjects that stream
//! holds rows for, read with `STREAM.INFO`. A missing stream lists nothing.
//! Each subagent carries the stream's own publish time of its newest row, so
//! a reader can tell a subagent that is running from one that has gone quiet.
//!
//! **No agent-status lookup.** Unlike [`crate::term_stream`], an agent that
//! has never reported gets an empty stream, not a 404.
//!
//! **Live tail only, output only.** The relay is a core subscription, so a
//! reader sees rows published while it is attached; unlike the agent's own
//! terminal route, it replays nothing from the stream. Nothing is sent toward a subagent: subagents take no
//! input from the swarm.

use std::collections::BTreeSet;
use std::convert::Infallible;
use std::sync::Arc;
use std::time::Duration;

use axum::Json;
use axum::extract::{Path, State};
use axum::http::StatusCode;
use axum::response::sse::{Event, KeepAlive, Sse};
use futures_util::{Stream, StreamExt as _, TryStreamExt as _};
use serde::Serialize;
use swarm_queue_client::{agent_term, subagent_term};
use utoipa::ToSchema;

use crate::AppState;
use crate::wanted::WantedWriter;

/// How often [`spawn`] checks every live agent's stream. Also how long a newly
/// declared agent's subagents can go unlisted.
const ENSURE_INTERVAL: Duration = Duration::from_mins(1);

/// How long [`list_subagents`] waits for one subagent's newest row before
/// listing it without a time. The lookups run concurrently, so a stream that
/// does not answer costs the list this long per [`LOOKUP_CONCURRENCY`]
/// subagents.
const LOOKUP_TIMEOUT: Duration = Duration::from_secs(2);

/// How many newest-row lookups [`list_subagents`] has in flight at once.
const LOOKUP_CONCURRENCY: usize = 16;

/// Create the subagent stream and the terminal stream (`agent_term`) of every
/// live agent that lacks one, now and every [`ENSURE_INTERVAL`] after. An
/// existing stream is left as it is.
///
/// The live agents are the ones some hive's wanted-state declaration names in
/// a state other than `Destroyed` (`agent_renewal::live_agents`). A pass skips
/// a hive whose declaration cannot be read and an agent whose stream cannot be
/// created, logging each; a pass while the queue is down does nothing. It
/// never stops the daemon.
pub(crate) fn spawn(client: async_nats::Client, wanted: Arc<WantedWriter>, hives: Vec<String>) {
    tokio::spawn(async move {
        let mut ticker = tokio::time::interval(ENSURE_INTERVAL);
        loop {
            ticker.tick().await;
            if swarm_queue_client::ensure_connected(&client).is_err() {
                tracing::debug!("terminal streams: queue not connected; next pass");
                continue;
            }
            for agent in live_agents(&wanted, &hives).await {
                if let Err(e) = subagent_term::open_or_create(&client, &agent).await {
                    tracing::warn!(
                        %agent,
                        error = %swarm_queue_client::chain(&e),
                        "subagent streams: creating the agent's stream failed; next pass"
                    );
                }
                if let Err(e) = agent_term::open_or_create(&client, &agent).await {
                    tracing::warn!(
                        %agent,
                        error = %swarm_queue_client::chain(&e),
                        "terminal streams: creating the agent's stream failed; next pass"
                    );
                }
            }
        }
    });
}

/// Every agent the readable declarations of `hives` keep alive.
pub(crate) async fn live_agents(wanted: &WantedWriter, hives: &[String]) -> BTreeSet<String> {
    let mut declarations = Vec::with_capacity(hives.len());
    for hive in hives {
        match wanted.view(hive).await {
            Ok(Some(d)) => declarations.push(d),
            Ok(None) => {}
            Err(e) => tracing::warn!(
                %hive,
                error = %format!("{e:#}"),
                "subagent streams: reading the hive's wanted-state declaration failed"
            ),
        }
    }
    crate::agent_renewal::live_agents(&declarations)
}

/// One subagent of an agent, as the list endpoint reports it.
#[derive(Debug, PartialEq, Eq, Serialize, ToSchema)]
pub(crate) struct SubagentEntry {
    /// The subagent's name, as it appears in its stream subject.
    pub name: String,
    /// When its newest row in the stream was published, unix seconds, by the
    /// stream's own clock. `None` when that row could not be read.
    pub last_row_unix: Option<i64>,
}

#[utoipa::path(
    get,
    path = "/api/agents/{name}/subagents",
    params(("name" = String, Path, description = "agent whose subagents to list")),
    responses(
        (status = 200, description = "the agent's subagents with terminal rows in the \
            swarm queue, sorted by name, each with the publish time of its newest row; \
            empty when the agent has published none",
            body = Vec<SubagentEntry>),
        (status = 400, description = "the agent name is not shaped like an identifier \
            (problem+json)", body = String),
        (status = 503, description = "no swarm queue is configured on this host, or the \
            queue could not be reached (problem+json)", body = String),
        (status = 500, description = "reading the agent's subagent stream failed \
            (problem+json)", body = String),
    ),
    tag = "agents"
)]
pub(crate) async fn list_subagents(
    State(state): State<AppState>,
    Path(agent): Path<String>,
) -> Result<Json<Vec<SubagentEntry>>, problem_details::ProblemDetails> {
    let agent = ident(&agent).map_err(bad_request)?;
    let client = connected_client(&state).map_err(|e| unavailable(&e))?;
    let js = async_nats::jetstream::new(client);
    let stream_name = subagent_term::stream_name(&agent);
    let stream = match js.get_stream(&stream_name).await {
        Ok(stream) => stream,
        Err(e) if is_stream_not_found(&e) => return Ok(Json(Vec::new())),
        Err(e) => return Err(read_failed(&agent, &e)),
    };
    let subjects: Vec<(String, usize)> = stream
        .info_with_subjects(subagent_term::stream_subjects(&agent))
        .await
        .map_err(|e| read_failed(&agent, &e))?
        .try_collect()
        .await
        .map_err(|e| read_failed(&agent, &e))?;
    let names = subagent_term::subagent_names(&agent, subjects.iter().map(|(s, _)| s.as_str()));
    // Owned names and `async move`, not a closure over `&names`: a borrowed
    // item type trips a "`FnOnce` is not general enough" inference error where
    // `routes!` checks this fn as a handler.
    let (stream_ref, agent_ref) = (&stream, &agent);
    let lookups: Vec<Result<i64, String>> = futures_util::stream::iter(names.clone())
        .map(|name| async move { last_row_unix(stream_ref, agent_ref, &name).await })
        .buffered(LOOKUP_CONCURRENCY)
        .collect()
        .await;
    let failed = lookups.iter().filter(|l| l.is_err()).count();
    if let Some(Err(error)) = lookups.iter().find(|l| l.is_err()) {
        tracing::warn!(
            %agent,
            failed,
            subagents = names.len(),
            %error,
            "listing subagents: reading the newest row failed (first error shown)"
        );
    }
    let entries = names
        .into_iter()
        .zip(lookups)
        .map(|(name, lookup)| SubagentEntry {
            name,
            last_row_unix: lookup.ok(),
        })
        .collect();
    Ok(Json(entries))
}

/// The publish time of a subagent's newest row, or why it could not be read.
/// A row that cannot be read, aged out between the listing and this lookup
/// say, is a missing time and not a failed list.
async fn last_row_unix(
    stream: &async_nats::jetstream::stream::Stream,
    agent: &str,
    subagent: &str,
) -> Result<i64, String> {
    let subject = subagent_term::subject(agent, subagent);
    let lookup = stream.get_last_raw_message_by_subject(&subject);
    match tokio::time::timeout(LOOKUP_TIMEOUT, lookup).await {
        Ok(Ok(row)) => Ok(row.time.unix_timestamp()),
        Ok(Err(e)) => Err(format!("{subagent}: {}", swarm_queue_client::chain(&e))),
        Err(_) => Err(format!("{subagent}: no answer in {LOOKUP_TIMEOUT:?}")),
    }
}

#[utoipa::path(
    get,
    path = "/api/agents/{name}/subagents/{subagent}/term/stream",
    params(
        ("name" = String, Path, description = "agent the subagent belongs to"),
        ("subagent" = String, Path, description = "subagent whose terminal to stream"),
    ),
    responses(
        (status = 200, description = "server-sent event stream; each event's `data` is \
            one already-classified TermMsg row, JSON as the agent's subagent daemon \
            published it (opaque to this daemon) — live only, no replay",
            body = String, content_type = "text/event-stream"),
        (status = 400, description = "the agent or subagent name is not shaped like an \
            identifier (problem+json)", body = String),
        (status = 503, description = "no swarm queue is configured on this host, or the \
            queue could not be reached (problem+json)", body = String),
        (status = 500, description = "the subscribe itself failed (problem+json)",
            body = String),
    ),
    tag = "agents"
)]
pub(crate) async fn stream_subagent_term(
    State(state): State<AppState>,
    Path((agent, subagent)): Path<(String, String)>,
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, problem_details::ProblemDetails> {
    let agent = ident(&agent).map_err(bad_request)?;
    let subagent = ident(&subagent).map_err(bad_request)?;
    let client = connected_client(&state).map_err(|e| unavailable(&e))?;
    let subscriber = crate::term_stream::subscribe(
        &client,
        subagent_term::subject(&agent, &subagent),
        "subagent term",
    )
    .await?;
    let stream =
        subscriber.map(|msg| Ok(Event::default().data(String::from_utf8_lossy(&msg.payload))));
    Ok(Sse::new(stream).keep_alive(KeepAlive::default()))
}

/// The queue client, or why there is none: the 503 every queue-backed route
/// answers.
fn connected_client(state: &AppState) -> Result<async_nats::Client, String> {
    let status = state
        .status
        .as_ref()
        .ok_or_else(|| "no swarm queue is configured on this host".to_owned())?;
    let client = status.queue_client();
    swarm_queue_client::ensure_connected(&client).map_err(|e| swarm_queue_client::chain(&e))?;
    Ok(client)
}

/// A path segment as a single subject token, or why it is not one.
fn ident(name: &str) -> Result<String, &'static str> {
    hive_types::Ident::parse(name).map(hive_types::Ident::into_string)
}

fn bad_request(reason: &str) -> problem_details::ProblemDetails {
    crate::error_problem(StatusCode::BAD_REQUEST, reason)
}

fn unavailable(detail: &str) -> problem_details::ProblemDetails {
    crate::error_problem(StatusCode::SERVICE_UNAVAILABLE, detail)
}

fn is_stream_not_found(e: &async_nats::jetstream::context::GetStreamError) -> bool {
    matches!(
        e.kind(),
        async_nats::jetstream::context::GetStreamErrorKind::JetStream(js)
            if js.error_code() == async_nats::jetstream::ErrorCode::STREAM_NOT_FOUND
    )
}

fn read_failed(agent: &str, e: &dyn std::error::Error) -> problem_details::ProblemDetails {
    let detail = swarm_queue_client::chain(e);
    tracing::warn!(%agent, error = %detail, "listing subagents: reading the stream failed");
    crate::error_problem(StatusCode::INTERNAL_SERVER_ERROR, &detail)
}

#[cfg(test)]
mod tests {
    use swarm_queue_client::subagent_term::subagent_names;

    /// The list is what `STREAM.INFO` reports for the stream's subjects, by
    /// name: a subagent that published once and one that published many rows
    /// list the same way, and the order is the name's, not the stream's.
    #[test]
    fn the_list_is_the_streams_subjects_by_name() {
        let reported = [
            ("$SWARM.term.atlas.sub.scout".to_owned(), 12),
            ("$SWARM.term.atlas.sub.builder".to_owned(), 1),
        ];
        assert_eq!(
            subagent_names("atlas", reported.iter().map(|(s, _)| s.as_str())),
            ["builder", "scout"]
        );
    }
}