vinzenz/hyperhive
vinzenz/hyperhive
Watch
0
0
Fork
You've already forked hyperhive
 0
Code Activity

review-classifier: ask a check once more when its answer isn't valid JSON

Browse source 
An answer that fails to validate is asked for a second time, for that
classification only, with a reminder to reply with the JSON object alone.
The second ask comes out of the classification's CLASSIFICATION_DEADLINE:
it gets what the first ask left, and is skipped when less than 20 seconds
remain, so RUN_DEADLINE's bound still holds. A failed ask (runtime error,
timeout) is not repeated. If the second answer does not validate either,
the row has no answer, as before.

Model::ask takes the deadline for the ask; the ACP runtime no longer holds
its own.

The parser also takes the first fenced block of a reply that has prose
around the fence; prose around a bare object is still rejected.
This commit is contained in:
atlas 2026-10-16 15:32:07 +02:00 • committed by mara
commit ce7b3e055f
8 changed files with 345 additions and 66 deletions
Download patch file Download diff file Expand all files Collapse all files

13
docs/swarm/review-classifier.md
View file

@ -338,9 +338,16 @@ An answer:
``` ```
The verdict field holds `match`, `no-match` or `not-applicable`, and the The verdict field holds `match`, `no-match` or `not-applicable`, and the
certainty field a whole number. A runtime error, a timeout, or an answer that certainty field a whole number. The classifier reads the object from the
isn't that JSON object leaves the classification without a verdict: it's listed reply, or from the first code fence in it. When an answer isn't that object,
with its error, and doesn't count. A run with no key for its runtime, or in which no the classifier asks that classification once more, with a reminder to reply
with the JSON object alone; the other classifications aren't asked again. The
second ask runs on what the first left of the classification's three minutes,
so a retry never lengthens the run, and it doesn't happen when the first ask
leaves less than 20 seconds. The classifier also doesn't ask again after a
runtime error or a timeout. A runtime error, a timeout, or an answer that
still isn't that JSON object after the second ask leaves the classification without a verdict: it's listed with its error, and
doesn't count. A run with no key for its runtime, or in which no
classification got an answer, fails as a whole, and shows its error instead of a classification got an answer, fails as a whole, and shows its error instead of a
verdict. verdict.

4
swarm-queue-client/src/review_classifier.rs
View file

@ -23,8 +23,8 @@ pub const QUEUE_GROUP: &str = "review-classifier";
/// The most classifications one run carries. /// The most classifications one run carries.
pub const MAX_CLASSIFICATIONS: usize = 10; pub const MAX_CLASSIFICATIONS: usize = 10;
/// The longest one classification's model run may take before it is stopped /// The longest one classification may take, its first ask and any second ask
/// and its row carries a timeout. /// together, before it is stopped and its row carries a timeout.
pub const CLASSIFICATION_DEADLINE: Duration = Duration::from_mins(3); pub const CLASSIFICATION_DEADLINE: Duration = Duration::from_mins(3);
/// The longest a whole run may take: every classification at its deadline, /// The longest a whole run may take: every classification at its deadline,

2
swarm-review-classifier/README.md
View file

@ -22,6 +22,8 @@ forge and holds no forge credential.
- An answer is a JSON object with a verdict, a certainty and a reason; the - An answer is a JSON object with a verdict, a certainty and a reason; the
prompt asks for a reason under 350 characters, and the classifier truncates one over the 500 cap to 499 plus `…`. A failed run, a prompt asks for a reason under 350 characters, and the classifier truncates one over the 500 cap to 499 plus `…`. A failed run, a
timeout or an answer that doesn't validate leaves that row without a verdict, and it doesn't count. timeout or an answer that doesn't validate leaves that row without a verdict, and it doesn't count.
An answer that doesn't validate is asked for once more first, with a
reminder to reply with the JSON object alone, within the same three minutes.
- A run with no API key, or with no classification answered, fails as a - A run with no API key, or with no classification answered, fails as a
whole: the caller sees an error, never a verdict. whole: the caller sees an error, never a verdict.

