hyperhive/swarm-notify-router/src/main.rs

//! The swarm notification router.
//!
//! Drains its durable consumer on the forge events stream
//! (`swarm_wire::forge_events`), runs every flow over each delivery
//! ([`flow`], [`route`]), and queues the todo ops the flows emit on each
//! recipient agent's subject of `ROUTER_TODOS` (`swarm_wire::router_todo`).
//! Each agent's harness applies them to its own todo list. The flows are the
//! built-in ones, changed by the `notify-flows` document ([`stored`]).
//!
//! A `request_reviewer` node is a request to swarm-controller, which holds the
//! forge token (`swarm_wire::router_sink`). It is sent once, after the event's
//! ops are queued and its snapshot is written, and the answer is only logged:
//! a controller that is down or slow must not hold up every agent's todos, and
//! the next push asks again.
//!
//! Per issue it keeps a snapshot in its state bucket ([`event::Snapshot`]):
//! the assignees it last saw and the agents it queued a todo for. An event is
//! acked once its ops are queued and the snapshot is written, so a crash
//! between the two redelivers it, and the stream drops the republished ops
//! by their message id.
//!
//! It connects to the queue as its own OIDC client, whose secret it reads
//! from the store under its own certificate, as `swarm-claude-login` does.
//! It holds no forge token and no webhook secret.

use std::collections::BTreeSet;
use std::path::PathBuf;
use std::time::{Duration, Instant};

use anyhow::{Context as _, Result, bail};
use async_nats::jetstream::{self, consumer::pull, kv, message::PublishMessage};
use futures_util::{StreamExt as _, TryStreamExt as _};
use swarm_queue_client::{ClientSecret, QueueConfig};
use swarm_secret_client::stored_secret::StoredSecret;
use swarm_wire::router_sink::{self, Outcome, RequestReviewer};
use swarm_wire::{forge_events, notify_router, router_todo};

mod event;
mod flow;
mod route;
mod settle;
mod stored;

/// The cert-auth role this service logs in to the store as. `swarm-bao.nix`
/// writes it, matching the leaf's subject.
const CERT_ROLE: &str = "swarm-notify-router";

/// The environment prefix `swarm-notify-router.nix` sets.
const PREFIX: &str = "SWARM_NOTIFY_ROUTER";

/// How many events one pull asks for, and how long it waits for them.
const BATCH: usize = 16;
const PULL_WAIT: Duration = Duration::from_secs(30);

/// How long an event may stay unacked before the stream redelivers it.
const EVENT_ACK_WAIT: Duration = Duration::from_mins(1);

/// How long the router waits after a failed pull or event before it pulls
/// again, when the failed event comes back first.
const RETRY_DELAY: Duration = Duration::from_secs(30);

/// How long the list of agents is reused before it is read again.
const ROSTER_TTL: Duration = Duration::from_mins(1);

#[tokio::main]
async fn main() -> Result<()> {
    hive_log::init();
    swarm_secret_client::install_crypto_provider();
    let var = |name: &str| {
        let full = format!("{PREFIX}_{name}");
        std::env::var(&full).with_context(|| format!("{full} is unset"))
    };
    let client_id = var("OIDC_CLIENT_ID")?;
    let secret =
        swarm_secret_client::read_with_startup_retry("reading the queue client secret", || {
            client_secret(&client_id)
        })
        .await?;
    let reread_id = client_id.clone();
    let secret = StoredSecret::new(client_id.clone(), secret, move || {
        let client_id = reread_id.clone();
        async move { client_secret(&client_id).await.map_err(Into::into) }
    });
    let queue = swarm_queue_client::connect(QueueConfig {
        url: var("NATS_URL")?,
        token_endpoint: var("OIDC_TOKEN_ENDPOINT")?,
        client_id,
        client_secret: ClientSecret::Stored(secret),
        ca_file: std::env::var(format!("{PREFIX}_OIDC_CA_FILE"))
            .ok()
            .map(PathBuf::from),
    })
    .await?;

    let mut router = Router::open(&queue).await?;
    tracing::info!(
        stream = forge_events::STREAM,
        "notify router: routing forge events"
    );
    loop {
        if let Err(e) = router.pull_once().await {
            tracing::warn!(
                error = format!("{e:#}"),
                "notify router: pulling events failed; retrying"
            );
            tokio::time::sleep(RETRY_DELAY).await;
        }
    }
}

