997 lines
29 KiB
TypeScript
997 lines
29 KiB
TypeScript
import "server-only";
|
||
|
||
import { and, asc, desc, eq, lt, notInArray, sql } from "drizzle-orm";
|
||
import { db } from "@/lib/db";
|
||
import { canvasImageTask } from "@/lib/db/schema";
|
||
import {
|
||
getCanvasImageRunningCountForModel,
|
||
getCanvasImageSingleModelConcurrency,
|
||
getCanvasImageTaskQueueConfig,
|
||
normalizeCanvasImageModelKey,
|
||
} from "./image-concurrency";
|
||
import {
|
||
generateCanvasImage,
|
||
type CanvasGenerateImageRequest,
|
||
type CanvasGenerateImageResult,
|
||
} from "./generate-image";
|
||
import {
|
||
executeStyleProTask,
|
||
getStyleProTaskPrimaryModelKey,
|
||
type StyleProTaskRequest,
|
||
type StyleProTaskResult,
|
||
} from "@/lib/style-pro/task";
|
||
import {
|
||
executeCanvasLayoutReplaceTask,
|
||
isCanvasLayoutReplaceTaskRequest,
|
||
type CanvasLayoutReplaceTaskRequest,
|
||
} from "@/lib/canvas/layout-replace-task";
|
||
import {
|
||
estimateImageDispatchWaitMs,
|
||
getAdaptiveImageQueueOwnerLimit,
|
||
} from "./image-queue-fairness";
|
||
|
||
type CanvasImageTaskStatus = "pending" | "running" | "succeeded" | "failed";
|
||
type CanvasImageTaskRequest =
|
||
| CanvasGenerateImageRequest
|
||
| StyleProTaskRequest
|
||
| CanvasLayoutReplaceTaskRequest;
|
||
type CanvasImageTaskResult = CanvasGenerateImageResult | StyleProTaskResult;
|
||
|
||
export type CanvasGenerateImageTaskUpdate = {
|
||
taskId: string;
|
||
status: CanvasImageTaskStatus;
|
||
queueStatus: CanvasImageTaskQueueStatus;
|
||
error?: string | null;
|
||
};
|
||
|
||
export type CanvasImageTaskQueueStatus = {
|
||
status: CanvasImageTaskStatus;
|
||
queuePosition: number;
|
||
sameModelPendingCount: number;
|
||
sameModelRunningCount: number;
|
||
concurrencyLimit: number;
|
||
dispatchIntervalMs: number;
|
||
estimatedWaitMs: number;
|
||
source: "model_queue" | "batch_queue";
|
||
};
|
||
|
||
export type CanvasDetailReplacementTaskContextSnapshot = {
|
||
id: string;
|
||
status: CanvasImageTaskStatus;
|
||
modeType: "detail_replacement";
|
||
referenceId: string;
|
||
messageId: string | null;
|
||
error: string | null;
|
||
result: unknown;
|
||
attempts: number;
|
||
startedAt: Date | null;
|
||
completedAt: Date | null;
|
||
createdAt: Date;
|
||
updatedAt: Date;
|
||
queueStatus: CanvasImageTaskQueueStatus | null;
|
||
};
|
||
|
||
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>>;
|
||
runningByModelUser: Map<string, Set<string>>;
|
||
nextBatchEligibleAt: Map<string, number>;
|
||
nextBatchTimers: Map<string, ReturnType<typeof setTimeout>>;
|
||
awaitingNextBatch: Set<string>;
|
||
lastStaleRunningCleanupAt: number;
|
||
draining: boolean;
|
||
}
|
||
|
||
const DEFAULT_RUNNING_TIMEOUT_MS = 10 * 60 * 1000;
|
||
const DETAIL_REPLACEMENT_RUNNING_TIMEOUT_MS = 18 * 60 * 1000;
|
||
const SCALE_ADJUST_RUNNING_TIMEOUT_MS = 15 * 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 QUEUE_STATUS_PENDING_LOOKAHEAD_LIMIT = 500;
|
||
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>>(),
|
||
runningByModelUser: new Map<string, Set<string>>(),
|
||
nextBatchEligibleAt: new Map<string, number>(),
|
||
nextBatchTimers: new Map<string, ReturnType<typeof setTimeout>>(),
|
||
awaitingNextBatch: new Set<string>(),
|
||
lastStaleRunningCleanupAt: 0,
|
||
draining: false,
|
||
};
|
||
}
|
||
return globalObj[queueStateKey] as QueueState;
|
||
}
|
||
|
||
function normalizeTaskModelKey(request: unknown) {
|
||
if (isStyleProTaskRequest(request)) {
|
||
return normalizeCanvasImageModelKey(getStyleProTaskPrimaryModelKey(request));
|
||
}
|
||
|
||
if (isCanvasLayoutReplaceTaskRequest(request)) {
|
||
return normalizeCanvasImageModelKey(request.model || "__default__");
|
||
}
|
||
|
||
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 isDetailReplacement =
|
||
debugLabel === "agent-detail-replacement" ||
|
||
operation === "agent-detail-replacement";
|
||
const isScaleAdjust =
|
||
debugLabel === "canvas-scale-adjust" ||
|
||
operation === "scale-adjust" ||
|
||
Boolean(debugMeta?.scaleAdjustPatchWorkflowApplied) ||
|
||
Boolean(debugMeta?.scaleAdjustPatchForcedForShrink) ||
|
||
Boolean(record?.scaleAdjustPatchComposite);
|
||
|
||
if (isScaleAdjust) {
|
||
return SCALE_ADJUST_RUNNING_TIMEOUT_MS;
|
||
}
|
||
return isDetailReplacement
|
||
? DETAIL_REPLACEMENT_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"
|
||
);
|
||
}
|
||
|
||
function getModelUserKey(modelKey: string, userId: string) {
|
||
return `${modelKey}::${userId || "__anonymous__"}`;
|
||
}
|
||
|
||
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 getRunningCountForModel(modelKey: string) {
|
||
const state = getQueueState();
|
||
return state.runningByModel.get(modelKey)?.size || 0;
|
||
}
|
||
|
||
function getRunningCountForModelUser(modelKey: string, userId: string) {
|
||
const state = getQueueState();
|
||
return state.runningByModelUser.get(getModelUserKey(modelKey, userId))?.size || 0;
|
||
}
|
||
|
||
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);
|
||
|
||
const modelUserKey = getModelUserKey(modelKey, userId);
|
||
const userTasks = state.runningByModelUser.get(modelUserKey) || new Set<string>();
|
||
userTasks.add(taskId);
|
||
state.runningByModelUser.set(modelUserKey, userTasks);
|
||
}
|
||
|
||
function releaseRunningTask(taskId: string) {
|
||
const state = getQueueState();
|
||
const entry = state.running.get(taskId);
|
||
if (!entry) {
|
||
return null;
|
||
}
|
||
|
||
state.running.delete(taskId);
|
||
const { modelKey, userId } = entry;
|
||
const tasks = state.runningByModel.get(modelKey);
|
||
if (tasks) {
|
||
tasks.delete(taskId);
|
||
if (tasks.size === 0) {
|
||
state.runningByModel.delete(modelKey);
|
||
}
|
||
}
|
||
|
||
const modelUserKey = getModelUserKey(modelKey, userId);
|
||
const userTasks = state.runningByModelUser.get(modelUserKey);
|
||
if (userTasks) {
|
||
userTasks.delete(taskId);
|
||
if (userTasks.size === 0) {
|
||
state.runningByModelUser.delete(modelUserKey);
|
||
}
|
||
}
|
||
|
||
return modelKey;
|
||
}
|
||
|
||
function isModelBatchDispatchDelayed(modelKey: string) {
|
||
const state = getQueueState();
|
||
const delayedUntil = state.nextBatchEligibleAt.get(modelKey);
|
||
if (!delayedUntil) {
|
||
return false;
|
||
}
|
||
|
||
if (delayedUntil <= Date.now()) {
|
||
state.nextBatchEligibleAt.delete(modelKey);
|
||
const timer = state.nextBatchTimers.get(modelKey);
|
||
if (timer) {
|
||
clearTimeout(timer);
|
||
state.nextBatchTimers.delete(modelKey);
|
||
}
|
||
return false;
|
||
}
|
||
|
||
return true;
|
||
}
|
||
|
||
function markModelAwaitingNextBatch(modelKey: string) {
|
||
getQueueState().awaitingNextBatch.add(modelKey);
|
||
}
|
||
|
||
function scheduleModelNextBatch(modelKey: string, delayMs: number) {
|
||
const state = getQueueState();
|
||
if (delayMs <= 0) {
|
||
state.awaitingNextBatch.delete(modelKey);
|
||
return false;
|
||
}
|
||
|
||
if (state.nextBatchTimers.has(modelKey)) {
|
||
return true;
|
||
}
|
||
|
||
state.awaitingNextBatch.delete(modelKey);
|
||
state.nextBatchEligibleAt.set(modelKey, Date.now() + delayMs);
|
||
const timer = setTimeout(() => {
|
||
const latestState = getQueueState();
|
||
latestState.nextBatchTimers.delete(modelKey);
|
||
latestState.nextBatchEligibleAt.delete(modelKey);
|
||
void drainCanvasImageTaskQueue();
|
||
}, delayMs);
|
||
state.nextBatchTimers.set(modelKey, timer);
|
||
return true;
|
||
}
|
||
|
||
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);
|
||
|
||
const maxConcurrentPerModel = await getCanvasImageSingleModelConcurrency();
|
||
const hasCompetingUser = state.pending.some(
|
||
(entry) =>
|
||
entry.modelKey === resolvedModelKey && entry.userId !== resolvedUserId
|
||
);
|
||
const ownerLimit = getAdaptiveImageQueueOwnerLimit({
|
||
maxConcurrent: maxConcurrentPerModel,
|
||
hasCompetingOwner: hasCompetingUser,
|
||
});
|
||
if (
|
||
getRunningCountForModel(resolvedModelKey) >= maxConcurrentPerModel ||
|
||
getRunningCountForModelUser(resolvedModelKey, resolvedUserId) >= ownerLimit
|
||
) {
|
||
markModelAwaitingNextBatch(resolvedModelKey);
|
||
}
|
||
}
|
||
|
||
function hasModelCapacity(
|
||
modelKey: string,
|
||
userId: string,
|
||
maxConcurrentPerModel: number,
|
||
pendingEntries: PendingTaskEntry[]
|
||
) {
|
||
const hasCompetingUser = pendingEntries.some(
|
||
(entry) => entry.modelKey === modelKey && entry.userId !== userId
|
||
);
|
||
const ownerLimit = getAdaptiveImageQueueOwnerLimit({
|
||
maxConcurrent: maxConcurrentPerModel,
|
||
hasCompetingOwner: hasCompetingUser,
|
||
});
|
||
return (
|
||
getRunningCountForModel(modelKey) < maxConcurrentPerModel &&
|
||
getRunningCountForModelUser(modelKey, userId) < ownerLimit
|
||
);
|
||
}
|
||
|
||
async function claimNextPendingTask(
|
||
maxConcurrentPerModel: number
|
||
): 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;
|
||
if (isModelBatchDispatchDelayed(entry.modelKey)) continue;
|
||
if (
|
||
!hasModelCapacity(
|
||
entry.modelKey,
|
||
entry.userId,
|
||
maxConcurrentPerModel,
|
||
state.pending
|
||
)
|
||
) {
|
||
markModelAwaitingNextBatch(entry.modelKey);
|
||
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);
|
||
|
||
const databasePendingEntries = rows.map((row) => ({
|
||
taskId: row.id,
|
||
modelKey: normalizeTaskModelKey(row.request),
|
||
userId: String(row.userId || "__anonymous__"),
|
||
}));
|
||
const capacityCandidates = [...state.pending, ...databasePendingEntries];
|
||
|
||
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);
|
||
if (isModelBatchDispatchDelayed(modelKey)) {
|
||
continue;
|
||
}
|
||
const userId = String(row.userId || "__anonymous__");
|
||
if (
|
||
!hasModelCapacity(
|
||
modelKey,
|
||
userId,
|
||
maxConcurrentPerModel,
|
||
capacityCandidates
|
||
)
|
||
) {
|
||
markModelAwaitingNextBatch(modelKey);
|
||
continue;
|
||
}
|
||
|
||
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 dispatchSlotReleased = false;
|
||
const releaseDispatchSlot = async () => {
|
||
if (dispatchSlotReleased) return;
|
||
dispatchSlotReleased = true;
|
||
releaseRunningTask(entry.taskId);
|
||
setTimeout(() => {
|
||
void drainCanvasImageTaskQueue();
|
||
}, 0);
|
||
};
|
||
|
||
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;
|
||
const runningTimeoutMs = getCanvasImageRunningTimeoutMs(request);
|
||
const execution: Promise<CanvasImageTaskResult> = isStyleProTaskRequest(request)
|
||
? executeStyleProTask({
|
||
request,
|
||
userId: String(task.userId || entry.userId),
|
||
onRequestDispatched: releaseDispatchSlot,
|
||
})
|
||
: isCanvasLayoutReplaceTaskRequest(request)
|
||
? executeCanvasLayoutReplaceTask({
|
||
request,
|
||
userId: String(task.userId || entry.userId),
|
||
onRequestDispatched: releaseDispatchSlot,
|
||
})
|
||
: generateCanvasImage({
|
||
...(request as CanvasGenerateImageRequest),
|
||
ownerKey: task.userId,
|
||
onRequestDispatched: releaseDispatchSlot,
|
||
});
|
||
const result = await withCanvasImageTaskExecutionTimeout(
|
||
execution,
|
||
runningTimeoutMs,
|
||
entry.taskId
|
||
);
|
||
await db
|
||
.update(canvasImageTask)
|
||
.set({
|
||
status: "succeeded",
|
||
result: result as CanvasImageTaskResult,
|
||
error: null,
|
||
completedAt: new Date(),
|
||
updatedAt: new Date(),
|
||
})
|
||
.where(
|
||
and(eq(canvasImageTask.id, entry.taskId), eq(canvasImageTask.status, "running"))
|
||
);
|
||
} catch (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 {
|
||
await releaseDispatchSlot();
|
||
}
|
||
}
|
||
|
||
export async function drainCanvasImageTaskQueue() {
|
||
const state = getQueueState();
|
||
if (state.draining) return;
|
||
|
||
state.draining = true;
|
||
try {
|
||
await markStaleDatabaseRunningTasksFailed();
|
||
const { concurrency: maxConcurrentPerModel, batchDispatchIntervalMs } =
|
||
await getCanvasImageTaskQueueConfig();
|
||
while (true) {
|
||
const nextTask = await claimNextPendingTask(maxConcurrentPerModel);
|
||
if (!nextTask) break;
|
||
if (state.running.has(nextTask.taskId)) continue;
|
||
markRunningTask(nextTask.taskId, nextTask.modelKey, nextTask.userId);
|
||
scheduleModelNextBatch(nextTask.modelKey, batchDispatchIntervalMs);
|
||
void runSingleTask(nextTask);
|
||
}
|
||
} finally {
|
||
state.draining = false;
|
||
}
|
||
}
|
||
|
||
export async function enqueueCanvasImageTask(taskId: string) {
|
||
await enqueueLocalTask(taskId);
|
||
await drainCanvasImageTaskQueue();
|
||
}
|
||
|
||
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;
|
||
}
|
||
|
||
function isCanvasGenerateImageResult(
|
||
value: unknown
|
||
): value is CanvasGenerateImageResult {
|
||
if (!value || typeof value !== "object" || Array.isArray(value)) {
|
||
return false;
|
||
}
|
||
|
||
const result = value as Partial<CanvasGenerateImageResult>;
|
||
return (
|
||
Array.isArray(result.images) &&
|
||
result.images.length > 0 &&
|
||
result.images.every(
|
||
(image) =>
|
||
Boolean(image) &&
|
||
typeof image.originalUrl === "string" &&
|
||
typeof image.thumbnailUrl === "string"
|
||
)
|
||
);
|
||
}
|
||
|
||
export async function generateCanvasImageWithPersistentQueue(params: {
|
||
userId: string;
|
||
request: CanvasGenerateImageRequest;
|
||
timeoutMs?: number;
|
||
onTaskUpdate?: (
|
||
update: CanvasGenerateImageTaskUpdate
|
||
) => void | Promise<void>;
|
||
}): Promise<CanvasGenerateImageResult> {
|
||
const {
|
||
onQueueStatusChange,
|
||
onRequestDispatched: _requestDispatchCallback,
|
||
ownerKey: _requestOwnerKey,
|
||
...serializableRequest
|
||
} = params.request;
|
||
const task = await createCanvasImageTask({
|
||
userId: params.userId,
|
||
request: serializableRequest as unknown as Record<string, unknown>,
|
||
});
|
||
await enqueueCanvasImageTask(task.id);
|
||
|
||
const timeoutMs = Math.max(60_000, params.timeoutMs ?? 30 * 60 * 1000);
|
||
const startedAt = Date.now();
|
||
let lastStatusSignature = "";
|
||
|
||
while (Date.now() - startedAt < timeoutMs) {
|
||
const currentTask = await getCanvasImageTaskForUser({
|
||
taskId: task.id,
|
||
userId: params.userId,
|
||
});
|
||
if (!currentTask) {
|
||
throw new Error(`图片生成任务不存在（${task.id}）`);
|
||
}
|
||
|
||
const status = currentTask.status as CanvasImageTaskStatus;
|
||
const queueStatus = currentTask.queueStatus;
|
||
const statusSignature = [
|
||
status,
|
||
queueStatus.queuePosition,
|
||
queueStatus.sameModelPendingCount,
|
||
queueStatus.sameModelRunningCount,
|
||
queueStatus.concurrencyLimit,
|
||
queueStatus.dispatchIntervalMs,
|
||
].join(":");
|
||
if (statusSignature !== lastStatusSignature) {
|
||
lastStatusSignature = statusSignature;
|
||
await params.onTaskUpdate?.({
|
||
taskId: task.id,
|
||
status,
|
||
queueStatus,
|
||
error: currentTask.error,
|
||
});
|
||
if (status === "pending" || status === "running") {
|
||
await onQueueStatusChange?.({
|
||
...queueStatus,
|
||
modelName: serializableRequest.model || null,
|
||
modelRef: serializableRequest.model || null,
|
||
});
|
||
}
|
||
}
|
||
|
||
if (status === "succeeded") {
|
||
if (!isCanvasGenerateImageResult(currentTask.result)) {
|
||
throw new Error(`图片生成任务已完成但结果无效（${task.id}）`);
|
||
}
|
||
return currentTask.result;
|
||
}
|
||
if (status === "failed") {
|
||
throw new Error(currentTask.error || "图片生成任务失败");
|
||
}
|
||
|
||
const estimatedWaitMs = Number(queueStatus.estimatedWaitMs || 0);
|
||
const pollAfterMs =
|
||
status === "pending" && estimatedWaitMs > 0
|
||
? Math.min(4000, Math.max(1000, estimatedWaitMs))
|
||
: 1000;
|
||
await new Promise((resolve) => setTimeout(resolve, pollAfterMs));
|
||
}
|
||
|
||
throw new Error(`图片生成任务等待超时（${task.id}）`);
|
||
}
|
||
|
||
async function buildCanvasImageTaskQueueStatus(task: {
|
||
id: string;
|
||
status: CanvasImageTaskStatus;
|
||
request: unknown;
|
||
createdAt: Date;
|
||
}) {
|
||
const modelKey = normalizeTaskModelKey(task.request);
|
||
const {
|
||
concurrency: concurrencyLimit,
|
||
batchDispatchIntervalMs: dispatchIntervalMs,
|
||
} = await getCanvasImageTaskQueueConfig();
|
||
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)
|
||
);
|
||
|
||
if (task.status !== "pending") {
|
||
return {
|
||
status: task.status,
|
||
queuePosition: 0,
|
||
sameModelPendingCount: 0,
|
||
sameModelRunningCount,
|
||
concurrencyLimit,
|
||
dispatchIntervalMs,
|
||
estimatedWaitMs: 0,
|
||
source: "model_queue" as const,
|
||
};
|
||
}
|
||
|
||
const pendingRows = await db
|
||
.select({
|
||
id: canvasImageTask.id,
|
||
request: canvasImageTask.request,
|
||
createdAt: canvasImageTask.createdAt,
|
||
})
|
||
.from(canvasImageTask)
|
||
.where(eq(canvasImageTask.status, "pending"))
|
||
.orderBy(asc(canvasImageTask.createdAt))
|
||
.limit(QUEUE_STATUS_PENDING_LOOKAHEAD_LIMIT);
|
||
|
||
const sameModelPendingRows = pendingRows.filter(
|
||
(row) => normalizeTaskModelKey(row.request) === modelKey
|
||
);
|
||
const queuePosition = Math.max(
|
||
sameModelPendingRows.findIndex((row) => row.id === task.id),
|
||
0
|
||
);
|
||
const sameModelPendingCount = sameModelPendingRows.length;
|
||
const estimatedWaitMs = estimateImageDispatchWaitMs({
|
||
queuePosition,
|
||
dispatchIntervalMs,
|
||
delayRemainingMs: Math.max(
|
||
0,
|
||
Number(getQueueState().nextBatchEligibleAt.get(modelKey) || 0) - Date.now()
|
||
),
|
||
});
|
||
|
||
return {
|
||
status: task.status,
|
||
queuePosition,
|
||
sameModelPendingCount,
|
||
sameModelRunningCount,
|
||
concurrencyLimit,
|
||
dispatchIntervalMs,
|
||
estimatedWaitMs,
|
||
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,
|
||
}),
|
||
};
|
||
}
|
||
|
||
export async function getCanvasDetailReplacementTasksForUserContext(params: {
|
||
userId: string;
|
||
chatId: string;
|
||
projectId?: string | null;
|
||
limit?: number;
|
||
}): Promise<CanvasDetailReplacementTaskContextSnapshot[]> {
|
||
const chatId = String(params.chatId || "").trim();
|
||
const projectId = String(params.projectId || "").trim();
|
||
if (!chatId) {
|
||
return [];
|
||
}
|
||
|
||
const contextConditions = [
|
||
eq(canvasImageTask.userId, params.userId),
|
||
sql`coalesce(${canvasImageTask.request}->'debugMeta'->>'chatId', '') = ${chatId}`,
|
||
sql`coalesce(${canvasImageTask.request}->'debugMeta'->>'operation', '') = 'agent-detail-replacement'`,
|
||
];
|
||
if (projectId) {
|
||
contextConditions.push(
|
||
sql`coalesce(${canvasImageTask.request}->'debugMeta'->>'projectId', '') = ${projectId}`
|
||
);
|
||
}
|
||
|
||
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(...contextConditions))
|
||
.orderBy(desc(canvasImageTask.createdAt))
|
||
.limit(Math.max(1, Math.min(128, Math.floor(params.limit || 128))));
|
||
|
||
const snapshots: Array<CanvasDetailReplacementTaskContextSnapshot | null> = await Promise.all(
|
||
rows.map(async (row) => {
|
||
const request = getRecord(row.request);
|
||
const debugMeta = getRecord(request?.debugMeta);
|
||
const referenceId = String(debugMeta?.referenceId || "").trim();
|
||
const messageId = String(debugMeta?.messageId || "").trim() || null;
|
||
if (!referenceId) {
|
||
return null;
|
||
}
|
||
|
||
const activeTask =
|
||
row.status === "pending" || row.status === "running"
|
||
? await getCanvasImageTaskForUser({
|
||
taskId: row.id,
|
||
userId: params.userId,
|
||
})
|
||
: null;
|
||
const status = (activeTask?.status || row.status) as CanvasImageTaskStatus;
|
||
|
||
return {
|
||
id: row.id,
|
||
status,
|
||
modeType: "detail_replacement" as const,
|
||
referenceId,
|
||
messageId,
|
||
error: activeTask?.error || row.error || null,
|
||
result: activeTask?.result || row.result || null,
|
||
attempts: activeTask?.attempts || row.attempts,
|
||
startedAt: activeTask?.startedAt || row.startedAt,
|
||
completedAt: activeTask?.completedAt || row.completedAt,
|
||
createdAt: activeTask?.createdAt || row.createdAt,
|
||
updatedAt: activeTask?.updatedAt || row.updatedAt,
|
||
queueStatus: activeTask?.queueStatus || null,
|
||
} satisfies CanvasDetailReplacementTaskContextSnapshot;
|
||
})
|
||
);
|
||
|
||
return snapshots.filter(
|
||
(task): task is CanvasDetailReplacementTaskContextSnapshot => task !== null
|
||
);
|
||
}
|