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

//! Answering one try-out: read the API key, ask the model one classification
//! at a time, and turn the answers into the reply.
//!
//! Each classification is one `claude --print` with no session, no tools and
//! no MCP server, in a directory made for it and removed after it, so nothing
//! carries from one run to the next. The API key is read from the store for
//! every run and reaches only the child's environment, which keeps nothing of
//! the daemon's beyond [`crate::child_env`]'s allowlist.

use std::future::Future;
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;

use hive_claude::{Attach, Config, Sink};
use swarm_queue_client::agent_inference::Effort;
use swarm_queue_client::review_classifier::{CLASSIFICATION_DEADLINE, Reply, Row, TryOut};
use swarm_secret_client::review_classifier::{ApiKey, api_key_path};

use crate::{aggregate, answer, prompt};

/// How long `claude` may print nothing before it is stopped.
const IDLE_TIMEOUT: Duration = Duration::from_mins(2);

/// Where the run's API key comes from: the store in production, a fake in
/// the tests.
pub trait KeySource {
    /// The key, or why there is none, in words the operator acts on.
    fn api_key(&self) -> impl Future<Output = Result<String, String>> + Send;
}

/// Asks the model one classification: `claude` in production, a fake in the
/// tests.
pub trait Model {
    /// The final message of `model` at `effort` for `prompt`, or why there
    /// is none.
    fn ask(
        &self,
        prompt: &str,
        api_key: &str,
        resource_attributes: &str,
        model: &str,
        effort: Option<Effort>,
    ) -> impl Future<Output = Result<String, String>> + Send;
}

/// Answer `run`. A missing key, or every classification failing, is a failed
/// run; otherwise each classification is a row, with an error where it has no
/// answer.
pub async fn answer(keys: &impl KeySource, model: &impl Model, run: &TryOut) -> Reply {
    let api_key = match keys.api_key().await {
        Ok(key) => key,
        Err(error) => return Reply::Failed { error },
    };
    let shared = prompt::shared_context(run);
    let attributes = resource_attributes(std::env::var("OTEL_RESOURCE_ATTRIBUTES").ok(), run);
    let mut rows: Vec<Row> = Vec::with_capacity(run.classifications.len());
    for c in &run.classifications {
        let outcome = model
            .ask(
                &prompt::classification_prompt(&shared, c),
                &api_key,
                &attributes,
                &run.model,
                run.effort,
            )
            .await
            .and_then(|text| answer::parse(&text));
        rows.push(aggregate::row(c, outcome));
    }
    if let Some(first) = rows
        .first()
        .filter(|_| rows.iter().all(|r| r.error.is_some()))
    {
        return Reply::Failed {
            error: format!(
                "no classification got an answer; the first failed with: {}",
                first.error.as_deref().unwrap_or_default()
            ),
        };
    }
    let overall = aggregate::overall(&rows);
    Reply::Done { rows, overall }
}

/// The resource attributes a run's `claude` exports with: the container's
/// own, then the repo and the PR, so its cost reads back by both.
fn resource_attributes(base: Option<String>, run: &TryOut) -> String {
    let own = format!("repo={},pr={}", run.repo, run.pr);
    match base.filter(|b| !b.trim().is_empty()) {
        Some(base) => format!("{base},{own}"),
        None => own,
    }
}

/// The store, logged in to for each read: a run can come long after the last
/// one's token expired.
pub struct Store;

impl KeySource for Store {
    async fn api_key(&self) -> Result<String, String> {
        let path = api_key_path().map_err(|e| e.to_string())?;
        let store = swarm_secret_client::SecretStore::from_env(crate::CERT_ROLE)
            .await
            .map_err(|e| format!("logging in to the secret store to read the API key: {e}"))?;
        let stored: Option<ApiKey> = store
            .read_optional(&path)
            .await
            .map_err(|e| format!("reading secret/{path}: {e}"))?;
        match stored {
            Some(k) if !k.api_key.trim().is_empty() => Ok(k.api_key.trim().to_owned()),
            _ => Err(format!(
                "no API key at secret/{path}: write it in the secret store's UI, field `api_key`"
            )),
        }
    }
}

/// The `claude` binary the classifications run on, started through the
/// trampoline that clears its environment.
pub struct Claude {
    pub program: PathBuf,
    /// This binary, which execs [`Self::program`] when started as the
    /// trampoline.
    pub trampoline: PathBuf,
}

