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

//! Answering one try-out: read the key of the run's runtime, ask the model one
//! classification at a time, and turn the answers into the reply.
//!
//! On the API key runtime each classification is one `claude --print` with
//! no session, no tools and no MCP server; on the ACP runtime it is one
//! prompt on a new ACP session ([`crate::acp`]). Either way its files live
//! in a directory made for it and removed after it, so nothing carries from
//! one run to the next. The 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, Runtime, TryOut};
use swarm_secret_client::acp::ProviderKey;
use swarm_secret_client::review_classifier::{ApiKey, acp_provider_key_path, api_key_path};

use crate::prompt::Prompt;
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 key comes from: the store in production, a fake in the
/// tests.
pub trait KeySource {
    /// The key `runtime` runs on, or why there is none, in words the
    /// operator acts on.
    fn key(&self, runtime: &Runtime) -> impl Future<Output = Result<String, String>> + Send;
}

/// Asks the model one classification: [`Runtimes`] in production, a fake in
/// the tests.
pub trait Model {
    /// The final message for `prompt` on `run`'s runtime, model and effort,
    /// run with `key`, or why there is none.
    fn ask(
        &self,
        prompt: &Prompt<'_>,
        key: &str,
        run: &TryOut,
    ) -> 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 key = match keys.key(&run.runtime).await {
        Ok(key) => key,
        Err(error) => return Reply::Failed { error },
    };
    let shared = prompt::shared_context(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), &key, run)
            .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,
    }
}

/// Each classification on the model its run's runtime names.
pub struct Runtimes<C, A> {
    pub claude: C,
    pub acp: A,
}

impl<C: Model + Sync, A: Model + Sync> Model for Runtimes<C, A> {
    async fn ask(&self, prompt: &Prompt<'_>, key: &str, run: &TryOut) -> Result<String, String> {
        match run.runtime {
            Runtime::ApiKey => self.claude.ask(prompt, key, run).await,
            Runtime::Acp(_) => self.acp.ask(prompt, key, run).await,
        }
    }
}

/// 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 key(&self, runtime: &Runtime) -> Result<String, String> {
        match runtime {
            Runtime::ApiKey => {
                let path = api_key_path().map_err(|e| e.to_string())?;
                let key = read_key(&path, "API key", |k: ApiKey| k.api_key).await?;
                key.ok_or_else(|| {
                    format!(
                        "no API key at secret/{path}: write it in the secret store's UI, field \
                         `api_key`"
                    )
                })
            }
            Runtime::Acp(_) => {
                let path = acp_provider_key_path().map_err(|e| e.to_string())?;
                let key = read_key(&path, "ACP provider key", |k: ProviderKey| k.api_key).await?;
                key.ok_or_else(|| {
                    format!(
                        "no ACP provider key at secret/{path}: write it in the secret store's \
                         UI, field `api_key`"
                    )
                })
            }
        }
    }
}

/// The key `field` takes from the object at `path`, trimmed; `None` when
/// nothing or a blank is stored.
async fn read_key<T: serde::de::DeserializeOwned>(
    path: &str,
    what: &str,
    field: impl FnOnce(T) -> String,
) -> Result<Option<String>, 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 {what}: {e}"))?;
    let stored: Option<T> = store
        .read_optional(path)
        .await
        .map_err(|e| format!("reading secret/{path}: {e}"))?;
    Ok(stored
        .map(field)
        .map(|key| key.trim().to_owned())
        .filter(|key| !key.is_empty()))
}

/// 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: &Prompt<'_>,
        api_key: &str,
        run: &TryOut,
    ) -> Result<String, String> {
        let attributes = resource_attributes(std::env::var("OTEL_RESOURCE_ATTRIBUTES").ok(), run);
        let dir = scratch_dir();
        let outcome = self
            .ask_in(&dir, prompt, api_key, &attributes, &run.model, run.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: &Prompt<'_>,
        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 message = stream_json_message(prompt);
        let wait = running.wait(&message, &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()),
            // The prompt is [`stream_json_message`], not text.
            extra_args: vec!["--input-format".to_owned(), "stream-json".to_owned()],
            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()
        }
    }
}

/// `prompt` as the one user message `--input-format stream-json` reads: the
/// shared context and the question as separate text blocks, with a cache
/// breakpoint on the shared one.
///
/// `claude` sets a breakpoint only on the last block, after the question, so
/// without this one no cache entry ends where the run's classifications stop
/// agreeing, and none of them reads the shared context from the cache. With
/// the three `claude` sets (two in the system prompt, one on the question)
/// this is four, the most one request may carry.
fn stream_json_message(prompt: &Prompt<'_>) -> String {
    let message = serde_json::json!({
        "type": "user",
        "message": {
            "role": "user",
            "content": [
                {"type": "text", "text": prompt.shared, "cache_control": {"type": "ephemeral"}},
                {"type": "text", "text": prompt.question},
            ],
        },
    });
    format!("{message}\n")
}

