hyperhive/swarm-queue-client/src/agent_term.rs

//! An agent's own terminal: the stream that keeps its rows, and the consumer
//! a reader replays them through.
//!
//! The agent publishes its classified terminal rows on
//! [`crate::term_subject`] with a plain core publish, under its own queue
//! credential, which may publish there and nothing else of these streams.
//! `swarm-controller` creates one stream per agent with the fixed config
//! [`crate::agent_term::open_or_create`] spells, so the queue keeps the rows a
//! reader who attaches later replays.
//!
//! The stream captures exactly one literal subject, `$SWARM.term.<agent>`.
//! [`crate::subagent_term`]'s stream captures `$SWARM.term.<agent>.sub.>`,
//! which only matches subjects with at least two more tokens, so no subject
//! is in both and the two streams never overlap. The `JetStream` server refuses
//! to create a stream whose subjects overlap another's.
//!
//! The names live here because three crates must agree on them: the
//! publisher in `hive-agent`, `swarm-controller`, and the reader's consumer
//! grant in `swarm-nats-auth`.

/// The token between the prefix and the agent in [`stream_name`]. It differs
/// from `term-sub-` at the same position, so no agent name makes this
/// stream's name equal another agent's subagent stream.
const STREAM_KIND: &str = "main";

/// How long the stream keeps a row. The same window as the subagent stream.
pub const MAX_AGE: std::time::Duration = std::time::Duration::from_hours(24);

/// How large one agent's stream may grow before the discard policy reclaims
/// space. The same cap as the subagent stream, for the same reason: it
/// bounds the disk one agent can claim.
#[cfg(feature = "agent-term")]
const MAX_BYTES: i64 = 64 * 1024 * 1024;

/// How many of the newest rows a reader that attaches fresh is replayed
/// before the live tail. The agent's own web UI loads the same number of
/// events as its first history page.
pub const REPLAY_ROWS: u64 = 100;

/// How often the server sends an idle replay consumer a heartbeat. Two missed
/// in a row end the reader's row stream, so a consumer the server dropped,
/// as on a restart, ends the SSE response and the browser reconnects.
pub const IDLE_HEARTBEAT: std::time::Duration = std::time::Duration::from_secs(5);

/// The stream holding `agent`'s terminal rows: `term-main-<agent>`.
#[must_use]
pub fn stream_name(agent: &str) -> String {
    format!("term-{STREAM_KIND}-{agent}")
}

/// The one subject [`stream_name`]'s stream captures.
#[must_use]
pub fn stream_subject(agent: &str) -> String {
    crate::term_subject(agent)
}

/// The reader's grant to create a replay consumer: on any stream, under any
/// consumer name, in the extended create form whose last tokens are the
/// consumer's filter.
///
/// The server refuses a create whose payload filter differs from those
/// tokens, so a consumer made under this grant delivers only rows on
/// subjects below `$SWARM.term`. The grant's last `*` also matches a literal
/// `*` or `>` token, so besides one agent's [`crate::term_subject`] it admits
/// `$SWARM.term.*` and `$SWARM.term.>`, which take in every agent's and every
/// subagent's terminal. The bare `CONSUMER.CREATE.<stream>` form, whose
/// filter rides only in the payload, is not covered.
#[must_use]
pub fn consumer_create_grant() -> String {
    format!(
        "$JS.API.CONSUMER.CREATE.*.*.{}",
        crate::term_subject_wildcard()
    )
}

/// The first stream sequence to replay from.
///
/// `resume_after` is the last sequence the reader already holds, from a
/// reconnecting browser's `Last-Event-ID`; the replay continues right after
/// it. A fresh reader gets the newest [`REPLAY_ROWS`] below `last_sequence`,
/// the stream's newest row when it was read. A sequence the stream has since
/// discarded starts the consumer at the oldest row it still has.
///
/// A `resume_after` above `last_sequence` names no row this stream ever held,
/// as after the stream was created again, and is read as a fresh reader: the
/// server accepts a start past the end and delivers nothing until the stream
/// reaches it.
#[must_use]
pub fn replay_start(last_sequence: u64, resume_after: Option<u64>) -> u64 {
    match resume_after {
        Some(seen) if seen <= last_sequence => seen + 1,
        _ => last_sequence.saturating_sub(REPLAY_ROWS - 1).max(1),
    }
}

