From 2c5eef73cddb751b5efefea34b396c4eaad2a81e Mon Sep 17 00:00:00 2001 From: atlas Date: Mon, 12 Oct 2026 19:21:33 +0200 Subject: [PATCH] review classifier: model, effort and review concurrency as swarm-wide settings The classifier ran every classification on a hard-coded `haiku` with no `--effort`, and the controller ran at most two requested reviews at once, also hard-coded, so changing any of them took a code change and a redeploy. swarm-controller now keeps one settings document, `{model, effort, max_concurrent_reviews}`, in its own `review-classifier-settings` KV bucket under `all`, the way it keeps the schedule definitions. No document means today's values: haiku, no `--effort`, two reviews. `GET /api/review-classifier/settings` reads it; `PUT` replaces it whole, needs the gateway's `Remote-User` like the try-out POST, and refuses a bad model (the inference route's check), an effort outside the five levels or a limit outside 1..=8 with problem+json 400. Every run request (try-out and requested review) now carries the model and effort, and the classifier passes them to claude instead of its constant, which is gone. A try-out reads the settings when it is posted; a review reads them for the limit before it waits for a place and again for the model and effort once it has one. The review semaphore is replaced by a running count under a resizable limit, so a PUT, or the next review's read, changes how many run without a restart, and a lower limit lets the runs already going finish. MAX_WAITING_REVIEWS is unchanged. Closes #5151, part of #5096. --- docs/swarm/review-classifier.md | 46 ++- swarm-controller/README.md | 5 + swarm-controller/src/inference.rs | 2 +- swarm-controller/src/main.rs | 6 +- swarm-controller/src/review_classifier.rs | 83 +++- .../src/review_classifier/review.rs | 312 ++++++++++++-- .../src/review_classifier/settings.rs | 391 ++++++++++++++++++ swarm-queue-client/src/review_classifier.rs | 6 + swarm-review-classifier/src/prompt.rs | 2 + swarm-review-classifier/src/run.rs | 84 +++- 10 files changed, 868 insertions(+), 69 deletions(-) create mode 100644 swarm-controller/src/review_classifier/settings.rs diff --git a/docs/swarm/review-classifier.md b/docs/swarm/review-classifier.md index 4e8b86e9..1a1ff9ac 100644 --- a/docs/swarm/review-classifier.md +++ b/docs/swarm/review-classifier.md @@ -128,11 +128,49 @@ credential: `swarm-controller` reads the PR and sends it along. ### Swarm-level settings -The operator sets the classifier's settings for the whole swarm: +The model, the effort and how many requested reviews run at once are one +document for the whole swarm. `swarm-controller` keeps it in its +`review-classifier-settings` KV bucket under the key `all`, and no other +service reads it. With no document stored, the classifier runs on `haiku` at +claude's default effort, two reviews at a time. -- the runtime (see [Runtimes](#runtimes)) and its credentials -- the model and effort, and related inference settings -- the maximum number of concurrent runs +`GET /api/review-classifier/settings` answers the settings in force. +`PUT /api/review-classifier/settings` replaces them whole and answers what it +stored. Both take and give one shape: + +```json +{ "model": "haiku", "effort": null, "max_concurrent_reviews": 2 } +``` + +- `model`: the model every classification runs on, as claude's `--model` + takes it: one word with no whitespace, not starting with `-` +- `effort`: `low`, `medium`, `high`, `xhigh` or `max`; `null`, or left out of + a PUT, passes no `--effort` +- `max_concurrent_reviews`: 1 to 8, how many requested reviews run at once. + Try-outs don't count toward it. Up to eight more reviews wait for a place + +The controller reads the settings for every run and sends the model and +effort in the run request, so a change applies from the next run with no +restart. A try-out reads them when it's posted, a review when it gets its +place. Lowering the limit lets the runs already going finish. + +The PUT needs the `Remote-User` header the gateway sets, as the try-out POST +does. Every error is problem+json with the reason in `detail`: + +- **400** for a body that isn't JSON or isn't this shape (a missing `model` + or `max_concurrent_reviews`, an unknown field, an effort outside the five + levels, a `max_concurrent_reviews` that isn't a whole number), an empty + `model`, one with whitespace or one starting with `-`, or a + `max_concurrent_reviews` outside 1 to 8 +- **401** for a PUT with no `Remote-User` +- **503** when this host has no swarm queue, or it isn't connected +- **500** when the controller can't read or write the settings + +When the controller can't read the settings, a try-out answers the same 503 +or 500, and a review posts a no-verdict comment. The runtime (see +[Runtimes](#runtimes)) isn't a setting yet, and the operator writes the API +key by hand in the secret store +([`services.md`](services.md#review-classifier)). The per-agent equivalent is [`inference.md`](inference.md), where the controller writes model and effort to a queue KV bucket diff --git a/swarm-controller/README.md b/swarm-controller/README.md index 6af0b1f3..dd8bf723 100644 --- a/swarm-controller/README.md +++ b/swarm-controller/README.md @@ -63,6 +63,11 @@ The swarm's control plane. The swarm UI and `swarmctl` are its clients. (`POST /api/review-classifier/try-out`, `GET /api/review-classifier/try-out/{id}`, `GET /api/review-classifier/rules`, `GET /api/review-classifier/runs`) → [`docs/swarm/review-classifier.md`](../docs/swarm/review-classifier.md). +- **Review classifier settings** — the classifier's model, effort and how + many requested reviews run at once, one document in the + `review-classifier-settings` KV bucket, which this daemon alone reads and + writes (`GET`/`PUT /api/review-classifier/settings`) → + [`docs/swarm/review-classifier.md`](../docs/swarm/review-classifier.md#swarm-level-settings). - **Relays** — each agent's terminal and turn-state header as SSE, agent icons, the cross-repo issue report, and the UI's quick links. diff --git a/swarm-controller/src/inference.rs b/swarm-controller/src/inference.rs index 7b9a89b5..25184ed1 100644 --- a/swarm-controller/src/inference.rs +++ b/swarm-controller/src/inference.rs @@ -61,7 +61,7 @@ impl From for InferenceSettings { /// A model name as the agent passes it to claude: one non-empty word that /// cannot be read as a command-line flag. -fn check_model(name: &str, what: &str) -> Result<(), String> { +pub(crate) fn check_model(name: &str, what: &str) -> Result<(), String> { if name.is_empty() { return Err(format!("the {what} is empty")); } diff --git a/swarm-controller/src/main.rs b/swarm-controller/src/main.rs index 3e823c96..c57244af 100644 --- a/swarm-controller/src/main.rs +++ b/swarm-controller/src/main.rs @@ -769,7 +769,7 @@ fn socket_path() -> PathBuf { (name = "repos", description = "forge repo browsing (swarm-ui's issue-report page)"), (name = "forge", description = "forge site admins (`swarmctl forge make-admin`)"), (name = "schedules", description = "scheduled prompts (`swarmctl schedule`)"), - (name = "review-classifier", description = "trying review classifications against a PR (swarm-ui's try-out page)"), + (name = "review-classifier", description = "trying review classifications against a PR (swarm-ui's try-out page), and the classifier's swarm-wide settings"), (name = "pulls", description = "open pull requests, for the try-out page's dropdowns"), ) )] @@ -3351,6 +3351,10 @@ fn build_app(state: AppState) -> axum::Router { .routes(routes!(review_classifier::get_try_out)) .routes(routes!(review_classifier::get_rules)) .routes(routes!(review_classifier::runs::get_runs)) + .routes(routes!( + review_classifier::settings::get_settings, + review_classifier::settings::put_settings + )) .routes(routes!(webhook::post_webhook_forge)) .split_for_parts(); // Just the JSON, not the UI — Swagger UI itself is nginx-hosted from diff --git a/swarm-controller/src/review_classifier.rs b/swarm-controller/src/review_classifier.rs index 5a308741..b679516d 100644 --- a/swarm-controller/src/review_classifier.rs +++ b/swarm-controller/src/review_classifier.rs @@ -35,11 +35,13 @@ use swarm_queue_client::review_classifier::{ }; use utoipa::{IntoParams, ToSchema}; +use self::settings::{Settings, SettingsSource}; use crate::forge::{self, DefaultBranchFiles, PullForReview}; use crate::{AppState, error_problem}; pub(crate) mod review; pub(crate) mod runs; +pub(crate) mod settings; /// How long a finished run's result stays readable. pub(crate) const RESULT_TTL: Duration = Duration::from_hours(1); @@ -321,6 +323,8 @@ impl RunSink for async_nats::Client { (status = 503, description = "this host has no forge token yet, no swarm queue \ configured, or a queue that is not connected; or the forge failed or did not \ answer within two minutes (problem+json)", body = String), + (status = 500, description = "the review classifier's settings could not be read \ + (problem+json)", body = String), ), tag = "review-classifier" )] @@ -340,6 +344,7 @@ pub(crate) async fn post_try_out( let run = Run::checked(req).map_err(|d| error_problem(StatusCode::BAD_REQUEST, &d))?; #[cfg(test)] if let Some(fake) = state.try_outs.fake.clone() { + let settings = fake.settings().await.map_err(|e| settings_problem(&e))?; return Box::pin(start( &state.try_outs, &*fake, @@ -347,6 +352,7 @@ pub(crate) async fn post_try_out( state.runs.clone(), operator, run, + settings, )) .await; } @@ -370,6 +376,7 @@ pub(crate) async fn post_try_out( ) })?; let runs = state.runs.clone(); + let settings = queue.settings().await.map_err(|e| settings_problem(&e))?; Box::pin(start( &state.try_outs, forge.as_ref(), @@ -377,6 +384,7 @@ pub(crate) async fn post_try_out( runs, operator, run, + settings, )) .await } @@ -617,8 +625,9 @@ fn operator(headers: &HeaderMap) -> Option { .map(str::to_owned) } -/// Read the PR, record the run and send it; the reply settles the run when -/// it arrives, and the settled run goes into `runs`. +/// Read the PR, record the run and send it on `settings`' model and effort; +/// the reply settles the run when it arrives, and the settled run goes into +/// `runs`. async fn start( try_outs: &TryOuts, pulls: &impl PullSource, @@ -626,6 +635,7 @@ async fn start( runs: runs::Runs, operator: String, run: Run, + settings: Settings, ) -> Result<(StatusCode, Json), ProblemDetails> { let started_unix = crate::agent_renewal::unix_now(); let repo = format!("{}/{}", run.owner, run.name); @@ -647,6 +657,8 @@ async fn start( commits: pull.commits, diff: pull.diff, classifications: run.classifications, + model: settings.model, + effort: settings.effort, }; let payload = serde_json::to_vec(&request).map_err(|e| { error_problem( @@ -692,6 +704,14 @@ async fn start( Ok((StatusCode::ACCEPTED, Json(TryOutAccepted { id }))) } +/// The settings could not be read: 503 when the queue is not connected, 500 +/// otherwise. +fn settings_problem(e: &anyhow::Error) -> ProblemDetails { + let detail = format!("{e:#}"); + tracing::warn!(error = %detail, "reading the review classifier's settings for a try-out failed"); + error_problem(crate::wanted_error_status(e), &detail) +} + /// 404 when the forge knows no such repo or PR, 503 for anything else. fn pull_problem(e: &anyhow::Error, repo: &str, pr: i64) -> ProblemDetails { match e.downcast_ref::() { @@ -743,18 +763,31 @@ pub(crate) mod tests { use super::*; /// A forge holding one PR, `o/r` 7, with `files` on `o/r`'s default - /// branch, and a classifier answering every run with `reply`. `o/empty` - /// has no commit, `o/down` fails with something other than a 404, and - /// `o/hang` never answers. + /// branch, `settings` stored, and a classifier answering every run with + /// `reply` and keeping its request. `o/empty` has no commit, `o/down` + /// fails with something other than a 404, and `o/hang` never answers. #[derive(Clone)] pub(crate) struct Fake { reply: Reply, files: Vec<(&'static str, String)>, + pub(crate) settings: Settings, + pub(crate) sent: Arc>>, } impl Fake { pub(crate) fn new(reply: Reply, files: Vec<(&'static str, String)>) -> Self { - Self { reply, files } + Self { + reply, + files, + settings: Settings::default(), + sent: Arc::default(), + } + } + } + + impl SettingsSource for Fake { + async fn settings(&self) -> anyhow::Result { + Ok(self.settings.clone()) } } @@ -805,7 +838,10 @@ pub(crate) mod tests { } impl RunSink for Fake { - async fn run(&self, _request: Vec) -> Reply { + async fn run(&self, request: Vec) -> Reply { + if let Ok(sent) = serde_json::from_slice(&request) { + self.sent.lock().expect("unpoisoned").push(sent); + } self.reply.clone() } } @@ -833,7 +869,7 @@ pub(crate) mod tests { fn app_with_files(reply: Reply, files: Vec<(&'static str, String)>) -> axum::Router { let (mut state, _) = crate::tests::state_with_roster(); - state.try_outs.fake = Some(Arc::new(Fake { reply, files })); + state.try_outs.fake = Some(Arc::new(Fake::new(reply, files))); crate::build_app(state) } @@ -903,6 +939,32 @@ pub(crate) mod tests { assert!(run.get("error").is_none(), "{run}"); } + #[tokio::test] + async fn a_try_out_runs_on_the_stored_model_and_effort() { + use swarm_queue_client::agent_inference::Effort; + + let (mut state, _) = crate::tests::state_with_roster(); + let fake = Arc::new(Fake { + settings: Settings { + model: "opus".to_owned(), + effort: Some(Effort::High), + max_concurrent_reviews: 2, + }, + ..Fake::new(done(), Vec::new()) + }); + state.try_outs.fake = Some(Arc::clone(&fake)); + let app = crate::build_app(state); + let (status, accepted) = post(&app, body("o/r", 7)).await; + assert_eq!(status, StatusCode::ACCEPTED, "{accepted}"); + settled(&app, accepted["id"].as_str().expect("an id")).await; + let sent = fake.sent.lock().expect("unpoisoned"); + assert_eq!(sent.len(), 1); + assert_eq!( + (sent[0].model.as_str(), sent[0].effort), + ("opus", Some(Effort::High)) + ); + } + #[tokio::test] async fn a_failed_run_reports_its_error_and_no_verdict() { let app = app(Reply::Failed { @@ -929,10 +991,7 @@ pub(crate) mod tests { ] { let doc = Arc::new(runs::tests::FakeDoc::default()); let (mut state, _) = crate::tests::state_with_roster(); - state.try_outs.fake = Some(Arc::new(Fake { - reply: reply.clone(), - files: Vec::new(), - })); + state.try_outs.fake = Some(Arc::new(Fake::new(reply.clone(), Vec::new()))); state.runs = runs::Runs::Fake(Arc::clone(&doc)); let app = crate::build_app(state); let (_, accepted) = post(&app, body("o/r", 7)).await; diff --git a/swarm-controller/src/review_classifier/review.rs b/swarm-controller/src/review_classifier/review.rs index c673ae60..621537b5 100644 --- a/swarm-controller/src/review_classifier/review.rs +++ b/swarm-controller/src/review_classifier/review.rs @@ -8,29 +8,28 @@ //! official reviews count toward approvals or block a merge. //! //! A run with no clean verdict never approves: see [`no_clean_verdict_event`]. -//! At most [`MAX_CONCURRENT_REVIEWS`] runs go at once and at most +//! At most `max_concurrent_reviews` runs go at once, as the swarm's +//! [settings](super::settings) say when a run asks for its place, and at most //! [`MAX_WAITING_REVIEWS`] more wait here, in memory, so a restart drops them. //! See [`Reviews`] for what happens to a request beyond that. use std::collections::HashSet; use std::fmt::Write as _; use std::future::Future; -use std::sync::{Arc, Mutex, PoisonError}; +use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; use serde::Deserialize; use swarm_queue_client::review_classifier::{Overall, Reply, Row, TryOut, Verdict}; -use tokio::sync::Semaphore; +use tokio::sync::Notify; use tokio::task::JoinHandle; use super::runs::{self, RunOutcome, RunRecord, Runs, Trigger}; +use super::settings::{Settings, SettingsSource}; use super::{FileSource, PullSource, RunSink, split_repo}; use crate::AppState; use crate::agent_renewal::unix_now; use crate::forge::{self, classifier_account::LOGIN, classifier_account::Review}; -/// The most review runs in flight at once. -const MAX_CONCURRENT_REVIEWS: usize = 2; - /// The most review runs waiting for a slot. const MAX_WAITING_REVIEWS: usize = 8; @@ -144,17 +143,29 @@ type Key = (String, String, i64, String); /// The review runs this controller has accepted and not yet finished. /// /// A request whose [`Key`] is already running or waiting starts nothing. A -/// request beyond [`MAX_CONCURRENT_REVIEWS`] + [`MAX_WAITING_REVIEWS`] starts -/// no run either; the account posts a no-verdict comment saying so. +/// request beyond the run limit + [`MAX_WAITING_REVIEWS`] starts no run +/// either; the account posts a no-verdict comment saying so. pub(crate) struct Reviews { - slots: Arc, + slots: Mutex, + /// Woken when a run ends or the limit changes. + changed: Notify, accepted: Mutex>, } +/// How many runs may go at once, and how many do. +struct Slots { + limit: usize, + running: usize, +} + impl Default for Reviews { fn default() -> Self { Self { - slots: Arc::new(Semaphore::new(MAX_CONCURRENT_REVIEWS)), + slots: Mutex::new(Slots { + limit: Settings::default().max_concurrent_reviews, + running: 0, + }), + changed: Notify::new(), accepted: Mutex::default(), } } @@ -165,18 +176,34 @@ impl Default for Reviews { enum Refused { /// A run for the same key is running or waiting. Duplicate, - /// Every run and wait place is taken. - Full, + /// Every run and wait place is taken, by this many runs. + Full(usize), } impl Reviews { + fn slots(&self) -> MutexGuard<'_, Slots> { + self.slots.lock().unwrap_or_else(PoisonError::into_inner) + } + + /// Let `limit` runs go at once from now on. Runs already going when it + /// shrinks finish; no new one starts until fewer than `limit` run. + pub(crate) fn resize(&self, limit: usize) { + self.slots().limit = limit; + self.changed.notify_waiters(); + } + + #[cfg(test)] + pub(crate) fn limit(&self) -> usize { + self.slots().limit + } + fn accept(self: &Arc, key: Key) -> Result { let mut accepted = self.accepted.lock().unwrap_or_else(PoisonError::into_inner); if accepted.contains(&key) { return Err(Refused::Duplicate); } - if accepted.len() >= MAX_CONCURRENT_REVIEWS + MAX_WAITING_REVIEWS { - return Err(Refused::Full); + if accepted.len() >= self.slots().limit + MAX_WAITING_REVIEWS { + return Err(Refused::Full(accepted.len())); } accepted.insert(key.clone()); Ok(Accepted { @@ -184,6 +211,23 @@ impl Reviews { key, }) } + + /// A place to run in, once fewer runs go than the limit allows. + async fn slot(self: &Arc) -> Slot { + loop { + let changed = self.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + { + let mut slots = self.slots(); + if slots.running < slots.limit { + slots.running += 1; + return Slot(Arc::clone(self)); + } + } + changed.await; + } + } } /// One accepted run; dropping it frees its key and its place. @@ -202,6 +246,16 @@ impl Drop for Accepted { } } +/// One run holding a place; dropping it lets the next one start. +struct Slot(Arc); + +impl Drop for Slot { + fn drop(&mut self) { + self.0.slots().running -= 1; + self.0.changed.notify_waiters(); + } +} + /// Start a review run for a verified `review-request` delivery that asks for /// [`LOGIN`], and return at once. pub(crate) fn on_delivery(state: &AppState, body: &[u8]) { @@ -231,6 +285,9 @@ pub(crate) fn on_delivery(state: &AppState, body: &[u8]) { /// Accept `request` into `reviews` and spawn its run, or the no-verdict /// comment when there's no place for it, and record it in `runs` when it /// ends. The task, if one was spawned. +/// +/// The run reads the settings for the limit before it waits for a place, and +/// again for the model and effort once it has one. fn start( reviews: &Arc, forge: Arc, @@ -241,18 +298,25 @@ fn start( ) -> Option> where S: FileSource + PullSource + Send + Sync + 'static, - Q: RunSink, + Q: RunSink + SettingsSource, P: ReviewPoster + Send + Sync + 'static, { let repo = format!("{}/{}", request.owner, request.name); match reviews.accept(request.key()) { Ok(accepted) => { - let slots = Arc::clone(&reviews.slots); + let reviews = Arc::clone(reviews); Some(tokio::spawn(async move { - let Ok(_slot) = slots.acquire_owned().await else { - return; - }; + match queue.settings().await { + Ok(settings) => reviews.resize(settings.max_concurrent_reviews), + Err(e) => tracing::warn!( + error = %format!("{e:#}"), + "reading the review classifier's settings failed; the run waits under \ + the last limit read" + ), + } + let slot = reviews.slot().await; let run = Box::pin(run(&*forge, &*forge, &queue, &*poster, &request)).await; + drop(slot); drop(accepted); runs.record(run).await; })) @@ -264,14 +328,14 @@ where ); None } - Err(Refused::Full) => { + Err(Refused::Full(taken)) => { tracing::warn!( %repo, pr = request.pr, head_sha = request.head.as_str(), "review requested, but every run and wait place is taken; posting a no-verdict comment" ); Some(tokio::spawn(async move { let started_unix = unix_now(); - let ruling = queue_full_ruling(); + let ruling = queue_full_ruling(taken); let reason = ruling.reason.clone(); let sent = post(&*poster, &request, None, ruling.review()).await; let ended = ended(RunOutcome::QueueFull, reason, sent); @@ -282,11 +346,12 @@ where } } -/// The ruling for a request [`Reviews`] had no place for. -fn queue_full_ruling() -> Ruling { +/// The ruling for a request [`Reviews`] had no place for, with `taken` runs +/// running or waiting. +fn queue_full_ruling(taken: usize) -> Ruling { let reason = format!( - "{MAX_CONCURRENT_REVIEWS} reviews are running and {MAX_WAITING_REVIEWS} are waiting, so \ - this request starts no run; request the review again later" + "{taken} reviews are running or waiting, so this request starts no run; \ + request the review again later" ); Ruling { event: no_clean_verdict_event(), @@ -342,6 +407,16 @@ impl RunSink for Queue { } } +impl SettingsSource for Queue { + async fn settings(&self) -> anyhow::Result { + let Some(client) = &self.0 else { + anyhow::bail!("no swarm queue is configured on this host"); + }; + swarm_queue_client::ensure_connected(client)?; + client.settings().await + } +} + /// How far a run got. #[derive(Debug, Clone, PartialEq, Eq)] enum Outcome { @@ -356,7 +431,7 @@ enum Outcome { async fn run( files: &impl FileSource, pulls: &impl PullSource, - sink: &impl RunSink, + sink: &(impl RunSink + SettingsSource), poster: &impl ReviewPoster, req: &Requested, ) -> RunRecord { @@ -479,7 +554,7 @@ fn record( async fn classify( files: &impl FileSource, pulls: &impl PullSource, - sink: &impl RunSink, + sink: &(impl RunSink + SettingsSource), req: &Requested, ) -> (Option, Outcome) { let repo = format!("{}/{}", req.owner, req.name); @@ -514,6 +589,15 @@ async fn classify( ); } }; + let settings = match sink.settings().await { + Ok(settings) => settings, + Err(e) => { + return ( + head, + Outcome::NotRun(format!("reading the classifier's settings failed: {e:#}")), + ); + } + }; let request = TryOut { operator: req.requester.clone(), repo, @@ -522,6 +606,8 @@ async fn classify( commits: pull.commits, diff: pull.diff, classifications: rules.classifications, + model: settings.model, + effort: settings.effort, }; let sent = request.classifications.len(); let outcome = match serde_json::to_vec(&request) { @@ -686,6 +772,7 @@ fn cell(text: &str) -> String { #[cfg(test)] mod tests { use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use swarm_queue_client::review_classifier::Polarity; @@ -838,19 +925,44 @@ mod tests { assert_eq!(ruling.event, Event::Comment); } - /// A classifier that answers approve once [`Gate::open`] lets it through. + /// A classifier that counts the runs reaching it and answers approve once + /// [`Gate::open`] lets each through, beside a store holding `settings`. #[derive(Clone)] - struct Gate(Arc); + struct Gate { + permits: Arc, + reached: Arc, + settings: Arc>, + } impl Gate { - fn open(&self) { - self.0.add_permits(1); + fn closed() -> Self { + Self { + permits: Arc::new(tokio::sync::Semaphore::new(0)), + reached: Arc::default(), + settings: Arc::default(), + } + } + + fn open(&self, runs: usize) { + self.permits.add_permits(runs); + } + + fn reached(&self) -> usize { + self.reached.load(Ordering::SeqCst) + } + + fn store(&self, max_concurrent_reviews: usize) { + self.settings + .lock() + .expect("unpoisoned") + .max_concurrent_reviews = max_concurrent_reviews; } } impl RunSink for Gate { async fn run(&self, _request: Vec) -> Reply { - let _permit = self.0.acquire().await.expect("never closed"); + self.reached.fetch_add(1, Ordering::SeqCst); + let _permit = self.permits.acquire().await.expect("never closed"); Reply::Done { rows: vec![row(None, true, true)], overall: Overall::Approve, @@ -858,6 +970,19 @@ mod tests { } } + impl SettingsSource for Gate { + async fn settings(&self) -> anyhow::Result { + Ok(self.settings.lock().expect("unpoisoned").clone()) + } + } + + /// Let every spawned task run until it waits on something. + async fn settle() { + for _ in 0..200 { + tokio::task::yield_now().await; + } + } + fn requested_at(head: &str) -> Requested { Requested { head: head.to_owned(), @@ -874,7 +999,7 @@ mod tests { }, vec![rules_json()], )); - let gate = Gate(Arc::new(tokio::sync::Semaphore::new(0))); + let gate = Gate::closed(); let poster = Arc::new(Posted::default()); let doc = Arc::new(FakeDoc::default()); let store = Runs::Fake(Arc::clone(&doc)); @@ -896,7 +1021,7 @@ mod tests { for head in ["a", "b", "c"] { assert!(start_at(head).is_none(), "{head}"); } - gate.open(); + gate.open(1); for task in tasks.drain(..) { task.await.expect("ran"); } @@ -922,7 +1047,7 @@ mod tests { }, vec![rules_json()], )); - let gate = Gate(Arc::new(tokio::sync::Semaphore::new(0))); + let gate = Gate::closed(); let poster = Arc::new(Posted::default()); let doc = Arc::new(FakeDoc::default()); let store = Runs::Fake(Arc::clone(&doc)); @@ -938,7 +1063,7 @@ mod tests { .expect("a run or a comment") }; - let places = MAX_CONCURRENT_REVIEWS + MAX_WAITING_REVIEWS; + let places = Settings::default().max_concurrent_reviews + MAX_WAITING_REVIEWS; let runs: Vec<_> = (0..places).map(|i| start_at(format!("h{i}"))).collect(); tokio::time::timeout(Duration::from_secs(5), start_at("over".to_owned())) .await @@ -952,9 +1077,9 @@ mod tests { assert_eq!(*commit, None); assert_eq!(review.event, "COMMENT"); assert!( - review - .body - .starts_with("**No verdict:** 2 reviews are running and 8 are waiting"), + review.body.starts_with( + "**No verdict:** 10 reviews are running or waiting, so this request" + ), "{}", review.body ); @@ -971,12 +1096,12 @@ mod tests { refused[0] .detail .as_deref() - .is_some_and(|d| d.starts_with("2 reviews are running and 8 are waiting")), + .is_some_and(|d| d.starts_with("10 reviews are running or waiting")), "{:?}", refused[0].detail ); - gate.open(); + gate.open(1); for run in runs { run.await.expect("ran"); } @@ -1166,6 +1291,111 @@ mod tests { assert_eq!(all[0].2.event, "APPROVED"); } + #[tokio::test] + async fn a_new_limit_applies_to_the_next_run_without_a_restart() { + let reviews = Arc::new(Reviews::default()); + let fake = Arc::new(Fake::new( + Reply::Failed { + error: String::new(), + }, + vec![rules_json()], + )); + let gate = Gate::closed(); + let poster = Arc::new(Posted::default()); + let start_at = |head: &str| { + start( + &reviews, + Arc::clone(&fake), + gate.clone(), + Arc::clone(&poster), + Runs::Fake(Arc::default()), + requested_at(head), + ) + .expect("accepted") + }; + + // The stored limit is read before a run waits, not the default 2. + gate.store(1); + let mut tasks = vec![start_at("a"), start_at("b")]; + settle().await; + assert_eq!(gate.reached(), 1); + + // Raised, as a PUT does: the waiting run starts at once. + gate.store(3); + reviews.resize(3); + settle().await; + assert_eq!(gate.reached(), 2); + tasks.push(start_at("c")); + settle().await; + assert_eq!(gate.reached(), 3); + + // Lowered: the three keep running and the next one waits. + gate.store(1); + tasks.push(start_at("d")); + settle().await; + assert_eq!(gate.reached(), 3); + + gate.open(1); + for task in tasks { + tokio::time::timeout(Duration::from_secs(5), task) + .await + .expect("a finished run frees its place") + .expect("ran"); + } + assert_eq!(gate.reached(), 4); + assert_eq!(reviews.slots().running, 0, "every place is freed"); + assert_eq!(poster.0.lock().expect("unpoisoned").len(), 4); + } + + #[tokio::test] + async fn a_review_runs_on_the_stored_model_and_effort() { + use swarm_queue_client::agent_inference::Effort; + + let mut fake = Fake::new( + Reply::Done { + rows: vec![row(None, true, true)], + overall: Overall::Approve, + }, + vec![rules_json()], + ); + fake.settings = Settings { + model: "sonnet".to_owned(), + effort: Some(Effort::Medium), + max_concurrent_reviews: 2, + }; + let poster = Posted::default(); + run(&fake, &fake, &fake, &poster, &request("r")).await; + let sent = fake.sent.lock().expect("unpoisoned"); + assert_eq!(sent.len(), 1); + assert_eq!( + (sent[0].model.as_str(), sent[0].effort), + ("sonnet", Some(Effort::Medium)) + ); + } + + #[tokio::test] + async fn unreadable_settings_run_nothing_and_post_a_comment() { + let fake = Fake::new( + Reply::Done { + rows: vec![row(None, true, true)], + overall: Overall::Approve, + }, + vec![rules_json()], + ); + let poster = Posted::default(); + run(&fake, &fake, &Queue(None), &poster, &request("r")).await; + let (_, commit, review) = poster.0.into_inner().expect("unpoisoned").remove(0); + assert_eq!(review.event, "COMMENT"); + assert!( + review + .body + .contains("reading the classifier's settings failed: no swarm queue is configured"), + "{}", + review.body + ); + assert_eq!(commit.as_deref(), Some("abc123")); + } + #[tokio::test] async fn a_clean_run_posts_its_verdict_on_the_head_commit() { let (target, commit, review) = posted( diff --git a/swarm-controller/src/review_classifier/settings.rs b/swarm-controller/src/review_classifier/settings.rs new file mode 100644 index 00000000..c361b4d2 --- /dev/null +++ b/swarm-controller/src/review_classifier/settings.rs @@ -0,0 +1,391 @@ +//! `GET` and `PUT /api/review-classifier/settings`: the model and effort +//! every run uses, and how many requested reviews run at once, for the whole +//! swarm. +//! +//! The settings are one document under one key ([`KEY`] in [`BUCKET`]), +//! which only this daemon reads or writes, as the schedule definitions are. +//! No document is [`Settings::default`]. A try-out reads them when it starts +//! and a review when it gets its place, and each sends the model and effort +//! in its run request, so a change applies to the next run with no restart. + +use std::future::Future; +use std::ops::RangeInclusive; + +use anyhow::Context as _; +use async_nats::jetstream::kv::Store; +use axum::Json; +use axum::extract::State; +use axum::extract::rejection::JsonRejection; +use axum::http::{HeaderMap, StatusCode}; +use problem_details::ProblemDetails; +use serde::{Deserialize, Serialize}; +use swarm_queue_client::agent_inference::Effort; +use utoipa::ToSchema; + +use crate::schedules::{queue, read, refused, write}; +use crate::{AppState, error_problem, wanted_error_status}; + +/// The bucket holding the settings. No other service holds a grant on it. +pub const BUCKET: &str = "review-classifier-settings"; + +/// The one key in [`BUCKET`]. +pub const KEY: &str = "all"; + +/// How many times a losing writer re-reads and re-applies, as in +/// `wanted::MAX_ATTEMPTS`. +const MAX_ATTEMPTS: usize = 5; + +/// What `max_concurrent_reviews` may be. Each running review is one `claude` +/// at a time in the classifier's container, which has no bound of its own. +pub(crate) const CONCURRENT_REVIEWS: RangeInclusive = 1..=8; + +/// The review classifier's settings: the value under [`KEY`], and the body of +/// both routes. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, ToSchema)] +#[serde(deny_unknown_fields)] +pub(crate) struct Settings { + /// The model every classification runs on, as claude's `--model` takes + /// it: one word, not starting with `-`. + pub(crate) model: String, + /// One of `low`, `medium`, `high`, `xhigh`, `max`; `null` or absent runs + /// at claude's own default. + #[serde(default)] + #[schema(value_type = Option)] + pub(crate) effort: Option, + /// How many requested reviews run at once, 1 to 8. Try-outs don't count. + pub(crate) max_concurrent_reviews: usize, +} + +impl Default for Settings { + fn default() -> Self { + Self { + model: "haiku".to_owned(), + effort: None, + max_concurrent_reviews: 2, + } + } +} + +impl Settings { + /// What is wrong with these settings, if anything. + fn check(&self) -> Result<(), String> { + crate::inference::check_model(&self.model, "model")?; + if !CONCURRENT_REVIEWS.contains(&self.max_concurrent_reviews) { + return Err(format!( + "max_concurrent_reviews {} is not {} to {}", + self.max_concurrent_reviews, + CONCURRENT_REVIEWS.start(), + CONCURRENT_REVIEWS.end() + )); + } + Ok(()) + } +} + +/// Where a run reads the settings: the queue in production, a fake in the +/// tests. +pub(crate) trait SettingsSource { + fn settings(&self) -> impl Future> + Send; +} + +impl SettingsSource for async_nats::Client { + async fn settings(&self) -> anyhow::Result { + let (settings, _) = read::(&open(self).await?, KEY) + .await + .context("reading the review classifier's settings")?; + Ok(settings) + } +} + +/// Open [`BUCKET`], creating it if nothing has yet. +async fn open(client: &async_nats::Client) -> anyhow::Result { + let js = async_nats::jetstream::new(client.clone()); + if let Ok(store) = js.get_key_value(BUCKET).await { + return Ok(store); + } + js.create_key_value(async_nats::jetstream::kv::Config { + bucket: BUCKET.to_owned(), + description: "The review classifier's swarm-wide settings".to_owned(), + history: 1, + ..Default::default() + }) + .await + .with_context(|| format!("creating the {BUCKET} bucket")) +} + +fn queue_problem(e: &anyhow::Error) -> ProblemDetails { + tracing::warn!(error = %format!("{e:#}"), "review classifier settings: queue failure"); + error_problem(wanted_error_status(e), &format!("{e:#}")) +} + +/// Read the settings. +#[utoipa::path( + get, + path = "/api/review-classifier/settings", + responses( + (status = 200, description = "the settings in force; with none stored, \ + `{\"model\":\"haiku\",\"effort\":null,\"max_concurrent_reviews\":2}`", + body = Settings), + (status = 503, description = "no swarm queue is configured on this host, or it is \ + not connected (problem+json)", body = String), + (status = 500, description = "the settings could not be read (problem+json)", + body = String), + ), + tag = "review-classifier" +)] +pub(crate) async fn get_settings( + State(state): State, +) -> Result, ProblemDetails> { + let client = queue(&state).map_err(refused)?; + client + .settings() + .await + .map(Json) + .map_err(|e| queue_problem(&e)) +} + +/// Replace the settings. +#[utoipa::path( + put, + path = "/api/review-classifier/settings", + request_body = Settings, + responses( + (status = 200, description = "the settings as stored; they replace the previous \ + ones whole and apply from the next run", body = Settings), + (status = 400, description = "not JSON or not this shape: a missing `model` or \ + `max_concurrent_reviews`, an unknown field, an effort other than the five \ + levels, a negative or non-integer `max_concurrent_reviews`; or a model that is \ + empty, has whitespace or starts with `-`; or `max_concurrent_reviews` outside \ + 1 to 8 (problem+json)", body = String), + (status = 401, description = "the request carries no `Remote-User`, so there is no \ + operator to name (problem+json)", body = String), + (status = 503, description = "no swarm queue is configured on this host, or it is \ + not connected (problem+json)", body = String), + (status = 500, description = "the settings could not be written (problem+json)", + body = String), + ), + tag = "review-classifier" +)] +pub(crate) async fn put_settings( + State(state): State, + headers: HeaderMap, + body: Result, JsonRejection>, +) -> Result, ProblemDetails> { + let operator = super::operator(&headers).ok_or_else(|| { + error_problem( + StatusCode::UNAUTHORIZED, + "no Remote-User header: the gateway names the operator, and this request did not \ + come through it", + ) + })?; + let Json(value) = body.map_err(|e| error_problem(StatusCode::BAD_REQUEST, &e.body_text()))?; + value + .check() + .map_err(|why| error_problem(StatusCode::BAD_REQUEST, &why))?; + let client = queue(&state).map_err(refused)?; + let store = open(&client).await.map_err(|e| queue_problem(&e))?; + for _ in 0..MAX_ATTEMPTS { + let (current, revision) = read::(&store, KEY) + .await + .map_err(|e| queue_problem(&e))?; + if current == value + || write(&store, KEY, &value, revision) + .await + .map_err(|e| queue_problem(&e))? + { + state.reviews.resize(value.max_concurrent_reviews); + tracing::info!(%operator, ?value, "set the review classifier's settings"); + return Ok(Json(value)); + } + } + Err(queue_problem(&anyhow::anyhow!( + "gave up writing the review classifier's settings after {MAX_ATTEMPTS} conflicting \ + writes" + ))) +} + +#[cfg(test)] +mod tests { + use axum::body::Body; + use axum::http::Request; + use tower::ServiceExt as _; + + use super::*; + use crate::test_nats::TestNats; + + fn settings(json: &str) -> serde_json::Result { + serde_json::from_str(json) + } + + #[test] + fn no_document_is_todays_values() { + assert_eq!( + serde_json::to_value(Settings::default()).unwrap(), + serde_json::json!({"model": "haiku", "effort": null, "max_concurrent_reviews": 2}) + ); + } + + #[test] + fn bad_settings_are_refused() { + for json in [ + r#"{"model":"","max_concurrent_reviews":2}"#, + r#"{"model":"opus 4","max_concurrent_reviews":2}"#, + r#"{"model":"--print","max_concurrent_reviews":2}"#, + r#"{"model":"haiku","max_concurrent_reviews":0}"#, + r#"{"model":"haiku","max_concurrent_reviews":9}"#, + ] { + assert!(settings(json).unwrap().check().is_err(), "{json}"); + } + // Control: the bounds themselves, and a full model id. + for json in [ + r#"{"model":"claude-haiku-4-5","max_concurrent_reviews":1}"#, + r#"{"model":"haiku","effort":"max","max_concurrent_reviews":8}"#, + ] { + assert!(settings(json).unwrap().check().is_ok(), "{json}"); + } + } + + #[test] + fn a_bad_shape_does_not_decode() { + for json in [ + r#"{"max_concurrent_reviews":2}"#, + r#"{"model":"haiku"}"#, + r#"{"model":"haiku","effort":"extreme","max_concurrent_reviews":2}"#, + r#"{"model":"haiku","effort":"HIGH","max_concurrent_reviews":2}"#, + r#"{"model":"haiku","max_concurrent_reviews":-1}"#, + r#"{"model":"haiku","max_concurrent_reviews":2.5}"#, + r#"{"model":"haiku","max_concurrent_reviews":2,"api_key":"x"}"#, + ] { + assert!(settings(json).is_err(), "{json}"); + } + // Control: effort may be left out or null. + assert_eq!( + settings(r#"{"model":"haiku","max_concurrent_reviews":2}"#).unwrap(), + Settings::default() + ); + assert_eq!( + settings(r#"{"model":"haiku","effort":null,"max_concurrent_reviews":2}"#).unwrap(), + Settings::default() + ); + } + + /// An app whose queue is `nats`. + async fn app(nats: &TestNats) -> (axum::Router, AppState) { + let (mut state, _) = crate::tests::state_with_roster(); + state.status = Some(std::sync::Arc::new(crate::status::StatusReader::new( + nats.client().await, + std::time::Duration::from_mins(1), + ))); + (crate::build_app(state.clone()), state) + } + + async fn send( + app: &axum::Router, + req: Request, + ) -> (StatusCode, serde_json::Value, Option) { + let res = app.clone().oneshot(req).await.expect("infallible"); + let status = res.status(); + let kind = res + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()) + .map(str::to_owned); + let bytes = axum::body::to_bytes(res.into_body(), usize::MAX) + .await + .expect("a body"); + ( + status, + serde_json::from_slice(&bytes).unwrap_or_default(), + kind, + ) + } + + fn get() -> Request { + Request::get("/api/review-classifier/settings") + .body(Body::empty()) + .expect("a request") + } + + fn put(body: &str) -> Request { + Request::put("/api/review-classifier/settings") + .header("content-type", "application/json") + .header("remote-user", "alice") + .body(Body::from(body.to_owned())) + .expect("a request") + } + + #[tokio::test] + async fn a_put_is_what_the_get_and_the_next_run_read() { + let nats = TestNats::start(); + let (app, state) = app(&nats).await; + let (status, before, _) = send(&app, get()).await; + assert_eq!(status, StatusCode::OK, "{before}"); + assert_eq!(before, serde_json::to_value(Settings::default()).unwrap()); + + let wanted = serde_json::json!({ + "model": "sonnet", "effort": "high", "max_concurrent_reviews": 5, + }); + let (status, stored, _) = send(&app, put(&wanted.to_string())).await; + assert_eq!(status, StatusCode::OK, "{stored}"); + assert_eq!(stored, wanted); + let (_, after, _) = send(&app, get()).await; + assert_eq!(after, wanted); + assert_eq!(state.reviews.limit(), 5); + + let client = nats.client().await; + let read = client.settings().await.expect("readable"); + assert_eq!( + (read.model.as_str(), read.effort), + ("sonnet", Some(Effort::High)) + ); + + // The same value again is no conflict. + let (status, _, _) = send(&app, put(&wanted.to_string())).await; + assert_eq!(status, StatusCode::OK); + } + + #[tokio::test] + async fn a_refused_put_is_a_problem_and_changes_nothing() { + let nats = TestNats::start(); + let (app, state) = app(&nats).await; + for bad in [ + "not json", + r#"{"model":"haiku"}"#, + r#"{"model":"haiku","effort":"extreme","max_concurrent_reviews":2}"#, + r#"{"model":"-x","max_concurrent_reviews":2}"#, + r#"{"model":"haiku","max_concurrent_reviews":9}"#, + ] { + let (status, problem, kind) = send(&app, put(bad)).await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{bad}: {problem}"); + assert!(problem["detail"].is_string(), "{bad}: {problem}"); + assert_eq!(kind.as_deref(), Some("application/problem+json"), "{bad}"); + } + let (_, after, _) = send(&app, get()).await; + assert_eq!(after, serde_json::to_value(Settings::default()).unwrap()); + assert_eq!(state.reviews.limit(), 2); + } + + #[tokio::test] + async fn a_put_without_an_operator_is_401() { + let (state, _) = crate::tests::state_with_roster(); + let req = Request::put("/api/review-classifier/settings") + .header("content-type", "application/json") + .body(Body::from( + serde_json::to_string(&Settings::default()).unwrap(), + )) + .expect("a request"); + let (status, _, _) = send(&crate::build_app(state), req).await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + } + + #[tokio::test] + async fn without_a_queue_both_verbs_are_503() { + let (state, _) = crate::tests::state_with_roster(); + let app = crate::build_app(state); + let (status, _, _) = send(&app, get()).await; + assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE); + let body = serde_json::to_string(&Settings::default()).unwrap(); + let (status, _, _) = send(&app, put(&body)).await; + assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE); + } +} diff --git a/swarm-queue-client/src/review_classifier.rs b/swarm-queue-client/src/review_classifier.rs index cb048ec7..d3d5620f 100644 --- a/swarm-queue-client/src/review_classifier.rs +++ b/swarm-queue-client/src/review_classifier.rs @@ -9,6 +9,8 @@ use std::time::Duration; use serde::{Deserialize, Serialize}; +use crate::agent_inference::Effort; + /// The subject every run goes to. The controller is its one /// publisher and the classifier its one subscriber; the classifier answers /// on the request's reply subject. @@ -48,6 +50,10 @@ pub struct TryOut { /// The PR's unified diff, as the forge serves it. pub diff: String, pub classifications: Vec, + /// The model every classification runs on, as claude's `--model` takes it. + pub model: String, + /// The effort every classification runs at; `None` passes no `--effort`. + pub effort: Option, } /// One commit of the PR. diff --git a/swarm-review-classifier/src/prompt.rs b/swarm-review-classifier/src/prompt.rs index 8fbabcc4..8148a86b 100644 --- a/swarm-review-classifier/src/prompt.rs +++ b/swarm-review-classifier/src/prompt.rs @@ -215,6 +215,8 @@ diff --git a/src/main.rs b/src/main.rs " .to_owned(), classifications: Vec::new(), + model: "haiku".to_owned(), + effort: None, } } diff --git a/swarm-review-classifier/src/run.rs b/swarm-review-classifier/src/run.rs index 8a70ab4b..f631e9ba 100644 --- a/swarm-review-classifier/src/run.rs +++ b/swarm-review-classifier/src/run.rs @@ -14,14 +14,12 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; use hive_claude::{Attach, Config, Sink}; +use swarm_queue_client::agent_inference::Effort; use swarm_queue_client::review_classifier::{CLASSIFICATION_DEADLINE, Reply, Row, TryOut}; use swarm_secret_client::review_classifier::{ApiKey, api_key_path}; use crate::{aggregate, answer, prompt}; -/// The model every classification runs on. -const MODEL: &str = "haiku"; - /// How long `claude` may print nothing before it is stopped. const IDLE_TIMEOUT: Duration = Duration::from_mins(2); @@ -35,12 +33,15 @@ pub trait KeySource { /// Asks the model one classification: `claude` in production, a fake in the /// tests. pub trait Model { - /// The model's final message for `prompt`, or why there is none. + /// The final message of `model` at `effort` for `prompt`, or why there + /// is none. fn ask( &self, prompt: &str, api_key: &str, resource_attributes: &str, + model: &str, + effort: Option, ) -> impl Future> + Send; } @@ -61,6 +62,8 @@ pub async fn answer(keys: &impl KeySource, model: &impl Model, run: &TryOut) -> &prompt::classification_prompt(&shared, c), &api_key, &attributes, + &run.model, + run.effort, ) .await .and_then(|text| answer::parse(&text)); @@ -124,9 +127,18 @@ pub struct Claude { } impl Model for Claude { - async fn ask(&self, prompt: &str, api_key: &str, attributes: &str) -> Result { + async fn ask( + &self, + prompt: &str, + api_key: &str, + attributes: &str, + model: &str, + effort: Option, + ) -> Result { let dir = scratch_dir(); - let outcome = self.ask_in(&dir, prompt, api_key, attributes).await; + let outcome = self + .ask_in(&dir, prompt, api_key, attributes, model, effort) + .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"); } @@ -141,13 +153,15 @@ impl Claude { prompt: &str, api_key: &str, attributes: &str, + model: &str, + effort: Option, ) -> 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); + let config = self.config(dir, system, api_key, attributes, model, effort); let running = hive_claude::Claude::spawn(&config, &Attach::OneOff).map_err(|e| e.to_string())?; let cancel = running.cancel_handle(); @@ -170,9 +184,18 @@ impl Claude { /// One classification's invocation: the trampoline as the program, with /// the real `claude` and the run's own variables in its environment. - fn config(&self, dir: &Path, system: PathBuf, api_key: &str, attributes: &str) -> Config { + fn config( + &self, + dir: &Path, + system: PathBuf, + api_key: &str, + attributes: &str, + model: &str, + effort: Option, + ) -> Config { Config { - model: Some(MODEL.to_owned()), + model: Some(model.to_owned()), + effort: effort.map(|e| e.as_str().to_owned()), cwd: Some(dir.to_owned()), system_prompt_file: Some(system), strict_mcp_config: true, @@ -257,8 +280,16 @@ mod tests { struct Scripted; impl Model for Scripted { - async fn ask(&self, prompt: &str, api_key: &str, attrs: &str) -> Result { + async fn ask( + &self, + prompt: &str, + api_key: &str, + attrs: &str, + model: &str, + effort: Option, + ) -> Result { assert_eq!(api_key, "k"); + assert_eq!((model, effort), ("sonnet", Some(Effort::Low))); assert!(attrs.ends_with("repo=o/r,pr=7"), "{attrs}"); if prompt.contains("ANSWER-FAIL") { Err("claude failed: exit 1".to_owned()) @@ -288,6 +319,8 @@ mod tests { threshold: 50, }) .collect(), + model: "sonnet".to_owned(), + effort: Some(Effort::Low), } } @@ -357,6 +390,8 @@ mod tests { PathBuf::from("/tmp/run/system.md"), "k", "a=b", + "sonnet", + None, ); assert_eq!(config.program, Some(claude.trampoline.clone())); assert!(config.env.contains(&( @@ -365,6 +400,35 @@ mod tests { ))); } + #[test] + fn the_runs_model_and_effort_are_claudes_flags() { + let claude = Claude { + program: PathBuf::from("claude"), + trampoline: PathBuf::from("trampoline"), + }; + let config = |model, effort| { + claude.config( + Path::new("/tmp/run"), + PathBuf::from("/tmp/run/system.md"), + "k", + "a=b", + model, + effort, + ) + }; + let set = config("opus", Some(Effort::Xhigh)); + assert_eq!( + (set.model.as_deref(), set.effort.as_deref()), + (Some("opus"), Some("xhigh")) + ); + let unset = config("haiku", None); + assert_eq!( + (unset.model.as_deref(), unset.effort.as_deref()), + (Some("haiku"), None), + "no effort passes no --effort" + ); + } + #[test] fn the_repo_and_pr_extend_the_containers_own_attributes() { let r = run(&[]);