528 lines
14 KiB
TypeScript
528 lines
14 KiB
TypeScript
import { validateGroupBatch, groupDependencyState } from "./group-execution";
|
|
import { canResumeCanvasTask, hasRecoverableCanvasTask } from "./task-recovery";
|
|
import "server-only";
|
|
import { and, asc, desc, eq, gte, lt, sql, inArray } 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,
|
|
beforeCreate?: () => Promise<void>
|
|
) {
|
|
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 }
|
|
);
|
|
await beforeCreate?.();
|
|
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;
|
|
const pending = await tx
|
|
.select()
|
|
.from(table)
|
|
.where(and(scope(userId, projectId), eq(table.status, "pending")))
|
|
.orderBy(asc(table.createdAt))
|
|
.limit(40);
|
|
const dependencyIds = [
|
|
...new Set(
|
|
pending.flatMap(
|
|
(row) =>
|
|
(row.request as CanvasTaskRequest).workflow?.inputs
|
|
.map((input) => input.requestId)
|
|
.filter((id): id is string => Boolean(id)) || []
|
|
)
|
|
),
|
|
];
|
|
const predecessors = dependencyIds.length
|
|
? await tx
|
|
.select()
|
|
.from(table)
|
|
.where(
|
|
and(
|
|
scope(userId, projectId),
|
|
inArray(table.requestId, dependencyIds)
|
|
)
|
|
)
|
|
: [];
|
|
const tasks = [
|
|
...new Map(
|
|
[...pending, ...predecessors].map((row) => [row.id, serialize(row)])
|
|
).values(),
|
|
];
|
|
let next: typeof table.$inferSelect | undefined;
|
|
for (const candidate of pending) {
|
|
const state = groupDependencyState(
|
|
candidate.request as CanvasTaskRequest,
|
|
tasks
|
|
);
|
|
if (state === "failed") {
|
|
await tx
|
|
.update(table)
|
|
.set({
|
|
status: "failed",
|
|
error: "上游任务失败或取消，未提交此节点",
|
|
updatedAt: new Date(),
|
|
completedAt: new Date(),
|
|
})
|
|
.where(eq(table.id, candidate.id));
|
|
tasks.find((t) => t.id === candidate.id)!.status = "failed";
|
|
} else if (state === "ready" && !next) next = candidate;
|
|
}
|
|
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);
|
|
}
|
|
|
|
export async function createCanvasGroupTasks(
|
|
userId: string,
|
|
requests: CanvasTaskRequest[],
|
|
beforeCreate?: () => Promise<void>
|
|
) {
|
|
validateGroupBatch(requests);
|
|
return db.transaction(async (tx) => {
|
|
const projectId = requests[0].projectId;
|
|
await tx.execute(
|
|
sql`select pg_advisory_xact_lock(hashtext(${capabilityId}), hashtext(${projectId}))`
|
|
);
|
|
const same = await tx
|
|
.select()
|
|
.from(table)
|
|
.where(
|
|
and(
|
|
scope(userId, projectId),
|
|
sql`${table.request}->'workflow'->>'runId' = ${requests[0].workflow!.runId}`
|
|
)
|
|
);
|
|
if (same.length) {
|
|
if (
|
|
same.length !== requests.length ||
|
|
requests.some((r) => {
|
|
const old = same.find((t) => t.requestId === r.requestId);
|
|
return (
|
|
!old ||
|
|
JSON.stringify(sortObject(old.request)) !==
|
|
JSON.stringify(sortObject(r))
|
|
);
|
|
})
|
|
)
|
|
throw Object.assign(Error("同一编组执行编号不能用于不同参数"), {
|
|
status: 409,
|
|
});
|
|
return requests.map((request) =>
|
|
serialize(same.find((row) => row.requestId === request.requestId)!)
|
|
);
|
|
}
|
|
const candidates = await tx
|
|
.select()
|
|
.from(table)
|
|
.where(
|
|
and(
|
|
scope(userId, projectId),
|
|
sql`(${table.status} in ('pending','running') or (${table.status}='interrupted' and ${table.createdAt} > now() - interval '24 hours'))`
|
|
)
|
|
);
|
|
const active = candidates.filter(
|
|
(t) =>
|
|
["pending", "running"].includes(t.status) ||
|
|
(t.status === "interrupted" && hasRecoverableCanvasTask(serialize(t)))
|
|
);
|
|
if (active.length + requests.length > 40)
|
|
throw Object.assign(Error("画布任务队列已满"), { status: 429 });
|
|
if (
|
|
active.some((t) =>
|
|
requests.some(
|
|
(r) => r.node.id === (t.request as CanvasTaskRequest).node.id
|
|
)
|
|
)
|
|
)
|
|
throw Object.assign(Error("编组中已有节点正在执行，请等待完成"), {
|
|
status: 409,
|
|
});
|
|
await beforeCreate?.();
|
|
const now = Date.now();
|
|
const rows = await tx
|
|
.insert(table)
|
|
.values(
|
|
requests.map((request, i) => ({
|
|
requestId: request.requestId,
|
|
userId,
|
|
projectId,
|
|
capabilityId,
|
|
source: "video-canvas",
|
|
request,
|
|
createdAt: new Date(now + i),
|
|
updatedAt: new Date(now),
|
|
}))
|
|
)
|
|
.returning();
|
|
return rows.map(serialize);
|
|
});
|
|
}
|
|
export async function groupTaskInputs(
|
|
userId: string,
|
|
request: CanvasTaskRequest
|
|
) {
|
|
const rows = await db
|
|
.select()
|
|
.from(table)
|
|
.where(
|
|
and(
|
|
scope(userId, request.projectId),
|
|
sql`${table.request}->'workflow'->>'runId' = ${request.workflow?.runId || ""}`
|
|
)
|
|
);
|
|
return rows
|
|
.map(serialize)
|
|
.filter((t) => t.request.workflow?.runId === request.workflow?.runId);
|
|
}
|