ai-chatbot/lib/canvas/capabilities/task-store.ts

import { and, eq, inArray, lt, sql } from "drizzle-orm";

import type { CanvasCapabilityClientCommand } from "@/lib/canvas/capabilities/client-command";
import { db } from "@/lib/db";
import { canvasCapabilityTask } from "@/lib/db/schema";

const STALE_RUNNING_TASK_MS = 15 * 60 * 1000;
const INTERRUPTED_TASK_MESSAGE = "任务执行中断，结果尚未确认。请先查看画布和历史记录，或联系管理员核对原任务，避免重复处理。";

export type CanvasCapabilityTaskResult = Record<string, unknown>;

export async function getCanvasCapabilityTask(params: {
  requestId: string;
  userId: string;
  projectId: string;
}) {
  const [task] = await db
    .select()
    .from(canvasCapabilityTask)
    .where(
      and(
        eq(canvasCapabilityTask.requestId, params.requestId),
        eq(canvasCapabilityTask.userId, params.userId),
        eq(canvasCapabilityTask.projectId, params.projectId)
      )
    )
    .limit(1);
  if (!task) return null;
  const staleBefore = new Date(Date.now() - STALE_RUNNING_TASK_MS);
  if (["pending", "running"].includes(task.status) && task.updatedAt < staleBefore) {
    const [interrupted] = await db.update(canvasCapabilityTask).set({
      status: "failed",
      error: INTERRUPTED_TASK_MESSAGE,
      completedAt: new Date(),
      updatedAt: new Date(),
    }).where(and(
      eq(canvasCapabilityTask.id, task.id),
      inArray(canvasCapabilityTask.status, ["pending", "running"]),
      lt(canvasCapabilityTask.updatedAt, staleBefore),
    )).returning();
    if (interrupted) return interrupted;
    const [latest] = await db.select().from(canvasCapabilityTask)
      .where(eq(canvasCapabilityTask.id, task.id)).limit(1);
    return latest || task;
  }
  return task;
}

export async function createOrRecoverCanvasCapabilityTask(params: {
  command: CanvasCapabilityClientCommand;
  userId: string;
  projectId: string;
}) {
  await db
    .insert(canvasCapabilityTask)
    .values({
      requestId: params.command.requestId,
      userId: params.userId,
      projectId: params.projectId,
      capabilityId: params.command.capabilityId,
      source: params.command.source,
      request: params.command,
      updatedAt: new Date(),
    })
    .onConflictDoNothing({ target: canvasCapabilityTask.requestId });

  const taskLookup = {
    requestId: params.command.requestId,
    userId: params.userId,
    projectId: params.projectId,
  };
  const task = await getCanvasCapabilityTask(taskLookup);
  if (!task) {
    throw new Error("无法创建画布能力任务");
  }

  return task;
}

export async function heartbeatCanvasCapabilityTask(taskId: string) {
  await db.update(canvasCapabilityTask).set({ updatedAt: new Date() })
    .where(and(eq(canvasCapabilityTask.id, taskId), inArray(canvasCapabilityTask.status, ["pending", "running"])));
}

export async function markCanvasCapabilityTaskRunning(taskId: string) {
  const now = new Date();
  const [task] = await db
    .update(canvasCapabilityTask)
    .set({
      status: "running",
      attempts: sql`${canvasCapabilityTask.attempts} + 1`,
      startedAt: now,
      completedAt: null,
      error: null,
      updatedAt: now,
    })
    .where(and(eq(canvasCapabilityTask.id, taskId), eq(canvasCapabilityTask.status, "pending")))
    .returning();

  return task || null;
}

export async function markCanvasCapabilityTaskSucceeded(
  taskId: string,
  result: CanvasCapabilityTaskResult
) {
  const now = new Date();
  const [task] = await db
    .update(canvasCapabilityTask)
    .set({
      status: "succeeded",
      result,
      error: null,
      completedAt: now,
      updatedAt: now,
    })
    .where(eq(canvasCapabilityTask.id, taskId))
    .returning();
  return task || null;
}

export async function markCanvasCapabilityTaskFailed(
  taskId: string,
  error: unknown
) {
  const now = new Date();
  const message = error instanceof Error ? error.message : String(error || "执行失败");
  const [task] = await db
    .update(canvasCapabilityTask)
    .set({
      status: "failed",
      error: message.slice(0, 4000),
      completedAt: now,
      updatedAt: now,
    })
    .where(eq(canvasCapabilityTask.id, taskId))
    .returning();
  return task || null;
}