use std::collections::BTreeMap;
use std::path::Path;
use serde::{Deserialize, Serialize};
use crate::hook::END_START_BUDGET;
pub const STEP: f64 = 1.5;
pub const HOLD: f64 = 0.85;
pub const REFERENCE_LESSONS: u64 = 30;
pub const MIN_CALLS: u64 = 5;
pub const EXPLORE_SHARE: f64 = 0.2;
pub const FAILED_EXPLORE_SHARE: f64 = EXPLORE_SHARE / 4.0;
const Z95: f64 = 1.96;
pub const CACHE: &str = "judge-length.json";
pub const VERSION: u32 = 2;
#[must_use]
pub fn edge(j: i32) -> usize {
(END_START_BUDGET as f64 * STEP.powi(j)).round() as usize
}
#[must_use]
pub fn step_of(chars: usize) -> i32 {
let mut j = 0;
while chars > edge(j) {
j += 1;
}
while chars <= edge(j - 1) && edge(j - 1) > 0 {
j -= 1;
}
j
}
#[derive(Debug, Clone, PartialEq)]
pub struct Call {
pub model: String,
pub chars: usize,
pub sent: u64,
pub answered: u64,
}
#[must_use]
pub fn parse(line: &str) -> Option<Call> {
let (rest, failed) = match line.strip_prefix("judge ") {
Some(rest) => (rest, false),
None => (line.strip_prefix("judge-failed ")?, true),
};
let field = |k: &str| rest.split_whitespace().find_map(|w| w.strip_prefix(k));
let num = |k: &str| field(k).and_then(|v| v.parse::<u64>().ok());
let sent = num("pending=").filter(|n| *n > 0)?;
if field("at=") != Some(crate::outcome::AT_END) {
return None;
}
let chars = num("chars=")? as usize;
let answered = match field("verdicts=") {
_ if failed => 0,
Some(v) => v
.split(',')
.filter_map(|p| p.split_once(':'))
.filter(|(word, _)| *word != crate::predict::Verdict::Unclear.word())
.filter_map(|(_, n)| n.parse::<u64>().ok())
.sum(),
None => num("resolved=").unwrap_or(0),
};
Some(Call { model: field("model=").unwrap_or_default().to_string(), chars, sent, answered: answered.min(sent) })
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct Tally {
pub sent: u64,
pub answered: u64,
pub calls: u64,
pub sent_sq: u64,
pub answered_sq: u64,
pub cross: u64,
}
impl Tally {
pub fn add(&mut self, sent: u64, answered: u64) {
self.sent += sent;
self.answered += answered;
self.calls += 1;
self.sent_sq += sent * sent;
self.answered_sq += answered * answered;
self.cross += sent * answered;
}
pub fn merge(&mut self, o: &Tally) {
self.sent += o.sent;
self.answered += o.answered;
self.calls += o.calls;
self.sent_sq += o.sent_sq;
self.answered_sq += o.answered_sq;
self.cross += o.cross;
}
#[must_use]
pub fn share(&self) -> f64 {
if self.sent == 0 { 0.0 } else { self.answered as f64 / self.sent as f64 }
}
#[must_use]
pub fn effective(&self) -> f64 {
let n = self.sent as f64;
let p = self.share();
if self.calls < 2 || p <= 0.0 || p >= 1.0 {
return n;
}
let spread = (self.answered_sq as f64 - 2.0 * p * self.cross as f64 + p * p * self.sent_sq as f64).max(0.0);
let c = self.calls as f64;
let var = c / (c - 1.0) * spread / (n * n);
if var <= 0.0 { n } else { (p * (1.0 - p) / var).min(n) }
}
#[must_use]
pub fn interval(&self) -> (f64, f64) {
wilson_of(self.share(), self.effective())
}
}
pub type Steps = BTreeMap<i32, Tally>;
#[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)]
struct Cache {
version: u32,
stamp: String,
models: BTreeMap<String, Vec<(i32, Tally)>>,
}
fn stamp(state: &Path) -> String {
let (mut files, mut bytes, mut newest) = (0u64, 0u64, 0u128);
for e in std::fs::read_dir(state).into_iter().flatten().flatten() {
if e.path().extension().is_some_and(|x| x == "cost") {
let Ok(m) = e.metadata() else { continue };
files += 1;
bytes += m.len();
let t = m.modified().ok().and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok()).map_or(0, |d| d.as_nanos());
newest = newest.max(t);
}
}
format!("{files}:{bytes}:{newest}")
}
fn count(state: &Path) -> BTreeMap<String, Steps> {
let mut out: BTreeMap<String, Steps> = BTreeMap::new();
for e in std::fs::read_dir(state).into_iter().flatten().flatten() {
if e.path().extension().is_none_or(|x| x != "cost") {
continue;
}
for c in std::fs::read_to_string(e.path()).unwrap_or_default().lines().filter_map(parse) {
out.entry(c.model).or_default().entry(step_of(c.chars)).or_default().add(c.sent, c.answered);
}
}
out
}
#[must_use]
pub fn counts(state: &Path) -> BTreeMap<String, Steps> {
let path = state.join(CACHE);
let now = stamp(state);
if let Some(c) = std::fs::read_to_string(&path).ok().and_then(|t| serde_json::from_str::<Cache>(&t).ok())
&& c.version == VERSION
&& c.stamp == now
{
return c.models.into_iter().map(|(m, v)| (m, v.into_iter().collect())).collect();
}
let counted = count(state);
let cache = Cache { version: VERSION, stamp: now, models: counted.iter().map(|(m, s)| (m.clone(), s.iter().map(|(j, t)| (*j, *t)).collect())).collect() };
if let Ok(text) = serde_json::to_string(&cache) {
static WRITER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
let tmp = state.join(format!("{CACHE}.{}.{}", std::process::id(), WRITER.fetch_add(1, std::sync::atomic::Ordering::Relaxed)));
if std::fs::write(&tmp, text).is_ok() && std::fs::rename(&tmp, &path).is_err() {
let _ = std::fs::remove_file(&tmp);
}
}
counted
}
#[must_use]
pub fn wilson(k: u64, n: u64) -> (f64, f64) {
wilson_of(if n == 0 { 0.0 } else { k as f64 / n as f64 }, n as f64)
}
#[must_use]
pub fn wilson_of(p: f64, n: f64) -> (f64, f64) {
if n <= 0.0 {
return (0.0, 1.0);
}
let z2 = Z95 * Z95;
let centre = (p + z2 / (2.0 * n)) / (1.0 + z2 / n);
let half = Z95 * (p * (1.0 - p) / n + z2 / (4.0 * n * n)).sqrt() / (1.0 + z2 / n);
((centre - half).max(0.0), (centre + half).min(1.0))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum Decision {
Reference,
Below,
Holds,
Fails,
Open,
}
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct Step {
pub step: i32,
pub upto: usize,
pub sent: u64,
pub answered: u64,
pub calls: u64,
pub lo: f64,
pub hi: f64,
pub decision: Decision,
}
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct Length {
pub model: String,
pub chars: usize,
pub start: bool,
pub line: Option<f64>,
pub steps: Vec<Step>,
pub explore: Option<usize>,
pub psi: f64,
}
impl Length {
#[must_use]
pub fn of(model: &str, steps: &Steps) -> Length {
let mut pooled = Tally::default();
let reference = steps.iter().find_map(|(j, t)| {
pooled.merge(t);
(pooled.sent >= REFERENCE_LESSONS && pooled.calls >= MIN_CALLS).then_some((*j, pooled.share()))
});
let line = reference.map(|(_, share)| HOLD * share);
let decide = |j: i32| -> Decision {
let Some((r, _)) = reference else { return Decision::Open };
match j.cmp(&r) {
std::cmp::Ordering::Less => return Decision::Below,
std::cmp::Ordering::Equal => return Decision::Reference,
std::cmp::Ordering::Greater => {}
}
let t = steps.get(&j).copied().unwrap_or_default();
if t.calls < MIN_CALLS {
return Decision::Open;
}
let (lo, hi) = t.interval();
let l = line.unwrap_or(0.0);
if lo >= l {
Decision::Holds
} else if hi < l {
Decision::Fails
} else {
Decision::Open
}
};
let holds = |j: i32| matches!(decide(j), Decision::Holds | Decision::Reference | Decision::Below);
let mut j = 0;
while decide(j) == Decision::Fails {
j -= 1;
}
if j == 0 {
while holds(j + 1) {
j += 1;
}
}
let (explore, psi) = match decide(j + 1) {
Decision::Open => (Some(edge(j + 1)), EXPLORE_SHARE),
Decision::Fails => (Some(edge(j + 1)), FAILED_EXPLORE_SHARE),
_ => (None, 0.0),
};
let mut shown: Vec<Step> = steps
.iter()
.map(|(&k, t)| {
let (lo, hi) = t.interval();
Step { step: k, upto: edge(k), sent: t.sent, answered: t.answered, calls: t.calls, lo, hi, decision: decide(k) }
})
.collect();
if explore.is_some() && !steps.contains_key(&(j + 1)) {
shown.push(Step { step: j + 1, upto: edge(j + 1), sent: 0, answered: 0, calls: 0, lo: 0.0, hi: 1.0, decision: Decision::Open });
}
Length { model: model.to_string(), chars: edge(j), start: j == 0, line, steps: shown, explore, psi }
}
}
fn of_model(all: &BTreeMap<String, Steps>, model: &str, unnamed: &str) -> Steps {
let mut out = all.get(model).cloned().unwrap_or_default();
if !unnamed.is_empty() && model == unnamed {
for (j, t) in all.get("").into_iter().flatten() {
out.entry(*j).or_default().merge(t);
}
}
out
}
#[must_use]
pub fn length(state: &Path, model: &str) -> Length {
let all = counts(state);
let unnamed = if all.contains_key("") { crate::hook::judge_model("") } else { String::new() };
Length::of(model, &of_model(&all, model, &unnamed))
}
#[must_use]
pub fn lengths(state: &Path, current: &str) -> Vec<Length> {
let all = counts(state);
let unnamed = crate::hook::judge_model("");
let mut models: std::collections::BTreeSet<String> = all.keys().filter(|m| !m.is_empty()).cloned().collect();
if all.contains_key("") && !unnamed.is_empty() {
models.insert(unnamed.clone());
}
if !current.is_empty() {
models.insert(current.to_string());
}
models.into_iter().map(|m| Length::of(&m, &of_model(&all, &m, &unnamed))).collect()
}
#[cfg(test)]
mod tests {
use super::*;
fn steps(rows: &[(usize, u64, u64)]) -> Steps {
rows.iter()
.map(|&(upto, n, a)| {
let c = n.min(10);
let mut t = Tally::default();
for i in 0..c {
t.add(n / c + u64::from(i < n % c), a / c + u64::from(i < a % c));
}
(step_of(upto), t)
})
.collect()
}
fn calls(rows: &[(u64, u64)]) -> Tally {
let mut t = Tally::default();
for &(n, a) in rows {
t.add(n, a);
}
t
}
#[test]
fn steps_are_half_again_apart_around_the_start_budget() {
assert_eq!((edge(-2), edge(-1), edge(0), edge(1), edge(2)), (32_000, 48_000, 72_000, 108_000, 162_000));
assert_eq!((step_of(72_000), step_of(72_001), step_of(48_001), step_of(48_000), step_of(1)), (0, 1, 0, -1, step_of(1)));
assert!(step_of(1) < -10 && edge(step_of(1)) >= 1);
for c in [1, 999, 31_999, 32_000, 100_000, 2_000_000] {
let j = step_of(c);
assert!(c <= edge(j) && c > edge(j - 1), "{c}: {j}");
}
}
#[test]
fn the_wilson_interval_matches_known_values() {
let (lo, hi) = wilson(50, 100);
assert!((lo - 0.4038).abs() < 1e-3 && (hi - 0.5962).abs() < 1e-3, "{lo} {hi}");
let (lo, hi) = wilson(0, 10);
assert!(lo == 0.0 && (hi - 0.2775).abs() < 1e-3, "{hi}");
assert_eq!(wilson(0, 0), (0.0, 1.0));
}
#[test]
fn every_cost_line_format_gives_what_it_holds() {
let check = "judge call=3 pending=4 resolved=2 chars=900 returned=a,b,c unmatched= verdicts=changed:1,ignored:1,not_arisen:0,consequence_remains:1,unclear:1 reason_chars=40 model=claude:sonnet prompt=30000";
assert_eq!(parse(check), None, "a check's line: another task");
let end = "judge call=5 pending=10 resolved=6 chars=70000 returned= unmatched= verdicts=changed:6,ignored:0,not_arisen:0,consequence_remains:0,unclear:4 reason_chars=9 model=codex:gpt-5 at=end batch=1/3 budget=72000 cut=after";
assert_eq!(parse(end), Some(Call { model: "codex:gpt-5".into(), chars: 70_000, sent: 10, answered: 6 }));
let end_075 = "judge call=2 pending=7 resolved=1 chars=40000 returned= unmatched= verdicts=changed:1,ignored:0,not_arisen:0,consequence_remains:2,unclear:4 reason_chars=9 at=end";
assert_eq!(parse(end_075), Some(Call { model: String::new(), chars: 40_000, sent: 7, answered: 3 }));
let failed = "judge-failed pending=12 chars=71000 model=claude:sonnet at=end batch=2/3 budget=72000 cut=after";
assert_eq!(parse(failed), Some(Call { model: "claude:sonnet".into(), chars: 71_000, sent: 12, answered: 0 }));
assert_eq!(parse("judge call=2 pending=3 resolved=2 chars=4100"), None, "a check's line before 0.7.6");
assert_eq!(parse("judge call=2 pending=3 resolved=2 chars=4100 at=end").map(|c| c.answered), Some(2), "before 0.7.5: resolved=");
assert_eq!(parse("judge call=2 pending=0 resolved=0 chars=4100 at=end"), None);
assert_eq!(parse("block questions=2 chars=600"), None);
}
#[test]
fn one_call_never_decides_a_step_and_disagreeing_calls_count_for_less() {
let mut base = steps(&[(4_000, 59, 15)]);
base.insert(step_of(100_000), calls(&[(2, 2)]));
let up = Length::of("m", &base);
assert_eq!((up.chars, up.start, up.explore), (72_000, true, Some(108_000)), "{up:?}");
let mut base = steps(&[(4_000, 59, 15)]);
base.insert(step_of(70_000), calls(&[(30, 0)]));
let down = Length::of("m", &base);
assert_eq!((down.chars, down.start, down.psi), (72_000, true, EXPLORE_SHARE), "{down:?}");
assert_eq!(down.steps.iter().find(|s| s.upto == 72_000).map(|s| s.decision), Some(Decision::Open));
let reference = (32_000, 400, 200);
let clustered = calls(&[(30, 30), (30, 0), (30, 0), (30, 24), (30, 0), (30, 0)]);
assert!(wilson(clustered.answered, clustered.sent).1 < HOLD * 0.5, "independent: would fail");
assert!(clustered.effective() < 20.0, "{}", clustered.effective());
let mut st = steps(&[reference, (72_000, 300, 150)]);
st.insert(step_of(100_000), clustered);
let l = Length::of("m", &st);
assert_eq!(l.steps.iter().find(|s| s.upto == 108_000).map(|s| s.decision), Some(Decision::Open), "{l:?}");
let mut st = steps(&[reference, (72_000, 300, 150)]);
st.insert(step_of(100_000), calls(&[(30, 9); 6]));
assert_eq!(Length::of("m", &st).steps.iter().find(|s| s.upto == 108_000).map(|s| s.decision), Some(Decision::Fails));
}
#[test]
fn the_reference_pools_the_smallest_steps() {
let l = Length::of("m", &steps(&[(10_000, 20, 10), (21_333, 20, 10), (72_000, 300, 150)]));
assert_eq!(l.steps.iter().filter(|s| s.sent > 0).map(|s| s.decision).collect::<Vec<_>>(), [Decision::Below, Decision::Reference, Decision::Holds]);
assert!((l.line.unwrap() - HOLD * 0.5).abs() < 1e-9);
let mut few = Steps::new();
few.insert(step_of(10_000), calls(&[(40, 20), (40, 20)]));
assert_eq!(Length::of("m", &few).line, None);
}
#[test]
fn thin_evidence_keeps_the_start_and_explores_the_step_above() {
let none = Length::of("m", &Steps::new());
assert_eq!((none.chars, none.start, none.explore, none.psi, none.line), (72_000, true, Some(108_000), EXPLORE_SHARE, None));
assert_eq!(none.steps.len(), 1, "the explored step shows: {:?}", none.steps);
let thin = Length::of("m", &steps(&[(32_000, 40, 20), (108_000, 5, 2)]));
assert_eq!((thin.chars, thin.start, thin.explore), (72_000, true, Some(108_000)));
assert_eq!(thin.steps.iter().map(|s| s.decision).collect::<Vec<_>>(), [Decision::Reference, Decision::Open]);
let small = Length::of("m", &steps(&[(32_000, 12, 12), (108_000, 10, 1)]));
assert_eq!((small.chars, small.line, small.explore), (72_000, None, Some(108_000)));
}
#[test]
fn flat_logs_step_the_budget_up() {
let flat = Length::of("m", &steps(&[(32_000, 400, 200), (72_000, 300, 150), (108_000, 300, 151), (162_000, 300, 149), (243_000, 3, 1)]));
assert_eq!((flat.chars, flat.start, flat.explore, flat.psi), (162_000, false, Some(243_000), EXPLORE_SHARE));
let d: Vec<Decision> = flat.steps.iter().map(|s| s.decision).collect();
assert_eq!(d, [Decision::Reference, Decision::Holds, Decision::Holds, Decision::Holds, Decision::Open]);
let capped = Length::of("m", &steps(&[(32_000, 400, 200), (108_000, 300, 150), (162_000, 400, 100)]));
assert_eq!((capped.chars, capped.explore, capped.psi), (108_000, Some(162_000), FAILED_EXPLORE_SHARE));
}
#[test]
fn degrading_logs_step_the_budget_down() {
let worse = Length::of("m", &steps(&[(21_333, 400, 300), (32_000, 300, 250), (48_000, 300, 230), (72_000, 300, 120)]));
assert_eq!((worse.chars, worse.start), (48_000, false), "{worse:?}");
assert_eq!((worse.explore, worse.psi), (Some(72_000), FAILED_EXPLORE_SHARE), "72k failed: explored at the floor, so it can hold again");
let d: Vec<Decision> = worse.steps.iter().map(|s| s.decision).collect();
assert_eq!(d, [Decision::Reference, Decision::Holds, Decision::Holds, Decision::Fails]);
let open = Length::of("m", &steps(&[(21_333, 400, 300), (32_000, 300, 250), (48_000, 300, 180), (72_000, 300, 120)]));
assert_eq!((open.chars, open.explore, open.psi), (48_000, Some(72_000), FAILED_EXPLORE_SHARE), "{open:?}");
assert_eq!(open.steps[2].decision, Decision::Open);
let deep = Length::of("m", &steps(&[(21_333, 400, 300), (32_000, 300, 250), (48_000, 300, 120), (72_000, 300, 120)]));
assert_eq!(deep.chars, 32_000);
let bottom = Length::of("m", &steps(&[(48_000, 300, 270), (72_000, 300, 30)]));
assert_eq!(bottom.chars, 48_000);
}
#[test]
fn the_counts_are_cached_until_the_cost_lines_change() {
let d = std::env::temp_dir().join(format!("borg-judge-length-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&d);
std::fs::create_dir_all(&d).unwrap();
let end = |n: u64, a: u64, model: &str| {
format!(
"judge call=1 pending={n} resolved={a} chars=60000 returned= unmatched= verdicts=changed:{a},ignored:0,not_arisen:0,consequence_remains:0,unclear:{} reason_chars=1{model} at=end\n",
n - a
)
};
std::fs::write(d.join("a.cost"), end(10, 4, " model=x:y") + &end(6, 6, "") + "block questions=1 chars=3\n").unwrap();
let c = counts(&d);
assert_eq!((c["x:y"][&0].sent, c["x:y"][&0].answered, c["x:y"][&0].calls), (10, 4, 1));
assert_eq!((c[""][&0].sent, c[""][&0].answered), (6, 6));
let cached: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(d.join(CACHE)).unwrap()).unwrap();
assert_eq!(cached["version"], VERSION);
let all = counts(&d);
let pair = |t: &Tally| (t.sent, t.answered, t.calls);
assert_eq!(pair(&of_model(&all, "x:y", "x:y")[&0]), (16, 10, 2));
assert_eq!(pair(&of_model(&all, "x:y", "z:w")[&0]), (10, 4, 1));
std::fs::write(d.join("b.cost"), end(5, 5, " model=x:y")).unwrap();
assert_eq!(pair(&counts(&d)["x:y"][&0]), (15, 9, 2));
let mut stale: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(d.join(CACHE)).unwrap()).unwrap();
stale["version"] = serde_json::json!(VERSION + 1);
stale["models"]["x:y"][0][1]["sent"] = serde_json::json!(1);
std::fs::write(d.join(CACHE), stale.to_string()).unwrap();
assert_eq!(pair(&counts(&d)["x:y"][&0]), (15, 9, 2));
std::fs::write(d.join("c.cost"), end(5, 1, " model=x:y")).unwrap();
std::thread::scope(|s| {
for _ in 0..8 {
s.spawn(|| assert_eq!(pair(&counts(&d)["x:y"][&0]), (20, 10, 3)));
}
});
let cache: Cache = serde_json::from_str(&std::fs::read_to_string(d.join(CACHE)).unwrap()).unwrap();
assert_eq!(cache.version, VERSION);
assert!(std::fs::read_dir(&d).unwrap().all(|e| !e.unwrap().file_name().to_string_lossy().starts_with(&format!("{CACHE}."))), "no temp file left");
assert!(lengths(&d, "q:r").iter().any(|l| l.model == "q:r" && l.start));
let _ = std::fs::remove_dir_all(&d);
}
}