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.

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<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");