876 lines
30 KiB
Rust
876 lines
30 KiB
Rust
//! 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, Instant};
|
|
|
|
use hive_claude::{Attach, Config, Sink};
|
|
use swarm_secret_client::acp::ProviderKey;
|
|
use swarm_secret_client::review_classifier::{ApiKey, acp_provider_key_path, api_key_path};
|
|
use swarm_wire::agent_inference::Effort;
|
|
use swarm_wire::review_classifier::{CLASSIFICATION_DEADLINE, Reply, Row, Runtime, TryOut};
|
|
|
|
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. The ask is stopped after
|
|
/// `deadline`.
|
|
fn ask(
|
|
&self,
|
|
prompt: &Prompt<'_>,
|
|
key: &str,
|
|
run: &TryOut,
|
|
deadline: Duration,
|
|
) -> 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 prompt = prompt::classification_prompt(&shared, c);
|
|
let outcome = classify(
|
|
model,
|
|
&prompt,
|
|
&key,
|
|
run,
|
|
CLASSIFICATION_DEADLINE,
|
|
RETRY_FLOOR,
|
|
)
|
|
.await;
|
|
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 least time left in a classification's budget for which an answer that
|
|
/// did not validate is asked for again.
|
|
const RETRY_FLOOR: Duration = Duration::from_secs(20);
|
|
|
|
/// One classification's answer, from at most `budget` of model time in all. An
|
|
/// answer that does not validate is asked for once more, with a reminder to
|
|
/// reply with the JSON object alone, and only the time the first ask left of
|
|
/// `budget` to do it in; with less than `floor` left, or after a failed ask,
|
|
/// it is not asked again.
|
|
async fn classify(
|
|
model: &impl Model,
|
|
prompt: &Prompt<'_>,
|
|
key: &str,
|
|
run: &TryOut,
|
|
budget: Duration,
|
|
floor: Duration,
|
|
) -> Result<answer::Answer, String> {
|
|
let started = Instant::now();
|
|
let first = model.ask(prompt, key, run, budget).await?;
|
|
let unparsed = match answer::parse(&first) {
|
|
Ok(answer) => return Ok(answer),
|
|
Err(unparsed) => unparsed,
|
|
};
|
|
let left = budget.saturating_sub(started.elapsed());
|
|
if left < floor {
|
|
return Err(unparsed);
|
|
}
|
|
let again = model
|
|
.ask(&prompt.with_json_reminder(), key, run, left)
|
|
.await?;
|
|
answer::parse(&again)
|
|
}
|
|
|
|
/// 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,
|
|
deadline: Duration,
|
|
) -> Result<String, String> {
|
|
match run.runtime {
|
|
Runtime::ApiKey => self.claude.ask(prompt, key, run, deadline).await,
|
|
Runtime::Acp(_) => self.acp.ask(prompt, key, run, deadline).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,
|
|
deadline: Duration,
|
|
) -> 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, deadline)
|
|
.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,
|
|
run: &TryOut,
|
|
deadline: Duration,
|
|
) -> 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, &run.model, run.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(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", 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_wire::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,
|
|
_: Duration,
|
|
) -> 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_wire::review_classifier::Acp {
|
|
base_url: "https://llm.example/v1".to_owned(),
|
|
model: "qwen".to_owned(),
|
|
prices: swarm_wire::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:?}"
|
|
);
|
|
}
|
|
|
|
/// Records each prompt it is asked. A prompt naming `FLAKY` is answered
|
|
/// with prose until its reminder arrives; one naming `JUNK` never is.
|
|
#[derive(Default)]
|
|
struct Recording(Mutex<Vec<String>>);
|
|
|
|
impl Recording {
|
|
fn asked(&self, marker: &str) -> usize {
|
|
let prompts = self.0.lock().expect("lock");
|
|
prompts.iter().filter(|p| p.contains(marker)).count()
|
|
}
|
|
}
|
|
|
|
impl Model for Recording {
|
|
async fn ask(
|
|
&self,
|
|
prompt: &Prompt<'_>,
|
|
_: &str,
|
|
_: &TryOut,
|
|
_: Duration,
|
|
) -> Result<String, String> {
|
|
let text = prompt.text();
|
|
self.0.lock().expect("lock").push(text.clone());
|
|
let json = r#"{"verdict":"no-match","certainty":90,"reason":"fine"}"#;
|
|
if text.contains("ANSWER-FAIL") {
|
|
Err("claude failed: exit 1".to_owned())
|
|
} else if text.contains("FLAKY") && !text.contains("Reply with the JSON") {
|
|
Ok(format!("Looking at the diff, {json}"))
|
|
} else if text.contains("JUNK") {
|
|
Ok("not json".to_owned())
|
|
} else {
|
|
Ok(json.to_owned())
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn an_answer_that_is_not_json_is_asked_again_for_that_classification_only() {
|
|
let model = Recording::default();
|
|
let reply = answer(&Key(Ok(())), &model, &run(&["STEADY", "FLAKY", "OTHER"])).await;
|
|
let Reply::Done { rows, overall } = reply else {
|
|
panic!("all answered: {reply:?}");
|
|
};
|
|
assert_eq!(overall, Overall::Approve);
|
|
assert!(
|
|
rows.iter().all(|r| r.error.is_none() && r.passed),
|
|
"{rows:?}"
|
|
);
|
|
assert_eq!(model.asked("STEADY"), 1);
|
|
assert_eq!(model.asked("FLAKY"), 2);
|
|
assert_eq!(model.asked("OTHER"), 1);
|
|
assert_eq!(model.asked("Reply with the JSON"), 1);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn an_answer_that_is_still_not_json_after_one_retry_has_no_answer() {
|
|
let model = Recording::default();
|
|
let reply = answer(&Key(Ok(())), &model, &run(&["STEADY", "JUNK"])).await;
|
|
let Reply::Done { rows, .. } = reply else {
|
|
panic!("one classification answered: {reply:?}");
|
|
};
|
|
assert!(rows[0].error.is_none());
|
|
assert!(
|
|
rows[1]
|
|
.error
|
|
.as_deref()
|
|
.is_some_and(|e| e.contains("not the JSON"))
|
|
);
|
|
assert!(!rows[1].counted && !rows[1].passed);
|
|
assert_eq!(model.asked("JUNK"), 2);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn a_failed_ask_is_not_retried() {
|
|
let model = Recording::default();
|
|
let reply = answer(&Key(Ok(())), &model, &run(&["STEADY", "ANSWER-FAIL"])).await;
|
|
let Reply::Done { rows, .. } = reply else {
|
|
panic!("one classification answered: {reply:?}");
|
|
};
|
|
assert_eq!(rows[1].error.as_deref(), Some("claude failed: exit 1"));
|
|
assert_eq!(model.asked("ANSWER-FAIL"), 1);
|
|
}
|
|
|
|
/// Answers with prose after taking `first` to do so, then with JSON, and
|
|
/// records the deadline each ask was given.
|
|
struct Slow {
|
|
first: Duration,
|
|
deadlines: Mutex<Vec<Duration>>,
|
|
}
|
|
|
|
impl Slow {
|
|
fn new(first: Duration) -> Self {
|
|
Self {
|
|
first,
|
|
deadlines: Mutex::default(),
|
|
}
|
|
}
|
|
|
|
fn deadlines(&self) -> Vec<Duration> {
|
|
self.deadlines.lock().expect("lock").clone()
|
|
}
|
|
}
|
|
|
|
impl Model for Slow {
|
|
async fn ask(
|
|
&self,
|
|
_: &Prompt<'_>,
|
|
_: &str,
|
|
_: &TryOut,
|
|
deadline: Duration,
|
|
) -> Result<String, String> {
|
|
let asked = {
|
|
let mut deadlines = self.deadlines.lock().expect("lock");
|
|
deadlines.push(deadline);
|
|
deadlines.len()
|
|
};
|
|
if asked == 1 {
|
|
tokio::time::sleep(self.first).await;
|
|
return Ok("prose".to_owned());
|
|
}
|
|
Ok(r#"{"verdict":"no-match","certainty":90,"reason":"fine"}"#.to_owned())
|
|
}
|
|
}
|
|
|
|
const BUDGET: Duration = Duration::from_millis(600);
|
|
const FLOOR: Duration = Duration::from_millis(100);
|
|
|
|
async fn classify_slow(model: &Slow) -> Result<answer::Answer, String> {
|
|
let run = run(&["c"]);
|
|
let prompt = Prompt::of("q");
|
|
classify(model, &prompt, "k", &run, BUDGET, FLOOR).await
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn the_retry_gets_only_the_time_the_first_ask_left() {
|
|
let model = Slow::new(Duration::from_millis(200));
|
|
let answer = classify_slow(&model).await.expect("the retry answers");
|
|
assert_eq!(
|
|
answer.verdict,
|
|
swarm_wire::review_classifier::Verdict::NoMatch
|
|
);
|
|
let [first, second] = model.deadlines().try_into().expect("two asks");
|
|
assert_eq!(first, BUDGET);
|
|
assert!(
|
|
second.saturating_add(Duration::from_millis(200)) <= BUDGET,
|
|
"{second:?}"
|
|
);
|
|
assert!(second >= FLOOR, "{second:?}");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn no_retry_when_the_first_ask_left_less_than_the_floor() {
|
|
let model = Slow::new(Duration::from_millis(550));
|
|
let error = classify_slow(&model).await.expect_err("no second ask");
|
|
assert!(error.contains("not the JSON"), "{error}");
|
|
assert_eq!(model.deadlines(), [BUDGET]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn no_retry_when_the_first_ask_used_the_whole_budget() {
|
|
let model = Slow::new(BUDGET + Duration::from_millis(50));
|
|
classify_slow(&model).await.expect_err("no second ask");
|
|
assert_eq!(model.deadlines(), [BUDGET]);
|
|
}
|
|
|
|
#[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,
|
|
_: Duration,
|
|
) -> 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(&[]), CLASSIFICATION_DEADLINE)
|
|
.await
|
|
.as_deref(),
|
|
Ok("claude")
|
|
);
|
|
assert_eq!(
|
|
runtimes
|
|
.ask(
|
|
&Prompt::of("p"),
|
|
"k",
|
|
&acp_run(&[]),
|
|
CLASSIFICATION_DEADLINE
|
|
)
|
|
.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",
|
|
&TryOut {
|
|
effort: None,
|
|
..run(&[])
|
|
},
|
|
CLASSIFICATION_DEADLINE,
|
|
)
|
|
.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");
|
|
}
|
|
}
|