ai-chatbot/lib/ai/canvas-agent-runtime/store.ts

import { canvasAgentActivities, canvasAgentProgressLabel } from "./progress";
import { and, desc, eq, gt, lt, sql } from "drizzle-orm";
import { db } from "@/lib/db";
import { canvasAgentRun, canvasProject, chat } from "@/lib/db/schema";
import {
  AgentRunControlChangedError,
  type AgentRunSnapshot,
  type AgentRunView,
  type CanvasAgentRunStatus,
} from "./protocol";
import { inferCanvasTaskRequirements } from "./requirements";
import { inferVideoTaskRequirements } from "./video-protocol";

const LEASE_STALE_MS = 90_000;
type Run = typeof canvasAgentRun.$inferSelect;
export class AgentRunError extends Error {
  status: number;
  constructor(message: string, status = 409) {
    super(message);
    this.status = status;
  }
}
export async function assertAgentRunOwnership(userId: string, projectId: string, chatId: string) {
  const [ownedProject] = await db
    .select({ id: canvasProject.id })
    .from(canvasProject)
    .where(and(eq(canvasProject.id, projectId), eq(canvasProject.userId, userId)))
    .limit(1);
  const [ownedChat] = await db
    .select({ id: chat.id })
    .from(chat)
    .where(and(eq(chat.id, chatId), eq(chat.userId, userId)))
    .limit(1);
  if (!ownedProject || !ownedChat) throw new AgentRunError("画布或对话不存在", 404);
}
export function agentRunView(run: Run): AgentRunView {
  const stale = run.status === "running" && Date.now() - run.updatedAt.getTime() > LEASE_STALE_MS;
  return {
    assistantMessageId: run.assistantMessageId,
    activities: canvasAgentActivities(run.snapshot),
    id: run.id,
    videoSettings: run.snapshot.videoSettings,
    chatId: run.chatId,
    projectId: run.projectId,
    status: stale ? "interrupted" : (run.status as CanvasAgentRunStatus),
    message: stale ? "任务连接已中断，已保留执行记录" : canvasAgentProgressLabel(run.snapshot.message),
    steps: run.snapshot.steps,
    toolCount: run.snapshot.calls.length,
    resultNodeIds: run.snapshot.resultNodeIds,
    inputs: [
      ...run.snapshot.inputs,
      ...run.pendingInputs.map((input) => ({ ...input, applied: false })),
    ],
    updatedAt: run.updatedAt.toISOString(),
    cancelRequested: run.cancelRequested,
    resumable:
      (stale || run.status !== "running") &&
      run.status !== "completed" &&
      !run.snapshot.calls.some((call) => call.status === "started" || call.uncertain),
  };
}
export async function getAgentRun(
  userId: string,
  chatId: string,
  projectId: string,
  runId?: string,
) {
  const [run] = await db
    .select()
    .from(canvasAgentRun)
    .where(
      and(
        eq(canvasAgentRun.userId, userId),
        eq(canvasAgentRun.chatId, chatId),
        eq(canvasAgentRun.projectId, projectId),
        ...(runId ? [eq(canvasAgentRun.id, runId)] : []),
      ),
    )
    .orderBy(desc(canvasAgentRun.createdAt))
    .limit(1);
  return run || null;
}
export async function acquireAgentRun(params: {
  id: string;
  userId: string;
  chatId: string;
  projectId: string;
  snapshot: AgentRunSnapshot;
  assistantMessageId: string;
  resumeRunId?: string;
}) {
  await assertAgentRunOwnership(params.userId, params.projectId, params.chatId);
  return db.transaction(async (tx) => {
    await tx.execute(
      sql`select pg_advisory_xact_lock(hashtext(${`canvas-agent:${params.userId}:${params.chatId}`}))`,
    );
    const staleBefore = new Date(Date.now() - LEASE_STALE_MS);
    const [active] = await tx
      .select({ id: canvasAgentRun.id })
      .from(canvasAgentRun)
      .where(
        and(
          eq(canvasAgentRun.userId, params.userId),
          eq(canvasAgentRun.chatId, params.chatId),
          eq(canvasAgentRun.status, "running"),
          gt(canvasAgentRun.updatedAt, staleBefore),
        ),
      )
      .limit(1);
    if (active) throw new AgentRunError("这个对话已有任务正在处理，请使用补充要求或停止操作。");
    const [existing] = await tx
      .select()
      .from(canvasAgentRun)
      .where(
        and(
          eq(canvasAgentRun.userId, params.userId),
          eq(canvasAgentRun.chatId, params.chatId),
          eq(canvasAgentRun.id, params.resumeRunId || params.id),
        ),
      )
      .limit(1);
    if (params.resumeRunId && (!existing || existing.projectId !== params.projectId)) {
      throw new AgentRunError("找不到需要继续的任务", 404);
    }
    if (
      existing &&
      (existing.snapshot.workspace || "canvas") !== (params.snapshot.workspace || "canvas")
    )
      throw new AgentRunError("任务与当前画布类型不一致", 400);
    if (existing && !params.resumeRunId)
      throw new AgentRunError("此消息已经处理，请刷新查看结果。");
    if (existing?.snapshot.calls.some((call) => call.status === "started" || call.uncertain)) {
      throw new AgentRunError(
        "有工具的执行结果尚未确认。为避免重复生图，请先核对画布结果，再创建新任务。",
      );
    }
    // Revoke expired workers before a new run acquires this conversation.
    await tx
      .update(canvasAgentRun)
      .set({ status: "interrupted" })
      .where(
        and(
          eq(canvasAgentRun.userId, params.userId),
          eq(canvasAgentRun.chatId, params.chatId),
          eq(canvasAgentRun.status, "running"),
          lt(canvasAgentRun.updatedAt, staleBefore),
        ),
      );
    const leaseId = crypto.randomUUID();
    const snapshot = existing?.snapshot ?? params.snapshot;
    if (existing) {
      if (existing.status === "completed") throw new AgentRunError("该任务已完成，无需重新执行。");
      snapshot.message = "正在从已完成的步骤继续";
      if (snapshot.workspace === "video" && params.snapshot.videoSettings) {
        snapshot.videoSettings = params.snapshot.videoSettings;
        snapshot.videoRequirements = {
          ...snapshot.videoRequirements!,
          duration: snapshot.videoSettings.duration,
          ratio: snapshot.videoSettings.ratio,
        };
      }
      snapshot.messages.push({ role: "user", content: params.snapshot.objective });
      snapshot.inputs.push({ id: params.id, text: params.snapshot.objective, applied: true });
      updateRequirements(snapshot, params.snapshot.objective);
      const [run] = await tx
        .update(canvasAgentRun)
        .set({
          status: "running",
          leaseId,
          cancelRequested: false,
          snapshot,
          assistantMessageId: params.assistantMessageId,
          updatedAt: new Date(),
        })
        .where(eq(canvasAgentRun.id, existing.id))
        .returning();
      return run;
    }
    const [run] = await tx
      .insert(canvasAgentRun)
      .values({
        id: params.id,
        userId: params.userId,
        chatId: params.chatId,
        projectId: params.projectId,
        snapshot,
        leaseId,
        assistantMessageId: params.assistantMessageId,
        createdAt: new Date(),
        updatedAt: new Date(),
      })
      .onConflictDoNothing({ target: canvasAgentRun.id })
      .returning();
    if (!run) throw new AgentRunError("此消息已经处理，请刷新查看结果。");
    return run;
  });
}
export async function checkpointAgentRun(
  run: Pick<Run, "id" | "leaseId">,
  snapshot: AgentRunSnapshot,
  status: CanvasAgentRunStatus = "running",
) {
  return db.transaction(async (tx) => {
    const [current] = await tx
      .select()
      .from(canvasAgentRun)
      .where(
        and(
          eq(canvasAgentRun.id, run.id),
          eq(canvasAgentRun.leaseId, run.leaseId),
          eq(canvasAgentRun.status, "running"),
        ),
      )
      .for("update")
      .limit(1);
    if (!current) throw new AgentRunError("任务执行权已变更，已停止本次执行。");
    if (status === "completed" && (current.cancelRequested || current.pendingInputs.length)) {
      throw new AgentRunControlChangedError("完成前收到新要求或停止请求");
    }
    const [saved] = await tx
      .update(canvasAgentRun)
      .set({ snapshot, status, updatedAt: new Date() })
      .where(eq(canvasAgentRun.id, run.id))
      .returning();
    return saved;
  });
}
export async function heartbeatAgentRun(run: Pick<Run, "id" | "leaseId">) {
  const [saved] = await db
    .update(canvasAgentRun)
    .set({ updatedAt: new Date() })
    .where(
      and(
        eq(canvasAgentRun.id, run.id),
        eq(canvasAgentRun.leaseId, run.leaseId),
        eq(canvasAgentRun.status, "running"),
      ),
    )
    .returning();
  return saved || null;
}
export async function consumeAgentRunInputs(
  run: Pick<Run, "id" | "leaseId">,
  snapshot: AgentRunSnapshot,
) {
  return db.transaction(async (tx) => {
    const [current] = await tx
      .select()
      .from(canvasAgentRun)
      .where(
        and(
          eq(canvasAgentRun.id, run.id),
          eq(canvasAgentRun.leaseId, run.leaseId),
          eq(canvasAgentRun.status, "running"),
        ),
      )
      .for("update")
      .limit(1);
    if (!current) throw new AgentRunError("任务已停止");
    if (current.cancelRequested) return { cancelled: true, inputs: [] };
    const inputs = current.pendingInputs;
    for (const input of inputs) {
      snapshot.messages.push({ role: "user", content: input.text });
      snapshot.inputs.push({ ...input, applied: true });
      updateRequirements(snapshot, input.text);
    }
    if (inputs.length)
      await tx
        .update(canvasAgentRun)
        .set({ snapshot, pendingInputs: [], updatedAt: new Date() })
        .where(eq(canvasAgentRun.id, run.id));
    return { cancelled: false, inputs };
  });
}

