mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
7e15f96019
GenerationHooks.afterRun asks, after a run answered, for durable tasks (runEndTask(task, input)). The answer commits and the run's inputs settle done first; the generation then creates the tasks as its children in a finish phase and waits for them in drain, keeping pi.live.run (run.afterRun) and the conversation busy, so idle waits include them, Stop aborts them and a crash resumes them. Queued input is placed when they are terminal. pi.live.run.first and runFailure.first name where a run's history starts, also for a continue after a failure. closes #11141
392 lines
16 KiB
TypeScript
392 lines
16 KiB
TypeScript
import { mkdtemp, rm } from "node:fs/promises";
|
|
import { tmpdir } from "node:os";
|
|
import { join } from "node:path";
|
|
import { DatabaseSync } from "node:sqlite";
|
|
import { type AssistantMessage, createModels, fauxAssistantMessage, type Message } from "@earendil-works/pi-ai";
|
|
import { AssistantEntry, type Harness, LiveDoc, type TaskId } from "@earendil-works/pi-durable";
|
|
import { afterEach, describe, expect, it } from "vitest";
|
|
import { openNodeSqliteStorage } from "../src/storage/sqlite/node.ts";
|
|
import { allEntries, type ChatSetup, chatSetup, openChat, textOf, unanswered, waitFor } from "./chat-support.ts";
|
|
import { addSection } from "./harness-support.ts";
|
|
import { context, onValue } from "./session-support.ts";
|
|
import { aborted, deferred, eventually } 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-generation-"));
|
|
directories.add(directory);
|
|
return join(directory, "session.sqlite");
|
|
}
|
|
|
|
async function open(path: string, setup: ChatSetup) {
|
|
return openChat(await openNodeSqliteStorage(path), setup);
|
|
}
|
|
|
|
async function runTaskId(harness: Harness): Promise<TaskId> {
|
|
const live = await harness.snapshot(LiveDoc, 1 as never, context);
|
|
return live!.run!.taskId;
|
|
}
|
|
|
|
async function checkpoint(harness: Harness, id: TaskId) {
|
|
const record = await harness.getTask(id, context);
|
|
return record?.state.status === "terminal" ? undefined : record?.state.checkpoint;
|
|
}
|
|
|
|
describe("generation recovery", () => {
|
|
it("reruns preparation interrupted before its commit", async () => {
|
|
const path = await sqlitePath();
|
|
const setup = chatSetup();
|
|
const reached = deferred();
|
|
let block = true;
|
|
addSection(setup.registry, "preamble", async (_input, ctx) => {
|
|
if (block) {
|
|
block = false;
|
|
reached.resolve();
|
|
await aborted(ctx.abortSignal!);
|
|
}
|
|
return "p";
|
|
});
|
|
setup.faux.setResponses([fauxAssistantMessage("answer")]);
|
|
let opened = await open(path, setup);
|
|
opened.harness.resume();
|
|
const id = (await opened.root.submit({ type: "input", content: "hi" }, context)).id;
|
|
await reached.promise;
|
|
const taskId = await runTaskId(opened.harness);
|
|
await opened.harness.close(context);
|
|
|
|
opened = await open(path, setup);
|
|
expect(await checkpoint(opened.harness, taskId)).toEqual({ phase: "prepare", attempt: 1 });
|
|
opened.harness.resume();
|
|
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
|
|
expect((await allEntries(opened.root)).map((entry) => entry.kind)).toEqual([
|
|
"pi.user",
|
|
"pi.system",
|
|
"pi.assistant",
|
|
]);
|
|
await opened.harness.close(context);
|
|
});
|
|
|
|
it("resends a request interrupted before any partial without repeating preparation", async () => {
|
|
const path = await sqlitePath();
|
|
const setup = chatSetup();
|
|
addSection(setup.registry, "preamble", () => "p", { tag: false });
|
|
const reached = deferred();
|
|
const sent: string[][] = [];
|
|
const timeouts: (number | undefined)[] = [];
|
|
setup.faux.setResponses([
|
|
async (_context, options) => {
|
|
reached.resolve();
|
|
return aborted(options!.signal!);
|
|
},
|
|
(request, options) => {
|
|
sent.push(request.messages.map((message) => message.role));
|
|
timeouts.push(options?.timeoutMs);
|
|
return fauxAssistantMessage("answer");
|
|
},
|
|
]);
|
|
let opened = await open(path, setup);
|
|
setup.settings.stream = { timeoutMs: 1234 };
|
|
opened.harness.resume();
|
|
const id = (await opened.root.submit({ type: "input", content: "hi" }, context)).id;
|
|
await reached.promise;
|
|
const taskId = await runTaskId(opened.harness);
|
|
await opened.harness.close(context);
|
|
|
|
opened = await open(path, setup);
|
|
expect(await checkpoint(opened.harness, taskId)).toMatchObject({
|
|
phase: "request",
|
|
attempt: 1,
|
|
thinkingLevel: "off",
|
|
streamOptions: { timeoutMs: 1234 },
|
|
});
|
|
// The resend uses the pinned request, not options changed after preparation.
|
|
setup.settings.stream = { timeoutMs: 999 };
|
|
expect(await opened.harness.snapshot(LiveDoc, opened.root.id, context)).toMatchObject({
|
|
generation: { attempt: 1 },
|
|
});
|
|
opened.harness.resume();
|
|
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
|
|
expect(sent).toEqual([["system", "user"]]);
|
|
expect(timeouts).toEqual([1234]);
|
|
expect((await allEntries(opened.root)).map((entry) => entry.kind)).toEqual([
|
|
"pi.user",
|
|
"pi.system",
|
|
"pi.assistant",
|
|
]);
|
|
await opened.harness.close(context);
|
|
});
|
|
|
|
it("converts a committed partial into an aborted entry and resends the same messages", async () => {
|
|
const path = await sqlitePath();
|
|
const slow = chatSetup({ tokensPerSecond: 20, tokenSize: { min: 1, max: 1 } });
|
|
slow.faux.setResponses([fauxAssistantMessage("z".repeat(400))]);
|
|
let opened = await open(path, slow);
|
|
opened.harness.resume();
|
|
const watch = (await opened.harness.watchDoc(LiveDoc, opened.root.id, context))!;
|
|
let watched: string | undefined;
|
|
watch.start(
|
|
onValue(async ({ value }) => {
|
|
watched = textOf(value.generation?.message as Message) ?? watched;
|
|
}),
|
|
);
|
|
const id = (await opened.root.submit({ type: "input", content: "hi" }, context)).id;
|
|
await waitFor(() => watched !== undefined);
|
|
await opened.harness.close(context);
|
|
|
|
const setup = chatSetup();
|
|
const sent: string[][] = [];
|
|
setup.faux.setResponses([
|
|
(request) => {
|
|
sent.push(request.messages.map((message) => message.role));
|
|
return fauxAssistantMessage("answer");
|
|
},
|
|
]);
|
|
opened = await open(path, setup);
|
|
const stored = await opened.harness.snapshot(LiveDoc, opened.root.id, context);
|
|
const partial = textOf(stored?.generation?.message as Message)!;
|
|
// Everything observers saw before the crash is durable.
|
|
expect(partial.startsWith(watched!)).toBe(true);
|
|
opened.harness.resume();
|
|
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
|
|
expect(sent).toEqual([["user"]]);
|
|
const entries = await allEntries(opened.root);
|
|
expect(entries.map((entry) => entry.kind)).toEqual(["pi.user", "pi.assistant", "pi.assistant"]);
|
|
const converted = entries[1]!.model![0] as AssistantMessage;
|
|
expect(converted.stopReason).toBe("aborted");
|
|
expect(textOf(converted)).toBe(partial);
|
|
expect(await opened.harness.snapshot(LiveDoc, opened.root.id, context)).toEqual({});
|
|
await opened.harness.close(context);
|
|
});
|
|
|
|
it("resumes a retry backoff after reopen", async () => {
|
|
const path = await sqlitePath();
|
|
const setup = chatSetup();
|
|
let now = 1_000;
|
|
setup.now = () => now;
|
|
setup.faux.setResponses([
|
|
fauxAssistantMessage([], { stopReason: "error", errorMessage: "503 Service Unavailable" }),
|
|
fauxAssistantMessage("recovered"),
|
|
]);
|
|
let opened = await open(path, setup);
|
|
opened.harness.resume();
|
|
setup.settings.retry = { enabled: true, maxRetries: 2, baseDelayMs: 60_000 };
|
|
const id = (await opened.root.submit({ type: "input", content: "hi" }, context)).id;
|
|
await waitFor(async () => {
|
|
const live = await opened.harness.snapshot(LiveDoc, opened.root.id, context);
|
|
return live?.generation?.retry !== undefined;
|
|
});
|
|
const taskId = await runTaskId(opened.harness);
|
|
await opened.harness.close(context);
|
|
|
|
opened = await open(path, setup);
|
|
expect((await opened.harness.getTask(taskId, context))?.state).toEqual({
|
|
status: "sleeping",
|
|
checkpoint: { phase: "retry", attempt: 1, until: 61_000 },
|
|
until: 61_000,
|
|
});
|
|
await eventually(() => !opened.harness.wake().active);
|
|
expect(opened.harness.wake()).toEqual({ active: false, at: 61_000 });
|
|
// An early wake retries nothing.
|
|
opened.harness.resume();
|
|
await new Promise((resolve) => setTimeout(resolve, 20));
|
|
expect(setup.faux.state.callCount).toBe(1);
|
|
expect(await opened.harness.snapshot(LiveDoc, opened.root.id, context)).toEqual({
|
|
run: { taskId, inputs: [id], turns: 0, first: expect.any(Number) },
|
|
generation: { attempt: 1, retry: { at: 61_000, error: "503 Service Unavailable" } },
|
|
});
|
|
now = 61_000;
|
|
opened.harness.resume();
|
|
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
|
|
expect((await allEntries(opened.root)).map((entry) => entry.kind)).toEqual([
|
|
"pi.user",
|
|
"pi.assistant",
|
|
"pi.assistant",
|
|
]);
|
|
await opened.harness.close(context);
|
|
});
|
|
|
|
it("sleeps again for a retry stored as a running phase by an earlier version", async () => {
|
|
const path = await sqlitePath();
|
|
const setup = chatSetup();
|
|
let now = 1_000;
|
|
setup.now = () => now;
|
|
setup.faux.setResponses([
|
|
fauxAssistantMessage([], { stopReason: "error", errorMessage: "503 Service Unavailable" }),
|
|
fauxAssistantMessage("recovered"),
|
|
]);
|
|
let opened = await open(path, setup);
|
|
opened.harness.resume();
|
|
setup.settings.retry = { enabled: true, maxRetries: 2, baseDelayMs: 60_000 };
|
|
const id = (await opened.root.submit({ type: "input", content: "hi" }, context)).id;
|
|
await waitFor(
|
|
async () =>
|
|
(await opened.harness.getTask(await runTaskId(opened.harness), context))?.state.status === "sleeping",
|
|
);
|
|
const taskId = await runTaskId(opened.harness);
|
|
await opened.harness.close(context);
|
|
// Earlier versions slept inside a running `retry` phase.
|
|
const database = new DatabaseSync(path);
|
|
database
|
|
.prepare(
|
|
"UPDATE tasks SET status = 'running', record = json_set(json_remove(record, '$.state.until'), '$.state.status', 'running') WHERE id = ?",
|
|
)
|
|
.run(taskId);
|
|
database.close();
|
|
|
|
opened = await open(path, setup);
|
|
opened.harness.resume();
|
|
await waitFor(async () => (await opened.harness.getTask(taskId, context))?.state.status === "sleeping");
|
|
expect((await opened.harness.getTask(taskId, context))?.state).toMatchObject({
|
|
status: "sleeping",
|
|
until: 61_000,
|
|
});
|
|
expect(setup.faux.state.callCount).toBe(1);
|
|
now = 61_000;
|
|
opened.harness.resume();
|
|
expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
|
|
expect(setup.faux.state.callCount).toBe(2);
|
|
await opened.harness.close(context);
|
|
});
|
|
|
|
it("resumes polling a deferred response after reopen", async () => {
|
|
const path = await sqlitePath();
|
|
const setup = chatSetup({ deferred: { pollAfterMs: 60_000 } });
|
|
let now = 1_000;
|
|
setup.now = () => now;
|
|
setup.faux.setResponses([fauxAssistantMessage("deferred answer")]);
|
|
let opened = await open(path, setup);
|
|
opened.harness.resume();
|
|
setup.settings.stream = { deferred: true };
|
|
const id = (await opened.root.submit({ type: "input", content: "hi" }, context)).id;
|
|
await waitFor(async () => {
|
|
const live = await opened.harness.snapshot(LiveDoc, opened.root.id, context);
|
|
return live?.generation?.deferred !== undefined;
|
|
});
|
|
const taskId = await runTaskId(opened.harness);
|
|
await opened.harness.close(context);
|
|
|
|
opened = await open(path, setup);
|
|
expect(await checkpoint(opened.harness, taskId)).toMatchObject({ phase: "poll", attempt: 1, pollAt: 61_000 });
|
|
expect((await opened.harness.getTask(taskId, context))?.state).toMatchObject({
|
|
status: "sleeping",
|
|
until: 61_000,
|
|
});
|
|
await eventually(() => !opened.harness.wake().active);
|
|
expect(opened.harness.wake()).toEqual({ active: false, at: 61_000 });
|
|
now = 61_000;
|
|
opened.harness.resume();
|
|
const settled = await (await opened.harness.submission(id, context))!.wait(context);
|
|
if (settled.status !== "done" || settled.type !== "input") throw new Error(`Unexpected ${settled.status}`);
|
|
const answer = await opened.root.commit((tx) => tx.entry(AssistantEntry, settled.answer), context);
|
|
expect(textOf(answer?.model?.[0])).toBe("deferred answer");
|
|
expect(setup.faux.state.deferredFetchCount).toBe(1);
|
|
await opened.harness.close(context);
|
|
});
|
|
|
|
it("polls a checkpoint stored without a thinking level and records none", async () => {
|
|
const path = await sqlitePath();
|
|
const setup = chatSetup({ deferred: { pollAfterMs: 60_000 } });
|
|
let now = 1_000;
|
|
setup.now = () => now;
|
|
setup.faux.setResponses([fauxAssistantMessage("deferred answer")]);
|
|
let opened = await open(path, setup);
|
|
await opened.root.configure({ thinkingLevel: "high" }, context);
|
|
opened.harness.resume();
|
|
setup.settings.stream = { deferred: true };
|
|
const id = (await opened.root.submit({ type: "input", content: "hi" }, context)).id;
|
|
await waitFor(async () => {
|
|
const live = await opened.harness.snapshot(LiveDoc, opened.root.id, context);
|
|
return live?.generation?.deferred !== undefined;
|
|
});
|
|
const taskId = await runTaskId(opened.harness);
|
|
expect(await checkpoint(opened.harness, taskId)).toMatchObject({ phase: "poll", thinkingLevel: "high" });
|
|
await opened.harness.close(context);
|
|
// Earlier versions stored poll checkpoints without the field.
|
|
const database = new DatabaseSync(path);
|
|
database
|
|
.prepare("UPDATE tasks SET record = json_remove(record, '$.state.checkpoint.thinkingLevel') WHERE id = ?")
|
|
.run(taskId);
|
|
database.close();
|
|
|
|
opened = await open(path, setup);
|
|
expect(await checkpoint(opened.harness, taskId)).not.toHaveProperty("thinkingLevel");
|
|
now = 61_000;
|
|
opened.harness.resume();
|
|
const settled = await (await opened.harness.submission(id, context))!.wait(context);
|
|
if (settled.status !== "done" || settled.type !== "input") throw new Error(`Unexpected ${settled.status}`);
|
|
const answer = (await opened.root.commit((tx) => tx.entry(AssistantEntry, settled.answer), context))?.model?.[0];
|
|
expect(textOf(answer)).toBe("deferred answer");
|
|
expect(answer).not.toHaveProperty("thinkingLevel");
|
|
await opened.harness.close(context);
|
|
});
|
|
|
|
it("fails no_model when the pinned model is gone after reopen, in request and in poll", async () => {
|
|
const path = await sqlitePath();
|
|
const setup = chatSetup({ deferred: { pollAfterMs: 60_000 } });
|
|
let now = 1_000;
|
|
setup.now = () => now;
|
|
const busy = unanswered();
|
|
setup.faux.setResponses([busy.step, fauxAssistantMessage("deferred")]);
|
|
let opened = await open(path, setup);
|
|
opened.harness.resume();
|
|
const requesting = (await opened.root.submit({ type: "input", content: "one" }, context)).id;
|
|
await busy.reached;
|
|
await opened.harness.close(context);
|
|
|
|
// Reopened without the faux provider: the request's pinned model is unknown.
|
|
const empty: ChatSetup = { ...setup, models: createModels() };
|
|
opened = await open(path, empty);
|
|
opened.harness.resume();
|
|
expect(await (await opened.harness.submission(requesting, context))!.wait(context)).toMatchObject({
|
|
status: "unanswered",
|
|
reason: "no_model",
|
|
});
|
|
await opened.harness.close(context);
|
|
|
|
opened = await open(path, setup);
|
|
setup.settings.stream = { deferred: true };
|
|
opened.harness.resume();
|
|
const polling = (await opened.root.submit({ type: "input", content: "two" }, context)).id;
|
|
await waitFor(async () => {
|
|
const live = await opened.harness.snapshot(LiveDoc, opened.root.id, context);
|
|
return live?.generation?.deferred !== undefined;
|
|
});
|
|
await opened.harness.close(context);
|
|
|
|
// The poll sleeps until it is due; then the model is resolved.
|
|
now = 61_000;
|
|
opened = await open(path, empty);
|
|
opened.harness.resume();
|
|
expect(await (await opened.harness.submission(polling, context))!.wait(context)).toMatchObject({
|
|
status: "unanswered",
|
|
reason: "no_model",
|
|
});
|
|
await opened.harness.close(context);
|
|
});
|
|
|
|
it("runs a print-style turn and reads the durable answer after reopen", async () => {
|
|
const path = await sqlitePath();
|
|
const setup = chatSetup();
|
|
addSection(setup.registry, "preamble", () => "You are terse.", { tag: false });
|
|
setup.faux.setResponses([fauxAssistantMessage("42")]);
|
|
let opened = await open(path, setup);
|
|
opened.harness.resume();
|
|
const submission = await opened.root.submit({ type: "input", content: "answer?" }, context);
|
|
const settled = await submission.wait(context);
|
|
if (settled.status !== "done" || settled.type !== "input") throw new Error(`Unexpected ${settled.status}`);
|
|
const answer = await opened.root.commit((tx) => tx.entry(AssistantEntry, settled.answer), context);
|
|
expect(textOf(answer?.model?.[0])).toBe("42");
|
|
await opened.harness.close(context);
|
|
|
|
opened = await open(path, setup);
|
|
expect(await (await opened.harness.submission(submission.id, context))!.status(context)).toEqual(settled);
|
|
expect(await opened.root.commit((tx) => tx.entry(AssistantEntry, settled.answer), context)).toEqual(answer);
|
|
await opened.harness.close(context);
|
|
});
|
|
});
|