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