hanzoai/dev
2
0
Fork 0
mirror of https://github.com/hanzoai/dev.git synced 2026-10-17 15:34:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

kai: the first call begins beside the search, and the tests around the failing ones wait for nothing

Browse Source 
The first state's first call is begun as soon as the state is ranked, against its revision. The
search settling the task, or a kept step moving the state, cancels it; a cancelled call closes its
stream and leaves an estimated line in the spend file. Past its cap the search goes on while that
call has not answered, to 8 s, and a visible test too slow for the cap is narrowed to the failing
tests it names instead of the search being skipped. The tests around the failing ones run beside
the visible test, and what failed among them before any edit is read only when they fail after
it. A test run in the tree leaves no file git neither tracks nor ignores, and Hypothesis keeps its
examples in the run's own temporary directory.

Model-free on 60 dev-split tasks on evo, A/B interleaved: search-settled tasks 1.74x faster (95%
1.64-1.85, 13 of 13), solves 13 to 19: five exact pyrsistent fixes no longer fail on the
.hypothesis files the runs wrote, and a lark task with an 18 s test file is searched.
This commit is contained in:
zeekay
2026-10-17 08:34:28 -07:00
parent 0934258016
commit eb83f6a1bf
8 changed files with 863 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files
LLM.md
+27 -1
View File
@@ -173,6 +173,14 @@ and the tools they run. Search, saved actions and the shell classifier are hanzo
transition is a ledger line with its reason (`fast`, `fallback`); no step waits on Kai when an
action already applies. Each program has its own outage cooldown: one capability's timeout never
disables another.
- **Nothing waits that need not.** The first state's first call is begun the moment the state is
ranked, beside the search (`flow::Ahead`, `zen::Zen::start`), against that state's revision. The
search settling the task, or a kept step moving the state, cancels it; when the search settles
nothing the call has had the search's time to answer, and the ledger's `generate` says how long
(`ahead_ms`). A cancelled call closes its stream (the server stops writing it) and leaves a
`cancelled` line in the spend file, its prompt counted at three characters a token at the ceiling
rate; the ledger has a `cancel` line, never a `generate` one. An answer asked ahead is tried
before any new call at its revision, and with one queued Kai is not asked `fix.turn@1`.
- **Revision-bound.** Preconditions are checked immediately before execution; an edit invalidates
every pending decision made against the old revision. A revision is the content of the files the
run touched, and a revert returns to `base`. A cached result is keyed on what it depends on
@@ -188,7 +196,25 @@ and the tools they run. Search, saved actions and the shell classifier are hanzo
- **Fast path.** The action `mutate` tries generic token transforms (comparison and operator flips,
constants off by one, arguments swapped, a condition inverted, a `return` dropped) at the best
sites against the one visible test command, eight at once, capped in time; it names no
repository, file or task. It solves single faults of that kind and nothing else.
repository, file or task. It solves single faults of that kind and nothing else. Past its 2.5 s
cap it goes on while the call asked beside it has not answered, to 8 s (`mutate::LONGEST`): that
wait is spent either way, and on a loaded host the cap cut 34 of 39 unsettled dev-split searches
short of their candidates, the fix in the top four sites in 9 of them. A visible test too slow for
the cap (its time x 1.5 > the cap) is narrowed to the failing tests it names (`mutate::screen`)
instead of the search being skipped: lark's 18 s file runs its failing tests in 1.7 s, and the
candidate picked is verified with the whole command in the tree as ever. A fast file narrowed
ran 5% faster (two measurements, 13 and 19 tasks), not worth a second command there.
- **The tests around the failing ones** (`wide`) run beside the visible test after an edit, and a
failing visible test stops them. What failed among them before any edit is read only when they
fail after it, with the files the run touched put back as they were for that read and then as
they are (`flow::before`, ledger `lazy`): a run that passes shows nothing new, so on most tasks
it is never read, and it never competes with the search for the CPU.
- **A test run leaves nothing.** Files a run in the tree writes that git neither tracks nor ignores
and that were not there before it are removed after it, with the directories they emptied
(`tools::Litter`, ledger `sweep`): they are the run's, never the edit's. What `.gitignore` names
is left alone, and so is a tree outside git. Hypothesis keeps its example database in the run's
own temporary directory (`HYPOTHESIS_STORAGE_DIRECTORY`, unless the caller set it), so no run
replays another's examples and none lands in the tree.
- **Jobs** (not built): long work runs in the background and reports events into the same ledger.
Start, progress, completion and verified outcome are four states; a finished process is not a
finished task.
crates/hanzo-kai/src/edit.rs
+5
View File
@@ -196,6 +196,11 @@ impl Applied {
pub fn files(&self) -> Vec<String> {
self.before.keys().cloned().collect()
}
/// Each file changed with its text before the change.
pub fn texts(&self) -> &BTreeMap<String, String> {
&self.before
}
}
fn safe(file: &str) -> bool {
crates/hanzo-kai/src/flow.rs
+271 -73
View File
@@ -35,6 +35,15 @@
//! the ledger and a line of the trace's `generations`. Neither is an attempt: the tests judge the
//! edit that results, as for any other.
//!
//! Nothing waits that need not. The tests around the failing ones run beside the failing command
//! after an edit, and what failed among them before any edit is read only when they fail then
//! ([`crate::wide`]). The first attempt's call is begun as soon as the first state is ranked,
//! beside the search, against that state's revision: when the search settles the task the call is
//! cancelled, and when it does not the call has had the search's time to answer. An answer asked
//! ahead answers an attempt only at the revision it was asked against; a state the loop leaves
//! cancels it. What a test run in the tree writes beside it is removed after it
//! ([`tools::Litter`]).
//!
//! Every step is timed, and the trace carries the times with Kai's answers.
use crate::decide;
@@ -73,6 +82,7 @@ use hanzo_config::kai::Mode;
use indexmap::IndexMap;
use serde_json::Value;
use serde_json::json;
use std::collections::BTreeMap;
use std::collections::HashSet;
use std::path::Path;
use std::path::PathBuf;
@@ -103,7 +113,8 @@ const SHOWN: u32 = 100;
const TEST: u32 = 30;
/// Characters of the failing output shown.
const OUTPUT: usize = 2000;
/// How long the search of candidate edits may run before the model writes the edit.
/// How long the search of candidate edits may run before the model writes the edit; past it, the
/// search goes on only while the call asked beside it has not answered ([`mutate::LONGEST`]).
const SEARCH: Duration = Duration::from_millis(2500);
/// A line of code is shown to this many characters.
const WIDE: usize = 300;
@@ -785,22 +796,160 @@ fn discover(finder: &dyn Finder, root: &Path, plan: &mut Plan, ledger: &Ledger,
(added, a.ms)
}
/// The tests around the failing ones, after an edit that passed the failing command: the run and
/// the tests failing that did not before. Recorded as an action of its own. `None` without a
/// wider command.
fn widen(wide: &Option<Wide>, root: &Path, ledger: &Ledger, revision: &str) -> Option<wide::Verdict> {
let w = wide.as_ref()?;
/// A model call begun before the loop came to it, and what it was asked against: the revision,
/// the functions shown and the digests of the files the model saw. It answers an attempt at that
/// revision and no other.
struct Ahead {
revision: String,
ask: Ask,
pending: crate::zen::Pending,
}
/// What an attempt asks: the model, the effort, the prompt, and what the prompt shows.
#[derive(Clone)]
struct Ask {
model: String,
effort: Effort,
/// The escalation model's.
escalated: bool,
user: String,
shown: Vec<String>,
saw: Vec<(String, String)>,
bundle_ms: f64,
}
/// What attempt `n` asks over the state `s`, after `last` and with `history` tried.
#[allow(clippy::too_many_arguments)]
fn ask(s: &mut Seen, setup: &Setup, n: u32, strategy: Strategy, stronger: bool, command: &str, goal: &str, last: Option<&Attempt>, history: &[String]) -> Ask {
// After an escalation the escalation model writes, at full effort over a wider view.
let escalation = stronger.then(|| setup.zen.escalation()).flatten();
let model = escalation.unwrap_or_else(|| setup.zen.model()).to_string();
let wider = strategy == Strategy::Full || escalation.is_some();
let effort = if wider { Effort::Full } else { s.order.effort };
let begun = Instant::now();
let (code, shown, saw) = excerpts(
&s.plan,
&mut s.tree,
&s.order.by_rank,
LINES * n.min(3) * if wider { 2 } else { 1 },
if n == 1 { COVER } else { 0.999 },
);
let user = prompt(command, goal, &s.shown, &code, last, history);
Ask { model, effort, escalated: escalation.is_some(), user, shown, saw, bundle_ms: ms(begun) }
}
/// The first answer in among those asked ahead at `revision`, waiting for one when none is in
/// yet; `None` when none was asked at it.
fn take(ahead: &mut Vec<Ahead>, revision: &str) -> Option<Ahead> {
if !ahead.iter().any(|a| a.revision == revision) {
return None;
}
loop {
if let Some(i) = ahead.iter().position(|a| a.revision == revision && a.pending.done()) {
return Some(ahead.remove(i));
}
std::thread::sleep(Duration::from_millis(5));
}
}
/// Cancels the calls asked ahead that `drop` says go, each a ledger line with why.
fn cancel(ahead: &mut Vec<Ahead>, ledger: &Ledger, why: &str, drop: impl Fn(&Ahead) -> bool) {
let (gone, kept): (Vec<Ahead>, Vec<Ahead>) = std::mem::take(ahead).into_iter().partition(|a| drop(a));
*ahead = kept;
for a in gone {
let age = a.pending.age();
let running = a.pending.cancel();
ledger.write(
"cancel",
json!({"id": ledger.next(), "revision": a.revision, "model": a.ask.model, "escalated": a.ask.escalated,
"running": running, "ms": age, "reason": why}),
);
}
}
/// The failing command and the tests around it, at once, on the tree at `revision`: the command's
/// run, and the wider verdict when the command passed. A failing command stops the wider run, the
/// edit being settled without it. What the wider run failed on before any edit is read only when
/// it fails now ([`before`]). An error when the tree cannot be put back as it was after that read.
fn check(
root: &Path,
command: &str,
limit: Duration,
wide: Option<&mut Wide>,
pristine: &BTreeMap<String, String>,
ledger: &Ledger,
revision: &str,
) -> Result<(Value, Option<wide::Verdict>), String> {
let litter = tools::Litter::before(root);
let Some(w) = wide else {
let test = tools::test(root, command, limit);
swept(root, &litter, ledger, revision);
return Ok((test, None));
};
let failed = std::sync::atomic::AtomicBool::new(false);
let shared: &Wide = w;
let (test, run) = std::thread::scope(|s| {
let around = s.spawn(|| shared.run(root, &|| failed.load(std::sync::atomic::Ordering::SeqCst)));
let test = tools::test(root, command, limit);
if test["passed"] != true {
failed.store(true, std::sync::atomic::Ordering::SeqCst);
}
(test, around.join().ok())
});
swept(root, &litter, ledger, revision);
let Some(run) = run.filter(|_| test["passed"] == true) else {
return Ok((test, None));
};
if Wide::needs_before(&run) && w.before.is_none() {
w.before = Some(before(root, w, pristine, ledger)?);
}
let mut a = ledger.action(revision, Operation::Test { command: w.command.clone() }, Source::Rule, "the tests around the failing ones, before done");
a.preconditions = vec![Precondition::Revision { revision: revision.into() }];
a.expected = "no test fails that did not fail before the edit".into();
a.advance(Status::Attempted);
let v = w.check(root);
let v = w.verdict(run);
a.ms = v.run["ms"].as_f64().unwrap_or(0.0);
a.advance(Status::Completed);
a.advance(if v.new.is_empty() { Status::Verified } else { Status::Rejected });
a.result = json!({"new": v.new, "timed_out": v.timed_out, "code": v.run["code"]});
a.result = json!({"new": v.new, "timed_out": v.timed_out, "code": v.run["code"], "beside": true});
ledger.record(&a);
Some(v)
Ok((test, Some(v)))
}
/// What a test run in the tree wrote beside it, removed: its files are not the edit's. A ledger
/// line names them.
fn swept(root: &Path, litter: &tools::Litter, ledger: &Ledger, revision: &str) {
let gone = litter.sweep(root);
if !gone.is_empty() {
ledger.write("sweep", json!({"id": ledger.next(), "revision": revision, "removed": gone.len(), "files": gone.iter().take(20).collect::<Vec<_>>()}));
}
}
/// What the tests around the failing ones fail on before any edit: the files the run touched are
/// put back as they were for the run, then as they are. An error when they cannot be put back.
fn before(root: &Path, w: &Wide, pristine: &BTreeMap<String, String>, ledger: &Ledger) -> Result<wide::Before, String> {
let mut a = ledger.action("base", Operation::Test { command: w.command.clone() }, Source::Rule, "the tests around the failing ones, before any edit: what already fails there");
a.advance(Status::Attempted);
let put = |files: &[(&String, &[u8])]| files.iter().all(|(f, text)| std::fs::write(root.join(f), text).is_ok());
let now: Vec<(&String, Vec<u8>)> = pristine.keys().map(|f| (f, std::fs::read(root.join(f)).unwrap_or_default())).collect();
let now: Vec<(&String, &[u8])> = now.iter().map(|(f, b)| (*f, b.as_slice())).collect();
let then: Vec<(&String, &[u8])> = pristine.iter().map(|(f, t)| (f, t.as_bytes())).collect();
if !put(&then) {
let _ = put(&now);
return Err("the files cannot be put back as they were to read the tests around the failing ones".into());
}
let litter = tools::Litter::before(root);
let run = w.run(root, &|| false);
swept(root, &litter, ledger, "base");
if !put(&now) {
return Err("the files cannot be put back after reading the tests around the failing ones, so the working tree lacks the edit; see `git diff`".into());
}
let read = Wide::read(&run);
a.ms = run["ms"].as_f64().unwrap_or(0.0);
a.advance(Status::Completed);
a.result = json!({"code": read.code, "failing": read.failing, "lazy": true});
ledger.record(&a);
Ok(read)
}
/// What the choice of site, action and effort is made from.
@@ -948,7 +1097,7 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
// tree; the failures give recover its memory.
let mut touched: Vec<String> = Vec::new();
// Each touched file's digest when the run began, so a revert is the revision it left.
let mut origin: std::collections::BTreeMap<String, String> = std::collections::BTreeMap::new();
let mut origin: BTreeMap<String, String> = std::collections::BTreeMap::new();
let mut kept: Vec<edit::Applied> = Vec::new();
let mut kept_summary: Vec<String> = Vec::new();
let mut failures: Vec<Failure> = Vec::new();
@@ -967,7 +1116,9 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
(given.clone(), false, 1)
}
None => {
let litter = tools::Litter::before(root);
let run = tools::test(root, command, options.limit);
swept(root, &litter, &ledger, "base");
(
run["output"].as_str().unwrap_or_default().to_string(),
run["passed"] == true,
@@ -992,18 +1143,12 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
let errors0 = diagnostic::errors(&output);
report.failing = failing;
let test_ms = steps["test0"];
// The tests around the failing ones, once before any edit: what fails there already.
let begun = Instant::now();
let wide = Wide::before(root, command, &output);
if let Some(w) = &wide {
let mut a = ledger.action("base", Operation::Test { command: w.command.clone() }, Source::Rule, "the tests around the failing ones, before any edit: what already fails there");
a.advance(Status::Attempted);
a.advance(Status::Completed);
a.ms = ms(begun);
a.result = json!({"code": w.before_code, "failing": w.before});
ledger.record(&a);
steps.insert("wide0".into(), a.ms);
}
// The tests around the failing ones, and the files the run touched as they were: what failed
// there before any edit is read from them only when a run after an edit fails.
let mut wide = Wide::new(command, &output);
let mut pristine: BTreeMap<String, String> = std::collections::BTreeMap::new();
// Calls begun before the loop came to them.
let mut ahead: Vec<Ahead> = Vec::new();
let mut seen: Option<Seen> = None;
let mut look = true;
@@ -1069,15 +1214,30 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
let cx = Context { goal, revision: &revision, fingerprint: &fingerprint, failing, previous };
let order = choose(&plan, &mut tree, kai, setup, &ledger, &cx, &mut report.missing);
*steps.entry("site".into()).or_insert(0.0) += ms(begun);
seen = Some(Seen { tree, plan, order, shown: tail(&output, OUTPUT), fingerprint });
// A call asked at a state the loop has left answers nothing.
cancel(&mut ahead, &ledger, "the state it was asked at is left", |a| a.revision != revision);
let mut s = Seen { tree, plan, order, shown: tail(&output, OUTPUT), fingerprint };
// The first state: its first attempt's call is begun now, beside the search, so a search
// that settles nothing has cost the call no time.
if n == 0 && kept.is_empty() && ahead.is_empty() {
let begun = Instant::now();
let first = ask(&mut s, setup, 1, Strategy::Retry, false, command, goal, None, &history);
let pending = setup.zen.start(&first.model, SYSTEM, &first.user, first.effort);
ahead.push(Ahead { revision: revision.clone(), ask: first, pending });
*steps.entry("ahead".into()).or_insert(0.0) += ms(begun);
}
seen = Some(s);
look = false;
}
let Some(s) = seen.as_mut() else { break };
let fingerprint = s.fingerprint.clone();
// The cheap action first: candidate edits at the best sites against the visible test.
// A test that takes longer than the cap holds cannot be tried candidate after candidate.
let slow = diagnostic::seconds(&output).is_some_and(|t| t * 1.5 > SEARCH.as_secs_f64());
// The cheap action first: candidate edits at the best sites against the visible test. A
// test that takes longer than the cap holds cannot be tried candidate after candidate:
// its failing tests alone are, where they are named, and else the search is skipped.
let long = diagnostic::seconds(&output).is_some_and(|t| t * 1.5 > SEARCH.as_secs_f64());
let screen = long.then(|| mutate::screen(root, command, &output)).flatten();
let slow = long && screen.is_none();
if slow && options.search && s.order.act == Act::Search && searched.insert(kept.len()) {
ledger.fallback(&revision, "search", "rule", "model", &format!("the visible test takes {:.1} s: longer than the search cap of {:.1} s holds", diagnostic::seconds(&output).unwrap_or(0.0), SEARCH.as_secs_f64()));
}
@@ -1098,7 +1258,12 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
let floor = if options.failure.is_some() { 10.0 } else { 3.0 };
let limit = Duration::from_secs_f64((test_ms / 1000.0 * 3.0 + 1.0).clamp(floor, options.limit.as_secs_f64().max(floor)));
let workers = std::thread::available_parallelism().map_or(2, usize::from);
let found = mutate::search(root, command, &candidates, failing, errors0, SEARCH, limit, workers);
// The search goes on past its cap while the call asked beside it has not answered.
let waiting = || {
let mut asked = ahead.iter().filter(|a| a.revision == revision).peekable();
asked.peek().is_some() && asked.all(|a| !a.pending.done())
};
let found = mutate::search(root, screen.as_deref().unwrap_or(command), &candidates, failing, errors0, SEARCH, &waiting, limit, workers);
m.ms = ms(begun);
fast.tried += found.tried;
fast.runs += found.runs;
@@ -1112,7 +1277,7 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
let pick = found.winner.or(found.progress);
m.advance(Status::Completed);
m.advance(if pick.is_some() { Status::Selected } else { Status::Rejected });
m.result = json!({"tried": found.tried, "runs": found.runs, "workers": found.workers, "winner": found.winner, "progress": found.progress, "remaining": found.remaining, "note": found.note, "copy_ms": found.copy});
m.result = json!({"tried": found.tried, "runs": found.runs, "workers": found.workers, "winner": found.winner, "progress": found.progress, "remaining": found.remaining, "note": found.note, "copy_ms": found.copy, "screen": screen});
ledger.record(&m);
if let Some(w) = pick {
let c = &candidates[w];
@@ -1129,6 +1294,9 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
for (f, d) in made.originals() {
origin.entry(f).or_insert(d);
}
for (f, t) in made.texts() {
pristine.entry(f.clone()).or_insert_with(|| t.clone());
}
for f in &files {
if !touched.contains(f) {
touched.push(f.clone());
@@ -1145,16 +1313,18 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
t.preconditions = vec![Precondition::Revision { revision: rev.clone() }];
t.expected = "the tests pass, or fewer fail".into();
t.advance(Status::Attempted);
// Passing the failing command is not done until the tests around it
// agree; both run at once.
let begun = Instant::now();
let test = tools::test(root, command, options.limit);
t.ms = ms(begun);
*steps.entry("verify".into()).or_insert(0.0) += t.ms;
let (test, around) = check(root, command, options.limit, wide.as_mut(), &pristine, &ledger, &rev)?;
let checked = ms(begun);
t.ms = test["ms"].as_f64().unwrap_or(checked);
*steps.entry("verify".into()).or_insert(0.0) += checked;
let mut out = test["output"].as_str().unwrap_or_default().to_string();
let mut passed = test["passed"] == true;
let mut test = test;
// Passing the failing command is not done until the tests around it agree.
let mut broke = None;
if passed && let Some(v) = widen(&wide, root, &ledger, &rev) {
if let Some(v) = around {
*steps.entry("wide".into()).or_insert(0.0) += v.run["ms"].as_f64().unwrap_or(0.0);
if !v.new.is_empty() {
broke = Some(format!("the candidate passed the failing command and broke {}", v.new.join(", ")));
@@ -1167,10 +1337,11 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
t.advance(Status::Completed);
let better = !passed && now >= 1 && now < failing && !diagnostic::broken(&out) && diagnostic::errors(&out) <= errors0;
t.advance(if passed || better { Status::Verified } else { Status::Rejected });
t.result = json!({"passed": passed, "code": test["code"], "failing": now});
t.result = json!({"passed": passed, "code": test["code"], "failing": now, "check_ms": checked});
ledger.record(&t);
ledger.verify(&t.id, &rev, passed, test["code"].as_i64().unwrap_or(-1), &ledger::fingerprint(&out), &tail(&out, 600), t.ms);
if passed {
cancel(&mut ahead, &ledger, "the search settled the task", |_| true);
fast.solved = true;
fast.kind = Some(c.kind.to_string());
fast.edit = made.summary.first().cloned();
@@ -1227,24 +1398,19 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
}
}
// The model writes the edit.
// The model writes the edit: an answer asked ahead at this revision before a new call.
n += 1;
// After an escalation the escalation model writes, at full effort over a wider view.
let escalation = stronger.then(|| setup.zen.escalation()).flatten();
let model = escalation.unwrap_or_else(|| setup.zen.model()).to_string();
let wider = strategy == Strategy::Full || escalation.is_some();
let effort = if wider { Effort::Full } else { s.order.effort };
let mut ms_of: IndexMap<String, f64> = IndexMap::new();
let begun = Instant::now();
let (code, shown, saw) = excerpts(
&s.plan,
&mut s.tree,
&s.order.by_rank,
LINES * n.min(3) * if wider { 2 } else { 1 },
if n == 1 { COVER } else { 0.999 },
);
ms_of.insert("bundle".into(), ms(begun));
let user = prompt(command, goal, &s.shown, &code, last.as_ref(), &history);
let mut early = take(&mut ahead, &revision);
let asked = match &early {
Some(a) => a.ask.clone(),
None => ask(s, setup, n, strategy, stronger, command, goal, last.as_ref(), &history),
};
ms_of.insert("bundle".into(), asked.bundle_ms);
let Ask { model, effort, escalated, user, shown, saw, .. } = asked;
// An attempt the escalation model writes makes the ones after it its own too.
stronger |= escalated;
let escalation = escalated.then_some(model.as_str());
let source = if n == 1 && kept.is_empty() { Source::Model } else { Source::Recover };
let mut attempt = Attempt {
n,
@@ -1277,7 +1443,12 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
model.clone(),
effort,
user.clone(),
if escalation.is_some() { "recover: the escalation model writes the attempt" } else if n == 1 { "the first edit" } else { "recover: another attempt" },
match (escalation.is_some(), n) {
(true, 1) => "the first edit, by the escalation model",
(true, _) => "recover: the escalation model writes the attempt",
(false, 1) => "the first edit",
(false, _) => "recover: another attempt",
},
"first",
);
let (g, e, after, model) = loop {
@@ -1287,15 +1458,22 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
g.expected = "one JSON edit that applies".into();
g.advance(Status::Attempted);
let begun = Instant::now();
let generated = setup.zen.generate(&model, SYSTEM, &user, effort);
// The call asked ahead answers the first stage: the loop waits only for what is left of it.
let (generated, ahead_ms) = match early.take() {
Some(a) => {
let age = a.pending.age();
(a.pending.wait(), Some(age))
}
None => (setup.zen.generate(&model, SYSTEM, &user, effort), None),
};
g.ms = ms(begun);
*ms_of.entry("generate".into()).or_insert(0.0) += g.ms;
let answer = match generated {
Ok(x) => {
g.result = json!({
"model": x.model, "served": x.served, "cached": x.cached, "escalated": escalation.is_some(), "stage": stage,
"first_ms": x.first, "total_ms": x.total, "prompt_tokens": x.prompt, "completion_tokens": x.completion,
"reasoning_tokens": x.reasoning, "finish": x.finish,
"first_ms": x.first, "last_ms": x.last, "total_ms": x.total, "prompt_tokens": x.prompt, "completion_tokens": x.completion,
"reasoning_tokens": x.reasoning, "finish": x.finish, "ahead_ms": ahead_ms,
});
let text = x.text.clone();
attempt.generations.push(x);
@@ -1304,7 +1482,7 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
}
Err(why) => {
g.fail(&why);
g.result = json!({"error": why, "stage": stage});
g.result = json!({"error": why, "stage": stage, "ahead_ms": ahead_ms});
Err(why)
}
};
@@ -1337,6 +1515,9 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
for (f, d) in made.originals() {
origin.entry(f).or_insert(d);
}
for (f, t) in made.texts() {
pristine.entry(f.clone()).or_insert_with(|| t.clone());
}
let rev = ledger::revision_over(root, &files.iter().chain(touched.iter()).cloned().collect::<Vec<_>>(), &origin);
// The same edit again is no new state: recover chooses another action.
if tried.contains(&rev) {
@@ -1369,6 +1550,11 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
let Err(why) = &after else {
break (g, e, after, model);
};
// An answer asked ahead at this revision is the next attempt: no call is made to say
// this one again.
if ahead.iter().any(|a| a.revision == revision) {
break (g, e, after, model);
}
asking = if let (true, Ok(text)) = (stage == "first" && edit::spelling(why), &answer) {
(
model.clone(),
@@ -1407,14 +1593,15 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
t.preconditions = vec![Precondition::Revision { revision: rev.clone() }];
t.expected = "the tests pass".into();
t.advance(Status::Attempted);
// Passing the failing command is not done until the tests around it agree; both run at
// once.
let begun = Instant::now();
test = tools::test(root, command, options.limit);
t.ms = ms(begun);
ms_of.insert("test".into(), t.ms);
// Passing the failing command is not done until the tests around it agree.
if test["passed"] == true
&& let Some(v) = widen(&wide, root, &ledger, rev)
{
let (run, around) = check(root, command, options.limit, wide.as_mut(), &pristine, &ledger, rev)?;
let checked = ms(begun);
test = run;
t.ms = test["ms"].as_f64().unwrap_or(checked);
ms_of.insert("test".into(), checked);
if let Some(v) = around {
ms_of.insert("wide".into(), v.run["ms"].as_f64().unwrap_or(0.0));
if !v.new.is_empty() {
test = v.run;
@@ -1427,17 +1614,33 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
after_fp = ledger::fingerprint(text);
t.advance(Status::Completed);
t.advance(if attempt.passed { Status::Verified } else { Status::Rejected });
t.result = json!({"passed": attempt.passed, "code": test["code"], "failing": now});
t.result = json!({"passed": attempt.passed, "code": test["code"], "failing": now, "check_ms": checked});
ledger.record(&t);
ledger.verify(&t.id, rev, attempt.passed, test["code"].as_i64().unwrap_or(-1), &after_fp, &attempt.after, t.ms);
tried.push(rev.clone());
}
// A run that passes needs no decision. A run that fails: Kai says what next, and reads
// the patch, in one request.
// Fewer tests failing is progress: the edit stays.
let progress = attempt.applied
&& !attempt.passed
&& now >= 1
&& now < failing
&& !diagnostic::broken(test["output"].as_str().unwrap_or_default())
&& diagnostic::errors(test["output"].as_str().unwrap_or_default()) <= errors0;
// An answer asked ahead at the revision a failed edit is undone to is tried next.
let queued = !progress && ahead.iter().any(|a| a.revision == revision);
// A run that passes needs no decision, nor one whose next attempt is already asked. A run
// that fails: Kai says what next, and reads the patch, in one request.
let mut turn: Option<Decision> = None;
let mut turn_action: Option<Action> = None;
if attempt.applied && !attempt.passed {
if attempt.applied && !attempt.passed && kai && queued {
ledger.write(
"fast",
json!({"id": ledger.next(), "revision": after_rev, "op": "turn", "source": "rule",
"reason": "an answer asked ahead at the revision the edit is undone to is the next attempt: no inference"}),
);
} else if attempt.applied && !attempt.passed {
if kai {
ledger.fallback(&after_rev, "turn", "rule", "kai", "the verifier failed: what next, and is the patch correct and grounded");
let begun = Instant::now();
@@ -1494,12 +1697,6 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
let kai_escalates = answer == Some("escalate") && !attempt.passed;
// recover: the failure's fingerprint and the state it happened in decide what is next.
let progress = attempt.applied
&& !attempt.passed
&& now >= 1
&& now < failing
&& !diagnostic::broken(test["output"].as_str().unwrap_or_default())
&& diagnostic::errors(test["output"].as_str().unwrap_or_default()) <= errors0;
let mut why = String::new();
let mut stuck = false;
if attempt.applied && !attempt.passed {
@@ -1638,6 +1835,7 @@ pub fn run(root: &Path, options: &Options, setup: &Setup) -> Result<Report, Stri
if report.outcome.is_empty() {
report.outcome = "escalate".into();
}
cancel(&mut ahead, &ledger, "the flow is over", |_| true);
// A flow that does not end done leaves the tree as it found it.
if report.outcome != "done" {
while let Some(made) = kept.pop() {
@@ -1680,7 +1878,7 @@ fn line(order: &Order, turn: Option<&Decision>, attempt: &Attempt, run: &str, wh
.generations
.iter()
.map(|g| {
json!({"model": g.model, "served": g.served, "cached": g.cached, "first_ms": g.first, "total_ms": g.total, "prompt_tokens": g.prompt,
json!({"model": g.model, "served": g.served, "cached": g.cached, "first_ms": g.first, "last_ms": g.last, "total_ms": g.total, "prompt_tokens": g.prompt,
"completion_tokens": g.completion, "reasoning_tokens": g.reasoning, "finish": g.finish})
})
.collect();
crates/hanzo-kai/src/flow_tests.rs
+82 -3
View File
@@ -153,14 +153,19 @@ async fn asked_zen(zen: &MockServer) -> Vec<String> {
}
#[tokio::test(flavor = "multi_thread")]
async fn a_single_fault_is_found_by_search_and_the_model_is_never_asked() {
async fn a_single_fault_is_found_by_search_and_the_call_begun_beside_it_is_cancelled() {
let repo = repo_with(SIMPLE);
let kai = honest().await;
let zen = testing::zen(|_| String::new()).await;
let report = run_flow(repo.path(), &kai.uri(), &zen, options(3, Control::Kai)).await;
assert_eq!(report.outcome, "done");
assert!(report.attempts.is_empty());
assert!(asked_zen(&zen).await.is_empty(), "no generation");
assert!(asked_zen(&zen).await.len() <= 1, "at most the call begun beside the search");
let led = ledger(repo.path());
assert!(ops(&led, "generate").is_empty(), "no answer was used: {led:?}");
let cancelled: Vec<&Value> = led.iter().filter(|l| l["kind"] == "cancel").collect();
assert_eq!(cancelled.len(), 1, "{led:?}");
assert_eq!((cancelled[0]["model"].as_str(), cancelled[0]["reason"].as_str()), (Some("zen-test"), Some("the search settled the task")));
let fast = report.fast.as_ref().unwrap();
assert!(fast.solved && fast.tried >= 1 && fast.runs >= 1, "{fast:?}");
assert_eq!(fast.kind.as_deref(), Some("operator"));
@@ -226,6 +231,32 @@ async fn the_model_writes_the_edit_when_search_finds_none_and_the_flow_is_done()
assert!(led.iter().any(|l| l["kind"] == "fallback" && l["op"] == "search"));
}
#[tokio::test(flavor = "multi_thread")]
async fn the_first_call_is_begun_before_the_search_and_answers_the_first_attempt() {
let repo = repo();
let kai = honest().await;
let asked: Arc<std::sync::Mutex<Vec<std::time::SystemTime>>> = Arc::default();
let at = asked.clone();
let zen = testing::zen(move |_| {
at.lock().unwrap().push(std::time::SystemTime::now());
fix()
})
.await;
let report = run_flow(repo.path(), &kai.uri(), &zen, options(3, Control::Kai)).await;
assert_eq!(report.outcome, "done", "{}", report.summary());
assert_eq!(report.attempts.len(), 1);
assert_eq!(asked_zen(&zen).await.len(), 1, "one call, asked once");
let led = ledger(repo.path());
let mutate = ops(&led, "mutate")[0];
let searched: std::time::SystemTime = chrono::DateTime::parse_from_rfc3339(mutate["time"].as_str().unwrap()).unwrap().into();
assert!(asked.lock().unwrap()[0] < searched, "the call was made before the search ended");
let generate = ops(&led, "generate")[0];
let ahead = generate["result"]["ahead_ms"].as_f64().unwrap();
assert!(ahead >= mutate["ms"].as_f64().unwrap(), "the call had the search's time: {generate}");
assert_eq!(generate["result"]["stage"], "first");
assert_eq!(read(repo.path()), FIXED);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_wrong_edit_is_undone_and_the_retry_sees_it() {
let repo = repo();
@@ -467,7 +498,7 @@ async fn a_step_that_fixes_one_of_two_faults_is_kept_and_the_state_looked_at_aga
options.task["test"] = json!("python3 tests/test_both.py");
let report = run_flow(repo.path(), &kai.uri(), &zen, options).await;
assert_eq!(report.outcome, "done", "{}", report.summary());
assert!(asked_zen(&zen).await.is_empty(), "no generation: the search did both");
assert!(asked_zen(&zen).await.len() <= 1, "at most the call begun beside the first search");
let fast = report.fast.as_ref().unwrap();
assert!(fast.solved && fast.kept == 1 && fast.rounds == 2, "{fast:?}");
assert_eq!(report.failing, 2);
@@ -479,6 +510,12 @@ async fn a_step_that_fixes_one_of_two_faults_is_kept_and_the_state_looked_at_aga
assert_eq!(mutates, 2);
let verifies: Vec<&Value> = led.iter().filter(|l| l["kind"] == "verify").collect();
assert_eq!(verifies.iter().filter(|v| v["passed"] == true).count(), 1, "one pass, the last");
// No answer was used: the call asked at the first state was cancelled when the step moved it.
assert!(ops(&led, "generate").is_empty(), "{led:?}");
let cancelled: Vec<&Value> = led.iter().filter(|l| l["kind"] == "cancel").collect();
assert_eq!(cancelled.len(), 1, "{led:?}");
assert_eq!(cancelled[0]["revision"], "base");
assert_eq!(cancelled[0]["reason"], "the state it was asked at is left");
}
#[tokio::test(flavor = "multi_thread")]
@@ -800,6 +837,48 @@ async fn passing_the_failing_command_is_not_done_while_a_test_beside_it_fails()
let wide: Vec<&Value> = ops(&led, "test").into_iter().filter(|l| l["operation"]["command"] == "sh tools/pytest tests").collect();
assert!(wide.iter().any(|l| l["status"] == "rejected" && l["result"]["new"][0].as_str().is_some_and(|n| n.contains("test_mul.py"))), "{wide:?}");
assert!(wide.iter().any(|l| l["status"] == "verified"), "and the last, clean: {wide:?}");
let first = wide.iter().find(|l| l["revision"] == "base").unwrap();
assert_eq!(first["result"]["lazy"], true, "what failed before was read when a run after the edit failed: {first}");
}
#[test]
fn the_failing_command_and_the_tests_around_it_run_at_once_and_a_failure_stops_them() {
let dir = tempfile::tempdir().unwrap();
let ledger = Ledger::new("t", dir.path().join("ledger.jsonl"));
let mut slow = Wide { command: "sleep 30".into(), before: None };
let begun = Instant::now();
let (test, around) = check(dir.path(), "exit 1", Duration::from_secs(60), Some(&mut slow), &BTreeMap::new(), &ledger, "r").unwrap();
assert_eq!(test["passed"], false);
assert!(around.is_none(), "a failing command settles the edit without them");
assert!(begun.elapsed() < Duration::from_secs(25), "they were stopped: {:?}", begun.elapsed());
// The failing command passes only if the wider run began while it ran: at once, not one after the other.
let mut both = Wide { command: "touch began && sleep 3".into(), before: None };
let (test, around) = check(dir.path(), "sleep 3 && test -e began", Duration::from_secs(60), Some(&mut both), &BTreeMap::new(), &ledger, "r").unwrap();
assert_eq!(test["passed"], true, "{test}");
assert!(around.is_some_and(|v| v.new.is_empty()));
assert!(both.before.is_none(), "a wider run that passes needs nothing from before any edit");
}
#[test]
fn what_failed_before_any_edit_is_read_with_the_files_as_they_were_and_the_edit_is_put_back() {
let dir = tempfile::tempdir().unwrap();
let root = dir.path();
std::fs::write(root.join("m.txt"), "edited").unwrap();
std::fs::write(root.join("old.txt"), "BAD").unwrap();
let ledger = Ledger::new("t", root.join("ledger.jsonl"));
// The "runner" fails each file that says BAD, and m.txt fails until it is edited.
let script = "for f in m.txt old.txt; do grep -q BAD $f && echo \"FAILED $f::case - boom\"; done; grep -q BAD m.txt old.txt && exit 1; exit 0";
let mut wide = Wide { command: script.into(), before: None };
let pristine: BTreeMap<String, String> = [("m.txt".to_string(), "BAD".to_string())].into();
let (test, around) = check(root, "true", Duration::from_secs(20), Some(&mut wide), &pristine, &ledger, "r").unwrap();
assert_eq!(test["passed"], true);
let v = around.unwrap();
assert!(v.new.is_empty(), "old.txt failed before the edit too: {:?}", v.new);
let before = wide.before.as_ref().unwrap();
assert_eq!(before.failing.iter().map(String::as_str).collect::<Vec<_>>(), ["m.txt::case", "old.txt::case"]);
assert_eq!(std::fs::read_to_string(root.join("m.txt")).unwrap(), "edited", "the edit is put back");
let led: Vec<Value> = std::fs::read_to_string(root.join("ledger.jsonl")).unwrap().lines().map(|l| serde_json::from_str(l).unwrap()).collect();
assert!(led.iter().any(|l| l["result"]["lazy"] == true && l["revision"] == "base"), "{led:?}");
}
#[tokio::test(flavor = "multi_thread")]
crates/hanzo-kai/src/mutate.rs
+120 -12
View File
@@ -6,7 +6,8 @@
//! tokens that name no repository, file or task, at the lines of the ranked sites, best first.
//! [`search`] tries them in order, several at once, each in its own copy of the tree, against the
//! one test command the task shows, and stops at the first that passes. It never sees a hidden
//! test, and it gives up at a time cap so the model writes the edit instead.
//! test, and it gives up at a time cap so the model writes the edit instead; while the model
//! asked beside it has not answered, that wait is the search's, to a longer cap.
//!
//! What this can solve is a single fault of that kind. It is a search, not a repair of anything
//! else, and a candidate that passes the visible test is only as right as that test.
@@ -38,6 +39,8 @@ const TREE: u64 = 128 * 1024 * 1024;
const FILE: u64 = 8 * 1024 * 1024;
/// Lines of a function considered.
const LINES: u32 = 150;
/// The longest a search runs past its cap while the model it stands beside has not answered.
pub const LONGEST: Duration = Duration::from_secs(8);
/// One change to one line.
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -497,6 +500,52 @@ fn has_assignment(body: &str) -> bool {
false
}
/// pytest's options whose value is the next word.
const VALUED: &[&str] = &[
"-p", "-k", "-m", "-o", "-c", "-W", "-r", "--rootdir", "--deselect", "--ignore", "--ignore-glob", "--confcutdir", "--basetemp",
"--override-ini", "--config-file", "--junitxml", "--junit-xml", "--tb", "--maxfail", "--durations", "--log-file", "--import-mode",
];
/// The command a candidate is tried with when the visible one is too slow to try candidate after
/// candidate: a plain pytest command narrowed to the tests the failing `output` names, its paths
/// and node ids in the tree at `root` replaced by theirs. Only those tests can tell a candidate
/// that mends the fault from one that does not, so the rest of their files need not run; the
/// candidate the search picks is verified with the whole command in the tree all the same.
/// `None` for a command that is not plain pytest or an output that names no failing test.
pub fn screen(root: &Path, command: &str, output: &str) -> Option<String> {
if command.contains(['"', '\'', '\\']) {
return None;
}
let mut tokens = command.split_whitespace();
let mut kept: Vec<&str> = Vec::new();
let mut pytest = false;
for t in tokens.by_ref() {
kept.push(t);
if t == "pytest" || t.ends_with("/pytest") {
pytest = true;
break;
}
}
if !pytest {
return None;
}
// After pytest: a path in the tree, or a node id in one, selects tests; every other word is an
// option or the value of one and stays.
let mut value = false;
for t in tokens {
let selects = !value && !t.starts_with('-') && root.join(t.split("::").next().unwrap_or(t)).exists();
value = VALUED.contains(&t);
if !selects {
kept.push(t);
}
}
let ids = crate::wide::failing(output);
if ids.is_empty() || ids.iter().any(|id| id.contains('\'') || id.starts_with('-')) {
return None;
}
Some(format!("{} {}", kept.join(" "), ids.iter().map(|id| format!("'{id}'")).collect::<Vec<_>>().join(" ")))
}
/// What a search found and what it cost.
#[derive(Debug, Clone, Default)]
pub struct Search {
@@ -580,11 +629,24 @@ fn with_line(text: &str, number: u32, new: &str) -> String {
}
/// Tries `candidates` in order against `command`, `workers` at a time in copies of the tree at
/// `root`, until one passes, `cap` has gone, or they are spent. A candidate's run is cut at
/// `limit` (a mutation can loop). `failing` is how many tests fail now: when none passes, the
/// candidate that leaves the fewest failing, if fewer than that and the code still reads, is the
/// search's progress; `errors` of the failing are tests that could not run, and may not grow.
pub fn search(root: &Path, command: &str, candidates: &[Candidate], failing: usize, errors: usize, cap: Duration, limit: Duration, workers: usize) -> Search {
/// `root`, until one passes, `cap` has gone, or they are spent. While `waiting` says the model
/// asked beside the search has not answered, the search goes on past `cap`, to [`LONGEST`] at
/// most: that time is spent waiting either way. A candidate's run is cut at `limit` (a mutation
/// can loop). `failing` is how many tests fail now: when none passes, the candidate that leaves
/// the fewest failing, if fewer than that and the code still reads, is the search's progress;
/// `errors` of the failing are tests that could not run, and may not grow.
#[allow(clippy::too_many_arguments)]
pub fn search(
root: &Path,
command: &str,
candidates: &[Candidate],
failing: usize,
errors: usize,
cap: Duration,
waiting: &(dyn Fn() -> bool + Sync),
limit: Duration,
workers: usize,
) -> Search {
let begun = Instant::now();
let mut found = Search::default();
if candidates.is_empty() {
@@ -614,9 +676,13 @@ pub fn search(root: &Path, command: &str, candidates: &[Candidate], failing: usi
let step: Mutex<Option<(usize, usize, String)>> = Mutex::new(None);
let copy_ms = Mutex::new(0.0_f64);
let beaten = |i: usize| winner.lock().unwrap_or_else(std::sync::PoisonError::into_inner).as_ref().is_some_and(|(w, _)| *w < i);
let over = || {
let t = begun.elapsed();
t > cap && (t > LONGEST.max(cap) || !waiting())
};
std::thread::scope(|s| {
for w in 0..workers {
let (next, tried, runs, winner, step, copy_ms, beaten) = (&next, &tried, &runs, &winner, &step, &copy_ms, &beaten);
let (next, tried, runs, winner, step, copy_ms, beaten, over) = (&next, &tried, &runs, &winner, &step, &copy_ms, &beaten, &over);
let dir: PathBuf = scratch.path().join(format!("w{w}"));
s.spawn(move || {
let copying = Instant::now();
@@ -629,7 +695,7 @@ pub fn search(root: &Path, command: &str, candidates: &[Candidate], failing: usi
}
let mut originals: std::collections::HashMap<String, String> = std::collections::HashMap::new();
loop {
if begun.elapsed() > cap {
if over() {
break;
}
let i = next.fetch_add(1, Ordering::SeqCst);
@@ -648,7 +714,7 @@ pub fn search(root: &Path, command: &str, candidates: &[Candidate], failing: usi
}
// A test still running when the cap is reached is stopped: a round of search costs
// the cap at most, however long the tests are.
let run = tools::shell_until(&dir, command, limit, &|| beaten(i) || begun.elapsed() > cap);
let run = tools::shell_until(&dir, command, limit, &|| beaten(i) || over());
let _ = std::fs::write(&path, &original);
let code = run["code"].as_i64().unwrap_or(-1);
if code == -2 || code == -3 {
@@ -808,13 +874,55 @@ mod tests {
.iter()
.map(|n| Candidate { file: "m.py".into(), line: 2, old: " return a - b".into(), new: (*n).into(), kind: "operator", site: 0 })
.collect();
let found = search(dir.path(), "python3 t.py", &cands, 1, 0, Duration::from_secs(20), Duration::from_secs(10), 2);
let found = search(dir.path(), "python3 t.py", &cands, 1, 0, Duration::from_secs(20), &|| false, Duration::from_secs(10), 2);
assert_eq!(found.winner, Some(1), "{found:?}");
assert_eq!(std::fs::read_to_string(dir.path().join("m.py")).unwrap(), "def f(a, b):\n return a - b\n");
let none = search(dir.path(), "python3 t.py", &cands[..1], 1, 0, Duration::from_secs(20), Duration::from_secs(10), 1);
let none = search(dir.path(), "python3 t.py", &cands[..1], 1, 0, Duration::from_secs(20), &|| false, Duration::from_secs(10), 1);
assert_eq!(none.winner, None);
}
#[test]
fn a_candidate_is_tried_on_the_failing_tests_alone() {
let dir = tempfile::tempdir().unwrap();
std::fs::create_dir_all(dir.path().join("tests")).unwrap();
std::fs::write(dir.path().join("tests/test_a.py"), "").unwrap();
let out = "FAILED tests/test_a.py::test_x[a&b?c] - assert 1 == 2\nFAILED tests/test_a.py::test_y - boom\n2 failed, 9 passed in 0.31s\n";
let base = "PYTHONPATH=src:. python -m pytest -q -p no:cacheprovider -o addopts=";
assert_eq!(
screen(dir.path(), &format!("{base} tests/test_a.py"), out).as_deref(),
Some(&*format!("{base} 'tests/test_a.py::test_x[a&b?c]' 'tests/test_a.py::test_y'"))
);
// A directory or a node id is a selector too; an option's value is not.
assert_eq!(
screen(dir.path(), &format!("{base} tests tests/test_a.py::test_y"), out).as_deref(),
Some(&*format!("{base} 'tests/test_a.py::test_x[a&b?c]' 'tests/test_a.py::test_y'"))
);
assert_eq!(
screen(dir.path(), "pytest -q --rootdir tests -c tests/test_a.py tests/test_a.py", out).as_deref(),
Some("pytest -q --rootdir tests -c tests/test_a.py 'tests/test_a.py::test_x[a&b?c]' 'tests/test_a.py::test_y'"),
"the value of an option that takes one stays"
);
assert_eq!(screen(dir.path(), "cargo test", out), None, "not pytest");
assert_eq!(screen(dir.path(), &format!("{base} tests"), "1 failed in 0.1s\n"), None, "no failing test named");
assert_eq!(screen(dir.path(), "pytest 'tests/test a.py'", out), None, "quoted words are left alone");
}
#[test]
fn a_search_goes_on_past_its_cap_while_the_model_beside_it_has_not_answered() {
let dir = tempfile::tempdir().unwrap();
std::fs::write(dir.path().join("m.py"), "def f(a, b):\n return a - b\n").unwrap();
std::fs::write(dir.path().join("t.py"), "import time\nfrom m import f\ntime.sleep(0.05)\nassert f(3, 2) == 5\n").unwrap();
let cands: Vec<Candidate> = (0..200)
.map(|k| Candidate { file: "m.py".into(), line: 2, old: " return a - b".into(), new: format!(" return a - b - {k}"), kind: "constant", site: 0 })
.collect();
let alone = search(dir.path(), "python3 t.py", &cands, 1, 0, Duration::ZERO, &|| false, Duration::from_secs(10), 1);
assert_eq!(alone.tried, 0, "past its cap with nothing to wait for, the search stops: {alone:?}");
let begun = Instant::now();
let beside = search(dir.path(), "python3 t.py", &cands, 1, 0, Duration::ZERO, &|| begun.elapsed() < Duration::from_millis(1500), Duration::from_secs(10), 1);
assert!(beside.tried >= 1 && beside.total >= 1400.0, "it went on while the model wrote: {beside:?}");
assert!(beside.tried < cands.len(), "and stopped short of its candidates: {beside:?}");
}
#[test]
fn with_two_faults_the_candidate_that_fixes_one_is_the_searchs_progress() {
let dir = tempfile::tempdir().unwrap();
@@ -829,7 +937,7 @@ mod tests {
.iter()
.map(|(line, n)| Candidate { file: "m.py".into(), line: *line, old: String::new(), new: (*n).into(), kind: "operator", site: 0 })
.collect();
let found = search(dir.path(), "python3 t.py", &cands, 2, 0, Duration::from_secs(20), Duration::from_secs(10), 2);
let found = search(dir.path(), "python3 t.py", &cands, 2, 0, Duration::from_secs(20), &|| false, Duration::from_secs(10), 2);
assert_eq!(found.winner, None, "{found:?}");
assert_eq!(found.progress, Some(1), "the f fix leaves one failing: {found:?}");
assert_eq!(found.remaining, 1);
crates/hanzo-kai/src/tools.rs
+108
View File
@@ -160,6 +160,14 @@ pub fn shell_until(root: &Path, script: &str, limit: Duration, stop: &dyn Fn() -
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::null());
// A cache a test framework keeps in the working directory is the run's own, kept outside the
// tree and gone with the run, so no run sees another's examples and none lands in the tree.
let caches = tempfile::tempdir().ok();
if let Some(dir) = &caches
&& std::env::var_os("HYPOTHESIS_STORAGE_DIRECTORY").is_none()
{
command.env("HYPOTHESIS_STORAGE_DIRECTORY", dir.path().join("hypothesis"));
}
// Its own process group, so a timeout stops everything the script started.
#[cfg(unix)]
std::os::unix::process::CommandExt::process_group(&mut command, 0);
@@ -204,6 +212,55 @@ pub fn shell_until(root: &Path, script: &str, limit: Duration, stop: &dyn Fn() -
}
/// Runs the tests: `{passed, code, output, ms}`.
/// The files under a repository that git neither tracks nor ignores, and the untracked
/// directories as git names them (`dir/`): what `git status` would show beside the change.
/// `None` outside a repository.
fn untracked(root: &Path, directories: bool) -> Option<HashSet<String>> {
let mut git = Command::new("git");
git.args(["ls-files", "--others", "--exclude-standard", "-z"]);
if directories {
git.arg("--directory");
}
let out = git.current_dir(root).stderr(Stdio::null()).output().ok().filter(|o| o.status.success())?;
Some(out.stdout.split(|b| *b == 0).filter(|p| !p.is_empty()).map(|p| String::from_utf8_lossy(p).into_owned()).collect())
}
/// What a repository holds beside its tracked files before a run in it: the run's own files are
/// what it adds to that ([`Litter::sweep`]).
pub struct Litter {
files: Option<HashSet<String>>,
entries: HashSet<String>,
}
impl Litter {
pub fn before(root: &Path) -> Litter {
Litter { files: untracked(root, false), entries: untracked(root, true).unwrap_or_default() }
}
/// Removes the files a run in `root` wrote that were not there before it and that git neither
/// tracks nor ignores, with the directories they emptied that were not there before either:
/// they are the run's, never an edit's. What `.gitignore` names is left alone, and so is
/// anything outside a repository. The paths removed.
pub fn sweep(&self, root: &Path) -> Vec<String> {
let (Some(before), Some(now)) = (&self.files, untracked(root, false)) else {
return Vec::new();
};
let mut gone: Vec<String> = now.difference(before).filter(|p| std::fs::remove_file(root.join(p)).is_ok()).cloned().collect();
gone.sort();
for p in &gone {
let mut dir = Path::new(p).parent();
while let Some(d) = dir.filter(|d| !d.as_os_str().is_empty()) {
let name = format!("{}/", d.to_string_lossy());
if self.entries.contains(&name) || std::fs::remove_dir(root.join(d)).is_err() {
break;
}
dir = d.parent();
}
}
gone
}
}
pub fn test(root: &Path, command: &str, limit: Duration) -> Value {
let mut run = shell(root, command, limit);
let passed = run["code"] == 0;
@@ -307,6 +364,57 @@ pub fn revert(root: &Path, answer: &str) -> bool {
mod tests {
use super::*;
fn git(root: &Path, args: &[&str]) {
let ok = Command::new("git")
.args(["-c", "user.name=t", "-c", "user.email=t@t", "-c", "commit.gpgsign=false"])
.args(args)
.current_dir(root)
.output()
.unwrap()
.status
.success();
assert!(ok, "git {args:?}");
}
#[test]
fn what_a_test_run_writes_beside_the_tree_is_swept_and_nothing_else() {
let dir = tempfile::tempdir().unwrap();
let root = dir.path();
std::fs::write(root.join("m.py"), "x = 1\n").unwrap();
std::fs::write(root.join(".gitignore"), "cache/\n").unwrap();
git(root, &["init", "-q"]);
git(root, &["add", "."]);
git(root, &["commit", "-q", "-m", "start"]);
// Already there before the run: an untracked file and an empty directory, both kept.
std::fs::write(root.join("notes.txt"), "mine").unwrap();
std::fs::create_dir_all(root.join("empty")).unwrap();
let litter = Litter::before(root);
let run = shell(root, "mkdir -p .hypothesis/examples/ab cache && echo e > .hypothesis/examples/ab/1 && echo o > out.log && echo c > cache/x && echo 'x = 2' > m.py", Duration::from_secs(10));
assert_eq!(run["code"], 0);
let gone = litter.sweep(root);
assert_eq!(gone, [".hypothesis/examples/ab/1", "out.log"]);
assert!(!root.join(".hypothesis").exists(), "the directories it emptied go too");
assert!(root.join("notes.txt").exists() && root.join("empty").is_dir(), "what was there before stays");
assert!(root.join("cache/x").exists(), "what .gitignore names is left alone");
assert_eq!(std::fs::read_to_string(root.join("m.py")).unwrap(), "x = 2\n", "a tracked file is never touched");
// Outside a repository nothing is swept.
let plain = tempfile::tempdir().unwrap();
let litter = Litter::before(plain.path());
std::fs::write(plain.path().join("new.txt"), "n").unwrap();
assert!(litter.sweep(plain.path()).is_empty());
}
#[test]
fn a_runs_test_framework_cache_is_its_own_and_outside_the_tree() {
let dir = tempfile::tempdir().unwrap();
let run = shell(dir.path(), "echo \"$HYPOTHESIS_STORAGE_DIRECTORY\"", Duration::from_secs(10));
let at = run["output"].as_str().unwrap().trim().to_string();
if std::env::var_os("HYPOTHESIS_STORAGE_DIRECTORY").is_none() {
assert!(!at.is_empty() && !Path::new(&at).starts_with(dir.path()), "{at}");
assert!(!Path::new(&at).parent().unwrap().exists(), "gone with the run: {at}");
}
}
#[test]
fn a_fenced_diff_is_taken_from_the_answer() {
let answer =
crates/hanzo-kai/src/wide.rs
+70 -37
View File
@@ -3,10 +3,11 @@
//! The failing command passing settles the failing tests and nothing near them. Before a flow
//! says `done`, the tests around the failing ones run too: the directory of every test file the
//! command or the failing output names, in place of the files and selectors that narrowed the
//! command. They ran once before any edit, so a failure that was there already is not the edit's;
//! a test failing now that was not then makes the edit not done. The run is capped at [`LIMIT`]; a
//! run that does not end by then says nothing. Only the command and the failing output are read,
//! never anything of the task beyond them.
//! command. A failure that was there before any edit is not the edit's; a test failing now that
//! was not then makes the edit not done. What failed before is read only when a run after an edit
//! fails: a run that passes shows no new failure whatever failed before, so on most tasks it is
//! never read. The run is capped at [`LIMIT`]; a run that does not end by then says nothing. Only
//! the command and the failing output are read, never anything of the task beyond them.
use crate::tools;
use serde_json::Value;
@@ -17,12 +18,19 @@ use std::time::Duration;
/// The longest a wider run may take.
pub const LIMIT: Duration = Duration::from_secs(20);
/// The wider command, and what failed in it before any edit.
/// The wider command, and what failed in it before any edit once that is read.
#[derive(Debug, Clone)]
pub struct Wide {
pub command: String,
pub before: BTreeSet<String>,
pub before_code: i64,
pub before: Option<Before>,
}
/// What the wider run said before any edit.
#[derive(Debug, Clone, Default)]
pub struct Before {
pub failing: BTreeSet<String>,
/// Its exit code; -2 when it did not end within [`LIMIT`], and then it says nothing.
pub code: i64,
}
/// What a wider run found.
@@ -102,34 +110,47 @@ pub fn failing(output: &str) -> BTreeSet<String> {
}
impl Wide {
/// The wider run on the tree as it is, before any edit, for what fails in it already. `None`
/// when there is no wider command or the run does not end within [`LIMIT`].
pub fn before(root: &Path, command_line: &str, output: &str) -> Option<Wide> {
let command = command(command_line, output)?;
let run = tools::test(root, &command, LIMIT);
let code = run["code"].as_i64().unwrap_or(-1);
if code == -2 || code == -1 {
return None;
}
Some(Wide {
before: failing(run["output"].as_str().unwrap_or_default()),
before_code: code,
command,
})
/// The wider command for `command_line` and its failing `output`; `None` when there is none.
pub fn new(command_line: &str, output: &str) -> Option<Wide> {
Some(Wide { command: command(command_line, output)?, before: None })
}
/// The wider run on the tree as it is, against what failed before.
pub fn check(&self, root: &Path) -> Verdict {
let run = tools::test(root, &self.command, LIMIT);
/// The wider run on the tree as it is, `{code, output, passed, ms}`; stopped as soon as `stop`
/// says so (code -3).
pub fn run(&self, root: &Path, stop: &dyn Fn() -> bool) -> Value {
let mut run = tools::shell_until(root, &self.command, LIMIT, stop);
run["passed"] = Value::Bool(run["code"] == 0);
run
}
/// Whether `run`, after an edit, needs what failed before to be judged: it failed, and ended.
pub fn needs_before(run: &Value) -> bool {
!matches!(run["code"].as_i64(), Some(0 | -2 | -3))
}
/// What a run before any edit says.
pub fn read(run: &Value) -> Before {
Before { failing: failing(run["output"].as_str().unwrap_or_default()), code: run["code"].as_i64().unwrap_or(-1) }
}
/// What `run`, after an edit, says against what failed before.
pub fn verdict(&self, run: Value) -> Verdict {
let code = run["code"].as_i64().unwrap_or(-1);
let timed_out = code == -2;
let now = failing(run["output"].as_str().unwrap_or_default());
let mut new: Vec<String> = now.difference(&self.before).cloned().collect();
// A run that dies without naming a test (a collection error in a file that read before).
if new.is_empty() && now.is_empty() && code != 0 && !timed_out && self.before_code == 0 {
new.push(format!("the wider run exited {code}"));
let mut new = Vec::new();
// A run that passes shows nothing new; one before any edit that said nothing leaves
// nothing to judge against.
if let Some(before) = self.before.as_ref().filter(|b| b.code >= 0)
&& Wide::needs_before(&run)
{
let now = failing(run["output"].as_str().unwrap_or_default());
new = now.difference(&before.failing).cloned().collect();
// A run that dies without naming a test (a collection error in a file that read before).
if new.is_empty() && now.is_empty() && before.code == 0 {
new.push(format!("the wider run exited {code}"));
}
}
Verdict { new: if timed_out { Vec::new() } else { new }, timed_out, run }
Verdict { new, timed_out, run }
}
}
@@ -182,16 +203,28 @@ mod tests {
let script = "for f in t/*.txt; do grep -q BAD $f && echo \"FAILED $f::case - boom\"; done; grep -rq BAD t && exit 1; exit 0";
std::fs::write(root.join("t/old.txt"), "BAD").unwrap();
std::fs::write(root.join("t/new.txt"), "good").unwrap();
let before = Wide {
command: script.to_string(),
before: ["t/old.txt::case".to_string()].into(),
before_code: 1,
};
let v = before.check(root);
let mut wide = Wide { command: script.to_string(), before: None };
let run = wide.run(root, &|| false);
assert!(Wide::needs_before(&run), "the run before fails");
wide.before = Some(Wide::read(&run));
let v = wide.verdict(wide.run(root, &|| false));
assert!(v.new.is_empty(), "{:?}", v.new);
std::fs::write(root.join("t/new.txt"), "BAD").unwrap();
let v = before.check(root);
let v = wide.verdict(wide.run(root, &|| false));
assert_eq!(v.new, ["t/new.txt::case"]);
assert!(!v.timed_out);
}
#[test]
fn a_run_that_passes_needs_nothing_from_before_and_a_stopped_one_says_nothing() {
let dir = tempfile::tempdir().unwrap();
let wide = Wide { command: "true".into(), before: None };
let run = wide.run(dir.path(), &|| false);
assert!(!Wide::needs_before(&run));
assert!(wide.verdict(run).new.is_empty());
let wide = Wide { command: "sleep 5".into(), before: None };
let run = wide.run(dir.path(), &|| true);
assert_eq!(run["code"], -3);
assert!(!Wide::needs_before(&run));
}
}
crates/hanzo-kai/src/zen.rs
+180 -4
View File
@@ -15,6 +15,11 @@
//!
//! An edit is asked for as a JSON schema (`response_format`, strict) where the server is Hanzo's,
//! which constrains the decoding: the answer cannot be prose, nor a string with a bare quote in it.
//!
//! A call may be begun before the loop needs it ([`Zen::start`]) and waited for later. Dropping
//! the [`Pending`] cancels it: the stream is closed, so the server stops writing and billing it,
//! and the spend file gets a line for what it was sent (`cancelled`, its prompt counted from the
//! request's characters at the ceiling rate, as every unpriced call is).
use hanzo_config::kai::Generate;
use serde_json::Value;
@@ -59,7 +64,9 @@ pub struct Generation {
pub text: String,
/// Request sent to the first token, milliseconds.
pub first: f64,
/// Request sent to the last token, milliseconds; retries included.
/// Request sent to the last token of the answer's text, milliseconds.
pub last: f64,
/// Request sent to the end of the stream, milliseconds; retries included.
pub total: f64,
pub prompt: u64,
pub completion: u64,
@@ -136,6 +143,7 @@ fn schema() -> Value {
}
/// Zen over `POST {url}/chat/completions` as `model`.
#[derive(Clone)]
pub struct Zen {
url: String,
model: String,
@@ -213,6 +221,35 @@ impl Zen {
"cost_basis": format!("ceiling at ${}/M in, ${}/M out", RATE.0, RATE.1),
"headers": headers,
});
Self::append(path, &line);
}
/// The line a call cancelled before its end leaves in the spend file: the prompt it was sent,
/// counted at three characters a token (more tokens than text takes, so the budget stays a
/// ceiling), and the pieces it had streamed, a token each.
fn record_cut(&self, model: &str, body: &Value, streamed: u64, ms: f64) {
let Some((path, _)) = &self.spend else {
return;
};
let chars: usize = body["messages"]
.as_array()
.map(|m| m.iter().filter_map(|x| x["content"].as_str()).map(str::len).sum())
.unwrap_or_default();
let prompt = chars.div_ceil(3) as u64;
let line = json!({
"time": crate::trace::now(),
"model": model,
"cancelled": true,
"prompt_tokens": prompt,
"completion_tokens": streamed,
"ms": ms,
"cost": (prompt as f64 * RATE.0 + streamed as f64 * RATE.1) / 1e6,
"cost_basis": format!("estimated: the prompt at 3 characters a token, a token a piece streamed, at ${}/M in, ${}/M out", RATE.0, RATE.1),
});
Self::append(path, &line);
}
fn append(path: &PathBuf, line: &Value) {
if let Some(dir) = path.parent() {
let _ = std::fs::create_dir_all(dir);
}
@@ -264,7 +301,16 @@ impl Zen {
request = request.bearer_auth(key);
}
let begun = Instant::now();
let mut response = request.send().await.map_err(|e| format!("zen: {}", e.without_url()))?;
// Until the stream has ended, a drop of this future is a cancelled call: it leaves its
// line in the spend file.
let mut cut = Cut { zen: self, model, body, streamed: 0, begun, armed: true };
let mut response = match request.send().await {
Ok(r) => r,
Err(e) => {
cut.armed = false;
return Err(format!("zen: {}", e.without_url()));
}
};
let status = response.status();
let headers = json!({
"served": response.headers().get("x-hanzo-served").and_then(|v| v.to_str().ok()),
@@ -272,6 +318,8 @@ impl Zen {
"paid_by": response.headers().get("x-hanzo-paid-by").and_then(|v| v.to_str().ok()),
});
if !status.is_success() {
// A request the server refused wrote nothing.
cut.armed = false;
let body = response.text().await.unwrap_or_default();
let why: Value = serde_json::from_str(&body).unwrap_or_default();
let msg = why["error"]["message"]
@@ -281,7 +329,7 @@ impl Zen {
.unwrap_or_else(|| crate::decide::cut(&body, 200));
return Err(format!("zen: {status}: {msg}"));
}
let (mut pending, mut answer, mut first) = (String::new(), String::new(), None);
let (mut pending, mut answer, mut first, mut last) = (String::new(), String::new(), None, 0.0);
let (mut prompt_tokens, mut completion_tokens, mut reasoning_tokens) = (0, 0, 0);
let (mut served, mut finish) = (String::new(), String::new());
while let Some(chunk) = response.chunk().await.map_err(|e| format!("zen: {}", e.without_url()))? {
@@ -302,11 +350,14 @@ impl Zen {
let thought = delta["reasoning_content"].as_str().or_else(|| delta["reasoning"].as_str());
if thought.is_some_and(|t| !t.is_empty()) {
first.get_or_insert_with(|| begun.elapsed().as_secs_f64() * 1000.0);
cut.streamed += 1;
}
if let Some(piece) = delta["content"].as_str()
&& !piece.is_empty()
{
first.get_or_insert_with(|| begun.elapsed().as_secs_f64() * 1000.0);
last = begun.elapsed().as_secs_f64() * 1000.0;
cut.streamed += 1;
answer.push_str(piece);
}
if let Some(why) = event["choices"][0]["finish_reason"].as_str() {
@@ -323,6 +374,7 @@ impl Zen {
}
}
let total = begun.elapsed().as_secs_f64() * 1000.0;
cut.armed = false;
if let Some(s) = headers["served"].as_str() {
served = s.to_string();
}
@@ -330,6 +382,7 @@ impl Zen {
Generation {
text: answer,
first: first.unwrap_or(total),
last: if last > 0.0 { last } else { total },
total,
prompt: prompt_tokens,
completion: completion_tokens,
@@ -351,6 +404,7 @@ impl Zen {
Some(Generation {
text: r["text"].as_str().filter(|t| !t.is_empty() && doc["error"].is_null())?.to_string(),
first: 0.0,
last: 0.0,
total: 0.0,
prompt: r["prompt_tokens"].as_u64()?,
completion: r["completion_tokens"].as_u64()?,
@@ -399,7 +453,7 @@ impl Zen {
let looked = Instant::now();
if let Some(mut g) = self.kept(&name) {
let took = looked.elapsed().as_secs_f64() * 1000.0;
(g.first, g.total) = (took, took);
(g.first, g.last, g.total) = (took, took, took);
return Ok(g);
}
if let Some((path, budget)) = &self.spend
@@ -450,6 +504,83 @@ impl Zen {
pub fn generate(&self, model: &str, system: &str, user: &str, effort: Effort) -> Result<Generation, String> {
self.runtime.block_on(self.complete(model, system, user, effort))
}
/// [`Zen::generate`] begun now and waited for later: the call runs on the runtime while the
/// caller does other work. Dropping the [`Pending`] cancels it.
pub fn start(&self, model: &str, system: &str, user: &str, effort: Effort) -> Pending {
let zen = self.clone();
let (m, s, u) = (model.to_string(), system.to_string(), user.to_string());
Pending {
handle: Some(self.runtime.spawn(async move { zen.complete(&m, &s, &u, effort).await })),
runtime: self.runtime.clone(),
begun: Instant::now(),
}
}
}
/// A call begun with [`Zen::start`].
pub struct Pending {
handle: Option<tokio::task::JoinHandle<Result<Generation, String>>>,
runtime: tokio::runtime::Handle,
begun: Instant,
}
impl Pending {
/// Whether the answer is in.
pub fn done(&self) -> bool {
self.handle.as_ref().is_none_or(tokio::task::JoinHandle::is_finished)
}
/// Milliseconds since the call was begun.
pub fn age(&self) -> f64 {
self.begun.elapsed().as_secs_f64() * 1000.0
}
/// The answer, waiting for it if it is not in. Blocks: call it off the async runtime.
pub fn wait(mut self) -> Result<Generation, String> {
let handle = self.handle.take().ok_or("zen: the call was already waited for")?;
self.runtime.block_on(handle).map_err(|e| format!("zen: the call did not finish: {e}"))?
}
/// Cancels the call and returns once it has stopped, its line in the spend file written;
/// `true` when it had not finished. Blocks: call it off the async runtime.
pub fn cancel(mut self) -> bool {
let Some(handle) = self.handle.take() else {
return false;
};
let running = !handle.is_finished();
handle.abort();
let _ = self.runtime.block_on(handle);
running
}
}
impl Drop for Pending {
fn drop(&mut self) {
if let Some(handle) = self.handle.take() {
handle.abort();
}
}
}
/// While armed, a call that is dropped before its stream ends: its drop records the call as
/// cancelled.
struct Cut<'a> {
zen: &'a Zen,
model: &'a str,
body: &'a Value,
/// Pieces of text the stream had written.
streamed: u64,
begun: Instant,
armed: bool,
}
impl Drop for Cut<'_> {
fn drop(&mut self) {
if self.armed {
self.zen.record_cut(self.model, self.body, self.streamed, self.begun.elapsed().as_secs_f64() * 1000.0);
}
}
}
#[cfg(test)]
@@ -619,6 +750,51 @@ mod tests {
}
}
#[tokio::test(flavor = "multi_thread")]
async fn a_call_begun_ahead_answers_when_waited_for_as_the_same_call_would() {
let server = testing::zen(|_| "{\"edits\":[]}".to_string()).await;
let home = tempfile::tempdir().unwrap();
let (calls, spend) = (home.path().join("calls"), home.path().join("spend.jsonl"));
let zen = zen_at(&server, &calls, &spend);
let pending = zen.start("base", "sys", "user", Effort::Low);
let g = tokio::task::block_in_place(|| pending.wait()).unwrap();
assert!(!g.cached && g.text == "{\"edits\":[]}" && g.first <= g.last && g.last <= g.total, "{g:?}");
// It was kept as any call is: the same request again is answered from there.
let again = tokio::task::block_in_place(|| zen.generate("base", "sys", "user", Effort::Low)).unwrap();
assert!(again.cached);
assert_eq!(server.received_requests().await.unwrap().len(), 1);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_cancelled_call_stops_and_leaves_its_line_in_the_spend_file() {
use wiremock::matchers::{method, path};
let server = wiremock::MockServer::start().await;
wiremock::Mock::given(method("POST"))
.and(path("/v1/chat/completions"))
.respond_with(wiremock::ResponseTemplate::new(200).set_delay(Duration::from_secs(30)))
.mount(&server)
.await;
let home = tempfile::tempdir().unwrap();
let (calls, spend) = (home.path().join("calls"), home.path().join("spend.jsonl"));
let mut zen = zen_at(&server, &calls, &spend);
// Paid, as every call to Hanzo's API is.
zen.spend = Some((spend.clone(), f64::INFINITY));
let pending = zen.start("base", "sys", "user", Effort::Low);
tokio::time::sleep(Duration::from_millis(300)).await;
assert!(!pending.done());
let begun = Instant::now();
assert!(tokio::task::block_in_place(|| pending.cancel()), "it was running");
assert!(begun.elapsed() < Duration::from_secs(5), "a cancel does not wait for the answer");
assert_eq!(server.received_requests().await.unwrap().len(), 1);
let lines: Vec<Value> = std::fs::read_to_string(&spend).unwrap().lines().map(|l| serde_json::from_str(l).unwrap()).collect();
assert_eq!(lines.len(), 1);
assert_eq!((lines[0]["cancelled"].as_bool(), lines[0]["model"].as_str()), (Some(true), Some("base")));
// "sys" and "user": seven characters, three tokens at three characters a token; nothing streamed.
assert_eq!((lines[0]["prompt_tokens"].as_u64(), lines[0]["completion_tokens"].as_u64()), (Some(3), Some(0)));
assert!(lines[0]["cost"].as_f64().unwrap() > 0.0);
assert!(!calls.exists() || std::fs::read_dir(&calls).unwrap().next().is_none(), "a cancelled call is not kept");
}
#[tokio::test(flavor = "multi_thread")]
async fn a_failed_call_is_not_kept_and_a_damaged_one_is_asked_again() {
let server = testing::zen(|_| "{}".to_string()).await;