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

import type { JobPublisher } from "@rakazo/adapter-kit";
import { runContinueJob } from "@rakazo/adapter-kit";
import type { AutopilotBlock, Autopilot as AutopilotView } from "@rakazo/contracts";
import { AUTOPILOT_LIMITS } from "@rakazo/contracts";
import {
  ACTIVE_RUN_STATUSES,
  advanceAutopilotGuard,
  autopilotContinuationPrompt,
  autopilotExpired,
  toolRequiresExplicitApproval,
} from "@rakazo/core";
import type { Prisma, PrismaClient, ThreadEvents } from "@rakazo/db";
import {
  appendEventInTransaction,
  endAutopilotLinesInTransaction,
  postAutopilotLineInTransaction,
  recordAutopilotCheckpointInTransaction,
} from "@rakazo/db";
import { getLogger } from "@rakazo/logging";

type AutopilotDeps = {
  prisma: PrismaClient;
  events: Pick<ThreadEvents, "notify">;
  jobs: Pick<JobPublisher, "enqueue">;
};

type AutopilotRow = NonNullable<Awaited<ReturnType<PrismaClient["autopilot"]["findFirst"]>>>;

/** Who the autopilot works for and where. */
export interface AutopilotOwner {
  spaceId: string;
  userId: string;
  botId: string;
  threadId: string;
}

const OPEN = ["active", "paused"];

function clip(value: string, max: number) {
  const flat = value.replace(/\s+/g, " ").trim();
  return flat.length > max ? `${flat.slice(0, max - 1)}…` : flat;
}

export function autopilotView(row: AutopilotRow): AutopilotView {
  return {
    id: row.id,
    botId: row.botId,
    threadId: row.threadId,
    goal: row.goal,
    focus: row.focus,
    until: row.until?.toISOString() ?? null,
    status: row.status as AutopilotView["status"],
    pauseReason: row.pauseReason,
    lastCheckpoint: row.lastCheckpoint,
    nextStep: row.nextStep,
    result: row.result,
    stepCount: row.stepCount,
    startedAt: row.startedAt.toISOString(),
    endedAt: row.endedAt?.toISOString() ?? null,
  };
}

async function notify(events: AutopilotDeps["events"], threadId: string, seq: number | null) {
  if (seq === null) return;
  await events.notify(threadId, seq).catch((error) => {
    getLogger().error("autopilot realtime notification", error);
  });
}

async function lockAutopilot(tx: Prisma.TransactionClient, id: string) {
  await tx.$queryRaw`SELECT id FROM autopilots WHERE id = ${id} FOR UPDATE`;
  return tx.autopilot.findUnique({ where: { id } });
}

/** The thread's open autopilot for this user, active or paused. */
export function findOpenAutopilot(
  prisma: PrismaClient | Prisma.TransactionClient,
  owner: Pick<AutopilotOwner, "userId" | "threadId">,
) {
  return prisma.autopilot.findFirst({
    where: { threadId: owner.threadId, userId: owner.userId, status: { in: OPEN } },
    orderBy: { startedAt: "desc" },
  });
}

/**
 * Start an autopilot, or give the open one in this thread a new order. One per user and space:
 * an open autopilot in another thread must be stopped first. Approvals granted for it become
 * rules that end with it.
 */
