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

//! The per-agent schedule KV bucket: one key per agent, holding the scheduled
//! prompts that fire in that agent.
//!
//! **swarm-controller is the only writer**, and creates the bucket on its
//! first write. It keeps each prompt once, with its name and every agent it
//! goes to, and writes each agent's key as the prompts naming that agent,
//! without the name or the other agents. Each agent's harness reads its own
//! key and fires the schedules itself; the fire times are a pure function of
//! each schedule ([`crate::agent_schedules::Schedule::latest_slot`]), so
//! rewriting the list never moves them.
//!
//! **An agent reads with one `DIRECT.GET` on its own key, not a watch.** A KV
//! watch is an unnamed consumer, whose create subject carries no filter, so
//! granting it would grant every agent's key. The direct get carries the key
//! in its subject ([`crate::agent_schedules::direct_get_subject`]), which is
//! what lets the queue grant scope the read to one agent. It is also the
//! agent's whole `JetStream` grant: `read_own` builds the stream handle
//! without a `STREAM.INFO`.
//!
//! Keyed by the agent alone, like [`crate::agent_icon`], and for the same
//! reason: an agent can move between hives.

/// The KV bucket schedules are published into, keyed by agent name.
pub const BUCKET: &str = "agent-schedules";

/// The stream backing [`BUCKET`]. `create_key_value` prefixes `KV_`.
pub const STREAM: &str = "KV_agent-schedules";

/// The subject `agent`'s entry is stored under.
#[must_use]
pub fn key_subject(agent: &str) -> String {
    format!("$KV.{BUCKET}.{agent}")
}

/// The request subject that reads `agent`'s current entry, and so the one
/// subject the queue grants that agent in this bucket.
#[must_use]
pub fn direct_get_subject(agent: &str) -> String {
    format!("$JS.API.DIRECT.GET.{STREAM}.{}", key_subject(agent))
}

/// The prompts that fire in one agent: the value under `BUCKET`/`<agent>`.
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct AgentSchedules {
    #[serde(default)]
    pub schedules: Vec<Schedule>,
}

/// One scheduled prompt.
///
/// Its fire times are `anchor_unix + k * interval_seconds` for `k >= 0`, or
/// `anchor_unix` alone when there is no interval.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct Schedule {
    /// Minted by the controller, the same in every agent the prompt names,
    /// and never reused, so an agent's record of what it last fired for an
    /// id cannot apply to a different schedule.
    pub id: String,
    /// The prompt, delivered verbatim as the todo's summary.
    pub body: String,
    /// The first fire time, unix seconds.
    pub anchor_unix: i64,
    /// Seconds between fires; `None` is a one-shot.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub interval_seconds: Option<u64>,
    /// A paused schedule does not fire, and the slots that pass while it is
    /// paused are not fired when it resumes.
    #[serde(default)]
    pub paused: bool,
}

impl Schedule {
    /// The latest fire time at or before `now`, or `None` while the first is
    /// still ahead.
    #[must_use]
    pub fn latest_slot(&self, now: i64) -> Option<i64> {
        if now < self.anchor_unix {
            return None;
        }
        let Some(interval) = self
            .interval_seconds
            .filter(|i| *i > 0)
            .and_then(|i| i64::try_from(i).ok())
        else {
            return Some(self.anchor_unix);
        };
        let elapsed = now - self.anchor_unix;
        Some(self.anchor_unix + elapsed - elapsed % interval)
    }
}

/// Open the schedule bucket for writing, creating it if nothing has yet.
///
/// **Controller-side only.** `history: 1`: an agent wants its current list.
#[cfg(feature = "kv")]
pub async fn open_or_create(
    client: &async_nats::Client,
) -> Result<async_nats::jetstream::kv::Store, crate::Error> {
    let js = async_nats::jetstream::new(client.clone());
    match js.get_key_value(BUCKET).await {
        Ok(store) => Ok(store),
        Err(e) => {
            tracing::info!(
                bucket = BUCKET,
                reason = %e,
                "agent-schedule bucket not available, creating it"
            );
            js.create_key_value(async_nats::jetstream::kv::Config {
                bucket: BUCKET.to_owned(),
                description: "Scheduled prompts the swarm controller declares for each agent"
                    .to_owned(),
                history: 1,
                ..Default::default()
            })
            .await
            .map_err(|source| crate::Error::CreateBucket {
                bucket: BUCKET.to_owned(),
                source,
            })
        }
    }
}

