From 3559e8adc6169514594317ca306ad267a10d4c3d Mon Sep 17 00:00:00 2001 From: Marcus <1922576605@qq.com> Date: Sun, 26 Apr 2026 13:20:14 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=E5=9B=BE=E5=BA=8A?= =?UTF-8?q?=E5=BD=92=E6=A1=A3=E4=BB=BB=E5=8A=A1=E6=94=AF=E6=8C=81=EF=BC=8C?= =?UTF-8?q?=E4=BC=98=E5=8C=96=E5=BD=92=E6=A1=A3=E9=80=BB=E8=BE=91=E5=92=8C?= =?UTF-8?q?=E9=94=99=E8=AF=AF=E5=A4=84=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- create_tables.sql | 8 + .../migration.sql | 33 +++ prisma/schema.prisma | 16 ++ server/api/images/generate.post.ts | 81 +----- server/plugins/imageArchiveQueue.ts | 6 + server/utils/imageArchiveQueue.ts | 255 ++++++++++++++++++ server/utils/imageGenerationRecords.ts | 178 +++++++++++- server/utils/index.ts | 1 + server/utils/lsky.ts | 141 ++++++++-- server/utils/openai.ts | 86 +++--- 10 files changed, 667 insertions(+), 138 deletions(-) create mode 100644 prisma/migrations/20260426143000_add_image_archive_queue/migration.sql create mode 100644 server/plugins/imageArchiveQueue.ts create mode 100644 server/utils/imageArchiveQueue.ts diff --git a/create_tables.sql b/create_tables.sql index 256b70f..1688819 100644 --- a/create_tables.sql +++ b/create_tables.sql @@ -28,12 +28,20 @@ CREATE TABLE IF NOT EXISTS `image_generations` ( `upstream_response` JSON NULL COMMENT '完整上游生图接口响应', `hosted_response` JSON NULL COMMENT '完整 Lsky 图床上传接口响应', `error_message` TEXT NULL COMMENT '失败原因或图床归档失败提示', + `archive_status` ENUM('NOT_REQUIRED', 'PENDING', 'RUNNING', 'SUCCEEDED', 'FAILED') NOT NULL DEFAULT 'NOT_REQUIRED' COMMENT '图床归档任务状态', + `archive_attempts` INT NOT NULL DEFAULT 0 COMMENT '图床归档已尝试次数', + `archive_next_run_at` DATETIME(3) NULL COMMENT '图床归档下次可执行时间', + `archive_locked_at` DATETIME(3) NULL COMMENT '图床归档任务被 worker 锁定的时间', + `archive_locked_by` VARCHAR(191) NULL COMMENT '图床归档任务 worker 标识', + `archive_last_error` TEXT NULL COMMENT '图床归档最近一次失败摘要', `created_at` DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3) COMMENT '创建时间', `updated_at` DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3) ON UPDATE CURRENT_TIMESTAMP(3) COMMENT '更新时间', `deleted_at` DATETIME(3) NULL COMMENT '软删除时间', PRIMARY KEY (`id`), INDEX `image_generations_user_id_created_at_idx` (`user_id`, `created_at`), INDEX `image_generations_status_idx` (`status`), + INDEX `image_generations_archive_status_archive_next_run_at_idx` (`archive_status`, `archive_next_run_at`), + INDEX `image_generations_archive_locked_at_idx` (`archive_locked_at`), INDEX `image_generations_deleted_at_idx` (`deleted_at`), CONSTRAINT `image_generations_user_id_fkey` FOREIGN KEY (`user_id`) REFERENCES `users` (`id`) diff --git a/prisma/migrations/20260426143000_add_image_archive_queue/migration.sql b/prisma/migrations/20260426143000_add_image_archive_queue/migration.sql new file mode 100644 index 0000000..2cf06d7 --- /dev/null +++ b/prisma/migrations/20260426143000_add_image_archive_queue/migration.sql @@ -0,0 +1,33 @@ +ALTER TABLE `image_generations` + ADD COLUMN `archive_status` ENUM('NOT_REQUIRED', 'PENDING', 'RUNNING', 'SUCCEEDED', 'FAILED') NOT NULL DEFAULT 'NOT_REQUIRED', + ADD COLUMN `archive_attempts` INT NOT NULL DEFAULT 0, + ADD COLUMN `archive_next_run_at` DATETIME(3) NULL, + ADD COLUMN `archive_locked_at` DATETIME(3) NULL, + ADD COLUMN `archive_locked_by` VARCHAR(191) NULL, + ADD COLUMN `archive_last_error` TEXT NULL; + +UPDATE `image_generations` +SET `archive_status` = 'SUCCEEDED' +WHERE `status` = 'SUCCEEDED' + AND `hosted_image_url` IS NOT NULL; + +UPDATE `image_generations` +SET + `archive_status` = 'PENDING', + `archive_next_run_at` = CURRENT_TIMESTAMP(3) +WHERE `status` = 'SUCCEEDED' + AND `hosted_image_url` IS NULL + AND `error_message` IS NULL + AND `image_url` IS NOT NULL; + +UPDATE `image_generations` +SET `archive_status` = 'FAILED' +WHERE `status` = 'SUCCEEDED' + AND `hosted_image_url` IS NULL + AND `error_message` IS NOT NULL; + +CREATE INDEX `image_generations_archive_status_archive_next_run_at_idx` + ON `image_generations`(`archive_status`, `archive_next_run_at`); + +CREATE INDEX `image_generations_archive_locked_at_idx` + ON `image_generations`(`archive_locked_at`); diff --git a/prisma/schema.prisma b/prisma/schema.prisma index 454c84e..a0e0fd5 100644 --- a/prisma/schema.prisma +++ b/prisma/schema.prisma @@ -19,6 +19,14 @@ enum ImageGenerationStatus { FAILED } +enum ImageArchiveStatus { + NOT_REQUIRED + PENDING + RUNNING + SUCCEEDED + FAILED +} + model User { id Int @id username String @default("") @db.VarChar(191) @@ -49,6 +57,12 @@ model ImageGeneration { upstreamResponse Json? @map("upstream_response") hostedResponse Json? @map("hosted_response") errorMessage String? @map("error_message") @db.Text + archiveStatus ImageArchiveStatus @default(NOT_REQUIRED) @map("archive_status") + archiveAttempts Int @default(0) @map("archive_attempts") + archiveNextRunAt DateTime? @map("archive_next_run_at") + archiveLockedAt DateTime? @map("archive_locked_at") + archiveLockedBy String? @map("archive_locked_by") @db.VarChar(191) + archiveLastError String? @map("archive_last_error") @db.Text createdAt DateTime @default(now()) @map("created_at") updatedAt DateTime @updatedAt @map("updated_at") deletedAt DateTime? @map("deleted_at") @@ -56,6 +70,8 @@ model ImageGeneration { @@index([userId, createdAt]) @@index([status]) + @@index([archiveStatus, archiveNextRunAt]) + @@index([archiveLockedAt]) @@index([deletedAt]) @@map("image_generations") } diff --git a/server/api/images/generate.post.ts b/server/api/images/generate.post.ts index b96cf42..2df4cb3 100644 --- a/server/api/images/generate.post.ts +++ b/server/api/images/generate.post.ts @@ -1,15 +1,9 @@ -// server/api/images/generate.post.ts - 图片生成接口：创建记录、流式调用上游生图并在响应后异步归档图床。 -import { setImmediate } from "node:timers"; +// server/api/images/generate.post.ts - 图片生成接口：创建记录、流式调用上游生图并将图床归档交给数据库队列。 import type { IImageGenerateData, IImageGenerateRequest } from "#shared/types/openai"; -import { - finishImageGenerationArchiveFailed, - finishImageGenerationArchiveSuccess -} from "~~/server/utils"; - -type ApiLogger = ReturnType; +import { wakeImageArchiveWorker } from "~~/server/utils"; /** * POST /api/images/generate @@ -20,8 +14,8 @@ type ApiLogger = ReturnType; * 3. 创建 RUNNING 生图记录，并递增全局请求/进行中统计。 * 4. 在服务端确保并读取 AIArtStudio 完整 key，完整 key 不返回前端。 * 5. 调用 Chat Completions 流式生图接口，累积 SSE delta content 并提取最终图片 URL。 - * 6. 上游返回图片 URL 后立刻把生图结果落库为 SUCCEEDED，并马上返回前端。 - * 7. 响应返回后再用后台异步任务上传 Lsky；归档失败只补写记录和日志，不影响本次响应。 + * 6. 上游返回图片 URL 后立刻把生图结果落库为 SUCCEEDED，并把图床归档任务标记为待执行。 + * 7. 唤醒数据库归档队列 worker；归档失败只补写记录和日志，不影响本次响应。 * 8. 主链路失败时把记录标记为 FAILED；鉴权失败会清理本地登录态并返回 401。 */ export default defineEventHandler(async (event) => { @@ -102,18 +96,7 @@ export default defineEventHandler(async (event) => { archiveScheduled: true }); - const finishedRecord = record; - // TODO: 归档任务现在仍跑在当前 Node 进程里；如果并发继续升高，应该迁移到独立队列/worker， - // 避免大量响应后任务在同一进程内叠加，拖慢其他请求的调度与内存回收。 - setImmediate(() => { - void archiveGeneratedImageInBackground({ - imageUrl: result.imageUrl, - userId, - recordId: finishedRecord.id, - createdAt: finishedRecord.startedAt, - logger - }); - }); + wakeImageArchiveWorker(); return createSuccessResponse( { @@ -152,60 +135,6 @@ export default defineEventHandler(async (event) => { } }); -/** 响应返回后异步归档图片，成功则补写图床字段，失败则补写本地错误文案 */ -const archiveGeneratedImageInBackground = async ({ - imageUrl, - userId, - recordId, - createdAt, - logger -}: { - imageUrl: string; - userId: number; - recordId: bigint; - createdAt: Date; - logger: ApiLogger; -}) => { - logger.info("后台归档开始", { - recordId: recordId.toString() - }); - - try { - const identity = await getUserArchiveIdentity(userId); - const uploaded = await uploadImageFromUrl({ - imageUrl, - userId, - username: identity.username, - recordId, - createdAt - }); - - await finishImageGenerationArchiveSuccess(recordId, { - hostedImageUrl: uploaded.publicUrl, - imageMimeType: uploaded.mimetype, - hostedResponse: uploaded.response - }); - - logger.info("后台归档成功", { - recordId: recordId.toString(), - mimeType: uploaded.mimetype, - hosted: true - }); - } catch (error) { - logger.error("后台归档失败", { - recordId: recordId.toString(), - error: toSafeLogError(error) - }); - - await finishImageGenerationArchiveFailed(recordId).catch((recordError) => { - logger.error("更新归档失败记录失败", { - recordId: recordId.toString(), - error: toSafeLogError(recordError) - }); - }); - } -}; - // 日志只读取流式聚合结果的元信息，不记录完整提示词、key 或图片内容。 const getStreamContentLength = (upstreamResponse: unknown) => { if ( diff --git a/server/plugins/imageArchiveQueue.ts b/server/plugins/imageArchiveQueue.ts new file mode 100644 index 0000000..053d65c --- /dev/null +++ b/server/plugins/imageArchiveQueue.ts @@ -0,0 +1,6 @@ +// server/plugins/imageArchiveQueue.ts - 服务启动时唤起图床归档队列 worker，恢复待处理或锁超时任务。 +import { startImageArchiveWorker } from "~~/server/utils"; + +export default defineNitroPlugin(() => { + startImageArchiveWorker(); +}); diff --git a/server/utils/imageArchiveQueue.ts b/server/utils/imageArchiveQueue.ts new file mode 100644 index 0000000..c687862 --- /dev/null +++ b/server/utils/imageArchiveQueue.ts @@ -0,0 +1,255 @@ +// server/utils/imageArchiveQueue.ts - 图床归档队列 worker：从数据库领取任务、限流上传并按退避策略重试。 +import { randomUUID } from "node:crypto"; +import { clearTimeout, setTimeout } from "node:timers"; +import { consola } from "consola"; +import { + claimRunnableImageArchiveTasks, + finishImageGenerationArchiveFailed, + finishImageGenerationArchiveSuccess, + getUserArchiveIdentity, + scheduleImageGenerationArchiveRetry, + type IImageArchiveTask +} from "~~/server/utils/imageGenerationRecords"; +import { uploadImageFromUrl } from "~~/server/utils/lsky"; + +interface IImageArchiveWorkerConfig { + concurrency: number; + maxAttempts: number; + lockTtlMs: number; + pollMs: number; +} + +interface IImageArchiveWorkerState { + activeCount: number; + draining: boolean; + started: boolean; + timer: ReturnType | null; + workerId: string; +} + +const DEFAULT_CONCURRENCY = 2; +const DEFAULT_MAX_ATTEMPTS = 3; +const DEFAULT_LOCK_TTL_MS = 10 * 60 * 1000; +const DEFAULT_POLL_MS = 15 * 1000; +const RETRY_DELAYS_MS = [30 * 1000, 2 * 60 * 1000, 10 * 60 * 1000]; + +const globalForArchiveWorker = globalThis as unknown as { + imageArchiveWorkerState?: IImageArchiveWorkerState; +}; + +const workerState = + globalForArchiveWorker.imageArchiveWorkerState ?? + (globalForArchiveWorker.imageArchiveWorkerState = { + activeCount: 0, + draining: false, + started: false, + timer: null, + workerId: `archive-${randomUUID()}` + }); + +/** 启动应用内图床归档 worker；重复调用只会唤醒同一个 worker */ +export const startImageArchiveWorker = () => { + if (!workerState.started) { + workerState.started = true; + consola.info("[imageArchiveQueue] worker 启动", { + workerId: workerState.workerId, + concurrency: getImageArchiveWorkerConfig().concurrency + }); + } + + scheduleDrain(0); +}; + +/** 生图成功入队后调用，用于尽快扫描新任务 */ +export const wakeImageArchiveWorker = () => { + if (!workerState.started) return; + scheduleDrain(0); +}; + +const drainImageArchiveQueue = async () => { + if (workerState.draining) return; + + workerState.draining = true; + let nextDelayMs = DEFAULT_POLL_MS; + + try { + const config = getImageArchiveWorkerConfig(); + const availableSlots = config.concurrency - workerState.activeCount; + + if (availableSlots <= 0) { + nextDelayMs = config.pollMs; + return; + } + + const tasks = await claimRunnableImageArchiveTasks({ + limit: availableSlots, + workerId: workerState.workerId, + lockTtlMs: config.lockTtlMs + }); + + if (tasks.length === 0) { + nextDelayMs = config.pollMs; + return; + } + + for (const task of tasks) { + runImageArchiveTask(task, config); + } + + nextDelayMs = + workerState.activeCount < config.concurrency ? 0 : config.pollMs; + } catch (error) { + consola.error("[imageArchiveQueue] 扫描归档任务失败", { + workerId: workerState.workerId, + error: toArchiveLogError(error) + }); + } finally { + workerState.draining = false; + + if (workerState.started) { + scheduleDrain(nextDelayMs); + } + } +}; + +const runImageArchiveTask = ( + task: IImageArchiveTask, + config: IImageArchiveWorkerConfig +) => { + workerState.activeCount += 1; + + void (async () => { + const startedAt = Date.now(); + + consola.info("[imageArchiveQueue] 归档开始", { + workerId: workerState.workerId, + recordId: task.id.toString(), + attempt: task.archiveAttempts + }); + + try { + const identity = await getUserArchiveIdentity(task.userId); + const uploaded = await uploadImageFromUrl({ + imageUrl: task.imageUrl, + userId: task.userId, + username: identity.username, + recordId: task.id, + createdAt: task.startedAt + }); + + await finishImageGenerationArchiveSuccess(task.id, { + hostedImageUrl: uploaded.publicUrl, + imageMimeType: uploaded.mimetype, + hostedResponse: uploaded.response + }); + + consola.info("[imageArchiveQueue] 归档成功", { + workerId: workerState.workerId, + recordId: task.id.toString(), + mimeType: uploaded.mimetype, + byteLength: uploaded.byteLength, + elapsedMs: Date.now() - startedAt + }); + } catch (error) { + await handleImageArchiveFailure(task, config, error); + } finally { + workerState.activeCount = Math.max(0, workerState.activeCount - 1); + wakeImageArchiveWorker(); + } + })(); +}; + +const handleImageArchiveFailure = async ( + task: IImageArchiveTask, + config: IImageArchiveWorkerConfig, + error: unknown +) => { + const safeError = toArchiveLogError(error); + const errorMessage = getArchiveErrorMessage(error); + const shouldStop = task.archiveAttempts >= config.maxAttempts; + + consola.error("[imageArchiveQueue] 归档失败", { + workerId: workerState.workerId, + recordId: task.id.toString(), + attempt: task.archiveAttempts, + terminal: shouldStop, + error: safeError + }); + + if (shouldStop) { + await finishImageGenerationArchiveFailed(task.id, errorMessage); + return; + } + + await scheduleImageGenerationArchiveRetry({ + recordId: task.id, + nextRunAt: new Date(Date.now() + getRetryDelayMs(task.archiveAttempts)), + errorMessage + }); +}; + +const scheduleDrain = (delayMs: number) => { + if (workerState.timer) { + clearTimeout(workerState.timer); + } + + workerState.timer = setTimeout(() => { + workerState.timer = null; + void drainImageArchiveQueue(); + }, delayMs); + workerState.timer.unref?.(); +}; + +const getImageArchiveWorkerConfig = (): IImageArchiveWorkerConfig => { + return { + concurrency: getPositiveIntegerEnv( + "IMAGE_ARCHIVE_CONCURRENCY", + DEFAULT_CONCURRENCY + ), + maxAttempts: getPositiveIntegerEnv( + "IMAGE_ARCHIVE_MAX_ATTEMPTS", + DEFAULT_MAX_ATTEMPTS + ), + lockTtlMs: DEFAULT_LOCK_TTL_MS, + pollMs: DEFAULT_POLL_MS + }; +}; + +const getRetryDelayMs = (attempt: number) => { + const lastDelayMs = RETRY_DELAYS_MS[RETRY_DELAYS_MS.length - 1] ?? 0; + return ( + RETRY_DELAYS_MS[Math.min(attempt - 1, RETRY_DELAYS_MS.length - 1)] ?? + lastDelayMs + ); +}; + +const getPositiveIntegerEnv = (name: string, fallback: number) => { + const value = Number.parseInt(process.env[name] ?? "", 10); + return Number.isInteger(value) && value > 0 ? value : fallback; +}; + +const getArchiveErrorMessage = (error: unknown) => { + if (error instanceof Error && error.message) return error.message; + if (typeof error === "string" && error) return error; + return "图片归档失败"; +}; + +const toArchiveLogError = (error: unknown) => { + const maybeError = error as { + message?: string; + name?: string; + response?: { + status?: number; + }; + status?: number; + statusCode?: number; + statusMessage?: string; + }; + + return { + name: maybeError.name, + message: maybeError.message ?? maybeError.statusMessage, + status: + maybeError.status ?? maybeError.statusCode ?? maybeError.response?.status + }; +}; diff --git a/server/utils/imageGenerationRecords.ts b/server/utils/imageGenerationRecords.ts index e9ca965..4ce133f 100644 --- a/server/utils/imageGenerationRecords.ts +++ b/server/utils/imageGenerationRecords.ts @@ -7,7 +7,11 @@ import type { IImageHistoryListData } from "#shared/types/openai"; import { consola } from "consola"; -import { ImageGenerationStatus, Prisma } from "~~/app/generated/prisma/client"; +import { + ImageArchiveStatus, + ImageGenerationStatus, + Prisma +} from "~~/app/generated/prisma/client"; import { prisma } from "~~/server/utils/prisma"; const GLOBAL_STATS_ID = "global"; @@ -31,6 +35,19 @@ interface IFinishImageGenerationArchiveSuccessInput { hostedResponse: unknown; } +export interface IImageArchiveTask { + /** 生图记录 ID */ + id: bigint; + /** NewAPI 用户 ID */ + userId: number; + /** 上游生成图片 URL */ + imageUrl: string; + /** 记录创建时间，用于构造图床文件名 */ + startedAt: Date; + /** 当前这次归档是第几次尝试 */ + archiveAttempts: number; +} + /** 登录或恢复登录时保存 NewAPI 用户快照 */ export const upsertUserSnapshot = (user: IUserBasicData) => { return prisma.user.upsert({ @@ -135,7 +152,13 @@ export const finishImageGenerationSuccess = async ( revisedPrompt: input.revisedPrompt || null, upstreamResponse: input.upstreamResponse as Prisma.InputJsonValue, hostedResponse: Prisma.DbNull, - errorMessage: null + errorMessage: null, + archiveStatus: ImageArchiveStatus.PENDING, + archiveAttempts: 0, + archiveNextRunAt: new Date(), + archiveLockedAt: null, + archiveLockedBy: null, + archiveLastError: null } }); @@ -182,6 +205,86 @@ export const finishImageGenerationFailed = async ( }); }; +/** 领取可执行的图床归档任务；通过 updateMany 条件锁避免多实例重复处理 */ +export const claimRunnableImageArchiveTasks = async ({ + limit, + workerId, + lockTtlMs +}: { + limit: number; + workerId: string; + lockTtlMs: number; +}): Promise => { + if (limit <= 0) return []; + + const now = new Date(); + const staleBefore = new Date(now.getTime() - lockTtlMs); + const runnableWhere = buildRunnableArchiveWhere(now, staleBefore); + const candidates = await prisma.imageGeneration.findMany({ + where: runnableWhere, + select: { + id: true + }, + orderBy: [ + { + archiveNextRunAt: "asc" + }, + { + createdAt: "asc" + } + ], + take: limit * 3 + }); + const tasks: IImageArchiveTask[] = []; + + for (const candidate of candidates) { + if (tasks.length >= limit) break; + + const claimed = await prisma.imageGeneration.updateMany({ + where: { + id: candidate.id, + ...runnableWhere + }, + data: { + archiveStatus: ImageArchiveStatus.RUNNING, + archiveAttempts: { + increment: 1 + }, + archiveLockedAt: now, + archiveLockedBy: workerId, + archiveLastError: null + } + }); + + if (claimed.count !== 1) continue; + + const task = await prisma.imageGeneration.findUnique({ + where: { + id: candidate.id + }, + select: { + id: true, + userId: true, + imageUrl: true, + startedAt: true, + archiveAttempts: true + } + }); + + if (task?.imageUrl) { + tasks.push({ + id: task.id, + userId: task.userId, + imageUrl: task.imageUrl, + startedAt: task.startedAt, + archiveAttempts: task.archiveAttempts + }); + } + } + + return tasks; +}; + /** 后台归档成功后回写图床地址、MIME 和图床响应 */ export const finishImageGenerationArchiveSuccess = async ( recordId: bigint, @@ -195,12 +298,42 @@ export const finishImageGenerationArchiveSuccess = async ( hostedImageUrl: input.hostedImageUrl, imageMimeType: input.imageMimeType, hostedResponse: input.hostedResponse as Prisma.InputJsonValue, + errorMessage: null, + archiveStatus: ImageArchiveStatus.SUCCEEDED, + archiveNextRunAt: null, + archiveLockedAt: null, + archiveLockedBy: null, + archiveLastError: null + } + }); +}; + +/** 归档可重试失败时释放任务锁并安排下一次执行，前端仍显示归档中 */ +export const scheduleImageGenerationArchiveRetry = async ({ + recordId, + nextRunAt, + errorMessage +}: { + recordId: bigint; + nextRunAt: Date; + errorMessage: string; +}) => { + await prisma.imageGeneration.update({ + where: { + id: recordId + }, + data: { + archiveStatus: ImageArchiveStatus.PENDING, + archiveNextRunAt: nextRunAt, + archiveLockedAt: null, + archiveLockedBy: null, + archiveLastError: errorMessage, errorMessage: null } }); }; -/** 后台归档失败时仅补写本地错误文案，不影响已成功的生图状态 */ +/** 后台归档最终失败时仅补写本地错误文案，不影响已成功的生图状态 */ export const finishImageGenerationArchiveFailed = async ( recordId: bigint, errorMessage: string = "图片归档失败" @@ -210,7 +343,12 @@ export const finishImageGenerationArchiveFailed = async ( id: recordId }, data: { - errorMessage + archiveStatus: ImageArchiveStatus.FAILED, + archiveNextRunAt: null, + archiveLockedAt: null, + archiveLockedBy: null, + archiveLastError: errorMessage, + errorMessage: "图片归档失败" } }); }; @@ -350,6 +488,38 @@ const buildStatsCreateInput = (update: Prisma.GenerationStatsUpdateInput) => { }; }; +/** 只领取成功生成、未删除、已到执行时间或锁超时的归档任务 */ +const buildRunnableArchiveWhere = (now: Date, staleBefore: Date) => { + return { + status: ImageGenerationStatus.SUCCEEDED, + deletedAt: null, + imageUrl: { + not: null + }, + OR: [ + { + archiveStatus: ImageArchiveStatus.PENDING, + OR: [ + { + archiveNextRunAt: null + }, + { + archiveNextRunAt: { + lte: now + } + } + ] + }, + { + archiveStatus: ImageArchiveStatus.RUNNING, + archiveLockedAt: { + lt: staleBefore + } + } + ] + } satisfies Prisma.ImageGenerationWhereInput; +}; + /** 从 Prisma increment 操作里取出增量；非 increment 操作在创建时按 0 处理 */ const getIncrementValue = (value: unknown): number => { if ( diff --git a/server/utils/index.ts b/server/utils/index.ts index 9466f68..b52bb25 100644 --- a/server/utils/index.ts +++ b/server/utils/index.ts @@ -1,6 +1,7 @@ // server/utils/index.ts - 服务端工具统一导出口，方便 API handler 从同一入口导入。 export * from "./createApiResponse"; export * from "./fetch"; +export * from "./imageArchiveQueue"; export * from "./imageGenerationRecords"; export * from "./logging"; export * from "./lsky"; diff --git a/server/utils/lsky.ts b/server/utils/lsky.ts index beda153..d9d984a 100644 --- a/server/utils/lsky.ts +++ b/server/utils/lsky.ts @@ -27,9 +27,15 @@ export interface ILskyUploadedImage { publicUrl: string; filename: string; mimetype: string; + byteLength: number; response: ILskyUploadResponse; } +/** 单张归档图片默认最大 20MB，避免多用户并发时把图片二进制堆满进程内存 */ +const DEFAULT_ARCHIVE_MAX_BYTES = 20 * 1024 * 1024; +/** 单次下载或上传默认 60 秒超时，避免慢连接长期占用归档 worker 槽位 */ +const DEFAULT_ARCHIVE_TIMEOUT_MS = 60 * 1000; + /** * 从远程图片 URL 下载图片并上传到 Lsky。 * @@ -42,7 +48,8 @@ export const uploadImageFromUrl = async ( input: ILskyUploadInput ): Promise => { const config = getLskyConfig(); - const downloadedImage = await downloadImage(input.imageUrl); + const protection = getArchiveProtectionConfig(); + const downloadedImage = await downloadImage(input.imageUrl, protection); const extension = getImageExtension(downloadedImage.mimeType, input.imageUrl); const filename = buildArchiveFilename({ userId: input.userId, @@ -68,14 +75,11 @@ export const uploadImageFromUrl = async ( formData.append("tags[]", tag); } - const response = await $fetch("/upload", { - baseURL: config.baseUrl, - method: "POST", - body: formData, - headers: { - Accept: "application/json", - Authorization: `Bearer ${config.token}` - } + const response = await uploadToLsky({ + baseUrl: config.baseUrl, + token: config.token, + timeoutMs: protection.timeoutMs, + formData }); const publicUrl = response.data?.public_url; @@ -90,6 +94,7 @@ export const uploadImageFromUrl = async ( publicUrl, filename: response.data?.filename || filename, mimetype: response.data?.mimetype || downloadedImage.mimeType, + byteLength: downloadedImage.byteLength, response }; }; @@ -119,21 +124,108 @@ const getLskyConfig = () => { }; }; -/** 下载上游图片二进制，保留 content-type 供上传和入库使用 */ -const downloadImage = async (imageUrl: string) => { - const response = await fetch(imageUrl); - if (!response.ok) { - throw new Error(`图片下载失败：${response.status}`); - } - - // TODO: 当前会把整张图片一次性读入当前 Node 进程内存，再转 Blob 上传到图床； - // 如果图片尺寸或并发继续增大，建议改成流式转发或落临时文件，降低内存峰值和主线程抖动。 +/** 读取归档保护参数，限制单次下载/上传的大小和耗时 */ +const getArchiveProtectionConfig = () => { return { - bytes: await response.arrayBuffer(), - mimeType: response.headers.get("content-type") || "image/png" + maxBytes: getPositiveIntegerEnv( + "IMAGE_ARCHIVE_MAX_BYTES", + DEFAULT_ARCHIVE_MAX_BYTES + ), + timeoutMs: getPositiveIntegerEnv( + "IMAGE_ARCHIVE_TIMEOUT_MS", + DEFAULT_ARCHIVE_TIMEOUT_MS + ) }; }; +/** 调用 Lsky 上传接口，超时后主动取消请求，避免 worker 长时间占位 */ +const uploadToLsky = async ({ + baseUrl, + token, + timeoutMs, + formData +}: { + baseUrl: string; + token: string; + timeoutMs: number; + formData: FormData; +}) => { + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), timeoutMs); + + try { + return await $fetch("/upload", { + baseURL: baseUrl, + method: "POST", + body: formData, + signal: controller.signal, + headers: { + Accept: "application/json", + Authorization: `Bearer ${token}` + } + }); + } catch (error) { + if (isAbortError(error)) { + throw new Error("图床上传超时"); + } + + throw error; + } finally { + clearTimeout(timeout); + } +}; + +/** 下载上游图片二进制，保留 content-type 供上传和入库使用 */ +const downloadImage = async ( + imageUrl: string, + protection: ReturnType +) => { + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), protection.timeoutMs); + + try { + const response = await fetch(imageUrl, { + signal: controller.signal + }); + if (!response.ok) { + throw new Error(`图片下载失败：${response.status}`); + } + + const contentLength = getContentLength(response.headers); + if (contentLength !== null && contentLength > protection.maxBytes) { + throw new Error("图片文件过大，归档已停止"); + } + + const bytes = await response.arrayBuffer(); + if (bytes.byteLength > protection.maxBytes) { + throw new Error("图片文件过大，归档已停止"); + } + + return { + bytes, + byteLength: bytes.byteLength, + mimeType: response.headers.get("content-type") || "image/png" + }; + } catch (error) { + if (isAbortError(error)) { + throw new Error("图片下载超时"); + } + + throw error; + } finally { + clearTimeout(timeout); + } +}; + +/** content-length 存在且合法时用于提前拒绝超大图片 */ +const getContentLength = (headers: Headers) => { + const value = headers.get("content-length"); + if (!value) return null; + + const parsed = Number.parseInt(value, 10); + return Number.isInteger(parsed) && parsed >= 0 ? parsed : null; +}; + /** 构造图床文件名：userId_username_timestamp_recordId.ext */ const buildArchiveFilename = ({ userId, @@ -196,3 +288,12 @@ const getImageExtension = (mimeType: string, imageUrl: string) => { const isSuccessStatus = (status: unknown) => { return status === true || status === "success" || status === "ok"; }; + +const getPositiveIntegerEnv = (name: string, fallback: number) => { + const value = Number.parseInt(process.env[name] ?? "", 10); + return Number.isInteger(value) && value > 0 ? value : fallback; +}; + +const isAbortError = (error: unknown) => { + return error instanceof Error && error.name === "AbortError"; +}; diff --git a/server/utils/openai.ts b/server/utils/openai.ts index c0f47ed..9cc9524 100644 --- a/server/utils/openai.ts +++ b/server/utils/openai.ts @@ -5,6 +5,8 @@ import type { BaseOptions, IImageGenerateData } from "#shared/types/openai"; /** 图像生成改走 Chat Completions 流式接口，通常能更快拿到中转平台返回的图片地址 */ const CHAT_COMPLETIONS_URL = "https://api.qflink.xyz/v1/chat/completions"; const IMAGE_GENERATION_MODEL = "gpt-image-2"; +const DEFAULT_MAX_SSE_DATA_LINE_BYTES = 1024 * 1024; +const textEncoder = new TextEncoder(); /** Chat Completions SSE 每个 data chunk 的最小结构 */ interface IChatCompletionStreamChunk { @@ -28,15 +30,9 @@ interface IChatCompletionStreamChunk { interface IImageStreamUpstreamResponse { content: string; imageUrl?: string; - chunks: Array<{ - id?: string; - created?: number; - model?: string; - content?: string; - finishReason?: string | null; - metadata?: unknown; - hasUsage: boolean; - }>; + contentLength: number; + chunkCount: number; + finishReason?: string | null; usage?: unknown; } @@ -221,11 +217,14 @@ const readChatCompletionStream = async ( ): Promise => { const reader = stream.getReader(); const decoder = new TextDecoder(); - // TODO: 这里会把完整 SSE 文本和 chunk 元数据都累计到当前 Node 进程内存里， - // 高并发或超长响应下可能带来明显的事件循环抖动；后续可改为限制缓存窗口或只保留摘要。 - const chunks: IImageStreamUpstreamResponse["chunks"] = []; + const maxDataLineBytes = getPositiveIntegerEnv( + "IMAGE_STREAM_MAX_SSE_LINE_BYTES", + DEFAULT_MAX_SSE_DATA_LINE_BYTES + ); let buffer = ""; let content = ""; + let chunkCount = 0; + let finishReason: string | null | undefined; let usage: unknown; while (true) { @@ -233,12 +232,14 @@ const readChatCompletionStream = async ( if (done) break; buffer += decoder.decode(value, { stream: true }); + assertSseBufferSize(buffer, maxDataLineBytes); const lines = buffer.split(/\r?\n/); buffer = lines.pop() ?? ""; for (const line of lines) { - const chunk = parseSseDataLine(line); + const chunk = parseSseDataLine(line, maxDataLineBytes); if (!chunk) continue; + chunkCount += 1; const deltaContent = collectDeltaContent(chunk); if (deltaContent) { @@ -249,23 +250,17 @@ const readChatCompletionStream = async ( usage = chunk.usage; } - chunks.push({ - id: chunk.id, - created: chunk.created, - model: chunk.model, - content: deltaContent || undefined, - finishReason: chunk.choices?.find((choice) => choice.finish_reason) - ?.finish_reason, - metadata: chunk.metadata, - hasUsage: Boolean(chunk.usage) - }); + finishReason = + chunk.choices?.find((choice) => choice.finish_reason) + ?.finish_reason ?? finishReason; } } const finalText = buffer + decoder.decode(); for (const line of finalText.split(/\r?\n/)) { - const chunk = parseSseDataLine(line); + const chunk = parseSseDataLine(line, maxDataLineBytes); if (!chunk) continue; + chunkCount += 1; const deltaContent = collectDeltaContent(chunk); if (deltaContent) { @@ -276,35 +271,35 @@ const readChatCompletionStream = async ( usage = chunk.usage; } - chunks.push({ - id: chunk.id, - created: chunk.created, - model: chunk.model, - content: deltaContent || undefined, - finishReason: chunk.choices?.find((choice) => choice.finish_reason) - ?.finish_reason, - metadata: chunk.metadata, - hasUsage: Boolean(chunk.usage) - }); + finishReason = + chunk.choices?.find((choice) => choice.finish_reason)?.finish_reason ?? + finishReason; } return { content, - chunks, + contentLength: content.length, + chunkCount, + finishReason, usage }; }; /** 解析单行 SSE data，跳过空行和 [DONE] */ -const parseSseDataLine = (line: string): IChatCompletionStreamChunk | null => { +const parseSseDataLine = ( + line: string, + maxDataLineBytes: number +): IChatCompletionStreamChunk | null => { const trimmed = line.trim(); if (!trimmed.startsWith("data:")) return null; const data = trimmed.slice("data:".length).trim(); if (!data || data === "[DONE]") return null; - // TODO: JSON.parse 仍运行在主线程；如果后续接入更大的 SSE chunk，需要增加单行大小保护， - // 避免异常大包在单个 tick 内长时间占用事件循环。 + if (textEncoder.encode(data).byteLength > maxDataLineBytes) { + throw new Error("图片生成失败：流式响应过大"); + } + return JSON.parse(data) as IChatCompletionStreamChunk; }; @@ -322,3 +317,18 @@ const collectDeltaContent = (chunk: IChatCompletionStreamChunk) => { const cleanupImageUrl = (url: string) => { return url.replace(/[,.!?，。！？]+$/u, ""); }; + +/** 限制尚未切分出换行的 SSE 缓冲区，避免异常长 data 行持续堆在内存里 */ +const assertSseBufferSize = (buffer: string, maxDataLineBytes: number) => { + if ( + !buffer.includes("\n") && + textEncoder.encode(buffer).byteLength > maxDataLineBytes + ) { + throw new Error("图片生成失败：流式响应过大"); + } +}; + +const getPositiveIntegerEnv = (name: string, fallback: number) => { + const value = Number.parseInt(process.env[name] ?? "", 10); + return Number.isInteger(value) && value > 0 ? value : fallback; +};