export async function startAutopilot(
  deps: AutopilotDeps,
  owner: AutopilotOwner,
  input: { goal: string; focus?: string; until?: Date; allow?: string[] },
): Promise<{ ok: true; autopilot: AutopilotRow } | { error: string }> {
  const goal = clip(input.goal, AUTOPILOT_LIMITS.goal);
  if (!goal) return { error: "goal is required." };
  const focus = input.focus?.trim() ? clip(input.focus, AUTOPILOT_LIMITS.focus) : null;
  if (input.until && input.until.getTime() <= Date.now()) {
    return { error: "until must be in the future." };
  }
  const allow = [...new Set((input.allow ?? []).map((name) => name.trim()).filter(Boolean))]
    .filter((name) => !toolRequiresExplicitApproval(name))
    .slice(0, 20);
  const result = await deps.prisma.$transaction(async (tx) => {
    // One open autopilot per user and space: serialize starts on the user's space row.
    await tx.$queryRaw`SELECT id FROM spaces WHERE id = ${owner.spaceId} FOR UPDATE`;
    const open = await tx.autopilot.findFirst({
      where: { spaceId: owner.spaceId, userId: owner.userId, status: { in: OPEN } },
    });
    if (open && open.threadId !== owner.threadId) {
      return { error: "Another autopilot is already running in this space. Stop it first." };
    }
    const data = {
      goal,
      focus,
      until: input.until ?? null,
      status: "active",
      pauseReason: null,
      idleCount: 0,
      failCount: 0,
    };
    const autopilot = open
      ? await tx.autopilot.update({ where: { id: open.id }, data: { ...data, botId: owner.botId } })
      : await tx.autopilot.create({ data: { ...owner, ...data } });
    if (open) {
      await tx.actionApprovalRule.deleteMany({ where: { autopilotId: open.id } });
      await endAutopilotLinesInTransaction(tx, { ...owner, autopilotId: open.id });
    }
    if (allow.length) {
      await tx.actionApprovalRule.createMany({
        data: allow.map((name) => ({
          spaceId: owner.spaceId,
          createdByUserId: owner.userId,
          effect: "always_allow",
          matchKind: "tool",
          matchValue: name,
          autopilotId: autopilot.id,
        })),
        // A standing rule of the user's own stays theirs and outlives the autopilot.
        skipDuplicates: true,
      });
    }
    const line = await postAutopilotLineInTransaction(tx, {
      ...owner,
      block: {
        kind: "autopilot",
        event: "autopilot.started",
        autopilotId: autopilot.id,
        goal,
        ...(focus ? { focus } : {}),
        ...(input.until ? { until: input.until.toISOString() } : {}),
        ...(allow.length ? { allowed: allow } : {}),
      },
    });
    const updated = await tx.autopilot.update({
      where: { id: autopilot.id },
      data: { startMessageId: line.messageId },
    });
    return { ok: true as const, autopilot: updated, seq: line.seq };
  });
  if ("error" in result) return result;
  await notify(deps.events, owner.threadId, result.seq);
  return { ok: true, autopilot: result.autopilot };
}

/** A step is done: remember where the work stands. */
export async function checkpointAutopilot(
  deps: Pick<AutopilotDeps, "prisma" | "events">,
  owner: Pick<AutopilotOwner, "userId" | "threadId">,
  input: { summary: string; next?: string },
): Promise<{ ok: true; steps: number } | { error: string }> {
  const open = await findOpenAutopilot(deps.prisma, owner);
  if (open?.status !== "active") return { error: "No autopilot is running in this thread." };
  const summary = clip(input.summary, AUTOPILOT_LIMITS.checkpoint);
  const recorded = await deps.prisma.$transaction(async (tx) => {
    const autopilot = await lockAutopilot(tx, open.id);
    if (autopilot?.status !== "active") return null;
    const updated = await tx.autopilot.update({
      where: { id: autopilot.id },
      data: {
        stepCount: { increment: 1 },
        idleCount: 0,
        lastCheckpoint: summary,
        nextStep: input.next?.trim() ? clip(input.next, AUTOPILOT_LIMITS.checkpoint) : null,
      },
    });
    const seq = autopilot.startMessageId
      ? await recordAutopilotCheckpointInTransaction(tx, {
          spaceId: autopilot.spaceId,
          threadId: autopilot.threadId,
          messageId: autopilot.startMessageId,
          autopilotId: autopilot.id,
          checkpoint: { text: summary, at: new Date().toISOString() },
          steps: updated.stepCount,
        })
      : null;
    return { steps: updated.stepCount, threadId: autopilot.threadId, seq };
  });
  if (!recorded) return { error: "No autopilot is running in this thread." };
  await notify(deps.events, recorded.threadId, recorded.seq);
  return { ok: true, steps: recorded.steps };
}

