alexion
6081e3e325
Imported as a linear vendor commit so forge rebases do not replay the upstream root tree. git-subtree-dir: third_party/pi git-subtree-split: 71023aec1546d6d677a65669c5665c0927ce0dd0
301 lines
12 KiB
TypeScript
301 lines
12 KiB
TypeScript
import { Type } from "@earendil-works/pi-ai";
|
|
import {
|
|
type AgentState,
|
|
CompactionTask,
|
|
createRegistry,
|
|
defineExtension,
|
|
defineTask,
|
|
GenerationTask,
|
|
type HarnessSettings,
|
|
hook,
|
|
type RegistrySnapshot,
|
|
section,
|
|
type ToolRegistration,
|
|
ToolTask,
|
|
wrapSection,
|
|
wrapTool,
|
|
} from "@earendil-works/pi-durable";
|
|
import { describe, expect, it } from "vitest";
|
|
import { agentHooks, resolveAgent, resolveSettings } from "../src/harness/agent.ts";
|
|
import { context } from "./session-support.ts";
|
|
|
|
type AppTool = ToolRegistration & { readonly snippet?: string };
|
|
|
|
function tool(name: string, extra: Partial<AppTool> = {}): AppTool {
|
|
return {
|
|
name,
|
|
description: `${name} tool`,
|
|
parameters: Type.Object({}),
|
|
execute: async () => ({ output: [] }),
|
|
...extra,
|
|
};
|
|
}
|
|
|
|
function task(name: string, version = 1) {
|
|
return defineTask<undefined, { phase: "run" }, null>({
|
|
name,
|
|
version,
|
|
initial: () => ({ phase: "run" }),
|
|
phases: { run: async () => {} },
|
|
abort: async () => {},
|
|
});
|
|
}
|
|
|
|
function names(items: readonly { readonly name: string }[]): string[] {
|
|
return items.map((item) => item.name);
|
|
}
|
|
|
|
/** Resolve `state` against `snapshot` and `settings`, collecting reports. */
|
|
function resolve(
|
|
state: AgentState | undefined,
|
|
snapshot: RegistrySnapshot<AppTool>,
|
|
settings: HarnessSettings = {},
|
|
reports: unknown[] = [],
|
|
) {
|
|
return resolveAgent(state, snapshot, resolveSettings(settings), (error) => reports.push(error));
|
|
}
|
|
|
|
async function rendered(agent: ReturnType<typeof resolve>): Promise<[string, string | undefined][]> {
|
|
const input = {
|
|
conversationId: 1 as never,
|
|
agent,
|
|
env: undefined,
|
|
shown: {},
|
|
read: { snapshot: async () => undefined, snapshotAsOf: async () => undefined },
|
|
};
|
|
return Promise.all(
|
|
agent.sections.map(async (item) => [item.key, await item.render(input, context)] as [string, string | undefined]),
|
|
);
|
|
}
|
|
|
|
describe("registry", () => {
|
|
it("installs, replaces in place, and uninstalls extensions by name", () => {
|
|
const registry = createRegistry<AppTool>();
|
|
const listener: string[] = [];
|
|
registry.subscribe(() => listener.push(names(registry.snapshot().installed()).join(",")));
|
|
const a = defineExtension<AppTool>({ name: "a", tools: [tool("read", { snippet: "Read files" })] });
|
|
const b = defineExtension<AppTool>({ name: "b", tools: [tool("read"), tool("bash")] });
|
|
registry.install(a);
|
|
registry.install(b);
|
|
const before = registry.snapshot();
|
|
// A new object with an installed name replaces it at its position.
|
|
const a2 = defineExtension<AppTool>({ name: "a", tools: [tool("grep")] });
|
|
registry.install(a2);
|
|
expect(names(registry.snapshot().installed())).toEqual(["a", "b"]);
|
|
expect(registry.snapshot().extension("a")).toBe(a2);
|
|
expect(
|
|
registry
|
|
.snapshot()
|
|
.tools()
|
|
.map(({ extension, tool }) => `${tool.name}@${extension.name}`),
|
|
).toEqual(["grep@a", "read@b", "bash@b"]);
|
|
// Old snapshots stay as they were.
|
|
expect(before.extension("a")).toBe(a);
|
|
expect(before.tools()[0]!.tool.snippet).toBe("Read files");
|
|
// Uninstall matches the name, whichever object; a later install appends.
|
|
registry.uninstall(a);
|
|
registry.uninstall(a);
|
|
expect(names(registry.snapshot().installed())).toEqual(["b"]);
|
|
registry.install(a);
|
|
expect(names(registry.snapshot().installed())).toEqual(["b", "a"]);
|
|
expect(listener).toEqual(["a", "a,b", "a,b", "b", "b,a"]);
|
|
});
|
|
|
|
it("validates the registry as it would be after an install and publishes nothing when invalid", () => {
|
|
const registry = createRegistry();
|
|
const Tasks = defineExtension({ name: "tasks", tasks: [task("app.index")] });
|
|
registry.install(Tasks);
|
|
const published: number[] = [];
|
|
registry.subscribe(() => published.push(1));
|
|
const before = registry.snapshot();
|
|
const invalid = [
|
|
[defineExtension({ name: "x", tools: [tool("read"), tool("read")] }), "two tools named read"],
|
|
[defineExtension({ name: "x", sections: [section("a", () => "1"), section("a", () => "2")] }), "two sections"],
|
|
[defineExtension({ name: "x", sections: [section("Bad Key", () => "")] }), "must match"],
|
|
[defineExtension({ name: "x", sections: [section("instructions", () => "")] }), "reserved"],
|
|
[defineExtension({ name: "x", tasks: [task("pi.generation")] }), "already installed"],
|
|
[defineExtension({ name: "x", tasks: [task("app.index")] }), "already installed"],
|
|
] as const;
|
|
for (const [extension, message] of invalid) expect(() => registry.install(extension)).toThrow(message);
|
|
expect(registry.snapshot()).toBe(before);
|
|
expect(published).toEqual([]);
|
|
// Replacing the extension that holds a task name is valid: the check runs on the state after replacement.
|
|
registry.install(defineExtension({ name: "tasks", tasks: [task("app.index", 2)] }));
|
|
expect(registry.snapshot().task("app.index")?.definition.version).toBe(2);
|
|
});
|
|
|
|
it("always holds the built-in tasks, which are not an extension", () => {
|
|
const registry = createRegistry();
|
|
expect(registry.snapshot().installed()).toEqual([]);
|
|
expect(registry.snapshot().tasks()).toEqual([GenerationTask, ToolTask, CompactionTask]);
|
|
const custom = task("app.custom");
|
|
registry.install(defineExtension({ name: "custom", tasks: [custom] }));
|
|
expect(registry.snapshot().tasks()).toEqual([GenerationTask, ToolTask, CompactionTask, custom]);
|
|
expect(registry.snapshot().task("app.custom")).toBe(custom);
|
|
registry.uninstall(defineExtension({ name: "custom" }));
|
|
expect(registry.snapshot().task("app.custom")).toBeUndefined();
|
|
});
|
|
});
|
|
|
|
describe("settings", () => {
|
|
it("keeps the default of a progress interval given as undefined", () => {
|
|
expect(resolveSettings({ progress: { partialIntervalMs: undefined, outputIntervalMs: 250 } }).progress).toEqual({
|
|
partialIntervalMs: 100,
|
|
outputIntervalMs: 250,
|
|
});
|
|
});
|
|
});
|
|
|
|
describe("agent resolution", () => {
|
|
const read = tool("read");
|
|
const bash = tool("bash");
|
|
const edit = tool("edit");
|
|
const Coding = defineExtension<AppTool>({
|
|
name: "coding",
|
|
tools: [read, bash, edit],
|
|
sections: [section("preamble", () => "You code.", { tag: false }), section("cwd", () => "/repo")],
|
|
});
|
|
const Skills = defineExtension<AppTool>({ name: "skills", sections: [section("skills", () => "S")] });
|
|
const Reviewer = defineExtension<AppTool>({ name: "reviewer", sections: [section("role", () => "Review.")] });
|
|
const registry = createRegistry<AppTool>();
|
|
for (const extension of [Coding, Skills, Reviewer]) registry.install(extension);
|
|
const snapshot = registry.snapshot();
|
|
|
|
it("selects the default, an array, or the default edited by add and remove", () => {
|
|
expect(names(resolve(undefined, snapshot).extensions)).toEqual(["coding", "skills", "reviewer"]);
|
|
const settings = { extensions: [Coding, Skills] };
|
|
expect(names(resolve({}, snapshot, settings).extensions)).toEqual(["coding", "skills"]);
|
|
expect(names(resolve({ extensions: ["reviewer", "coding"] }, snapshot, settings).extensions)).toEqual([
|
|
"reviewer",
|
|
"coding",
|
|
]);
|
|
// Add appends, remove drops, duplicates keep their first position, uninstalled names are skipped.
|
|
const edited = { extensions: { add: ["reviewer", "coding", "gone"], remove: ["skills"] } };
|
|
expect(names(resolve(edited, snapshot, settings).extensions)).toEqual(["coding", "reviewer"]);
|
|
// An old object stands for its name: the installed extension is selected.
|
|
const stale = { extensions: [defineExtension({ name: "skills" }), Skills] };
|
|
expect(resolve({}, snapshot, stale).extensions).toEqual([Skills]);
|
|
});
|
|
|
|
it("skips uninstalled names and resolves them again once they are installed", () => {
|
|
const local = createRegistry<AppTool>();
|
|
local.install(Coding);
|
|
const state: AgentState = { extensions: ["coding", "skills"] };
|
|
expect(names(resolve(state, local.snapshot()).extensions)).toEqual(["coding"]);
|
|
local.install(Skills);
|
|
expect(names(resolve(state, local.snapshot()).extensions)).toEqual(["coding", "skills"]);
|
|
});
|
|
|
|
it("replaces same-name tools in place, wraps the winner, then applies the filter", () => {
|
|
const local = createRegistry<AppTool>();
|
|
const venvBash = tool("bash", { description: "venv bash" });
|
|
const calls: string[] = [];
|
|
const Venv = defineExtension<AppTool>({ name: "venv", tools: [venvBash] });
|
|
const Timing = defineExtension<AppTool>({
|
|
name: "timing",
|
|
wraps: [
|
|
wrapTool(bash, (inner) => ({ ...inner, description: `${inner.description} (timed)` })),
|
|
wrapTool(bash, (inner) => ({ ...inner, description: `${inner.description} [2]` })),
|
|
// No `grep` is selected: the wrapper does nothing and reports nothing.
|
|
wrapTool(tool("grep"), () => {
|
|
calls.push("grep");
|
|
return tool("grep");
|
|
}),
|
|
],
|
|
});
|
|
for (const extension of [Coding, Venv, Timing]) local.install(extension);
|
|
const reports: unknown[] = [];
|
|
const agent = resolve(undefined, local.snapshot(), {}, reports);
|
|
expect(agent.tools.map((each) => [each.name, each.description])).toEqual([
|
|
["read", "read tool"],
|
|
["bash", "venv bash (timed) [2]"],
|
|
["edit", "edit tool"],
|
|
]);
|
|
expect(reports).toEqual([]);
|
|
expect(calls).toEqual([]);
|
|
|
|
// An array keeps exactly these names in its order, a repeated name at its first position.
|
|
const filtered = resolve({ tools: ["edit", "missing", "read", "edit"] }, local.snapshot());
|
|
expect(names(filtered.tools)).toEqual(["edit", "read"]);
|
|
expect(names(resolve({ tools: { remove: ["bash"] } }, local.snapshot()).tools)).toEqual(["read", "edit"]);
|
|
});
|
|
|
|
it("drops a tool or section whose wrapper throws or renames it and reports the failure", () => {
|
|
const local = createRegistry<AppTool>();
|
|
const Broken = defineExtension<AppTool>({
|
|
name: "broken",
|
|
wraps: [
|
|
wrapTool(read, () => {
|
|
throw new Error("wrapper failed");
|
|
}),
|
|
wrapTool(edit, (inner) => ({ ...inner, name: "renamed" })),
|
|
wrapSection("cwd", () => {
|
|
throw new Error("section wrapper failed");
|
|
}),
|
|
],
|
|
});
|
|
local.install(Coding);
|
|
local.install(Broken);
|
|
const reports: unknown[] = [];
|
|
const agent = resolve(undefined, local.snapshot(), {}, reports);
|
|
expect(names(agent.tools)).toEqual(["bash"]);
|
|
expect(agent.sections.map((each) => each.key)).toEqual(["preamble"]);
|
|
expect(reports.map((error) => (error as Error).message)).toEqual([
|
|
"wrapper failed",
|
|
"Wrapper renamed edit to renamed",
|
|
"section wrapper failed",
|
|
]);
|
|
});
|
|
|
|
it("orders sections by extension, replaces same keys in place, and renders instructions last and unwrapped", async () => {
|
|
const local = createRegistry<AppTool>();
|
|
const Override = defineExtension<AppTool>({
|
|
name: "override",
|
|
sections: [section("preamble", () => "You review.", { tag: false })],
|
|
wraps: [
|
|
wrapSection("cwd", (inner) => ({
|
|
...inner,
|
|
render: async (input, ctx) => `${await inner.render(input, ctx)}!`,
|
|
})),
|
|
// Instructions are not wrapped.
|
|
wrapSection("instructions", () => {
|
|
throw new Error("never");
|
|
}),
|
|
],
|
|
});
|
|
for (const extension of [Coding, Skills, Override]) local.install(extension);
|
|
const reports: unknown[] = [];
|
|
const agent = resolve({ instructions: "Be terse." }, local.snapshot(), {}, reports);
|
|
expect(await rendered(agent)).toEqual([
|
|
["preamble", "You review."],
|
|
["cwd", "/repo!"],
|
|
["skills", "S"],
|
|
["instructions", "Be terse."],
|
|
]);
|
|
expect(agent.sections.at(-1)!.tag).toBeUndefined();
|
|
expect(reports).toEqual([]);
|
|
});
|
|
|
|
it("collects hooks of the selected extensions in extension order and applies field defaults", () => {
|
|
const local = createRegistry<AppTool>();
|
|
const first = { beforeTool: () => undefined };
|
|
const second = { beforeTool: () => undefined };
|
|
const onYield = { onYield: () => undefined };
|
|
local.install(
|
|
defineExtension<AppTool>({ name: "a", hooks: [hook(ToolTask, first), hook(GenerationTask, onYield)] }),
|
|
);
|
|
local.install(defineExtension<AppTool>({ name: "b", hooks: [hook(ToolTask, second)] }));
|
|
expect(agentHooks(resolve(undefined, local.snapshot()), "pi.tool")).toEqual([first, second]);
|
|
expect(agentHooks(resolve({ extensions: ["b", "a"] }, local.snapshot()), "pi.tool")).toEqual([second, first]);
|
|
expect(agentHooks(resolve({ extensions: ["b"] }, local.snapshot()), "pi.generation")).toEqual([]);
|
|
|
|
const defaults = resolve(undefined, local.snapshot());
|
|
expect(defaults.model).toBeUndefined();
|
|
expect(defaults.thinkingLevel).toBe("off");
|
|
expect(defaults.cwd).toBeUndefined();
|
|
expect(
|
|
resolve({ model: { provider: "p", modelId: "m" }, thinkingLevel: "high", cwd: "/w" }, local.snapshot()),
|
|
).toMatchObject({ model: { provider: "p", modelId: "m" }, thinkingLevel: "high", cwd: "/w" });
|
|
});
|
|
});
|