返回 oh-my-ppt
handlers.ts
根目录 / src / main / image-generation / handlers.ts
1 import { ipcMain } from 'electron'
2 import fs from 'fs'
3 import path from 'path'
4 import { nanoid } from 'nanoid'
5 import log from 'electron-log/main.js'
6 import type {
7 GeneratedImageAsset,
8 ImageGenerationHistoryRecord,
9 ImageModelProvider
10 } from '@shared/image-generation'
11 import { resolveModelTimeoutMs } from '@shared/model-timeout'
12 import type { IpcContext } from '../ipc/context'
13 import { readAppLocale, uiText, type AppLocale } from '../config/locale-utils'
14 import {
15 resolveGlobalModelTimeouts,
16 resolveModelConfigForTask,
17 type ActiveModelConfig
18 } from '../config/model-config-utils'
19 import { allowLocalAssetRoot } from '../io/local-asset-roots'
20 import { extractModelText, resolveModel } from '../agent-runtime/model'
21 import { buildImagePromptGenerationMessages } from '../agent-runtime/prompt'
22 import {
23 resolveImageGenerationProvider,
24 type ResolvedImageModelConfig
25 } from '../agent-runtime/provider/image'
26 import {
27 imageHistoryLockKey,
28 JobCoordinator,
29 type TypedEventBus
30 } from '../agent-runtime'
31 import {
32 getImageRunState,
33 setImageRunState,
34 type ImageRunState,
35 type ImageRunStatus
36 } from './run-state'
37
38 const VALID_IMAGE_PROVIDERS = [
39 'jimeng',
40 'jimeng4',
41 'agnes',
42 'siliconflow',
43 'openaiCompatible',
44 'gemini',
45 'seedream'
46 ] as const
47
48 const resolveProvider = (provider: unknown): ImageModelProvider => {
49 if (VALID_IMAGE_PROVIDERS.includes(provider as ImageModelProvider)) {
50 return provider as ImageModelProvider
51 }
52 throw new Error('Unsupported image provider')
53 }
54
55 const readJsonObject = (value: unknown): Record<string, unknown> => {
56 const text = typeof value === 'string' ? value.trim() : ''
57 if (!text) return {}
58 try {
59 const parsed = JSON.parse(text)
60 return parsed && typeof parsed === 'object' && !Array.isArray(parsed)
61 ? (parsed as Record<string, unknown>)
62 : {}
63 } catch {
64 return {}
65 }
66 }
67
68 const readConfigString = (config: ResolvedImageModelConfig, key: string): string => {
69 const value = config.modelConfig[key]
70 return typeof value === 'string' ? value.trim() : ''
71 }
72
73 const resolveDisplayModel = (config: ResolvedImageModelConfig): string =>
74 readConfigString(config, 'model') || readConfigString(config, 'reqKey') || config.provider
75
76 const resolveImageModelConfig = async (
77 ctx: IpcContext,
78 modelConfigId?: string
79 ): Promise<ResolvedImageModelConfig> => {
80 const raw =
81 modelConfigId && modelConfigId.trim().length > 0
82 ? await ctx.db.getImageModelConfig(modelConfigId.trim())
83 : await ctx.db.getActiveImageModelConfig()
84 if (!raw) throw new Error('请先在设置中配置并启用生图模型。')
85 return {
86 id: raw.id,
87 name: raw.name,
88 provider: resolveProvider(raw.provider),
89 active: raw.active === 1,
90 modelConfig: readJsonObject(ctx.decryptApiKey(raw.modelConfig || '{}'))
91 }
92 }
93
94 const resolvePageContext = async (
95 ctx: IpcContext,
96 sessionId: string,
97 pageId: string
98 ): Promise<{ pageId: string; title: string; contentOutline: string; htmlPath: string }> => {
99 const pages = await ctx.db.listSessionPages(sessionId)
100 const page = pages.find(
101 (item) => item.id === pageId || item.file_slug === pageId || item.legacy_page_id === pageId
102 )
103 if (!page) throw new Error('请先选择一个可用页面。')
104 const snapshots = await ctx.db.listLatestGenerationPageSnapshot(sessionId)
105 const snapshot = snapshots.find(
106 (item) => item.page_id === page.id || item.page_id === page.file_slug || item.page_id === page.legacy_page_id
107 )
108 return {
109 pageId: page.id,
110 title: page.title || snapshot?.title || `Page ${page.page_number}`,
111 contentOutline: snapshot?.content_outline || '',
112 htmlPath: page.html_path
113 }
114 }
115
116 const sanitizeExt = (extension: string): string =>
117 /^\.[a-z0-9]{2,5}$/i.test(extension) ? extension.toLowerCase() : '.png'
118
119 const compactPageHtmlForPrompt = (html: string): string =>
120 html
121 .replace(/<script\b[^>]*>[\s\S]*?<\/script>/gi, '')
122 .replace(/<!--[\s\S]*?-->/g, '')
123 .replace(/\s+/g, ' ')
124 .trim()
125 .slice(0, 24000)
126
127 const normalizeGeneratedImagePrompt = (raw: string): string =>
128 raw
129 .replace(/^```(?:text|markdown|md)?\s*/i, '')
130 .replace(/```$/i, '')
131 .replace(/^\s*(?:prompt|提示词)\s*[::]\s*/i, '')
132 .trim()
133
134 const resolvePromptModelConfig = async (
135 ctx: IpcContext,
136 modelConfigId?: string
137 ): Promise<ActiveModelConfig> => {
138 return resolveModelConfigForTask(ctx, {
139 modelConfigId,
140 purpose: 'images:generatePrompt'
141 })
142 }
143
144 export function registerImageGenerationHandlers(
145 ctx: IpcContext,
146 coordinator: JobCoordinator,
147 runtimeEvents: TypedEventBus
148 ): void {
149 const emitImageProgress = (state: ImageRunState): void => {
150 runtimeEvents.emit({
151 type: 'image.progress',
152 payload: {
153 runId: state.runId,
154 sessionId: state.sessionId,
155 pageId: state.pageId,
156 progress: state.progress,
157 label: state.label,
158 status: state.status
159 },
160 jobId: state.runId,
161 domain: 'image',
162 owner: { sessionId: state.sessionId, imageHistoryOwner: state.sessionId },
163 audience: { kind: 'broadcast' },
164 occurredAt: state.updatedAt
165 })
166 }
167
168 const updateImageRunState = (state: ImageRunState): void => {
169 setImageRunState(state)
170 emitImageProgress(state)
171 }
172
173 const emitImageJobTerminalEvent = (args: {
174 runId: string
175 sessionId: string
176 status: Exclude<ImageRunStatus, 'running'>
177 errorMessage?: string
178 }): void => {
179 runtimeEvents.emit({
180 type:
181 args.status === 'completed'
182 ? 'job.completed'
183 : args.status === 'cancelled'
184 ? 'job.cancelled'
185 : 'job.failed',
186 payload:
187 args.status === 'completed'
188 ? {}
189 : args.status === 'cancelled'
190 ? { reason: 'user' }
191 : {
192 errorCode: 'image_generation_failed',
193 errorMessage: args.errorMessage || 'Image generation failed'
194 },
195 jobId: args.runId,
196 domain: 'image',
197 owner: { sessionId: args.sessionId, imageHistoryOwner: args.sessionId },
198 audience: { kind: 'broadcast' },
199 occurredAt: Date.now()
200 })
201 }
202
203 const throwIfCancelled = (signal: AbortSignal, locale: AppLocale): void => {
204 if (signal.aborted) {
205 throw new Error(uiText(locale, '已取消生图', 'Image generation cancelled'))
206 }
207 }
208
209 ipcMain.handle('images:generatePrompt', async (_event, payload) => {
210 const locale = await readAppLocale(ctx)
211 const record =
212 payload && typeof payload === 'object' ? (payload as Record<string, unknown>) : {}
213 const sessionId = typeof record.sessionId === 'string' ? record.sessionId.trim() : ''
214 const htmlPath = typeof record.htmlPath === 'string' ? record.htmlPath.trim() : ''
215 if (!sessionId) throw new Error(uiText(locale, '会话 ID 不能为空。', 'Session ID is required.'))
216 if (!htmlPath) {
217 throw new Error(
218 uiText(locale, '当前页文件地址不能为空。', 'Current page file path is required.')
219 )
220 }
221
222 const safeHtmlPath = await ctx.assertPathInAllowedRoots({
223 filePath: htmlPath,
224 mode: 'read',
225 sessionId,
226 htmlOnly: true
227 })
228 const pageHtml = compactPageHtmlForPrompt(await fs.promises.readFile(safeHtmlPath, 'utf-8'))
229 if (!pageHtml) {
230 throw new Error(uiText(locale, '当前页内容为空。', 'Current page content is empty.'))
231 }
232
233 const activeModel = await resolvePromptModelConfig(
234 ctx,
235 typeof record.modelConfigId === 'string' ? record.modelConfigId : undefined
236 )
237 const modelTimeouts = await resolveGlobalModelTimeouts(ctx)
238 const timeoutMs = resolveModelTimeoutMs(modelTimeouts.agent, 'agent')
239 const model = resolveModel(
240 activeModel.provider,
241 activeModel.apiKey,
242 activeModel.model,
243 activeModel.baseUrl,
244 0.45,
245 activeModel.maxTokens,
246 ctx.modelRuntime
247 )
248 const userPrompt = typeof record.userPrompt === 'string' ? record.userPrompt.trim() : ''
249 const pageTitle = typeof record.pageTitle === 'string' ? record.pageTitle.trim() : ''
250 const pageOutline = typeof record.pageOutline === 'string' ? record.pageOutline.trim() : ''
251 log.info('[images:generatePrompt] start', {
252 sessionId,
253 htmlPath: safeHtmlPath,
254 modelConfigId: activeModel.id,
255 model: activeModel.model,
256 htmlLength: pageHtml.length,
257 userPromptLength: userPrompt.length,
258 pageTitleLength: pageTitle.length,
259 pageOutlineLength: pageOutline.length
260 })
261
262 const response = await model.invoke(
263 buildImagePromptGenerationMessages({
264 locale,
265 userPrompt,
266 pageTitle,
267 pageOutline,
268 pageHtml
269 }),
270 { signal: AbortSignal.timeout(timeoutMs) }
271 )
272 const prompt = normalizeGeneratedImagePrompt(extractModelText(response))
273 if (!prompt) {
274 throw new Error(uiText(locale, '模型未返回提示词。', 'The model returned an empty prompt.'))
275 }
276 log.info('[images:generatePrompt] completed', {
277 sessionId,
278 promptLength: prompt.length
279 })
280 return { prompt }
281 })
282
283 ipcMain.handle('images:generate', async (_event, payload) => {
284 const locale = await readAppLocale(ctx)
285 const record = payload && typeof payload === 'object' ? (payload as Record<string, unknown>) : {}
286 const sessionId = typeof record.sessionId === 'string' ? record.sessionId.trim() : ''
287 const pageId = typeof record.pageId === 'string' ? record.pageId.trim() : ''
288 const prompt = typeof record.prompt === 'string' ? record.prompt.trim() : ''
289 if (!sessionId) throw new Error(uiText(locale, '会话 ID 不能为空。', 'Session ID is required.'))
290 if (!pageId) throw new Error(uiText(locale, '请先选择页面。', 'Select a page first.'))
291 if (!prompt) throw new Error(uiText(locale, '请先填写图片描述。', 'Enter an image prompt first.'))
292
293 const count =
294 typeof record.count === 'number' && record.count > 0
295 ? Math.min(Math.floor(record.count), 4)
296 : 1
297 const size = typeof record.size === 'string' && record.size.trim() ? record.size.trim() : '16:9'
298 const runId = nanoid(12)
299 const reservation = await coordinator.reserve({
300 jobId: runId,
301 domain: 'image',
302 owner: { kind: 'image-history', id: sessionId },
303 claims: { write: [imageHistoryLockKey(sessionId)] },
304 wait: 'fail'
305 })
306 if (reservation.status === 'busy') {
307 throw new Error(
308 uiText(
309 locale,
310 '当前会话已有图片生成任务正在进行。',
311 'An image generation task is already running for this session.'
312 )
313 )
314 }
315
316 const lease = reservation.lease
317 const startedAt = Date.now()
318 let resolvedPageId = pageId
319 let resolvedProvider: ImageModelProvider | undefined
320 const uncommittedImagePaths: string[] = []
321 let historyCommitted = false
322 const initialState: ImageRunState = {
323 runId,
324 sessionId,
325 pageId,
326 progress: 5,
327 label: uiText(locale, '准备生图', 'Preparing image generation'),
328 status: 'running',
329 updatedAt: Date.now()
330 }
331 setImageRunState(initialState)
332 runtimeEvents.emit({
333 type: 'job.started',
334 payload: {},
335 jobId: runId,
336 domain: 'image',
337 owner: { sessionId, imageHistoryOwner: sessionId },
338 audience: { kind: 'broadcast' },
339 occurredAt: startedAt
340 })
341 emitImageProgress(initialState)
342
343 try {
344 const modelConfig = await resolveImageModelConfig(
345 ctx,
346 typeof record.imageModelConfigId === 'string'
347 ? record.imageModelConfigId
348 : typeof record.modelConfigId === 'string'
349 ? record.modelConfigId
350 : undefined
351 )
352 throwIfCancelled(lease.signal, locale)
353 const pageContext = await resolvePageContext(ctx, sessionId, pageId)
354 throwIfCancelled(lease.signal, locale)
355 resolvedPageId = pageContext.pageId
356 resolvedProvider = modelConfig.provider
357 const displayModel = resolveDisplayModel(modelConfig)
358 log.info('[images:generate] start', {
359 runId,
360 sessionId,
361 pageId: pageContext.pageId,
362 requestedPageId: pageId,
363 pageTitle: pageContext.title,
364 modelConfigId: modelConfig.id,
365 modelConfigName: modelConfig.name,
366 provider: modelConfig.provider,
367 model: displayModel,
368 count,
369 size,
370 promptLength: prompt.length,
371 negativePromptLength:
372 typeof record.negativePrompt === 'string' ? record.negativePrompt.length : 0,
373 hasSeed: typeof record.seed === 'number'
374 })
375 const adapter = resolveImageGenerationProvider(modelConfig.provider)
376 const results = await adapter.generate(modelConfig, {
377 prompt,
378 count,
379 size,
380 negativePrompt: typeof record.negativePrompt === 'string' ? record.negativePrompt : undefined,
381 seed: typeof record.seed === 'number' ? record.seed : undefined,
382 signal: lease.signal
383 })
384 throwIfCancelled(lease.signal, locale)
385 log.info('[images:generate] provider returned', {
386 runId,
387 provider: modelConfig.provider,
388 resultCount: results.length,
389 elapsedMs: Date.now() - startedAt
390 })
391 updateImageRunState({
392 runId,
393 sessionId,
394 pageId: pageContext.pageId,
395 progress: 80,
396 label: uiText(locale, '正在保存图片', 'Saving images'),
397 status: 'running',
398 updatedAt: Date.now()
399 })
400 const projectDir = await ctx.resolveSessionProjectDir(sessionId)
401 throwIfCancelled(lease.signal, locale)
402 const imagesDir = path.join(projectDir, 'images')
403 log.info('[images:generate] save start', {
404 runId,
405 imagesDir,
406 resultCount: results.length
407 })
408 await fs.promises.mkdir(imagesDir, { recursive: true })
409 allowLocalAssetRoot(imagesDir)
410 const createdAt = Math.floor(Date.now() / 1000)
411 const assets: GeneratedImageAsset[] = []
412 for (const result of results) {
413 throwIfCancelled(lease.signal, locale)
414 const id = nanoid(10)
415 const extension = sanitizeExt(result.extension)
416 const fileName = `${ctx.toSafeAssetBaseName(`generated-${pageContext.title}`)}-${id}${extension}`
417 const absolutePath = path.join(imagesDir, fileName)
418 await fs.promises.writeFile(absolutePath, result.bytes)
419 uncommittedImagePaths.push(absolutePath)
420 const stat = await fs.promises.stat(absolutePath)
421 log.info('[images:generate] asset saved', {
422 runId,
423 assetId: id,
424 fileName,
425 mimeType: result.mimeType,
426 size: stat.size
427 })
428 assets.push({
429 id,
430 fileName,
431 originalName: fileName,
432 relativePath: `./images/${fileName}`,
433 absolutePath,
434 mimeType: result.mimeType,
435 size: stat.size,
436 prompt,
437 modelConfigId: modelConfig.id,
438 provider: modelConfig.provider,
439 model: displayModel,
440 pageId: pageContext.pageId,
441 createdAt
442 })
443 }
444 throwIfCancelled(lease.signal, locale)
445 const historyId = await ctx.db.insertImageGenerationHistory({
446 sessionId,
447 pageId: pageContext.pageId,
448 prompt,
449 imagePaths: assets.map((asset) => asset.relativePath),
450 modelConfigId: modelConfig.id,
451 provider: modelConfig.provider,
452 model: displayModel,
453 createdAt
454 })
455 historyCommitted = true
456 const history: ImageGenerationHistoryRecord = {
457 id: historyId,
458 sessionId,
459 pageId: pageContext.pageId,
460 prompt,
461 imagePaths: assets.map((asset) => asset.relativePath),
462 assets,
463 modelConfigId: modelConfig.id,
464 provider: modelConfig.provider,
465 model: displayModel,
466 createdAt
467 }
468 log.info('[images:generate] history saved', {
469 runId,
470 historyId,
471 imagePathCount: history.imagePaths.length
472 })
473 updateImageRunState({
474 runId,
475 sessionId,
476 pageId: pageContext.pageId,
477 progress: 100,
478 label: uiText(locale, '生图完成', 'Image generation completed'),
479 status: 'completed',
480 updatedAt: Date.now()
481 })
482 emitImageJobTerminalEvent({ runId, sessionId, status: 'completed' })
483 log.info('[images:generate] completed', {
484 runId,
485 sessionId,
486 pageId: pageContext.pageId,
487 assetCount: assets.length,
488 elapsedMs: Date.now() - startedAt
489 })
490 return { history }
491 } catch (error) {
492 if (!historyCommitted && uncommittedImagePaths.length > 0) {
493 const cleanupResults = await Promise.allSettled(
494 uncommittedImagePaths.map((filePath) => fs.promises.rm(filePath, { force: true }))
495 )
496 const cleanupFailures = cleanupResults.filter((result) => result.status === 'rejected')
497 if (cleanupFailures.length > 0) {
498 log.warn('[images:generate] failed to clean up uncommitted assets', {
499 runId,
500 failureCount: cleanupFailures.length
501 })
502 }
503 }
504 const message = error instanceof Error ? error.message : String(error)
505 const wasCancelled = lease.signal.aborted
506 const logPayload = {
507 runId,
508 sessionId,
509 pageId: resolvedPageId,
510 provider: resolvedProvider,
511 message,
512 elapsedMs: Date.now() - startedAt
513 }
514 if (wasCancelled) {
515 log.warn('[images:generate] cancelled', logPayload)
516 } else {
517 log.error('[images:generate] failed', logPayload)
518 }
519 updateImageRunState({
520 runId,
521 sessionId,
522 pageId: resolvedPageId,
523 progress: 100,
524 label: wasCancelled ? uiText(locale, '已取消生图', 'Image generation cancelled') : message,
525 status: wasCancelled ? 'cancelled' : 'failed',
526 error: message,
527 updatedAt: Date.now()
528 })
529 emitImageJobTerminalEvent({
530 runId,
531 sessionId,
532 status: wasCancelled ? 'cancelled' : 'failed',
533 errorMessage: message
534 })
535 throw error
536 } finally {
537 lease.release()
538 }
539 })
540
541 ipcMain.handle('images:cancel', async (_event, sessionId) => {
542 if (typeof sessionId !== 'string' || !sessionId.trim()) return { success: false }
543 const state = getImageRunState(sessionId.trim())
544 if (!state || state.status !== 'running') return { success: false }
545 log.warn('[images:cancel] abort requested', {
546 runId: state.runId,
547 sessionId: state.sessionId,
548 pageId: state.pageId,
549 progress: state.progress
550 })
551 return { success: coordinator.cancel(state.runId) }
552 })
553
554 ipcMain.handle('images:getState', async (_event, sessionId) => {
555 if (typeof sessionId !== 'string' || !sessionId.trim()) return null
556 const state = getImageRunState(sessionId.trim())
557 if (!state) return null
558 return {
559 runId: state.runId,
560 sessionId: state.sessionId,
561 pageId: state.pageId,
562 progress: state.progress,
563 label: state.label,
564 status: state.status,
565 error: state.error || null,
566 updatedAt: state.updatedAt
567 }
568 })
569 }
570
570 lines TYPESCRIPT