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
290 lines
12 KiB
TypeScript
290 lines
12 KiB
TypeScript
import { mkdtemp, rm } from "node:fs/promises";
|
|
import { tmpdir } from "node:os";
|
|
import { join } from "node:path";
|
|
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 } 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-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(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 checkpoint(opened.harness, taskId)).toEqual({ phase: "retry", attempt: 1, until: 61_000 });
|
|
expect(await opened.harness.snapshot(LiveDoc, opened.root.id, context)).toEqual({
|
|
run: { taskId, inputs: [id] },
|
|
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("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 });
|
|
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("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 } });
|
|
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);
|
|
|
|
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);
|
|
});
|
|
});
|