struct Router {
    js: jetstream::Context,
    nats: async_nats::Client,
    events: jetstream::consumer::Consumer<pull::Config>,
    state: kv::Store,
    flows: stored::Flows,
    roster: Option<(Instant, BTreeSet<String>)>,
}

impl Router {
    async fn open(queue: &async_nats::Client) -> Result<Self> {
        let js = jetstream::new(queue.clone());
        let events = swarm_queue_client::forge_events::open_or_create(queue)
            .await
            .context("opening the forge events stream")?;
        swarm_queue_client::router_todo::open_or_create(queue)
            .await
            .context("opening the router todo stream")?;
        let events = events
            .create_consumer(pull::Config {
                durable_name: Some(forge_events::ROUTER_CONSUMER.to_owned()),
                description: Some("The notification router".to_owned()),
                filter_subject: forge_events::subject_wildcard(),
                ack_wait: EVENT_ACK_WAIT,
                ..Default::default()
            })
            .await
            .context("creating the router's consumer")?;
        let state = match js.get_key_value(notify_router::STATE_BUCKET).await {
            Ok(store) => store,
            Err(e) => {
                tracing::info!(reason = %e, "notify router: state bucket not available, creating it");
                js.create_key_value(kv::Config {
                    bucket: notify_router::STATE_BUCKET.to_owned(),
                    description: "Per-issue assignees and todo holders".to_owned(),
                    history: 1,
                    ..Default::default()
                })
                .await
                .context("creating the state bucket")?
            }
        };
        Ok(Self {
            js,
            nats: queue.clone(),
            events,
            state,
            flows: stored::Flows::builtins(),
            roster: None,
        })
    }

    /// Pull one batch and handle each event in stream order.
    async fn pull_once(&mut self) -> Result<()> {
        let batch = self
            .events
            .batch()
            .max_messages(BATCH)
            .expires(PULL_WAIT)
            .messages()
            .await
            .context("requesting a batch")?;
        let batch = batch.map(|m| m.map_err(|e| anyhow::anyhow!("{e}")));
        settle::batch(batch, async |m: &jetstream::Message| self.handle(m).await).await
    }

    async fn handle(&mut self, message: &jetstream::Message) -> Result<(), Skip> {
        let event_type = forge_events::event_type(&message.subject)
            .ok_or_else(|| Skip::Unusable("not a forge event subject".to_owned()))?
            .to_owned();
        let Some(delivery) = event::Delivery::parse(&message.payload) else {
            // A delivery that names no issue or PR has nothing to route.
            return Ok(());
        };
        let seq = message
            .info()
            .map(|i| i.stream_sequence)
            .map_err(|e| Skip::Unusable(format!("no stream sequence: {e}")))?;
        let key = delivery.state_key();
        let snapshot: event::Snapshot = match self.state.get(&key).await.map_err(retry)? {
            Some(bytes) => serde_json::from_slice(&bytes).unwrap_or_else(|e| {
                tracing::warn!(%key, error = %e, "notify router: unreadable snapshot; starting afresh");
                event::Snapshot::default()
            }),
            None => event::Snapshot::default(),
        };
        self.flows.refresh(&self.js).await;
        let roster = self.roster().await.map_err(Skip::Retry)?;
        let routed = route::route(&event_type, &delivery, &snapshot, &self.flows.active, |a| {
            roster.contains(a)
        });
        for action in &routed.actions {
            self.queue(seq, action).await.map_err(Skip::Retry)?;
        }
        if !routed.skipped.is_empty() {
            tracing::info!(%event_type, repo = %delivery.repo, number = delivery.number, skipped = ?routed.skipped, "notify router: recipients that are not agents got nothing");
        }
        if routed.next != snapshot {
            let bytes = serde_json::to_vec(&routed.next).expect("a snapshot encodes");
            self.state.put(&key, bytes.into()).await.map_err(retry)?;
        }
        for request in &routed.requests {
            self.request_review(request).await;
        }
        Ok(())
    }

