//! The review classifier's try-out service. //! //! Connects to the swarm queue as its own OIDC client, takes try-out requests //! from swarm-controller on //! [`swarm_queue_client::review_classifier::SUBJECT`], answers each one on its //! reply subject, and logs one line per run. It holds no forge credential: //! the controller reads the PR and sends it along. //! //! The store identity comes from the `BAO_*` environment. The queue client //! secret is read once at start and the API key once per run, both from the //! store under that identity. use std::path::PathBuf; use std::time::Duration; use anyhow::{Context, Result, bail}; use futures_util::StreamExt as _; use swarm_queue_client::review_classifier::{QUEUE_GROUP, Reply, SUBJECT, TryOut}; use swarm_queue_client::{ClientSecret, QueueConfig}; mod aggregate; mod answer; mod prompt; mod run; /// 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-review-classifier"; /// The environment prefix `swarm-review-classifier.nix` sets. const PREFIX: &str = "SWARM_REVIEW_CLASSIFIER"; /// Seconds to wait after each failed read of the queue client secret. A /// longer outage fails the start, and the unit's `Restart=` tries again. const RETRY_DELAYS_S: [u64; 6] = [1, 2, 4, 8, 16, 30]; #[tokio::main] async fn main() -> Result<()> { hive_log::init(); swarm_queue_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 claude = run::Claude { program: PathBuf::from(var("CLAUDE")?), }; let secret = client_secret_with_retry(&client_id).await?; let queue = swarm_queue_client::connect(QueueConfig { url: var("NATS_URL")?, token_endpoint: var("OIDC_TOKEN_ENDPOINT")?, client_id, client_secret: ClientSecret::Value(secret), ca_file: std::env::var(format!("{PREFIX}_OIDC_CA_FILE")) .ok() .map(PathBuf::from), }) .await?; let mut requests = queue .queue_subscribe(SUBJECT, QUEUE_GROUP.to_owned()) .await .with_context(|| format!("subscribing to {SUBJECT}"))?; tracing::info!(subject = SUBJECT, "review classifier: awaiting try-outs"); let claude = std::sync::Arc::new(claude); while let Some(msg) = requests.next().await { let Some(reply_to) = msg.reply.clone() else { tracing::warn!("a try-out arrived with no reply subject; dropping it"); continue; }; let queue = queue.clone(); let claude = std::sync::Arc::clone(&claude); tokio::spawn(async move { let reply = handle(&msg.payload, &run::Store, &*claude).await; match serde_json::to_vec(&reply) { Ok(body) => { if let Err(e) = queue.publish(reply_to, body.into()).await { tracing::warn!(error = %e, "publishing a try-out's reply failed"); } } Err(e) => tracing::warn!(error = %e, "encoding a try-out's reply failed"), } }); } bail!("the subscription to {SUBJECT} ended") } /// Decode one request, answer it and log the run. async fn handle(payload: &[u8], keys: &impl run::KeySource, model: &impl run::Model) -> Reply { let run: TryOut = match serde_json::from_slice(payload) { Ok(run) => run, Err(e) => { tracing::warn!(error = %e, "an undecodable try-out; answering with the error"); return Reply::Failed { error: format!("the classifier could not decode the request: {e}"), }; } }; let reply = run::answer(keys, model, &run).await; log_run(&run, &reply); reply } /// The one journal line a run leaves. fn log_run(run: &TryOut, reply: &Reply) { match reply { Reply::Done { rows, overall } => tracing::info!( operator = %run.operator, repo = %run.repo, pr = run.pr, head_sha = %run.head_sha, rows = %serde_json::to_string(rows).unwrap_or_default(), overall = ?overall, "review classification run" ), Reply::Failed { error } => tracing::warn!( operator = %run.operator, repo = %run.repo, pr = run.pr, head_sha = %run.head_sha, %error, "review classification run failed" ), } } /// The queue client secret, read from the store, retried through a store that /// is restarting alongside this service. async fn client_secret_with_retry(client_id: &str) -> Result { for delay in RETRY_DELAYS_S.map(Duration::from_secs) { match client_secret(client_id).await { Ok(secret) => return Ok(secret), Err(e) => { tracing::warn!(error = %format!("{e:#}"), retry_in = ?delay, "reading the queue client secret failed"); tokio::time::sleep(delay).await; } } } client_secret(client_id).await } async fn client_secret(client_id: &str) -> Result { use swarm_secret_client::review_classifier::{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}: swarm-secret-publish, on the host that runs authelia, has not \ published this service's client secret" ), } }