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

import { and, eq, 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;
export const MAX_CANVAS_CAPABILITY_TASK_ATTEMPTS = 3;

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);
  return task || null;
}

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,
    })
    .onConflictDoNothing({ target: canvasCapabilityTask.requestId });

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

  const staleBefore = new Date(Date.now() - STALE_RUNNING_TASK_MS);
  if (
    task.status === "failed" &&
    task.attempts < MAX_CANVAS_CAPABILITY_TASK_ATTEMPTS
  ) {
    await db
      .update(canvasCapabilityTask)
      .set({
        status: "pending",
        error: null,
        completedAt: null,
        updatedAt: new Date(),
      })
      .where(eq(canvasCapabilityTask.id, task.id));
  } else if (
    task.status === "running" &&
    task.updatedAt < staleBefore &&
    task.attempts < MAX_CANVAS_CAPABILITY_TASK_ATTEMPTS
  ) {
    await db
      .update(canvasCapabilityTask)
      .set({
        status: "pending",
        error: "恢复超时的执行中任务",
        updatedAt: new Date(),
      })
      .where(
        and(
          eq(canvasCapabilityTask.id, task.id),
          lt(canvasCapabilityTask.updatedAt, staleBefore)
        )
      );
  } else if (task.status === "running" && task.updatedAt < staleBefore) {
    await db
      .update(canvasCapabilityTask)
      .set({
        status: "failed",
        error: "画布能力任务超过最大重试次数",
        completedAt: new Date(),
        updatedAt: new Date(),
      })
      .where(eq(canvasCapabilityTask.id, task.id));
  }

  task = (await getCanvasCapabilityTask(taskLookup)) || task;
  return task;
}

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(eq(canvasCapabilityTask.id, taskId))
    .returning();

  if (!task) throw new Error("画布能力任务不存在");
  return task;
}

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;
}