function updateRequirements(snapshot: AgentRunSnapshot, text: string) {
  if (snapshot.workspace === "video") {
    const inferred = inferVideoTaskRequirements(text);
    const current = snapshot.videoRequirements!;
    snapshot.videoRequirements = {
      ...current,
      expectsVideo: /不要(?:生成|制作|执行).{0,4}视频/.test(text)
        ? false
        : current.expectsVideo || inferred.expectsVideo,
      expectsScript: current.expectsScript || inferred.expectsScript,
      ...(inferred.expectedVideoCount
        ? {
            expectedVideoCount: /再|补/.test(text)
              ? (current.expectedVideoCount || 1) + inferred.expectedVideoCount
              : inferred.expectedVideoCount,
          }
        : {}),
      ...(inferred.duration ? { duration: inferred.duration } : {}),
      ...(inferred.ratio ? { ratio: inferred.ratio } : {}),
    };
    if (snapshot.videoSettings) {
      if (inferred.duration && inferred.duration <= 120)
        snapshot.videoSettings.duration = inferred.duration;
      const ratios = ["16:9", "9:16", "1:1", "4:3", "3:4", "21:9"];
      if (inferred.ratio && ratios.includes(inferred.ratio))
        snapshot.videoSettings.ratio = inferred.ratio as typeof snapshot.videoSettings.ratio;
    }
    return;
  }
  const inferred = inferCanvasTaskRequirements(text, snapshot.requirements?.sourceCount);
  const current = snapshot.requirements ?? {};
  const plan = snapshot.calls.filter((call) => call.name === "planCanvasTask").at(-1)?.output;
  const previousCount = current.expectedImageCount ?? Number(plan?.expectedImageCount || 0);
  snapshot.requirements = {
    ...current,
    ...(inferred.aspectRatio ? { aspectRatio: inferred.aspectRatio } : {}),
    ...(inferred.expectedImageCount
      ? {
          expectedImageCount: /再(?:出|做|生成)|补(?:出|做)?/.test(text)
            ? previousCount + inferred.expectedImageCount
            : inferred.expectedImageCount,
        }
      : {}),
  };
}
export async function controlAgentRun(params: {
  userId: string;
  projectId: string;
  chatId: string;
  runId: string;
  action: "steer" | "cancel";
  inputId?: string;
  text?: string;
}) {
  return db.transaction(async (tx) => {
    const [run] = await tx
      .select()
      .from(canvasAgentRun)
      .where(
        and(
          eq(canvasAgentRun.id, params.runId),
          eq(canvasAgentRun.userId, params.userId),
          eq(canvasAgentRun.projectId, params.projectId),
          eq(canvasAgentRun.chatId, params.chatId),
        ),
      )
      .for("update")
      .limit(1);
    if (!run) throw new AgentRunError("任务不存在", 404);
    if (agentRunView(run).status !== "running")
      throw new AgentRunError("任务已结束或中断，请继续任务后再补充。");
    if (params.action === "cancel") {
      const [updated] = await tx
        .update(canvasAgentRun)
        .set({ cancelRequested: true })
        .where(eq(canvasAgentRun.id, run.id))
        .returning();
      return agentRunView(updated);
    }
    if (run.cancelRequested) throw new AgentRunError("任务正在停止，不能再补充要求。");
    const inputId = params.inputId!;
    if (
      run.pendingInputs.some((input) => input.id === inputId) ||
      run.snapshot.inputs.some((input) => input.id === inputId)
    )
      return agentRunView(run);
    if (run.pendingInputs.length >= 10)
      throw new AgentRunError("待处理补充要求过多，请等待当前步骤完成。");
    const [updated] = await tx
      .update(canvasAgentRun)
      .set({ pendingInputs: [...run.pendingInputs, { id: inputId, text: params.text! }] })
      .where(eq(canvasAgentRun.id, run.id))
      .returning();
    return agentRunView(updated);
  });
}