/** End the autopilot: finished with a result, or stopped by the user. */
export async function endAutopilot(
  deps: AutopilotDeps,
  autopilotId: string,
  outcome: { status: "finished"; result: string } | { status: "stopped" },
): Promise<boolean> {
  const ended = await deps.prisma.$transaction(async (tx) => {
    const autopilot = await lockAutopilot(tx, autopilotId);
    if (!autopilot || !OPEN.includes(autopilot.status)) return null;
    const endedAt = new Date();
    const result =
      outcome.status === "finished" ? clip(outcome.result, AUTOPILOT_LIMITS.result) : null;
    await tx.autopilot.update({
      where: { id: autopilot.id },
      data: { status: outcome.status, endedAt, result, pauseReason: null },
    });
    await tx.actionApprovalRule.deleteMany({ where: { autopilotId: autopilot.id } });
    await endAutopilotLinesInTransaction(tx, { ...autopilot, autopilotId: autopilot.id });
    const block: AutopilotBlock = {
      kind: "autopilot",
      event: outcome.status === "finished" ? "autopilot.finished" : "autopilot.stopped",
      autopilotId: autopilot.id,
      goal: autopilot.goal,
      steps: autopilot.stepCount,
      durationMs: endedAt.getTime() - autopilot.startedAt.getTime(),
      ...(result ? { text: result } : {}),
    };
    const line = await postAutopilotLineInTransaction(tx, { ...autopilot, block });
    return { threadId: autopilot.threadId, seq: line.seq };
  });
  if (!ended) return false;
  await notify(deps.events, ended.threadId, ended.seq);
  return true;
}

/** Hold the autopilot and say why; the user continues or stops it from the line. */
async function pauseAutopilotInTransaction(
  tx: Prisma.TransactionClient,
  autopilot: AutopilotRow,
  reason: string,
) {
  await tx.autopilot.update({
    where: { id: autopilot.id },
    data: { status: "paused", pauseReason: reason },
  });
  return postAutopilotLineInTransaction(tx, {
    ...autopilot,
    block: {
      kind: "autopilot",
      event: "autopilot.paused",
      autopilotId: autopilot.id,
      goal: autopilot.goal,
      steps: autopilot.stepCount,
      text: reason,
    },
  });
}

/** Let a paused autopilot go on, with a fresh idle count, and start its next run. */
export async function resumeAutopilot(deps: AutopilotDeps, autopilotId: string): Promise<boolean> {
  const resumed = await deps.prisma.$transaction(async (tx) => {
    const autopilot = await lockAutopilot(tx, autopilotId);
    if (autopilot?.status !== "paused") return null;
    const seq = await endPausedLines(tx, autopilot);
    const active = await tx.autopilot.update({
      where: { id: autopilot.id },
      data: { status: "active", pauseReason: null, idleCount: 0, failCount: 0 },
    });
    const runId = await createContinuationRun(tx, active);
    return { threadId: autopilot.threadId, seq, runId };
  });
  if (!resumed) return false;
  await notify(deps.events, resumed.threadId, resumed.seq);
  await enqueue(deps, resumed.runId);
  return true;
}

/** Only the pause lines lose their Continue control; the start line keeps its Stop. */
async function endPausedLines(tx: Prisma.TransactionClient, autopilot: AutopilotRow) {
  const messages = await tx.message.findMany({
    where: {
      threadId: autopilot.threadId,
      blocks: {
        array_contains: [
          { kind: "autopilot", autopilotId: autopilot.id, event: "autopilot.paused" },
        ],
      },
    },
    select: { id: true, role: true, botId: true, blocks: true },
  });
  let seq: number | null = null;
  for (const message of messages) {
    const blocks = (message.blocks as AutopilotBlock[]).map((block) =>
      block.kind === "autopilot" && block.event === "autopilot.paused" && !block.ended
        ? { ...block, ended: true }
        : block,
    );
    await tx.message.update({ where: { id: message.id }, data: { blocks } });
    const event = await appendEventInTransaction(tx, {
      spaceId: autopilot.spaceId,
      threadId: autopilot.threadId,
      botId: message.botId ?? autopilot.botId,
      type: "thread.message.updated",
      payload: { messageId: message.id, role: message.role, blocks },
    });
    seq = event.seq;
  }
  return seq;
}

async function enqueue(deps: AutopilotDeps, runId: string | null) {
  if (!runId) return;
  await deps.jobs.enqueue(runContinueJob(runId)).catch((error) => {
    getLogger().error("autopilot continuation enqueue", error);
  });
}

