Files

27 lines
7.6 KiB
TypeScript

import { mkdtempSync, readFileSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { resolve } from "node:path";
import { describe, expect, it } from "vitest";
import { AgentTree, AgentTreeError } from "../src/agent-tree.js";
import { DEFAULT_CONFIG } from "../src/config.js";
import { RunStore } from "../src/run-store.js";
import { FakeRuntimeFactory } from "./helpers.js";
const request = (parentId: string, name: string, task: string, background = false) => ({ parentId, name, role: `${name} specialist`, task, expectedOutput: "Evidence report", completionCriteria: "Report is saved", background });
function setup(limits = DEFAULT_CONFIG, delay = 0) { const dir = mkdtempSync(resolve(tmpdir(), "hm-tree-")); const store = new RunStore(dir); const factory = new FakeRuntimeFactory(delay); const tree = new AgentTree(store, factory, limits, { goal: "Investigate a testable question" }); return { dir, store, factory, tree }; }
describe("AgentTree", () => {
it("builds parent → child → grandchild lineage and returns structured results", async () => { const { tree } = setup(); const child = await tree.spawn(request(tree.rootId, "Literature", "Survey primary literature for mechanism alpha")); const grandchild = await tree.spawn(request(child.id, "Verifier", "Verify the strongest primary source independently")); expect(grandchild.lineage).toEqual([tree.rootId, child.id]); expect(grandchild.depth).toBe(2); const result = await tree.start(grandchild.id); expect(result.structured).toEqual({ id: grandchild.id }); expect(tree.render()).toContain("└─ literature-"); expect(tree.render()).toContain(" └─ verifier-"); });
it("blocks duplicates unless explicitly independent replication", async () => { const { tree } = setup(); const first = await tree.spawn(request(tree.rootId, "One", "Reproduce published numerical result number one")); await expect(tree.spawn(request(tree.rootId, "Two", "Reproduce published numerical result number one"))).rejects.toThrow(/Duplicate/); await expect(tree.spawn({ ...request(tree.rootId, "Replica", "Reproduce published numerical result number one"), replicationOf: first.id, independentContext: true })).resolves.toBeTruthy(); });
it("enforces configurable depth and child limits", async () => { const { tree } = setup({ ...DEFAULT_CONFIG, max_depth: 1, max_children_per_agent: 1 }); const child = await tree.spawn(request(tree.rootId, "Only", "Complete a sufficiently concrete unique assignment")); await expect(tree.spawn(request(tree.rootId, "Extra", "Complete a different sufficiently concrete assignment"))).rejects.toThrow(/child limit/); await expect(tree.spawn(request(child.id, "Deep", "Explore a deeper sufficiently concrete assignment"))).rejects.toThrow(/depth/); });
it("supports parallel children, steering, follow-up, and recursive cancellation", async () => { const { tree, factory } = setup(DEFAULT_CONFIG, 40); const one = await tree.spawn(request(tree.rootId, "One", "Parallel investigation branch number one", true)); const two = await tree.spawn(request(tree.rootId, "Two", "Parallel investigation branch number two", true)); await new Promise((resolve) => setTimeout(resolve, 5)); await tree.steer(one.id, "focus"); await tree.followUp(two.id, "cite sources"); expect(factory.runtimes.get(one.id)?.steered).toEqual(["focus"]); await tree.cancelBranch(tree.rootId); expect(tree.inspect(one.id).status).toBe("cancelled"); expect(tree.inspect(two.id).status).toBe("cancelled"); });
it("restores relationships and marks active work interrupted", async () => { const { store, tree } = setup(); const child = await tree.spawn(request(tree.rootId, "Crash", "Investigate recovery behavior after process crash")); const path = store.manifestPath(tree.runId); const manifest = JSON.parse(readFileSync(path, "utf8")); manifest.agents[child.id].status = "running"; writeFileSync(path, JSON.stringify(manifest)); const restored = AgentTree.restore(store, new FakeRuntimeFactory(), DEFAULT_CONFIG, tree.runId); expect(restored.inspect(child.id).status).toBe("interrupted"); expect(restored.inspect(child.id).parentId).toBe(tree.rootId); });
it("rejects vague tasks", async () => { const { tree } = setup(); await expect(tree.spawn(request(tree.rootId, "Vague", "look"))).rejects.toThrow(AgentTreeError); });
it("bridges partial child results to the Supervisor root", async () => { const dir = mkdtempSync(resolve(tmpdir(), "hm-tree-message-")); const updates: string[] = []; const tree = new AgentTree(new RunStore(dir), new FakeRuntimeFactory(), DEFAULT_CONFIG, { goal: "Receive upward results", onRootMessage: (from, message) => updates.push(`${from}:${message}`) }); await tree.message(tree.rootId, "partial evidence", "child-1"); expect(updates).toEqual(["child-1:partial evidence"]); });
it("keeps the root spawnable after stop, and restart revives the run", async () => { const { tree } = setup(); const child = await tree.spawn(request(tree.rootId, "First", "First sufficiently concrete research assignment")); await tree.start(child.id); await tree.stop(); expect(tree.status).toBe("stopped"); expect(tree.inspect(tree.rootId).status).not.toBe("cancelled"); tree.restart("Entirely new research question to investigate"); expect(tree.status).toBe("active"); expect(tree.inspect(tree.rootId).status).toBe("created"); expect(tree.inspect(child.id).status).toBe("archived"); expect(tree.inspect(tree.rootId).children).toEqual([]); await expect(tree.spawn(request(tree.rootId, "Next", "Next sufficiently concrete research assignment"))).resolves.toBeTruthy(); });
it("does not count cancelled or failed children toward the child limit", async () => { const { tree } = setup({ ...DEFAULT_CONFIG, max_children_per_agent: 2 }); const first = await tree.spawn(request(tree.rootId, "Alpha", "First concrete assignment in a limited tree")); await tree.cancel(first.id); const second = await tree.spawn(request(tree.rootId, "Beta", "Second concrete assignment in a limited tree")); await tree.cancel(second.id); await expect(tree.spawn(request(tree.rootId, "Gamma", "Third concrete assignment in a limited tree"))).resolves.toBeTruthy(); });
it("cancel does not clobber an already recorded result", async () => { const { tree } = setup(); const child = await tree.spawn(request(tree.rootId, "Done", "Complete a concrete assignment and return")); const result = await tree.start(child.id); expect(result.status).toBe("completed"); await tree.cancel(child.id); expect(tree.inspect(child.id).status).toBe("cancelled"); expect(tree.inspect(child.id).result?.status).toBe("completed"); });
it("fails agents that exceed the configured timeout", async () => { const { tree } = setup({ ...DEFAULT_CONFIG, agent_timeout_seconds: 0.05 }, 200); const child = await tree.spawn(request(tree.rootId, "Slow", "Run a deliberately slow concrete assignment")); const result = await tree.start(child.id); expect(result.status).toBe("failed"); expect(result.summary).toMatch(/timed out/); expect(tree.inspect(child.id).error).toMatch(/timed out/); });
it("caps how many agent sessions stream at the same time", async () => { const { tree } = setup({ ...DEFAULT_CONFIG, agent_concurrency: 1 }, 40); const a = await tree.spawn(request(tree.rootId, "First", "Sequential streaming assignment number one")); const b = await tree.spawn(request(tree.rootId, "Second", "Sequential streaming assignment number two")); const [ra, rb] = await Promise.all([tree.start(a.id), tree.start(b.id)]); expect(ra.status).toBe("completed"); expect(rb.status).toBe("completed"); const startedB = Date.parse(tree.inspect(b.id).startedAt!); const finishedA = Date.parse(tree.inspect(a.id).finishedAt!); expect(startedB).toBeGreaterThanOrEqual(finishedA - 5); });
});