//! 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 and //! recorded: //! a controller that is down or slow must not hold up every agent's todos, and //! the next push asks again. //! //! Each delivery, made or failed, is a row of the deliveries log ([`log`]), and //! the router answers swarm-controller's requests for the newest rows and for a //! dry run of a flows document ([`api`], [`dry_run`]). //! //! 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_api::Kind; use swarm_wire::router_sink::{self, Outcome, RequestReviewer}; use swarm_wire::{forge_events, notify_router, router_todo}; mod api; mod dry_run; mod event; mod flow; mod log; mod route; mod settle; mod stored; #[cfg(test)] mod test_nats; /// 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?; tokio::spawn(api::serve(queue.clone())); 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, state: kv::Store, flows: stored::Flows, roster: Option<(Instant, BTreeSet)>, } impl Router { async fn open(queue: &async_nats::Client) -> Result { 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")?; log::open_or_create(&js) .await .context("opening the deliveries log 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) }); let flows: BTreeSet<&str> = routed .actions .iter() .map(|a| a.flow.as_str()) .chain(routed.requests.iter().map(|r| r.flow.as_str())) .collect(); let delivery_id = message .headers .as_ref() .and_then(|h| h.get(async_nats::header::NATS_MESSAGE_ID)) .map(async_nats::HeaderValue::as_str); log::delivery(&event_type, &delivery, delivery_id, flows.len()); let cause = log::Cause { event: event_type.clone(), repo: delivery.repo.clone(), number: u64::try_from(delivery.number).unwrap_or_default(), }; for action in &routed.actions { self.queue(seq, &cause, 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(&cause, request).await; } Ok(()) } /// Ask swarm-controller for one review, log what it answered and record /// the delivery. One attempt: nothing here fails the event. async fn request_review(&self, cause: &log::Cause, 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, reviewer) = (&request.flow, &request.reviewer); let result = 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 => Ok("requested".to_owned()), Outcome::Skipped(why) => Ok(format!("skipped: {why}")), Outcome::Error(why) => Err(why), }, Err(e) => Err(format!("no answer: {e}")), }; log::sink(cause, flow, Kind::RequestReviewer, reviewer, &result); let row = cause.row(flow, &request.node, Kind::RequestReviewer, reviewer, result); log::append(&self.js, &row).await; } /// Queue one op on its agent's subject, wait for the stream's ack and /// record the delivery, failed or not. async fn queue(&self, seq: u64, cause: &log::Cause, action: &flow::Action) -> Result<()> { let sent = self.publish(seq, action).await; let result = match &sent { Ok(()) => Ok(format!("{} {}", op_name(action.op.op), action.op.key)), Err(e) => Err(format!("{e:#}")), }; log::sink(cause, &action.flow, Kind::Todo, &action.agent, &result); let row = cause.row( &action.flow, &action.node, Kind::Todo, &action.agent, result, ); log::append(&self.js, &row).await; sent } async fn publish(&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"))?; Ok(()) } /// The agents of this swarm: the agent names in the agent-status bucket's /// keys (`.`), read from the stream's subject list. async fn roster(&mut self) -> Result> { 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 = 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()) } /// The word the log uses for an op. const fn op_name(op: router_todo::Op) -> &'static str { match op { router_todo::Op::Upsert => "upsert", router_todo::Op::Clear => "clear", } } async fn client_secret(client_id: &str) -> Result { 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 = 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}: authelia's mint, in its container, has not \ stored this service's client secret" ), } } #[cfg(test)] mod tests { use swarm_wire::notify_flows::builtins; use swarm_wire::router_api::Row; use super::*; use crate::test_nats::TestNats; pub(super) fn cause() -> log::Cause { log::Cause { event: "issue_assign".to_owned(), repo: "o/r".to_owned(), number: 9, } } pub(super) fn action() -> flow::Action { let flow = builtins().remove(0); let (agent, op) = ( "atlas".to_owned(), router_todo::RouterTodoOp { op: router_todo::Op::Upsert, key: "forge:o/r#9:assignment".to_owned(), prefix: false, summary: "s".to_owned(), source: "router:assigned".to_owned(), }, ); flow::Action { flow: flow.id, node: "sink".to_owned(), agent, op, } } pub(super) fn review() -> flow::ReviewRequest { flow::ReviewRequest { flow: "head-classifier".to_owned(), node: "request".to_owned(), repo: "o/r".to_owned(), number: 9, reviewer: "review-classifier".to_owned(), } } async fn rows(router: &Router) -> Vec { log::recent(&router.js, 10).await.expect("the log") } /// A todo op that was queued is a row that went through; one the stream /// would not take is a row with the error, and the delivery still fails. #[tokio::test] async fn a_queued_todo_is_a_row_and_a_failed_one_is_a_row_with_the_error() { let nats = TestNats::start(); let client = nats.client().await; let router = Router::open(&client).await.expect("opens"); router.queue(5, &cause(), &action()).await.expect("queued"); let [row] = &rows(&router).await[..] else { panic!("one row") }; assert!(row.ok, "{row:?}"); assert_eq!( (row.flow.as_str(), row.node.as_str(), row.target.as_str()), ("assigned", "sink", "atlas") ); assert_eq!( (row.event.as_str(), row.repo.as_str(), row.number), ("issue_assign", "o/r", 9) ); assert_eq!(row.kind, Kind::Todo); assert_eq!(row.detail, "upsert forge:o/r#9:assignment"); router .js .delete_stream(router_todo::STREAM) .await .expect("deleted"); router .queue(6, &cause(), &action()) .await .expect_err("no stream takes it"); let rows = rows(&router).await; assert_eq!(rows.len(), 2); assert!(!rows[0].ok && !rows[0].detail.is_empty(), "{:?}", rows[0]); } /// Each answer the controller can give is a row: requested and skipped /// went through, an error did not, and no answer at all did not either. #[tokio::test] async fn a_review_request_is_a_row_whatever_the_controller_answers() { let nats = TestNats::start(); let client = nats.client().await; let router = Router::open(&client).await.expect("opens"); // Nobody answers. router.request_review(&cause(), &review()).await; let first = rows(&router).await.remove(0); assert!( !first.ok && first.detail.starts_with("no answer"), "{first:?}" ); assert_eq!(first.kind, Kind::RequestReviewer); assert_eq!(first.target, "review-classifier"); let mut hears = client .subscribe(router_sink::REQUEST_REVIEWER_SUBJECT) .await .expect("subscribed"); let answering = client.clone(); tokio::spawn(async move { for line in [ "requested", "skipped:a review stands", "error:the forge said no", ] { let m = hears.next().await.expect("a request"); answering .publish(m.reply.expect("a reply subject"), line.into()) .await .expect("answered"); } }); for _ in 0..3 { router.request_review(&cause(), &review()).await; } let got: Vec<(bool, String)> = rows(&router) .await .into_iter() .take(3) .rev() .map(|r| (r.ok, r.detail)) .collect(); assert_eq!( got, [ (true, "requested".to_owned()), (true, "skipped: a review stands".to_owned()), (false, "the forge said no".to_owned()), ] ); } } #[cfg(test)] mod log_tests { use std::collections::BTreeMap; use std::sync::{Arc, Mutex}; use tracing_subscriber::layer::SubscriberExt as _; use super::tests::{action, cause, review}; use super::*; use crate::test_nats::TestNats; type Event = (tracing::Level, BTreeMap); /// The level and fields of every event, fields as their debug text. #[derive(Clone, Default)] struct Events(Arc>>); impl tracing_subscriber::Layer for Events { fn on_event( &self, event: &tracing::Event<'_>, _: tracing_subscriber::layer::Context<'_, S>, ) { let mut fields = BTreeMap::new(); event.record( &mut |field: &tracing::field::Field, value: &dyn std::fmt::Debug| { fields.insert(field.name().to_owned(), format!("{value:?}")); }, ); self.0 .lock() .expect("unpoisoned") .push((*event.metadata().level(), fields)); } } impl Events { /// The events whose message is `message`, with their levels. fn named(&self, message: &str) -> Vec { let want = format!("notify router: {message}"); self.0 .lock() .expect("unpoisoned") .iter() .filter(|(_, f)| f.get("message").is_some_and(|m| *m == want)) .cloned() .collect() } } fn capture() -> (Events, tracing::subscriber::DefaultGuard) { let events = Events::default(); let guard = tracing::subscriber::set_default(tracing_subscriber::registry().with(events.clone())); (events, guard) } fn body(action: &str) -> String { format!( r#"{{"action":"{action}","issue":{{"number":9,"title":"t","html_url":"u","assignees":[{{"login":"atlas"}}]}},"repository":{{"full_name":"o/r","id":1}},"sender":{{"login":"mara"}}}}"# ) } /// Publish one delivery and hand it to the router, as the pull would. async fn deliver(router: &mut Router, event_type: &str, id: &str, body: String) { let ack = router .js .send_publish( forge_events::subject(event_type), PublishMessage::build().payload(body.into()).message_id(id), ) .await .expect("sent") .await .expect("acked"); let mut batch = router .events .batch() .max_messages(1) .expires(Duration::from_secs(5)) .messages() .await .expect("batch"); let message = batch.next().await.expect("a message").expect("pulled"); assert_eq!(message.info().expect("info").stream_sequence, ack.sequence); assert!(router.handle(&message).await.is_ok()); } /// A delivery a flow wants is one info line with its type, action, repo, /// number, sender, delivery id and flow count; one nothing wants is a /// debug line, and neither carries the body. #[tokio::test] async fn a_delivery_is_logged_at_info_when_a_flow_matched_and_at_debug_when_none() { let nats = TestNats::start(); let client = nats.client().await; let js = jetstream::new(client.clone()); let bucket = swarm_wire::agent_status::BUCKET; js.create_key_value(kv::Config { bucket: bucket.to_owned(), ..Default::default() }) .await .expect("bucket") .put("hive.atlas", "{}".into()) .await .expect("an agent"); let mut router = Router::open(&client).await.expect("opens"); let (events, _guard) = capture(); // Control: nothing is logged before a delivery. assert!(events.named("forge event routed").is_empty()); deliver(&mut router, "issue_assign", "d-1", body("assigned")).await; let info = events.named("forge event routed"); let [(level, f)] = info.as_slice() else { panic!("one line, got {info:?}"); }; assert_eq!(*level, tracing::Level::INFO); let got = |k: &str| f.get(k).map(String::as_str); assert_eq!(got("event_type"), Some("\"issue_assign\"")); assert_eq!(got("action"), Some("assigned")); assert_eq!(got("repo"), Some("o/r")); assert_eq!(got("number"), Some("9")); assert_eq!(got("sender"), Some("mara")); assert_eq!(got("delivery_id"), Some("\"d-1\"")); assert_eq!(got("flows"), Some("1")); assert!( !format!("{f:?}").contains("html_url"), "the body is not logged" ); assert!(events.named("forge event matched no flow").is_empty()); deliver(&mut router, "issue_comment", "d-2", body("created")).await; let none = events.named("forge event matched no flow"); let [(level, f)] = none.as_slice() else { panic!("one line, got {none:?}"); }; assert_eq!(*level, tracing::Level::DEBUG); assert_eq!(f.get("flows").map(String::as_str), Some("0")); assert_eq!(events.named("forge event routed").len(), 1); } /// A sink that went through is an info line and one that failed is a warn /// carrying the error text, for a review request and for a todo. #[tokio::test] async fn a_failing_sink_is_a_warn_with_the_error_and_a_working_one_is_info() { let nats = TestNats::start(); let client = nats.client().await; let router = Router::open(&client).await.expect("opens"); let (events, _guard) = capture(); let mut hears = client .subscribe(router_sink::REQUEST_REVIEWER_SUBJECT) .await .expect("subscribed"); let answering = client.clone(); tokio::spawn(async move { for line in ["requested", "error:Doer can't choose reviewer"] { let m = hears.next().await.expect("a request"); answering .publish(m.reply.expect("a reply subject"), line.into()) .await .expect("answered"); } }); router.request_review(&cause(), &review()).await; let ran = events.named("sink ran"); let [(level, f)] = ran.as_slice() else { panic!("one line, got {ran:?}"); }; assert_eq!(*level, tracing::Level::INFO); assert_eq!(f["flow"], "head-classifier"); assert_eq!((f["repo"].as_str(), f["number"].as_str()), ("o/r", "9")); assert_eq!(f["kind"], "RequestReviewer"); assert_eq!(f["target"], "review-classifier"); assert_eq!(f["outcome"], "requested"); assert!(events.named("sink failed").is_empty()); router.request_review(&cause(), &review()).await; let failed = events.named("sink failed"); let [(level, f)] = failed.as_slice() else { panic!("one line, got {failed:?}"); }; assert_eq!(*level, tracing::Level::WARN); assert_eq!(f["target"], "review-classifier"); assert_eq!(f["error"], "Doer can't choose reviewer"); assert_eq!((f["repo"].as_str(), f["number"].as_str()), ("o/r", "9")); router.queue(5, &cause(), &action()).await.expect("queued"); assert_eq!(events.named("sink ran").len(), 2); router .js .delete_stream(router_todo::STREAM) .await .expect("deleted"); router .queue(6, &cause(), &action()) .await .expect_err("no stream takes it"); let failed = events.named("sink failed"); assert_eq!(failed.len(), 2); assert_eq!(failed[1].0, tracing::Level::WARN); assert_eq!(failed[1].1["kind"], "Todo"); assert_eq!(failed[1].1["repo"], "o/r"); assert_eq!(failed[1].1["number"], "9"); } }