Something went wrong. Try again.
Parametric learning in your code harness.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228import assert from "node:assert/strict";import { mkdir, mkdtemp, readFile, rm, symlink, writeFile } from "node:fs/promises";import os from "node:os";import path from "node:path";import { after, test } from "node:test";import { fileURLToPath } from "node:url";import { checkCase, diskStore, formatSpecialization, localSandboxes, machine, scriptedHarness, splitCases, type CaseDraft, type HarnessRequest } from "../src/index.ts";import { parseCases, parseDesign } from "../src/specialize.ts";
const repo = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "..");const dirs: string[] = [];after(() => Promise.all(dirs.map((d) => rm(d, { recursive: true, force: true }))));
async function tempDir() { const dir = await mkdtemp(path.join(os.tmpdir(), "machine-specialize-test-")); dirs.push(dir); return dir;}
const sandboxes = () => localSandboxes(path.join(os.tmpdir(), `machine-specialize-boxes-${process.pid}`), { isolation: "none" });
const BUGGY = "def mean(xs):\n return sum(xs) / len(xs)\n";const FIXED = "def mean(xs):\n if not xs:\n return 0\n return sum(xs) / len(xs)\n";const TEST = "import unittest\nfrom stats import mean\n\n\nclass T(unittest.TestCase):\n def test_empty(self):\n self.assertEqual(mean([]), 0)\n\n\nif __name__ == '__main__':\n unittest.main()\n";
function caseBlock(input: string, source: string, solution: string, test = TEST) { return `<case>\n<input>${input}</input>\n<file path="stats.py">\n${source}</file>\n<file path="test_stats.py">\n${test}</file>\n<solution path="stats.py">\n${solution}</solution>\n<test>python3 -m unittest -q</test>\n</case>`;}
test("parseDesign reads the tagged files and topics, and names what is missing", () => { const ok = parseDesign("<slug>Async</slug><rubric_ts>```ts\nr\n```</rubric_ts><machine_ts>m</machine_ts><instructions_md>i</instructions_md><topics>\n- races\n2. deadlocks\n</topics>"); assert.deepEqual(ok, { slug: "async", rubricTs: "r", machineTs: "m", instructions: "i", topics: ["races", "deadlocks"] }); assert.match(String(parseDesign("<slug>a</slug>")), /missing rubricTs, machineTs, instructions/); assert.match(String(parseDesign("<slug>Bad Slug!</slug><rubric_ts>r</rubric_ts><machine_ts>m</machine_ts><instructions_md>i</instructions_md>")), /slug/);});
test("parseCases reads files, the solution, and the test command", () => { const [c] = parseCases(caseBlock("mean([]) crashes", BUGGY, FIXED)); assert.equal(c.input, "mean([]) crashes"); assert.deepEqual(Object.keys(c.files!), ["stats.py", "test_stats.py"]); assert.equal(c.solution!["stats.py"], FIXED); assert.deepEqual(c.test, { command: "python3", args: ["-m", "unittest", "-q"] });});
test("checkCase keeps a test that fails before the fix and passes after, and discards the rest", async () => { const boxes = sandboxes(); const [good] = parseCases(caseBlock("mean([]) crashes", BUGGY, FIXED)); assert.equal(await checkCase(good, boxes, true), undefined); const [passing] = parseCases(caseBlock("already fine", FIXED, FIXED)); assert.match(String(await checkCase(passing, boxes, true)), /already passes/); const [wrongFix] = parseCases(caseBlock("bad fix", BUGGY, BUGGY)); assert.match(String(await checkCase(wrongFix, boxes, true)), /fail after the reference fix/); const cheat: CaseDraft = { ...good, solution: { "test_stats.py": "" } }; assert.match(String(await checkCase(cheat, boxes, true)), /changes a test file/); assert.match(String(await checkCase({ input: "x", files: { "a.py": "" } }, boxes, true)), /no test command/); assert.equal(await checkCase({ input: "no fixture needed" }, boxes, false), undefined); assert.match(String(await checkCase({ input: "x", files: { "../a.py": "" } }, boxes, false)), /relative path/);});
const HIDDEN = "import unittest\nfrom stats import mean\n\n\nclass H(unittest.TestCase):\n def test_one(self):\n self.assertEqual(mean([4]), 4)\n\n def test_empty_tuple(self):\n self.assertEqual(mean(()), 0)\n";
test("checkCase with hidden tests: they must be new test files that fail alone before the fix and pass with it", async () => { const boxes = sandboxes(); const [good] = parseCases(caseBlock("mean([]) crashes", BUGGY, FIXED).replace("<test>", `<hidden path="test_hidden_stats.py">\n${HIDDEN}</hidden>\n<test>`)); assert.deepEqual(Object.keys(good.hidden!), ["test_hidden_stats.py"]); assert.equal(await checkCase(good, boxes, true, true), undefined); const [none] = parseCases(caseBlock("mean([]) crashes", BUGGY, FIXED)); assert.match(String(await checkCase(none, boxes, true, true)), /no hidden tests/); const trivial: CaseDraft = { ...good, hidden: { "test_hidden_stats.py": "import unittest\n\n\nclass H(unittest.TestCase):\n def test_ok(self):\n pass\n" } }; assert.match(String(await checkCase(trivial, boxes, true, true)), /already pass before the fix/); const clobber: CaseDraft = { ...good, hidden: { "test_stats.py": HIDDEN } }; assert.match(String(await checkCase(clobber, boxes, true, true)), /new test file/);});
test("splitCases spreads held-out cases through the kept ones", () => { const { train, holdout } = splitCases(Array.from({ length: 28 }, (_, i) => i), 20, 8); assert.equal(train.length, 20); assert.equal(holdout.length, 8); assert.ok(holdout[0] < 4 && holdout.at(-1)! > 23, "spread from start to end"); const few = splitCases([1, 2, 3, 4, 5], 20, 8); assert.equal(few.train.length + few.holdout.length, 5); assert.ok(few.holdout.length >= 1);});
const PARENT = `import { check, defineEvaluator, defineMachine, defineRubric } from "../../src/index.ts";const rubric = defineRubric({ good: { description: 'Says "good".' } });export default defineMachine({ name: "toy", intent: "Fix small Python projects.", instructions: "Be good.", model: "Qwen/Qwen3-8B", rubric, evaluator: defineEvaluator(rubric, { good: check("good", (run) => run.output.includes("good")) }),});`;
const childRubric = (ids: string[]) => `import { defineRubric } from "../../../src/index.ts";export const rubric = defineRubric({${ids.map((id) => ` ${id}: { description: "Mentions ${id}.", weight: 2 },`).join("\n")}});`;
const CHILD = `import { check, defineChild, defineEvaluator, instructionsFrom, modelFrom } from "../../../src/index.ts";import toy from "../machine.ts";import { rubric } from "./rubric.ts";export default defineChild(toy, { name: "toy/empty", intent: "Handle empty inputs.", instructions: instructionsFrom(import.meta.url), model: modelFrom(import.meta.url, "thinkingmachines/Inkling-Small"), rubric, evaluator: defineEvaluator(rubric, { mentions_empty: check("mentions-empty", (run) => ({ score: run.output.includes("empty") ? 1 : 0, reason: "looked" })) }),});`;
const design = (rubricIds: string[]) => `<slug>empty</slug>\n<rubric_ts>${childRubric(rubricIds)}</rubric_ts>\n<machine_ts>${CHILD}</machine_ts>\n<instructions_md>Check empty inputs first.</instructions_md>\n<topics>\nempty lists\nmissing keys\n</topics>`;
test("specialize: the author's invalid child goes back to it, then the child and its checked cases are written", async () => { const dir = await tempDir(); await symlink(path.join(repo, "src"), path.join(dir, "src"), "dir"); await mkdir(path.join(dir, "machines", "toy"), { recursive: true }); await writeFile(path.join(dir, "machines", "toy", "machine.ts"), PARENT); await writeFile(path.join(dir, "machines", "toy", "cases.json"), JSON.stringify([{ input: "fix it", files: { "stats.py": BUGGY, "test_stats.py": TEST } }]));
const prompts: HarnessRequest[] = []; let batch = 0; const harness = scriptedHarness((req) => { prompts.push(req); if (req.input.startsWith("Write a child machine")) return design(["good"]); // redefines the inherited criterion if (req.input.startsWith("Your child machine failed validation")) return design(["mentions_empty"]); batch++; // Each batch: two good cases and one whose test already passes (discarded). return [ caseBlock(`mean of an empty list crashes (${batch}a)`, BUGGY, FIXED), caseBlock(`mean([]) should be 0 (${batch}b)`, BUGGY, FIXED), caseBlock(`nothing wrong (${batch}c)`, FIXED, FIXED), ].join("\n"); }); const toy = await machine.load("toy", { workspace: dir, harness, store: diskStore(path.join(dir, ".machine")), sandboxes: sandboxes(), events: { emit() {} }, }); const result = await toy.specialize({ focus: "empty inputs", model: "thinkingmachines/Inkling-Small", authorModel: "author-model", budget: 1, trainCases: 4, holdoutCases: 2, concurrency: 1 });
assert.equal(result.name, "toy/empty"); assert.equal(result.attempts, 2, "one retry after the failed validation"); const retry = prompts.find((p) => p.input.startsWith("Your child machine failed validation"))!; assert.match(retry.input, /does not load|redefines/); assert.ok(retry.history.length > 0, "the retry sees the author's first attempt"); assert.deepEqual(result.inherited.map((c) => c.id), ["good"]); assert.deepEqual(result.added.map((c) => [c.id, c.grader]), [["mentions_empty", "check:mentions-empty"]]); assert.deepEqual(result.topics, ["empty lists", "missing keys"]); assert.equal(result.cases, 4); assert.equal(result.holdout, 2); assert.ok(result.discarded.some((d) => /already passes/.test(d.reason)));
const childDir = path.join(dir, "machines", "toy", "empty"); const cases = JSON.parse(await readFile(path.join(childDir, "cases.json"), "utf8")); const holdout = JSON.parse(await readFile(path.join(childDir, "holdout.json"), "utf8")); assert.equal(cases.length + holdout.length, 6); assert.ok(cases.every((c: { solution?: unknown; files: Record<string, string> }) => !c.solution && c.files["stats.py"] === BUGGY), "the reference fix is not in the cases"); assert.ok(!holdout.some((h: { input: string }) => cases.some((c: { input: string }) => c.input === h.input)));
const child = await machine.load("toy/empty", { workspace: dir, harness, store: diskStore(path.join(dir, ".machine")), events: { emit() {} } }); assert.equal(child.model, "thinkingmachines/Inkling-Small"); assert.deepEqual(Object.keys(child.definition.rubric).sort(), ["good", "mentions_empty"]); assert.match(child.definition.instructions, /Be good\.\n\nCheck empty inputs first\./);
// Rewriting the child keeps its cases. const again = await toy.specialize({ focus: "empty inputs", model: "thinkingmachines/Inkling-Small", authorModel: "author-model", budget: 1, slug: "empty", reuseCases: true }); assert.equal(again.cases, 4); assert.equal(again.holdout, 2); assert.deepEqual(JSON.parse(await readFile(path.join(childDir, "cases.json"), "utf8")), cases);
const text = formatSpecialization(result); assert.match(text, /inherits\s+good/); assert.match(text, /adds\s+mentions_empty\s+check/); assert.ok(text.split("\n").length <= 14);});
test("validation rejects a check that grades only the files the run started with", async () => { const dir = await tempDir(); await symlink(path.join(repo, "src"), path.join(dir, "src"), "dir"); await mkdir(path.join(dir, "machines", "toy"), { recursive: true }); await writeFile(path.join(dir, "machines", "toy", "machine.ts"), PARENT); const fixtureOnly = CHILD.replace('run.output.includes("empty")', 'Object.keys(run.files ?? {}).length > 0'); let calls = 0; const harness = scriptedHarness(() => (calls++ === 0 ? design(["mentions_empty"]).replace(CHILD, fixtureOnly) : "<case><input>x</input></case><case><input>y</input></case>")); const toy = await machine.load("toy", { workspace: dir, harness, store: diskStore(path.join(dir, ".machine")), events: { emit() {} } }); await assert.rejects(toy.specialize({ focus: "x", authorModel: "a", budget: 1, retries: 0 }), /reads run\.files/);});
test("specialize gives up after its retries when the child never validates", async () => { const dir = await tempDir(); await symlink(path.join(repo, "src"), path.join(dir, "src"), "dir"); await mkdir(path.join(dir, "machines", "toy"), { recursive: true }); await writeFile(path.join(dir, "machines", "toy", "machine.ts"), PARENT); let calls = 0; const harness = scriptedHarness(() => (calls++, design(["good"]))); const toy = await machine.load("toy", { workspace: dir, harness, store: diskStore(path.join(dir, ".machine")), events: { emit() {} } }); await assert.rejects(toy.specialize({ focus: "x", authorModel: "a", budget: 1, retries: 1 }), /failed validation after 2 attempts/); assert.equal(calls, 2); await assert.rejects(readFile(path.join(dir, "machines", "toy", "empty", "machine.ts")), "a child that never validated is removed");});
test("specialize asks the author once to reconsider a check that scores every sample run the same, then accepts", async () => { const dir = await tempDir(); await symlink(path.join(repo, "src"), path.join(dir, "src"), "dir"); await mkdir(path.join(dir, "machines", "toy"), { recursive: true }); await writeFile(path.join(dir, "machines", "toy", "machine.ts"), PARENT); await writeFile(path.join(dir, "machines", "toy", "cases.json"), JSON.stringify(["a", "b", "c"])); const inputs: string[] = []; let n = 0; const harness = scriptedHarness((req) => { inputs.push(req.input); if (req.input.startsWith("Write ") && req.input.includes("cases for")) return `<case><input>case ${++n}</input></case><case><input>case ${++n}</input></case>`; return design(["mentions_empty"]); // the check scores "Done." runs 0 every time }); const toy = await machine.load("toy", { workspace: dir, harness, store: diskStore(path.join(dir, ".machine")), events: { emit() {} } }); const result = await toy.specialize({ focus: "empty inputs", model: "thinkingmachines/Inkling-Small", authorModel: "a", budget: 1, trainCases: 3, holdoutCases: 1, concurrency: 1 }); assert.equal(result.attempts, 2); assert.ok(inputs.some((i) => i.startsWith("Your child machine is valid, but check these") && /mentions_empty gives every one of 3 sample runs/.test(i))); assert.equal(result.testedFixtures, false); assert.equal(result.cases + result.holdout, 4);});