evOo-local/packages/adapters/src/autopilot.postgres.test.ts

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<void>;
  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" });
  });
});