/// A directory of its own for one classification.
pub 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::*;

    /// The API key runtime's key is `k` and the ACP runtime's `p`, or both
    /// fail with the same error.
    struct Key(Result<(), &'static str>);

    impl KeySource for Key {
        async fn key(&self, runtime: &Runtime) -> Result<String, String> {
            self.0.map_err(str::to_owned)?;
            Ok(match runtime {
                Runtime::ApiKey => "k",
                Runtime::Acp(_) => "p",
            }
            .to_owned())
        }
    }

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

    impl Model for Scripted {
        async fn ask(
            &self,
            prompt: &Prompt<'_>,
            key: &str,
            run: &TryOut,
        ) -> Result<String, String> {
            let prompt = prompt.text();
            let wanted = match run.runtime {
                Runtime::ApiKey => "k",
                Runtime::Acp(_) => "p",
            };
            assert_eq!(key, wanted, "the key of the run's runtime");
            assert_eq!(
                (run.model.as_str(), run.effort),
                ("sonnet", Some(Effort::Low))
            );
            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),
            runtime: Runtime::ApiKey,
        }
    }

    fn acp_run(descriptions: &[&str]) -> TryOut {
        TryOut {
            runtime: Runtime::Acp(swarm_queue_client::review_classifier::Acp {
                base_url: "https://llm.example/v1".to_owned(),
                model: "qwen".to_owned(),
                prices: swarm_queue_client::review_classifier::Prices {
                    input: 1.0,
                    output: 1.0,
                    cache_read: None,
                    cache_write: None,
                },
            }),
            ..run(descriptions)
        }
    }

    #[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(())),
            &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(())),
            &Scripted,
            &run(&["ANSWER-FAIL", "ANSWER-JUNK"]),
        )
        .await;
        assert!(
            matches!(&reply, Reply::Failed { error } if error.contains("exit 1")),
            "{reply:?}"
        );
    }

    #[tokio::test]
    async fn an_acp_run_is_asked_with_the_acp_key() {
        let reply = answer(&Key(Ok(())), &Scripted, &acp_run(&["ANSWER-MATCH"])).await;
        assert!(matches!(reply, Reply::Done { .. }), "{reply:?}");
    }

    /// Answers with its own name.
    struct Named(&'static str);

    impl Model for Named {
        async fn ask(&self, _: &Prompt<'_>, _: &str, _: &TryOut) -> Result<String, String> {
            Ok(self.0.to_owned())
        }
    }

    #[tokio::test]
    async fn each_run_goes_to_the_model_its_runtime_names() {
        let runtimes = Runtimes {
            claude: Named("claude"),
            acp: Named("acp"),
        };
        assert_eq!(
            runtimes
                .ask(&Prompt::of("p"), "k", &run(&[]))
                .await
                .as_deref(),
            Ok("claude")
        );
        assert_eq!(
            runtimes
                .ask(&Prompt::of("p"), "k", &acp_run(&[]))
                .await
                .as_deref(),
            Ok("acp")
        );
    }

    #[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");
    }

    /// The two content blocks of `line`, a stream-json user message.
    fn blocks(line: &str) -> Vec<serde_json::Value> {
        let message: serde_json::Value = serde_json::from_str(line).expect("one JSON message");
        assert_eq!(
            (&message["type"], &message["message"]["role"]),
            (&"user".into(), &"user".into())
        );
        let content = message["message"]["content"]
            .as_array()
            .expect("content blocks");
        assert_eq!(content.len(), 2, "{message}");
        content.clone()
    }

    #[test]
    fn the_shared_context_is_its_own_block_with_the_cache_breakpoint() {
        let prompt = Prompt {
            shared: "# Pull request o/r#7\n\n## Diff\n",
            question: "## The question\n\nIs it?\n".to_owned(),
        };
        let line = stream_json_message(&prompt);
        assert!(
            line.ends_with('\n') && line.trim_end().lines().count() == 1,
            "{line}"
        );
        let [shared, question] = blocks(&line).try_into().expect("two blocks");
        assert_eq!(shared["text"], prompt.shared);
        assert_eq!(
            shared["cache_control"],
            serde_json::json!({"type": "ephemeral"})
        );
        assert_eq!(question["text"], prompt.question.as_str());
        assert!(question.get("cache_control").is_none(), "{question}");
    }

    /// Stands in for `claude`: keeps its arguments and stdin beside itself,
    /// and answers.
    const CLAUDE: &str = r#"#!/bin/sh
printf '%s\n' "$@" > "$0.args"
cat > "$0.stdin"
printf '{"type":"result","subtype":"success","is_error":false,"result":"ok"}\n'
"#;

    #[tokio::test]
    async fn claude_reads_the_shared_context_and_the_question_as_separate_blocks() {
        use std::os::unix::fs::PermissionsExt as _;

        let tmp = tempfile::tempdir().expect("a dir");
        let script = tmp.path().join("claude.sh");
        std::fs::write(&script, CLAUDE).expect("written");
        std::fs::set_permissions(&script, std::fs::Permissions::from_mode(0o755)).expect("chmod");
        let claude = Claude {
            program: PathBuf::from("claude"),
            trampoline: script.clone(),
        };
        let prompt = Prompt {
            shared: "# Pull request o/r#7\n",
            question: "## The question\n".to_owned(),
        };
        let answer = claude
            .ask_in(&tmp.path().join("run"), &prompt, "k", "a=b", "sonnet", None)
            .await;
        assert_eq!(answer.as_deref(), Ok("ok"));
        let args = std::fs::read_to_string(tmp.path().join("claude.sh.args")).expect("args");
        assert!(args.contains("--input-format\nstream-json\n"), "{args}");
        let stdin = std::fs::read_to_string(tmp.path().join("claude.sh.stdin")).expect("stdin");
        let [shared, question] = blocks(&stdin).try_into().expect("two blocks");
        assert_eq!(
            (&shared["text"], &question["text"]),
            (&prompt.shared.into(), &prompt.question.as_str().into())
        );
        assert_eq!(shared["cache_control"]["type"], "ephemeral");
    }
}