impl Model for Claude {
    async fn ask(
        &self,
        prompt: &str,
        api_key: &str,
        attributes: &str,
        model: &str,
        effort: Option<Effort>,
    ) -> Result<String, String> {
        let dir = scratch_dir();
        let outcome = self
            .ask_in(&dir, prompt, api_key, attributes, model, effort)
            .await;
        if let Err(e) = tokio::fs::remove_dir_all(&dir).await {
            tracing::warn!(dir = %dir.display(), error = %e, "removing a run's directory failed");
        }
        outcome
    }
}

impl Claude {
    async fn ask_in(
        &self,
        dir: &Path,
        prompt: &str,
        api_key: &str,
        attributes: &str,
        model: &str,
        effort: Option<Effort>,
    ) -> Result<String, String> {
        let system = dir.join("system.md");
        tokio::fs::create_dir_all(dir.join("config"))
            .await
            .and(tokio::fs::write(&system, prompt::SYSTEM).await)
            .map_err(|e| format!("preparing the run's directory: {e}"))?;
        let config = self.config(dir, system, api_key, attributes, model, effort);
        let running =
            hive_claude::Claude::spawn(&config, &Attach::OneOff).map_err(|e| e.to_string())?;
        let cancel = running.cancel_handle();
        let sink = ResultSink::default();
        let wait = running.wait(prompt, &sink);
        tokio::pin!(wait);
        match tokio::time::timeout(CLASSIFICATION_DEADLINE, &mut wait).await {
            Ok(Ok(())) => sink.take(),
            Ok(Err(e)) => Err(format!("claude failed: {e}")),
            Err(_) => {
                cancel.cancel(true);
                let _ = wait.await;
                Err(format!(
                    "no answer within {} seconds",
                    CLASSIFICATION_DEADLINE.as_secs()
                ))
            }
        }
    }

    /// One classification's invocation: the trampoline as the program, with
    /// the real `claude` and the run's own variables in its environment.
    fn config(
        &self,
        dir: &Path,
        system: PathBuf,
        api_key: &str,
        attributes: &str,
        model: &str,
        effort: Option<Effort>,
    ) -> Config {
        Config {
            model: Some(model.to_owned()),
            effort: effort.map(|e| e.as_str().to_owned()),
            cwd: Some(dir.to_owned()),
            system_prompt_file: Some(system),
            strict_mcp_config: true,
            // `--tools ""`: no built-in tool at all.
            tools: Some(String::new()),
            program: Some(self.trampoline.clone()),
            idle_timeout: Some(IDLE_TIMEOUT),
            claude_config_dir: Some(dir.join("config")),
            env: vec![
                (
                    crate::child_env::EXEC_CLAUDE.to_owned(),
                    self.program.display().to_string(),
                ),
                ("ANTHROPIC_API_KEY".to_owned(), api_key.to_owned()),
                ("HOME".to_owned(), dir.display().to_string()),
                ("OTEL_RESOURCE_ATTRIBUTES".to_owned(), attributes.to_owned()),
            ],
            ..Config::default()
        }
    }
}

/// A directory of its own for one classification.
fn scratch_dir() -> PathBuf {
    static NEXT: AtomicU64 = AtomicU64::new(0);
    std::env::temp_dir().join(format!(
        "review-classifier-{}-{}",
        std::process::id(),
        NEXT.fetch_add(1, Ordering::Relaxed)
    ))
}

/// Keeps the `result` event's text: the model's final message.
#[derive(Default)]
struct ResultSink {
    result: Mutex<Option<Result<String, String>>>,
}

impl ResultSink {
    fn take(&self) -> Result<String, String> {
        self.result
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .take()
            .unwrap_or_else(|| Err("claude ended without a result".to_owned()))
    }
}

impl Sink for ResultSink {
    fn on_event(&self, event: &serde_json::Value) {
        if event["type"] != "result" {
            return;
        }
        let text = event["result"].as_str().unwrap_or_default().to_owned();
        let outcome = if event["is_error"].as_bool().unwrap_or(false) {
            Err(format!("claude reported an error: {text}"))
        } else {
            Ok(text)
        };
        *self
            .result
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(outcome);
    }
}

#[cfg(test)]
mod tests {
    use swarm_queue_client::review_classifier::{Classification, Overall, Polarity};

    use super::*;

    struct Key(Result<&'static str, &'static str>);

    impl KeySource for Key {
        async fn api_key(&self) -> Result<String, String> {
            self.0.map(str::to_owned).map_err(str::to_owned)
        }
    }

    /// Answers each classification with the text its description names.
    struct Scripted;

