Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
41 kB · 1138 lines
TypeScript
at main
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139import { appendFileSync, chmodSync, mkdtempSync, readdirSync, rmSync, statSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import type { AgentMessage } from "@earendil-works/pi-agent-core";import type * as PiAi from "@earendil-works/pi-ai";import type { AssistantMessage, Model } from "@earendil-works/pi-ai";import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";import { appendGlobalRefinement, applyRefinementProposal, formatHarnessStateForPrompt, getGlobalHarnessStateDir, getHarnessStatePath, getLocalHarnessStateDir, getRefinementHistory, getRefinementHistoryPath, type HarnessEntry, type HarnessState, harnessQueryTerms, inferRefinementResultScope, isAddressableHarnessId, loadGlobalRefinementHistory, loadHarnessState, mergeHarnessStates, mergeRefinementHistory, planRefinement, type RefinementAction, type RefinementKind, type RefinementProposal, type RefinementResult, recordHarnessViolation, refineHarness, saveHarnessState, scoreHarnessEntryForQuery, selectOverviewEntries,} from "../src/core/refinement/index.js";import type { CustomEntry } from "../src/core/session-manager.js";
const { completeSimpleMock } = vi.hoisted(() => ({ completeSimpleMock: vi.fn(),}));
vi.mock("@earendil-works/pi-ai", async (importOriginal) => { const actual = await importOriginal<typeof PiAi>(); return { ...actual, completeSimple: completeSimpleMock, };});
let tempDir: string | undefined;
beforeEach(() => { completeSimpleMock.mockReset();});
afterEach(() => { if (tempDir) { rmSync(tempDir, { recursive: true, force: true }); tempDir = undefined; }});
function makeTempDir(): string { tempDir = mkdtempSync(join(tmpdir(), "prime-agent-refinement-test-")); return tempDir;}
const kinds = ["prompt", "memory", "skill", "subagent"] as const satisfies readonly RefinementKind[];const skillReference = { type: "python", import: "agent_skills.example", callable: "run", call_pattern: "await run(...)",};const skillContract = { reference: skillReference, arguments: { input: { type: "string", required: true, description: "Task input" } },};
function proposal(summary: string, edits: RefinementProposal["edits"]): RefinementProposal { return { summary, rationale: `${summary} rationale`, expectedOutcome: `${summary} outcome`, edits, };}
function createRefineModel(reasoning: boolean): Model<"openai-completions"> { return { id: "openai/gpt-5.5", name: "GPT 5.5", api: "openai-completions", provider: "prime-inference", baseUrl: "https://inference.primeintellect.ai/v1", reasoning, input: ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 200000, maxTokens: 8192, };}
function assistantText(text: string): AssistantMessage { return { role: "assistant", content: [{ type: "text", text }], api: "openai-completions", provider: "prime-inference", model: "openai/gpt-5.5", usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 2, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop", timestamp: Date.now(), };}
function seedEntry(state: HarnessState, kind: RefinementKind, id = `${kind}_entry`): void { applyRefinementProposal( state, proposal(`seed ${kind}`, [ { action: "create", kind, id, title: `${kind} title`, content: `${kind} content`, path: `${kind}/path`, ...(kind === "skill" ? skillContract : {}), metadata: { seeded: true }, }, ]), { id: `seed_${kind}_${id}` }, );}
describe("harness refinement", () => { it("rejects an edit when the target entry changed after planning", () => { const harnessStateDir = makeTempDir(); const baselineState = loadHarnessState(harnessStateDir); seedEntry(baselineState, "memory"); saveHarnessState(harnessStateDir, baselineState); const currentState = loadHarnessState(harnessStateDir); currentState.entries.memory.memory_entry.content = "concurrent kernel content"; currentState.entries.memory.memory_entry.version++;
const result = applyRefinementProposal( currentState, proposal("Update memory", [ { action: "update", kind: "memory", id: "memory_entry", title: "Planned title", content: "stale planned content", }, ]), { id: "refine_conflict", baselineState }, );
expect(result.appliedEdits).toMatchObject([ { applied: false, error: "entry changed during refinement planning" }, ]); expect(currentState.entries.memory.memory_entry.content).toBe("concurrent kernel content"); });
it("allows sequential edits to the same entry after the baseline matches once", () => { const state = loadHarnessState(makeTempDir(), "local"); seedEntry(state, "memory"); const baselineState = structuredClone(state);
const result = applyRefinementProposal( state, proposal("Update memory twice", [ { action: "update", kind: "memory", id: "memory_entry", title: "First", content: "first" }, { action: "update", kind: "memory", id: "memory_entry", title: "Second", content: "second" }, ]), { id: "refine_same_entry", baselineState, scope: "local" }, );
expect(result.appliedEdits.map((edit) => edit.applied)).toEqual([true, true]); expect(state.entries.memory.memory_entry.content).toBe("second"); expect(state.entries.memory.memory_entry.version).toBe(3); });
it("atomically replaces harness state without leaving temporary files", () => { const harnessStateDir = makeTempDir(); const state = loadHarnessState(harnessStateDir); seedEntry(state, "memory");
const statePath = saveHarnessState(harnessStateDir, state);
expect(loadHarnessState(harnessStateDir).entries.memory.memory_entry).toBeDefined(); expect(readdirSync(harnessStateDir)).toEqual([statePath.split("/").at(-1)]); chmodSync(statePath, 0o600); saveHarnessState(harnessStateDir, state); expect(statSync(statePath).mode & 0o777).toBe(0o600); });
it.each(kinds)("applies the create/update/delete lifecycle for %s entries", (kind) => { const state = loadHarnessState(makeTempDir()); const id = `${kind}_entry`; const skillFields = kind === "skill" ? skillContract : {}; const apply = (edits: RefinementProposal["edits"], refinementId: string) => applyRefinementProposal(state, proposal(`${refinementId} ${kind}`, edits), { id: refinementId });
const created = apply( [ { action: "create", kind, id, title: `${kind} title`, content: `${kind} content`, path: `${kind}/created`, ...skillFields, metadata: { kind }, }, ], `refine_create_${kind}`, );
expect(created.appliedEdits[0]).toMatchObject({ applied: true, after: { version: 1 } }); expect(created.appliedEdits[0].before).toBeUndefined(); expect(state.entries[kind][id]).toMatchObject({ id, kind, title: `${kind} title`, content: `${kind} content`, path: `${kind}/created`, metadata: { kind }, source: "refine", version: 1, }); expect(state.refinements.at(-1)?.changes).toEqual([`create ${kind}:${id}`]);
const updated = apply( [ { action: "update", kind, id, title: `${kind} title updated`, content: `${kind} content updated`, path: `${kind}/updated`, ...skillFields, metadata: { updated: kind }, }, ], `refine_update_${kind}`, );
expect(updated.appliedEdits[0]).toMatchObject({ applied: true, before: { version: 1 }, after: { version: 2 } }); expect(state.entries[kind][id]).toMatchObject({ title: `${kind} title updated`, content: `${kind} content updated`, path: `${kind}/updated`, metadata: { updated: kind }, version: 2, });
const deleted = apply([{ action: "delete", kind, id }], `refine_delete_${kind}`);
expect(deleted.appliedEdits[0]).toMatchObject({ applied: true, before: { version: 2 } }); expect(deleted.appliedEdits[0].after).toBeUndefined(); expect(state.entries[kind][id]).toBeUndefined(); expect(state.refinements.at(-1)?.changes).toEqual([`delete ${kind}:${id}`]); });
it("creates ids from titles and uses default path and metadata when omitted", () => { const state = loadHarnessState(makeTempDir());
const result = applyRefinementProposal( state, proposal("Create with generated id", [ { action: "create", kind: "skill", title: "Native Check!", content: "Run checks.", ...skillContract }, ]), { id: "refine_generated_id" }, );
expect(result.appliedEdits[0]).toMatchObject({ applied: true, id: "native_check", after: { id: "native_check", path: "general", reference: skillReference, metadata: {}, version: 1 }, }); });
// Every rejected edit must leave harness state and the refinement log untouched. type InvalidCase = { label: string; edit: RefinementProposal["edits"][number]; error: string; seed?: RefinementKind; }; const skillFieldsFor = (kind: RefinementKind) => (kind === "skill" ? skillContract : {}); it.each<InvalidCase>([ ...kinds.map( (kind): InvalidCase => ({ label: `duplicate create for ${kind}`, seed: kind, edit: { action: "create", kind, id: `${kind}_entry`, title: "t", content: "c", ...skillFieldsFor(kind) }, error: "entry already exists", }), ), ...kinds.map( (kind): InvalidCase => ({ label: `update of a missing ${kind}`, edit: { action: "update", kind, id: `${kind}_missing`, title: "t", content: "c", ...skillFieldsFor(kind) }, error: "entry not found", }), ), ...kinds.map( (kind): InvalidCase => ({ label: `delete of a missing ${kind}`, edit: { action: "delete", kind, id: `${kind}_missing` }, error: "entry not found", }), ), { label: "create without content", edit: { action: "create", kind: "memory", id: "missing_fields", title: "t" }, error: "create requires title and content", }, { label: "update without content", seed: "memory", edit: { action: "update", kind: "memory", id: "memory_entry", title: "t" }, error: "update requires title and content", }, { label: "update without an id", edit: { action: "update", kind: "skill", title: "t", content: "c" }, error: "update requires id", }, { label: "delete without an id", edit: { action: "delete", kind: "skill" }, error: "delete requires id" }, { label: "an unsupported action", edit: { action: "rename" as RefinementAction, kind: "memory", id: "bad_action", title: "t", content: "c" }, error: "unsupported action rename", }, { label: "an unsupported kind", edit: { action: "create", kind: "tool" as RefinementKind, id: "bad_kind", title: "t", content: "c" }, error: "unsupported kind tool", }, { label: "a skill without an argument contract", edit: { action: "create", kind: "skill", id: "argumentless", title: "t", content: "c", reference: skillReference, }, error: "create skill requires arguments", }, { label: "a skill without a Python reference", edit: { action: "create", kind: "skill", id: "unbacked", title: "t", content: "c", arguments: {} }, error: "create skill requires python reference", }, { label: "a skill backed by a non-Python reference", edit: { action: "create", kind: "skill", id: "shell_skill", title: "t", content: "c", reference: { type: "shell", command: "edit" }, arguments: {}, }, error: "create skill reference.type must be python", }, { label: "an update of the base system prompt", edit: { action: "update", kind: "prompt", id: "base_system_prompt", title: "t", content: "c" }, error: "base system prompt", }, { label: "a create whose title derives the base system prompt id", edit: { action: "create", kind: "prompt", title: "Base System Prompt", content: "c" }, error: "base system prompt", }, ])("rejects $label without mutating state", ({ seed, edit, error }) => { const state = loadHarnessState(makeTempDir()); if (seed) seedEntry(state, seed); const before = structuredClone(state.entries);
const result = applyRefinementProposal(state, proposal("Invalid edit", [edit]), { id: "refine_invalid" });
expect(result.appliedEdits).toHaveLength(1); expect(result.appliedEdits[0].applied).toBe(false); expect(result.appliedEdits[0].error).toContain(error); expect(state.entries).toEqual(before); expect(state.refinements.at(-1)?.changes).toEqual([]); }); it("keeps global and local harness state separate when merging and persisting", () => { const root = makeTempDir(); const globalState = loadHarnessState(join(root, "global"), "global"); const localState = loadHarnessState(join(root, "local"), "local"); const note = (content: string): RefinementProposal => proposal(content, [{ action: "create", kind: "memory", id: "shared", title: "Shared", content }]); applyRefinementProposal(globalState, note("Global content."), { id: "refine_global", scope: "global" }); applyRefinementProposal(localState, note("Local content."), { id: "refine_local", scope: "local" });
const merged = mergeHarnessStates(globalState, localState);
// A colliding local id is namespaced instead of shadowed. expect(merged.entries.memory.shared).toMatchObject({ content: "Global content.", scope: "global" }); expect(merged.entries.memory["local:shared"]).toMatchObject({ id: "shared", content: "Local content.", scope: "local", }); expect(globalState.entries.memory.shared.scope).toBe("global");
// Scope stored inside a shared file wins over the file's own default scope. applyRefinementProposal( globalState, proposal("Session-local note in shared file", [ { action: "create", kind: "memory", id: "session_note", title: "Session note", content: "Local store." }, ]), { id: "refine_local_in_global_file", scope: "local" }, ); expect(mergeHarnessStates(globalState).entries.memory.session_note.scope).toBe("local"); });
it("resolves the harness state directories and persists state in the selected one", () => { const dir = makeTempDir(); expect(getGlobalHarnessStateDir(dir)).toBe(join(dir, "harness")); expect(getHarnessStatePath(getGlobalHarnessStateDir(dir))).toBe(join(dir, "harness", "harness_state.json")); expect(getLocalHarnessStateDir(dir)).toBe(join(dir, "harness")); expect(getLocalHarnessStateDir(undefined)).toBeUndefined();
const state = loadHarnessState(dir, "local"); applyRefinementProposal( state, proposal("Add prompt note", [ { action: "create", kind: "prompt", id: "focused_edits", title: "Focused edits", content: "Prefer small harness edits.", }, ]), { id: "refine_1" }, );
const statePath = saveHarnessState(dir, state); const reloaded = loadHarnessState(dir, "local");
expect(statePath).toBe(getHarnessStatePath(dir)); expect(reloaded.entries.prompt.focused_edits).toMatchObject({ content: "Prefer small harness edits.", scope: "local", }); expect(reloaded.refinements[0]).toMatchObject({ id: "refine_1", trigger: "Add prompt note", changes: ["create prompt:focused_edits"], }); });
it.each(["not json at all", "null", "[]", '"a string"', "123"])( "loads empty harness state from a corrupt or non-object file (%s)", (payload) => { const dir = makeTempDir(); writeFileSync(getHarnessStatePath(dir), payload, "utf8");
const state = loadHarnessState(dir);
expect(state.entries).toEqual({ prompt: {}, memory: {}, skill: {}, subagent: {} }); expect(state.refinements).toEqual([]); applyRefinementProposal( state, proposal("Recover", [ { action: "create", kind: "memory", id: "recovered", title: "Recovered", content: "ok" }, ]), { id: "refine_recover" }, ); saveHarnessState(dir, state); expect(loadHarnessState(dir).entries.memory.recovered.content).toBe("ok"); }, );
it("extracts well-formed refinement history from custom session entries", () => { const result: RefinementResult = { id: "refine_1", summary: "Add skill", rationale: "Repeated failure.", expectedOutcome: "Better validation.", appliedEdits: [], harnessStatePath: "/tmp/harness_state.json", }; const entry = (customType: string, data: object, id: string): CustomEntry => ({ type: "custom", customType, data, id, parentId: null, timestamp: new Date().toISOString(), });
expect( getRefinementHistory([ entry("other", {}, "custom_1"), entry("prime-agent.refinement", result, "custom_2"), entry("prime-agent.refinement", { id: "malformed" }, "custom_malformed"), ]), ).toEqual([result]); });
it("requests a JSON refinement within the model's own output budget and applies it", async () => { const state = loadHarnessState(makeTempDir()); const edit = { action: "create", kind: "memory", id: "native_validation", title: "Native validation", content: "Run validation through the target project environment.", }; completeSimpleMock.mockResolvedValueOnce( assistantText(JSON.stringify({ summary: "s", rationale: "r", expectedOutcome: "o", edits: [edit] })), );
const result = await refineHarness( [{ role: "user", content: "Use native validation.", timestamp: Date.now() } satisfies AgentMessage], state, [], createRefineModel(true), "api-key", {}, { "x-test-header": "1" }, undefined, "xhigh", );
// Budget is derived from the model (8192) rather than a fixed literal. expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({ reasoning: "low", maxTokens: 8192, apiKey: "api-key", headers: { "x-test-header": "1" }, }); expect(result.appliedEdits[0]).toMatchObject({ kind: "memory", id: "native_validation", applied: true }); expect(state.entries.memory.native_validation.content).toBe(edit.content); });
it("caps the refinement output budget by the policy ceiling for wide models", async () => { const state = loadHarnessState(makeTempDir()); completeSimpleMock.mockResolvedValueOnce( assistantText(JSON.stringify({ summary: "s", rationale: "r", expectedOutcome: "o", edits: [] })), );
await refineHarness([], state, [], { ...createRefineModel(false), maxTokens: 128_000 }, "api-key", {});
expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({ maxTokens: 32_000 }); });
// A cut-off reply must be reported as truncation, and a complete-but-invalid // reply as a formatting failure: swapping the two sends callers the wrong fix. it.each<{ label: string; reply: string; stopReason?: "length"; error: RegExp }>([ { label: "an exhausted output budget", reply: `{ "summary": "s", "edits": [ { "action": "create", "kind": "memory", "id": "a", "title": "t", "content": "first" }, { "action": "create", "kind": "memory", "id": "b", "title": "t2", "content": "second`, stopReason: "length", error: /output budget was exhausted/, }, { label: "a truncated reply that never reports a length stop reason", reply: `{ "summary": "s", "edits": [ { "action": "create", "kind": "memory", "id": "a", "title": "t", "content": "first" }, { "action": "create", "kind": "memory", "id": "b", "title": "t2", "content": "second`, error: /stopped before completing its JSON object/, }, { // Ends on "}" so it takes the startsWith/endsWith fast path rather than // the brace-slicing fallback, but is still an incomplete object. label: "a JSON-only reply cut after a nested closing brace", reply: `{ "summary": "s", "edits": [ { "action": "create", "kind": "memory", "id": "a", "title": "t", "content": "first" }`, error: /stopped before completing its JSON object/, }, { label: "a complete but malformed reply", reply: 'Here is the result: {"edits": [oops]}', error: /did not return valid JSON/, }, ])("reports $label", async ({ reply, stopReason, error }) => { const state = loadHarnessState(makeTempDir()); const message = assistantText(reply); completeSimpleMock.mockResolvedValueOnce(stopReason ? { ...message, stopReason } : message);
await expect(refineHarness([], state, [], createRefineModel(false), "api-key", {})).rejects.toThrow(error); });
it("rolls back created, updated, and deleted entries from refinement history", async () => { const state = loadHarnessState(makeTempDir()); seedEntry(state, "memory", "kept_memory"); seedEntry(state, "skill", "deleted_skill");
const target = applyRefinementProposal( state, proposal("Target refinement", [ { action: "create", kind: "prompt", id: "created_prompt", title: "Created", content: "Created content" }, { action: "update", kind: "memory", id: "kept_memory", title: "Updated memory", content: "Updated memory content", path: "updated/path", metadata: { updated: true }, }, { action: "delete", kind: "skill", id: "deleted_skill" }, ]), { id: "refine_target" }, );
const rollback = await refineHarness([], state, [target], {} as never, "api-key", { rollbackId: "refine_target", });
expect(rollback).toMatchObject({ rollbackOf: "refine_target", scope: "local" }); // Rollback undoes the edits in reverse order and restores the pre-edit snapshots. expect(rollback.appliedEdits.map((edit) => `${edit.action} ${edit.kind}:${edit.id}`)).toEqual([ "create skill:deleted_skill", "update memory:kept_memory", "delete prompt:created_prompt", ]); expect(state.entries.prompt.created_prompt).toBeUndefined(); expect(state.entries.memory.kept_memory).toMatchObject({ content: "memory content", path: "memory/path", metadata: { seeded: true }, version: 3, }); expect(state.entries.skill.deleted_skill).toMatchObject({ content: "skill content", reference: skillReference, arguments: skillContract.arguments, version: 1, }); expect(state.refinements.at(-1)?.trigger).toBe("Rollback refinement refine_target"); });
it("throws when rollback target is missing", async () => { const state = loadHarnessState(makeTempDir());
await expect( refineHarness([], state, [], {} as never, "api-key", { rollbackId: "missing_refinement" }), ).rejects.toThrow("Refinement missing_refinement not found"); });});
describe("global refinement history", () => { function sampleResult(id: string, overrides: Partial<RefinementResult> = {}): RefinementResult { return { id, summary: `${id} summary`, rationale: `${id} rationale`, expectedOutcome: `${id} outcome`, appliedEdits: [], harnessStatePath: "/tmp/harness_state.json", ...overrides, }; }
it("appends, reloads, and skips malformed history lines", () => { const dir = makeTempDir(); expect(loadGlobalRefinementHistory(dir)).toEqual([]);
const first = sampleResult("refine_1"); const second = sampleResult("refine_2"); const historyPath = appendGlobalRefinement(dir, first); appendGlobalRefinement(dir, second); appendFileSync(historyPath, "not json\n", "utf8"); appendFileSync(historyPath, `${JSON.stringify({ id: "x" })}\n`, "utf8");
expect(historyPath).toBe(getRefinementHistoryPath(dir)); expect(loadGlobalRefinementHistory(dir)).toEqual([ { ...first, scope: "global" }, { ...second, scope: "global" }, ]); });
// History written before result.scope existed must still resolve to global. it("defaults legacy global history entries to global scope", () => { const dir = makeTempDir(); const legacyEdit = { action: "create" as const, kind: "memory" as const, id: "legacy_global_memory", title: "Legacy global memory", content: "created globally", applied: true, after: { id: "legacy_global_memory", kind: "memory" as const, title: "Legacy global memory", content: "created globally", path: "general", scope: "global" as const, reference: {}, arguments: {}, metadata: {}, source: "refine", created_at: "2026-01-01T00:00:00.000Z", updated_at: "2026-01-01T00:00:00.000Z", version: 1, }, }; const scopeless = sampleResult("refine_legacy", { scope: undefined }); const editScoped = sampleResult("refine_legacy_edits", { scope: undefined, appliedEdits: [legacyEdit] }); for (const legacy of [scopeless, editScoped]) { appendFileSync(getRefinementHistoryPath(dir), `${JSON.stringify(legacy)}\n`, "utf8"); }
// Scope is recovered from the edit snapshots and written back onto the loaded results. expect(loadGlobalRefinementHistory(dir).map((entry) => entry.scope)).toEqual(["global", "global"]); expect(inferRefinementResultScope(editScoped)).toBe("global"); });
it("merges global and session history, preferring session entries but keeping global scope", () => { const globalShared = sampleResult("refine_shared", { scope: "global", summary: "global version" }); const globalOnly = sampleResult("refine_global_only", { scope: "global" }); const sessionShared = sampleResult("refine_shared", { scope: undefined, summary: "session version" }); const sessionOnly = sampleResult("refine_session_only");
const merged = mergeRefinementHistory([globalShared, globalOnly], [sessionShared, sessionOnly]);
expect(merged).toHaveLength(3); expect(merged.find((item) => item.id === "refine_shared")).toMatchObject({ summary: "session version", scope: "global", }); expect(merged.map((item) => item.id)).toEqual( expect.arrayContaining(["refine_shared", "refine_global_only", "refine_session_only"]), ); });
it("plans a proposal without mutating harness state", async () => { const state = loadHarnessState(makeTempDir()); const edit = { action: "create", kind: "memory", id: "planned_memory", title: "Planned memory", content: "Created only when applied.", }; completeSimpleMock.mockResolvedValueOnce( assistantText(JSON.stringify({ summary: "s", rationale: "r", expectedOutcome: "o", edits: [edit] })), );
const plan = await planRefinement( [{ role: "user", content: "remember this", timestamp: Date.now() } satisfies AgentMessage], state, [], createRefineModel(false), "api-key", {}, );
// planRefinement must not touch state: the host re-reads the file before applying, // so applying must be the only thing that mutates state. expect(plan.proposal.edits).toHaveLength(1); expect(plan.id).toMatch(/^refine_/); const request = completeSimpleMock.mock.calls[0][1]; const userPrompt = request.messages[0].content[0].text; expect(userPrompt).toContain("Requested refinement scope: local"); expect(userPrompt).toContain("Global entries in the overview are read-only context"); expect(request.systemPrompt).toContain('handle = await rlm.spawn("sub-task", name="worker", model="<model>")'); expect(request.systemPrompt).toContain("never the child's answer"); expect(request.systemPrompt).toContain('receiver_role="parent"'); expect(request.systemPrompt).toContain("await rlm.list_subagents()"); expect(request.systemPrompt).toContain('receiver_role="child"'); expect(request.systemPrompt).not.toContain("asyncio.create_task(rlm"); expect(request.systemPrompt).not.toContain("asyncio.gather(rlm"); expect(state.entries.memory.planned_memory).toBeUndefined(); expect(state.refinements).toHaveLength(0);
const result = applyRefinementProposal(state, plan.proposal, { id: plan.id }); expect(result.appliedEdits[0]).toMatchObject({ id: "planned_memory", applied: true }); expect(state.entries.memory.planned_memory).toBeDefined(); });
it.each<{ label: string; scope: "local" | "global" }>([ { label: "local", scope: "local" }, { label: "global", scope: "global" }, ])("plans a $label rollback against the recorded scope without mutating state", async ({ scope }) => { const dir = makeTempDir(); const state = loadHarnessState(dir, scope); const target = applyRefinementProposal( state, proposal("Target", [ { action: "create", kind: "memory", id: "rollback_me", title: "Rollback me", content: "content" }, ]), { id: "refine_rollback_target", scope }, ); expect(target.scope).toBe(scope);
const plan = await planRefinement([], state, [target], {} as never, "api-key", { rollbackId: "refine_rollback_target", });
expect(plan.rollbackOf).toBe("refine_rollback_target"); expect(plan.rollbackScope).toBe(scope); expect(state.entries.memory.rollback_me).toBeDefined();
const rollback = applyRefinementProposal(state, plan.proposal, { id: plan.id, rollbackOf: plan.rollbackOf, scope: plan.rollbackScope, }); expect(rollback.scope).toBe(scope); expect(state.entries.memory.rollback_me).toBeUndefined(); });
it("rolls back a refinement recorded in a different session via global history", async () => { const dir = makeTempDir(); const sessionAState = loadHarnessState(dir); const applied = applyRefinementProposal( sessionAState, proposal("Session A refinement", [ { action: "create", kind: "memory", id: "session_a_memory", title: "Session A memory", content: "Created in session A.", }, ]), { id: "refine_session_a" }, ); applied.harnessStatePath = saveHarnessState(dir, sessionAState); appendGlobalRefinement(dir, applied);
// A fresh session loads the global state and the global history (its own session // has no record of refine_session_a) and can still roll it back. const sessionBState = loadHarnessState(dir); expect(sessionBState.entries.memory.session_a_memory).toBeDefined();
const globalHistory = mergeRefinementHistory(loadGlobalRefinementHistory(dir), getRefinementHistory([])); const rollback = await refineHarness([], sessionBState, globalHistory, {} as never, "api-key", { rollbackId: "refine_session_a", });
expect(rollback).toMatchObject({ rollbackOf: "refine_session_a", scope: "local" }); expect(sessionBState.entries.memory.session_a_memory).toBeUndefined(); });});
describe("harness digest relevance ranking", () => { function makeEntry(id: string, title: string, content: string, updatedAt: string): HarnessEntry { return { id, kind: "memory", title, content, path: "memory", scope: "global", reference: {}, arguments: {}, metadata: {}, source: "test", created_at: updatedAt, updated_at: updatedAt, version: 1, }; }
function rank(entries: Record<string, HarnessEntry>, query: string | Map<string, number>): string { const state = loadHarnessState(join(makeTempDir(), "h"), "local"); Object.assign(state.entries.memory, entries); return formatHarnessStateForPrompt(state, { maxEntriesPerKind: 1, // These cases are about which entry the ranking puts in the detailed tier. // The shipped default also lists the rest as addressable stubs, which the // "without hiding entries behind a black hole" cases below cover. maxListedEntries: 0, queryTerms: typeof query === "string" ? new Map(harnessQueryTerms(query).map((term) => [term, 1])) : query, }); }
it("scores weighted term overlap with field coverage", () => { const entry = makeEntry( "repo", "Repository facts", "The checkout lives at ~/repo with worktrees.", "2026-09-01T00:00:00.000Z", ); const terms = new Map([ ["repository", 2], ["worktree", 1], ["absent-term", 5], ]); // "repository" matches the title only -> 2 * 1 = 2. // "worktree" matches content only (1 field) -> 1 * 1 = 1. Total 3. expect(scoreHarnessEntryForQuery(entry, terms)).toBe(3); expect(scoreHarnessEntryForQuery(entry, new Map())).toBe(0); });
it.each<{ label: string; query: string; winner: string; loser: string; entries: Record<string, HarnessEntry> }>([ { label: "relevance over alphabetical order", query: "worktree", winner: "zzz", loser: "aaa", entries: { aaa: makeEntry("aaa", "Alphabetical first", "Unrelated content about tea.", "2026-08-01T00:00:00.000Z"), zzz: makeEntry("zzz", "Zebra note", "The worktree workflow matters.", "2026-08-02T00:00:00.000Z"), }, }, { label: "recency as the tie-break for equal scores", query: "worktree", winner: "newer", loser: "older", entries: { older: makeEntry("older", "Worktree policy", "Same worktree signal.", "2026-08-01T00:00:00.000Z"), newer: makeEntry("newer", "Worktree policy 2", "Same worktree signal.", "2026-09-01T00:00:00.000Z"), }, }, { label: "whitespace-free CJK matches through bigram terms", query: "修复登录", winner: "login", loser: "tea", entries: { login: makeEntry("login", "Login fix", "登录故障排查记录。", "2026-08-01T00:00:00.000Z"), tea: makeEntry("tea", "Tea notes", "All about oolong brewing.", "2026-08-02T00:00:00.000Z"), }, }, { label: "real terms rather than incidental query punctuation", query: "worktree?", winner: "worktree", loser: "question", entries: { worktree: makeEntry("worktree", "Branch hygiene", "Use git worktrees.", "2026-08-01T00:00:00.000Z"), question: makeEntry("question", "Question", "Anything else left open?", "2026-08-02T00:00:00.000Z"), }, }, ])("ranks by $label", ({ query, winner, loser, entries }) => { const ranked = rank(entries, query);
expect(ranked).toContain(`[global:${winner}]`); expect(ranked).not.toContain(`[global:${loser}]`); });
it("falls back to alphabetical selection and tolerates non-string persisted fields", () => { const state = loadHarnessState(join(makeTempDir(), "h"), "local"); state.entries.memory.aaa = makeEntry("aaa", "Alphabetical first", "Tea.", "2026-08-01T00:00:00.000Z"); state.entries.memory.zzz = makeEntry("zzz", "Zebra note", "Worktree workflow.", "2026-08-02T00:00:00.000Z"); (state.entries.memory.zzz as unknown as { title: null }).title = null;
expect(formatHarnessStateForPrompt(state, { maxEntriesPerKind: 1 })).toContain("[global:aaa]"); expect(() => formatHarnessStateForPrompt(state, { queryTerms: new Map([["worktree", 1]]) })).not.toThrow(); });
it("records violation counts and surfaces them in the harness digest", () => { const state = loadHarnessState(join(makeTempDir(), "h"), "local"); state.entries.prompt.note = makeEntry( "note", "Negative control check", "Verify before claiming pass.", "2026-08-01T00:00:00.000Z", );
const initialOverview = formatHarnessStateForPrompt(state); expect(initialOverview).toContain("[global:note] Negative control check (memory, v1):"); expect(initialOverview).not.toContain("v1,");
recordHarnessViolation(state, "note", 1); const oneViolationOverview = formatHarnessStateForPrompt(state); expect(oneViolationOverview).toContain("[global:note] Negative control check (memory, v1, 1 violation):");
recordHarnessViolation(state, "note", 1); const twoViolationsOverview = formatHarnessStateForPrompt(state); expect(twoViolationsOverview).toContain("[global:note] Negative control check (memory, v1, 2 violations):"); });
it.each<[string, string[]]>([ ["Worktree?", ["worktree"]], ["path/to/skill", ["path", "skill"]], ["harness_search", ["harness", "search"]], ["??? / . ,", []], // Short ASCII runs stay noise; other non-ASCII scripts stay whole. ["Fix the LOGIN bug", ["login"]], ["Привет мир", ["привет"]], // Combining marks stay in their run: mark-heavy scripts spell whole words. ["किताब notes", ["किताब", "notes"]], ["naïve approach", ["naïve", "approach"]], // CJK has no spaces between words: runs become overlapping bigrams, so // partial matches stay findable and single characters count. ["修复login", ["修复", "login"]], ["修复登录", ["修复", "复登", "登录"]], ["修复登录？", ["修复", "复登", "登录"]], ["東京会議 login", ["東京", "京会", "会議", "login"]], ["登", ["登"]], // Supplementary-plane ideographs count as CJK. ["𠀀", ["𠀀"]], ["𠀀𠀁𠀂", ["𠀀𠀁", "𠀁𠀂"]], ])("tokenizes %j into %j", (query, expected) => { expect(harnessQueryTerms(query)).toEqual(expected); });});
describe("selectOverviewEntries and two-tier overview", () => { const budget = { detailBudget: 100_000, maxContentLength: 180, reservedRecentSlots: 6 };
function makeEntry(overrides: Partial<HarnessEntry> & Pick<HarnessEntry, "id">): HarnessEntry { return { kind: "memory", title: overrides.id, content: `content for ${overrides.id}`, path: `memory/${overrides.id}.md`, scope: "global", reference: {}, arguments: {}, metadata: {}, source: "test", created_at: "2026-06-01T00:00:00.000Z", updated_at: "2026-06-01T00:00:00.000Z", version: 1, ...overrides, }; }
it("ranks pinned and high priority entries first", () => { const entries = [ makeEntry({ id: "recent", updated_at: "2026-06-10T00:00:00.000Z" }), makeEntry({ id: "pinned-note", updated_at: "2026-06-01T00:00:00.000Z", metadata: { pinned: true } }), makeEntry({ id: "priority-note", updated_at: "2026-06-05T00:00:00.000Z", metadata: { priority: 10 } }), ]; const result = selectOverviewEntries(entries, budget); expect(result.detailed.map((e) => e.id)).toEqual(["pinned-note", "priority-note", "recent"]); });
it("ranks recently updated entries before older entries regardless of path", () => { const entries = [ makeEntry({ id: "a-old", path: "a/first.md", updated_at: "2026-06-01T00:00:00.000Z" }), makeEntry({ id: "z-new", path: "z/last.md", updated_at: "2026-06-09T00:00:00.000Z" }), ]; const result = selectOverviewEntries(entries, budget); expect(result.detailed.map((e) => e.id)).toEqual(["z-new", "a-old"]); });
it("breaks recency ties using version desc then path tiebreak", () => { const entries = [ makeEntry({ id: "v1", path: "memory/b.md", version: 1 }), makeEntry({ id: "v2", path: "memory/b.md", version: 2 }), makeEntry({ id: "a-path", path: "memory/a.md", version: 1 }), ]; const result = selectOverviewEntries(entries, budget); expect(result.detailed.map((e) => e.id)).toEqual(["v2", "a-path", "v1"]); });
it("separates top entries into detailed and remainder into addressable stubs", () => { const entries = Array.from({ length: 8 }, (_, i) => makeEntry({ id: `entry_${i}`, updated_at: `2026-06-0${i + 1}T00:00:00.000Z` }), ); const result = selectOverviewEntries(entries, { reservedRecentSlots: 6 }); expect(result.detailed).toHaveLength(6); expect(result.listed).toHaveLength(2); // entry_7 and entry_6 are most recent, so they are detailed; entry_0 and entry_1 are oldest, so listed expect(result.listed.map((e) => e.id)).toEqual(["entry_1", "entry_0"]); });
it("formats harness prompt with two-tier detail and stubs without hiding entries behind a black hole", () => { const memoryRecords: Record<string, HarnessEntry> = {}; for (let i = 0; i < 8; i++) { memoryRecords[`mem_${i}`] = makeEntry({ id: `mem_${i}`, title: `Memory Title ${i}`, content: `Content for memory ${i}`, updated_at: `2026-06-0${i + 1}T00:00:00.000Z`, }); } const state: HarnessState = { schema: 1, entries: { prompt: {}, memory: memoryRecords, skill: {}, subagent: {} }, refinements: [], };
const prompt = formatHarnessStateForPrompt(state, { maxEntriesPerKind: 6 }); // 6 entries should have detailed content expect(prompt).toContain("- [global:mem_7] Memory Title 7 (memory/mem_7.md, v1): Content for memory 7"); expect(prompt).toContain("- [global:mem_2] Memory Title 2 (memory/mem_2.md, v1): Content for memory 2"); // older 2 entries should be addressable stubs without content expect(prompt).toContain("- [global:mem_1] Memory Title 1 (memory/mem_1.md, v1)"); expect(prompt).toContain("- [global:mem_0] Memory Title 0 (memory/mem_0.md, v1)"); expect(prompt).not.toContain("Content for memory 0"); // no overflow line because stubs fit in maxListedEntries expect(prompt).not.toContain("more memory entries"); });
it("emits overflow line when stubs exceed maxListedEntries", () => { const memoryRecords: Record<string, HarnessEntry> = {}; for (let i = 0; i < 10; i++) { memoryRecords[`mem_${i}`] = makeEntry({ id: `mem_${i}` }); } const state: HarnessState = { schema: 1, entries: { prompt: {}, memory: memoryRecords, skill: {}, subagent: {} }, refinements: [], };
// 2 detailed, 3 stubs, 5 overflow const prompt = formatHarnessStateForPrompt(state, { maxEntriesPerKind: 2, maxListedEntries: 3 }); expect(prompt).toContain("- +5 more memory entries"); });
it("validates addressable entry ids for overview stubs", () => { expect(isAddressableHarnessId("valid-id_123")).toBe(true); expect(isAddressableHarnessId("")).toBe(false); expect(isAddressableHarnessId("has[bracket]")).toBe(false); expect(isAddressableHarnessId("has\nnewline")).toBe(false); });});