ai-chatbot/lib/canvas/image-task-queue.ts

709 lines
21 KiB
TypeScript
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import "server-only";
import { resolveArchivedImageUrls } from "@/lib/storage/image-archive-db";
import { and, asc, eq, lt, notInArray, sql } from "drizzle-orm";
import { db } from "@/lib/db";
import { canvasImageTask } from "@/lib/db/schema";
import {
getCanvasImageRunningCountForModel,
getCanvasImageSingleModelConcurrency,
normalizeCanvasImageModelKey,
} from "./image-concurrency";
import {
generateCanvasImage,
type CanvasGenerateImageRequest,
type CanvasGenerateImageResult,
} from "./generate-image";
import { applyScaleAdjustModelRelight } from "./scale-adjust-model-relight";
import {
executeStyleProTask,
getStyleProTaskPrimaryModelKey,
type StyleProTaskRequest,
type StyleProTaskResult,
} from "@/lib/style-pro/task";
import {
executeCanvasLayoutReplaceTask,
isCanvasLayoutReplaceTaskRequest,
type CanvasLayoutReplaceTaskRequest,
} from "@/lib/canvas/layout-replace-task";
import {
completeCanvasLayoutReplaceBatchItem,
markCanvasLayoutReplaceBatchItemRunning,
} from "@/lib/canvas/layout-replace-batch";
import { persistAgentProjectGeneratedResult } from "@/lib/agents/agent-project-result-persistence";
type CanvasImageTaskStatus =
| "waiting"
| "pending"
| "running"
| "succeeded"
| "failed";
type CanvasImageTaskRequest =
| CanvasGenerateImageRequest
| StyleProTaskRequest
| CanvasLayoutReplaceTaskRequest;
type CanvasImageTaskResult = CanvasGenerateImageResult | StyleProTaskResult;
export type CanvasImageTaskQueueStatus = {
status: CanvasImageTaskStatus;
queuePosition: number;
sameModelPendingCount: number;
sameModelRunningCount: number;
concurrencyLimit: number;
estimatedWaitMs: number;
source: "model_queue" | "batch_queue";
};
interface PendingTaskEntry {
taskId: string;
modelKey: string;
userId: string;
}
interface RunningTaskEntry {
modelKey: string;
userId: string;
}
interface QueueState {
pending: PendingTaskEntry[];
pendingSet: Set<string>;
running: Map<string, RunningTaskEntry>;
runningByModel: Map<string, Set<string>>;
lastStaleRunningCleanupAt: number;
draining: boolean;
}
const DEFAULT_RUNNING_TIMEOUT_MS = 10 * 60 * 1000;
const SCALE_ADJUST_RUNNING_TIMEOUT_MS = 12 * 60 * 1000;
const LOCAL_ACTIVE_RUNNING_HARD_TIMEOUT_MS = 30 * 60 * 1000;
const STALE_RUNNING_CLEANUP_INTERVAL_MS = 60 * 1000;
const PENDING_LOOKAHEAD_LIMIT = 100;
const queueStateKey = "__canvas_image_task_queue_state__";
function getQueueState(): QueueState {
const globalObj = globalThis as typeof globalThis & {
[queueStateKey]?: QueueState;
};
if (!globalObj[queueStateKey]) {
globalObj[queueStateKey] = {
pending: [],
pendingSet: new Set<string>(),
running: new Map<string, RunningTaskEntry>(),
runningByModel: new Map<string, Set<string>>(),
lastStaleRunningCleanupAt: 0,
draining: false,
};
}
return globalObj[queueStateKey] as QueueState;
}
function normalizeTaskModelKey(request: unknown) {
if (isStyleProTaskRequest(request)) {
return normalizeCanvasImageModelKey(getStyleProTaskPrimaryModelKey(request));
}
if (!request || typeof request !== "object" || Array.isArray(request)) {
return "__default__";
}
const rawModel = (request as Record<string, unknown>).model;
return normalizeCanvasImageModelKey(rawModel);
}
function getRecord(value: unknown): Record<string, unknown> | null {
return value && typeof value === "object" && !Array.isArray(value)
? (value as Record<string, unknown>)
: null;
}
function getCanvasImageRunningTimeoutMs(request: unknown) {
const record = getRecord(request);
const debugMeta = getRecord(record?.debugMeta);
const debugLabel = String(record?.debugLabel || "");
const operation = String(debugMeta?.operation || "");
const isScaleAdjust =
debugLabel === "canvas-scale-adjust" ||
operation === "scale-adjust" ||
Boolean(debugMeta?.scaleAdjustPatchWorkflowApplied) ||
Boolean(debugMeta?.scaleAdjustPatchForcedForShrink) ||
Boolean(record?.scaleAdjustPatchComposite);
return isScaleAdjust
? SCALE_ADJUST_RUNNING_TIMEOUT_MS
: DEFAULT_RUNNING_TIMEOUT_MS;
}
async function withCanvasImageTaskExecutionTimeout<T>(
operation: Promise<T>,
timeoutMs: number,
taskId: string
): Promise<T> {
if (!Number.isFinite(timeoutMs) || timeoutMs <= 0) {
return operation;
}
let timeoutId: ReturnType<typeof setTimeout> | null = null;
try {
return await Promise.race([
operation,
new Promise<T>((_, reject) => {
timeoutId = setTimeout(() => {
reject(
new Error(
`图片生成任务执行超时(${Math.round(timeoutMs / 1000)}秒,${taskId})`
)
);
}, timeoutMs);
}),
]);
} finally {
if (timeoutId) {
clearTimeout(timeoutId);
}
}
}
function isStyleProTaskRequest(request: unknown): request is StyleProTaskRequest {
return Boolean(
request &&
typeof request === "object" &&
!Array.isArray(request) &&
(request as { taskKind?: unknown }).taskKind === "style-pro"
);
}
async function getTaskQueueMeta(taskId: string) {
const rows = await db
.select({
request: canvasImageTask.request,
userId: canvasImageTask.userId,
})
.from(canvasImageTask)
.where(eq(canvasImageTask.id, taskId))
.limit(1);
const row = rows[0];
return {
modelKey: normalizeTaskModelKey(row?.request),
userId: String(row?.userId || "__anonymous__"),
};
}
function markRunningTask(taskId: string, modelKey: string, userId: string) {
const state = getQueueState();
state.running.set(taskId, { modelKey, userId });
const tasks = state.runningByModel.get(modelKey) || new Set<string>();
tasks.add(taskId);
state.runningByModel.set(modelKey, tasks);
}
function releaseRunningTask(taskId: string) {
const state = getQueueState();
const entry = state.running.get(taskId);
if (!entry) {
return null;
}
state.running.delete(taskId);
const { modelKey } = entry;
const tasks = state.runningByModel.get(modelKey);
if (tasks) {
tasks.delete(taskId);
if (tasks.size === 0) {
state.runningByModel.delete(modelKey);
}
}
return modelKey;
}
async function enqueueLocalTask(
taskId: string,
meta?: { modelKey?: string; userId?: string } | string
) {
const state = getQueueState();
if (!taskId || state.running.has(taskId) || state.pendingSet.has(taskId)) {
return;
}
const resolvedMeta =
typeof meta === "string"
? { modelKey: meta, userId: undefined }
: meta || {};
const dbMeta =
resolvedMeta.modelKey && resolvedMeta.userId
? null
: await getTaskQueueMeta(taskId);
const resolvedModelKey = resolvedMeta.modelKey || dbMeta?.modelKey || "__default__";
const resolvedUserId = resolvedMeta.userId || dbMeta?.userId || "__anonymous__";
state.pending.push({
taskId,
modelKey: resolvedModelKey,
userId: resolvedUserId,
});
state.pendingSet.add(taskId);
}
async function claimNextPendingTask(): Promise<PendingTaskEntry | null> {
const state = getQueueState();
for (let index = 0; index < state.pending.length; index += 1) {
const entry = state.pending[index];
if (!entry) continue;
if (state.running.has(entry.taskId)) continue;
state.pending.splice(index, 1);
state.pendingSet.delete(entry.taskId);
return entry;
}
const runningIds = Array.from(state.running.keys());
const rows = await db
.select({
id: canvasImageTask.id,
request: canvasImageTask.request,
userId: canvasImageTask.userId,
})
.from(canvasImageTask)
.where(
runningIds.length > 0
? and(eq(canvasImageTask.status, "pending"), notInArray(canvasImageTask.id, runningIds))
: eq(canvasImageTask.status, "pending")
)
.orderBy(asc(canvasImageTask.createdAt))
.limit(PENDING_LOOKAHEAD_LIMIT);
for (const row of rows) {
if (!row?.id) continue;
if (state.pendingSet.has(row.id) || state.running.has(row.id)) continue;
const modelKey = normalizeTaskModelKey(row.request);
const userId = String(row.userId || "__anonymous__");
return {
taskId: row.id,
modelKey,
userId,
};
}
return null;
}
async function markStaleDatabaseRunningTasksFailed() {
const state = getQueueState();
const now = Date.now();
if (now - state.lastStaleRunningCleanupAt < STALE_RUNNING_CLEANUP_INTERVAL_MS) {
return;
}
state.lastStaleRunningCleanupAt = now;
const localRunningIds = Array.from(state.running.keys());
const staleBefore = new Date(now - LOCAL_ACTIVE_RUNNING_HARD_TIMEOUT_MS);
const staleWhere =
localRunningIds.length > 0
? and(
eq(canvasImageTask.status, "running"),
lt(canvasImageTask.startedAt, staleBefore),
notInArray(canvasImageTask.id, localRunningIds)
)
: and(
eq(canvasImageTask.status, "running"),
lt(canvasImageTask.startedAt, staleBefore)
);
const completedAt = new Date(now);
const staleTasks = await db
.update(canvasImageTask)
.set({
status: "failed",
error: "图片生成任务已超时,请重新发起生成",
completedAt,
updatedAt: completedAt,
})
.where(staleWhere)
.returning({ id: canvasImageTask.id });
if (staleTasks.length > 0) {
console.warn("[Canvas Image Task Queue] Marked stale running tasks failed:", {
count: staleTasks.length,
taskIds: staleTasks.slice(0, 10).map((task) => task.id),
});
}
}
async function runSingleTask(entry: PendingTaskEntry) {
let completedBatchRequest: CanvasLayoutReplaceTaskRequest | null = null;
let completedBatchResult: CanvasImageTaskResult | null = null;
let completedBatchError: string | null = null;
try {
const [task] = await db
.update(canvasImageTask)
.set({
status: "running",
startedAt: new Date(),
updatedAt: new Date(),
completedAt: null,
error: null,
attempts: sql`${canvasImageTask.attempts} + 1`,
})
.where(and(eq(canvasImageTask.id, entry.taskId), eq(canvasImageTask.status, "pending")))
.returning({
id: canvasImageTask.id,
request: canvasImageTask.request,
userId: canvasImageTask.userId,
});
if (!task) {
return;
}
const request = task.request as CanvasImageTaskRequest;
if (isCanvasLayoutReplaceTaskRequest(request) && request.batchId) {
completedBatchRequest = request;
try {
await markCanvasLayoutReplaceBatchItemRunning(request);
} catch (error) {
console.error("[Canvas Image Task Queue] Failed to mark batch item running:", {
batchId: request.batchId,
error: error instanceof Error ? error.message : String(error),
taskId: entry.taskId,
});
}
}
const runningTimeoutMs = getCanvasImageRunningTimeoutMs(request);
const execution: Promise<CanvasImageTaskResult> = isStyleProTaskRequest(request)
? executeStyleProTask({
request,
userId: String(task.userId || entry.userId),
})
: isCanvasLayoutReplaceTaskRequest(request)
? executeCanvasLayoutReplaceTask({
request,
userId: String(task.userId || entry.userId),
})
: (async () => {
const imageRequest = request as CanvasGenerateImageRequest;
const ownerKey = String(task.userId || entry.userId);
const initialResult = await generateCanvasImage({
...imageRequest,
ownerKey,
});
return await applyScaleAdjustModelRelight({
generateImage: generateCanvasImage,
initialResult,
ownerKey,
request: imageRequest,
});
})();
const result = await withCanvasImageTaskExecutionTimeout(
execution,
runningTimeoutMs,
entry.taskId
);
completedBatchResult = result;
if (isCanvasLayoutReplaceTaskRequest(request) && request.agentProjectId) {
const layoutReplaceResult = result as CanvasGenerateImageResult;
const generatedImageUrl = String(
layoutReplaceResult.images?.[0]?.originalUrl ||
layoutReplaceResult.images?.[0]?.thumbnailUrl ||
""
).trim();
try {
await persistAgentProjectGeneratedResult({
projectId: request.agentProjectId,
userId: String(task.userId || entry.userId),
imageUrl: generatedImageUrl,
modeType: "scene_generation",
});
} catch (error) {
console.error("[Canvas Image Task Queue] Failed to persist agent AI template result:", {
error: error instanceof Error ? error.message : String(error),
projectId: request.agentProjectId,
taskId: entry.taskId,
});
}
}
await db
.update(canvasImageTask)
.set({
status: "succeeded",
result: await resolveArchivedImageUrls(result) as CanvasImageTaskResult,
error: null,
completedAt: new Date(),
updatedAt: new Date(),
})
.where(
and(eq(canvasImageTask.id, entry.taskId), eq(canvasImageTask.status, "running"))
);
} catch (error) {
completedBatchError = error instanceof Error ? error.message : String(error);
await db
.update(canvasImageTask)
.set({
status: "failed",
error: error instanceof Error ? error.message : String(error),
completedAt: new Date(),
updatedAt: new Date(),
})
.where(
and(eq(canvasImageTask.id, entry.taskId), eq(canvasImageTask.status, "running"))
);
} finally {
releaseRunningTask(entry.taskId);
let nextBatchTaskId: string | null = null;
if (completedBatchRequest?.batchId) {
try {
nextBatchTaskId = await completeCanvasLayoutReplaceBatchItem({
request: completedBatchRequest,
result: completedBatchResult as CanvasGenerateImageResult | null,
error: completedBatchError,
});
} catch (error) {
console.error("[Canvas Image Task Queue] Failed to advance layout batch:", {
batchId: completedBatchRequest.batchId,
error: error instanceof Error ? error.message : String(error),
taskId: entry.taskId,
});
}
}
if (nextBatchTaskId) {
await enqueueLocalTask(nextBatchTaskId);
}
void drainCanvasImageTaskQueue();
}
}
export async function drainCanvasImageTaskQueue() {
const state = getQueueState();
if (state.draining) return;
state.draining = true;
try {
await markStaleDatabaseRunningTasksFailed();
while (true) {
const nextTask = await claimNextPendingTask();
if (!nextTask) break;
if (state.running.has(nextTask.taskId)) continue;
markRunningTask(nextTask.taskId, nextTask.modelKey, nextTask.userId);
void runSingleTask(nextTask);
}
} finally {
state.draining = false;
// An arrival can occur while the final database scan is in flight. Dispatch it
// immediately rather than relying on a later status poll to wake the queue.
if (state.pending.some((entry) => !state.running.has(entry.taskId))) {
void drainCanvasImageTaskQueue();
}
}
}
export async function enqueueCanvasImageTask(taskId: string) {
await enqueueLocalTask(taskId);
await drainCanvasImageTaskQueue();
}
export async function ensureCanvasImageTaskActive(params: {
taskId: string;
userId: string;
}) {
const [task] = await db
.select({
id: canvasImageTask.id,
status: canvasImageTask.status,
request: canvasImageTask.request,
result: canvasImageTask.result,
error: canvasImageTask.error,
})
.from(canvasImageTask)
.where(
and(
eq(canvasImageTask.id, params.taskId),
eq(canvasImageTask.userId, params.userId)
)
)
.limit(1);
if (!task || task.status === "waiting") return false;
const request = task.request as CanvasImageTaskRequest;
if (
(task.status === "succeeded" || task.status === "failed") &&
isCanvasLayoutReplaceTaskRequest(request) &&
request.batchId
) {
const nextTaskId = await completeCanvasLayoutReplaceBatchItem({
request,
result: task.result as CanvasGenerateImageResult | null,
error: task.status === "failed" ? task.error || "排版替换失败" : null,
});
if (nextTaskId) {
await enqueueLocalTask(nextTaskId);
await drainCanvasImageTaskQueue();
}
return true;
}
if (task.status === "running" && !getQueueState().running.has(task.id)) {
await db
.update(canvasImageTask)
.set({
status: "pending",
startedAt: null,
completedAt: null,
error: null,
updatedAt: new Date(),
})
.where(
and(
eq(canvasImageTask.id, task.id),
eq(canvasImageTask.status, "running")
)
);
} else if (task.status !== "pending" && task.status !== "running") {
return false;
}
await enqueueLocalTask(task.id, { userId: params.userId });
await drainCanvasImageTaskQueue();
return true;
}
export async function createCanvasImageTask(params: {
userId: string;
request: Record<string, unknown>;
}) {
const [task] = await db
.insert(canvasImageTask)
.values({
userId: params.userId,
status: "pending",
request: params.request as unknown as Record<string, unknown>,
attempts: 0,
createdAt: new Date(),
updatedAt: new Date(),
})
.returning({
id: canvasImageTask.id,
status: canvasImageTask.status,
createdAt: canvasImageTask.createdAt,
});
return task;
}
async function buildCanvasImageTaskQueueStatus(task: {
id: string;
status: CanvasImageTaskStatus;
request: unknown;
createdAt: Date;
}) {
const modelKey = normalizeTaskModelKey(task.request);
const concurrencyLimit = await getCanvasImageSingleModelConcurrency();
const runningRows = await db
.select({
request: canvasImageTask.request,
})
.from(canvasImageTask)
.where(eq(canvasImageTask.status, "running"));
const sameModelRunningCount = Math.max(
runningRows.filter((row) => normalizeTaskModelKey(row.request) === modelKey).length,
getCanvasImageRunningCountForModel(modelKey)
);
return {
status: task.status,
queuePosition: 0,
sameModelPendingCount: 0,
sameModelRunningCount,
concurrencyLimit,
estimatedWaitMs: 0,
source: "model_queue" as const,
};
}
export async function getCanvasImageTaskForUser(params: {
taskId: string;
userId: string;
}) {
const rows = await db
.select({
id: canvasImageTask.id,
status: canvasImageTask.status,
request: canvasImageTask.request,
result: canvasImageTask.result,
error: canvasImageTask.error,
attempts: canvasImageTask.attempts,
startedAt: canvasImageTask.startedAt,
completedAt: canvasImageTask.completedAt,
createdAt: canvasImageTask.createdAt,
updatedAt: canvasImageTask.updatedAt,
})
.from(canvasImageTask)
.where(and(eq(canvasImageTask.id, params.taskId), eq(canvasImageTask.userId, params.userId)))
.limit(1);
const task = rows[0];
if (!task) return null;
if (task.status === "running" && task.startedAt) {
const now = Date.now();
const runningAgeMs = now - task.startedAt.getTime();
const runningTimeoutMs = getCanvasImageRunningTimeoutMs(task.request);
const locallyRunning = getQueueState().running.has(task.id);
const shouldMarkStale =
runningAgeMs >= runningTimeoutMs &&
(!locallyRunning || runningAgeMs >= LOCAL_ACTIVE_RUNNING_HARD_TIMEOUT_MS);
if (shouldMarkStale) {
const staleBefore = new Date(now - runningTimeoutMs);
const timeoutMessage = "图片生成任务已超时,请重新发起生成";
const completedAt = new Date();
const [timedOutTask] = await db
.update(canvasImageTask)
.set({
status: "failed",
error: timeoutMessage,
completedAt,
updatedAt: completedAt,
})
.where(
and(
eq(canvasImageTask.id, task.id),
eq(canvasImageTask.status, "running"),
lt(canvasImageTask.startedAt, staleBefore)
)
)
.returning({
status: canvasImageTask.status,
error: canvasImageTask.error,
completedAt: canvasImageTask.completedAt,
});
if (timedOutTask) {
releaseRunningTask(task.id);
task.status = timedOutTask.status as CanvasImageTaskStatus;
task.error = timedOutTask.error;
task.completedAt = timedOutTask.completedAt;
}
}
}
if (task.status === "pending") {
await enqueueLocalTask(task.id, {
modelKey: normalizeTaskModelKey(task.request),
userId: params.userId,
});
await drainCanvasImageTaskQueue();
}
return {
...task,
queueStatus: await buildCanvasImageTaskQueueStatus({
id: task.id,
status: task.status as CanvasImageTaskStatus,
request: task.request,
createdAt: task.createdAt,
}),
};
}