535 lines
16 KiB
TypeScript
535 lines
16 KiB
TypeScript
import "server-only";
|
|
|
|
import { randomUUID } from "node:crypto";
|
|
import { and, asc, desc, eq, inArray } from "drizzle-orm";
|
|
import { db } from "@/lib/db";
|
|
import {
|
|
canvasImageTask,
|
|
canvasLayoutReplaceBatch,
|
|
canvasLayoutReplaceBatchItem,
|
|
} from "@/lib/db/schema";
|
|
import { deriveLayoutReplaceBatchAdvance } from "./layout-replace-batch-state";
|
|
import type { CanvasLayoutReplaceTaskRequest } from "./layout-replace-task";
|
|
|
|
export const CANVAS_LAYOUT_REPLACE_BATCH_LIMIT = 5;
|
|
|
|
export type CanvasLayoutReplaceBatchSource = {
|
|
sourceImageUrl: string;
|
|
sourceKind: "canvas" | "upload";
|
|
sourceNodeId?: string;
|
|
sourceTitle?: string;
|
|
};
|
|
|
|
export type CanvasLayoutReplaceBatchItemSnapshot = {
|
|
id: string;
|
|
taskId: string;
|
|
position: number;
|
|
sourceNodeId: string | null;
|
|
sourceImageUrl: string;
|
|
sourceTitle: string | null;
|
|
sourceKind: string;
|
|
status: string;
|
|
taskStatus: string;
|
|
resultImageUrl: string | null;
|
|
resultThumbnailUrl: string | null;
|
|
resultModelId: string | null;
|
|
error: string | null;
|
|
addedToCanvas: boolean;
|
|
addedCanvasNodeId: string | null;
|
|
};
|
|
|
|
export type CanvasLayoutReplaceBatchSnapshot = {
|
|
id: string;
|
|
projectId: string;
|
|
status: string;
|
|
referenceImageUrl: string;
|
|
supplementalPrompt: string | null;
|
|
imageSize: string;
|
|
totalCount: number;
|
|
completedCount: number;
|
|
failedCount: number;
|
|
completedAt: Date | null;
|
|
createdAt: Date;
|
|
updatedAt: Date;
|
|
items: CanvasLayoutReplaceBatchItemSnapshot[];
|
|
};
|
|
|
|
type LayoutReplaceTaskResult = {
|
|
images?: Array<{
|
|
originalUrl?: string;
|
|
thumbnailUrl?: string;
|
|
}>;
|
|
model?: string;
|
|
modelRef?: string;
|
|
};
|
|
|
|
function normalizeText(value: unknown, maxLength: number) {
|
|
return String(value || "").trim().slice(0, maxLength);
|
|
}
|
|
|
|
function normalizeHttpUrl(value: unknown) {
|
|
const url = normalizeText(value, 4096);
|
|
return /^https?:\/\//i.test(url) ? url : "";
|
|
}
|
|
|
|
function isTerminalItemStatus(status: string) {
|
|
return status === "succeeded" || status === "failed";
|
|
}
|
|
|
|
function getResultFields(result: LayoutReplaceTaskResult | null | undefined) {
|
|
const image = result?.images?.[0];
|
|
const resultImageUrl = normalizeHttpUrl(image?.originalUrl || image?.thumbnailUrl);
|
|
return {
|
|
resultImageUrl: resultImageUrl || null,
|
|
resultThumbnailUrl:
|
|
normalizeHttpUrl(image?.thumbnailUrl || image?.originalUrl) || resultImageUrl || null,
|
|
resultModelId: normalizeText(result?.model || result?.modelRef, 160) || null,
|
|
};
|
|
}
|
|
|
|
export async function createCanvasLayoutReplaceBatch(params: {
|
|
userId: string;
|
|
projectId: string;
|
|
sources: CanvasLayoutReplaceBatchSource[];
|
|
referenceImageUrl: string;
|
|
supplementalPrompt?: string;
|
|
imageSize?: string;
|
|
model?: string;
|
|
}) {
|
|
const sources = params.sources.slice(0, CANVAS_LAYOUT_REPLACE_BATCH_LIMIT);
|
|
if (sources.length === 0) {
|
|
throw new Error("请至少添加一张待处理图片");
|
|
}
|
|
|
|
const batchId = randomUUID();
|
|
const referenceImageUrl = normalizeHttpUrl(params.referenceImageUrl);
|
|
const supplementalPrompt = normalizeText(params.supplementalPrompt, 600);
|
|
const imageSize = normalizeText(params.imageSize || "2K", 16) || "2K";
|
|
const model = normalizeText(params.model, 160) || undefined;
|
|
const now = new Date();
|
|
const records = sources.map((source, index) => {
|
|
const taskId = randomUUID();
|
|
const itemId = randomUUID();
|
|
const position = index + 1;
|
|
const taskRequest: CanvasLayoutReplaceTaskRequest = {
|
|
taskKind: "canvas-layout-replace",
|
|
sourceImageUrl: normalizeHttpUrl(source.sourceImageUrl),
|
|
referenceImageUrl,
|
|
supplementalPrompt: supplementalPrompt || undefined,
|
|
model,
|
|
imageSize,
|
|
sourceNodeId: normalizeText(source.sourceNodeId, 160) || undefined,
|
|
referenceLabel: "画布批量排版参考图",
|
|
batchId,
|
|
batchItemId: itemId,
|
|
batchPosition: position,
|
|
batchTotal: sources.length,
|
|
canvasProjectId: params.projectId,
|
|
debugMeta: {
|
|
operation: "layout-replace-batch",
|
|
batchId,
|
|
batchItemId: itemId,
|
|
batchPosition: position,
|
|
batchTotal: sources.length,
|
|
},
|
|
};
|
|
return {
|
|
itemId,
|
|
taskId,
|
|
position,
|
|
source,
|
|
taskRequest,
|
|
};
|
|
});
|
|
|
|
if (!referenceImageUrl || records.some((record) => !record.taskRequest.sourceImageUrl)) {
|
|
throw new Error("图片地址不可用，请重新上传后再试");
|
|
}
|
|
|
|
await db.transaction(async (tx) => {
|
|
await tx.insert(canvasLayoutReplaceBatch).values({
|
|
id: batchId,
|
|
userId: params.userId,
|
|
projectId: params.projectId,
|
|
status: "pending",
|
|
referenceImageUrl,
|
|
supplementalPrompt: supplementalPrompt || null,
|
|
imageSize,
|
|
totalCount: records.length,
|
|
completedCount: 0,
|
|
failedCount: 0,
|
|
createdAt: now,
|
|
updatedAt: now,
|
|
});
|
|
|
|
await tx.insert(canvasImageTask).values(
|
|
records.map((record, index) => ({
|
|
id: record.taskId,
|
|
userId: params.userId,
|
|
status: index === 0 ? "pending" : "waiting",
|
|
request: record.taskRequest as unknown as Record<string, unknown>,
|
|
attempts: 0,
|
|
createdAt: new Date(now.getTime() + index),
|
|
updatedAt: now,
|
|
}))
|
|
);
|
|
|
|
await tx.insert(canvasLayoutReplaceBatchItem).values(
|
|
records.map((record, index) => ({
|
|
id: record.itemId,
|
|
batchId,
|
|
taskId: record.taskId,
|
|
position: record.position,
|
|
sourceNodeId: normalizeText(record.source.sourceNodeId, 160) || null,
|
|
sourceImageUrl: record.taskRequest.sourceImageUrl,
|
|
sourceTitle: normalizeText(record.source.sourceTitle, 256) || null,
|
|
sourceKind: record.source.sourceKind === "upload" ? "upload" : "canvas",
|
|
status: index === 0 ? "pending" : "queued",
|
|
createdAt: new Date(now.getTime() + index),
|
|
updatedAt: now,
|
|
}))
|
|
);
|
|
});
|
|
|
|
return { batchId, firstTaskId: records[0]?.taskId || "" };
|
|
}
|
|
|
|
export async function markCanvasLayoutReplaceBatchItemRunning(
|
|
request: CanvasLayoutReplaceTaskRequest
|
|
) {
|
|
if (!request.batchId || !request.batchItemId) return;
|
|
const now = new Date();
|
|
await db.transaction(async (tx) => {
|
|
await tx
|
|
.update(canvasLayoutReplaceBatchItem)
|
|
.set({ status: "running", error: null, updatedAt: now })
|
|
.where(
|
|
and(
|
|
eq(canvasLayoutReplaceBatchItem.id, request.batchItemId as string),
|
|
eq(canvasLayoutReplaceBatchItem.batchId, request.batchId as string)
|
|
)
|
|
);
|
|
await tx
|
|
.update(canvasLayoutReplaceBatch)
|
|
.set({ status: "running", completedAt: null, updatedAt: now })
|
|
.where(eq(canvasLayoutReplaceBatch.id, request.batchId as string));
|
|
});
|
|
}
|
|
|
|
export async function completeCanvasLayoutReplaceBatchItem(params: {
|
|
request: CanvasLayoutReplaceTaskRequest;
|
|
result?: LayoutReplaceTaskResult | null;
|
|
error?: string | null;
|
|
}) {
|
|
const { request } = params;
|
|
if (!request.batchId || !request.batchItemId) return null;
|
|
const now = new Date();
|
|
const candidateResultFields = getResultFields(params.result);
|
|
const failed = Boolean(params.error) || !candidateResultFields.resultImageUrl;
|
|
const resultFields = failed ? getResultFields(null) : candidateResultFields;
|
|
|
|
return db.transaction(async (tx) => {
|
|
const [completedItem] = await tx
|
|
.update(canvasLayoutReplaceBatchItem)
|
|
.set({
|
|
status: failed ? "failed" : "succeeded",
|
|
...resultFields,
|
|
error: failed
|
|
? normalizeText(params.error, 4000) || "排版替换未返回结果图片"
|
|
: null,
|
|
completedAt: now,
|
|
updatedAt: now,
|
|
})
|
|
.where(
|
|
and(
|
|
eq(canvasLayoutReplaceBatchItem.id, request.batchItemId as string),
|
|
eq(canvasLayoutReplaceBatchItem.batchId, request.batchId as string),
|
|
inArray(canvasLayoutReplaceBatchItem.status, ["pending", "running"])
|
|
)
|
|
)
|
|
.returning({ id: canvasLayoutReplaceBatchItem.id });
|
|
|
|
// Completion can be replayed after a process restart. Only the first state
|
|
// transition may promote the next item, otherwise one batch could fan out.
|
|
if (!completedItem) return null;
|
|
|
|
const items = await tx
|
|
.select({
|
|
id: canvasLayoutReplaceBatchItem.id,
|
|
taskId: canvasLayoutReplaceBatchItem.taskId,
|
|
status: canvasLayoutReplaceBatchItem.status,
|
|
position: canvasLayoutReplaceBatchItem.position,
|
|
})
|
|
.from(canvasLayoutReplaceBatchItem)
|
|
.where(eq(canvasLayoutReplaceBatchItem.batchId, request.batchId as string))
|
|
.orderBy(asc(canvasLayoutReplaceBatchItem.position));
|
|
|
|
const advance = deriveLayoutReplaceBatchAdvance(items);
|
|
const nextItem = advance.nextTaskId
|
|
? items.find((item) => item.taskId === advance.nextTaskId)
|
|
: null;
|
|
|
|
if (nextItem && !advance.isComplete) {
|
|
await tx
|
|
.update(canvasLayoutReplaceBatchItem)
|
|
.set({ status: "pending", updatedAt: now })
|
|
.where(eq(canvasLayoutReplaceBatchItem.id, nextItem.id));
|
|
await tx
|
|
.update(canvasImageTask)
|
|
.set({ status: "pending", updatedAt: now })
|
|
.where(
|
|
and(
|
|
eq(canvasImageTask.id, nextItem.taskId),
|
|
eq(canvasImageTask.status, "waiting")
|
|
)
|
|
);
|
|
}
|
|
|
|
await tx
|
|
.update(canvasLayoutReplaceBatch)
|
|
.set({
|
|
status: advance.status,
|
|
completedCount: advance.succeededCount,
|
|
failedCount: advance.failedCount,
|
|
completedAt: advance.isComplete ? now : null,
|
|
updatedAt: now,
|
|
})
|
|
.where(eq(canvasLayoutReplaceBatch.id, request.batchId as string));
|
|
|
|
return nextItem && !advance.isComplete ? nextItem.taskId : null;
|
|
});
|
|
}
|
|
|
|
async function buildBatchSnapshot(batch: typeof canvasLayoutReplaceBatch.$inferSelect) {
|
|
const items = await db
|
|
.select()
|
|
.from(canvasLayoutReplaceBatchItem)
|
|
.where(eq(canvasLayoutReplaceBatchItem.batchId, batch.id))
|
|
.orderBy(asc(canvasLayoutReplaceBatchItem.position));
|
|
const taskIds = items.map((item) => item.taskId);
|
|
const tasks = taskIds.length
|
|
? await db
|
|
.select({
|
|
id: canvasImageTask.id,
|
|
status: canvasImageTask.status,
|
|
error: canvasImageTask.error,
|
|
})
|
|
.from(canvasImageTask)
|
|
.where(inArray(canvasImageTask.id, taskIds))
|
|
: [];
|
|
const taskById = new Map(tasks.map((task) => [task.id, task]));
|
|
|
|
return {
|
|
id: batch.id,
|
|
projectId: batch.projectId,
|
|
status: batch.status,
|
|
referenceImageUrl: batch.referenceImageUrl,
|
|
supplementalPrompt: batch.supplementalPrompt,
|
|
imageSize: batch.imageSize,
|
|
totalCount: batch.totalCount,
|
|
completedCount: batch.completedCount,
|
|
failedCount: batch.failedCount,
|
|
completedAt: batch.completedAt,
|
|
createdAt: batch.createdAt,
|
|
updatedAt: batch.updatedAt,
|
|
items: items.map((item) => {
|
|
const task = taskById.get(item.taskId);
|
|
return {
|
|
id: item.id,
|
|
taskId: item.taskId,
|
|
position: item.position,
|
|
sourceNodeId: item.sourceNodeId,
|
|
sourceImageUrl: item.sourceImageUrl,
|
|
sourceTitle: item.sourceTitle,
|
|
sourceKind: item.sourceKind,
|
|
status: item.status,
|
|
taskStatus: task?.status || item.status,
|
|
resultImageUrl: item.resultImageUrl,
|
|
resultThumbnailUrl: item.resultThumbnailUrl,
|
|
resultModelId: item.resultModelId,
|
|
error: item.error || task?.error || null,
|
|
addedToCanvas: item.addedToCanvas,
|
|
addedCanvasNodeId: item.addedCanvasNodeId,
|
|
} satisfies CanvasLayoutReplaceBatchItemSnapshot;
|
|
}),
|
|
} satisfies CanvasLayoutReplaceBatchSnapshot;
|
|
}
|
|
|
|
export async function getCanvasLayoutReplaceBatchForUser(params: {
|
|
batchId: string;
|
|
userId: string;
|
|
}) {
|
|
const [batch] = await db
|
|
.select()
|
|
.from(canvasLayoutReplaceBatch)
|
|
.where(
|
|
and(
|
|
eq(canvasLayoutReplaceBatch.id, params.batchId),
|
|
eq(canvasLayoutReplaceBatch.userId, params.userId)
|
|
)
|
|
)
|
|
.limit(1);
|
|
return batch ? buildBatchSnapshot(batch) : null;
|
|
}
|
|
|
|
export async function getLatestCanvasLayoutReplaceBatchForUser(params: {
|
|
projectId: string;
|
|
userId: string;
|
|
sourceNodeId?: string;
|
|
}) {
|
|
const batches = await db
|
|
.select()
|
|
.from(canvasLayoutReplaceBatch)
|
|
.where(
|
|
and(
|
|
eq(canvasLayoutReplaceBatch.projectId, params.projectId),
|
|
eq(canvasLayoutReplaceBatch.userId, params.userId)
|
|
)
|
|
)
|
|
.orderBy(desc(canvasLayoutReplaceBatch.createdAt))
|
|
.limit(12);
|
|
|
|
for (const batch of batches) {
|
|
const snapshot = await buildBatchSnapshot(batch);
|
|
if (
|
|
!params.sourceNodeId ||
|
|
snapshot.items.some((item) => item.sourceNodeId === params.sourceNodeId)
|
|
) {
|
|
return snapshot;
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
export async function markCanvasLayoutReplaceBatchItemsAdded(params: {
|
|
batchId: string;
|
|
userId: string;
|
|
items: Array<{ itemId: string; canvasNodeId: string }>;
|
|
}) {
|
|
const [ownedBatch] = await db
|
|
.select({ id: canvasLayoutReplaceBatch.id })
|
|
.from(canvasLayoutReplaceBatch)
|
|
.where(
|
|
and(
|
|
eq(canvasLayoutReplaceBatch.id, params.batchId),
|
|
eq(canvasLayoutReplaceBatch.userId, params.userId)
|
|
)
|
|
)
|
|
.limit(1);
|
|
if (!ownedBatch) return false;
|
|
|
|
const now = new Date();
|
|
await db.transaction(async (tx) => {
|
|
for (const item of params.items.slice(0, CANVAS_LAYOUT_REPLACE_BATCH_LIMIT)) {
|
|
await tx
|
|
.update(canvasLayoutReplaceBatchItem)
|
|
.set({
|
|
addedToCanvas: true,
|
|
addedCanvasNodeId: normalizeText(item.canvasNodeId, 160) || null,
|
|
updatedAt: now,
|
|
})
|
|
.where(
|
|
and(
|
|
eq(canvasLayoutReplaceBatchItem.id, item.itemId),
|
|
eq(canvasLayoutReplaceBatchItem.batchId, params.batchId),
|
|
eq(canvasLayoutReplaceBatchItem.status, "succeeded")
|
|
)
|
|
);
|
|
}
|
|
});
|
|
return true;
|
|
}
|
|
|
|
export async function retryCanvasLayoutReplaceBatchItem(params: {
|
|
batchId: string;
|
|
itemId: string;
|
|
userId: string;
|
|
}) {
|
|
const [ownedBatch] = await db
|
|
.select({
|
|
id: canvasLayoutReplaceBatch.id,
|
|
status: canvasLayoutReplaceBatch.status,
|
|
})
|
|
.from(canvasLayoutReplaceBatch)
|
|
.where(
|
|
and(
|
|
eq(canvasLayoutReplaceBatch.id, params.batchId),
|
|
eq(canvasLayoutReplaceBatch.userId, params.userId)
|
|
)
|
|
)
|
|
.limit(1);
|
|
if (!ownedBatch || !isCanvasLayoutReplaceBatchTerminal(ownedBatch.status)) {
|
|
return null;
|
|
}
|
|
|
|
return db.transaction(async (tx) => {
|
|
const now = new Date();
|
|
const [retriedItem] = await tx
|
|
.update(canvasLayoutReplaceBatchItem)
|
|
.set({
|
|
status: "pending",
|
|
resultImageUrl: null,
|
|
resultThumbnailUrl: null,
|
|
resultModelId: null,
|
|
error: null,
|
|
completedAt: null,
|
|
updatedAt: now,
|
|
})
|
|
.where(
|
|
and(
|
|
eq(canvasLayoutReplaceBatchItem.id, params.itemId),
|
|
eq(canvasLayoutReplaceBatchItem.batchId, params.batchId),
|
|
eq(canvasLayoutReplaceBatchItem.status, "failed")
|
|
)
|
|
)
|
|
.returning({ taskId: canvasLayoutReplaceBatchItem.taskId });
|
|
if (!retriedItem) return null;
|
|
|
|
await tx
|
|
.update(canvasImageTask)
|
|
.set({
|
|
status: "pending",
|
|
result: null,
|
|
error: null,
|
|
startedAt: null,
|
|
completedAt: null,
|
|
updatedAt: now,
|
|
})
|
|
.where(eq(canvasImageTask.id, retriedItem.taskId));
|
|
|
|
const remainingItems = await tx
|
|
.select({ status: canvasLayoutReplaceBatchItem.status })
|
|
.from(canvasLayoutReplaceBatchItem)
|
|
.where(eq(canvasLayoutReplaceBatchItem.batchId, params.batchId));
|
|
await tx
|
|
.update(canvasLayoutReplaceBatch)
|
|
.set({
|
|
status: "running",
|
|
completedCount: remainingItems.filter((item) => item.status === "succeeded")
|
|
.length,
|
|
failedCount: remainingItems.filter((item) => item.status === "failed").length,
|
|
completedAt: null,
|
|
updatedAt: now,
|
|
})
|
|
.where(eq(canvasLayoutReplaceBatch.id, params.batchId));
|
|
|
|
return retriedItem.taskId;
|
|
});
|
|
}
|
|
|
|
export function getCanvasLayoutReplaceBatchActiveTaskId(
|
|
snapshot: CanvasLayoutReplaceBatchSnapshot
|
|
) {
|
|
if (["succeeded", "failed", "partial_failed"].includes(snapshot.status)) {
|
|
return null;
|
|
}
|
|
const active = snapshot.items.find((item) =>
|
|
["pending", "running"].includes(item.status)
|
|
);
|
|
return active?.taskId || null;
|
|
}
|
|
|
|
export function isCanvasLayoutReplaceBatchTerminal(status: string) {
|
|
return ["succeeded", "failed", "partial_failed"].includes(status);
|
|
}
|