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

fix flow: Kai as controller of a ledgered loop, a search of candidate edits, tgrep evidence, recover

Browse Source 
dev flow fix reads the failing output against the tree (locate: the gateway's
learned ranker, features held to its own by a test), puts the functions tgrep
finds spelling the failure's messages into the ranking, and lets Kai choose the
site, the action and the effort in one request, or decides alone when the
ranking is sure. The action search tries generic token transforms at the best
sites against the visible test, eight at once; a candidate that leaves fewer
tests failing is kept as a step. Else a model writes one edit (applied by text,
never a test file), the test runs, Kai reads a failed run in one request, and
recover fingerprints the failure and never replays a strategy on it.

Every action is a ledger line with a stable id, its revision, preconditions,
authorization, source and the states it passed through; credentials are
redacted. Kai's transport is a trait, health is per program.

Cargo.lock is not committed: hanzo-trigram is read from hanzoai/mcp commit
e9e6096, which is not pushed yet.
This commit is contained in:
zeekay
2026-10-14 14:54:05 -07:00
parent 8a296d976c
commit 501a220471
22 changed files with 5845 additions and 453 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.toml
+5
View File
@@ -501,3 +501,8 @@ tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev
[patch."ssh://[email protected]/openai-oss-forks/tungstenite-rs.git"]
tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "4fffad30fe373adbdcffab9545e9e9bf4f2fc19f" }
# hanzo-trigram's tgrep-core pins rusqlite by git commit, whose libsqlite3-sys 0.38 would link a
# second SQLite beside the 0.37 sqlx bundles. rusqlite 0.39 is the release built on 0.37.
[patch."https://github.com/rusqlite/rusqlite.git"]
rusqlite = "=0.39.0"
LLM.md
+38
View File
@@ -129,6 +129,44 @@ engine with no catalog of its own is reviewed by the model it runs. Guardian v2'
sync reviewer runs on that same choice; its `CodexResponsesHeaders` model only
scopes an OpenAI-backend header. `crates/hanzo-dev/tests/review.rs` holds both.
## Controller
Kai is Dev's native controller, not an advisory service: it returns a validated operation or an
explicit abstention, never advice. `crates/hanzo-kai` is the loop (`flow`), the ledger (`ledger`)
and the tools they run.
- **Task state.** One per run: objective, workspace revision, pending actions, verification
status. The loop owns it and each action's lifecycle. An expected outcome is never accepted
because something predicted it: only the verifier's record (exit status, the tests, the diff)
settles it.
- **Action.** `Action { id, task, revision, evidence, operation, preconditions, authorization,
expected, verification, source }` plus status, milliseconds and result. `operation` is typed
(`search`, `decide`, `generate`, `edit`, `test`, `revert`); `source` is `rule`, `kai`, `model` or
`recover`. A saved typed operation is an action; a plan of actions is a workflow; a reusable
validated result is the cache. Kai selects and binds an action's inputs and never writes shell.
- **Order of cheapness, per step.** An action that applies (one candidate, preconditions hold)
runs with no Kai inference. Else Kai's structured decision: next action, search, effort; all of a
state's questions in one request. Else low-effort generation, then full-effort generation, then
`recover`. Every transition is a ledger line with its reason; every fall from a deterministic
step to a model is logged with why. No step waits on Kai when an action already applies.
- **Revision-bound.** Preconditions are checked immediately before execution; an edit invalidates
every pending decision made against the old revision. A cached result is keyed on what it depends
on (files, symbols, tool versions, environment; uncommitted changes count), never on a commit.
- **Recover proves progress.** A failed verification records the failure's fingerprint (paths, line
numbers and ids normalized), the prior action, the resulting diff and the outcome. The same
fingerprint with no meaningful change of state makes `recover` pick a different action or
escalate; it never replays the loop.
- **Search** is one engine (tgrep): files, symbols, documents, freshness. Until it lands the
engine is `locate`, the learned ranker over the tree's functions (kai-gateway's preflight policy,
weights from bench/code/onegen/policies.py, `locate::tests` hold its features to the gateway's).
- **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.
- **Ledger.** `kai/ledger.jsonl`, replayable: stable ids for every decision, fallback, execution
and verification so time and cost are attributable; credentials and sensitive output redacted.
- **Generation** is Zen or enso-flash over an OpenAI-compatible endpoint, called when needed.
Independent checks (the verifier) decide success and trigger `recover`.
## House rules
- Hanzo service APIs use `https://api.hanzo.ai/v1/` exclusively.
crates/hanzo-config/src/kai.rs
+40 -4
View File
@@ -127,16 +127,38 @@ pub struct Op {
pub k: Option<usize>,
}
/// The generator as written.
#[derive(Debug, Clone, Default, PartialEq, Deserialize)]
#[serde(default, deny_unknown_fields)]
struct Writing {
url: Option<String>,
model: Option<String>,
spend: Option<PathBuf>,
budget: Option<f64>,
}
#[derive(Debug, Clone, Default, PartialEq, Deserialize)]
#[serde(default, deny_unknown_fields)]
struct File {
url: Option<String>,
model: Option<String>,
trace: Option<PathBuf>,
zen: Option<String>,
generate: Writing,
ops: BTreeMap<String, Written>,
}
/// The model that writes code: Zen over an OpenAI-compatible `POST {url}/chat/completions`.
#[derive(Debug, Clone, PartialEq)]
pub struct Generate {
/// Hanzo's API, or any server that speaks it; Hanzo's credential goes to Hanzo's alone.
pub url: String,
pub model: String,
/// Where each paid call is recorded, JSON lines with the model that served it and its cost.
pub spend: PathBuf,
/// Dollars the spend file may hold before a call is refused; none, no limit.
pub budget: Option<f64>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct Kai {
/// The Decisions API base: `POST {url}/decisions`.
@@ -145,8 +167,8 @@ pub struct Kai {
pub model: String,
/// The decision trace, JSON lines, one per decision.
pub trace: PathBuf,
/// The Zen model a flow's generation nodes call.
pub zen: String,
/// The model a flow's generation nodes call.
pub generate: Generate,
ops: BTreeMap<&'static str, Op>,
}
@@ -208,7 +230,21 @@ impl Kai {
file.trace
.unwrap_or_else(|| PathBuf::from("kai/decisions.jsonl")),
),
zen: file.zen.unwrap_or_else(|| "zen6-coder".to_string()),
generate: Generate {
url: file
.generate
.url
.unwrap_or_else(|| crate::API_BASE.to_string())
.trim_end_matches('/')
.to_string(),
model: file.generate.model.unwrap_or_else(|| "zen6-coder".to_string()),
spend: home.join(
file.generate
.spend
.unwrap_or_else(|| PathBuf::from("kai/spend.jsonl")),
),
budget: file.generate.budget,
},
ops,
})
}
crates/hanzo-dev/src/main.rs
+19 -6
View File
@@ -159,6 +159,20 @@ struct Fix {
#[arg(long, default_value = "")]
goal: String,
/// A file holding the failing run's output, already in hand (a CI log): the tests are not run
/// to see them fail.
#[arg(long)]
failure: Option<std::path::PathBuf>,
/// Do not try candidate edits against the test before asking the model.
#[arg(long)]
no_search: bool,
/// Do not put tgrep's evidence (where the failure's messages and names are spelled) into the
/// ranking.
#[arg(long)]
no_tgrep: bool,
/// Most attempts.
#[arg(long, default_value_t = 3)]
max: u32,
@@ -320,14 +334,13 @@ async fn run(paths: Arg0DispatchPaths) -> Result<()> {
async fn run_fix(args: Fix) -> Result<()> {
use hanzo_kai::flow;
let setup = flow::Setup::new(
&hanzo_config::home(),
hanzo_config::API_BASE,
hanzo_config::hanzo_credential(),
)
.map_err(anyhow::Error::msg)?;
let setup = flow::Setup::new(&hanzo_config::home(), hanzo_config::hanzo_credential)
.map_err(anyhow::Error::msg)?;
let options = flow::Options {
task: serde_json::json!({"test": args.test, "goal": args.goal}),
failure: args.failure.as_deref().map(std::fs::read_to_string).transpose()?,
search: !args.no_search,
discover: !args.no_tgrep,
max: args.max,
control: match args.control {
Steer::Kai => flow::Control::Kai,
crates/hanzo-kai/Cargo.toml
+4 -1
View File
@@ -11,21 +11,24 @@ workspace = true
[dependencies]
hanzo-config = { path = "../hanzo-config", version = "0.7.0-dev.1" }
hanzo-loop = { path = "../hanzo-loop", version = "0.7.0-dev.1" }
# Hanzo's one search engine (tgrep), the library hanzo-mcp's search actions run on.
hanzo-trigram = { git = "https://github.com/hanzoai/mcp", rev = "e9e6096ce46c3d3bf61e6830de6487a5b5b5b849" }
chrono = { workspace = true }
codex-core = { workspace = true }
codex-extension-api = { workspace = true }
codex-protocol = { workspace = true }
futures = { workspace = true, features = ["std"] }
indexmap = { workspace = true, features = ["serde"] }
regex = { workspace = true }
reqwest = { workspace = true, features = ["json"] }
serde = { workspace = true, features = ["derive"] }
# Option order is a question's meaning: a choice's labels are its options in order.
serde_json = { workspace = true, features = ["preserve_order"] }
sha2 = { workspace = true }
tempfile = { workspace = true }
tokio = { workspace = true, features = ["macros", "process", "rt", "rt-multi-thread", "sync", "time"] }
tracing = { workspace = true }
[dev-dependencies]
pretty_assertions = { workspace = true }
tempfile = { workspace = true }
wiremock = { workspace = true }
crates/hanzo-kai/programs/[email protected]
+43
View File
@@ -0,0 +1,43 @@
{
"id": "fix.turn@1",
"description": "After an attempt to fix a failing test, every read a turn needs at once: what happens next, and whether the attempt is done, correct, grounded, complete and safe. State: {goal, attempt, of, test: {passed, code, output}, patch}.",
"questions": {
"next": {
"type": "choice",
"instructions": "An attempt to fix a failing test ran the tests again. What should happen next?",
"criteria": {
"done": "the tests pass: the fix is finished",
"retry": "the tests still fail, and another attempt could fix them",
"escalate": "the tests still fail, and more attempts will not help: a person must look"
}
},
"done": {
"type": "noul",
"instructions": "Has the assistant finished everything the user asked, with nothing left to do or to check?"
},
"correct": {
"type": "noul",
"instructions": "Does the patch fix the cause of the failure the test reported, rather than only silencing the test?"
},
"grounded": {
"type": "noul",
"instructions": "Is the patch grounded in the failing output: does it change the code that output points to, or code that code depends on?"
},
"complete": {
"type": "noul",
"instructions": "Is the coding task complete: the requested change is made and the test that was failing now passes?"
},
"risk": {
"type": "choice",
"instructions": "The assistant applied this patch to the repository. Is it safe to keep without asking the user?",
"criteria": {
"allow": "Safe to keep: a small change to the code the failure points to.",
"ask": "Ask the user first: it changes more than the failure needs, or code unrelated to it.",
"deny": "Do not keep it: it edits tests, deletes code, weakens a check, or touches secrets."
}
}
},
"thresholds": {
"next": 0.5
}
}
crates/hanzo-kai/programs/locate.json
+145
View File
@@ -0,0 +1,145 @@
{
"names": [
"frame_file",
"frame_sym",
"called",
"name_in_diag",
"name_tok",
"test_tok",
"import_file",
"import_name",
"bm25",
"bm25_test",
"bm25_text",
"bm25_name",
"bm25_file",
"callee",
"callee2",
"class_called",
"ctor_called",
"dunder_op",
"parent_reached",
"dups",
"log_len",
"is_method",
"is_private",
"is_dunder",
"file_syms"
],
"mu": [
0.10470207500475918,
0.006980138333650613,
0.015578399644647503,
0.017291706326543563,
0.08850467911004764,
0.02007888916084238,
0.23101085094231869,
0.009708737864077669,
0.11943689918726805,
0.10303597976614734,
0.10910335264460333,
0.10002115004954805,
0.4750272907113233,
0.06145694523764198,
0.03735960403578907,
0.030776064471095883,
0.0016181229773462784,
0.03780379465702138,
0.002284408909194746,
1.1805330081355132,
2.0953102770785965,
0.746256107621042,
0.15625991496922392,
0.31616853861285615,
3.624674256606432
],
"sd": [
0.3061691534013262,
0.08325512598326555,
0.12383744631226044,
0.13035606322244755,
0.25369142940598965,
0.12345953815249955,
0.42147934432100514,
0.09805344600354622,
0.1702175095299602,
0.17252822092059408,
0.1788862047153644,
0.1816513331747839,
0.30243403621621956,
0.24016658618500525,
0.1896414090384921,
0.1727104464900932,
0.04019334093325854,
0.190721440238371,
0.04774086703370928,
1.3425269818289152,
0.8716236510486551,
0.43515276565757743,
0.36310157524194475,
0.4649795627815279,
0.8624830013428469
],
"w": [
0.025080951083837937,
0.2143809813786809,
0.04389633467581985,
-0.04211248328550059,
-0.009134958430156184,
0.14019169154899197,
0.06744148030140684,
-0.02301121977488332,
0.19530294743955018,
0.2714289097814333,
0.1281686122066849,
0.31622485908088965,
0.25750896395106243,
0.09844103617673333,
-0.005356574171962836,
0.1182462174797702,
0.025634486867147463,
0.15156383864643988,
0.05066298689129384,
-0.0207863118306674,
0.17221209738205603,
0.031599191377387495,
0.10064604210683334,
0.10068217190576402,
0.03259999314405524
],
"gamma": 0.99,
"escalate": {
"mu": [
0.3566452269710746,
0.23798092654887001,
1.4307183975561137,
0.45,
0.51,
1.572578390650942,
2.471884809273506,
5.739261113585301
],
"sd": [
0.2336897599888503,
0.2459166790215805,
0.5071126687200702,
0.4974937185533104,
0.4998999899979993,
0.6485649885129046,
0.42527839254627436,
0.3214577670160329
],
"w": [
0.6628252428377599,
0.09136153313352263,
-0.4951253770645405,
-0.07527733064688087,
-0.1703926001502168,
-0.05721521075769468,
0.7432816870565339,
-0.053098548954508246
],
"b": 1.4413764798936688,
"tau": 0.698
}
}
crates/hanzo-kai/src/decide.rs
+99 -50
View File
@@ -1,11 +1,13 @@
//! Kai answering over the Decisions API: `POST {url}/decisions`, one request per state, the
//! states of a decision in parallel.
//! Kai answering over the Decisions API: one request per state, the states of a decision in
//! parallel, over a [`Transport`] (`POST {url}/decisions` over HTTP, or the local engine's ZAP
//! socket when it serves Kai).
//!
//! Kai is optional. When it cannot be reached — no connection, a timeout, a refusal, a server
//! error, or no state of a decision answered — the decision fails at once, Kai is marked down
//! for [`COOLDOWN`], and every decision in that window fails without a request, so the loop runs
//! as it would without Kai. The first failure is logged; nothing else is until a decision asked
//! after it is answered.
//! error, or no state of a decision answered — the decision fails at once, that program is marked
//! down for [`COOLDOWN`], and every decision of it in that window fails without a request, so the
//! loop runs as it would without Kai. A program's outage is its own: one capability's timeout never
//! disables another (`site` timing out leaves `turn` asked). The first failure is logged; nothing
//! else is until a decision asked after it is answered.
use crate::program::Program;
use crate::program::Signal;
@@ -15,6 +17,10 @@ use hanzo_loop::Verdict;
use indexmap::IndexMap;
use serde_json::Value;
use serde_json::json;
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::Mutex;
use std::time::Duration;
use std::time::Instant;
@@ -26,6 +32,51 @@ const CONNECT: Duration = Duration::from_secs(2);
/// The longest one request may take.
const REQUEST: Duration = Duration::from_secs(10);
/// How a request reaches Kai and its answer comes back: the HTTP status and the body. The status
/// is the Decisions API's whatever the wire: 200, 400 or 422 for a request Kai refused, anything
/// else for a server that cannot answer. An error is a connection that did not hold.
pub trait Transport: Send + Sync {
fn send<'a>(&'a self, body: &'a Value) -> Pin<Box<dyn Future<Output = Result<(u16, String), String>> + Send + 'a>>;
}
/// `POST {url}/decisions` over HTTP, signed with `key`.
pub struct Http {
url: String,
key: Option<String>,
client: reqwest::Client,
}
impl Http {
pub fn new(url: &str, key: Option<String>) -> Http {
Http {
url: format!("{}/decisions", url.trim_end_matches('/')),
key,
client: reqwest::Client::builder()
.connect_timeout(CONNECT)
.timeout(REQUEST)
.pool_idle_timeout(Duration::from_secs(300))
.tcp_nodelay(true)
.build()
.unwrap_or_default(),
}
}
}
impl Transport for Http {
fn send<'a>(&'a self, body: &'a Value) -> Pin<Box<dyn Future<Output = Result<(u16, String), String>> + Send + 'a>> {
Box::pin(async move {
let mut request = self.client.post(&self.url).json(body);
if let Some(key) = &self.key {
request = request.bearer_auth(key);
}
let response = request.send().await.map_err(|e| e.without_url().to_string())?;
let status = response.status().as_u16();
let text = response.text().await.map_err(|e| e.without_url().to_string())?;
Ok((status, text))
})
}
}
/// What a decision is asked with: the operation's mode, how many states a selection keeps,
/// and the deterministic verdict a verdict rule joins.
#[derive(Debug, Clone, Copy, PartialEq)]
@@ -96,67 +147,79 @@ struct Health {
logged: bool,
}
/// Kai over HTTP, answering from any task.
/// Kai, answering from any task.
pub struct Decider {
url: String,
model: String,
key: Option<String>,
client: reqwest::Client,
health: Mutex<Health>,
transport: Arc<dyn Transport>,
/// Per program: one capability's outage is not another's.
health: Mutex<HashMap<String, Health>>,
}
impl Decider {
/// Kai at `url` (`https://api.hanzo.ai/v1`) as `model`, signed with `key`.
/// Kai at `url` (`https://api.hanzo.ai/v1`) as `model`, signed with `key`, over HTTP.
pub fn new(url: &str, model: &str, key: Option<String>) -> Decider {
Decider::over(url, model, Arc::new(Http::new(url, key)))
}
/// Kai as `model` over `transport`; `url` names it in the log.
pub fn over(url: &str, model: &str, transport: Arc<dyn Transport>) -> Decider {
Decider {
url: format!("{}/decisions", url.trim_end_matches('/')),
url: url.to_string(),
model: model.to_string(),
key,
client: reqwest::Client::builder()
.connect_timeout(CONNECT)
.timeout(REQUEST)
.build()
.unwrap_or_default(),
health: Mutex::new(Health::default()),
transport,
health: Mutex::new(HashMap::new()),
}
}
fn health(&self) -> std::sync::MutexGuard<'_, Health> {
fn health(&self) -> std::sync::MutexGuard<'_, HashMap<String, Health>> {
self.health
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
/// Whether Kai is down: a decision now fails without a request.
/// Whether some program of Kai's is down: the agent loop, whose programs fail together, skips
/// Kai while it is.
pub fn down(&self) -> bool {
let now = Instant::now();
self.health().values().any(|h| h.until.is_some_and(|until| now < until))
}
/// Whether `program` is down: a decision of it now fails without a request.
pub fn down_for(&self, program: &str) -> bool {
self.health()
.until
.get(program)
.and_then(|h| h.until)
.is_some_and(|until| Instant::now() < until)
}
fn fell(&self, why: &str) {
let mut h = self.health();
fn fell(&self, program: &str, why: &str) {
let mut all = self.health();
let h = all.entry(program.to_string()).or_default();
let now = Instant::now();
h.until = Some(now + COOLDOWN);
h.fell = Some(now);
if !h.logged {
h.logged = true;
tracing::warn!(
"kai: {} does not answer ({why}); deciding without Kai",
"kai: {} does not answer {program} ({why}); deciding without Kai",
self.url
);
}
}
/// A decision asked at `started` was answered. One asked before the last failure says
/// nothing about Kai now.
fn answered(&self, started: Instant) {
let mut h = self.health();
/// A decision of `program` asked at `started` was answered. One asked before the last failure
/// says nothing about Kai now.
fn answered(&self, program: &str, started: Instant) {
let mut all = self.health();
let Some(h) = all.get_mut(program) else {
return;
};
if h.fell.is_some_and(|fell| started <= fell) {
return;
}
if h.logged {
tracing::info!("kai: {} answers again", self.url);
tracing::info!("kai: {} answers {program} again", self.url);
}
*h = Health::default();
}
@@ -168,22 +231,8 @@ impl Decider {
"state": state,
"questions": program.questions,
});
let mut request = self.client.post(&self.url).json(&body);
if let Some(key) = &self.key {
request = request.bearer_auth(key);
}
let response = request
.send()
.await
.map_err(|e| Fault::Down(e.without_url().to_string()))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|e| Fault::Down(e.without_url().to_string()))?;
if status == reqwest::StatusCode::BAD_REQUEST
|| status == reqwest::StatusCode::UNPROCESSABLE_ENTITY
{
let (status, text) = self.transport.send(&body).await.map_err(Fault::Down)?;
if status == 400 || status == 422 {
let why: Value = serde_json::from_str(&text).unwrap_or_default();
let message = why["error"]["message"]
.as_str()
@@ -192,7 +241,7 @@ impl Decider {
.to_string();
return Err(Fault::Row(format!("{status}: {message}")));
}
if !status.is_success() {
if !(200..300).contains(&status) {
return Err(Fault::Down(status.to_string()));
}
let answer: Value = serde_json::from_str(&text)
@@ -227,7 +276,7 @@ impl Decider {
states: Vec<Value>,
call: Call,
) -> Result<Decision, String> {
if self.down() {
if self.down_for(&program.id) {
return Err("kai is down".into());
}
let started = Instant::now();
@@ -245,10 +294,10 @@ impl Decider {
_ => "no answer".to_string(),
});
if let Some(why) = down.or(silent) {
self.fell(&why);
self.fell(&program.id, &why);
return Err(format!("kai: {why}"));
}
self.answered(started);
self.answered(&program.id, started);
let mut model = String::new();
let results: Vec<Row> = asked
.into_iter()
crates/hanzo-kai/src/diagnostic.rs
+444
View File
@@ -0,0 +1,444 @@
//! What a failing run says, read out of its output: pytest's short and long tracebacks, Python
//! tracebacks, rustc and cargo test, and any `path:line[:col]`.
//!
//! The reading is kai-gateway's `diagnostic.ts`, line for line, so the ranker that was fit on its
//! output ([`crate::locate`]) sees the same frames, tests and names here.
use regex::Regex;
use std::collections::HashSet;
use std::sync::LazyLock;
/// Most identifiers kept.
const NAMES: usize = 64;
/// Most characters of error text kept.
const TEXT: usize = 2_000;
/// One place the output points at.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Frame {
pub file: String,
/// 1-based.
pub line: u32,
pub symbol: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Kind {
/// An assertion is the only evidence.
Assert,
/// A raised error other than an assertion.
Exception,
/// rustc errors.
Compile,
Unknown,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Diagnostic {
pub kind: Kind,
/// In the order the output names them; relative to the tree when under it.
pub frames: Vec<Frame>,
/// Failing test ids and names, as the runner printed them.
pub tests: Vec<String>,
/// Identifiers on the failing source lines and in the error messages, first spelling kept.
pub names: Vec<String>,
/// The error lines themselves, at most 2,000 characters.
pub text: String,
}
impl Diagnostic {
/// Whether the output holds a diagnostic at all: a failure of a known kind, a frame, or a
/// failing test.
pub fn found(&self) -> bool {
self.kind != Kind::Unknown || !self.frames.is_empty() || !self.tests.is_empty()
}
}
fn re(pattern: &str) -> Regex {
Regex::new(pattern).expect("a fixed pattern")
}
static PYTEST_FRAME: LazyLock<Regex> = LazyLock::new(|| {
re(r"^(\S+\.py):(\d+):(?:\s+(?:in\s+(\S+)|([A-Za-z_][\w.]*)))?\s*$")
});
static TRACEBACK: LazyLock<Regex> =
LazyLock::new(|| re(r#"^\s*File "([^"]+)", line (\d+)(?:, in (\S+))?"#));
static RUSTC_LEVEL: LazyLock<Regex> =
LazyLock::new(|| re(r"^(error|warning|help|note)(?:\[(E\d+)\])?: (.*)$"));
static RUSTC_SPAN: LazyLock<Regex> = LazyLock::new(|| re(r"^\s*--> (\S+?):(\d+):(\d+)\s*$"));
static RUSTC_CODE: LazyLock<Regex> = LazyLock::new(|| re(r"^\s*\d+\s*\|(.*)$"));
static RUSTC_LABEL: LazyLock<Regex> = LazyLock::new(|| re(r"^\s*(?:=|\|)"));
static PANIC: LazyLock<Regex> = LazyLock::new(|| {
re(r"^thread '([^']+)'(?: \(\d+\))? panicked at (?:'(.*)', )?(\S+?):(\d+):(\d+):?\s*$")
});
static CARGO_TEST: LazyLock<Regex> = LazyLock::new(|| {
re(r"^(?:---- (\S+) stdout ----|(?:test )?(\S+) (?:\.\.\.|---) FAILED)\s*$")
});
static PYTEST_SUMMARY: LazyLock<Regex> =
LazyLock::new(|| re(r"^(?:FAILED|ERROR) (\S+?)(?: - (.*))?\s*$"));
static PYTEST_HEADER: LazyLock<Regex> = LazyLock::new(|| re(r"^_{3,} (\S+) _{3,}\s*$"));
static E_LINE: LazyLock<Regex> = LazyLock::new(|| re(r"^E\s+(.*)$"));
static RAISED: LazyLock<Regex> = LazyLock::new(|| {
re(r"^([A-Za-z_][\w.]*(?:Error|Exception|Exit|Interrupt|Failure))(?::\s.*)?$")
});
static GENERIC: LazyLock<Regex> = LazyLock::new(|| {
re(r#"(?:^|[\s("'`=])((?:[\w.-]+/)*[\w.-]*[\w-]\.[A-Za-z][A-Za-z0-9]{0,7}):(\d+)(?::\d+)?"#)
});
static IDENT: LazyLock<Regex> = LazyLock::new(|| re(r"[A-Za-z_][A-Za-z0-9_]*"));
static ASSERTIVE: LazyLock<Regex> =
LazyLock::new(|| re(r"^(?:assert\b|AssertionError\b|assertion\b)"));
static ERROR_CODE: LazyLock<Regex> = LazyLock::new(|| re(r"^e\d+$"));
static INDENTED: LazyLock<Regex> = LazyLock::new(|| re(r"^\s+\S"));
static COUNT: LazyLock<Regex> = LazyLock::new(|| re(r"\b(\d+) (?:failed|failures?|errors?)\b"));
/// How many tests a run reports failing: the numbers on the runner's last summary line that
/// counts them (pytest's `2 failed, 3 passed`, cargo's `test result: FAILED. 3 passed; 2 failed`),
/// else the failing tests it names. Zero when it says neither.
pub fn failures(output: &str) -> usize {
let counted = output.lines().rev().find(|l| COUNT.is_match(l));
match counted {
Some(line) => COUNT.captures_iter(line).filter_map(|c| c[1].parse::<usize>().ok()).sum(),
None => parse(output, None).tests.iter().filter(|t| t.contains("::") || t.contains(' ')).count(),
}
}
/// Keywords and the runners' own words: on every failing line, and no help in finding the fix.
const NOISE: &[&str] = &[
"and", "as", "assert", "async", "await", "break", "class", "continue", "def", "del", "elif",
"else", "except", "finally", "for", "from", "global", "if", "import", "in", "is", "lambda",
"nonlocal", "not", "or", "pass", "raise", "return", "try", "while", "with", "yield", "none",
"true", "false", "self", "cls", "fn", "let", "mut", "impl", "pub", "use", "mod", "struct",
"enum", "match", "ref", "move", "crate", "super", "loop", "const", "static", "type", "unsafe",
"dyn", "trait", "where", "assertionerror", "assertion", "failed", "left", "right", "thread",
"panicked", "error", "warning", "note", "the", "an", "to", "of", "at", "by", "be", "on", "it",
"this", "that", "but", "was", "were", "are",
];
fn assertive(message: &str) -> bool {
ASSERTIVE.is_match(message.trim())
}
/// A path as the tree spells it: the tree's own absolute prefix and `./` dropped.
fn relative(file: &str, root: Option<&str>) -> String {
let mut path = file;
if let Some(root) = root {
let base = format!("{}/", root.trim_end_matches('/'));
if let Some(rest) = path.strip_prefix(&base) {
path = rest;
}
}
let mut path = path;
while let Some(rest) = path.strip_prefix("./") {
path = rest;
}
path.to_string()
}
/// The parse in progress.
#[derive(Default)]
struct Reading {
frames: Vec<Frame>,
tests: Vec<String>,
names: Vec<String>,
messages: Vec<String>,
seen_frames: HashSet<String>,
seen_tests: HashSet<String>,
seen_names: HashSet<String>,
seen_messages: HashSet<String>,
asserts: bool,
exceptions: bool,
compile: bool,
}
impl Reading {
fn frame(&mut self, file: &str, line: &str, symbol: Option<&str>, root: Option<&str>) {
let line: u32 = line.parse().unwrap_or(0);
let at = Frame {
file: relative(file, root),
line,
symbol: symbol.map(str::to_string),
};
let key = format!("{}:{}", at.file, at.line);
if at.line < 1 || !self.seen_frames.insert(key) {
return;
}
self.frames.push(at);
}
fn test(&mut self, id: &str) {
if self.seen_tests.insert(id.to_string()) {
self.tests.push(id.to_string());
}
}
fn named(&mut self, line: &str) {
for found in IDENT.find_iter(line) {
let name = found.as_str();
let lower = name.to_lowercase();
// rustc's error codes (E0308) name no code of the tree's.
if self.names.len() >= NAMES
|| name.len() < 2
|| NOISE.contains(&lower.as_str())
|| self.seen_names.contains(&lower)
|| ERROR_CODE.is_match(&lower)
{
continue;
}
self.seen_names.insert(lower);
self.names.push(name.to_string());
}
}
fn message(&mut self, line: &str) {
let trimmed = line.trim();
if trimmed.is_empty() {
return;
}
self.named(trimmed);
if self.seen_messages.insert(trimmed.to_string()) {
self.messages.push(trimmed.to_string());
}
}
/// An error the run reported: an assertion, or any other raised type.
fn classify(&mut self, line: &str) {
if assertive(line) {
self.asserts = true;
} else if RAISED.is_match(line.trim()) {
self.exceptions = true;
}
}
}
/// The diagnostic in `text`; paths under `root` are made relative to it.
pub fn parse(text: &str, root: Option<&str>) -> Diagnostic {
let mut r = Reading::default();
// rustc's level for the spans and code lines under its header: only an error's are the
// failure. An `error:` without a code counts once a span follows it.
let mut level: Option<bool> = None; // Some(true): error, Some(false): other
let mut pending: Option<String> = None;
// pytest and Python print the failing source line right after a frame.
let mut code_next = false;
// A panic's message follows it, up to a blank line; its first line says what kind it was.
let mut panic: Option<bool> = None; // Some(true): first, Some(false): rest
for line in text.split('\n') {
let line = line.strip_suffix('\r').unwrap_or(line);
let after_frame = code_next;
code_next = false;
if let Some(m) = RUSTC_LEVEL.captures(line) {
let error = &m[1] == "error";
level = Some(error);
pending = None;
if error && m.get(2).is_some() {
r.compile = true;
r.message(line);
} else if error {
pending = Some(line.to_string());
}
continue;
}
if let Some(m) = RUSTC_SPAN.captures(line) {
if level == Some(true) {
r.compile = true;
r.frame(&m[1], &m[2], None, root);
if let Some(p) = pending.take() {
r.message(&p);
}
}
continue;
}
if level.is_some()
&& let Some(m) = RUSTC_CODE.captures(line)
{
if level == Some(true) {
r.named(&m[1]);
}
continue;
}
if level.is_some() && RUSTC_LABEL.is_match(line) {
continue;
}
if line.trim().is_empty() {
level = None;
pending = None;
}
if let Some(m) = PANIC.captures(line) {
if &m[1] != "main" {
r.test(&m[1]);
}
r.frame(&m[3], &m[4], None, root);
match m.get(2) {
None => panic = Some(true),
Some(msg) => {
if assertive(msg.as_str()) {
r.asserts = true;
} else {
r.exceptions = true;
}
r.message(msg.as_str());
}
}
continue;
}
if let Some(first) = panic {
if line.trim().is_empty() || line.starts_with("note:") {
panic = None;
} else {
if first {
if assertive(line) {
r.asserts = true;
} else {
r.exceptions = true;
}
}
panic = Some(false);
r.message(line);
continue;
}
}
if let Some(m) = CARGO_TEST.captures(line) {
let id = m.get(1).or_else(|| m.get(2)).map_or("", |g| g.as_str());
r.test(id);
continue;
}
if let Some(m) = PYTEST_FRAME.captures(line) {
r.frame(&m[1], &m[2], m.get(3).map(|g| g.as_str()), root);
if let Some(kind) = m.get(4) {
r.classify(kind.as_str());
}
code_next = m.get(3).is_some();
continue;
}
if let Some(m) = TRACEBACK.captures(line) {
r.frame(&m[1], &m[2], m.get(3).map(|g| g.as_str()), root);
code_next = true;
continue;
}
if let Some(m) = PYTEST_SUMMARY.captures(line) {
r.test(&m[1]);
if let Some(msg) = m.get(2) {
r.classify(msg.as_str());
r.message(msg.as_str());
}
continue;
}
if let Some(m) = PYTEST_HEADER.captures(line) {
r.test(&m[1]);
continue;
}
if let Some(m) = E_LINE.captures(line) {
r.classify(&m[1]);
r.message(&m[1]);
continue;
}
// pytest's long traceback marks the failing line with `>`; the short one and Python's
// put it after the frame.
if line.starts_with('>') || (after_frame && INDENTED.is_match(line)) {
r.named(line.strip_prefix('>').unwrap_or(line));
continue;
}
if RAISED.is_match(line) {
r.classify(line);
r.message(line);
continue;
}
for m in GENERIC.captures_iter(line) {
// The path is not followed by a word character or a slash.
let end = m.get(0).map_or(0, |g| g.end());
let next = line[end..].chars().next();
if next.is_some_and(|c| c.is_alphanumeric() || c == '_' || c == '/') {
continue;
}
r.frame(&m[1], &m[2], None, root);
}
}
let kind = if r.compile {
Kind::Compile
} else if r.exceptions {
Kind::Exception
} else if r.asserts {
Kind::Assert
} else {
Kind::Unknown
};
let mut text = r.messages.join("\n");
if let Some((at, _)) = text.char_indices().nth(TEXT) {
text.truncate(at);
}
Diagnostic {
kind,
frames: r.frames,
tests: r.tests,
names: r.names,
text,
}
}
#[cfg(test)]
mod tests {
use super::*;
use pretty_assertions::assert_eq;
#[test]
fn failures_are_the_runners_own_count() {
assert_eq!(failures("....F.\n== 2 failed, 3 passed in 0.1s =="), 2);
assert_eq!(failures("1 failed, 4 passed, 1 error in 0.16s"), 2);
assert_eq!(failures("test result: FAILED. 3 passed; 2 failed; 0 ignored"), 2);
assert_eq!(failures("all good"), 0);
}
#[test]
fn a_pytest_short_traceback_names_its_frames_tests_and_error() {
let out = "....F [100%]\n=== FAILURES ===\n___ test_histogram_zero_interquartile_range ___\ntests/test_statsutils.py:44: in test_histogram_zero_interquartile_range\n counts = Stats(data).get_histogram_counts()\n ^^^^^^^^^^^^^^^^^^^^^^\nboltons/statsutils.py:631: in get_histogram_counts\n idxs = [bisect.bisect(d, bins) - 1 for d in self.data]\nE TypeError: object of type 'int' has no len()\n=== short test summary info ===\nFAILED tests/test_statsutils.py::test_histogram_zero_interquartile_range - Ty...\n1 failed, 4 passed in 0.16s\n";
let d = parse(out, None);
assert_eq!(d.kind, Kind::Exception);
let frames: Vec<_> = d.frames.iter().map(|f| (f.file.as_str(), f.line)).collect();
assert_eq!(
frames,
vec![("tests/test_statsutils.py", 44), ("boltons/statsutils.py", 631)]
);
assert_eq!(
d.tests,
vec![
"test_histogram_zero_interquartile_range",
"tests/test_statsutils.py::test_histogram_zero_interquartile_range"
]
);
assert!(d.names.iter().any(|n| n == "get_histogram_counts"));
assert!(d.text.contains("TypeError: object of type 'int' has no len()"));
}
#[test]
fn rustc_errors_are_a_compile_failure_at_their_span() {
let out = "error[E0308]: mismatched types\n --> src/lib.rs:7:5\n |\n7 | 1u8\n | ^^^ expected `u32`\n\nwarning: unused\n --> src/other.rs:1:1\n";
let d = parse(out, None);
assert_eq!(d.kind, Kind::Compile);
assert_eq!(d.frames.len(), 1);
assert_eq!((d.frames[0].file.as_str(), d.frames[0].line), ("src/lib.rs", 7));
}
#[test]
fn a_panic_is_an_exception_unless_it_asserts() {
let out = "---- add::tests::sums stdout ----\nthread 'add::tests::sums' (12) panicked at src/add.rs:9:5:\nassertion `left == right` failed\n left: 2\n right: 3\n";
let d = parse(out, None);
assert_eq!(d.kind, Kind::Assert);
assert_eq!(d.tests, vec!["add::tests::sums"]);
assert_eq!((d.frames[0].file.as_str(), d.frames[0].line), ("src/add.rs", 9));
}
#[test]
fn a_path_and_line_anywhere_is_a_frame_but_not_a_longer_word() {
let d = parse("see src/x.py:12 and src/y.py:3x\n", None);
let frames: Vec<_> = d.frames.iter().map(|f| (f.file.as_str(), f.line)).collect();
assert_eq!(frames, vec![("src/x.py", 12)]);
}
#[test]
fn the_trees_absolute_prefix_is_dropped() {
let d = parse("/w/t/pkg/a.py:5: in f\n", Some("/w/t"));
assert_eq!(d.frames[0].file, "pkg/a.py");
}
}
crates/hanzo-kai/src/edit.rs
+302
View File
@@ -0,0 +1,302 @@
//! Edits: the smallest change a model can state, and the tool that makes it.
//!
//! An edit names a file, the line its `old` text starts on, that text exactly (a part of a line, or
//! several lines) and its replacement. It is checked against the file, never against a diff's
//! line counts, so a model that miscounts a hunk header costs nothing. A test file is never
//! edited, and a path leaves the tree never.
use crate::symbols::is_test;
use serde_json::Value;
use std::collections::BTreeMap;
use std::path::Path;
/// Most files one answer may change.
const FILES: usize = 10;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Edit {
pub file: String,
/// The line `old` starts on, 1-based; a hint when the text stands in several places.
pub line: u32,
pub old: String,
pub new: String,
}
/// The edits in a model's answer: `{"edits": [...]}`, one edit object, or a list of them, as the
/// first JSON value in the text (fences and prose around it are ignored).
pub fn parse(answer: &str) -> Result<Vec<Edit>, String> {
let at = answer
.find(['{', '['])
.ok_or_else(|| "the answer holds no JSON".to_string())?;
let value: Value = serde_json::Deserializer::from_str(&answer[at..])
.into_iter::<Value>()
.next()
.ok_or_else(|| "the answer holds no JSON".to_string())?
.map_err(|e| format!("the answer's JSON does not parse: {e}"))?;
let list: Vec<Value> = match &value {
Value::Array(a) => a.clone(),
Value::Object(o) if o.contains_key("edits") => value["edits"]
.as_array()
.cloned()
.ok_or_else(|| "`edits` is not a list".to_string())?,
Value::Object(_) => vec![value.clone()],
_ => return Err("the answer's JSON is not an edit".into()),
};
if list.is_empty() {
return Err("no edits".into());
}
list.iter().map(edit).collect()
}
fn edit(v: &Value) -> Result<Edit, String> {
let text = |key: &str| {
v[key]
.as_str()
.map(str::to_string)
.ok_or_else(|| format!("an edit lacks `{key}`"))
};
let line = v["line"]
.as_u64()
.or_else(|| v["start"].as_u64())
.or_else(|| v["line"].as_str().and_then(|s| s.parse().ok()))
.unwrap_or(0);
Ok(Edit {
file: text("file")?,
line: u32::try_from(line).unwrap_or(0),
old: text("old")?,
new: text("new")?,
})
}
/// A change made to the tree, with what it replaced.
#[derive(Debug, Clone, Default)]
pub struct Applied {
/// File -> its text before the change.
before: BTreeMap<String, String>,
/// One line an edit: `file:line -old +new`, for a retry's prompt and the report.
pub summary: Vec<String>,
}
impl Applied {
pub fn files(&self) -> Vec<String> {
self.before.keys().cloned().collect()
}
}
fn safe(file: &str) -> bool {
!file.is_empty()
&& !file.starts_with('/')
&& !file.contains('\0')
&& !file.split('/').any(|p| p == ".." || p.is_empty() || p.starts_with('.'))
}
/// The byte offsets `old` starts at in `text`.
fn occurrences(text: &str, old: &str) -> Vec<usize> {
text.match_indices(old).map(|(i, _)| i).collect()
}
fn line_of(text: &str, offset: usize) -> u32 {
text[..offset].bytes().filter(|b| *b == b'\n').count() as u32 + 1
}
/// Where `old` stands in `text`: its one place, or of several the nearest to `line` when that
/// is clear. Trailing whitespace on a line is not compared.
fn locate(text: &str, edit: &Edit) -> Result<(usize, usize), String> {
let mut at = occurrences(text, &edit.old);
let mut len = edit.old.len();
if at.is_empty() {
// The model's lines may differ from the file's in trailing space only.
let strip = |s: &str| s.split('\n').map(str::trim_end).collect::<Vec<_>>().join("\n");
let want = strip(&edit.old);
if want.is_empty() {
return Err(format!("the old text of {} is empty", edit.file));
}
let lines: Vec<&str> = text.split('\n').collect();
let n = want.split('\n').count();
let mut found = Vec::new();
let mut offset = 0;
for i in 0..lines.len() {
if i + n <= lines.len() && strip(&lines[i..i + n].join("\n")) == want {
found.push((offset, lines[i..i + n].join("\n").len()));
}
offset += lines[i].len() + 1;
}
return match found.as_slice() {
[one] => Ok(*one),
[] => Err(format!(
"the old text of {}:{} does not stand in the file",
edit.file, edit.line
)),
_ => nearest(text, edit, found.iter().map(|f| f.0).collect(), found[0].1),
};
}
if at.len() == 1 {
return Ok((at.remove(0), len));
}
len = edit.old.len();
nearest(text, edit, at, len)
}
fn nearest(text: &str, edit: &Edit, at: Vec<usize>, len: usize) -> Result<(usize, usize), String> {
let mut by: Vec<(u32, usize)> = at
.iter()
.map(|o| (line_of(text, *o).abs_diff(edit.line), *o))
.collect();
by.sort();
if by.len() > 1 && by[0].0 == by[1].0 {
return Err(format!(
"the old text of {}:{} stands in {} places; name a longer text",
edit.file,
edit.line,
by.len()
));
}
Ok((by[0].1, len))
}
fn clip(s: &str) -> String {
let one = s.replace('\n', "\\n");
match one.char_indices().nth(160) {
Some((i, _)) => format!("{}...", &one[..i]),
None => one,
}
}
/// Makes `edits` in the tree at `root`, all or none: a change that cannot be made leaves every
/// file as it was.
pub fn apply(root: &Path, edits: &[Edit]) -> Result<Applied, String> {
let mut by: BTreeMap<&str, Vec<&Edit>> = BTreeMap::new();
for e in edits {
if !safe(&e.file) {
return Err(format!("{:?} is not a path in the tree", e.file));
}
if is_test(&e.file) {
return Err(format!("{} is a test file; tests are not edited", e.file));
}
if e.old.is_empty() {
return Err(format!("the old text of {}:{} is empty", e.file, e.line));
}
by.entry(&e.file).or_default().push(e);
}
if by.len() > FILES {
return Err(format!("{} files in one answer; at most {FILES}", by.len()));
}
let mut applied = Applied::default();
let mut after: BTreeMap<String, String> = BTreeMap::new();
for (file, es) in by {
let path = root.join(file);
let meta = std::fs::symlink_metadata(&path).map_err(|_| format!("no such file {file}"))?;
if !meta.is_file() {
return Err(format!("{file} is not a regular file"));
}
let original = std::fs::read_to_string(&path).map_err(|e| format!("{file}: {e}"))?;
let mut text = original.clone();
// Edits of one file from the bottom up, so an earlier edit's lines do not move.
let mut placed = Vec::new();
for e in &es {
let (start, len) = locate(&original, e)?;
placed.push((start, len, *e));
}
placed.sort_by_key(|p| std::cmp::Reverse(p.0));
for pair in placed.windows(2) {
if pair[1].0 + pair[1].1 > pair[0].0 {
return Err(format!("two edits of {file} overlap"));
}
}
for (start, len, e) in placed {
applied.summary.push(format!(
"{file}:{} -{} +{}",
line_of(&original, start),
clip(&e.old),
clip(&e.new)
));
text.replace_range(start..start + len, &e.new);
}
applied.before.insert(file.to_string(), original);
after.insert(file.to_string(), text);
}
for (file, text) in &after {
std::fs::write(root.join(file), text).map_err(|e| format!("{file}: {e}"))?;
}
Ok(applied)
}
/// Puts the files `applied` changed back. False when a file cannot be written.
pub fn revert(root: &Path, applied: &Applied) -> bool {
applied
.before
.iter()
.all(|(file, text)| std::fs::write(root.join(file), text).is_ok())
}
#[cfg(test)]
mod tests {
use super::*;
use pretty_assertions::assert_eq;
fn tree(files: &[(&str, &str)]) -> tempfile::TempDir {
let dir = tempfile::tempdir().unwrap();
for (name, text) in files {
let path = dir.path().join(name);
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
std::fs::write(path, text).unwrap();
}
dir
}
#[test]
fn an_answer_with_prose_and_fences_around_its_json_parses() {
let a = "Here:\n```json\n{\"edits\":[{\"file\":\"a.py\",\"line\":3,\"old\":\"x - 1\",\"new\":\"x + 1\"}]}\n```\n";
assert_eq!(
parse(a).unwrap(),
vec![Edit { file: "a.py".into(), line: 3, old: "x - 1".into(), new: "x + 1".into() }]
);
let single = "{\"file\":\"a.py\",\"line\":\"3\",\"old\":\"a\",\"new\":\"b\"}";
assert_eq!(parse(single).unwrap()[0].line, 3);
assert!(parse("no json here").is_err());
assert!(parse("{\"edits\":[]}").is_err());
}
#[test]
fn a_part_of_a_line_is_replaced_and_reverted_exactly() {
let dir = tree(&[("pkg/a.py", "def f(x):\n return x - 1\n")]);
let e = Edit { file: "pkg/a.py".into(), line: 2, old: "x - 1".into(), new: "x + 1".into() };
let applied = apply(dir.path(), &[e]).unwrap();
assert_eq!(std::fs::read_to_string(dir.path().join("pkg/a.py")).unwrap(), "def f(x):\n return x + 1\n");
assert_eq!(applied.summary, vec!["pkg/a.py:2 -x - 1 +x + 1".to_string()]);
assert!(revert(dir.path(), &applied));
assert_eq!(std::fs::read_to_string(dir.path().join("pkg/a.py")).unwrap(), "def f(x):\n return x - 1\n");
}
#[test]
fn the_nearest_place_wins_and_a_tie_is_refused() {
let dir = tree(&[("a.py", "a = 1\nb = 2\na = 1\n")]);
let near = Edit { file: "a.py".into(), line: 3, old: "a = 1".into(), new: "a = 9".into() };
apply(dir.path(), &[near]).unwrap();
assert_eq!(std::fs::read_to_string(dir.path().join("a.py")).unwrap(), "a = 1\nb = 2\na = 9\n");
let dir = tree(&[("a.py", "a = 1\nb = 2\na = 1\n")]);
let tie = Edit { file: "a.py".into(), line: 2, old: "a = 1".into(), new: "a = 9".into() };
let e = apply(dir.path(), &[tie]).unwrap_err();
assert!(e.contains("2 places"), "{e}");
}
#[test]
fn trailing_space_is_not_compared_and_nothing_changes_on_a_refusal() {
let dir = tree(&[("a.py", "if x: \n y = 1\n"), ("b.py", "z = 0\n")]);
let ok = Edit { file: "a.py".into(), line: 1, old: "if x:\n y = 1".into(), new: "if x:\n y = 2".into() };
let bad = Edit { file: "b.py".into(), line: 1, old: "nope".into(), new: "z = 1".into() };
assert!(apply(dir.path(), &[ok.clone(), bad]).is_err());
assert_eq!(std::fs::read_to_string(dir.path().join("a.py")).unwrap(), "if x: \n y = 1\n");
apply(dir.path(), &[ok]).unwrap();
assert_eq!(std::fs::read_to_string(dir.path().join("a.py")).unwrap(), "if x:\n y = 2\n");
}
#[test]
fn tests_and_paths_outside_the_tree_are_never_edited() {
let dir = tree(&[("tests/test_a.py", "assert 1\n"), ("a.py", "x\n")]);
for file in ["tests/test_a.py", "../a.py", "/etc/passwd", ".git/config"] {
let e = Edit { file: file.into(), line: 1, old: "x".into(), new: "y".into() };
assert!(apply(dir.path(), &[e]).is_err(), "{file}");
}
}
}
crates/hanzo-kai/src/find.rs
+81
View File
@@ -0,0 +1,81 @@
//! Evidence from search: where a message, a name or a string in a failing run is in the tree.
//!
//! The loop asks a [`Finder`] and nothing else; the engine behind it is Hanzo's one search, tgrep
//! (`hanzo-trigram`, the library hanzo-mcp's `fs search_text` and `code search_symbol` run on),
//! linked in process. A search is a call, not a spawn and not an MCP hop.
use std::path::Path;
/// A line of the tree that spells what was asked for.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Found {
/// Relative to the root.
pub file: String,
/// 1-based.
pub line: u32,
/// The line reads like a declaration of the symbol asked for.
pub definition: bool,
pub excerpt: String,
}
/// Text and symbol search over a tree, best first.
pub trait Finder: Send + Sync {
/// Lines matching the literal `text`.
fn text(&self, root: &Path, text: &str, limit: usize) -> Result<Vec<Found>, String>;
/// Lines spelling the identifier `name`, its definitions first.
fn symbol(&self, root: &Path, name: &str, limit: usize) -> Result<Vec<Found>, String>;
}
/// tgrep, in process.
pub struct Tgrep;
fn found(hits: Vec<hanzo_trigram::Hit>) -> Vec<Found> {
hits.into_iter()
.filter(|h| h.entry.is_none() && h.place.is_none())
.map(|h| Found {
file: h.path.to_string_lossy().into_owned(),
line: u32::try_from(h.line).unwrap_or(0),
definition: h.mark == hanzo_trigram::Mark::Definition,
excerpt: h.excerpt,
})
.collect()
}
impl Finder for Tgrep {
fn text(&self, root: &Path, text: &str, limit: usize) -> Result<Vec<Found>, String> {
let mut q = hanzo_trigram::Query::text(root, text);
q.fixed = true;
q.limit = Some(limit);
hanzo_trigram::search(&q).map(found).map_err(|e| e.to_string())
}
fn symbol(&self, root: &Path, name: &str, limit: usize) -> Result<Vec<Found>, String> {
let mut q = hanzo_trigram::Query::symbol(root, name);
q.limit = Some(limit);
hanzo_trigram::search(&q).map(found).map_err(|e| e.to_string())
}
}
#[cfg(test)]
mod tests {
use super::*;
use pretty_assertions::assert_eq;
#[test]
fn tgrep_finds_the_line_that_spells_a_message_and_the_definition_of_a_name() {
let dir = tempfile::tempdir().unwrap();
std::fs::create_dir_all(dir.path().join("pkg")).unwrap();
std::fs::write(
dir.path().join("pkg/parse.py"),
"def parse_flag(raw):\n if not raw:\n raise ValueError(\"flag must not be empty\")\n return raw\n",
)
.unwrap();
std::fs::write(dir.path().join("pkg/use.py"), "from pkg.parse import parse_flag\nparse_flag('x')\n").unwrap();
let text = Tgrep.text(dir.path(), "flag must not be empty", 5).unwrap();
assert_eq!(text.len(), 1, "{text:?}");
assert_eq!((text[0].file.as_str(), text[0].line), ("pkg/parse.py", 3));
let symbol = Tgrep.symbol(dir.path(), "parse_flag", 5).unwrap();
assert!(symbol[0].definition && symbol[0].file == "pkg/parse.py", "{symbol:?}");
assert!(symbol.iter().any(|f| !f.definition && f.file == "pkg/use.py"));
}
}
crates/hanzo-kai/src/flow.rs
+1291 -197
View File
File diff suppressed because it is too large Load Diff
crates/hanzo-kai/src/flow_tests.rs
+338 -143
View File
@@ -6,32 +6,25 @@ use pretty_assertions::assert_eq;
use std::process::Command;
use wiremock::MockServer;
const BUGGY: &str = "echo $(( $1 - $2 ))\n";
const FIXED: &str = "echo $(( $1 + $2 ))\n";
/// A fault no single change of the generic transforms mends (`^` is not among their operators): the
/// search finds nothing and the model writes the edit.
const BUGGY: &str = "def add(a, b):\n return a ^ b\n\n\ndef mul(a, b):\n return a * b\n";
const FIXED: &str = "def add(a, b):\n return a + b\n\n\ndef mul(a, b):\n return a * b\n";
/// A fault one change mends.
const SIMPLE: &str = "def add(a, b):\n return a - b\n\n\ndef mul(a, b):\n return a * b\n";
const TEST: &str = "from add import add\n\n\ndef test_add():\n assert add(2, 3) == 5\n\n\ntest_add()\n";
/// A repository whose test fails: `src/add.sh` subtracts.
fn repo() -> tempfile::TempDir {
/// A repository whose test fails: `src/add.py` is wrong.
fn repo_with(source: &str) -> tempfile::TempDir {
let dir = tempfile::tempdir().unwrap();
let root = dir.path();
std::fs::create_dir_all(root.join("src")).unwrap();
std::fs::write(root.join("src/add.sh"), BUGGY).unwrap();
std::fs::write(root.join("src/mul.sh"), "echo $(( $1 * $2 ))\n").unwrap();
std::fs::write(root.join("README.md"), "Shell arithmetic.\n").unwrap();
std::fs::write(
root.join("test.sh"),
"r=$(sh src/add.sh 2 3)\n[ \"$r\" = 5 ] || { echo \"FAIL src/add.sh: add 2 3 gave $r, want 5\"; exit 1; }\necho ok\n",
)
.unwrap();
std::fs::create_dir_all(root.join("tests")).unwrap();
std::fs::write(root.join("src/add.py"), source).unwrap();
std::fs::write(root.join("tests/test_add.py"), TEST).unwrap();
let git = |args: &[&str]| {
let ok = Command::new("git")
.args([
"-c",
"user.name=t",
"-c",
"user.email=t@t",
"-c",
"commit.gpgsign=false",
])
.args(["-c", "user.name=t", "-c", "user.email=t@t", "-c", "commit.gpgsign=false"])
.args(args)
.current_dir(root)
.output()
@@ -46,37 +39,37 @@ fn repo() -> tempfile::TempDir {
dir
}
fn diff(from: &str, to: &str) -> String {
format!(
"```diff\n--- a/src/add.sh\n+++ b/src/add.sh\n@@ -1 +1 @@\n-{}\n+{}\n```\n",
from.trim_end(),
to.trim_end()
)
fn repo() -> tempfile::TempDir {
repo_with(BUGGY)
}
/// Kai rating `src/add.sh` the file to change, and answering `next` by `next(test output)`.
async fn kai(next: impl Fn(&Value) -> (&'static str, f64) + Send + Sync + 'static) -> MockServer {
testing::kai(move |_, q, state| match q.kind {
Kind::Noul => {
let p = if state["file"]["path"] == "src/add.sh" {
0.9
} else {
0.2
};
vec![1.0 - p, p]
}
Kind::Choice => {
fn edit_json(old: &str, new: &str) -> String {
json!({"edits": [{"file": "src/add.py", "line": 2, "old": old, "new": new}]}).to_string()
}
fn fix() -> String {
edit_json("a ^ b", "a + b")
}
/// Kai choosing site `c{site}` and answering `next` by `next(test)`.
async fn kai(site: usize, next: impl Fn(&Value) -> (&'static str, f64) + Send + Sync + 'static) -> MockServer {
testing::kai(move |id, q, state| match (id, q.kind) {
("site", _) => on(q, &format!("c{site}"), 0.9),
("action", _) => on(q, "search", 0.9),
("effort", _) => on(q, "low", 0.9),
(_, Kind::Noul) => vec![0.5, 0.5],
(_, Kind::Choice) => {
let (label, p) = next(&state["test"]);
on(q, label, p)
}
Kind::Score => on(q, "0", 0.9),
(_, Kind::Score) => on(q, "0", 0.9),
})
.await
}
/// Kai as an honest judge of the tests.
async fn honest() -> MockServer {
kai(|test| {
kai(0, |test| {
if test["passed"] == true {
("done", 0.9)
} else {
@@ -86,21 +79,27 @@ async fn honest() -> MockServer {
.await
}
/// Kai and Zen at their servers, the trace in the repository's `.git`.
/// Kai and Zen at their servers, the trace and the ledger in the repository's `.git`.
fn setup(root: &Path, kai: &str, zen: &MockServer) -> Setup {
let home = root.join(".git/dev");
std::fs::create_dir_all(&home).unwrap();
std::fs::write(
home.join("kai.toml"),
format!("url = \"{kai}/v1\"\ntrace = \"kai.jsonl\"\n"),
format!(
"url = \"{kai}/v1\"\ntrace = \"kai.jsonl\"\n[generate]\nurl = \"{}/v1\"\nmodel = \"zen-test\"\n",
zen.uri()
),
)
.unwrap();
Setup::new(&home, &format!("{}/v1", zen.uri()), None).unwrap()
Setup::new(&home, || None).unwrap()
}
fn options(max: u32, control: Control) -> Options {
Options {
task: json!({"test": "sh test.sh", "goal": "make the test pass"}),
task: json!({"test": "PYTHONPATH=src python3 tests/test_add.py", "goal": "make the test pass"}),
failure: None,
search: true,
discover: true,
max,
control,
limit: Duration::from_secs(20),
@@ -117,118 +116,203 @@ fn trace(root: &Path) -> PathBuf {
root.join(".git/dev/kai.jsonl")
}
fn read(root: &Path) -> String {
std::fs::read_to_string(root.join("src/add.sh")).unwrap()
fn ledger(root: &Path) -> Vec<Value> {
std::fs::read_to_string(root.join(".git/dev/ledger.jsonl"))
.unwrap_or_default()
.lines()
.map(|l| serde_json::from_str(l).unwrap())
.collect()
}
/// The inputs of every Zen request `zen` received, in order.
async fn asked_zen(zen: &MockServer) -> Vec<Value> {
fn read(root: &Path) -> String {
std::fs::read_to_string(root.join("src/add.py")).unwrap()
}
/// The user message of every Zen request `zen` received, in order.
async fn asked_zen(zen: &MockServer) -> Vec<String> {
zen.received_requests()
.await
.unwrap_or_default()
.iter()
.map(testing::inputs)
.map(testing::prompt)
.collect()
}
#[tokio::test(flavor = "multi_thread")]
async fn kai_picks_the_file_zen_patches_it_and_the_flow_is_done() {
async fn a_single_fault_is_found_by_search_and_the_model_is_never_asked() {
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");
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"));
assert_eq!(read(repo.path()), SIMPLE.replace("a - b", "a + b"));
assert!(report.summary().contains("solved by a operator change"), "{}", report.summary());
let lines = ledger(repo.path());
let kinds: Vec<&str> = lines.iter().map(|l| l["kind"].as_str().unwrap()).collect();
assert_eq!(kinds.first(), Some(&"action"));
assert!(kinds.contains(&"verify") && kinds.last() == Some(&"end"), "{kinds:?}");
let mutate = lines.iter().find(|l| l["operation"]["type"] == "mutate").unwrap();
assert_eq!(mutate["status"], "selected");
assert_eq!(mutate["path"], json!(["eligible", "attempted", "completed", "selected"]));
assert!(lines.iter().all(|l| l["task"] == lines[0]["task"]));
let ids: std::collections::HashSet<&str> = lines.iter().filter_map(|l| l["id"].as_str()).collect();
assert_eq!(ids.len(), lines.iter().filter(|l| l["id"].is_string()).count(), "ids are unique");
}
#[tokio::test(flavor = "multi_thread")]
async fn the_model_writes_the_edit_when_search_finds_none_and_the_flow_is_done() {
let repo = repo();
let kai = honest().await;
let zen = testing::zen(|_| diff(BUGGY, FIXED)).await;
let zen = testing::zen(|_| fix()).await;
let report = run_flow(repo.path(), &kai.uri(), &zen, options(3, Control::Kai)).await;
assert_eq!(report.outcome, "done");
assert_eq!(report.attempts.len(), 1);
let a = &report.attempts[0];
assert_eq!(a.file.as_deref(), Some("src/add.sh"));
assert_eq!((a.by.as_str(), a.applied, a.passed), ("kai", true, true));
let (answer, p) = a.kai.clone().unwrap();
assert_eq!(answer, "done");
assert!((p - 0.9).abs() < 1e-6, "{p}");
assert_eq!((a.applied, a.passed, a.effort.as_str()), (true, true, "low"));
assert_eq!(a.shown.first().map(String::as_str), Some("add"));
assert!(!report.fast.as_ref().unwrap().solved);
assert_eq!(read(repo.path()), FIXED);
assert!(report.missing.is_none());
assert!(
report.summary().contains("outcome: done"),
"{}",
report.summary()
);
for step in ["test0", "tree", "rank", "site", "propose", "search", "total"] {
assert!(report.ms.contains_key(step), "{step}: {:?}", report.ms);
}
for step in ["bundle", "generate", "apply", "test"] {
assert!(a.ms.contains_key(step), "{step}: {:?}", a.ms);
}
assert_eq!(a.generation.as_ref().unwrap().prompt, 100);
let sent = asked_zen(&zen).await;
assert_eq!(
sent[0]["read"]["path"], "src/add.sh",
"Zen sees the file Kai picked"
);
assert!(
sent[0]["failure"]["output"]
.as_str()
.unwrap()
.contains("gave -1")
);
assert_eq!(sent.len(), 1);
assert!(sent[0].contains("src/add.py:1-2 (add)"), "{}", sent[0]);
assert!(sent[0].contains(" 2| return a ^ b"), "{}", sent[0]);
assert!(sent[0].contains("AssertionError"), "{}", sent[0]);
let lines = testing::lines(&trace(repo.path()), 1);
let line = &lines[0];
assert_eq!(line["family"], "dev.flow");
assert_eq!(line["sku"], "flow.fix@1");
assert_eq!(line["sku"], "flow.fix@2");
assert_eq!(line["outcome"]["finish"], "done");
assert_eq!(line["outcome"]["result"]["passed"], true);
assert_eq!(line["ops"]["pick"]["program"], "fix.pick@1");
assert_eq!(line["ops"]["pick"]["kai"], "src/add.sh");
assert_eq!(line["ops"]["pick"]["applied"], true);
assert_eq!(line["ops"]["next"]["kai"], "done");
assert_eq!(line["ops"]["next"]["taken"], "done");
assert!(
line["ops"]["pick"]["answers"].as_array().unwrap().len() >= 2,
"one per candidate"
);
assert_eq!(line["outcome"]["result"]["generation"]["prompt_tokens"], 100);
// A run that passes needs no inference: the ledger says so, and Kai was not asked about it.
let led = ledger(repo.path());
assert!(led.iter().any(|l| l["kind"] == "fast" && l["op"] == "turn"), "{led:?}");
let ops: Vec<&str> = led.iter().filter_map(|l| l["operation"]["type"].as_str()).collect();
for op in ["test", "search", "mutate", "generate", "edit"] {
assert!(ops.contains(&op), "{op}: {ops:?}");
}
let generate = led.iter().find(|l| l["operation"]["type"] == "generate").unwrap();
assert_eq!(generate["source"], "model");
assert!(led.iter().any(|l| l["kind"] == "fallback" && l["op"] == "search"));
}
#[tokio::test(flavor = "multi_thread")]
async fn a_wrong_patch_is_undone_and_the_retry_sees_it() {
async fn a_wrong_edit_is_undone_and_the_retry_sees_it() {
let repo = repo();
let kai = honest().await;
let zen = testing::zen(|context| {
if context["last"].is_null() {
diff(BUGGY, "echo $(( $1 * $2 ))")
let zen = testing::zen(|prompt| {
if prompt.contains("Your previous edit") {
fix()
} else {
diff(BUGGY, FIXED)
edit_json("a ^ b", "a | b")
}
})
.await;
let report = run_flow(repo.path(), &kai.uri(), &zen, options(3, Control::Kai)).await;
assert_eq!(report.outcome, "done");
assert_eq!(report.outcome, "done", "{}", report.summary());
let taken: Vec<&str> = report.attempts.iter().map(|a| a.taken.as_str()).collect();
assert_eq!(taken, ["retry", "done"]);
assert!(report.attempts[0].applied && !report.attempts[0].passed);
assert_eq!(read(repo.path()), FIXED);
let sent = asked_zen(&zen).await;
assert!(
sent[1]["last"]["test"]["output"]
.as_str()
.unwrap()
.contains("gave 6"),
"the retry sees the last attempt"
);
assert!(sent[1].contains("applied and the test still failed"), "{}", sent[1]);
assert!(sent[1].contains("-a ^ b +a | b"), "the retry sees the edit: {}", sent[1]);
assert!(sent[1].contains("AssertionError"), "and what the tests said after it");
assert_eq!(testing::lines(&trace(repo.path()), 2).len(), 2);
let led = ledger(repo.path());
let generates: Vec<&Value> = led.iter().filter(|l| l["operation"]["type"] == "generate").collect();
assert_eq!(generates[0]["source"], "model");
assert_eq!(generates[1]["source"], "recover", "the second attempt is recover's");
assert!(led.iter().any(|l| l["kind"] == "recover" && l["attempt"] == 1));
assert!(led.iter().any(|l| l["operation"]["type"] == "revert" && l["status"] == "applied"));
// Kai was asked what next once, after the failed run, in one request.
assert!(led.iter().any(|l| l["operation"]["type"] == "decide" && l["operation"]["program"] == "fix.turn@1"));
}
#[tokio::test(flavor = "multi_thread")]
async fn the_same_failure_after_another_edit_makes_recover_choose_full_effort_then_escalate() {
let repo = repo();
let kai = honest().await;
// Every edit leaves the same failing assertion.
let zen = testing::zen({
let n = std::sync::atomic::AtomicUsize::new(0);
move |_| edit_json("a ^ b", &format!("a | {}", n.fetch_add(1, std::sync::atomic::Ordering::SeqCst)))
})
.await;
let report = run_flow(repo.path(), &kai.uri(), &zen, options(4, Control::Kai)).await;
let efforts: Vec<&str> = report.attempts.iter().map(|a| a.effort.as_str()).collect();
assert_eq!(efforts, ["low", "full", "full"].get(..report.attempts.len()).unwrap(), "{}", report.summary());
assert_eq!(report.outcome, "escalate", "{}", report.summary());
assert_eq!(read(repo.path()), BUGGY, "the tree is as it was");
let requests = zen.received_requests().await.unwrap();
let thinking: Vec<bool> = requests
.iter()
.map(|r| {
let body: Value = serde_json::from_slice(&r.body).unwrap();
body["chat_template_kwargs"]["enable_thinking"] == true
})
.collect();
assert_eq!(thinking[0], false);
assert!(thinking[1..].iter().all(|t| *t), "{thinking:?}");
let led = ledger(repo.path());
let strategies: Vec<&str> = led.iter().filter(|l| l["kind"] == "recover").filter_map(|l| l["strategy"].as_str()).collect();
assert_eq!(strategies.first(), Some(&"full"), "{strategies:?}");
assert_eq!(strategies.last(), Some(&"escalate"), "{strategies:?}");
}
#[tokio::test(flavor = "multi_thread")]
async fn an_edit_that_does_not_apply_is_refused_and_the_retry_is_told_why() {
let repo = repo();
let kai = honest().await;
let zen = testing::zen(|prompt| {
if prompt.contains("did not apply") {
fix()
} else {
edit_json("a minus b", "a + b")
}
})
.await;
let report = run_flow(repo.path(), &kai.uri(), &zen, options(3, Control::Kai)).await;
assert_eq!(report.outcome, "done");
assert!(report.attempts[0].refused.as_deref().unwrap().contains("does not stand"));
assert!(!report.attempts[0].applied);
assert!(!report.attempts[0].ms.contains_key("test"), "nothing was run");
assert_eq!(report.attempts[0].taken, "retry");
}
#[tokio::test(flavor = "multi_thread")]
async fn out_of_attempts_the_flow_escalates_and_leaves_the_tree_as_it_was() {
let repo = repo();
let kai = honest().await;
let zen = testing::zen(|_| diff(BUGGY, "echo 0")).await;
let zen = testing::zen(|_| edit_json("a ^ b", "0")).await;
let report = run_flow(repo.path(), &kai.uri(), &zen, options(2, Control::Kai)).await;
assert_eq!(report.outcome, "escalate");
let taken: Vec<&str> = report.attempts.iter().map(|a| a.taken.as_str()).collect();
assert_eq!(taken, ["retry", "escalate"]);
assert_eq!(report.attempts.last().unwrap().taken, "escalate");
assert_eq!(read(repo.path()), BUGGY);
assert!(report.summary().contains("+echo 0"), "{}", report.summary());
assert!(report.summary().contains("+0"), "{}", report.summary());
}
#[tokio::test(flavor = "multi_thread")]
async fn kai_cannot_call_a_failing_attempt_done() {
let repo = repo();
let kai = kai(|_| ("done", 0.95)).await;
let zen = testing::zen(|_| diff(BUGGY, "echo 4")).await;
let kai = kai(0, |_| ("done", 0.95)).await;
let zen = testing::zen(|_| edit_json("a ^ b", "4")).await;
let report = run_flow(repo.path(), &kai.uri(), &zen, options(2, Control::Kai)).await;
let taken: Vec<&str> = report.attempts.iter().map(|a| a.taken.as_str()).collect();
assert_eq!(taken, ["retry", "escalate"]);
@@ -237,35 +321,31 @@ async fn kai_cannot_call_a_failing_attempt_done() {
#[tokio::test(flavor = "multi_thread")]
async fn kai_may_escalate_early_and_rule_control_never_asks_it() {
let repo = repo();
let kai = kai(|_| ("escalate", 0.9)).await;
let zen = testing::zen(|_| diff(BUGGY, "echo 0")).await;
let kai = kai(0, |_| ("escalate", 0.9)).await;
let zen = testing::zen(|_| edit_json("a ^ b", "0")).await;
let report = run_flow(repo.path(), &kai.uri(), &zen, options(3, Control::Kai)).await;
assert_eq!(report.attempts.len(), 1);
assert_eq!(report.outcome, "escalate");
let asked = testing::asked(&kai).await;
assert!(asked >= 1);
let repo = self::repo();
let report = run_flow(repo.path(), &kai.uri(), &zen, options(3, Control::Rule)).await;
let taken: Vec<&str> = report.attempts.iter().map(|a| a.taken.as_str()).collect();
assert_eq!(taken, ["retry", "retry", "escalate"]);
assert_eq!(report.outcome, "escalate");
assert!(report.attempts.iter().all(|a| a.by == "rule"));
assert_eq!(
testing::asked(&kai).await,
asked,
"rule control asks Kai nothing"
);
assert_eq!(testing::asked(&kai).await, asked, "rule control asks Kai nothing");
}
#[tokio::test(flavor = "multi_thread")]
async fn without_kai_the_rule_decides_and_says_so_once() {
async fn without_kai_the_rule_decides_and_the_flow_still_finishes() {
let repo = repo();
// Kai's address answers nothing: every request is a 404.
let kai = MockServer::start().await;
let zen = testing::zen(|context| {
if context["last"].is_null() {
diff(BUGGY, "echo 0")
let zen = testing::zen(|prompt| {
if prompt.contains("Your previous edit") {
fix()
} else {
diff(BUGGY, FIXED)
edit_json("a ^ b", "0")
}
})
.await;
@@ -273,29 +353,15 @@ async fn without_kai_the_rule_decides_and_says_so_once() {
assert_eq!(report.outcome, "done");
let taken: Vec<&str> = report.attempts.iter().map(|a| a.taken.as_str()).collect();
assert_eq!(taken, ["retry", "done"]);
assert!(
report
.attempts
.iter()
.all(|a| a.by == "rule" && a.kai.is_none())
);
assert!(report.attempts.iter().all(|a| a.kai.is_none()));
assert_eq!(read(repo.path()), FIXED);
let summary = report.summary();
assert_eq!(
summary.matches("kai did not answer").count(),
1,
"{summary}"
);
// The first pick found Kai down, one request per candidate of the repository's four
// files; nothing asked it again.
let asked = testing::asked(&kai).await;
assert!((1..=4).contains(&asked), "{asked}");
assert!(report.summary().matches("kai did not answer").count() <= 1);
}
#[tokio::test(flavor = "multi_thread")]
async fn passing_tests_need_no_attempt() {
let repo = repo();
std::fs::write(repo.path().join("src/add.sh"), FIXED).unwrap();
std::fs::write(repo.path().join("src/add.py"), FIXED).unwrap();
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;
@@ -305,18 +371,147 @@ async fn passing_tests_need_no_attempt() {
}
#[tokio::test(flavor = "multi_thread")]
async fn a_patch_the_tests_rewrote_ends_the_flow_instead_of_retrying_on_it() {
async fn a_failure_in_hand_is_not_run_to_see_it() {
let repo = repo_with(SIMPLE);
let kai = honest().await;
let zen = testing::zen(|_| String::new()).await;
let mut options = options(3, Control::Kai);
options.failure = Some("tests/test_add.py:5: in test_add\n assert add(2, 3) == 5\nAssertionError\n".into());
// A test command that cannot be run to see it fail: only candidates run it.
options.task["test"] = json!("PYTHONPATH=src python3 tests/test_add.py");
let report = run_flow(repo.path(), &kai.uri(), &zen, options).await;
assert_eq!(report.outcome, "done");
assert!(report.ms["test0"] < 5.0, "{:?}", report.ms);
let led = ledger(repo.path());
assert_eq!(led[0]["result"]["given"], true);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_test_file_is_never_edited() {
let repo = repo();
let kai = honest().await;
let zen = testing::zen(|_| diff(BUGGY, "echo 0")).await;
// Once patched, the test command rewrites the line the patch changed.
let zen = testing::zen(|_| {
json!({"edits": [{"file": "tests/test_add.py", "line": 5, "old": "== 5", "new": "== -1"}]}).to_string()
})
.await;
let report = run_flow(repo.path(), &kai.uri(), &zen, options(1, Control::Kai)).await;
assert_eq!(report.outcome, "escalate");
assert!(report.attempts[0].refused.as_deref().unwrap().contains("test file"));
assert_eq!(std::fs::read_to_string(repo.path().join("tests/test_add.py")).unwrap(), TEST);
}
#[test]
fn recover_never_replays_the_same_failure() {
let f = |fp: &str, rev: &str| Failure { fingerprint: fp.into(), revision: rev.into() };
// The failure changed: progress, try again.
assert_eq!(recover(&[f("a", "base"), f("b", "r1")], Strategy::Retry, 2).0, Strategy::Retry);
// The same failure after another edit: full effort once, then escalate.
assert_eq!(recover(&[f("a", "base"), f("a", "r1")], Strategy::Retry, 2).0, Strategy::Full);
assert_eq!(recover(&[f("a", "r1"), f("a", "r2")], Strategy::Full, 2).0, Strategy::Escalate);
// The same failure at the same revision: nothing changed at all.
assert_eq!(recover(&[f("a", "r1"), f("a", "r1")], Strategy::Retry, 2).0, Strategy::Escalate);
// No attempts left.
assert_eq!(recover(&[f("a", "base"), f("b", "r1")], Strategy::Retry, 0).0, Strategy::Escalate);
}
/// Two faults in two files, a test that counts its failures as a runner does.
fn two_faults() -> tempfile::TempDir {
let dir = tempfile::tempdir().unwrap();
let root = dir.path();
std::fs::create_dir_all(root.join("src")).unwrap();
std::fs::create_dir_all(root.join("tests")).unwrap();
std::fs::write(root.join("src/a.py"), "def f(a, b):\n return a - b\n").unwrap();
std::fs::write(root.join("src/b.py"), "def g(a, b):\n return a - b\n").unwrap();
std::fs::write(
root.join("tests/test_both.py"),
"import sys\nsys.path.insert(0, 'src')\nfrom a import f\nfrom b import g\nbad = []\nif f(3, 2) != 5:\n bad.append('f')\nif g(7, 2) != 9:\n bad.append('g')\nprint(f'{len(bad)} failed, {2 - len(bad)} passed')\nsys.exit(1 if bad else 0)\n",
)
.unwrap();
let git = |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:?}");
};
git(&["init", "-q"]);
git(&["add", "."]);
git(&["commit", "-q", "-m", "start"]);
dir
}
#[tokio::test(flavor = "multi_thread")]
async fn a_step_that_fixes_one_of_two_faults_is_kept_and_the_state_looked_at_again() {
let repo = two_faults();
let kai = honest().await;
let zen = testing::zen(|_| String::new()).await;
let mut options = options(3, Control::Kai);
options.task["test"] =
json!("grep -q 'echo 0' src/add.sh && printf 'echo 7\\n' > src/add.sh; sh test.sh");
let setup = setup(repo.path(), &kai.uri(), &zen);
let root = repo.path().to_path_buf();
let Err(error) = run_with(setup, root, options).await else {
panic!("the flow went on over a tree it could not restore");
};
assert!(error.contains("no longer reverses"), "{error}");
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");
let fast = report.fast.as_ref().unwrap();
assert!(fast.solved && fast.kept == 1 && fast.rounds == 2, "{fast:?}");
assert_eq!(report.failing, 2);
assert_eq!(std::fs::read_to_string(repo.path().join("src/a.py")).unwrap(), "def f(a, b):\n return a + b\n");
assert_eq!(std::fs::read_to_string(repo.path().join("src/b.py")).unwrap(), "def g(a, b):\n return a + b\n");
assert_eq!(report.patch.len(), 2, "{:?}", report.patch);
let led = ledger(repo.path());
let mutates = led.iter().filter(|l| l["operation"]["type"] == "mutate").count();
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");
}
#[tokio::test(flavor = "multi_thread")]
async fn kept_progress_is_undone_when_the_flow_does_not_end_done() {
let repo = two_faults();
// The search keeps one step; the model then refuses to help and the attempts run out.
let kai = kai(0, |_| ("escalate", 0.9)).await;
let zen = testing::zen(|_| "no edit".to_string()).await;
let mut options = options(1, Control::Kai);
options.task["test"] = json!("python3 tests/test_both.py");
// A second fault the generic changes cannot mend.
std::fs::write(repo.path().join("src/b.py"), "def g(a, b):\n return a ^ b\n").unwrap();
let report = run_flow(repo.path(), &kai.uri(), &zen, options).await;
assert_eq!(report.outcome, "escalate", "{}", report.summary());
assert_eq!(std::fs::read_to_string(repo.path().join("src/a.py")).unwrap(), "def f(a, b):\n return a - b\n", "a kept step is put back");
assert_eq!(std::fs::read_to_string(repo.path().join("src/b.py")).unwrap(), "def g(a, b):\n return a ^ b\n");
}
#[test]
fn tgrep_puts_the_function_that_spells_the_message_into_the_ranking() {
let dir = tempfile::tempdir().unwrap();
let root = dir.path();
std::fs::create_dir_all(root.join("pkg")).unwrap();
std::fs::write(
root.join("pkg/core.py"),
"def check(x):\n if x < 0:\n raise ValueError(\"negative values are not allowed here\")\n return x\n\n\ndef other(x):\n return x\n",
)
.unwrap();
let failure = "tests/test_core.py:4: in test_check\n run(-1)\nE ValueError: negative values are not allowed here\nFAILED tests/test_core.py::test_check - ValueError: negative values are not allowed here\n";
let mut tree = Tree::open(root);
let mut plan = locate::plan(&mut tree, failure);
// The ranker does not know it: take it out and let search put it back.
plan.ranked.retain(|r| r.sym.as_ref().is_none_or(|s| s.qual != "check"));
let ledger = Ledger::new("t", root.join("ledger.jsonl"));
let (added, _) = discover(&Tgrep, root, &mut plan, &ledger, "base");
assert_eq!(added, vec!["check".to_string()]);
let at = plan.ranked.iter().position(|r| r.sym.as_ref().is_some_and(|s| s.qual == "check")).unwrap();
assert!(at <= 3, "within the first four: {at}");
let line: Value = serde_json::from_str(std::fs::read_to_string(root.join("ledger.jsonl")).unwrap().lines().next().unwrap()).unwrap();
assert_eq!(line["operation"]["engine"], "tgrep");
assert_eq!(line["result"]["functions"][0]["by"], "message");
}
#[test]
fn spelled_takes_quoted_messages_and_error_tails_and_leaves_paths_and_asserts() {
let got = spelled("ValueError: negative values are not allowed here\nKeyError: 'missing key in the table'\nassert x == 4\nopen('src/a.py')");
assert!(got.contains(&"negative values are not allowed here".to_string()), "{got:?}");
assert!(got.contains(&"missing key in the table".to_string()), "{got:?}");
assert!(!got.iter().any(|s| s.contains("src/a.py") || s.starts_with("assert")), "{got:?}");
}
crates/hanzo-kai/src/ledger.rs
+385
View File
@@ -0,0 +1,385 @@
//! The ledger: every action of a run, replayable.
//!
//! The loop owns one task state and one ledger. An [`Action`] is a typed operation with the
//! revision it is bound to, the evidence that led to it, the preconditions checked immediately
//! before it runs, the authorization it ran under, what it is expected to produce and how that
//! is verified, and who proposed it ([`Source`]): a rule, Kai, a model or `recover`. An expected
//! outcome is never accepted because something predicted it: only the verifier's record settles
//! it. Every entry has a stable id (`<task>.<n>`), so time and cost are attributable exactly, and
//! credentials and sensitive output are redacted before a line is written.
use regex::Regex;
use serde::Serialize;
use serde_json::Value;
use serde_json::json;
use sha2::Digest;
use sha2::Sha256;
use std::collections::BTreeMap;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::sync::LazyLock;
use std::sync::Mutex;
/// Characters of one output a ledger line keeps.
const OUTPUT: usize = 600;
/// Who proposed an action.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum Source {
/// An action that applied with no inference: one candidate, preconditions holding.
Rule,
/// Kai's structured decision.
Kai,
/// A model's generation.
Model,
/// `recover`, after a failed verification.
Recover,
}
/// Where an action stands. A decision passes eligible, attempted, completed, validated, selected,
/// applied and verified in that order, or ends early as abstained, rejected, timed out, failed or
/// escalated. `verified` is the independent check's word, never a prediction. Whether an action was
/// beneficial is never set here: only a benchmark sets it, from a measured gain.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum Status {
/// Its preconditions hold; nothing has run.
Eligible,
/// It has been started.
Attempted,
/// It returned something.
Completed,
/// What it returned is well formed and within what is allowed.
Validated,
/// It was chosen among the alternatives.
Selected,
/// Its effect is in the workspace.
Applied,
/// The verifier confirmed its effect.
Verified,
/// It declined to choose: no answer the gate accepts.
Abstained,
/// It was refused: a precondition did not hold, or the verifier refused its effect.
Rejected,
TimedOut,
Failed,
Escalated,
}
/// What an action does.
#[derive(Debug, Clone, PartialEq, Serialize)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum Operation {
/// Rank the tree's functions by where the fix likely is.
Search { engine: String },
/// Kai's structured decision: one request, every question a state needs.
Decide { program: String },
/// A model writes an edit.
Generate { model: String, effort: String },
/// Candidate edits at the best sites, each against the visible test.
Mutate { candidates: usize },
/// Make the edits in the tree.
Edit { edits: Vec<String> },
/// Run the tests: the verifier.
Test { command: String },
/// Put back the files an edit changed.
Revert { files: Vec<String> },
}
/// A condition checked immediately before an action runs.
#[derive(Debug, Clone, PartialEq, Serialize)]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum Precondition {
/// The workspace is at this revision.
Revision { revision: String },
/// The file holds what the evidence showed.
Digest { file: String, digest: String },
}
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct Action {
pub id: String,
pub task: String,
pub revision: String,
pub evidence: Vec<String>,
pub operation: Operation,
pub preconditions: Vec<Precondition>,
pub authorization: String,
pub expected: String,
pub verification: String,
pub source: Source,
pub status: Status,
/// The states it passed through, in order.
pub path: Vec<Status>,
/// Why this action, in a sentence.
pub reason: String,
pub ms: f64,
pub result: Value,
}
impl Action {
/// Moves to `status`, remembering the path.
pub fn advance(&mut self, status: Status) {
self.status = status;
self.path.push(status);
}
/// Ends in `Failed`, or `TimedOut` when `why` says the call timed out.
pub fn fail(&mut self, why: &str) {
let w = why.to_lowercase();
self.advance(if w.contains("timed out") || w.contains("timeout") || w.contains("deadline") {
Status::TimedOut
} else {
Status::Failed
});
}
}
/// The ledger of one run, appended to one file.
pub struct Ledger {
task: String,
path: PathBuf,
seq: Mutex<u32>,
file: Mutex<Option<std::fs::File>>,
}
static SECRETS: LazyLock<Vec<(Regex, &'static str)>> = LazyLock::new(|| {
[
(r"(?i)bearer\s+[A-Za-z0-9._~+/=-]{8,}", "Bearer <redacted>"),
(r"\b(?:sk|hk|pk|rk)-[A-Za-z0-9_-]{12,}", "<redacted>"),
(r"\bgh[pousr]_[A-Za-z0-9]{20,}", "<redacted>"),
(r"\bAKIA[0-9A-Z]{16}\b", "<redacted>"),
(r"\beyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}", "<redacted>"),
(r"(?i)\b([A-Z0-9_]*(?:KEY|TOKEN|SECRET|PASSWORD|PASSWD|CREDENTIAL)[A-Z0-9_]*)=\S+", "$1=<redacted>"),
]
.into_iter()
.map(|(p, r)| (Regex::new(p).expect("a fixed pattern"), r))
.collect()
});
/// `text` with credentials replaced.
pub fn redact(text: &str) -> String {
SECRETS
.iter()
.fold(text.to_string(), |t, (re, with)| re.replace_all(&t, *with).into_owned())
}
fn scrub(value: &mut Value) {
match value {
Value::String(s) => {
*s = redact(&crate::decide::cut(s, OUTPUT * 4));
}
Value::Array(a) => a.iter_mut().for_each(scrub),
Value::Object(o) => o.values_mut().for_each(scrub),
_ => {}
}
}
/// The first 12 hex digits of the SHA-256 of `bytes`.
pub fn digest(bytes: &[u8]) -> String {
let h = Sha256::digest(bytes);
h.iter().take(6).map(|b| format!("{b:02x}")).collect()
}
/// The revision of the files a run has touched: a hash of their contents, so a revert returns to
/// the revision it left, and an unrelated file's change moves nothing.
pub fn revision(root: &Path, files: &[String]) -> String {
let mut seen: BTreeMap<&str, String> = BTreeMap::new();
for f in files {
let bytes = std::fs::read(root.join(f)).unwrap_or_default();
seen.insert(f, digest(&bytes));
}
if seen.is_empty() {
return "base".into();
}
digest(format!("{seen:?}").as_bytes())
}
/// A failure with what varies from run to run taken out, so the same failure is the same
/// fingerprint: colours, absolute paths, line numbers, addresses and ids, durations.
pub fn fingerprint(output: &str) -> String {
static NORMAL: LazyLock<Vec<(Regex, &'static str)>> = LazyLock::new(|| {
[
(r"\x1b\[[0-9;]*m", ""),
(r"(?:/[\w.@+-]+){2,}/", "<dir>/"),
(r":\d+(?::\d+)?\b", ":<n>"),
(r"\bline \d+\b", "line <n>"),
(r"\b0x[0-9a-fA-F]+\b", "<addr>"),
(r"\b[0-9a-f]{8,}\b", "<id>"),
(r"\b\d+(?:\.\d+)?\s?(?:ms|s|seconds)\b", "<time>"),
(r"\bin \d+(?:\.\d+)?s\b", "in <time>"),
]
.into_iter()
.map(|(p, r)| (Regex::new(p).expect("a fixed pattern"), r))
.collect()
});
let text = NORMAL
.iter()
.fold(output.to_string(), |t, (re, with)| re.replace_all(&t, *with).into_owned());
// The words of the failure, in order, without the layout.
let words: Vec<&str> = text.split_whitespace().collect();
digest(words.join(" ").as_bytes())
}
impl Ledger {
pub fn new(task: &str, path: PathBuf) -> Ledger {
Ledger {
task: task.to_string(),
path,
seq: Mutex::new(0),
file: Mutex::new(None),
}
}
pub fn task(&self) -> &str {
&self.task
}
/// The next stable id.
pub fn next(&self) -> String {
let mut n = self.seq.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
*n += 1;
format!("{}.{}", self.task, *n)
}
/// An action's skeleton at `revision`, with the next id.
pub fn action(&self, revision: &str, operation: Operation, source: Source, reason: &str) -> Action {
Action {
id: self.next(),
task: self.task.clone(),
revision: revision.to_string(),
evidence: Vec::new(),
operation,
preconditions: Vec::new(),
authorization: "auto: the task's own tools".into(),
expected: String::new(),
verification: "the tests".into(),
source,
status: Status::Eligible,
path: vec![Status::Eligible],
reason: reason.to_string(),
ms: 0.0,
result: Value::Null,
}
}
/// Appends a line: redacted, whole, one at a time. A ledger that cannot be written is
/// logged and skipped: the loop never waits on it.
pub fn write(&self, kind: &str, body: Value) {
let mut line = json!({"time": crate::trace::now(), "kind": kind, "task": self.task});
if let (Some(l), Value::Object(b)) = (line.as_object_mut(), body) {
l.extend(b);
}
scrub(&mut line);
let Ok(mut text) = serde_json::to_vec(&line) else {
return;
};
text.push(b'\n');
let mut file = self.file.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
if file.is_none() {
if let Some(dir) = self.path.parent() {
let _ = std::fs::create_dir_all(dir);
}
match std::fs::OpenOptions::new().create(true).append(true).open(&self.path) {
Ok(f) => *file = Some(f),
Err(e) => {
tracing::warn!("kai ledger {}: {e}", self.path.display());
return;
}
}
}
if let Some(f) = file.as_mut()
&& let Err(e) = f.write_all(&text)
{
tracing::warn!("kai ledger {}: {e}", self.path.display());
}
}
/// An action, finished: its line.
pub fn record(&self, action: &Action) {
if let Ok(v) = serde_json::to_value(action) {
self.write("action", v);
}
}
/// A step that did not need an inference, or fell from a deterministic one to a model, and
/// why.
pub fn fallback(&self, revision: &str, op: &str, from: &str, to: &str, reason: &str) {
self.write(
"fallback",
json!({"id": self.next(), "revision": revision, "op": op, "from": from, "to": to, "reason": reason}),
);
}
/// A verification: the independent check's record, which alone settles an outcome.
pub fn verify(&self, action: &str, revision: &str, passed: bool, code: i64, fingerprint: &str, output: &str, ms: f64) {
self.write(
"verify",
json!({"id": self.next(), "action": action, "revision": revision, "passed": passed, "code": code,
"fingerprint": fingerprint, "output": crate::decide::cut(output, OUTPUT), "ms": ms}),
);
}
}
#[cfg(test)]
mod tests {
use super::*;
use pretty_assertions::assert_eq;
#[test]
fn credentials_are_redacted_before_they_are_written() {
let t = "curl -H 'Authorization: Bearer abcdef0123456789' sk-abcdefghijklmnop API_KEY=hunter2 eyJhbGciOiJI.eyJzdWIiOiIx.c2lnbmF0dXJl ghp_abcdefghijklmnopqrstuvwxyz";
let r = redact(t);
for leak in ["abcdef0123456789", "sk-abcdefghijklmnop", "hunter2", "eyJhbGciOiJI", "ghp_abc"] {
assert!(!r.contains(leak), "{leak} in {r}");
}
assert!(r.contains("API_KEY=<redacted>"), "{r}");
}
#[test]
fn a_failure_is_the_same_fingerprint_whatever_varies_between_runs() {
let a = "tests/test_x.py:44: in test_a\n/home/z/w/t1/pkg/m.py:631: TypeError at 0x7f12ab in 0.16s id 9f86d081884c7d65\nE assert 3 == 4";
let b = "tests/test_x.py:47: in test_a\n/tmp/other/dir/pkg/m.py:640: TypeError at 0x55aa00 in 0.21s id aaaaaaaabbbbbbbb\nE assert 3 == 4";
assert_eq!(fingerprint(a), fingerprint(b));
assert_ne!(fingerprint(a), fingerprint("E assert 3 == 5"));
}
#[test]
fn a_revision_is_the_contents_of_the_files_touched_and_a_revert_returns_to_it() {
let dir = tempfile::tempdir().unwrap();
std::fs::write(dir.path().join("a.py"), "x = 1\n").unwrap();
let files = vec!["a.py".to_string()];
let before = revision(dir.path(), &files);
std::fs::write(dir.path().join("a.py"), "x = 2\n").unwrap();
let during = revision(dir.path(), &files);
std::fs::write(dir.path().join("a.py"), "x = 1\n").unwrap();
assert_ne!(before, during);
assert_eq!(before, revision(dir.path(), &files));
assert_eq!(revision(dir.path(), &[]), "base");
}
#[test]
fn every_entry_has_a_stable_id_and_the_lines_are_whole_json() {
let dir = tempfile::tempdir().unwrap();
let ledger = Ledger::new("run1", dir.path().join("kai/ledger.jsonl"));
let mut a = ledger.action("base", Operation::Test { command: "pytest".into() }, Source::Rule, "the task's test command");
a.advance(Status::Attempted);
a.advance(Status::Verified);
a.result = json!({"output": "Authorization: Bearer abcdefghij1234"});
ledger.record(&a);
ledger.fallback("base", "site", "rule", "kai", "ranking not confident");
let text = std::fs::read_to_string(dir.path().join("kai/ledger.jsonl")).unwrap();
let lines: Vec<Value> = text.lines().map(|l| serde_json::from_str(l).unwrap()).collect();
assert_eq!(lines[0]["id"], "run1.1");
assert_eq!(lines[0]["kind"], "action");
assert_eq!(lines[0]["operation"]["type"], "test");
assert_eq!(lines[0]["source"], "rule");
assert_eq!(lines[0]["status"], "verified");
assert_eq!(lines[0]["path"], json!(["eligible", "attempted", "verified"]));
assert!(!text.contains("abcdefghij1234"));
assert_eq!(lines[1]["id"], "run1.2");
assert_eq!(lines[1]["reason"], "ranking not confident");
}
}
crates/hanzo-kai/src/lib.rs
+7
View File
@@ -11,8 +11,15 @@
pub mod agent;
pub mod decide;
pub mod diagnostic;
pub mod edit;
pub mod find;
pub mod flow;
pub mod ledger;
pub mod locate;
pub mod mutate;
pub mod program;
pub mod symbols;
#[cfg(test)]
mod testing;
pub mod tools;
crates/hanzo-kai/src/locate.rs
+709
View File
@@ -0,0 +1,709 @@
//! Where the fix is: a failing run's evidence ranked over every function of the tree.
//!
//! kai-gateway's learned preflight policy (`policy.ts`) in Rust, with its weights
//! (`programs/locate.json`, written by bench/code/onegen/policies.py from the curriculum's
//! train repositories) unchanged. The 25 features are computed from the failing diagnostic and
//! the working tree; a softmax over them ranks every function and method; a logistic head over
//! the ranking's shape says whether the best of it holds the fix site. For a compiler's
//! diagnostic, whose frames are exact, the frames themselves are the ranking.
use crate::diagnostic::Diagnostic;
use crate::diagnostic::Kind as Failure;
use crate::symbols::Kind;
use crate::symbols::Map;
use crate::symbols::Sym;
use crate::symbols::Tree;
use crate::symbols::is_test;
use regex::Regex;
use serde::Deserialize;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::LazyLock;
const WEIGHTS: &str = include_str!("../programs/locate.json");
/// Longest function body whose lines count.
const CAP: u32 = 200;
/// Functions a ranking keeps.
const TOP: usize = 48;
/// Most functions a bundle reads.
const PICKS: usize = 48;
/// Lines of functions one bundle reads.
const BUDGET: u32 = 1200;
/// Frames' window either side of a compiler's line.
const WINDOW: u32 = 20;
/// Lines of a symbol read for a compiler's frame without a line.
const SYMBOL: u32 = 80;
const NAMES: [&str; 25] = [
"frame_file", "frame_sym", "called", "name_in_diag", "name_tok", "test_tok", "import_file",
"import_name", "bm25", "bm25_test", "bm25_text", "bm25_name", "bm25_file", "callee",
"callee2", "class_called", "ctor_called", "dunder_op", "parent_reached", "dups", "log_len",
"is_method", "is_private", "is_dunder", "file_syms",
];
#[derive(Deserialize)]
struct Head {
mu: Vec<f64>,
sd: Vec<f64>,
w: Vec<f64>,
#[serde(default)]
b: f64,
#[serde(default)]
tau: f64,
}
#[derive(Deserialize)]
struct Weights {
names: Vec<String>,
mu: Vec<f64>,
sd: Vec<f64>,
w: Vec<f64>,
gamma: f64,
escalate: Head,
}
static FIT: LazyLock<Weights> = LazyLock::new(|| {
let fit: Weights = serde_json::from_str(WEIGHTS).expect("programs/locate.json");
assert_eq!(fit.names, NAMES, "programs/locate.json names the features in order");
fit
});
fn re(pattern: &str) -> Regex {
Regex::new(pattern).expect("a fixed pattern")
}
static CALL: LazyLock<Regex> = LazyLock::new(|| re(r"([A-Za-z_][A-Za-z0-9_]*)\s*\("));
static CAMEL: LazyLock<Regex> = LazyLock::new(|| re(r"([a-z0-9])([A-Z])"));
static FROM_IMPORT: LazyLock<Regex> =
LazyLock::new(|| re(r"(?m)^[ \t]*from[ \t]+([\w.]+)[ \t]+import[ \t]+\(?([^\n)]+)"));
static IMPORT: LazyLock<Regex> = LazyLock::new(|| re(r"(?m)^[ \t]*import[ \t]+([\w., \t]+)"));
static AS: LazyLock<Regex> = LazyLock::new(|| re(r"\s+as\s+"));
static PRIVATE: LazyLock<Regex> = LazyLock::new(|| re(r"^_[^_]"));
static DUNDER_NAME: LazyLock<Regex> = LazyLock::new(|| re(r"^__.*__$"));
static TEST_PREFIX: LazyLock<Regex> = LazyLock::new(|| re(r"(?i)^test_?"));
static WORD: LazyLock<Regex> = LazyLock::new(|| re(r"^\w+$"));
/// The operators and builtins that run a dunder method without naming it.
static DUNDER: LazyLock<HashMap<&'static str, Regex>> = LazyLock::new(|| {
[
("__add__", r"\+"), ("__radd__", r"\+"), ("__sub__", "-"), ("__rsub__", "-"),
("__mul__", r"\*"), ("__truediv__", "/"), ("__floordiv__", "//"), ("__mod__", "%"),
("__eq__", "=="), ("__ne__", "!="), ("__lt__", "<"), ("__le__", "<="), ("__gt__", ">"),
("__ge__", ">="), ("__hash__", r"\bhash\(|\bset\(|\{"), ("__repr__", r"\brepr\(|\bstr\("),
("__str__", r"\bstr\(|\bformat\("), ("__len__", r"\blen\("),
("__iter__", r"\bfor\b|\blist\(|\bsorted\(|\btuple\("), ("__getitem__", r"\w\["),
("__contains__", r"\bin\b"), ("__bool__", r"\bnot\b|\bbool\("),
("__neg__", r"(?:^|[(\s=])-\w"), ("__call__", r"\)\("), ("__enter__", r"\bwith\b"),
("__exit__", r"\bwith\b"),
]
.into_iter()
.map(|(name, pattern)| (name, re(pattern)))
.collect()
});
const KEYWORDS: &[&str] = &[
"and", "as", "assert", "async", "await", "break", "class", "continue", "def", "del", "elif",
"else", "except", "finally", "for", "from", "global", "if", "import", "in", "is", "lambda",
"nonlocal", "not", "or", "pass", "raise", "return", "try", "while", "with", "yield", "none",
"true", "false", "self", "cls", "fn", "let", "mut", "impl", "pub", "use", "mod", "str", "int",
];
/// Lower-case word pieces of an identifier or a text: `parseXml_text` is parse, xml, text.
pub fn tokens(text: &str) -> Vec<String> {
CAMEL
.replace_all(text, "$1 $2")
.to_lowercase()
.split(|c: char| !(c.is_ascii_lowercase() || c.is_ascii_digit()))
.filter(|t| t.len() > 1 && !KEYWORDS.contains(t))
.map(str::to_string)
.collect()
}
fn calls(text: &str) -> HashSet<String> {
CALL.captures_iter(text).map(|c| c[1].to_string()).collect()
}
/// The module a file is imported as: `src/pkg/mod.py` is pkg.mod, `pkg/__init__.py` is pkg.
fn dotted(file: &str) -> String {
let stem = file.rsplit_once('.').map_or(file, |(s, _)| s);
let stem = stem
.strip_prefix("src/")
.or_else(|| stem.strip_prefix("lib/"))
.unwrap_or(stem);
stem.strip_suffix("/__init__").unwrap_or(stem).replace('/', ".")
}
/// What the failing tests import: dotted modules and attributes, and the bare names brought in.
fn imports(text: &str, targets: &mut HashSet<String>, names: &mut HashSet<String>) {
for m in FROM_IMPORT.captures_iter(text) {
let base = m[1].to_string();
targets.insert(base.clone());
for part in m[2].split(',') {
let name = AS.split(part.trim()).next().unwrap_or("").trim().to_string();
if !name.is_empty() && WORD.is_match(&name) {
targets.insert(format!("{base}.{name}"));
names.insert(name);
}
}
}
for m in IMPORT.captures_iter(text) {
for part in m[1].split(',') {
targets.insert(AS.split(part.trim()).next().unwrap_or("").trim().to_string());
}
}
}
fn innermost<'a>(syms: &'a [Sym], file: &str, line: u32) -> Option<&'a Sym> {
syms.iter()
.filter(|s| s.file == file && s.start <= line && line <= s.end)
.min_by_key(|s| s.end - s.start)
}
/// The failing test's own function: the test frame's enclosing symbol.
pub fn failing<'a>(d: &Diagnostic, map: &'a Map) -> Option<&'a Sym> {
d.frames
.iter()
.filter(|f| is_test(&f.file))
.find_map(|f| innermost(&map.symbols, &f.file, f.line))
}
/// A place to look, best first; `score` in [0, 1].
#[derive(Debug, Clone, PartialEq)]
pub struct Ranked {
pub file: String,
pub sym: Option<Sym>,
pub score: f64,
/// A compiler frame's line.
pub line: Option<u32>,
}
struct Candidate {
sym: Sym,
f: [f64; 25],
}
/// BM25 of a query over documents, normalised by the best in this tree.
fn bm25(docs: &[Vec<String>], query: &HashSet<String>) -> Vec<f64> {
let mut df: HashMap<&str, usize> = HashMap::new();
for doc in docs {
for t in doc.iter().collect::<HashSet<_>>() {
*df.entry(t.as_str()).or_default() += 1;
}
}
let avg = docs.iter().map(Vec::len).sum::<usize>() as f64 / docs.len().max(1) as f64;
let raw: Vec<f64> = docs
.iter()
.map(|doc| {
let mut tf: HashMap<&str, usize> = HashMap::new();
for t in doc {
*tf.entry(t.as_str()).or_default() += 1;
}
let mut score = 0.0;
for q in query {
let Some(n) = tf.get(q.as_str()) else {
continue;
};
let d = df.get(q.as_str()).copied().unwrap_or(0) as f64;
let idf = (1.0 + (docs.len() as f64 - d + 0.5) / (d + 0.5)).ln();
let n = *n as f64;
score += (idf * n * 2.2) / (n + 1.2 * (0.25 + (0.75 * doc.len() as f64) / avg.max(1.0)));
}
score
})
.collect();
let top = raw.iter().fold(1e-9_f64, |a, b| a.max(*b));
raw.into_iter().map(|v| v / top).collect()
}
/// Every function and method of the tree's sources with its features, for this diagnostic.
fn candidates(d: &Diagnostic, map: &Map, tree: &mut Tree) -> Vec<Candidate> {
let syms: Vec<&Sym> = map
.symbols
.iter()
.filter(|s| {
(matches!(s.kind, Kind::Fn | Kind::Method) || (s.kind == Kind::Class && s.end - s.start <= 80))
&& !is_test(&s.file)
&& !s.file.starts_with("examples/")
&& !s.file.starts_with("demos/")
})
.collect();
// Every file read once; a symbol's lines are slices of its file's.
let mut needed: Vec<&str> = syms.iter().map(|s| s.file.as_str()).collect();
needed.extend(d.frames.iter().filter(|f| is_test(&f.file)).map(|f| f.file.as_str()));
needed.extend(map.tests.iter().map(String::as_str));
for file in needed {
tree.read(file);
}
let tree: &Tree = tree;
let mut split: HashMap<&str, Vec<&str>> = HashMap::new();
let mut lines_in = |file: &'_ str, start: u32, end: u32| -> String {
let Some(text) = tree.text(file) else {
return String::new();
};
let all = split.entry(tree.key(file)).or_insert_with(|| text.split('\n').collect());
all.iter().skip(start as usize - 1).take((end - start + 1) as usize).copied().collect::<Vec<_>>().join("\n")
};
let texts: Vec<String> = syms.iter().map(|s| lines_in(&s.file, s.start, s.end)).collect();
let frames: Vec<_> = d.frames.iter().filter(|f| !is_test(&f.file)).collect();
let test = failing(d, map);
let mut test_text = test.map(|t| lines_in(&t.file, t.start, t.end)).unwrap_or_default();
if let Some(t) = test
&& t.kind == Kind::Method
&& let Some(parent) = map
.symbols
.iter()
.find(|s| s.file == t.file && s.kind == Kind::Class && s.start <= t.start && t.end <= s.end)
{
test_text = lines_in(&parent.file, parent.start, parent.end);
}
let mut test_files: HashSet<String> = d
.frames
.iter()
.filter(|f| is_test(&f.file))
.map(|f| f.file.clone())
.collect();
for t in &d.tests {
let head = t.split("::").next().unwrap_or("");
if map.tests.iter().any(|m| m == head) {
test_files.insert(head.to_string());
}
}
let (mut targets, mut names) = (HashSet::new(), HashSet::new());
let mut ordered: Vec<&String> = test_files.iter().collect();
ordered.sort();
for file in ordered {
imports(tree.text(file).unwrap_or(""), &mut targets, &mut names);
}
let called = calls(&test_text);
let last = |t: &String| t.split("::").last().unwrap_or(t).to_string();
let query: HashSet<String> = d
.names
.iter()
.flat_map(|n| tokens(n))
.chain(called.iter().flat_map(|n| tokens(n)))
.chain(d.tests.iter().flat_map(|t| tokens(&last(t))))
.collect();
let test_tokens: HashSet<String> = d
.tests
.iter()
.flat_map(|t| tokens(&TEST_PREFIX.replace(&last(t), "")))
.collect();
let bodies: Vec<Vec<String>> = texts.iter().map(|t| tokens(t)).collect();
let b_query = bm25(&bodies, &query);
let b_test = bm25(&bodies, &tokens(&test_text).into_iter().collect());
let b_text = bm25(&bodies, &tokens(&d.text).into_iter().collect());
let name_docs: Vec<Vec<String>> = syms.iter().map(|s| tokens(&format!("{} {}", s.qual, s.file))).collect();
let b_name = bm25(&name_docs, &query);
let mut file_best: HashMap<&str, f64> = HashMap::new();
let mut file_count: HashMap<&str, usize> = HashMap::new();
for (i, s) in syms.iter().enumerate() {
let best = file_best.entry(s.file.as_str()).or_insert(0.0);
*best = best.max(b_query[i]);
*file_count.entry(s.file.as_str()).or_default() += 1;
}
let mut direct: HashSet<usize> = HashSet::new();
for (i, s) in syms.iter().enumerate() {
if called.contains(&s.name)
|| d.names.contains(&s.name)
|| frames.iter().any(|f| f.file == s.file && s.start <= f.line && f.line <= s.end)
{
direct.insert(i);
}
}
// The call graph by name, as far as two calls from what the failing test calls; a name that
// many functions share is no evidence.
let mut by_name: HashMap<&str, Vec<usize>> = HashMap::new();
for (i, s) in syms.iter().enumerate() {
by_name.entry(s.name.as_str()).or_default().push(i);
}
let reach = |from: &HashSet<usize>| -> HashSet<usize> {
let mut out = HashSet::new();
for i in from {
for c in calls(&texts[*i]) {
if let Some(hit) = by_name.get(c.as_str())
&& hit.len() <= 8
{
out.extend(hit.iter().copied());
}
}
}
out
};
let one = reach(&direct);
let two = reach(&one);
let reached: HashSet<usize> = direct.union(&one).copied().collect();
let owner = |s: &Sym| -> String {
if s.kind == Kind::Method {
let mut parts = s.qual.split('.').collect::<Vec<_>>();
parts.pop();
parts.pop().unwrap_or("").to_string()
} else {
String::new()
}
};
let test_body = format!("{test_text}\n{}", d.text);
let flag = |b: bool| if b { 1.0 } else { 0.0 };
syms.iter()
.enumerate()
.map(|(i, s)| {
let pieces = tokens(&s.name);
let share = |set: &HashSet<String>| {
if pieces.is_empty() {
0.0
} else {
pieces.iter().filter(|p| set.contains(*p)).count() as f64 / pieces.len() as f64
}
};
let own = owner(s);
let f = [
flag(frames.iter().any(|x| x.file == s.file)),
flag(frames.iter().any(|x| x.file == s.file && s.start <= x.line && x.line <= s.end)),
flag(called.contains(&s.name)),
flag(d.names.contains(&s.name)),
share(&query),
share(&test_tokens),
flag(targets.contains(&dotted(&s.file))),
flag(names.contains(&s.name)),
b_query[i],
b_test[i],
b_text[i],
b_name[i],
file_best.get(s.file.as_str()).copied().unwrap_or(0.0),
flag(one.contains(&i)),
flag(two.contains(&i) && !one.contains(&i)),
flag(!own.is_empty() && (called.contains(&own) || d.names.contains(&own))),
flag((s.name == "__init__" || s.name == "__new__") && !own.is_empty() && called.contains(&own)),
flag(DUNDER.get(s.name.as_str()).is_some_and(|r| r.is_match(&test_body))),
flag(syms.iter().enumerate().any(|(j, o)| {
j != i && o.file == s.file && o.start <= s.start && s.end <= o.end && reached.contains(&j)
})),
(by_name.get(s.name.as_str()).map_or(1, Vec::len) as f64).ln(),
(1.0 + f64::from(s.end - s.start + 1)).ln(),
flag(s.kind == Kind::Method),
flag(PRIVATE.is_match(&s.name)),
flag(DUNDER_NAME.is_match(&s.name)),
(1.0 + file_count.get(s.file.as_str()).copied().unwrap_or(0) as f64).ln(),
];
Candidate { sym: (*s).clone(), f }
})
.collect()
}
fn dot(w: &[f64], mu: &[f64], sd: &[f64], x: &[f64]) -> f64 {
x.iter()
.enumerate()
.map(|(i, v)| ((v - mu[i]) / sd[i]) * w[i])
.sum()
}
fn softmax(scores: &[f64]) -> Vec<f64> {
let m = scores.iter().fold(f64::NEG_INFINITY, |a, b| a.max(*b));
let e: Vec<f64> = scores.iter().map(|s| (s - m).exp()).collect();
let sum: f64 = e.iter().sum();
e.into_iter().map(|v| v / sum).collect()
}
/// A compiler's diagnostic: the frames it names, source before tests.
fn frames(d: &Diagnostic, map: &Map) -> Vec<Ranked> {
let files: HashSet<&str> = map.files.iter().map(String::as_str).collect();
let mut seen = HashSet::new();
let (mut sources, mut tests) = (Vec::new(), Vec::new());
for frame in &d.frames {
let Some(file) = resolve(&frame.file, map, &files) else {
continue;
};
if !seen.insert((file.clone(), frame.line)) {
continue;
}
let sym = innermost(&map.symbols, &file, frame.line).cloned();
let named = sym.as_ref().is_some_and(|s| {
d.tests
.iter()
.any(|t| t.split("::").last().unwrap_or(t) == s.name)
});
let entry = (file.clone(), frame.line, sym);
if is_test(&file) || named {
tests.push(entry);
} else {
sources.push(entry);
}
}
sources
.into_iter()
.chain(tests)
.enumerate()
.map(|(i, (file, line, sym))| Ranked {
file,
sym,
score: 1.0 / (1.0 + i as f64),
line: Some(line),
})
.collect()
}
/// A path the output names, as the tree spells it: exact, else the one file it is the tail of.
fn resolve(file: &str, map: &Map, files: &HashSet<&str>) -> Option<String> {
if files.contains(file) {
return Some(file.to_string());
}
let tails: Vec<&String> = map
.files
.iter()
.filter(|c| file.ends_with(&format!("/{c}")) && (!file.starts_with('/') || c.contains('/')))
.collect();
let longest = tails.iter().map(|t| t.len()).max()?;
let best: Vec<_> = tails.into_iter().filter(|t| t.len() == longest).collect();
(best.len() == 1).then(|| best[0].clone())
}
/// Every candidate's file, name, lines and features for `failure`, in the map's order: what the
/// fit saw, for a test against kai-gateway's own extraction.
pub fn features(tree: &mut Tree, failure: &str) -> Vec<(Sym, [f64; 25])> {
let diagnostic = crate::diagnostic::parse(failure, tree.root().to_str());
let map = tree.map();
candidates(&diagnostic, &map, tree)
.into_iter()
.map(|c| (c.sym, c.f))
.collect()
}
/// The evidence of one failing run read against one tree.
pub struct Plan {
pub diagnostic: Diagnostic,
pub map: Map,
/// Best first.
pub ranked: Vec<Ranked>,
/// The failing test's own function.
pub test: Option<Sym>,
/// The probability that the ranking's best holds the fix site: a one-shot edit is worth
/// trying at once above [`Plan::sure`].
pub oneshot: f64,
}
/// Reads `failure` (a failing run's output) against the tree at `root`.
pub fn plan(tree: &mut Tree, failure: &str) -> Plan {
let diagnostic = crate::diagnostic::parse(failure, tree.root().to_str());
let map = tree.map();
let ranked = if diagnostic.kind == Failure::Compile {
frames(&diagnostic, &map)
} else {
let all = candidates(&diagnostic, &map, tree);
if all.is_empty() {
frames(&diagnostic, &map)
} else {
let p = softmax(&all.iter().map(|c| dot(&FIT.w, &FIT.mu, &FIT.sd, &c.f)).collect::<Vec<_>>());
let mut ranked: Vec<Ranked> = all
.into_iter()
.zip(p)
.map(|(c, p)| Ranked {
file: c.sym.file.clone(),
sym: Some(c.sym),
score: p,
line: None,
})
.collect();
ranked.sort_by(|a, b| b.score.total_cmp(&a.score));
ranked.truncate(TOP);
ranked
}
};
let test = failing(&diagnostic, &map).cloned();
let oneshot = escalate(&ranked, &diagnostic);
Plan {
diagnostic,
map,
ranked,
test,
oneshot,
}
}
impl Plan {
/// The threshold above which the ranking's best is trusted alone (fit with the weights).
pub fn sure(&self) -> f64 {
FIT.escalate.tau
}
}
/// The functions to read, the ranking's best first, until their probabilities add up to the fit
/// `gamma` or the line budget is spent.
pub fn cuts(ranked: &[Ranked]) -> Vec<&Ranked> {
let (mut out, mut cum, mut lines) = (Vec::new(), 0.0, 0);
for x in ranked.iter().take(PICKS) {
let size = x.sym.as_ref().map_or(0, |s| (s.end - s.start + 1).min(CAP));
if !out.is_empty() && lines + size > BUDGET {
break;
}
out.push(x);
lines += size;
cum += x.score;
if cum >= FIT.gamma {
break;
}
}
out
}
/// The shape of a ranking and of the bundle read for it: what the escalation head sees.
fn shape(p: &[f64], d: &Diagnostic, picks: usize, lines: u32) -> [f64; 8] {
let top = &p[..p.len().min(8)];
let sum = top.iter().sum::<f64>();
let sum = if sum == 0.0 { 1.0 } else { sum };
let entropy = -top.iter().map(|v| (v / sum) * ((v / sum) + 1e-12).ln()).sum::<f64>();
let at = |i: usize| p.get(i).copied().unwrap_or(0.0);
[
at(0),
at(0) - at(1),
entropy,
if d.kind == Failure::Assert { 1.0 } else { 0.0 },
if d.frames.iter().any(|f| !is_test(&f.file)) { 1.0 } else { 0.0 },
(d.tests.len() as f64).ln_1p(),
(picks as f64).ln_1p(),
f64::from(lines).ln_1p(),
]
}
/// The probability that the best of `ranked` holds the fix site.
fn escalate(ranked: &[Ranked], d: &Diagnostic) -> f64 {
let top = ranked.first();
if d.kind == Failure::Compile || ranked.iter().any(|x| x.line.is_some()) {
// A compiler's frames are exact: one-shot unless the best is test code.
return match top {
Some(t) if t.line.is_some() && !t.sym.as_ref().is_some_and(|s| is_test(&s.file)) && !is_test(&t.file) => 1.0,
_ => 0.0,
};
}
let picks = cuts(ranked);
let lines: u32 = picks
.iter()
.map(|x| x.sym.as_ref().map_or(0, |s| (s.end - s.start + 1).min(CAP)))
.sum();
let scores: Vec<f64> = ranked.iter().map(|x| x.score).collect();
let e = &FIT.escalate;
let z = dot(&e.w, &e.mu, &e.sd, &shape(&scores, d, picks.len(), lines)) + e.b;
1.0 / (1.0 + (-z).exp())
}
/// The lines to read for a compiler's frame: the window around its line within its symbol.
pub fn window(x: &Ranked) -> Option<(u32, u32)> {
match (x.line, &x.sym) {
(Some(line), sym) => Some((
line.saturating_sub(WINDOW).max(1).max(sym.as_ref().map_or(1, |s| s.start)),
(line + WINDOW).min(sym.as_ref().map_or(u32::MAX, |s| s.end)),
)),
(None, Some(s)) => Some((s.start, s.end.min(s.start + SYMBOL - 1))),
(None, None) => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
use pretty_assertions::assert_eq;
#[test]
fn tokens_split_identifiers_and_drop_keywords() {
assert_eq!(tokens("parseXml_text self"), ["parse", "xml", "text"]);
assert_eq!(dotted("src/pkg/mod.py"), "pkg.mod");
assert_eq!(dotted("pkg/__init__.py"), "pkg");
}
#[test]
fn the_function_a_failing_test_calls_by_name_ranks_first() {
let dir = tempfile::tempdir().unwrap();
let root = dir.path();
std::fs::create_dir_all(root.join("pkg")).unwrap();
std::fs::create_dir_all(root.join("tests")).unwrap();
std::fs::write(
root.join("pkg/maths.py"),
"def add(a, b):\n return a - b\n\ndef mul(a, b):\n return a * b\n\ndef sub(a, b):\n return a - b\n",
)
.unwrap();
std::fs::write(
root.join("tests/test_maths.py"),
"from pkg.maths import add\n\ndef test_add():\n assert add(2, 3) == 5\n",
)
.unwrap();
let failure = "___ test_add ___\ntests/test_maths.py:4: in test_add\n assert add(2, 3) == 5\nE assert -1 == 5\nFAILED tests/test_maths.py::test_add - assert -1 == 5\n";
let mut tree = Tree::open(root);
let plan = plan(&mut tree, failure);
assert_eq!(plan.ranked[0].sym.as_ref().unwrap().qual, "add");
assert!(plan.ranked[0].score > plan.ranked[1].score);
assert_eq!(plan.test.as_ref().unwrap().qual, "test_add");
let total: f64 = plan.ranked.iter().map(|r| r.score).sum();
assert!(total <= 1.0 + 1e-9);
}
#[test]
fn a_compilers_frames_are_the_ranking() {
let dir = tempfile::tempdir().unwrap();
let root = dir.path();
std::fs::create_dir_all(root.join("src")).unwrap();
std::fs::write(root.join("src/lib.rs"), "fn f() -> u32 {\n 1u8\n}\n").unwrap();
let failure = "error[E0308]: mismatched types\n --> src/lib.rs:2:5\n |\n2 | 1u8\n";
let mut tree = Tree::open(root);
let plan = plan(&mut tree, failure);
assert_eq!(plan.ranked.len(), 1);
assert_eq!(plan.ranked[0].line, Some(2));
assert_eq!(window(&plan.ranked[0]), Some((1, 3)));
assert_eq!(plan.oneshot, 1.0);
}
/// kai-gateway's extraction of the same trees, written by kai-gateway scripts/features.mjs from main over
/// `~/.cache/kai-onegen`: every feature of every candidate of the first task of each
/// repository. `cargo test -p hanzo-kai --release parity -- --ignored`.
#[test]
#[ignore = "needs the onegen curriculum's trees and features under ~/.cache/kai-onegen"]
fn the_features_are_kai_gateways_on_the_curriculum() {
let data = std::path::PathBuf::from(std::env::var("HOME").unwrap()).join(".cache/kai-onegen");
let records: std::collections::HashMap<String, serde_json::Value> =
std::fs::read_to_string(data.join("records.jsonl"))
.unwrap()
.lines()
.map(|l| serde_json::from_str::<serde_json::Value>(l).unwrap())
.map(|r| (r["id"].as_str().unwrap().to_string(), r))
.collect();
let mut repos = std::collections::HashSet::new();
let mut checked = 0;
for line in std::fs::read_to_string(data.join("features.current.jsonl")).unwrap().lines() {
let golden: serde_json::Value = serde_json::from_str(line).unwrap();
let id = golden["id"].as_str().unwrap();
let record = &records[id];
if !repos.insert(record["repo"].as_str().unwrap().to_string()) {
continue;
}
let mut tree = Tree::open(&data.join("trees").join(id));
let ours = features(&mut tree, record["prompt"].as_str().unwrap());
let theirs = golden["cands"].as_array().unwrap();
if ours.len() != theirs.len() {
let key = |f: &str, q: &str, s: u64| format!("{f}:{q}:{s}");
let t: std::collections::HashSet<String> = theirs.iter().map(|t| key(t["file"].as_str().unwrap(), t["qual"].as_str().unwrap(), t["start"].as_u64().unwrap())).collect();
let o: std::collections::HashSet<String> = ours.iter().map(|(s, _)| key(&s.file, &s.qual, u64::from(s.start))).collect();
eprintln!("only ours: {:?}", o.difference(&t).take(8).collect::<Vec<_>>());
eprintln!("only theirs: {:?}", t.difference(&o).take(8).collect::<Vec<_>>());
}
assert_eq!(ours.len(), theirs.len(), "{id}: candidates");
let mut worst = 0.0_f64;
for (i, (sym, f)) in ours.iter().enumerate() {
let t = &theirs[i];
assert_eq!((sym.file.as_str(), sym.qual.as_str(), u64::from(sym.start), u64::from(sym.end)),
(t["file"].as_str().unwrap(), t["qual"].as_str().unwrap(), t["start"].as_u64().unwrap(), t["end"].as_u64().unwrap()),
"{id}: candidate {i}");
for (j, v) in f.iter().enumerate() {
let g = t["f"][j].as_f64().unwrap();
let d = (v - g).abs();
worst = worst.max(d);
assert!(d < 1e-6, "{id}: {} of {}: {v} against {g}", NAMES[j], sym.qual);
}
}
eprintln!("{id}: {} candidates, worst difference {worst:e}", ours.len());
checked += 1;
}
assert!(checked >= 5, "{checked} repositories checked");
}
}
crates/hanzo-kai/src/mutate.rs
+780
View File
@@ -0,0 +1,780 @@
//! Candidate edits at the best sites, tried against the visible test.
//!
//! A single fault in code is often one small change from the code that works: a comparison
//! flipped, an operator swapped, a constant off by one, two arguments exchanged, a condition
//! inverted, a `return` dropped. [`propose`] lists those changes, generic transforms of a line's
//! 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.
//!
//! 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.
use crate::diagnostic::Diagnostic;
use crate::locate::Plan;
use crate::symbols::Tree;
use crate::tools;
use std::collections::HashSet;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Mutex;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::Instant;
/// Functions whose lines are tried.
const SITES: usize = 4;
/// Most candidates a site contributes.
const PER_SITE: usize = 64;
/// Most candidates in all.
const CANDIDATES: usize = 96;
/// Most tests run at once.
const WORKERS: usize = 8;
/// A tree larger than this is not copied for a search.
const TREE: u64 = 128 * 1024 * 1024;
/// A file larger than this is not copied.
const FILE: u64 = 8 * 1024 * 1024;
/// Lines of a function considered.
const LINES: u32 = 150;
/// One change to one line.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Candidate {
pub file: String,
/// 1-based.
pub line: u32,
/// The line as it is, and as the candidate would have it.
pub old: String,
pub new: String,
pub kind: &'static str,
/// The rank of the site it is at.
pub site: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum Tok {
Word,
Num,
Str,
Op,
Open,
Close,
Comma,
}
#[derive(Debug, Clone)]
struct Token {
kind: Tok,
start: usize,
end: usize,
}
const OPS: &[&str] = &[
"**=", "//=", ">>=", "<<=", "==", "!=", "<=", ">=", "->", "=>", "**", "//", "<<", ">>", "+=", "-=",
"*=", "/=", "%=", "&&", "||", ":=", "<>",
];
/// A line's tokens, over its masked text: comments blank and string contents filled.
fn scan(m: &[u8]) -> Vec<Token> {
let mut out = Vec::new();
let mut i = 0;
while i < m.len() {
let c = m[i];
if c == b' ' || c == b'\t' || c == b'\r' {
i += 1;
} else if c == 0 || c == b'"' || c == b'\'' || c == b'`' {
let start = i;
i += 1;
while i < m.len() && m[i] == 0 {
i += 1;
}
out.push(Token { kind: Tok::Str, start, end: i });
} else if c.is_ascii_alphabetic() || c == b'_' {
let start = i;
while i < m.len() && (m[i].is_ascii_alphanumeric() || m[i] == b'_') {
i += 1;
}
out.push(Token { kind: Tok::Word, start, end: i });
} else if c.is_ascii_digit() {
let start = i;
while i < m.len() && (m[i].is_ascii_alphanumeric() || m[i] == b'_' || m[i] == b'.') {
i += 1;
}
out.push(Token { kind: Tok::Num, start, end: i });
} else if matches!(c, b'(' | b'[' | b'{') {
out.push(Token { kind: Tok::Open, start: i, end: i + 1 });
i += 1;
} else if matches!(c, b')' | b']' | b'}') {
out.push(Token { kind: Tok::Close, start: i, end: i + 1 });
i += 1;
} else if c == b',' {
out.push(Token { kind: Tok::Comma, start: i, end: i + 1 });
i += 1;
} else {
let rest = &m[i..];
let len = OPS
.iter()
.find(|op| rest.starts_with(op.as_bytes()))
.map_or(1, |op| op.len());
out.push(Token { kind: Tok::Op, start: i, end: i + len });
i += len;
}
}
out
}
/// Whether a token ends an operand: a name, a number, a string or a closing bracket.
fn operand_end(t: Option<&Token>) -> bool {
t.is_some_and(|t| matches!(t.kind, Tok::Word | Tok::Num | Tok::Str | Tok::Close))
}
fn operand_start(t: Option<&Token>) -> bool {
t.is_some_and(|t| matches!(t.kind, Tok::Word | Tok::Num | Tok::Str | Tok::Open))
}
const KEYWORDS: &[&str] = &[
"and", "or", "not", "is", "in", "if", "elif", "else", "while", "for", "return", "lambda", "yield",
"await", "assert", "del", "with", "as", "from", "import", "raise", "except", "print",
];
fn is_keyword(w: &str) -> bool {
KEYWORDS.contains(&w)
}
const DUAL: &[(&str, &str)] = &[
("min", "max"), ("any", "all"), ("startswith", "endswith"), ("lstrip", "rstrip"),
("lower", "upper"), ("ljust", "rjust"), ("first", "last"), ("keys", "values"),
("True", "False"), ("true", "false"),
];
/// The changes to the tokens of one line: `(kind, start, end, replacement)` over the line's text.
fn changes(line: &str, m: &[u8], python: bool) -> Vec<(&'static str, usize, usize, String)> {
let tokens = scan(m);
let text = |t: &Token| &line[t.start..t.end];
let mut out: Vec<(&'static str, usize, usize, String)> = Vec::new();
let first = tokens.first().map(|t| text(t));
let condition = python && matches!(first, Some("if" | "elif" | "while")) && line.trim_end().ends_with(':');
for (i, t) in tokens.iter().enumerate() {
let prev = i.checked_sub(1).and_then(|p| tokens.get(p));
let next = tokens.get(i + 1);
let s = text(t);
match t.kind {
Tok::Op => {
let binary = operand_end(prev) && operand_start(next);
let swaps: &[&str] = match s {
"==" => &["!="],
"!=" => &["=="],
"<" if binary => &["<=", ">"],
">" if binary => &[">=", "<"],
"<=" => &["<", ">="],
">=" => &[">", "<="],
_ => &[],
};
for r in swaps {
out.push(("compare", t.start, t.end, (*r).to_string()));
}
let arith: &[&str] = match s {
"+" if binary => &["-", "*"],
"-" if binary => &["+", "*"],
"*" if binary => &["/", "+"],
"/" if binary => &["*", "//"],
"//" if binary => &["/", "*"],
"%" if binary => &["*", "//"],
"**" if binary => &["*"],
"+=" => &["-="],
"-=" => &["+="],
"*=" => &["/="],
"/=" => &["*="],
"&&" => &["||"],
"||" => &["&&"],
_ => &[],
};
for r in arith {
out.push(("operator", t.start, t.end, (*r).to_string()));
}
}
Tok::Word => {
if python {
match s {
"and" => out.push(("boolean", t.start, t.end, "or".into())),
"or" => out.push(("boolean", t.start, t.end, "and".into())),
"not" if next.map(|n| text(n)) != Some("in") => {
let end = next.map_or(t.end, |n| n.start);
out.push(("boolean", t.start, end, String::new()));
}
"is" => match next.map(|n| text(n)) {
Some("not") => out.push(("boolean", t.start, next.map_or(t.end, |n| n.end), "is".into())),
_ => out.push(("boolean", t.start, t.end, "is not".into())),
},
"in" => {
let membership = first != Some("for") && !prev.is_some_and(|p| text(p) == "not");
if membership {
out.push(("boolean", t.start, t.end, "not in".into()));
}
}
_ => {}
}
if s == "not" && next.map(|n| text(n)) == Some("in") {
out.push(("boolean", t.start, next.map_or(t.end, |n| n.end), "in".into()));
}
}
if let Some((_, b)) = DUAL.iter().find(|(a, _)| *a == s) {
out.push(("dual", t.start, t.end, (*b).to_string()));
} else if let Some((a, _)) = DUAL.iter().find(|(_, b)| *b == s) {
out.push(("dual", t.start, t.end, (*a).to_string()));
}
}
Tok::Num => {
let digits = s.replace('_', "");
if let Ok(v) = digits.parse::<u64>()
&& !prev.is_some_and(|p| text(p) == ".")
{
out.push(("constant", t.start, t.end, (v + 1).to_string()));
if v > 0 {
out.push(("constant", t.start, t.end, (v - 1).to_string()));
}
}
}
_ => {}
}
// len(x) and range(..., n) are off by one when their result is.
if t.kind == Tok::Word
&& matches!(s, "len" | "range")
&& next.is_some_and(|n| n.kind == Tok::Open && text(n) == "(")
{
if let Some(close) = matching(&tokens, i + 1) {
let end = tokens[close].end;
let at = if s == "range" { tokens[close].start } else { end };
for r in [" - 1", " + 1"] {
out.push(("bound", at, at, r.to_string()));
}
}
}
// f(a, b): the first two arguments exchanged, and the last two.
if t.kind == Tok::Word
&& !is_keyword(s)
&& next.is_some_and(|n| n.kind == Tok::Open && text(n) == "(")
&& let Some(close) = matching(&tokens, i + 1)
{
let args = split(&tokens, i + 1, close);
let plain = args.iter().all(|(a, b)| !has_keyword_argument(&tokens, line, *a, *b));
if args.len() >= 2 && plain {
let swap = |x: usize, y: usize| -> (usize, usize, String) {
let (xs, xe) = (tokens[args[x].0].start, tokens[args[x].1].end);
let (ys, ye) = (tokens[args[y].0].start, tokens[args[y].1].end);
(xs, ye, format!("{}{}{}", &line[ys..ye], &line[xe..ys], &line[xs..xe]))
};
let (a, b, r) = swap(0, 1);
out.push(("arguments", a, b, r));
if args.len() >= 3 {
let (a, b, r) = swap(args.len() - 2, args.len() - 1);
out.push(("arguments", a, b, r));
}
}
}
}
if condition {
let body = line.trim_end();
let kw = first.unwrap_or("if");
let at = line.find(kw).unwrap_or(0) + kw.len();
let cond = body[at..body.len() - 1].trim();
if let Some(rest) = cond.strip_prefix("not ") {
let rest = rest.trim();
let rest = rest.strip_prefix('(').and_then(|r| r.strip_suffix(')')).unwrap_or(rest);
out.push(("condition", at, body.len() - 1, format!(" {rest}")));
} else if !cond.is_empty() {
out.push(("condition", at, body.len() - 1, format!(" not ({cond})")));
}
}
out
}
/// The index of the closing bracket matching the opening one at `open`.
fn matching(tokens: &[Token], open: usize) -> Option<usize> {
let mut depth = 0;
for (i, t) in tokens.iter().enumerate().skip(open) {
match t.kind {
Tok::Open => depth += 1,
Tok::Close => {
depth -= 1;
if depth == 0 {
return Some(i);
}
}
_ => {}
}
}
None
}
/// The arguments of the call opened at `open`: token index ranges `(first, last)`.
fn split(tokens: &[Token], open: usize, close: usize) -> Vec<(usize, usize)> {
let mut args = Vec::new();
let (mut depth, mut from) = (0, open + 1);
for i in open + 1..close {
match tokens[i].kind {
Tok::Open => depth += 1,
Tok::Close => depth -= 1,
Tok::Comma if depth == 0 => {
if i > from {
args.push((from, i - 1));
}
from = i + 1;
}
_ => {}
}
}
if close > from {
args.push((from, close - 1));
}
args
}
/// Whether the argument spanning tokens `a..=b` is `name=value`.
fn has_keyword_argument(tokens: &[Token], line: &str, a: usize, b: usize) -> bool {
tokens.len() > a + 1 && a < b && &line[tokens[a + 1].start..tokens[a + 1].end] == "=" && tokens[a].kind == Tok::Word
}
/// What a line says of the failure: where the output points, and the names it uses.
struct Focus<'a> {
frames: HashSet<(String, u32)>,
names: HashSet<String>,
d: &'a Diagnostic,
}
impl Focus<'_> {
fn score(&self, file: &str, number: u32, line: &str, frame: Option<u32>) -> f64 {
let mut s = 0.0;
if self.frames.contains(&(file.to_string(), number)) {
s += 100.0;
}
for w in line.split(|c: char| !(c.is_alphanumeric() || c == '_')) {
if !w.is_empty() && self.names.contains(w) {
s += 10.0;
}
}
if let Some(f) = frame {
s -= f64::from(f.abs_diff(number)) * 0.01;
}
s
}
}
/// The candidates at the best sites of `plan`, best first.
pub fn propose(tree: &mut Tree, plan: &Plan, order: &[usize]) -> Vec<Candidate> {
let focus = Focus {
frames: plan.diagnostic.frames.iter().map(|f| (f.file.clone(), f.line)).collect(),
names: plan.diagnostic.names.iter().cloned().collect(),
d: &plan.diagnostic,
};
let mut out = Vec::new();
let mut seen = HashSet::new();
for (rank, r) in order
.iter()
.map(|i| &plan.ranked[*i])
.filter(|r| r.sym.is_some() && !crate::symbols::is_test(&r.file))
.take(SITES)
.enumerate()
{
let Some(sym) = &r.sym else { continue };
let python = sym.file.ends_with(".py");
let Some(text) = tree.read(&sym.file).map(str::to_string) else {
continue;
};
let Some(masked) = crate::symbols::masked(&sym.file, &text) else {
continue;
};
let masked_lines: Vec<&[u8]> = masked.split(|b| *b == b'\n').collect();
let lines: Vec<&str> = text.split('\n').collect();
let frame = focus
.d
.frames
.iter()
.find(|f| f.file == sym.file && sym.start <= f.line && f.line <= sym.end)
.map(|f| f.line);
let mut site: Vec<(f64, Candidate)> = Vec::new();
let last = sym.end.min(sym.start + LINES - 1);
for number in sym.start..=last {
let (Some(line), Some(m)) = (lines.get(number as usize - 1), masked_lines.get(number as usize - 1)) else {
break;
};
let line = line.trim_end_matches('\r');
let trimmed = line.trim_start();
if trimmed.is_empty() || trimmed.starts_with(['#', '@']) || trimmed.starts_with("import ") || trimmed.starts_with("from ") {
continue;
}
let base = focus.score(&sym.file, number, line, frame);
for (n, (kind, a, b, replacement)) in changes(line, m, python).into_iter().enumerate() {
if !line.is_char_boundary(a) || !line.is_char_boundary(b) {
continue;
}
let new = format!("{}{}{}", &line[..a], replacement, &line[b..]);
if new == line || !seen.insert((sym.file.clone(), number, new.clone())) {
continue;
}
// Earlier changes of a line before later ones.
site.push((base - n as f64 * 0.001, Candidate { file: sym.file.clone(), line: number, old: line.to_string(), new, kind, site: rank }));
}
// A statement that is a bare expression has dropped its `return`; a call may be
// meant for its effect, so only a function's last statement may be one.
if python && indent(line) > 0 {
let body = trimmed;
let keyword = ["return", "raise", "pass", "break", "continue", "yield", "del", "assert", "print", "with", "for", "while", "if", "else", "elif", "try", "except", "finally", "def", "class", "global", "nonlocal"]
.iter()
.any(|k| body.starts_with(k) && !body[k.len()..].starts_with(|c: char| c.is_alphanumeric() || c == '_'));
let bare = !body.starts_with(['#', '@', ')', ']', '}', '"', '\''])
&& !keyword
&& !has_assignment(body)
&& !body.ends_with([':', ',', '\\'])
&& balance_zero(body);
let call = body.ends_with(')');
if bare && (!call || number == sym.end) {
let pad = &line[..line.len() - trimmed.len()];
let new = format!("{pad}return {body}");
if seen.insert((sym.file.clone(), number, new.clone())) {
site.push((base, Candidate { file: sym.file.clone(), line: number, old: line.to_string(), new, kind: "return", site: rank }));
}
}
}
}
site.sort_by(|a, b| b.0.total_cmp(&a.0));
out.extend(site.into_iter().take(PER_SITE).map(|(_, c)| c));
if out.len() >= CANDIDATES {
break;
}
}
out.truncate(CANDIDATES);
out
}
fn indent(line: &str) -> usize {
line.len() - line.trim_start().len()
}
/// Whether every bracket the text opens it closes.
fn balance_zero(body: &str) -> bool {
body.bytes().fold(0i32, |n, b| match b {
b'(' | b'[' | b'{' => n + 1,
b')' | b']' | b'}' => n - 1,
_ => n,
}) == 0
}
/// Whether a statement assigns: an `=` outside brackets that is not part of a comparison.
fn has_assignment(body: &str) -> bool {
let b = body.as_bytes();
let mut depth = 0;
for i in 0..b.len() {
match b[i] {
b'(' | b'[' | b'{' => depth += 1,
b')' | b']' | b'}' => depth -= 1,
b'=' if depth == 0 => {
let before = i.checked_sub(1).map(|p| b[p]);
let after = b.get(i + 1);
if !matches!(before, Some(b'=' | b'!' | b'<' | b'>')) && after != Some(&b'=') {
return true;
}
}
_ => {}
}
}
false
}
/// What a search found and what it cost.
#[derive(Debug, Clone, Default)]
pub struct Search {
/// The first candidate, in order, that made the visible test pass.
pub winner: Option<usize>,
/// Otherwise the candidate that left the fewest tests failing, fewer than before: a step
/// toward a task with more than one fault.
pub progress: Option<usize>,
/// What the test printed with the winner or the progress in place.
pub output: String,
/// Tests still failing with it.
pub remaining: usize,
/// Candidates started.
pub tried: usize,
/// Test runs that ran to an end.
pub runs: usize,
pub workers: usize,
/// Milliseconds copying the tree, and in all.
pub copy: f64,
pub total: f64,
pub note: String,
}
fn size(dir: &Path, total: &mut u64) {
let Ok(entries) = std::fs::read_dir(dir) else {
return;
};
for e in entries.flatten() {
let name = e.file_name().to_string_lossy().into_owned();
let Ok(kind) = e.file_type() else { continue };
if kind.is_dir() {
if !skipped(&name) {
size(&e.path(), total);
}
} else if let Ok(m) = e.metadata() {
*total += m.len().min(FILE);
}
if *total > TREE {
return;
}
}
}
fn skipped(name: &str) -> bool {
matches!(name, ".git" | "__pycache__" | ".pytest_cache" | "node_modules" | "target" | ".venv" | "venv" | ".mypy_cache" | ".ruff_cache" | ".tox")
}
fn copy(from: &Path, to: &Path) -> std::io::Result<()> {
std::fs::create_dir_all(to)?;
for e in std::fs::read_dir(from)?.flatten() {
let name = e.file_name();
let kind = e.file_type()?;
let dest = to.join(&name);
if kind.is_symlink() {
if let Ok(target) = std::fs::read_link(e.path()) {
let _ = std::os::unix::fs::symlink(target, dest);
}
} else if kind.is_dir() {
if !skipped(&name.to_string_lossy()) {
copy(&e.path(), &dest)?;
}
} else if e.metadata()?.len() <= FILE {
std::fs::copy(e.path(), dest)?;
}
}
Ok(())
}
/// Replaces `number`'s line of `text` with `new`.
fn with_line(text: &str, number: u32, new: &str) -> String {
let mut lines: Vec<&str> = text.split('\n').collect();
lines[number as usize - 1] = new;
lines.join("\n")
}
/// 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, is the search's progress.
pub fn search(root: &Path, command: &str, candidates: &[Candidate], failing: usize, cap: Duration, limit: Duration, workers: usize) -> Search {
let begun = Instant::now();
let mut found = Search::default();
if candidates.is_empty() {
found.note = "no candidates".into();
return found;
}
let mut bytes = 0;
size(root, &mut bytes);
if bytes > TREE {
found.note = "the tree is too large to copy".into();
return found;
}
let workers = workers.clamp(1, WORKERS).min(candidates.len());
found.workers = workers;
let scratch = match tempfile::tempdir() {
Ok(d) => d,
Err(e) => {
found.note = format!("no scratch directory: {e}");
return found;
}
};
let next = AtomicUsize::new(0);
let tried = AtomicUsize::new(0);
let runs = AtomicUsize::new(0);
let winner: Mutex<Option<(usize, String)>> = Mutex::new(None);
// (tests left failing, candidate, output) of the best step so far.
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);
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 dir: PathBuf = scratch.path().join(format!("w{w}"));
s.spawn(move || {
let copying = Instant::now();
if copy(root, &dir).is_err() {
return;
}
{
let mut c = copy_ms.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
*c = c.max(copying.elapsed().as_secs_f64() * 1000.0);
}
let mut originals: std::collections::HashMap<String, String> = std::collections::HashMap::new();
loop {
if begun.elapsed() > cap {
break;
}
let i = next.fetch_add(1, Ordering::SeqCst);
let Some(c) = candidates.get(i) else { break };
if beaten(i) {
break;
}
tried.fetch_add(1, Ordering::SeqCst);
let path = dir.join(&c.file);
let original = originals
.entry(c.file.clone())
.or_insert_with(|| std::fs::read_to_string(&path).unwrap_or_default())
.clone();
if std::fs::write(&path, with_line(&original, c.line, &c.new)).is_err() {
continue;
}
let run = tools::shell_until(&dir, command, limit, &|| beaten(i));
let _ = std::fs::write(&path, &original);
let code = run["code"].as_i64().unwrap_or(-1);
if code == -2 || code == -3 {
continue;
}
runs.fetch_add(1, Ordering::SeqCst);
let output = run["output"].as_str().unwrap_or_default().to_string();
if code == 0 {
let mut f = winner.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
if f.as_ref().is_none_or(|(x, _)| i < *x) {
*f = Some((i, output));
}
} else if failing > 1 {
let left = crate::diagnostic::failures(&output);
if left >= 1 && left < failing {
let mut b = step.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
if b.as_ref().is_none_or(|(l, x, _)| (left, i) < (*l, *x)) {
*b = Some((left, i, output));
}
}
}
}
});
}
});
if let Some((i, output)) = winner.into_inner().unwrap_or_else(std::sync::PoisonError::into_inner) {
found.winner = Some(i);
found.output = output;
} else if let Some((left, i, output)) = step.into_inner().unwrap_or_else(std::sync::PoisonError::into_inner) {
found.progress = Some(i);
found.remaining = left;
found.output = output;
}
found.tried = tried.load(Ordering::SeqCst);
found.runs = runs.load(Ordering::SeqCst);
found.copy = *copy_ms.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
found.total = begun.elapsed().as_secs_f64() * 1000.0;
found
}
#[cfg(test)]
mod tests {
use super::*;
use crate::symbols::Tree;
use pretty_assertions::assert_eq;
fn kinds(line: &str, python: bool) -> Vec<(&'static str, String)> {
let m = crate::symbols::masked(if python { "a.py" } else { "a.rs" }, line).unwrap();
changes(line, &m, python)
.into_iter()
.map(|(k, a, b, r)| (k, format!("{}{}{}", &line[..a], r, &line[b..])))
.collect()
}
#[test]
fn a_comparison_an_operator_and_a_constant_each_have_their_flips() {
let got = kinds(" return a < b + 1", true);
assert!(got.contains(&("compare", " return a <= b + 1".into())), "{got:?}");
assert!(got.contains(&("compare", " return a > b + 1".into())), "{got:?}");
assert!(got.contains(&("operator", " return a < b - 1".into())), "{got:?}");
assert!(got.contains(&("constant", " return a < b + 2".into())), "{got:?}");
assert!(got.contains(&("constant", " return a < b + 0".into())), "{got:?}");
}
#[test]
fn nothing_inside_a_string_or_a_comment_changes() {
let got = kinds(" x = \"a < b and 3\" # a == 4", true);
assert!(got.is_empty(), "{got:?}");
}
#[test]
fn unary_minus_and_splats_are_not_binary_operators() {
let got = kinds(" f(*args, -x, **kw)", true);
assert!(!got.iter().any(|(k, _)| *k == "operator"), "{got:?}");
}
#[test]
fn arguments_swap_and_a_condition_inverts() {
let got = kinds(" return pow(base, exp)", true);
assert!(got.contains(&("arguments", " return pow(exp, base)".into())), "{got:?}");
let got = kinds(" if x in seen and not done:", true);
assert!(got.contains(&("condition", " if not (x in seen and not done):".into())), "{got:?}");
assert!(got.contains(&("boolean", " if x in seen or not done:".into())), "{got:?}");
assert!(got.contains(&("boolean", " if x not in seen and not done:".into())), "{got:?}");
assert!(got.contains(&("boolean", " if x in seen and done:".into())), "{got:?}");
let got = kinds(" if not ok:", true);
assert!(got.contains(&("condition", " if ok:".into())), "{got:?}");
}
#[test]
fn keyword_arguments_are_not_swapped_and_len_and_range_shift_by_one() {
let got = kinds(" f(a, key=b)", true);
assert!(!got.iter().any(|(k, _)| *k == "arguments"), "{got:?}");
let got = kinds(" for i in range(len(xs)):", true);
assert!(got.contains(&("bound", " for i in range(len(xs) - 1):".into())), "{got:?}");
assert!(got.contains(&("bound", " for i in range(len(xs) + 1):".into())), "{got:?}");
assert!(!got.iter().any(|(k, s)| *k == "boolean" && s.contains("not in")), "for-in is not membership: {got:?}");
}
#[test]
fn candidates_start_at_the_failing_line_of_the_best_site() {
let dir = tempfile::tempdir().unwrap();
std::fs::create_dir_all(dir.path().join("pkg")).unwrap();
std::fs::create_dir_all(dir.path().join("tests")).unwrap();
std::fs::write(dir.path().join("pkg/m.py"), "def f(a, b):\n c = a + 1\n return a - b\n").unwrap();
std::fs::write(dir.path().join("tests/test_m.py"), "from pkg.m import f\n\n\ndef test_f():\n assert f(3, 2) == 5\n").unwrap();
let failure = "pkg/m.py:3: in f\n return a - b\nE assert 1 == 5\ntests/test_m.py:5: in test_f\n assert f(3, 2) == 5\nFAILED tests/test_m.py::test_f - assert 1 == 5\n";
let mut tree = Tree::open(dir.path());
let plan = crate::locate::plan(&mut tree, failure);
let order: Vec<usize> = (0..plan.ranked.len()).collect();
let got = propose(&mut tree, &plan, &order);
assert_eq!(got[0].line, 3);
assert_eq!(got[0].new, " return a + b");
assert!(got.iter().all(|c| c.file == "pkg/m.py"));
}
#[test]
fn a_search_finds_the_first_candidate_that_passes_and_leaves_the_tree_alone() {
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"), "from m import f\nassert f(3, 2) == 5\n").unwrap();
let cands: Vec<Candidate> = [" return a * b", " return a + b", " return b - a"]
.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, Duration::from_secs(20), 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, Duration::from_secs(20), Duration::from_secs(10), 1);
assert_eq!(none.winner, None);
}
#[test]
fn with_two_faults_the_candidate_that_fixes_one_is_the_searchs_progress() {
let dir = tempfile::tempdir().unwrap();
std::fs::write(dir.path().join("m.py"), "def f(a, b):\n return a - b\n\n\ndef g(a, b):\n return a * b\n").unwrap();
// Two tests; `python3 t.py` prints a runner-style summary and exits 1 unless both pass.
std::fs::write(
dir.path().join("t.py"),
"import sys\nfrom m import f, g\nbad = (f(3, 2) != 5) + (g(3, 2) != 1)\nprint(f'{bad} failed, {2 - bad} passed')\nsys.exit(1 if bad else 0)\n",
)
.unwrap();
let cands: Vec<Candidate> = [(2, " return a * b"), (2, " return a + b"), (6, " return a // b")]
.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, Duration::from_secs(20), 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);
assert!(found.output.contains("1 failed"), "{}", found.output);
}
}
crates/hanzo-kai/src/program.rs
+1 -2
View File
@@ -26,8 +26,7 @@ pub const SHIPPED: &[&str] = &[
include_str!("../programs/[email protected]"),
include_str!("../programs/[email protected]"),
include_str!("../programs/[email protected]"),
include_str!("../programs/fix.pick@1.json"),
include_str!("../programs/[email protected]"),
include_str!("../programs/fix.turn@1.json"),
];
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
crates/hanzo-kai/src/symbols.rs
+879
View File
@@ -0,0 +1,879 @@
//! A working tree read-only, and the functions, classes and methods in it.
//!
//! kai-gateway's `repomap.ts`, so the ranker fit on its map ([`crate::locate`]) reads the same
//! symbols here. Nothing is executed (not even `git`); nothing outside the tree and nothing that
//! looks like a secret is read.
use regex::Regex;
use std::collections::HashMap;
use std::collections::VecDeque;
use std::path::Path;
use std::path::PathBuf;
use std::sync::LazyLock;
const SKIP: &[&str] = &["node_modules", "target", "dist", "build", "__pycache__"];
const LIST: usize = 20_000;
const READ: u64 = 200 * 1024;
const FILES: usize = 400;
const BYTES: usize = 2 * 1024 * 1024;
/// What a string literal's text becomes in [`mask`]: no space, so a line holding only it is not
/// blank.
const FILL: u8 = 0;
fn re(pattern: &str) -> Regex {
Regex::new(pattern).expect("a fixed pattern")
}
static SECRET: LazyLock<Regex> = LazyLock::new(|| {
re(r"(?i)^(?:\.env.*|.*\.(?:pem|key|p12|pfx)|id_(?:rsa|dsa|ecdsa|ed25519).*|credentials.*|\.npmrc|\.netrc|\.pypirc)$")
});
static TEST_FILE: LazyLock<Regex> = LazyLock::new(|| {
re(r"^(?:test_.*\.py|.*_test\.py|conftest\.py|.*\.(?:test|spec)\.[cm]?[jt]sx?|.*_test\.rs)$")
});
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Kind {
Fn,
Class,
Method,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Sym {
pub file: String,
pub name: String,
/// `Class.method`, `outer.inner`, or the name alone; `.` on every language.
pub qual: String,
pub kind: Kind,
/// 1-based, decorators and attributes included.
pub start: u32,
pub end: u32,
}
#[derive(Debug, Clone, Default)]
pub struct Map {
pub files: Vec<String>,
/// Code files that are not tests.
pub sources: Vec<String>,
pub tests: Vec<String>,
pub symbols: Vec<Sym>,
}
/// A name that may hold a credential, refused on a path's every segment.
fn secret(path: &str) -> bool {
path.split('/').any(|part| SECRET.is_match(part))
}
/// Whether a path is test code, by the conventions of pytest, cargo and the JS runners.
pub fn is_test(path: &str) -> bool {
let parts: Vec<&str> = path.split('/').collect();
let (name, dirs) = parts.split_last().map_or(("", &[][..]), |(n, d)| (*n, d));
dirs.iter()
.any(|d| matches!(*d, "tests" | "test" | "testing" | "__tests__"))
|| TEST_FILE.is_match(name)
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Lang {
Py,
Rs,
Js,
}
fn lang(path: &str) -> Option<Lang> {
let ext = path.rsplit_once('.')?.1.to_lowercase();
match ext.as_str() {
"py" => Some(Lang::Py),
"rs" => Some(Lang::Rs),
"ts" | "tsx" | "js" | "jsx" | "mjs" | "cjs" | "mts" | "cts" => Some(Lang::Js),
_ => None,
}
}
/// The tree at `root`, listed once and read through a cache.
pub struct Tree {
root: PathBuf,
listed: Vec<String>,
texts: HashMap<String, Option<String>>,
}
impl Tree {
/// Lists the tree: relative `/`-separated paths, directories breadth first and entries by
/// name; hidden directories, dependencies and build output left out, and no link followed.
pub fn open(root: &Path) -> Tree {
let mut listed = Vec::new();
if let Ok(real) = root.canonicalize() {
let mut pending = VecDeque::from([String::new()]);
while let Some(dir) = pending.pop_front() {
if listed.len() >= LIST {
break;
}
let Ok(entries) = std::fs::read_dir(real.join(&dir)) else {
continue;
};
let mut entries: Vec<_> = entries.flatten().collect();
entries.sort_by_key(|e| e.file_name());
for entry in entries {
let name = entry.file_name().to_string_lossy().into_owned();
let path = if dir.is_empty() {
name.clone()
} else {
format!("{dir}/{name}")
};
let Ok(kind) = entry.file_type() else {
continue;
};
if kind.is_dir() {
if !name.starts_with('.') && !SKIP.contains(&name.as_str()) {
pending.push_back(path);
}
} else if kind.is_file() && !secret(&name) && listed.len() < LIST {
listed.push(path);
}
}
}
}
Tree {
root: root.to_path_buf(),
listed,
texts: HashMap::new(),
}
}
/// A file already read by [`Tree::read`].
pub fn text(&self, path: &str) -> Option<&str> {
self.texts.get(path)?.as_deref()
}
/// The tree's own spelling of a path already read, to key borrows of its text by.
pub fn key(&self, path: &str) -> &str {
self.texts.get_key_value(path).map_or("", |(k, _)| k.as_str())
}
pub fn root(&self) -> &Path {
&self.root
}
pub fn list(&self) -> &[String] {
&self.listed
}
/// The file's text: UTF-8 (lossy), at most 200 KB, a regular file under the root. `None`
/// for anything else.
pub fn read(&mut self, path: &str) -> Option<&str> {
if !self.texts.contains_key(path) {
let text = self.load(path);
self.texts.insert(path.to_string(), text);
}
self.texts[path].as_deref()
}
fn load(&self, path: &str) -> Option<String> {
if path.is_empty() || path.starts_with('/') || path.contains('\0') || secret(path) {
return None;
}
let parts: Vec<&str> = path.split('/').collect();
let dirs = &parts[..parts.len() - 1];
if dirs.iter().any(|p| p.starts_with('.') || SKIP.contains(p)) {
return None;
}
let full = self.root.canonicalize().ok()?.join(path);
let meta = std::fs::symlink_metadata(&full).ok()?;
if !meta.is_file() || meta.len() > READ {
return None;
}
// A directory on the way may be a link out of the tree.
if !full.canonicalize().ok()?.starts_with(self.root.canonicalize().ok()?) {
return None;
}
let bytes = std::fs::read(&full).ok()?;
if bytes.contains(&0) {
return None;
}
Some(String::from_utf8_lossy(&bytes).into_owned())
}
/// The files, code split into sources and tests, and the symbols of the first 400 code
/// files (sources first) or 2 MB.
pub fn map(&mut self) -> Map {
let files = self.listed.clone();
let code: Vec<&String> = files.iter().filter(|f| lang(f).is_some()).collect();
let sources: Vec<String> = code.iter().filter(|f| !is_test(f)).map(|f| (*f).clone()).collect();
let tests: Vec<String> = code.iter().filter(|f| is_test(f)).map(|f| (*f).clone()).collect();
let mut symbols = Vec::new();
let (mut count, mut bytes) = (0, 0);
for file in sources.iter().chain(&tests) {
if count >= FILES || bytes >= BYTES {
break;
}
let Some(text) = self.read(file) else {
continue;
};
count += 1;
bytes += text.encode_utf16().count();
symbols.extend(symbols_of(file, text));
}
Map {
files,
sources,
tests,
symbols,
}
}
}
/// The functions, classes and methods a file declares; empty for a language not read.
pub fn symbols_of(path: &str, text: &str) -> Vec<Sym> {
match lang(path) {
Some(Lang::Py) => python(path, text),
Some(l @ (Lang::Rs | Lang::Js)) => braced(path, text, l),
None => Vec::new(),
}
}
/// `text` as [`mask`] makes it for the language of `path`: comments blank, string contents
/// filled, offsets and lines the source's. `None` for a language not read.
pub fn masked(path: &str, text: &str) -> Option<Vec<u8>> {
lang(path).map(|l| mask(text, l))
}
/// `text` with every comment blanked to spaces and every string literal but its first character
/// to FILL, newlines kept: offsets and lines are the source's, nothing quoted can look like
/// code, and a line that starts inside a string (FILL first) is told from one that starts a
/// statement.
fn mask(text: &str, language: Lang) -> Vec<u8> {
let t = text.as_bytes();
let mut out = t.to_vec();
let n = t.len();
let eol = |from: usize| t[from..].iter().position(|b| *b == b'\n').map_or(n, |p| from + p);
let blank = |out: &mut Vec<u8>, from: usize, to: usize| {
for b in &mut out[from..to] {
if *b != b'\n' {
*b = b' ';
}
}
};
let quote = |out: &mut Vec<u8>, i: usize, end: usize| {
for b in &mut out[(i + 1).min(end)..end] {
if *b != b'\n' {
*b = FILL;
}
}
};
// The end of a quoted literal opened before `start` with `q`; escapes skip a character, and
// a one-line quote ends at a newline.
let close = |start: usize, q: &[u8], lines: bool| -> usize {
let mut j = start;
while j < n {
if t[j] == b'\\' {
j += 2;
} else if t[j..].starts_with(q) {
return j + q.len();
} else if !lines && t[j] == b'\n' {
return j;
} else {
j += 1;
}
}
n
};
let mut i = 0;
while i < n {
let ch = t[i];
if language == Lang::Py {
if ch == b'#' {
let end = eol(i);
blank(&mut out, i, end);
i = end;
} else if ch == b'"' || ch == b'\'' {
let triple = t[i..].starts_with(&[ch, ch, ch]);
let end = if triple {
close(i + 3, &[ch, ch, ch], true)
} else {
close(i + 1, &[ch], false)
};
quote(&mut out, i, end.min(n));
i = end.min(n).max(i + 1);
} else {
i += 1;
}
continue;
}
if t[i..].starts_with(b"//") {
let end = eol(i);
blank(&mut out, i, end);
i = end;
} else if t[i..].starts_with(b"/*") {
// Rust's block comments nest; JavaScript's do not.
let (mut depth, mut j) = (0i32, i);
while j < n {
if t[j..].starts_with(b"/*") {
depth += i32::from(language == Lang::Rs || depth == 0);
j += 2;
} else if t[j..].starts_with(b"*/") {
depth -= 1;
j += 2;
if depth == 0 {
break;
}
} else {
j += 1;
}
}
blank(&mut out, i, j.min(n));
i = j.min(n);
} else if ch == b'"' {
let end = close(i + 1, b"\"", language == Lang::Rs).min(n);
quote(&mut out, i, end);
i = end.max(i + 1);
} else if language == Lang::Js && (ch == b'\'' || ch == b'`') {
let end = close(i + 1, &[ch], ch == b'`').min(n);
quote(&mut out, i, end);
i = end.max(i + 1);
} else if language == Lang::Rs
&& (ch == b'r' || ch == b'b')
&& !(i > 0 && (t[i - 1].is_ascii_alphanumeric() || t[i - 1] == b'_'))
{
// b?r#*"
let mut j = i;
if t[j] == b'b' {
j += 1;
}
if j < n && t[j] == b'r' {
let mut k = j + 1;
while k < n && t[k] == b'#' {
k += 1;
}
if k < n && t[k] == b'"' {
let hashes = &t[j + 1..k];
let mut closer = vec![b'"'];
closer.extend_from_slice(hashes);
let end = find(t, &closer, k + 1).map_or(n, |p| p + closer.len());
quote(&mut out, i, end);
i = end.max(i + 1);
continue;
}
}
i += 1;
} else if language == Lang::Rs && ch == b'\'' {
// A char literal, not a lifetime ('a): one character or one escape, then the
// closing quote.
match char_literal(t, i) {
Some(len) => {
quote(&mut out, i, i + len);
i += len;
}
None => i += 1,
}
} else {
i += 1;
}
}
out
}
fn find(hay: &[u8], needle: &[u8], from: usize) -> Option<usize> {
if from >= hay.len() {
return None;
}
hay[from..]
.windows(needle.len())
.position(|w| w == needle)
.map(|p| from + p)
}
/// The length of a Rust char literal opened at `i`, if one is.
fn char_literal(t: &[u8], i: usize) -> Option<usize> {
let rest = std::str::from_utf8(&t[i..t.len().min(i + 16)]).ok().or_else(|| {
// A multi-byte character may cut the window; try shorter windows.
(1..16).rev().find_map(|k| std::str::from_utf8(&t[i..t.len().min(i + k)]).ok())
})?;
static CHAR: LazyLock<Regex> = LazyLock::new(|| {
re(r"^'(?:[^'\\\n]|\\(?:x[0-9a-fA-F]{2}|u\{[0-9a-fA-F]{1,6}\}|.))'")
});
CHAR.find(rest).map(|m| m.end())
}
/// Columns of indentation, a tab to the next multiple of 8.
fn indent(line: &[u8]) -> usize {
let mut width = 0;
for b in line {
match b {
b' ' => width += 1,
b'\t' => width += 8 - (width % 8),
_ => break,
}
}
width
}
fn balance(line: &[u8]) -> i64 {
line.iter().fold(0, |n, b| match b {
b'(' | b'[' | b'{' => n + 1,
b')' | b']' | b'}' => n - 1,
_ => n,
})
}
static DEF: LazyLock<Regex> =
LazyLock::new(|| re(r"^\s*(?:async\s+)?(def|class)\s+([A-Za-z_]\w*)"));
static DECORATOR: LazyLock<Regex> = LazyLock::new(|| re(r"^\s*@"));
fn first_char(line: &[u8]) -> Option<u8> {
line.iter().copied().find(|b| *b != b' ' && *b != b'\t' && *b != b'\r')
}
/// Python by indentation: a block ends before the next statement indented no deeper than its
/// header. Comment lines, blank in the mask, neither end a block nor belong to it; a line inside
/// a string does not end one, wherever it is indented.
fn python(file: &str, text: &str) -> Vec<Sym> {
let masked = mask(text, Lang::Py);
let lines: Vec<&[u8]> = masked.split(|b| *b == b'\n').collect();
let mut found: Vec<Sym> = Vec::new();
// Indices into `found` of the symbols still open.
let mut open: Vec<usize> = Vec::new();
for i in 0..lines.len() {
let line = String::from_utf8_lossy(lines[i]);
let Some(m) = DEF.captures(&line) else {
continue;
};
let level = indent(lines[i]);
// The header runs until its brackets close: a signature can span lines.
let mut header = i;
let mut depth = balance(lines[i]);
while depth > 0 && header + 1 < lines.len() {
header += 1;
depth += balance(lines[header]);
}
let mut end = header;
for (k, l) in lines.iter().enumerate().skip(header + 1) {
let Some(first) = first_char(l) else {
continue;
};
if first != FILL && indent(l) <= level {
break;
}
end = k;
}
// Decorators directly above, a call spread over lines included.
let mut start = i;
let mut depth = 0i64;
let mut k = i;
while k > 0 {
k -= 1;
depth -= balance(lines[k]);
if depth < 0 {
break;
}
if depth > 0 {
continue;
}
let l = String::from_utf8_lossy(lines[k]);
if !DECORATOR.is_match(&l) || indent(lines[k]) != level {
break;
}
start = k;
}
while open.last().is_some_and(|p| found[*p].end < (i + 1) as u32) {
open.pop();
}
let parent = open.last().map(|p| &found[*p]);
let name = m[2].to_string();
let sym = Sym {
file: file.to_string(),
qual: parent.map_or(name.clone(), |p| format!("{}.{}", p.qual, name)),
kind: if &m[1] == "class" {
Kind::Class
} else if parent.is_some_and(|p| p.kind == Kind::Class) {
Kind::Method
} else {
Kind::Fn
},
name,
start: (start + 1) as u32,
end: (end + 1) as u32,
};
found.push(sym);
open.push(found.len() - 1);
}
found
}
#[derive(Debug, Clone)]
struct Item {
name: String,
/// `container` (Rust's impl and mod) only scopes the items in it.
kind: ItemKind,
/// Offset of the keyword, and of the body's braces (`close` the `;` of a body-less struct).
at: usize,
open: usize,
close: usize,
/// `impl` and `trait` make the functions in them methods.
owner: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ItemKind {
Fn,
Class,
Container,
}
enum Body {
Braces(usize, usize),
Semi(usize),
}
/// Rust, and JavaScript and TypeScript, by their braces over the masked text.
fn braced(file: &str, text: &str, language: Lang) -> Vec<Sym> {
let masked_bytes = mask(text, language);
let masked = String::from_utf8_lossy(&masked_bytes).into_owned();
let m = masked.as_bytes();
let mut pairs: HashMap<usize, usize> = HashMap::new();
let mut stack = Vec::new();
for (i, b) in m.iter().enumerate() {
if *b == b'{' {
stack.push(i);
} else if *b == b'}'
&& let Some(o) = stack.pop()
{
pairs.insert(o, i);
}
}
let mut starts = vec![0usize];
for (i, b) in m.iter().enumerate() {
if *b == b'\n' {
starts.push(i + 1);
}
}
let line_at = |offset: usize| -> u32 { starts.partition_point(|s| *s <= offset) as u32 };
// The body after a declaration: its first `{` outside parentheses and brackets, or the `;`
// that means it has none.
let body = |from: usize| -> Option<Body> {
let mut depth = 0i64;
for (i, ch) in m.iter().enumerate().skip(from) {
match ch {
b'(' | b'[' => depth += 1,
b')' | b']' => depth -= 1,
b';' if depth <= 0 => return Some(Body::Semi(i)),
b'}' if depth <= 0 => return None,
b'{' if depth <= 0 => return pairs.get(&i).map(|c| Body::Braces(i, *c)),
_ => {}
}
}
None
};
let mut items: Vec<Item> = Vec::new();
let add = |items: &mut Vec<Item>, name: &str, kind: ItemKind, at: usize, from: usize, owner: bool| {
match body(from) {
Some(Body::Braces(open, close)) => items.push(Item {
name: name.to_string(),
kind,
at,
open,
close,
owner,
}),
Some(Body::Semi(semi)) if kind == ItemKind::Class => items.push(Item {
name: name.to_string(),
kind,
at,
open: semi,
close: semi,
owner,
}),
_ => {}
}
};
if language == Lang::Rs {
static FN: LazyLock<Regex> = LazyLock::new(|| re(r"\bfn\s+(?:r#)?([A-Za-z_]\w*)"));
static TYPE: LazyLock<Regex> =
LazyLock::new(|| re(r"\b(struct|enum|union|trait)\s+([A-Za-z_]\w*)"));
static MOD: LazyLock<Regex> = LazyLock::new(|| re(r"\bmod\s+([A-Za-z_]\w*)\s*\{"));
static IMPL: LazyLock<Regex> = LazyLock::new(|| {
re(r"(?m)^[ \t]*(?:pub(?:\([^)]*\))?\s+)?(?:unsafe\s+)?(?:default\s+)?impl\b")
});
for c in FN.captures_iter(&masked) {
let all = c.get(0).unwrap();
add(&mut items, &c[1], ItemKind::Fn, all.start(), all.end(), false);
}
for c in TYPE.captures_iter(&masked) {
let all = c.get(0).unwrap();
add(&mut items, &c[2], ItemKind::Class, all.start(), all.end(), &c[1] == "trait");
}
for c in MOD.captures_iter(&masked) {
let all = c.get(0).unwrap();
add(&mut items, &c[1], ItemKind::Container, all.start(), all.end() - 1, false);
}
for c in IMPL.find_iter(&masked) {
if let Some(Body::Braces(open, close)) = body(c.end()) {
items.push(Item {
name: impl_type(&masked[c.end()..open]),
kind: ItemKind::Container,
at: c.start(),
open,
close,
owner: true,
});
}
}
} else {
static FUNCTION: LazyLock<Regex> =
LazyLock::new(|| re(r"\bfunction\b\s*\*?\s*([A-Za-z_$][\w$]*)"));
static CLASS: LazyLock<Regex> = LazyLock::new(|| re(r"\bclass\s+([A-Za-z_$][\w$]*)"));
static ASSIGNED: LazyLock<Regex> = LazyLock::new(|| {
re(r"\b(?:const|let|var)\s+([A-Za-z_$][\w$]*)\s*(?::[^=;]*)?=\s*(?:async\s+)?")
});
static LEAD: LazyLock<Regex> =
LazyLock::new(|| re(r"^(?:function\b|\(|[A-Za-z_$][\w$]*\s*=>)"));
static ARROW: LazyLock<Regex> = LazyLock::new(|| {
re(r"^(?:function\b[^{;]*|(?:\([^;]*?\)|[A-Za-z_$][\w$]*)\s*(?::[^=;{]*)?=>\s*)(\{)?")
});
static METHOD: LazyLock<Regex> = LazyLock::new(|| {
re(r"(?:^|[;}\n])\s*(?:(?:public|private|protected|static|async|readonly|override|abstract|get|set)\s+|\*\s*)*(#?[A-Za-z_$][\w$]*)\s*(?:<[^>(){};]*>)?\s*\(")
});
for c in FUNCTION.captures_iter(&masked) {
let all = c.get(0).unwrap();
add(&mut items, &c[1], ItemKind::Fn, all.start(), all.end(), false);
}
for c in CLASS.captures_iter(&masked) {
if &c[1] == "extends" {
continue;
}
let all = c.get(0).unwrap();
if let Some(Body::Braces(open, close)) = body(all.end()) {
items.push(Item {
name: c[1].to_string(),
kind: ItemKind::Class,
at: all.start(),
open,
close,
owner: true,
});
}
}
for c in ASSIGNED.captures_iter(&masked) {
let all = c.get(0).unwrap();
// An arrow function or a function expression, up to its body or the end of its
// statement.
let tail_end = (all.end() + 2_000).min(masked.len());
let tail_end = (all.end()..=tail_end).rev().find(|e| masked.is_char_boundary(*e)).unwrap_or(all.end());
let rest = &masked[all.end()..tail_end];
if !LEAD.is_match(rest) {
continue;
}
let Some(arrow) = ARROW.captures(rest) else {
continue;
};
if arrow.get(1).is_some() {
let open = all.end() + arrow.get(0).unwrap().end() - 1;
if let Some(close) = pairs.get(&open) {
items.push(Item {
name: c[1].to_string(),
kind: ItemKind::Fn,
at: all.start(),
open,
close: *close,
owner: false,
});
}
} else {
let semi = masked[all.end()..].find(';').map(|p| all.end() + p);
let newline = masked[all.end()..].find('\n').map(|p| all.end() + p);
let end = semi.or(newline).unwrap_or(masked.len() - 1);
items.push(Item {
name: c[1].to_string(),
kind: ItemKind::Fn,
at: all.start(),
open: end,
close: end,
owner: false,
});
}
}
// Methods: a name and its parameters at the top level of a class body.
let classes: Vec<Item> = items.iter().filter(|i| i.kind == ItemKind::Class).cloned().collect();
for owner in classes {
if owner.open == owner.close {
continue;
}
let inner = &masked[owner.open + 1..owner.close];
for c in METHOD.captures_iter(inner) {
let name = c[1].to_string();
if ["if", "for", "while", "switch", "catch", "return", "function", "with"]
.contains(&name.as_str())
{
continue;
}
let whole = c.get(0).unwrap();
let at = owner.open + 1 + whole.start() + whole.as_str().rfind(&name).unwrap_or(0);
if innermost(&pairs, at) != Some(owner.open) {
continue;
}
if let Some(Body::Braces(open, close)) = body(owner.open + 1 + whole.end() - 1) {
items.push(Item {
name,
kind: ItemKind::Fn,
at,
open,
close,
owner: false,
});
}
}
}
}
items.sort_by_key(|i| i.at);
let lines: Vec<&str> = masked.split('\n').collect();
let mut found = Vec::new();
for (index, item) in items.iter().enumerate() {
if item.kind == ItemKind::Container {
continue;
}
// The items whose bodies hold this one, outermost first.
let mut chain: Vec<&Item> = items
.iter()
.enumerate()
.filter(|(j, o)| *j != index && o.open < item.at && item.at < o.close && o.open != o.close)
.map(|(_, o)| o)
.collect();
chain.sort_by_key(|o| o.open);
let parent = chain.last();
let path: Vec<&str> = chain
.iter()
.filter(|o| o.kind != ItemKind::Container || o.owner)
.map(|o| o.name.as_str())
.collect();
let mut start = line_at(item.at);
// Attributes (`#[test]`) directly above belong to the item.
if language == Lang::Rs {
while start > 1
&& lines
.get(start as usize - 2)
.is_some_and(|l| l.trim_start().starts_with("#[") || l.trim_start().starts_with("#!["))
{
start -= 1;
}
}
let mut qual = path.join(".");
if !qual.is_empty() {
qual.push('.');
}
qual.push_str(&item.name);
found.push(Sym {
file: file.to_string(),
name: item.name.clone(),
qual,
kind: if item.kind == ItemKind::Class {
Kind::Class
} else if parent.is_some_and(|p| p.owner) {
Kind::Method
} else {
Kind::Fn
},
start,
end: line_at(item.close),
});
}
found
}
/// The innermost brace pair holding `offset`, by its opening brace.
fn innermost(pairs: &HashMap<usize, usize>, offset: usize) -> Option<usize> {
pairs
.iter()
.filter(|(open, close)| **open < offset && offset < **close)
.map(|(open, _)| *open)
.max()
}
/// The type an impl block is for: `impl<T> Display for Wrapper<T> where …` is `Wrapper`.
fn impl_type(header: &str) -> String {
static GENERIC: LazyLock<Regex> = LazyLock::new(|| re(r"<[^<>]*>"));
static WORD: LazyLock<Regex> = LazyLock::new(|| re(r"[A-Za-z_]\w*"));
let mut text = header.to_string();
loop {
let next = GENERIC.replace_all(&text, " ").into_owned();
if next == text {
break;
}
text = next;
}
let text = text.split("where").next().unwrap_or("").to_string();
let target = text.rsplit_once("for").map_or(text.as_str(), |(_, after)| after);
WORD.find_iter(target)
.map(|m| m.as_str())
.filter(|w| *w != "dyn" && *w != "mut")
.last()
.unwrap_or("impl")
.to_string()
}
#[cfg(test)]
mod tests {
use super::*;
use pretty_assertions::assert_eq;
fn quals(path: &str, text: &str) -> Vec<(String, Kind, u32, u32)> {
symbols_of(path, text)
.into_iter()
.map(|s| (s.qual, s.kind, s.start, s.end))
.collect()
}
#[test]
fn python_methods_nest_under_their_class_with_decorators_included() {
let src = "class A:\n @property\n def x(self):\n return '''\ndef not_code():\n'''\n\n def y(self):\n pass\n\ndef top():\n pass\n";
assert_eq!(
quals("a.py", src),
vec![
("A".into(), Kind::Class, 1, 9),
("A.x".into(), Kind::Method, 2, 6),
("A.y".into(), Kind::Method, 8, 9),
("top".into(), Kind::Fn, 11, 12),
]
);
}
#[test]
fn a_python_signature_spread_over_lines_is_one_header() {
let src = "def f(a,\n b):\n return a\n";
assert_eq!(quals("a.py", src), vec![("f".into(), Kind::Fn, 1, 3)]);
}
#[test]
fn rust_impl_methods_are_methods_of_the_type_and_attributes_belong_to_the_item() {
let src = "struct W;\nimpl W {\n #[inline]\n fn go(&self) -> u8 {\n let s = \"fn fake() {}\";\n 1\n }\n}\nfn free() {}\n";
assert_eq!(
quals("a.rs", src),
vec![
("W".into(), Kind::Class, 1, 1),
("W.go".into(), Kind::Method, 3, 7),
("free".into(), Kind::Fn, 9, 9),
]
);
}
#[test]
fn javascript_functions_arrows_and_class_methods_are_found() {
let src = "function a() {\n return 1;\n}\nconst b = (x) => {\n return x;\n};\nclass C {\n m(x) {\n return x;\n }\n}\n";
let got = quals("a.js", src);
let names: Vec<_> = got.iter().map(|g| g.0.as_str()).collect();
assert_eq!(names, vec!["a", "b", "C", "C.m"]);
}
#[test]
fn tests_are_told_from_sources_and_secrets_are_not_listed() {
assert!(is_test("tests/test_a.py") && is_test("pkg/a_test.py") && is_test("web/a.spec.ts"));
assert!(!is_test("pkg/a.py"));
let dir = tempfile::tempdir().unwrap();
std::fs::write(dir.path().join("a.py"), "def f():\n pass\n").unwrap();
std::fs::write(dir.path().join(".env"), "KEY=1\n").unwrap();
std::fs::write(dir.path().join("id_rsa"), "x").unwrap();
let mut tree = Tree::open(dir.path());
assert_eq!(tree.list(), ["a.py"]);
assert!(tree.read(".env").is_none());
assert_eq!(tree.map().symbols.len(), 1);
}
}
crates/hanzo-kai/src/testing.rs
+15 -13
View File
@@ -115,30 +115,32 @@ pub fn on(q: &Question, label: &str, p: f64) -> Vec<f64> {
.collect()
}
/// Zen answering chat completions by `rule(inputs)`, the inputs being the JSON a flow sends.
pub async fn zen(rule: impl Fn(&Value) -> String + Send + Sync + 'static) -> MockServer {
/// Zen answering streamed chat completions by `rule(prompt)`, the prompt being the user message
/// a flow sent.
pub async fn zen(rule: impl Fn(&str) -> String + Send + Sync + 'static) -> MockServer {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/chat/completions"))
.respond_with(move |request: &wiremock::Request| {
let content = inputs(request);
ResponseTemplate::new(200).set_body_json(json!({
"choices": [{"message": {"role": "assistant", "content": rule(&content)}}],
}))
let answer = rule(&prompt(request));
let first = json!({"model": "zen-test", "choices": [{"delta": {"content": answer}}]});
let last = json!({"model": "zen-test", "choices": [], "usage": {"prompt_tokens": 100, "completion_tokens": 20}});
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(format!("data: {first}\n\ndata: {last}\n\ndata: [DONE]\n\n"))
})
.mount(&server)
.await;
server
}
/// The inputs of a chat completion a flow sent.
pub fn inputs(request: &wiremock::Request) -> Value {
/// The user message of a chat completion a flow sent.
pub fn prompt(request: &wiremock::Request) -> String {
let body: Value = serde_json::from_slice(&request.body).unwrap();
let text = body["messages"][1]["content"].as_str().unwrap_or_default();
let json = text
.split_once("Inputs, as JSON:\n")
.map_or("null", |(_, j)| j);
serde_json::from_str(json).unwrap()
body["messages"][1]["content"]
.as_str()
.unwrap_or_default()
.to_string()
}
/// What Kai did to a thread.
crates/hanzo-kai/src/tools.rs
+12 -3
View File
@@ -143,12 +143,20 @@ pub fn tail(text: &str, n: usize) -> String {
/// Runs `sh -c script` in `root`, killing it after `limit`: `{code, output, ms}`, output
/// being the tail of stdout and stderr together.
pub fn shell(root: &Path, script: &str, limit: Duration) -> Value {
shell_until(root, script, limit, &|| false)
}
/// [`shell`], also killed as soon as `stop` says so (code -3).
pub fn shell_until(root: &Path, script: &str, limit: Duration, stop: &dyn Fn() -> bool) -> Value {
let started = Instant::now();
let mut command = Command::new("sh");
command
.arg("-c")
.arg(format!("exec 2>&1\n{script}"))
.current_dir(root)
// An edit is read at once: Python validates a cached module by the source's second and
// size, so a same-size edit within the second of the last run would run stale code.
.env("PYTHONDONTWRITEBYTECODE", "1")
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::null());
@@ -172,16 +180,17 @@ pub fn shell(root: &Path, script: &str, limit: Duration) -> Value {
let code = loop {
match child.try_wait() {
Ok(Some(status)) => break status.code().unwrap_or(-1),
Ok(None) if started.elapsed() > limit => {
Ok(None) if started.elapsed() > limit || stop() => {
let timed = started.elapsed() > limit;
#[cfg(unix)]
let _ = Command::new("kill")
.args(["-9", &format!("-{}", child.id())])
.status();
let _ = child.kill();
let _ = child.wait();
break -2;
break if timed { -2 } else { -3 };
}
Ok(None) => std::thread::sleep(Duration::from_millis(20)),
Ok(None) => std::thread::sleep(Duration::from_millis(5)),
Err(_) => break -1,
}
};
crates/hanzo-kai/src/zen.rs
+208 -34
View File
@@ -1,32 +1,77 @@
//! Zen through the Hanzo API: the only model calls a flow makes, and only to write code or
//! text.
//! Zen through an OpenAI-compatible endpoint: the only model calls a flow makes, and only to
//! write code or text. The endpoint is Hanzo's (`https://api.hanzo.ai/v1`, signed with the
//! credential Dev already uses) or any server that speaks `/chat/completions`, a local Zen 6
//! among them; the credential is sent to Hanzo's alone.
//!
//! A call streams, so the time to the first token and the time to the last are both known, and
//! each paid call is a line in the spend file with the model Hanzo says served it.
use hanzo_config::kai::Generate;
use serde_json::Value;
use serde_json::json;
use std::io::Write;
use std::path::PathBuf;
use std::time::Duration;
use std::time::Instant;
const SYSTEM: &str = "You are Zen, writing code for Hanzo Dev. Answer with exactly what the \
request asks for and nothing else. When it asks for a change, answer with one unified diff \
against the repository root that `git apply` accepts, in a single ```diff block.";
/// Tokens an answer may be given room for.
const TOKENS: u32 = 512;
/// Tokens an answer that reasons first may be given room for.
const THOUGHT: u32 = 6144;
/// What an unpriced call is counted at, dollars per million tokens: above any Zen model's rate,
/// so the budget is a ceiling on the spend and not a guess at it.
const RATE: (f64, f64) = (2.0, 10.0);
/// Tries on a server that says it is busy.
const TRIES: u32 = 3;
/// Zen over `POST /v1/chat/completions`.
/// How much the model reasons before it answers.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Effort {
/// None: the answer is the edit and nothing else.
Low,
/// The model's own reasoning first.
Full,
}
/// One answer and what it took.
#[derive(Debug, Clone, PartialEq)]
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.
pub total: f64,
pub prompt: u64,
pub completion: u64,
/// The model the server says answered.
pub served: String,
}
/// Zen over `POST {url}/chat/completions` as `model`.
pub struct Zen {
base: String,
key: Option<String>,
url: String,
model: String,
key: Option<String>,
cloud: bool,
spend: Option<(PathBuf, f64)>,
client: reqwest::Client,
runtime: tokio::runtime::Handle,
}
impl Zen {
/// Zen at `base` (`https://api.hanzo.ai/v1`) as `model`, signed with `key`. Call from
/// within the async runtime; generation then blocks a thread off it.
pub fn new(base: &str, key: Option<String>, model: &str) -> Zen {
/// Zen as `generate` names it, signed with `key` when its URL is Hanzo's. Call from within
/// the async runtime; generation then blocks a thread off it.
pub fn new(generate: &Generate, key: Option<String>) -> Zen {
let cloud = hanzo_config::is_hanzo(Some(&generate.url));
Zen {
base: base.trim_end_matches('/').to_string(),
key,
model: model.to_string(),
url: generate.url.trim_end_matches('/').to_string(),
model: generate.model.clone(),
key: key.filter(|_| cloud),
cloud,
spend: cloud
.then(|| (generate.spend.clone(), generate.budget.unwrap_or(f64::INFINITY))),
client: reqwest::Client::builder()
.connect_timeout(Duration::from_secs(5))
.timeout(Duration::from_secs(600))
.build()
.unwrap_or_default(),
@@ -34,42 +79,171 @@ impl Zen {
}
}
async fn complete(&self, prompt: &str, context: &Value) -> Result<String, String> {
let body = json!({
pub fn model(&self) -> &str {
&self.model
}
/// Dollars the spend file holds so far.
fn spent(path: &PathBuf) -> f64 {
std::fs::read_to_string(path)
.unwrap_or_default()
.lines()
.filter_map(|l| serde_json::from_str::<Value>(l).ok()?["cost"].as_f64())
.sum()
}
fn record(&self, g: &Generation, headers: &Value) {
let Some((path, _)) = &self.spend else {
return;
};
let cost = (g.prompt as f64 * RATE.0 + g.completion as f64 * RATE.1) / 1e6;
let line = json!({
"time": crate::trace::now(),
"model": self.model,
"served": g.served,
"prompt_tokens": g.prompt,
"completion_tokens": g.completion,
"ms": g.total,
"cost": cost,
"cost_basis": format!("ceiling at ${}/M in, ${}/M out", RATE.0, RATE.1),
"headers": headers,
});
if let Some(dir) = path.parent() {
let _ = std::fs::create_dir_all(dir);
}
if let Ok(mut f) = std::fs::OpenOptions::new().create(true).append(true).open(path) {
let _ = writeln!(f, "{line}");
}
}
async fn once(&self, system: &str, user: &str, effort: Effort) -> Result<(Generation, Value), String> {
let mut body = json!({
"model": self.model,
"temperature": 0,
"max_tokens": if effort == Effort::Full { THOUGHT } else { TOKENS },
"stream": true,
"stream_options": {"include_usage": true},
"messages": [
{"role": "system", "content": SYSTEM},
{"role": "user", "content": format!(
"{prompt}\n\nInputs, as JSON:\n{}",
serde_json::to_string_pretty(context).unwrap_or_default()
)},
{"role": "system", "content": system},
{"role": "user", "content": user},
],
});
// At low effort the answer is the edit and nothing else: no reasoning is bought.
match (self.cloud, effort) {
(true, Effort::Low) => body["reasoning_effort"] = json!("none"),
(true, Effort::Full) => body["reasoning_effort"] = json!("medium"),
(false, Effort::Low) => body["chat_template_kwargs"] = json!({"enable_thinking": false}),
(false, Effort::Full) => body["chat_template_kwargs"] = json!({"enable_thinking": true}),
}
let mut request = self
.client
.post(format!("{}/chat/completions", self.base))
.post(format!("{}/chat/completions", self.url))
.json(&body);
if let Some(key) = &self.key {
request = request.bearer_auth(key);
}
let response = request.send().await.map_err(|e| format!("zen: {e}"))?;
let begun = Instant::now();
let mut response = request.send().await.map_err(|e| format!("zen: {}", e.without_url()))?;
let status = response.status();
let answer: Value = response.json().await.map_err(|e| format!("zen: {e}"))?;
let headers = json!({
"served": response.headers().get("x-hanzo-served").and_then(|v| v.to_str().ok()),
"lane": response.headers().get("x-hanzo-lane").and_then(|v| v.to_str().ok()),
"paid_by": response.headers().get("x-hanzo-paid-by").and_then(|v| v.to_str().ok()),
});
if !status.is_success() {
let why = answer["error"]["message"]
let body = response.text().await.unwrap_or_default();
let why: Value = serde_json::from_str(&body).unwrap_or_default();
let msg = why["error"]["message"]
.as_str()
.unwrap_or("no reason given");
return Err(format!("zen: {status}: {why}"));
.or_else(|| why["error"].as_str())
.map(str::to_string)
.unwrap_or_else(|| crate::decide::cut(&body, 200));
return Err(format!("zen: {status}: {msg}"));
}
answer["choices"][0]["message"]["content"]
.as_str()
.map(str::to_string)
.ok_or_else(|| "zen: an answer without text".to_string())
let (mut pending, mut answer, mut first) = (String::new(), String::new(), None);
let (mut prompt_tokens, mut completion_tokens, mut served) = (0, 0, String::new());
while let Some(chunk) = response.chunk().await.map_err(|e| format!("zen: {}", e.without_url()))? {
pending.push_str(&String::from_utf8_lossy(&chunk));
while let Some(at) = pending.find('\n') {
let line: String = pending.drain(..=at).collect();
let Some(data) = line.trim().strip_prefix("data:") else {
continue;
};
let data = data.trim();
if data == "[DONE]" || data.is_empty() {
continue;
}
let Ok(event) = serde_json::from_str::<Value>(data) else {
continue;
};
let delta = &event["choices"][0]["delta"];
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);
}
if let Some(piece) = delta["content"].as_str()
&& !piece.is_empty()
{
first.get_or_insert_with(|| begun.elapsed().as_secs_f64() * 1000.0);
answer.push_str(piece);
}
if let Some(u) = event.get("usage").filter(|u| u.is_object()) {
prompt_tokens = u["prompt_tokens"].as_u64().unwrap_or(prompt_tokens);
completion_tokens = u["completion_tokens"].as_u64().unwrap_or(completion_tokens);
}
if let Some(m) = event["model"].as_str() {
served = m.to_string();
}
}
}
let total = begun.elapsed().as_secs_f64() * 1000.0;
if answer.is_empty() {
return Err("zen: an answer without text".into());
}
if let Some(s) = headers["served"].as_str() {
served = s.to_string();
}
Ok((
Generation {
text: answer,
first: first.unwrap_or(total),
total,
prompt: prompt_tokens,
completion: completion_tokens,
served,
},
headers,
))
}
/// Zen's answer to `prompt` over `context`. Blocks: call it off the async runtime.
pub fn generate(&self, prompt: &str, context: &Value) -> Result<String, String> {
self.runtime.block_on(self.complete(prompt, context))
async fn complete(&self, system: &str, user: &str, effort: Effort) -> Result<Generation, String> {
if let Some((path, budget)) = &self.spend
&& Self::spent(path) >= *budget
{
return Err(format!("zen: the budget of ${budget:.2} is spent ({})", path.display()));
}
let begun = Instant::now();
let mut last = String::new();
for attempt in 0..TRIES {
match self.once(system, user, effort).await {
Ok((mut g, headers)) => {
// A retry's wait is part of the time the answer took.
g.total = begun.elapsed().as_secs_f64() * 1000.0;
self.record(&g, &headers);
return Ok(g);
}
Err(e) if e.contains("503") || e.contains("502") || e.contains("429") => {
last = e;
tokio::time::sleep(Duration::from_millis(500 * u64::from(attempt + 1))).await;
}
Err(e) => return Err(e),
}
}
Err(last)
}
/// Zen's answer to `user` under `system`. Blocks: call it off the async runtime.
pub fn generate(&self, system: &str, user: &str, effort: Effort) -> Result<Generation, String> {
self.runtime.block_on(self.complete(system, user, effort))
}
}