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.
This commit is contained in:
atlas
mara
parent
78412b3d3b
commit
2c5eef73cd
10 changed files with 868 additions and 69 deletions
|
|
@ -61,7 +61,7 @@ impl From<AgentInference> 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"));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<String> {
|
|||
.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<TryOutAccepted>), 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::<forgejo_api::ForgejoError>() {
|
||||
|
|
@ -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<std::sync::Mutex<Vec<TryOut>>>,
|
||||
}
|
||||
|
||||
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<Settings> {
|
||||
Ok(self.settings.clone())
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -805,7 +838,10 @@ pub(crate) mod tests {
|
|||
}
|
||||
|
||||
impl RunSink for Fake {
|
||||
async fn run(&self, _request: Vec<u8>) -> Reply {
|
||||
async fn run(&self, request: Vec<u8>) -> 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;
|
||||
|
|
|
|||
|
|
@ -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<Semaphore>,
|
||||
slots: Mutex<Slots>,
|
||||
/// Woken when a run ends or the limit changes.
|
||||
changed: Notify,
|
||||
accepted: Mutex<HashSet<Key>>,
|
||||
}
|
||||
|
||||
/// 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<Self>, key: Key) -> Result<Accepted, Refused> {
|
||||
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<Self>) -> 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<Reviews>);
|
||||
|
||||
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<S, Q, P>(
|
||||
reviews: &Arc<Reviews>,
|
||||
forge: Arc<S>,
|
||||
|
|
@ -241,18 +298,25 @@ fn start<S, Q, P>(
|
|||
) -> Option<JoinHandle<()>>
|
||||
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<Settings> {
|
||||
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<String>, 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<tokio::sync::Semaphore>);
|
||||
struct Gate {
|
||||
permits: Arc<tokio::sync::Semaphore>,
|
||||
reached: Arc<AtomicUsize>,
|
||||
settings: Arc<Mutex<Settings>>,
|
||||
}
|
||||
|
||||
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<u8>) -> 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<Settings> {
|
||||
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(
|
||||
|
|
|
|||
391
swarm-controller/src/review_classifier/settings.rs
Normal file
391
swarm-controller/src/review_classifier/settings.rs
Normal file
|
|
@ -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<usize> = 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<String>)]
|
||||
pub(crate) effort: Option<Effort>,
|
||||
/// 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<Output = anyhow::Result<Settings>> + Send;
|
||||
}
|
||||
|
||||
impl SettingsSource for async_nats::Client {
|
||||
async fn settings(&self) -> anyhow::Result<Settings> {
|
||||
let (settings, _) = read::<Settings>(&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<Store> {
|
||||
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<AppState>,
|
||||
) -> Result<Json<Settings>, 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<AppState>,
|
||||
headers: HeaderMap,
|
||||
body: Result<Json<Settings>, JsonRejection>,
|
||||
) -> Result<Json<Settings>, 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::<Settings>(&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<Settings> {
|
||||
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<Body>,
|
||||
) -> (StatusCode, serde_json::Value, Option<String>) {
|
||||
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<Body> {
|
||||
Request::get("/api/review-classifier/settings")
|
||||
.body(Body::empty())
|
||||
.expect("a request")
|
||||
}
|
||||
|
||||
fn put(body: &str) -> Request<Body> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue