Watch
0
0
Fork
You've already forked hyperhive
0

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

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

View file

@ -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.

View file

@ -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,

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
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.

View file

@ -65,8 +65,6 @@ pub struct Acp<S> {
/// 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<S: CostSink> crate::run::Model for Acp<S> {
prompt: &crate::prompt::Prompt<'_>,
key: &str,
run: &TryOut,
deadline: Duration,
) -> Result<String, String> {
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<S: CostSink> Acp<S> {
key: &str,
run: &TryOut,
provider: &Provider,
deadline: Duration,
) -> Result<String, String> {
let system = dir.join(SYSTEM_FILE);
let config_file = dir.join("opencode.json");
@ -132,7 +134,7 @@ impl<S: CostSink> Acp<S> {
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<S: CostSink> Acp<S> {
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 `<log>.*`. It writes its
/// environment and copies of its config and the system prompt file
/// 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`.
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 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");

View file

@ -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 [

View file

@ -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"),

View file

@ -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)]

View file

@ -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<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 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));
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<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 {
@ -97,10 +138,16 @@ pub struct 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 {
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<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)
.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<Effort>,
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, 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<String, String> {
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<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]
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<String, String> {
async fn ask(
&self,
_: &Prompt<'_>,
_: &str,
_: &TryOut,
_: Duration,
) -> Result<String, String> {
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");