chore(vendor): import pi 1.3.0
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
This commit is contained in:
alexion committed
1 parent
719370c8f7
commit
6081e3e325
1743 files changed
+205459
-131704
No files matched your search
@@ -0,0 +1,405 @@
|
||||
import { mkdtemp, rm } from "node:fs/promises";
|
||||
import { tmpdir } from "node:os";
|
||||
import { dirname, join } from "node:path";
|
||||
import type { JsonValue } from "@earendil-works/chord";
|
||||
import {
|
||||
type AssistantMessage,
|
||||
fauxAssistantMessage,
|
||||
fauxText,
|
||||
fauxToolCall,
|
||||
type ToolResultMessage,
|
||||
Type,
|
||||
} from "@earendil-works/pi-ai";
|
||||
import {
|
||||
defineTool,
|
||||
type EntryRecord,
|
||||
GenerationTask,
|
||||
type Harness,
|
||||
LiveDoc,
|
||||
type TaskId,
|
||||
type ToolRegistration,
|
||||
ToolResultEntry,
|
||||
ToolTask,
|
||||
} from "@earendil-works/pi-durable";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import type { ExecutionEnv } from "../src/env/index.ts";
|
||||
import { NodeExecutionEnv } from "../src/env/node.ts";
|
||||
import { openNodeSqliteStorage } from "../src/storage/sqlite/node.ts";
|
||||
import { createBashTool } from "../src/tools/index.ts";
|
||||
import { allEntries, type ChatSetup, chatSetup, openChat, waitFor } from "./chat-support.ts";
|
||||
import { addHooks, addTool, type Installed } from "./harness-support.ts";
|
||||
import { context } from "./session-support.ts";
|
||||
import { aborted, deferred } from "./task-support.ts";
|
||||
|
||||
const directories = new Set<string>();
|
||||
|
||||
afterEach(async () => {
|
||||
for (const directory of directories) await rm(directory, { recursive: true, force: true });
|
||||
directories.clear();
|
||||
});
|
||||
|
||||
async function sqlitePath(): Promise<string> {
|
||||
const directory = await mkdtemp(join(tmpdir(), "pi-durable-tools-"));
|
||||
directories.add(directory);
|
||||
return join(directory, "session.sqlite");
|
||||
}
|
||||
|
||||
async function open(path: string, setup: ChatSetup, env?: ExecutionEnv) {
|
||||
const opened = await openChat(await openNodeSqliteStorage(path), setup, env === undefined ? {} : { env });
|
||||
opened.harness.resume();
|
||||
return opened;
|
||||
}
|
||||
|
||||
function tool(
|
||||
name: string,
|
||||
execute: ToolRegistration["execute"],
|
||||
extra: Partial<ToolRegistration> = {},
|
||||
): ToolRegistration {
|
||||
return defineTool({ name, description: name, parameters: Type.Object({}), execute, ...extra });
|
||||
}
|
||||
|
||||
function call(name: string, id = "c1"): AssistantMessage {
|
||||
return fauxAssistantMessage([fauxToolCall(name, {}, { id })], { stopReason: "toolUse" });
|
||||
}
|
||||
|
||||
const DONE = fauxAssistantMessage([fauxText("done")]);
|
||||
|
||||
function results(entries: readonly EntryRecord[]): ToolResultMessage[] {
|
||||
return entries.filter((entry) => ToolResultEntry.is(entry)).map((entry) => entry.model![0] as ToolResultMessage);
|
||||
}
|
||||
|
||||
function text(message: ToolResultMessage | undefined): string {
|
||||
return (message?.content ?? []).map((item) => (item.type === "text" ? item.text : "")).join("|");
|
||||
}
|
||||
|
||||
async function toolTaskId(harness: Harness): Promise<TaskId> {
|
||||
let id: TaskId | undefined;
|
||||
await waitFor(async () => {
|
||||
id = (await harness.snapshot(LiveDoc, 1 as never, context))?.tools?.[0]?.taskId;
|
||||
return id !== undefined;
|
||||
});
|
||||
return id!;
|
||||
}
|
||||
|
||||
/**
|
||||
* A tool that writes output, then blocks until its invocation is cancelled the first `blocking` times it runs.
|
||||
* `started` resolves once the output is durable.
|
||||
*/
|
||||
function blockingTool(name: string, extra: Partial<ToolRegistration> = {}) {
|
||||
const started = deferred();
|
||||
const state = { runs: 0, blocking: 1 };
|
||||
const registration = tool(
|
||||
name,
|
||||
async (_args, api, callContext) => {
|
||||
state.runs++;
|
||||
api.output(`run ${state.runs}\n`);
|
||||
await api.details({ run: state.runs }, callContext);
|
||||
if (state.runs <= state.blocking) {
|
||||
started.resolve();
|
||||
await aborted(callContext.abortSignal!);
|
||||
}
|
||||
return {};
|
||||
},
|
||||
extra,
|
||||
);
|
||||
return { registration, started, state };
|
||||
}
|
||||
|
||||
describe("tool recovery", () => {
|
||||
it("answers an unsafe tool interrupted after intent with its durable partial output", async () => {
|
||||
const path = await sqlitePath();
|
||||
const setup = chatSetup();
|
||||
const { registration, started, state } = blockingTool("work");
|
||||
addTool(setup.registry, registration);
|
||||
setup.faux.setResponses([call("work"), DONE]);
|
||||
let opened = await open(path, setup);
|
||||
const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
|
||||
await started.promise;
|
||||
const taskId = await toolTaskId(opened.harness);
|
||||
await opened.harness.close(context);
|
||||
|
||||
opened = await open(path, setup);
|
||||
expect((await opened.harness.getTask(taskId, context))?.state).toMatchObject({
|
||||
checkpoint: { phase: "execute", arguments: {}, replay: "unsafe" },
|
||||
});
|
||||
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
|
||||
expect(state.runs).toBe(1);
|
||||
const [result] = results(await allEntries(opened.root));
|
||||
expect(result).toMatchObject({ isError: true, details: { run: 1 } });
|
||||
expect(text(result)).toBe(
|
||||
"run 1\n|<harness>\n[error] Tool work was interrupted and may have partially run\n</harness>",
|
||||
);
|
||||
expect(await opened.harness.snapshot(LiveDoc, opened.root.id, context)).toEqual({});
|
||||
await opened.harness.close(context);
|
||||
});
|
||||
|
||||
it("reruns a tool only when both the stored and the current replay policy are safe", async () => {
|
||||
const cases: [stored: "safe" | "unsafe", current: "safe" | "unsafe", reruns: boolean][] = [
|
||||
["safe", "safe", true],
|
||||
["safe", "unsafe", false],
|
||||
["unsafe", "safe", false],
|
||||
];
|
||||
for (const [stored, current, reruns] of cases) {
|
||||
const path = await sqlitePath();
|
||||
const setup = chatSetup();
|
||||
const blocking = blockingTool("work", { replay: stored });
|
||||
const registration: Installed = addTool(setup.registry, blocking.registration);
|
||||
setup.faux.setResponses([call("work"), DONE]);
|
||||
let opened = await open(path, setup);
|
||||
const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
|
||||
await blocking.started.promise;
|
||||
await opened.harness.close(context);
|
||||
|
||||
registration.dispose();
|
||||
addTool(setup.registry, { ...blocking.registration, replay: current });
|
||||
opened = await open(path, setup);
|
||||
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
|
||||
const [result] = results(await allEntries(opened.root));
|
||||
expect(blocking.state.runs).toBe(reruns ? 2 : 1);
|
||||
expect(result!.isError).toBe(!reruns);
|
||||
if (reruns) expect(text(result)).toBe("run 2\n");
|
||||
await opened.harness.close(context);
|
||||
}
|
||||
});
|
||||
|
||||
it("reruns a safe tool with the environment of the conversation's cwd at rerun, and not once it is deselected", async () => {
|
||||
for (const change of ["cwd", "deselect"] as const) {
|
||||
const path = await sqlitePath();
|
||||
const setup = chatSetup();
|
||||
const cwds: (string | undefined)[] = [];
|
||||
const started = deferred();
|
||||
const work = tool(
|
||||
"work",
|
||||
async (_args, api, callContext) => {
|
||||
cwds.push(api.env?.cwd);
|
||||
if (cwds.length === 1) {
|
||||
started.resolve();
|
||||
await aborted(callContext.abortSignal!);
|
||||
}
|
||||
return {};
|
||||
},
|
||||
{ replay: "safe" },
|
||||
);
|
||||
addTool(setup.registry, work);
|
||||
setup.faux.setResponses([call("work"), DONE]);
|
||||
const env = ({ cwd = "/" }: { readonly cwd?: string }) => new NodeExecutionEnv({ cwd });
|
||||
const reopen = async () => {
|
||||
const opened = await openChat(await openNodeSqliteStorage(path), setup, { env });
|
||||
opened.harness.resume();
|
||||
return opened;
|
||||
};
|
||||
let opened = await reopen();
|
||||
await opened.root.configure({ cwd: "/one" }, context);
|
||||
const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
|
||||
await started.promise;
|
||||
await opened.root.configure(change === "cwd" ? { cwd: "/two" } : { extensions: [] }, context);
|
||||
await opened.harness.close(context);
|
||||
|
||||
opened = await reopen();
|
||||
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
|
||||
const [result] = results(await allEntries(opened.root));
|
||||
if (change === "cwd") {
|
||||
expect(cwds).toEqual(["/one", "/two"]);
|
||||
expect(result!.isError).toBe(false);
|
||||
} else {
|
||||
// A tool that no longer resolves is treated as unsafe: interrupted, not rerun.
|
||||
expect(cwds).toEqual(["/one"]);
|
||||
expect(text(result)).toContain("was interrupted");
|
||||
}
|
||||
await opened.harness.close(context);
|
||||
}
|
||||
});
|
||||
|
||||
it("reruns beforeTool when interrupted before intent, and executes once", async () => {
|
||||
const path = await sqlitePath();
|
||||
const setup = chatSetup();
|
||||
let runs = 0;
|
||||
addTool(
|
||||
setup.registry,
|
||||
tool("work", async () => {
|
||||
runs++;
|
||||
return { output: [] };
|
||||
}),
|
||||
);
|
||||
const reached = deferred();
|
||||
let asked = 0;
|
||||
const decisions: string[] = [];
|
||||
addHooks(setup.registry, ToolTask, {
|
||||
beforeTool: async (_call, api, callContext) => {
|
||||
asked++;
|
||||
// A durable first-writer-wins decision survives the rerun.
|
||||
decisions.push(await api.memo("test:decision", `attempt ${asked}`, callContext));
|
||||
if (asked === 1) {
|
||||
reached.resolve();
|
||||
await aborted(callContext.abortSignal!);
|
||||
}
|
||||
return undefined;
|
||||
},
|
||||
});
|
||||
setup.faux.setResponses([call("work"), DONE]);
|
||||
let opened = await open(path, setup);
|
||||
const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
|
||||
await reached.promise;
|
||||
await opened.harness.close(context);
|
||||
|
||||
opened = await open(path, setup);
|
||||
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
|
||||
expect([asked, runs]).toEqual([2, 1]);
|
||||
expect(decisions).toEqual(["attempt 1", "attempt 1"]);
|
||||
await opened.harness.close(context);
|
||||
});
|
||||
|
||||
it("reruns the generation tools phase interrupted before its commit", async () => {
|
||||
const path = await sqlitePath();
|
||||
const setup = chatSetup();
|
||||
addTool(
|
||||
setup.registry,
|
||||
tool("work", async () => ({ output: [] })),
|
||||
);
|
||||
const reached = deferred();
|
||||
let observed = 0;
|
||||
addHooks(setup.registry, GenerationTask, {
|
||||
afterTools: async (_assistant, _results, _api, callContext) => {
|
||||
observed++;
|
||||
if (observed === 1) {
|
||||
reached.resolve();
|
||||
await aborted(callContext.abortSignal!);
|
||||
}
|
||||
},
|
||||
});
|
||||
setup.faux.setResponses([call("work"), DONE]);
|
||||
let opened = await open(path, setup);
|
||||
const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
|
||||
await reached.promise;
|
||||
await opened.harness.close(context);
|
||||
|
||||
opened = await open(path, setup);
|
||||
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
|
||||
expect(observed).toBe(2);
|
||||
expect((await allEntries(opened.root)).map((entry) => entry.kind)).toEqual([
|
||||
"pi.user",
|
||||
"pi.system",
|
||||
"pi.assistant",
|
||||
"pi.tool-result",
|
||||
"pi.assistant",
|
||||
]);
|
||||
await opened.harness.close(context);
|
||||
});
|
||||
|
||||
it("answers an aborted tool with its partial output and continues the run", async () => {
|
||||
const path = await sqlitePath();
|
||||
const setup = chatSetup();
|
||||
const { registration, started } = blockingTool("work");
|
||||
addTool(setup.registry, registration);
|
||||
setup.faux.setResponses([call("work"), DONE]);
|
||||
const opened = await open(path, setup);
|
||||
const submission = await opened.root.submit({ type: "input", content: "go" }, context);
|
||||
await started.promise;
|
||||
const taskId = await toolTaskId(opened.harness);
|
||||
expect(await opened.harness.abortTask(taskId, context)).toBe("marked");
|
||||
expect((await opened.harness.waitForTask(taskId, context)).state.outcome).toMatchObject({
|
||||
status: "aborted",
|
||||
result: { entryId: expect.any(Number) },
|
||||
});
|
||||
expect((await submission.wait(context)).status).toBe("done");
|
||||
const [result] = results(await allEntries(opened.root));
|
||||
expect(text(result)).toBe("run 1\n|<harness>\n[error] Tool work was aborted\n</harness>");
|
||||
await opened.harness.close(context);
|
||||
});
|
||||
|
||||
it("lets context derivation answer a faulted tool and continues the run", async () => {
|
||||
const path = await sqlitePath();
|
||||
const setup = chatSetup();
|
||||
// A result that is not strict JSON makes the result commit throw, so the scheduler faults the task.
|
||||
addTool(
|
||||
setup.registry,
|
||||
tool("bad", async () => ({ output: [], details: { fn: (() => 1) as unknown as JsonValue } })),
|
||||
);
|
||||
const requests: string[] = [];
|
||||
setup.faux.setResponses([
|
||||
call("bad"),
|
||||
(request) => {
|
||||
const result = request.messages.find((message) => message.role === "toolResult");
|
||||
requests.push(result?.role === "toolResult" ? text(result) : "none");
|
||||
return DONE;
|
||||
},
|
||||
]);
|
||||
const opened = await open(path, setup);
|
||||
const settled = await (await opened.root.submit({ type: "input", content: "go" }, context)).wait(context);
|
||||
expect(settled.status).toBe("done");
|
||||
expect(results(await allEntries(opened.root))).toEqual([]);
|
||||
expect(requests).toEqual(["Tool result unavailable: history ends before this call completed."]);
|
||||
const tasks = await opened.harness.commit((tx) => tx.scanTasks({ conversationId: opened.root.id }, 20), context);
|
||||
const faulted = tasks.items.find((task) => task.kind === "pi.tool")!;
|
||||
expect(faulted.state).toMatchObject({ status: "terminal", outcome: { status: "faulted" } });
|
||||
await opened.harness.close(context);
|
||||
});
|
||||
|
||||
it("answers a real bash command interrupted by close and reopen, then finishes the run", async () => {
|
||||
const path = await sqlitePath();
|
||||
const env = new NodeExecutionEnv({ cwd: dirname(path) });
|
||||
const setup = chatSetup();
|
||||
addTool(setup.registry, createBashTool());
|
||||
setup.faux.setResponses([
|
||||
fauxAssistantMessage([fauxToolCall("bash", { command: "echo started; sleep 30" }, { id: "b" })], {
|
||||
stopReason: "toolUse",
|
||||
}),
|
||||
DONE,
|
||||
]);
|
||||
let opened = await open(path, setup, env);
|
||||
const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
|
||||
await waitFor(async () => {
|
||||
const slot = (await opened.harness.snapshot(LiveDoc, opened.root.id, context))?.tools?.[0];
|
||||
return slot?.output === "started\n";
|
||||
});
|
||||
await opened.harness.close(context);
|
||||
|
||||
opened = await open(path, setup, env);
|
||||
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
|
||||
const [result] = results(await allEntries(opened.root));
|
||||
expect(text(result)).toBe(
|
||||
"started\n|<harness>\n[error] Tool bash was interrupted and may have partially run\n</harness>",
|
||||
);
|
||||
await opened.harness.close(context);
|
||||
});
|
||||
|
||||
it("clears the interrupted attempt's progress before a safe rerun", async () => {
|
||||
const path = await sqlitePath();
|
||||
const setup = chatSetup();
|
||||
const started = [deferred(), deferred()];
|
||||
let runs = 0;
|
||||
addTool(
|
||||
setup.registry,
|
||||
tool(
|
||||
"work",
|
||||
async (_args, api, callContext) => {
|
||||
const run = runs++;
|
||||
if (run === 0) {
|
||||
api.diagnostic({ severity: "info", message: "first a" });
|
||||
api.diagnostic({ severity: "info", message: "first b" });
|
||||
await api.details({ run: 1, extra: true }, callContext);
|
||||
} else {
|
||||
api.diagnostic({ severity: "info", message: "second" });
|
||||
await api.details({ run: 2 }, callContext);
|
||||
}
|
||||
started[run]!.resolve();
|
||||
await aborted(callContext.abortSignal!);
|
||||
return {};
|
||||
},
|
||||
{ replay: "safe" },
|
||||
),
|
||||
);
|
||||
setup.faux.setResponses([call("work"), DONE]);
|
||||
let opened = await open(path, setup);
|
||||
await opened.root.submit({ type: "input", content: "go" }, context);
|
||||
await started[0]!.promise;
|
||||
await opened.harness.close(context);
|
||||
|
||||
opened = await open(path, setup);
|
||||
await started[1]!.promise;
|
||||
const slot = (await opened.harness.snapshot(LiveDoc, opened.root.id, context))?.tools?.[0];
|
||||
expect(slot?.diagnostics).toEqual([{ severity: "info", message: "second" }]);
|
||||
expect(slot?.details).toEqual({ run: 2 });
|
||||
await opened.harness.close(context);
|
||||
});
|
||||
});
|
||||
Reference in new issue
Block a user