hyperhive/swarm-review-classifier/src/main.rs

//! 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<String> {
    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<String> {
    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<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, on the host that runs authelia, has not \
             published this service's client secret"
        ),
    }
}