69
swarm-review-classifier/src/acp.rs
View file

@ -65,8 +65,6 @@ pub struct Acp<S> {
/// This binary, which execs [`Self::program`] when started as the /// This binary, which execs [`Self::program`] when started as the
/// trampoline. /// trampoline.
pub trampoline: PathBuf, pub trampoline: PathBuf,
/// How long one classification may take before it is cancelled.
pub deadline: Duration,
/// [`CANCEL_WAIT`] outside the tests. /// [`CANCEL_WAIT`] outside the tests.
pub cancel_wait: Duration, pub cancel_wait: Duration,
pub cost: S, pub cost: S,
@ -83,12 +81,15 @@ impl<S: CostSink> crate::run::Model for Acp<S> {
prompt: &crate::prompt::Prompt<'_>, prompt: &crate::prompt::Prompt<'_>,
key: &str, key: &str,
run: &TryOut, run: &TryOut,
deadline: Duration,
) -> Result<String, String> { ) -> Result<String, String> {
let Runtime::Acp(provider) = &run.runtime else { let Runtime::Acp(provider) = &run.runtime else {
return Err("not an ACP run".to_owned()); return Err("not an ACP run".to_owned());
}; };
let dir = crate::run::scratch_dir(); let dir = crate::run::scratch_dir();
let outcome = self.ask_in(&dir, &prompt.text(), key, run, provider).await; let outcome = self
.ask_in(&dir, &prompt.text(), key, run, provider, deadline)
.await;
if let Err(e) = tokio::fs::remove_dir_all(&dir).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"); tracing::warn!(dir = %dir.display(), error = %e, "removing a run's directory failed");
} }
@ -104,6 +105,7 @@ impl<S: CostSink> Acp<S> {
key: &str, key: &str,
run: &TryOut, run: &TryOut,
provider: &Provider, provider: &Provider,
deadline: Duration,
) -> Result<String, String> { ) -> Result<String, String> {
let system = dir.join(SYSTEM_FILE); let system = dir.join(SYSTEM_FILE);
let config_file = dir.join("opencode.json"); let config_file = dir.join("opencode.json");
@ -132,7 +134,7 @@ impl<S: CostSink> Acp<S> {
let sink = FinalText::default(); let sink = FinalText::default();
let turn = runtime.run(&config, prompt, &sink); let turn = runtime.run(&config, prompt, &sink);
tokio::pin!(turn); tokio::pin!(turn);
match tokio::time::timeout(self.deadline, &mut turn).await { match tokio::time::timeout(deadline, &mut turn).await {
Ok(Ok(progress)) => { Ok(Ok(progress)) => {
self.record_cost(run, provider, &progress.telemetry.cost); self.record_cost(run, provider, &progress.telemetry.cost);
sink.take() sink.take()
@ -141,10 +143,7 @@ impl<S: CostSink> Acp<S> {
Err(_) => { Err(_) => {
let _ = canceller.cancel(); let _ = canceller.cancel();
let _ = tokio::time::timeout(self.cancel_wait, turn).await; let _ = tokio::time::timeout(self.cancel_wait, turn).await;
Err(format!( Err(format!("no answer within {} seconds", deadline.as_secs()))
"no answer within {} seconds",
self.deadline.as_secs()
))
} }
} }
} }
@ -283,6 +282,11 @@ mod tests {
use crate::prompt::Prompt; use crate::prompt::Prompt;
use crate::run::Model as _; use crate::run::Model as _;
/// A deadline no test reaches.
const LONG: Duration = Duration::from_mins(1);
/// A deadline the stalled agents run into.
const SHORT: Duration = Duration::from_millis(500);
/// An ACP agent in plain `sh`, logging to `<log>.*`. It writes its /// An ACP agent in plain `sh`, logging to `<log>.*`. It writes its
/// environment and copies of its config and the system prompt file /// environment and copies of its config and the system prompt file
/// beside it and appends its `$HOME` to `<log>.homes` at start, appends /// beside it and appends its `$HOME` to `<log>.homes` at start, appends
@ -340,7 +344,7 @@ done
} }
/// An [`Acp`] whose trampoline is [`AGENT`], logging under `dir`. /// An [`Acp`] whose trampoline is [`AGENT`], logging under `dir`.
fn acp(dir: &Path, deadline: Duration) -> (Acp<Arc<Costs>>, Arc<Costs>) { fn acp(dir: &Path) -> (Acp<Arc<Costs>>, Arc<Costs>) {
let script = dir.join("agent.sh"); let script = dir.join("agent.sh");
let log = dir.join("log"); let log = dir.join("log");
std::fs::write( std::fs::write(
@ -353,7 +357,6 @@ done
let acp = Acp { let acp = Acp {
program: PathBuf::from("/nix/store/x-opencode/bin/opencode"), program: PathBuf::from("/nix/store/x-opencode/bin/opencode"),
trampoline: script, trampoline: script,
deadline,
cancel_wait: Duration::from_millis(500), cancel_wait: Duration::from_millis(500),
cost: Arc::clone(&costs), cost: Arc::clone(&costs),
cwd: dir.join("cwd"), cwd: dir.join("cwd"),
@ -392,9 +395,9 @@ done
#[tokio::test] #[tokio::test]
async fn the_answer_is_the_final_text_and_its_cost_is_labelled_with_the_pr() { async fn the_answer_is_the_final_text_and_its_cost_is_labelled_with_the_pr() {
let dir = tempfile::tempdir().expect("a dir"); let dir = tempfile::tempdir().expect("a dir");
let (acp, costs) = acp(dir.path(), Duration::from_mins(1)); let (acp, costs) = acp(dir.path());
let answer = acp.ask(&Prompt::of("classify"), "p", &run()).await; let answer = acp.ask(&Prompt::of("classify"), "p", &run(), LONG).await;
assert_eq!( assert_eq!(
answer.as_deref(), answer.as_deref(),
Ok(r#"{"verdict":"match","certainty":90,"reason":"ok"}"#) Ok(r#"{"verdict":"match","certainty":90,"reason":"ok"}"#)
@ -416,9 +419,9 @@ done
#[tokio::test] #[tokio::test]
async fn opencode_gets_the_key_no_tools_and_the_runs_provider() { async fn opencode_gets_the_key_no_tools_and_the_runs_provider() {
let dir = tempfile::tempdir().expect("a dir"); let dir = tempfile::tempdir().expect("a dir");
let (acp, _) = acp(dir.path(), Duration::from_mins(1)); let (acp, _) = acp(dir.path());
acp.ask(&Prompt::of("classify"), "p", &run()) acp.ask(&Prompt::of("classify"), "p", &run(), LONG)
.await .await
.expect("an answer"); .expect("an answer");
let env = read(dir.path(), "env"); let env = read(dir.path(), "env");
@ -448,9 +451,9 @@ done
#[tokio::test] #[tokio::test]
async fn the_system_prompt_is_the_agents_and_the_prompt_is_the_pr_alone() { async fn the_system_prompt_is_the_agents_and_the_prompt_is_the_pr_alone() {
let dir = tempfile::tempdir().expect("a dir"); let dir = tempfile::tempdir().expect("a dir");
let (acp, _) = acp(dir.path(), Duration::from_mins(1)); let (acp, _) = acp(dir.path());
acp.ask(&Prompt::of("classify"), "p", &run()) acp.ask(&Prompt::of("classify"), "p", &run(), LONG)
.await .await
.expect("an answer"); .expect("an answer");
let config: Value = serde_json::from_str(&read(dir.path(), "config")).expect("JSON"); let config: Value = serde_json::from_str(&read(dir.path(), "config")).expect("JSON");
@ -470,9 +473,9 @@ done
#[tokio::test] #[tokio::test]
async fn a_permission_request_is_refused() { async fn a_permission_request_is_refused() {
let dir = tempfile::tempdir().expect("a dir"); let dir = tempfile::tempdir().expect("a dir");
let (acp, _) = acp(dir.path(), Duration::from_mins(1)); let (acp, _) = acp(dir.path());
acp.ask(&Prompt::of("ASK-READ"), "p", &run()) acp.ask(&Prompt::of("ASK-READ"), "p", &run(), LONG)
.await .await
.expect("an answer"); .expect("an answer");
let reply: Value = let reply: Value =
@ -483,9 +486,9 @@ done
#[tokio::test] #[tokio::test]
async fn the_deadline_cancels_the_turn() { async fn the_deadline_cancels_the_turn() {
let dir = tempfile::tempdir().expect("a dir"); let dir = tempfile::tempdir().expect("a dir");
let (acp, costs) = acp(dir.path(), Duration::from_millis(500)); let (acp, costs) = acp(dir.path());
let answer = acp.ask(&Prompt::of("STALL"), "p", &run()).await; let answer = acp.ask(&Prompt::of("STALL"), "p", &run(), SHORT).await;
assert_eq!(answer, Err("no answer within 0 seconds".to_owned())); assert_eq!(answer, Err("no answer within 0 seconds".to_owned()));
assert!( assert!(
read(dir.path(), "methods") read(dir.path(), "methods")
@ -498,7 +501,7 @@ done
#[tokio::test] #[tokio::test]
async fn an_agent_stuck_before_its_prompt_is_dropped_at_the_deadline() { async fn an_agent_stuck_before_its_prompt_is_dropped_at_the_deadline() {
let dir = tempfile::tempdir().expect("a dir"); let dir = tempfile::tempdir().expect("a dir");
let (acp, _) = acp(dir.path(), Duration::from_millis(500)); let (acp, _) = acp(dir.path());
let script = dir.path().join("agent.sh"); let script = dir.path().join("agent.sh");
let unanswered = std::fs::read_to_string(&script).expect("written").replacen( let unanswered = std::fs::read_to_string(&script).expect("written").replacen(
" initialize)", " initialize)",
@ -509,7 +512,7 @@ done
let answer = tokio::time::timeout( let answer = tokio::time::timeout(
Duration::from_secs(10), Duration::from_secs(10),
acp.ask(&Prompt::of("one"), "p", &run()), acp.ask(&Prompt::of("one"), "p", &run(), SHORT),
) )
.await .await
.expect("the deadline ends the ask"); .expect("the deadline ends the ask");
@ -519,9 +522,9 @@ done
#[tokio::test] #[tokio::test]
async fn no_token_counts_record_no_cost() { async fn no_token_counts_record_no_cost() {
let dir = tempfile::tempdir().expect("a dir"); let dir = tempfile::tempdir().expect("a dir");
let (acp, costs) = acp(dir.path(), Duration::from_mins(1)); let (acp, costs) = acp(dir.path());
acp.ask(&Prompt::of("MUTE"), "p", &run()) acp.ask(&Prompt::of("MUTE"), "p", &run(), LONG)
.await .await
.expect("an answer"); .expect("an answer");
assert!(costs.0.lock().expect("unpoisoned").is_empty()); assert!(costs.0.lock().expect("unpoisoned").is_empty());
@ -530,12 +533,12 @@ done
#[tokio::test] #[tokio::test]
async fn a_claude_run_is_refused() { async fn a_claude_run_is_refused() {
let dir = tempfile::tempdir().expect("a dir"); let dir = tempfile::tempdir().expect("a dir");
let (acp, _) = acp(dir.path(), Duration::from_mins(1)); let (acp, _) = acp(dir.path());
let mut claude = run(); let mut claude = run();
claude.runtime = Runtime::ApiKey; claude.runtime = Runtime::ApiKey;
assert!( assert!(
acp.ask(&Prompt::of("classify"), "p", &claude) acp.ask(&Prompt::of("classify"), "p", &claude, LONG)
.await .await
.is_err() .is_err()
); );
@ -556,17 +559,17 @@ done
#[tokio::test] #[tokio::test]
async fn every_run_gets_the_same_cwd_even_at_once() { async fn every_run_gets_the_same_cwd_even_at_once() {
let dir = tempfile::tempdir().expect("a dir"); let dir = tempfile::tempdir().expect("a dir");
let (acp, _) = acp(dir.path(), Duration::from_mins(1)); let (acp, _) = acp(dir.path());
let (run_one, run_two) = (run(), run()); let (run_one, run_two) = (run(), run());
let (prompt_one, prompt_two) = (Prompt::of("one"), Prompt::of("two")); let (prompt_one, prompt_two) = (Prompt::of("one"), Prompt::of("two"));
let (one, two) = tokio::join!( let (one, two) = tokio::join!(
acp.ask(&prompt_one, "p", &run_one), acp.ask(&prompt_one, "p", &run_one, LONG),
acp.ask(&prompt_two, "p", &run_two) acp.ask(&prompt_two, "p", &run_two, LONG)
); );
one.expect("an answer"); one.expect("an answer");
two.expect("an answer"); two.expect("an answer");
acp.ask(&Prompt::of("three"), "p", &run()) acp.ask(&Prompt::of("three"), "p", &run(), LONG)
.await .await
.expect("an answer"); .expect("an answer");
let cwd = dir.path().join("cwd").display().to_string(); let cwd = dir.path().join("cwd").display().to_string();
@ -576,12 +579,12 @@ done
#[tokio::test] #[tokio::test]
async fn each_run_has_a_home_of_its_own_removed_after_it() { async fn each_run_has_a_home_of_its_own_removed_after_it() {
let dir = tempfile::tempdir().expect("a dir"); let dir = tempfile::tempdir().expect("a dir");
let (acp, _) = acp(dir.path(), Duration::from_mins(1)); let (acp, _) = acp(dir.path());
acp.ask(&Prompt::of("one"), "p", &run()) acp.ask(&Prompt::of("one"), "p", &run(), LONG)
.await .await
.expect("an answer"); .expect("an answer");
acp.ask(&Prompt::of("two"), "p", &run()) acp.ask(&Prompt::of("two"), "p", &run(), LONG)
.await .await
.expect("an answer"); .expect("an answer");
let homes = read(dir.path(), "homes"); let homes = read(dir.path(), "homes");

33
swarm-review-classifier/src/answer.rs
View file

@ -33,8 +33,8 @@ struct Raw {
reason: String, reason: String,
} }
/// Parse `text`, the model's final message. A surrounding code fence is /// Parse `text`, the model's final message. A code fence around the object is
/// tolerated; anything else around the object is not. /// tolerated, with prose around the fence; prose around a bare object is not.
/// ///
/// # Errors /// # Errors
/// A message describing what is wrong with the answer, without quoting it. /// A message describing what is wrong with the answer, without quoting it.
@ -80,10 +80,19 @@ fn cap_reason(reason: String) -> String {
cut cut
} }
/// `text` without a surrounding code fence, with or without a language tag. /// `text` without its code fence, with or without a language tag. A fence
/// around the whole message is stripped; otherwise the first fenced block in
/// the message is taken, so prose before or after it is dropped.
fn unfence(text: &str) -> &str { fn unfence(text: &str) -> &str {
let Some(inner) = text.strip_prefix("```").and_then(|t| t.strip_suffix("```")) else { let inner = match text.strip_prefix("```").and_then(|t| t.strip_suffix("```")) {
return text; Some(inner) => inner,
None => match text
.split_once("```")
.and_then(|(_, t)| t.split_once("```"))
{
Some((block, _)) => block,
None => return text,
},
}; };
let inner = inner.strip_prefix("json").unwrap_or(inner); let inner = inner.strip_prefix("json").unwrap_or(inner);
inner.trim() inner.trim()
@ -113,6 +122,20 @@ mod tests {
); );
} }
#[test]
fn a_fenced_answer_inside_prose_parses() {
let a = parse(
"Here is my answer:\n```json\n{\"verdict\":\"match\",\"certainty\":80,\"reason\":\"x\"}\n```\nHope that helps.",
)
.expect("fence in prose");
assert_eq!((a.verdict, a.certainty), (Verdict::Match, 80));
let untagged =
"Answer: ```{\"verdict\":\"no-match\",\"certainty\":70,\"reason\":\"\"}``` done";
assert_eq!(parse(untagged).expect("untagged").verdict, Verdict::NoMatch);
assert!(parse("Here is my answer:\n```json\nnot json\n```").is_err());
assert!(parse("an unclosed ```json\n{\"verdict\":\"match\"").is_err());
}
#[test] #[test]
fn a_malformed_answer_is_an_error() { fn a_malformed_answer_is_an_error() {
for bad in [ for bad in [

1
swarm-review-classifier/src/main.rs
View file

@ -67,7 +67,6 @@ async fn serve() -> Result<()> {
acp: acp::Acp { acp: acp::Acp {
program: PathBuf::from(var("OPENCODE")?), program: PathBuf::from(var("OPENCODE")?),
trampoline, trampoline,
deadline: swarm_queue_client::review_classifier::CLASSIFICATION_DEADLINE,
cancel_wait: acp::CANCEL_WAIT, cancel_wait: acp::CANCEL_WAIT,
cost: cost::Otel, cost: cost::Otel,
cwd: std::env::temp_dir().join("review-classifier-acp-cwd"), cwd: std::env::temp_dir().join("review-classifier-acp-cwd"),

14
swarm-review-classifier/src/prompt.rs
View file

@ -54,6 +54,10 @@ pub fn shared_context(run: &TryOut) -> String {
out out
} }
/// Appended to the question when it is asked a second time.
const JSON_REMINDER: &str = "Your previous answer was not a JSON object. Reply with the JSON \
object and nothing else: no text before or after it, and no code fence.";
/// One classification's prompt: the run's shared context unchanged, then the /// One classification's prompt: the run's shared context unchanged, then the
/// question. /// question.
pub struct Prompt<'a> { pub struct Prompt<'a> {
@ -67,6 +71,16 @@ impl Prompt<'_> {
pub fn text(&self) -> String { pub fn text(&self) -> String {
format!("{}{}", self.shared, self.question) format!("{}{}", self.shared, self.question)
} }
/// The same prompt with a reminder to answer with the JSON object alone,
/// for asking again after an answer that was not that object.
#[must_use]
pub fn with_json_reminder(&self) -> Self {
Self {
shared: self.shared,
question: format!("{}\n{JSON_REMINDER}\n", self.question.trim_end()),
}
}
} }
#[cfg(test)] #[cfg(test)]

275
swarm-review-classifier/src/run.rs
View file

@ -13,7 +13,7 @@ use std::future::Future;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use std::sync::Mutex; use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration; use std::time::{Duration, Instant};
use hive_claude::{Attach, Config, Sink}; use hive_claude::{Attach, Config, Sink};
use swarm_queue_client::agent_inference::Effort; use swarm_queue_client::agent_inference::Effort;
@ -39,12 +39,14 @@ pub trait KeySource {
/// the tests. /// the tests.
pub trait Model { pub trait Model {
/// The final message for `prompt` on `run`'s runtime, model and effort, /// The final message for `prompt` on `run`'s runtime, model and effort,
/// run with `key`, or why there is none. /// run with `key`, or why there is none. The ask is stopped after
/// `deadline`.
fn ask( fn ask(
&self, &self,
prompt: &Prompt<'_>, prompt: &Prompt<'_>,
key: &str, key: &str,
run: &TryOut, run: &TryOut,
deadline: Duration,
) -> impl Future<Output = Result<String, String>> + Send; ) -> impl Future<Output = Result<String, String>> + Send;
} }
@ -59,10 +61,16 @@ pub async fn answer(keys: &impl KeySource, model: &impl Model, run: &TryOut) ->
let shared = prompt::shared_context(run); let shared = prompt::shared_context(run);
let mut rows: Vec<Row> = Vec::with_capacity(run.classifications.len()); let mut rows: Vec<Row> = Vec::with_capacity(run.classifications.len());
for c in &run.classifications { for c in &run.classifications {
let outcome = model let prompt = prompt::classification_prompt(&shared, c);
.ask(&prompt::classification_prompt(&shared, c), &key, run) let outcome = classify(
.await model,
.and_then(|text| answer::parse(&text)); &prompt,
&key,
run,
CLASSIFICATION_DEADLINE,
RETRY_FLOOR,
)
.await;
rows.push(aggregate::row(c, outcome)); rows.push(aggregate::row(c, outcome));
} }
if let Some(first) = rows if let Some(first) = rows
@ -80,6 +88,39 @@ pub async fn answer(keys: &impl KeySource, model: &impl Model, run: &TryOut) ->
Reply::Done { rows, overall } 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 /// 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. /// own, then the repo and the PR, so its cost reads back by both.
fn resource_attributes(base: Option<String>, run: &TryOut) -> String { fn resource_attributes(base: Option<String>, run: &TryOut) -> String {
@ -97,10 +138,16 @@ pub struct Runtimes<C, A> {
} }
impl<C: Model + Sync, A: Model + Sync> Model for Runtimes<C, 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> { async fn ask(
&self,
prompt: &Prompt<'_>,
key: &str,
run: &TryOut,
deadline: Duration,
) -> Result<String, String> {
match run.runtime { match run.runtime {
Runtime::ApiKey => self.claude.ask(prompt, key, run).await, Runtime::ApiKey => self.claude.ask(prompt, key, run, deadline).await,
Runtime::Acp(_) => self.acp.ask(prompt, key, run).await, Runtime::Acp(_) => self.acp.ask(prompt, key, run, deadline).await,
} }
} }
} }
@ -171,11 +218,12 @@ impl Model for Claude {
prompt: &Prompt<'_>, prompt: &Prompt<'_>,
api_key: &str, api_key: &str,
run: &TryOut, run: &TryOut,
deadline: Duration,
) -> Result<String, String> { ) -> Result<String, String> {
let attributes = resource_attributes(std::env::var("OTEL_RESOURCE_ATTRIBUTES").ok(), run); let attributes = resource_attributes(std::env::var("OTEL_RESOURCE_ATTRIBUTES").ok(), run);
let dir = scratch_dir(); let dir = scratch_dir();
let outcome = self let outcome = self
.ask_in(&dir, prompt, api_key, &attributes, &run.model, run.effort) .ask_in(&dir, prompt, api_key, &attributes, run, deadline)
.await; .await;
if let Err(e) = tokio::fs::remove_dir_all(&dir).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"); tracing::warn!(dir = %dir.display(), error = %e, "removing a run's directory failed");
@ -191,15 +239,15 @@ impl Claude {
prompt: &Prompt<'_>, prompt: &Prompt<'_>,
api_key: &str, api_key: &str,
attributes: &str, attributes: &str,
model: &str, run: &TryOut,
effort: Option<Effort>, deadline: Duration,
) -> Result<String, String> { ) -> Result<String, String> {
let system = dir.join("system.md"); let system = dir.join("system.md");
tokio::fs::create_dir_all(dir.join("config")) tokio::fs::create_dir_all(dir.join("config"))
.await .await
.and(tokio::fs::write(&system, prompt::SYSTEM).await) .and(tokio::fs::write(&system, prompt::SYSTEM).await)
.map_err(|e| format!("preparing the run's directory: {e}"))?; .map_err(|e| format!("preparing the run's directory: {e}"))?;
let config = self.config(dir, system, api_key, attributes, model, effort); let config = self.config(dir, system, api_key, attributes, &run.model, run.effort);
let running = let running =
hive_claude::Claude::spawn(&config, &Attach::OneOff).map_err(|e| e.to_string())?; hive_claude::Claude::spawn(&config, &Attach::OneOff).map_err(|e| e.to_string())?;
let cancel = running.cancel_handle(); let cancel = running.cancel_handle();
@ -207,16 +255,13 @@ impl Claude {
let message = stream_json_message(prompt); let message = stream_json_message(prompt);
let wait = running.wait(&message, &sink); let wait = running.wait(&message, &sink);
tokio::pin!(wait); tokio::pin!(wait);
match tokio::time::timeout(CLASSIFICATION_DEADLINE, &mut wait).await { match tokio::time::timeout(deadline, &mut wait).await {
Ok(Ok(())) => sink.take(), Ok(Ok(())) => sink.take(),
Ok(Err(e)) => Err(format!("claude failed: {e}")), Ok(Err(e)) => Err(format!("claude failed: {e}")),
Err(_) => { Err(_) => {
cancel.cancel(true); cancel.cancel(true);
let _ = wait.await; let _ = wait.await;
Err(format!( Err(format!("no answer within {} seconds", deadline.as_secs()))
"no answer within {} seconds",
CLASSIFICATION_DEADLINE.as_secs()
))
} }
} }
} }
@ -356,6 +401,7 @@ mod tests {
prompt: &Prompt<'_>, prompt: &Prompt<'_>,
key: &str, key: &str,
run: &TryOut, run: &TryOut,
_: Duration,
) -> Result<String, String> { ) -> Result<String, String> {
let prompt = prompt.text(); let prompt = prompt.text();
let wanted = match run.runtime { let wanted = match run.runtime {
@ -470,6 +516,170 @@ mod tests {
); );
} }
/// 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_queue_client::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] #[tokio::test]
async fn an_acp_run_is_asked_with_the_acp_key() { async fn an_acp_run_is_asked_with_the_acp_key() {
let reply = answer(&Key(Ok(())), &Scripted, &acp_run(&["ANSWER-MATCH"])).await; let reply = answer(&Key(Ok(())), &Scripted, &acp_run(&["ANSWER-MATCH"])).await;
@ -480,7 +690,13 @@ mod tests {
struct Named(&'static str); struct Named(&'static str);
impl Model for Named { impl Model for Named {
async fn ask(&self, _: &Prompt<'_>, _: &str, _: &TryOut) -> Result<String, String> { async fn ask(
&self,
_: &Prompt<'_>,
_: &str,
_: &TryOut,
_: Duration,
) -> Result<String, String> {
Ok(self.0.to_owned()) Ok(self.0.to_owned())
} }
} }
@ -493,14 +709,19 @@ mod tests {
}; };
assert_eq!( assert_eq!(
runtimes runtimes
.ask(&Prompt::of("p"), "k", &run(&[])) .ask(&Prompt::of("p"), "k", &run(&[]), CLASSIFICATION_DEADLINE)
.await .await
.as_deref(), .as_deref(),
Ok("claude") Ok("claude")
); );
assert_eq!( assert_eq!(
runtimes runtimes
.ask(&Prompt::of("p"), "k", &acp_run(&[])) .ask(
&Prompt::of("p"),
"k",
&acp_run(&[]),
CLASSIFICATION_DEADLINE
)
.await .await
.as_deref(), .as_deref(),
Ok("acp") Ok("acp")
@ -629,7 +850,17 @@ printf '{"type":"result","subtype":"success","is_error":false,"result":"ok"}\n'
question: "## The question\n".to_owned(), question: "## The question\n".to_owned(),
}; };
let answer = claude let answer = claude
.ask_in(&tmp.path().join("run"), &prompt, "k", "a=b", "sonnet", None) .ask_in(
&tmp.path().join("run"),
&prompt,
"k",
"a=b",
&TryOut {
effort: None,
..run(&[])
},
CLASSIFICATION_DEADLINE,
)
.await; .await;
assert_eq!(answer.as_deref(), Ok("ok")); assert_eq!(answer.as_deref(), Ok("ok"));
let args = std::fs::read_to_string(tmp.path().join("claude.sh.args")).expect("args"); let args = std::fs::read_to_string(tmp.path().join("claude.sh.args")).expect("args");