    impl Model for Scripted {
        async fn ask(
            &self,
            prompt: &str,
            api_key: &str,
            attrs: &str,
            model: &str,
            effort: Option<Effort>,
        ) -> Result<String, String> {
            assert_eq!(api_key, "k");
            assert_eq!((model, effort), ("sonnet", Some(Effort::Low)));
            assert!(attrs.ends_with("repo=o/r,pr=7"), "{attrs}");
            if prompt.contains("ANSWER-FAIL") {
                Err("claude failed: exit 1".to_owned())
            } else if prompt.contains("ANSWER-MATCH") {
                Ok(r#"{"verdict":"match","certainty":90,"reason":"yes"}"#.to_owned())
            } else {
                Ok("not json".to_owned())
            }
        }
    }

    fn run(descriptions: &[&str]) -> TryOut {
        TryOut {
            operator: "alice".to_owned(),
            repo: "o/r".to_owned(),
            pr: 7,
            head_sha: "abc".to_owned(),
            commits: Vec::new(),
            diff: String::new(),
            classifications: descriptions
                .iter()
                .enumerate()
                .map(|(i, d)| Classification {
                    name: format!("c{i}"),
                    description: (*d).to_owned(),
                    polarity: Polarity::Negative,
                    threshold: 50,
                })
                .collect(),
            model: "sonnet".to_owned(),
            effort: Some(Effort::Low),
        }
    }

    #[tokio::test]
    async fn a_missing_key_fails_the_run_instead_of_answering() {
        let reply = answer(
            &Key(Err("no API key at x")),
            &Scripted,
            &run(&["ANSWER-MATCH"]),
        )
        .await;
        assert_eq!(
            reply,
            Reply::Failed {
                error: "no API key at x".to_owned()
            }
        );
    }

    #[tokio::test]
    async fn a_failed_or_malformed_answer_is_an_excluded_row_beside_the_others() {
        let reply = answer(
            &Key(Ok("k")),
            &Scripted,
            &run(&["ANSWER-MATCH", "ANSWER-FAIL", "ANSWER-JUNK"]),
        )
        .await;
        let Reply::Done { rows, overall } = reply else {
            panic!("one classification answered: {reply:?}");
        };
        assert_eq!(overall, Overall::RequestChanges);
        assert!(rows[0].counted && rows[0].error.is_none());
        assert_eq!(rows[1].error.as_deref(), Some("claude failed: exit 1"));
        assert!(
            rows[2]
                .error
                .as_deref()
                .is_some_and(|e| e.contains("not the JSON"))
        );
        assert!(rows[1..].iter().all(|r| !r.counted && !r.passed));
    }

    #[tokio::test]
    async fn no_answer_at_all_fails_the_run_instead_of_approving() {
        let reply = answer(
            &Key(Ok("k")),
            &Scripted,
            &run(&["ANSWER-FAIL", "ANSWER-JUNK"]),
        )
        .await;
        assert!(
            matches!(&reply, Reply::Failed { error } if error.contains("exit 1")),
            "{reply:?}"
        );
    }

    #[test]
    fn claude_is_started_through_the_trampoline() {
        let claude = Claude {
            program: PathBuf::from("/nix/store/x-claude-code/bin/claude"),
            trampoline: PathBuf::from(
                "/nix/store/y-swarm-review-classifier/bin/swarm-review-classifier",
            ),
        };
        let config = claude.config(
            Path::new("/tmp/run"),
            PathBuf::from("/tmp/run/system.md"),
            "k",
            "a=b",
            "sonnet",
            None,
        );
        assert_eq!(config.program, Some(claude.trampoline.clone()));
        assert!(config.env.contains(&(
            crate::child_env::EXEC_CLAUDE.to_owned(),
            "/nix/store/x-claude-code/bin/claude".to_owned()
        )));
    }

    #[test]
    fn the_runs_model_and_effort_are_claudes_flags() {
        let claude = Claude {
            program: PathBuf::from("claude"),
            trampoline: PathBuf::from("trampoline"),
        };
        let config = |model, effort| {
            claude.config(
                Path::new("/tmp/run"),
                PathBuf::from("/tmp/run/system.md"),
                "k",
                "a=b",
                model,
                effort,
            )
        };
        let set = config("opus", Some(Effort::Xhigh));
        assert_eq!(
            (set.model.as_deref(), set.effort.as_deref()),
            (Some("opus"), Some("xhigh"))
        );
        let unset = config("haiku", None);
        assert_eq!(
            (unset.model.as_deref(), unset.effort.as_deref()),
            (Some("haiku"), None),
            "no effort passes no --effort"
        );
    }

    #[test]
    fn the_repo_and_pr_extend_the_containers_own_attributes() {
        let r = run(&[]);
        assert_eq!(
            resource_attributes(Some("service.name=x".to_owned()), &r),
            "service.name=x,repo=o/r,pr=7"
        );
        assert_eq!(resource_attributes(None, &r), "repo=o/r,pr=7");
    }
}