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