/// The config [`open_or_create`] creates `agent`'s stream with.
///
/// `max_message_size` tracks the queue's live `max_payload`, as the subagent
/// stream's does, so the stream never rejects a row the queue accepted.
#[cfg(feature = "agent-term")]
fn stream_config(agent: &str, max_message_size: i32) -> async_nats::jetstream::stream::Config {
    async_nats::jetstream::stream::Config {
        name: stream_name(agent),
        description: Some(format!("Terminal rows of {agent}")),
        subjects: vec![stream_subject(agent)],
        max_age: MAX_AGE,
        max_bytes: MAX_BYTES,
        discard: async_nats::jetstream::stream::DiscardPolicy::Old,
        max_message_size,
        ..Default::default()
    }
}

/// Open `agent`'s terminal stream, creating it if it does not exist yet. An
/// existing stream is opened as it is.
#[cfg(feature = "agent-term")]
pub async fn open_or_create(
    client: &async_nats::Client,
    agent: &str,
) -> Result<async_nats::jetstream::stream::Stream, crate::Error> {
    let js = async_nats::jetstream::new(client.clone());
    let name = stream_name(agent);
    match js.get_stream(&name).await {
        Ok(stream) => Ok(stream),
        Err(e) => {
            tracing::info!(stream = %name, reason = %e, "agent terminal stream not available, creating it");
            let max_message_size = i32::try_from(crate::max_payload(client)).unwrap_or(i32::MAX);
            js.create_stream(stream_config(agent, max_message_size))
                .await
                .map_err(|source| crate::Error::CreateAgentStream {
                    stream: name,
                    source,
                })
        }
    }
}

