diff --git a/docs/swarm/review-classifier.md b/docs/swarm/review-classifier.md index 3e99943f..0ab14773 100644 --- a/docs/swarm/review-classifier.md +++ b/docs/swarm/review-classifier.md @@ -338,9 +338,16 @@ An answer: ``` 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 -isn't that JSON object 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 +certainty field a whole number. The classifier reads the object from the +reply, or from the first code fence in it. When an answer isn't that object, +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 verdict. diff --git a/swarm-queue-client/src/review_classifier.rs b/swarm-queue-client/src/review_classifier.rs index 7e4b9a94..c1c817b5 100644 --- a/swarm-queue-client/src/review_classifier.rs +++ b/swarm-queue-client/src/review_classifier.rs @@ -23,8 +23,8 @@ pub const QUEUE_GROUP: &str = "review-classifier"; /// The most classifications one run carries. pub const MAX_CLASSIFICATIONS: usize = 10; -/// The longest one classification's model run may take before it is stopped -/// and its row carries a timeout. +/// The longest one classification may take, its first ask and any second ask +/// together, before it is stopped and its row carries a timeout. pub const CLASSIFICATION_DEADLINE: Duration = Duration::from_mins(3); /// The longest a whole run may take: every classification at its deadline, diff --git a/swarm-review-classifier/README.md b/swarm-review-classifier/README.md index 70a29dcd..6a84d283 100644 --- a/swarm-review-classifier/README.md +++ b/swarm-review-classifier/README.md @@ -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 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. + 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 whole: the caller sees an error, never a verdict. diff --git a/swarm-review-classifier/src/acp.rs b/swarm-review-classifier/src/acp.rs index 1b19b59e..d9c369dd 100644 --- a/swarm-review-classifier/src/acp.rs +++ b/swarm-review-classifier/src/acp.rs @@ -65,8 +65,6 @@ pub struct Acp { /// This binary, which execs [`Self::program`] when started as the /// trampoline. pub trampoline: PathBuf, - /// How long one classification may take before it is cancelled. - pub deadline: Duration, /// [`CANCEL_WAIT`] outside the tests. pub cancel_wait: Duration, pub cost: S, @@ -83,12 +81,15 @@ impl crate::run::Model for Acp { prompt: &crate::prompt::Prompt<'_>, key: &str, run: &TryOut, + deadline: Duration, ) -> Result { let Runtime::Acp(provider) = &run.runtime else { return Err("not an ACP run".to_owned()); }; 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 { tracing::warn!(dir = %dir.display(), error = %e, "removing a run's directory failed"); } @@ -104,6 +105,7 @@ impl Acp { key: &str, run: &TryOut, provider: &Provider, + deadline: Duration, ) -> Result { let system = dir.join(SYSTEM_FILE); let config_file = dir.join("opencode.json"); @@ -132,7 +134,7 @@ impl Acp { let sink = FinalText::default(); let turn = runtime.run(&config, prompt, &sink); tokio::pin!(turn); - match tokio::time::timeout(self.deadline, &mut turn).await { + match tokio::time::timeout(deadline, &mut turn).await { Ok(Ok(progress)) => { self.record_cost(run, provider, &progress.telemetry.cost); sink.take() @@ -141,10 +143,7 @@ impl Acp { Err(_) => { let _ = canceller.cancel(); let _ = tokio::time::timeout(self.cancel_wait, turn).await; - Err(format!( - "no answer within {} seconds", - self.deadline.as_secs() - )) + Err(format!("no answer within {} seconds", deadline.as_secs())) } } } @@ -283,6 +282,11 @@ mod tests { use crate::prompt::Prompt; 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 `.*`. It writes its /// environment and copies of its config and the system prompt file /// beside it and appends its `$HOME` to `.homes` at start, appends @@ -340,7 +344,7 @@ done } /// An [`Acp`] whose trampoline is [`AGENT`], logging under `dir`. - fn acp(dir: &Path, deadline: Duration) -> (Acp>, Arc) { + fn acp(dir: &Path) -> (Acp>, Arc) { let script = dir.join("agent.sh"); let log = dir.join("log"); std::fs::write( @@ -353,7 +357,6 @@ done let acp = Acp { program: PathBuf::from("/nix/store/x-opencode/bin/opencode"), trampoline: script, - deadline, cancel_wait: Duration::from_millis(500), cost: Arc::clone(&costs), cwd: dir.join("cwd"), @@ -392,9 +395,9 @@ done #[tokio::test] async fn the_answer_is_the_final_text_and_its_cost_is_labelled_with_the_pr() { 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!( answer.as_deref(), Ok(r#"{"verdict":"match","certainty":90,"reason":"ok"}"#) @@ -416,9 +419,9 @@ done #[tokio::test] async fn opencode_gets_the_key_no_tools_and_the_runs_provider() { 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 .expect("an answer"); let env = read(dir.path(), "env"); @@ -448,9 +451,9 @@ done #[tokio::test] async fn the_system_prompt_is_the_agents_and_the_prompt_is_the_pr_alone() { 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 .expect("an answer"); let config: Value = serde_json::from_str(&read(dir.path(), "config")).expect("JSON"); @@ -470,9 +473,9 @@ done #[tokio::test] async fn a_permission_request_is_refused() { 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 .expect("an answer"); let reply: Value = @@ -483,9 +486,9 @@ done #[tokio::test] async fn the_deadline_cancels_the_turn() { 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!( read(dir.path(), "methods") @@ -498,7 +501,7 @@ done #[tokio::test] async fn an_agent_stuck_before_its_prompt_is_dropped_at_the_deadline() { 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 unanswered = std::fs::read_to_string(&script).expect("written").replacen( " initialize)", @@ -509,7 +512,7 @@ done let answer = tokio::time::timeout( Duration::from_secs(10), - acp.ask(&Prompt::of("one"), "p", &run()), + acp.ask(&Prompt::of("one"), "p", &run(), SHORT), ) .await .expect("the deadline ends the ask"); @@ -519,9 +522,9 @@ done #[tokio::test] async fn no_token_counts_record_no_cost() { 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 .expect("an answer"); assert!(costs.0.lock().expect("unpoisoned").is_empty()); @@ -530,12 +533,12 @@ done #[tokio::test] async fn a_claude_run_is_refused() { 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(); claude.runtime = Runtime::ApiKey; assert!( - acp.ask(&Prompt::of("classify"), "p", &claude) + acp.ask(&Prompt::of("classify"), "p", &claude, LONG) .await .is_err() ); @@ -556,17 +559,17 @@ done #[tokio::test] async fn every_run_gets_the_same_cwd_even_at_once() { 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 (prompt_one, prompt_two) = (Prompt::of("one"), Prompt::of("two")); let (one, two) = tokio::join!( - acp.ask(&prompt_one, "p", &run_one), - acp.ask(&prompt_two, "p", &run_two) + acp.ask(&prompt_one, "p", &run_one, LONG), + acp.ask(&prompt_two, "p", &run_two, LONG) ); one.expect("an answer"); two.expect("an answer"); - acp.ask(&Prompt::of("three"), "p", &run()) + acp.ask(&Prompt::of("three"), "p", &run(), LONG) .await .expect("an answer"); let cwd = dir.path().join("cwd").display().to_string(); @@ -576,12 +579,12 @@ done #[tokio::test] async fn each_run_has_a_home_of_its_own_removed_after_it() { 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 .expect("an answer"); - acp.ask(&Prompt::of("two"), "p", &run()) + acp.ask(&Prompt::of("two"), "p", &run(), LONG) .await .expect("an answer"); let homes = read(dir.path(), "homes"); diff --git a/swarm-review-classifier/src/answer.rs b/swarm-review-classifier/src/answer.rs index ecb69e5a..8d39d70e 100644 --- a/swarm-review-classifier/src/answer.rs +++ b/swarm-review-classifier/src/answer.rs @@ -33,8 +33,8 @@ struct Raw { reason: String, } -/// Parse `text`, the model's final message. A surrounding code fence is -/// tolerated; anything else around the object is not. +/// Parse `text`, the model's final message. A code fence around the object is +/// tolerated, with prose around the fence; prose around a bare object is not. /// /// # Errors /// A message describing what is wrong with the answer, without quoting it. @@ -80,10 +80,19 @@ fn cap_reason(reason: String) -> String { 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 { - let Some(inner) = text.strip_prefix("```").and_then(|t| t.strip_suffix("```")) else { - return text; + let inner = match text.strip_prefix("```").and_then(|t| t.strip_suffix("```")) { + 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); 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] fn a_malformed_answer_is_an_error() { for bad in [ diff --git a/swarm-review-classifier/src/main.rs b/swarm-review-classifier/src/main.rs index 5e07e176..c182a725 100644 --- a/swarm-review-classifier/src/main.rs +++ b/swarm-review-classifier/src/main.rs @@ -67,7 +67,6 @@ async fn serve() -> Result<()> { acp: acp::Acp { program: PathBuf::from(var("OPENCODE")?), trampoline, - deadline: swarm_queue_client::review_classifier::CLASSIFICATION_DEADLINE, cancel_wait: acp::CANCEL_WAIT, cost: cost::Otel, cwd: std::env::temp_dir().join("review-classifier-acp-cwd"), diff --git a/swarm-review-classifier/src/prompt.rs b/swarm-review-classifier/src/prompt.rs index c72eb469..082b25c5 100644 --- a/swarm-review-classifier/src/prompt.rs +++ b/swarm-review-classifier/src/prompt.rs @@ -54,6 +54,10 @@ pub fn shared_context(run: &TryOut) -> String { 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 /// question. pub struct Prompt<'a> { @@ -67,6 +71,16 @@ impl Prompt<'_> { pub fn text(&self) -> String { 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)] diff --git a/swarm-review-classifier/src/run.rs b/swarm-review-classifier/src/run.rs index 0c4a44ae..288b4734 100644 --- a/swarm-review-classifier/src/run.rs +++ b/swarm-review-classifier/src/run.rs @@ -13,7 +13,7 @@ use std::future::Future; use std::path::{Path, PathBuf}; use std::sync::Mutex; use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::Duration; +use std::time::{Duration, Instant}; use hive_claude::{Attach, Config, Sink}; use swarm_queue_client::agent_inference::Effort; @@ -39,12 +39,14 @@ pub trait KeySource { /// 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. + /// 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> + Send; } @@ -59,10 +61,16 @@ pub async fn answer(keys: &impl KeySource, model: &impl Model, run: &TryOut) -> let shared = prompt::shared_context(run); let mut rows: Vec = 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)); + 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 @@ -80,6 +88,39 @@ pub async fn answer(keys: &impl KeySource, model: &impl Model, run: &TryOut) -> 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 { + 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, run: &TryOut) -> String { @@ -97,10 +138,16 @@ pub struct Runtimes { } impl Model for Runtimes { - async fn ask(&self, prompt: &Prompt<'_>, key: &str, run: &TryOut) -> Result { + async fn ask( + &self, + prompt: &Prompt<'_>, + key: &str, + run: &TryOut, + deadline: Duration, + ) -> Result { match run.runtime { - Runtime::ApiKey => self.claude.ask(prompt, key, run).await, - Runtime::Acp(_) => self.acp.ask(prompt, key, run).await, + Runtime::ApiKey => self.claude.ask(prompt, key, run, deadline).await, + Runtime::Acp(_) => self.acp.ask(prompt, key, run, deadline).await, } } } @@ -171,11 +218,12 @@ impl Model for Claude { prompt: &Prompt<'_>, api_key: &str, run: &TryOut, + deadline: Duration, ) -> Result { 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) + .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"); @@ -191,15 +239,15 @@ impl Claude { prompt: &Prompt<'_>, api_key: &str, attributes: &str, - model: &str, - effort: Option, + run: &TryOut, + deadline: Duration, ) -> Result { 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 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(); @@ -207,16 +255,13 @@ impl Claude { let message = stream_json_message(prompt); let wait = running.wait(&message, &sink); 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(Err(e)) => Err(format!("claude failed: {e}")), Err(_) => { cancel.cancel(true); let _ = wait.await; - Err(format!( - "no answer within {} seconds", - CLASSIFICATION_DEADLINE.as_secs() - )) + Err(format!("no answer within {} seconds", deadline.as_secs())) } } } @@ -356,6 +401,7 @@ mod tests { prompt: &Prompt<'_>, key: &str, run: &TryOut, + _: Duration, ) -> Result { let prompt = prompt.text(); 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>); + + 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 { + 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>, + } + + impl Slow { + fn new(first: Duration) -> Self { + Self { + first, + deadlines: Mutex::default(), + } + } + + fn deadlines(&self) -> Vec { + self.deadlines.lock().expect("lock").clone() + } + } + + impl Model for Slow { + async fn ask( + &self, + _: &Prompt<'_>, + _: &str, + _: &TryOut, + deadline: Duration, + ) -> Result { + 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 { + 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] async fn an_acp_run_is_asked_with_the_acp_key() { let reply = answer(&Key(Ok(())), &Scripted, &acp_run(&["ANSWER-MATCH"])).await; @@ -480,7 +690,13 @@ mod tests { struct Named(&'static str); impl Model for Named { - async fn ask(&self, _: &Prompt<'_>, _: &str, _: &TryOut) -> Result { + async fn ask( + &self, + _: &Prompt<'_>, + _: &str, + _: &TryOut, + _: Duration, + ) -> Result { Ok(self.0.to_owned()) } } @@ -493,14 +709,19 @@ mod tests { }; assert_eq!( runtimes - .ask(&Prompt::of("p"), "k", &run(&[])) + .ask(&Prompt::of("p"), "k", &run(&[]), CLASSIFICATION_DEADLINE) .await .as_deref(), Ok("claude") ); assert_eq!( runtimes - .ask(&Prompt::of("p"), "k", &acp_run(&[])) + .ask( + &Prompt::of("p"), + "k", + &acp_run(&[]), + CLASSIFICATION_DEADLINE + ) .await .as_deref(), Ok("acp") @@ -629,7 +850,17 @@ printf '{"type":"result","subtype":"success","is_error":false,"result":"ok"}\n' question: "## The question\n".to_owned(), }; 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; assert_eq!(answer.as_deref(), Ok("ok")); let args = std::fs::read_to_string(tmp.path().join("claude.sh.args")).expect("args");