/// What reading one agent's key found.
#[cfg(feature = "agent-schedules")]
#[derive(Debug)]
pub enum Read {
    /// The bucket answered; an absent or deleted key is an empty list.
    Current(AgentSchedules),
    /// Nothing answers for the bucket's stream: no schedule has been written
    /// yet, or `JetStream` is not serving it right now. Not an empty list,
    /// since the second case says nothing about the agent's schedules.
    NoBucket,
}

/// Read `agent`'s current list with the one subject its grant covers.
///
/// # Errors
///
/// A failed request or an undecodable value. A request the queue refused for
/// permissions reaches the client as a timeout, not as a refusal.
#[cfg(feature = "agent-schedules")]
pub async fn read_own(
    client: &async_nats::Client,
    agent: &str,
) -> Result<Read, Box<dyn std::error::Error + Send + Sync>> {
    use async_nats::jetstream::stream::DirectGetErrorKind;

    let js = async_nats::jetstream::new(client.clone());
    let stream = js.get_stream_no_info(STREAM).await?;
    let message = match stream.direct_get_last_for_subject(key_subject(agent)).await {
        Ok(message) => message,
        Err(e) if matches!(e.kind(), DirectGetErrorKind::NotFound) => {
            return Ok(Read::Current(AgentSchedules::default()));
        }
        Err(e) if is_no_responders(&e) => return Ok(Read::NoBucket),
        Err(e) => return Err(e.into()),
    };
    // A KV delete or purge is stored as an empty marker carrying this header.
    let removed = message
        .headers
        .get("KV-Operation")
        .is_some_and(|op| matches!(op.as_str(), "DEL" | "PURGE"));
    if removed {
        return Ok(Read::Current(AgentSchedules::default()));
    }
    Ok(Read::Current(serde_json::from_slice(&message.payload)?))
}

#[cfg(feature = "agent-schedules")]
fn is_no_responders(e: &async_nats::jetstream::stream::DirectGetError) -> bool {
    matches!(
        e.kind(),
        async_nats::jetstream::stream::DirectGetErrorKind::ErrorResponse(status, _)
            if status == async_nats::StatusCode::NO_RESPONDERS
    )
}

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

    fn every(anchor_unix: i64, interval: Option<u64>) -> Schedule {
        Schedule {
            id: "abc".to_owned(),
            body: "hi".to_owned(),
            anchor_unix,
            interval_seconds: interval,
            paused: false,
        }
    }

    /// The names the grant in `swarm-nats.nix` spells with `{agent}`.
    #[test]
    fn names_match_the_queue_grant() {
        assert_eq!(STREAM, format!("KV_{BUCKET}"));
        assert_eq!(key_subject("iris"), "$KV.agent-schedules.iris");
        assert_eq!(
            direct_get_subject("iris"),
            "$JS.API.DIRECT.GET.KV_agent-schedules.$KV.agent-schedules.iris"
        );
    }

    #[test]
    fn slots_sit_on_the_anchor_grid() {
        let s = every(1000, Some(60));
        assert_eq!(s.latest_slot(999), None);
        assert_eq!(s.latest_slot(1000), Some(1000));
        assert_eq!(s.latest_slot(1059), Some(1000));
        assert_eq!(s.latest_slot(1060), Some(1060));
        assert_eq!(s.latest_slot(1000 + 60 * 7 + 5), Some(1000 + 60 * 7));
    }

    #[test]
    fn a_one_shot_has_one_slot() {
        let s = every(1000, None);
        assert_eq!(s.latest_slot(999), None);
        assert_eq!(s.latest_slot(1000), Some(1000));
        assert_eq!(s.latest_slot(1_000_000), Some(1000));
    }

    /// Fields a newer controller adds must not cost an older agent its list,
    /// and a list written without the optional fields decodes to a recurring
    /// or one-shot schedule that is not paused.
    #[test]
    fn decoding_tolerates_missing_and_unknown_fields() {
        let doc = r#"{"schedules":[{"id":"a","body":"b","anchor_unix":5,"note":"x"}],"rev":2}"#;
        let decoded: AgentSchedules = serde_json::from_str(doc).expect("decodes");
        let mut want = every(5, None);
        "a".clone_into(&mut want.id);
        "b".clone_into(&mut want.body);
        assert_eq!(decoded.schedules, vec![want]);
    }
}