311 lines
9.6 KiB
TypeScript
311 lines
9.6 KiB
TypeScript
import { canResumeCanvasTask, hasRecoverableCanvasTask } from "./task-recovery";
|
|
import "server-only";
|
|
import { and, asc, desc, eq, gte, lt, sql } from "drizzle-orm";
|
|
import { db } from "@/lib/db";
|
|
import { canvasCapabilityTask as table, canvasProject } from "@/lib/db/schema";
|
|
import { type CanvasTaskRequest, type CanvasGenerationTask } from "./task-contract";
|
|
|
|
const capabilityId = "video-canvas.generation";
|
|
const scope = (userId: string, projectId: string) =>
|
|
and(
|
|
eq(table.userId, userId),
|
|
eq(table.projectId, projectId),
|
|
eq(table.capabilityId, capabilityId),
|
|
);
|
|
const serialize = (row: typeof table.$inferSelect): CanvasGenerationTask => ({
|
|
id: row.id,
|
|
requestId: row.requestId,
|
|
status: row.status,
|
|
createdAt: row.createdAt.toISOString(),
|
|
updatedAt: row.updatedAt.toISOString(),
|
|
attempts: row.attempts,
|
|
request: row.request as CanvasTaskRequest,
|
|
result: row.result as Record<string, any> | null,
|
|
error: row.error,
|
|
});
|
|
export async function ownsTaskProject(userId: string, projectId: string) {
|
|
return Boolean(
|
|
(
|
|
await db
|
|
.select({ id: canvasProject.id })
|
|
.from(canvasProject)
|
|
.where(and(eq(canvasProject.id, projectId), eq(canvasProject.userId, userId)))
|
|
.limit(1)
|
|
)[0],
|
|
);
|
|
}
|
|
export async function createCanvasGenerationTask(userId: string, request: CanvasTaskRequest) {
|
|
return db.transaction(async (tx) => {
|
|
await tx.execute(
|
|
sql`select pg_advisory_xact_lock(hashtext(${capabilityId}), hashtext(${request.projectId}))`,
|
|
);
|
|
const existing = (
|
|
await tx
|
|
.select()
|
|
.from(table)
|
|
.where(and(scope(userId, request.projectId), eq(table.requestId, request.requestId)))
|
|
.limit(1)
|
|
)[0];
|
|
if (existing) {
|
|
if (
|
|
JSON.stringify(existing.request) !== JSON.stringify(JSON.parse(JSON.stringify(request)))
|
|
) {
|
|
// JSONB property order is not stable. Compare the immutable logical request.
|
|
const old = existing.request as CanvasTaskRequest;
|
|
if (
|
|
old.node.id !== request.node.id ||
|
|
old.kind !== request.kind ||
|
|
JSON.stringify(sortObject(old.attachment)) !==
|
|
JSON.stringify(sortObject(request.attachment)) ||
|
|
JSON.stringify(sortObject(old.params)) !== JSON.stringify(sortObject(request.params))
|
|
)
|
|
throw Object.assign(Error("同一任务 ID 不能用于不同请求"), { status: 409 });
|
|
}
|
|
return serialize(existing);
|
|
}
|
|
const active = await tx
|
|
.select()
|
|
.from(table)
|
|
.where(and(scope(userId, request.projectId), sql`${table.status} in ('pending','running')`));
|
|
const duplicate = active.find(
|
|
(row) => (row.request as CanvasTaskRequest).node.id === request.node.id,
|
|
);
|
|
if (duplicate) return serialize(duplicate);
|
|
const interrupted = await tx
|
|
.select()
|
|
.from(table)
|
|
.where(and(scope(userId, request.projectId), eq(table.status, "interrupted")));
|
|
const recovering = interrupted.find(
|
|
(row) =>
|
|
(row.request as CanvasTaskRequest).node.id === request.node.id &&
|
|
hasRecoverableCanvasTask(serialize(row)),
|
|
);
|
|
if (recovering) return serialize(recovering);
|
|
|
|
if (active.length >= 40)
|
|
throw Object.assign(Error("当前画布排队任务已达 40 个，请等待完成后继续"), { status: 429 });
|
|
const [created] = await tx
|
|
.insert(table)
|
|
.values({
|
|
requestId: request.requestId,
|
|
userId,
|
|
projectId: request.projectId,
|
|
capabilityId,
|
|
source: "video-canvas",
|
|
request,
|
|
createdAt: new Date(),
|
|
updatedAt: new Date(),
|
|
})
|
|
.returning();
|
|
return serialize(created);
|
|
});
|
|
}
|
|
function sortObject(value: any): any {
|
|
return Array.isArray(value)
|
|
? value.map(sortObject)
|
|
: value && typeof value === "object"
|
|
? Object.fromEntries(
|
|
Object.keys(value)
|
|
.sort()
|
|
.map((key) => [key, sortObject(value[key])]),
|
|
)
|
|
: value;
|
|
}
|
|
export async function listCanvasGenerationTasks(userId: string, projectId: string) {
|
|
// A worker lost to process shutdown must never silently re-submit a paid request.
|
|
await db
|
|
.update(table)
|
|
.set({
|
|
status: "interrupted",
|
|
error:
|
|
"服务端执行连接中断，供应商任务可能仍在运行。请核实结果后再决定是否重新生成，系统不会自动重复提交。",
|
|
updatedAt: new Date(),
|
|
})
|
|
.where(
|
|
and(
|
|
scope(userId, projectId),
|
|
eq(table.status, "running"),
|
|
lt(table.updatedAt, new Date(Date.now() - 120000)),
|
|
),
|
|
);
|
|
const rows = await db
|
|
.select()
|
|
.from(table)
|
|
.where(scope(userId, projectId))
|
|
.orderBy(desc(table.createdAt))
|
|
.limit(240);
|
|
return rows.map(serialize).reverse();
|
|
}
|
|
export async function claimCanvasGenerationTask(userId: string, projectId: string) {
|
|
return db.transaction(async (tx) => {
|
|
await tx.execute(
|
|
sql`select pg_advisory_xact_lock(hashtext(${capabilityId}), hashtext(${projectId}))`,
|
|
);
|
|
if (
|
|
(
|
|
await tx
|
|
.select({ id: table.id })
|
|
.from(table)
|
|
.where(and(scope(userId, projectId), eq(table.status, "running")))
|
|
.limit(1)
|
|
).length
|
|
)
|
|
return null;
|
|
let [next] = await tx
|
|
.select()
|
|
.from(table)
|
|
.where(and(scope(userId, projectId), eq(table.status, "pending")))
|
|
.orderBy(asc(table.createdAt))
|
|
.limit(1);
|
|
if (!next) {
|
|
const interrupted = await tx
|
|
.select()
|
|
.from(table)
|
|
.where(
|
|
and(
|
|
scope(userId, projectId),
|
|
eq(table.status, "interrupted"),
|
|
gte(table.createdAt, new Date(Date.now() - 86400000)),
|
|
lt(table.attempts, 20),
|
|
),
|
|
)
|
|
.orderBy(asc(table.updatedAt));
|
|
next = interrupted.find((row) => canResumeCanvasTask(serialize(row)))!;
|
|
}
|
|
if (!next) return null;
|
|
const [claimed] = await tx
|
|
.update(table)
|
|
.set({
|
|
status: "running",
|
|
error: null,
|
|
startedAt: new Date(),
|
|
updatedAt: new Date(),
|
|
attempts: sql`${table.attempts}+1`,
|
|
})
|
|
.where(and(eq(table.id, next.id), eq(table.status, next.status)))
|
|
.returning();
|
|
return claimed ? serialize(claimed) : null;
|
|
});
|
|
}
|
|
export async function heartbeatCanvasGenerationTask(id: string, attempt?: number) {
|
|
await db
|
|
.update(table)
|
|
.set({ updatedAt: new Date() })
|
|
.where(
|
|
and(
|
|
eq(table.id, id),
|
|
eq(table.status, "running"),
|
|
attempt === undefined ? undefined : eq(table.attempts, attempt),
|
|
),
|
|
);
|
|
}
|
|
export async function finishCanvasGenerationTask(
|
|
id: string,
|
|
result?: Record<string, any>,
|
|
error?: string,
|
|
attempt?: number,
|
|
) {
|
|
// A late real result can recover an interrupted task, but cannot replace a newer request ID.
|
|
await db
|
|
.update(table)
|
|
.set({
|
|
status: result ? "succeeded" : "failed",
|
|
result: result || null,
|
|
error: error || null,
|
|
updatedAt: new Date(),
|
|
completedAt: new Date(),
|
|
})
|
|
.where(
|
|
and(
|
|
eq(table.id, id),
|
|
sql`${table.status} in ('running','interrupted')`,
|
|
attempt === undefined ? undefined : eq(table.attempts, attempt),
|
|
),
|
|
);
|
|
}
|
|
export async function cancelCanvasGenerationTask(userId: string, projectId: string, id: string) {
|
|
const [cancelled] = await db
|
|
.update(table)
|
|
.set({
|
|
status: "cancelled",
|
|
error: "已取消排队，未提交供应商",
|
|
updatedAt: new Date(),
|
|
completedAt: new Date(),
|
|
})
|
|
.where(and(scope(userId, projectId), eq(table.id, id), eq(table.status, "pending")))
|
|
.returning();
|
|
if (cancelled) return { ok: true, task: serialize(cancelled) };
|
|
const [current] = await db
|
|
.select()
|
|
.from(table)
|
|
.where(and(scope(userId, projectId), eq(table.id, id)))
|
|
.limit(1);
|
|
return { ok: current?.status === "cancelled", task: current ? serialize(current) : null };
|
|
}
|
|
export async function hideCanvasGenerationNode(
|
|
userId: string,
|
|
projectId: string,
|
|
id: string,
|
|
nodeId: string,
|
|
) {
|
|
return db.transaction(async (tx) => {
|
|
const [task] = await tx
|
|
.select()
|
|
.from(table)
|
|
.where(and(scope(userId, projectId), eq(table.id, id)))
|
|
.for("update");
|
|
if (!task) return false;
|
|
const request = task.request as CanvasTaskRequest & { hiddenNodeIds?: string[] };
|
|
if (nodeId !== request.node.id && !nodeId.startsWith(`${request.node.id}-result-`))
|
|
return false;
|
|
await tx
|
|
.update(table)
|
|
.set({
|
|
request: {
|
|
...request,
|
|
hiddenNodeIds: [...new Set([...(request.hiddenNodeIds || []), nodeId])],
|
|
},
|
|
})
|
|
.where(eq(table.id, id));
|
|
return true;
|
|
});
|
|
}
|
|
|
|
export async function checkpointCanvasGenerationTask(
|
|
id: string,
|
|
result: Record<string, unknown>,
|
|
attempt?: number,
|
|
) {
|
|
await db
|
|
.update(table)
|
|
.set({ result, updatedAt: new Date() })
|
|
.where(
|
|
and(
|
|
eq(table.id, id),
|
|
eq(table.status, "running"),
|
|
attempt === undefined ? undefined : eq(table.attempts, attempt),
|
|
),
|
|
);
|
|
}
|
|
export async function interruptCanvasGenerationTask(id: string, error: string, attempt?: number) {
|
|
await db
|
|
.update(table)
|
|
.set({ status: "interrupted", error, updatedAt: new Date() })
|
|
.where(
|
|
and(
|
|
eq(table.id, id),
|
|
eq(table.status, "running"),
|
|
attempt === undefined ? undefined : eq(table.attempts, attempt),
|
|
),
|
|
);
|
|
}
|
|
|
|
export async function ownsCanvasTaskLease(id: string, attempt: number) {
|
|
const [row] = await db
|
|
.select({ id: table.id })
|
|
.from(table)
|
|
.where(and(eq(table.id, id), eq(table.attempts, attempt), eq(table.status, "running")))
|
|
.limit(1);
|
|
return Boolean(row);
|
|
}
|