/// The consumer that replays `agent`'s rows from `start_sequence` and then
/// follows new ones, delivered to `deliver_subject` in stream order.
///
/// Ephemeral and unacknowledged: the server drops it once nothing listens
/// on `deliver_subject`, so a reader that goes away needs no delete grant.
/// `name` must be set, because the server only takes the filter from the
/// subject [`consumer_create_grant`] covers when the create names a consumer.
#[cfg(feature = "agent-term")]
#[must_use]
pub fn replay_consumer_config(
    agent: &str,
    name: String,
    deliver_subject: String,
    start_sequence: u64,
) -> async_nats::jetstream::consumer::push::Config {
    async_nats::jetstream::consumer::push::Config {
        name: Some(name),
        deliver_subject,
        filter_subject: stream_subject(agent),
        deliver_policy: async_nats::jetstream::consumer::DeliverPolicy::ByStartSequence {
            start_sequence,
        },
        ack_policy: async_nats::jetstream::consumer::AckPolicy::None,
        idle_heartbeat: IDLE_HEARTBEAT,
        ..Default::default()
    }
}

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

    /// Token-wise NATS subject overlap: whether some literal subject matches
    /// both patterns.
    fn overlaps(a: &str, b: &str) -> bool {
        let (mut a, mut b) = (a.split('.'), b.split('.'));
        loop {
            match (a.next(), b.next()) {
                (Some(">"), Some(_)) | (Some(_), Some(">")) | (None, None) => return true,
                (None, _) | (_, None) => return false,
                (Some(x), Some(y)) if x != "*" && y != "*" && x != y => return false,
                _ => {}
            }
        }
    }

    #[test]
    fn names_match_the_publisher_and_the_grant() {
        assert_eq!(stream_name("iris"), "term-main-iris");
        assert_eq!(stream_subject("iris"), "$SWARM.term.iris");
        assert_eq!(
            consumer_create_grant(),
            "$JS.API.CONSUMER.CREATE.*.*.$SWARM.term.*"
        );
    }

    /// The server refuses to create a stream whose subjects overlap
    /// another's, so this stream must share no subject with any agent's
    /// subagent stream, its own agent's included.
    #[test]
    fn the_subject_is_disjoint_from_every_subagent_stream() {
        for other in ["iris", "argus", "sub", "iris-2"] {
            assert!(
                !overlaps(
                    &stream_subject("iris"),
                    &crate::subagent_term::stream_subjects(other)
                ),
                "term-main-iris overlaps term-sub-{other}"
            );
        }
        // Control: the checker does report a real overlap.
        assert!(overlaps(
            "$SWARM.term.iris.>",
            &crate::subagent_term::stream_subjects("iris")
        ));
        assert!(overlaps(&stream_subject("iris"), "$SWARM.term.*"));
    }

    /// No agent name makes this stream's name another agent's subagent
    /// stream: an agent may be called `sub-x`.
    #[test]
    fn the_name_never_collides_with_a_subagent_stream() {
        for (agent, other) in [("sub-x", "x"), ("x", "sub-x"), ("main-x", "x")] {
            assert_ne!(stream_name(agent), crate::subagent_term::stream_name(other));
        }
    }

    /// The create the controller sends is one the grant covers, and every
    /// filter the grant admits stays below `$SWARM.term`.
    #[test]
    fn the_grant_admits_only_filters_below_the_terminal_prefix() {
        let grant = consumer_create_grant();
        let create = |filter: &str| format!("$JS.API.CONSUMER.CREATE.term-main-iris.r1.{filter}");
        assert!(overlaps(&grant, &create(&stream_subject("iris"))));
        // The width the doc states: a literal wildcard as the last token.
        assert!(overlaps(&grant, &create("$SWARM.term.>")));
        for other in [
            crate::subagent_term::subject("iris", "scout"),
            crate::subagent_term::stream_subjects("iris"),
            crate::agent_state_subject("iris"),
        ] {
            assert!(!overlaps(&grant, &create(&other)), "{other}");
        }
        assert!(!overlaps(&grant, "$JS.API.CONSUMER.CREATE.term-main-iris"));
    }

    #[test]
    fn a_fresh_reader_replays_the_newest_rows() {
        assert_eq!(replay_start(1000, None), 1000 - REPLAY_ROWS + 1);
        assert_eq!(replay_start(5, None), 1);
        assert_eq!(replay_start(0, None), 1);
    }

    #[test]
    fn a_reconnecting_reader_resumes_after_its_last_row() {
        assert_eq!(replay_start(1000, Some(990)), 991);
        assert_eq!(replay_start(1000, Some(1000)), 1001);
    }

    /// An id past the stream's newest row, as after the stream was created
    /// again, replays like a fresh reader instead of waiting for that number.
    #[test]
    fn an_id_past_the_newest_row_replays_as_a_fresh_reader() {
        assert_eq!(replay_start(5, Some(99)), replay_start(5, None));
        assert_eq!(replay_start(1000, Some(1001)), replay_start(1000, None));
        assert_eq!(replay_start(0, Some(7)), 1);
    }

    /// The consumer starts at a stored sequence, not at the next new row, so
    /// the replay reaches the reader before the live tail, in stream order,
    /// through the one subscription.
    #[cfg(feature = "agent-term")]
    #[test]
    fn the_consumer_replays_before_it_follows() {
        let cfg = replay_consumer_config("iris", "r1".to_owned(), "_INBOX.x".to_owned(), 901);
        assert_eq!(
            cfg.deliver_policy,
            async_nats::jetstream::consumer::DeliverPolicy::ByStartSequence {
                start_sequence: 901
            }
        );
        assert_eq!(cfg.filter_subject, "$SWARM.term.iris");
        assert_eq!(cfg.name.as_deref(), Some("r1"));
        assert_eq!(
            cfg.ack_policy,
            async_nats::jetstream::consumer::AckPolicy::None
        );
        assert!(cfg.durable_name.is_none());
        // Without it a dropped consumer is never noticed.
        assert_eq!(cfg.idle_heartbeat, IDLE_HEARTBEAT);
        assert!(!cfg.idle_heartbeat.is_zero());
    }

    /// Every limit is read from the subagent stream's own config, so the two
    /// fail this test the moment they drift apart.
    #[cfg(all(feature = "agent-term", feature = "subagent-term"))]
    #[test]
    fn stream_config_matches_the_subagent_streams_limits() {
        let main = stream_config("iris", 8_388_608);
        let sub = crate::subagent_term::stream_config("iris", 8_388_608);
        assert_eq!(main.subjects, ["$SWARM.term.iris"]);
        assert_eq!(main.max_age, sub.max_age);
        assert_eq!(main.max_bytes, sub.max_bytes);
        assert_eq!(main.discard, sub.discard);
        assert_eq!(main.max_message_size, sub.max_message_size);
        assert_eq!(main.retention, sub.retention);
        assert_eq!(main.storage, sub.storage);
        assert_eq!(main.num_replicas, sub.num_replicas);
    }
}