    /// Ask swarm-controller for one review and log what it answered. One
    /// attempt: nothing here fails the event.
    async fn request_review(&self, request: &flow::ReviewRequest) {
        let wire = RequestReviewer {
            repo: request.repo.clone(),
            number: request.number,
            reviewer: request.reviewer.clone(),
        };
        let body = serde_json::to_vec(&wire).expect("a review request encodes");
        let sent = async_nats::Request::new()
            .payload(body.into())
            .timeout(Some(router_sink::REPLY_DEADLINE));
        let (flow, repo, number, reviewer) = (
            &request.flow,
            &request.repo,
            request.number,
            &request.reviewer,
        );
        match self
            .nats
            .send_request(router_sink::REQUEST_REVIEWER_SUBJECT, sent)
            .await
        {
            Ok(reply) => match Outcome::decode(&String::from_utf8_lossy(&reply.payload)) {
                Outcome::Requested => {
                    tracing::info!(%flow, %repo, number, %reviewer, "notify router: review requested");
                }
                Outcome::Skipped(why) => {
                    tracing::info!(%flow, %repo, number, %reviewer, %why, "notify router: review not requested");
                }
                Outcome::Error(why) => {
                    tracing::warn!(%flow, %repo, number, %reviewer, %why, "notify router: the controller could not request the review");
                }
            },
            Err(e) => {
                tracing::warn!(%flow, %repo, number, %reviewer, error = %e, "notify router: no answer to a review request; not retried");
            }
        }
    }

    /// Queue one op on its agent's subject and wait for the stream's ack.
    async fn queue(&self, seq: u64, action: &flow::Action) -> Result<()> {
        let subject = router_todo::subject(&action.agent);
        let body = serde_json::to_vec(&action.op).expect("an op encodes");
        let id = format!("{seq}:{}:{}:{}", action.flow, action.node, action.agent);
        self.js
            .send_publish(
                subject.clone(),
                PublishMessage::build().payload(body.into()).message_id(id),
            )
            .await
            .with_context(|| format!("publishing to {subject}"))?
            .await
            .with_context(|| format!("waiting for {subject}'s ack"))?;
        tracing::info!(agent = %action.agent, flow = %action.flow, key = %action.op.key, op = ?action.op.op, "notify router: todo op queued");
        Ok(())
    }

    /// The agents of this swarm: the agent names in the agent-status bucket's
    /// keys (`<hive>.<agent>`), read from the stream's subject list.
    async fn roster(&mut self) -> Result<BTreeSet<String>> {
        if let Some((at, agents)) = &self.roster
            && at.elapsed() < ROSTER_TTL
        {
            return Ok(agents.clone());
        }
        let bucket = swarm_wire::agent_status::BUCKET;
        let stream = self
            .js
            .get_stream(format!("KV_{bucket}"))
            .await
            .context("opening the agent-status bucket")?;
        let subjects: Vec<(String, usize)> = stream
            .info_with_subjects(format!("$KV.{bucket}.>"))
            .await
            .context("listing agent-status keys")?
            .try_collect()
            .await
            .context("listing agent-status keys")?;
        let prefix = format!("$KV.{bucket}.");
        let agents: BTreeSet<String> = subjects
            .into_iter()
            .filter_map(|(subject, _)| {
                let key = subject.strip_prefix(&prefix)?;
                swarm_wire::agent_status::split_key(key).map(|(_, agent)| agent.to_owned())
            })
            .collect();
        self.roster = Some((Instant::now(), agents.clone()));
        Ok(agents)
    }
}

/// Why an event was not handled.
enum Skip {
    /// It can never be handled; dropped for good.
    Unusable(String),
    /// The queue or the store failed; tried again.
    Retry(anyhow::Error),
}

fn retry(e: impl std::error::Error + Send + Sync + 'static) -> Skip {
    Skip::Retry(e.into())
}

async fn client_secret(client_id: &str) -> Result<String> {
    use swarm_secret_client::builder::{ClientSecret, client_secret_path};

    let path = client_secret_path(client_id)?;
    let store = swarm_secret_client::SecretStore::from_env(CERT_ROLE)
        .await
        .context("logging in to the secret store")?;
    let stored: Option<ClientSecret> = store
        .read_optional(&path)
        .await
        .with_context(|| format!("reading {path}"))?;
    match stored {
        Some(s) if !s.value.trim().is_empty() => Ok(s.value.trim().to_owned()),
        _ => bail!(
            "nothing at {path}: swarm-secret-publish, in authelia's container, has not \
             published this service's client secret"
        ),
    }
}