/**
 * The next run of an active autopilot, unless the bot is already busy in the thread. Each
 * continuation has its own nonce, so two workers finishing at once start it only once.
 */
async function createContinuationRun(
  tx: Prisma.TransactionClient,
  autopilot: AutopilotRow,
): Promise<string | null> {
  const busy = await tx.run.findFirst({
    where: {
      threadId: autopilot.threadId,
      botId: autopilot.botId,
      status: { in: [...ACTIVE_RUN_STATUSES] },
    },
    select: { id: true },
  });
  if (busy) return null;
  const continuation = autopilot.continuations + 1;
  const claimed = await tx.autopilot.updateMany({
    where: { id: autopilot.id, continuations: autopilot.continuations },
    data: { continuations: continuation },
  });
  if (claimed.count === 0) return null;
  const prompt = autopilotContinuationPrompt(autopilot);
  const task = await tx.task.create({
    data: {
      spaceId: autopilot.spaceId,
      botId: autopilot.botId,
      threadId: autopilot.threadId,
      userId: autopilot.userId,
      prompt,
      status: "queued",
    },
  });
  const run = await tx.run.create({
    data: {
      spaceId: autopilot.spaceId,
      botId: autopilot.botId,
      threadId: autopilot.threadId,
      taskId: task.id,
      userId: autopilot.userId,
      status: "queued",
      trigger: "autopilot",
      clientNonce: `autopilot:${autopilot.id}:${continuation}`,
    },
  });
  return run.id;
}

/** Start the next run of a just-started autopilot when nothing in the thread runs yet. */
export async function kickAutopilot(deps: AutopilotDeps, autopilotId: string): Promise<void> {
  const runId = await deps.prisma.$transaction(async (tx) => {
    const autopilot = await lockAutopilot(tx, autopilotId);
    return autopilot?.status === "active" ? createContinuationRun(tx, autopilot) : null;
  });
  await enqueue(deps, runId);
}

/**
 * After a run of the autopilot's bot ends in its thread: finish at the chosen time, pause when it
 * stopped making progress or keeps failing, and otherwise start the next run.
 */
export async function continueAutopilotAfterRun(
  deps: AutopilotDeps,
  run: { userId: string; threadId: string; botId: string },
  outcome: { checkpointed: boolean; failed: boolean; error?: string },
): Promise<{ paused: string } | null> {
  const open = await deps.prisma.autopilot.findFirst({
    where: {
      threadId: run.threadId,
      userId: run.userId,
      botId: run.botId,
      status: "active",
    },
    select: { id: true },
  });
  if (!open) return null;
  const next = await deps.prisma.$transaction(async (tx) => {
    const autopilot = await lockAutopilot(tx, open.id);
    if (autopilot?.status !== "active") return null;
    if (autopilotExpired(autopilot.until)) {
      return { expired: { id: autopilot.id, lastCheckpoint: autopilot.lastCheckpoint } };
    }
    const guard = advanceAutopilotGuard(autopilot, outcome);
    if (guard.pause) {
      const reason =
        guard.pause === "failing"
          ? `Runs keep failing${outcome.error ? `: ${clip(outcome.error, 200)}` : "."}`
          : `No progress in ${guard.idleCount} runs${autopilot.lastCheckpoint ? `. Last: ${autopilot.lastCheckpoint}` : "."}`;
      const line = await pauseAutopilotInTransaction(tx, autopilot, reason);
      return { threadId: autopilot.threadId, seq: line.seq, runId: null, paused: reason };
    }
    const updated = await tx.autopilot.update({
      where: { id: autopilot.id },
      data: { idleCount: guard.idleCount, failCount: guard.failCount },
    });
    return {
      threadId: autopilot.threadId,
      seq: null,
      runId: await createContinuationRun(tx, updated),
      paused: null,
    };
  });
  if (!next) return null;
  if ("expired" in next && next.expired) {
    const autopilot = next.expired;
    await endAutopilot(deps, autopilot.id, {
      status: "finished",
      result: autopilot.lastCheckpoint ?? "Time is up.",
    });
    return null;
  }
  await notify(deps.events, next.threadId, next.seq);
  await enqueue(deps, next.runId);
  return next.paused ? { paused: next.paused } : null;
}