import { randomUUID } from "node:crypto"; import type { AutopilotBlock, MessageBlock } from "@rakazo/contracts"; import type { PrismaClient } from "@rakazo/db"; import { bootstrapUserSpace, createDb, ensureMainConversation } from "@rakazo/db"; import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; import { checkpointAutopilot, continueAutopilotAfterRun, endAutopilot, findOpenAutopilot, resumeAutopilot, startAutopilot, } from "./autopilot.js"; const databaseUrl = process.env.DATABASE_URL; const describePostgres = process.env.VERIFY_DATABASE && databaseUrl ? describe.sequential : describe.skip; const signupEnv = { signupsEnabled: "true", signupAllowlist: "" }; describePostgres("autopilot (PostgreSQL)", () => { let prisma: PrismaClient; let close: () => Promise; const userIds: string[] = []; beforeAll(() => { const db = createDb(databaseUrl!); prisma = db.prisma; close = async () => { await db.prisma.$disconnect(); await db.pool.end(); }; }); afterAll(async () => { if (!prisma) return; const memberships = await prisma.member.findMany({ where: { userId: { in: userIds } }, select: { organizationId: true }, }); await prisma.organization.deleteMany({ where: { id: { in: memberships.map((member) => member.organizationId) } }, }); await prisma.user.deleteMany({ where: { id: { in: userIds } } }); await close(); }); async function newOwner() { const id = `ap${randomUUID().replaceAll("-", "")}`; userIds.push(id); await prisma.user.create({ data: { id, name: "Ada", email: `${id}@rakazo.test` } }); await bootstrapUserSpace(prisma, { id }, signupEnv, { claimDeploymentOwner: false }); const conversation = await ensureMainConversation(prisma, id); const enqueue = vi.fn(async () => undefined); const deps = { prisma, events: { notify: vi.fn(async () => undefined) }, jobs: { enqueue }, }; const owner = { spaceId: conversation.spaceId, userId: id, botId: conversation.leadBotId, threadId: conversation.threadId, }; return { deps, owner, enqueue }; } async function lines(threadId: string) { const messages = await prisma.message.findMany({ where: { threadId }, orderBy: { seq: "asc" }, select: { blocks: true }, }); return messages.flatMap((message) => (message.blocks as MessageBlock[]).filter( (block): block is AutopilotBlock => block.kind === "autopilot", ), ); } async function finishQueuedRuns(threadId: string) { await prisma.run.updateMany({ where: { threadId, status: "queued" }, data: { status: "completed" }, }); } it("works step by step until it finishes, then takes its approvals away", async () => { const { deps, owner, enqueue } = await newOwner(); const started = await startAutopilot(deps, owner, { goal: "Tidy up the tickets", allow: ["github_commit", "create_space"], }); if (!("ok" in started)) throw new Error(started.error); const rules = await prisma.actionApprovalRule.findMany({ where: { autopilotId: started.autopilot.id }, }); // A security-boundary tool is never granted in advance. expect(rules.map((rule) => rule.matchValue)).toEqual(["github_commit"]); for (let step = 1; step <= 4; step += 1) { await checkpointAutopilot(deps, owner, { summary: `Ticket ${step} tidied` }); await continueAutopilotAfterRun(deps, owner, { checkpointed: true, failed: false }); await finishQueuedRuns(owner.threadId); } const runs = await prisma.run.findMany({ where: { threadId: owner.threadId, trigger: "autopilot" }, include: { task: true }, orderBy: { createdAt: "asc" }, }); expect(runs).toHaveLength(4); expect(runs.map((run) => run.clientNonce)).toEqual( [1, 2, 3, 4].map((n) => `autopilot:${started.autopilot.id}:${n}`), ); expect(runs[3]?.task.prompt).toContain("Last checkpoint: Ticket 4 tidied"); expect(enqueue).toHaveBeenCalledTimes(4); expect( await endAutopilot(deps, started.autopilot.id, { status: "finished", result: "Done." }), ).toBe(true); const row = await prisma.autopilot.findUniqueOrThrow({ where: { id: started.autopilot.id } }); expect(row).toMatchObject({ status: "finished", stepCount: 4, result: "Done." }); expect( await prisma.actionApprovalRule.count({ where: { autopilotId: started.autopilot.id } }), ).toBe(0); const timeline = await lines(owner.threadId); expect(timeline.map((line) => [line.event, Boolean(line.ended)])).toEqual([ ["autopilot.started", true], ["autopilot.finished", false], ]); // The start line carries the live progress. expect(timeline[0]?.steps).toBe(4); expect(timeline[0]?.checkpoints?.map((checkpoint) => checkpoint.text)).toEqual( [1, 2, 3, 4].map((step) => `Ticket ${step} tidied`), ); // Nothing runs after the end. await continueAutopilotAfterRun(deps, owner, { checkpointed: false, failed: false }); expect(enqueue).toHaveBeenCalledTimes(4); }); it("pauses after three runs without progress and goes on when the user continues", async () => { const { deps, owner, enqueue } = await newOwner(); const started = await startAutopilot(deps, owner, { goal: "Hunt bugs" }); if (!("ok" in started)) throw new Error(started.error); for (let run = 0; run < 3; run += 1) { await continueAutopilotAfterRun(deps, owner, { checkpointed: false, failed: false }); await finishQueuedRuns(owner.threadId); } expect(enqueue).toHaveBeenCalledTimes(2); const paused = await findOpenAutopilot(prisma, owner); expect(paused).toMatchObject({ status: "paused" }); expect(paused?.pauseReason).toContain("No progress in 3 runs"); expect((await lines(owner.threadId)).at(-1)).toMatchObject({ event: "autopilot.paused" }); expect(await resumeAutopilot(deps, paused!.id)).toBe(true); expect(enqueue).toHaveBeenCalledTimes(3); expect(await findOpenAutopilot(prisma, owner)).toMatchObject({ status: "active", idleCount: 0, }); expect((await lines(owner.threadId)).at(-1)).toMatchObject({ event: "autopilot.paused", ended: true, }); }); it("runs one autopilot per space, ends at its time and stops at once", async () => { const { deps, owner, enqueue } = await newOwner(); const first = await startAutopilot(deps, owner, { goal: "Work through the night", until: new Date(Date.now() + 60_000), }); if (!("ok" in first)) throw new Error(first.error); const elsewhere = await prisma.thread.findFirstOrThrow({ where: { spaceId: owner.spaceId, id: { not: owner.threadId } }, }); expect( await startAutopilot(deps, { ...owner, threadId: elsewhere.id }, { goal: "Second" }), ).toEqual({ error: expect.stringContaining("Stop it first") }); await prisma.autopilot.update({ where: { id: first.autopilot.id }, data: { until: new Date(Date.now() - 1000), lastCheckpoint: "Half the tickets" }, }); await continueAutopilotAfterRun(deps, owner, { checkpointed: true, failed: false }); expect(enqueue).not.toHaveBeenCalled(); expect( await prisma.autopilot.findUniqueOrThrow({ where: { id: first.autopilot.id } }), ).toMatchObject({ status: "finished", result: "Half the tickets" }); const second = await startAutopilot(deps, owner, { goal: "Again" }); if (!("ok" in second)) throw new Error(second.error); expect(await endAutopilot(deps, second.autopilot.id, { status: "stopped" })).toBe(true); expect(await endAutopilot(deps, second.autopilot.id, { status: "stopped" })).toBe(false); expect((await lines(owner.threadId)).at(-1)).toMatchObject({ event: "autopilot.stopped" }); }); });