| 1 | import fs from 'fs' |
| 2 | import os from 'os' |
| 3 | import path from 'path' |
| 4 | import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' |
| 5 | import { JobCoordinator } from '../../../src/main/agent-runtime/job/coordinator' |
| 6 | import { TypedEventBus } from '../../../src/main/agent-runtime/events/bus' |
| 7 | |
| 8 | const state = vi.hoisted(() => { |
| 9 | const handlers = new Map<string, (...args: any[]) => Promise<unknown>>() |
| 10 | const generate = vi.fn() |
| 11 | return { |
| 12 | handlers, |
| 13 | generate, |
| 14 | resolveImageGenerationProvider: vi.fn(() => ({ generate })), |
| 15 | ipcMain: { |
| 16 | handle: vi.fn((channel: string, handler: (...args: any[]) => Promise<unknown>) => { |
| 17 | handlers.set(channel, handler) |
| 18 | }) |
| 19 | } |
| 20 | } |
| 21 | }) |
| 22 | |
| 23 | vi.mock('electron', () => ({ ipcMain: state.ipcMain })) |
| 24 | vi.mock('electron-log/main.js', () => ({ |
| 25 | default: { info: vi.fn(), warn: vi.fn(), error: vi.fn() } |
| 26 | })) |
| 27 | vi.mock('../../../src/main/agent-runtime/provider/image', () => ({ |
| 28 | resolveImageGenerationProvider: state.resolveImageGenerationProvider |
| 29 | })) |
| 30 | vi.mock('../../../src/main/io/assets-handlers', () => ({ allowLocalAssetRoot: vi.fn() })) |
| 31 | |
| 32 | const imagePayload = { |
| 33 | sessionId: 'session-1', |
| 34 | pageId: 'page-1', |
| 35 | prompt: 'A calm sunrise over a mountain lake', |
| 36 | count: 1, |
| 37 | size: '16:9' |
| 38 | } |
| 39 | |
| 40 | const createContext = (projectDir: string) => ({ |
| 41 | db: { |
| 42 | getSetting: vi.fn().mockResolvedValue('en'), |
| 43 | getActiveImageModelConfig: vi.fn().mockResolvedValue({ |
| 44 | id: 'image-model-1', |
| 45 | name: 'Image model', |
| 46 | provider: 'openaiCompatible', |
| 47 | active: 1, |
| 48 | modelConfig: JSON.stringify({ model: 'image-model' }) |
| 49 | }), |
| 50 | listSessionPages: vi.fn().mockResolvedValue([ |
| 51 | { |
| 52 | id: 'page-1', |
| 53 | file_slug: 'page-1', |
| 54 | legacy_page_id: null, |
| 55 | title: 'Cover', |
| 56 | page_number: 1, |
| 57 | html_path: path.join(projectDir, 'pages', 'page-1.html') |
| 58 | } |
| 59 | ]), |
| 60 | listLatestGenerationPageSnapshot: vi.fn().mockResolvedValue([]), |
| 61 | insertImageGenerationHistory: vi.fn().mockResolvedValue('history-1') |
| 62 | }, |
| 63 | decryptApiKey: vi.fn((value: string) => value), |
| 64 | resolveSessionProjectDir: vi.fn().mockResolvedValue(projectDir), |
| 65 | toSafeAssetBaseName: vi.fn(() => 'generated-cover') |
| 66 | }) |
| 67 | |
| 68 | const register = async (projectDir: string) => { |
| 69 | const { registerImageGenerationHandlers } = await import( |
| 70 | '../../../src/main/image-generation/handlers' |
| 71 | ) |
| 72 | const coordinator = new JobCoordinator() |
| 73 | const runtimeEvents = new TypedEventBus() |
| 74 | const events: Array<{ type: string; payload: unknown; jobId: string }> = [] |
| 75 | runtimeEvents.subscribe({}, (event) => events.push(event)) |
| 76 | registerImageGenerationHandlers(createContext(projectDir) as never, coordinator, runtimeEvents) |
| 77 | const generate = state.handlers.get('images:generate') |
| 78 | const cancel = state.handlers.get('images:cancel') |
| 79 | const getState = state.handlers.get('images:getState') |
| 80 | if (!generate || !cancel || !getState) throw new Error('Image IPC handlers were not registered') |
| 81 | return { coordinator, events, generate, cancel, getState } |
| 82 | } |
| 83 | |
| 84 | describe('registerImageGenerationHandlers', () => { |
| 85 | const temporaryDirs: string[] = [] |
| 86 | |
| 87 | beforeEach(() => { |
| 88 | vi.resetModules() |
| 89 | state.handlers.clear() |
| 90 | state.generate.mockReset() |
| 91 | state.resolveImageGenerationProvider.mockClear() |
| 92 | state.ipcMain.handle.mockClear() |
| 93 | }) |
| 94 | |
| 95 | afterEach(async () => { |
| 96 | await Promise.all( |
| 97 | temporaryDirs.splice(0).map((directory) => fs.promises.rm(directory, { recursive: true, force: true })) |
| 98 | ) |
| 99 | }) |
| 100 | |
| 101 | it('uses one image-history lease for generation, persistence, and terminal events', async () => { |
| 102 | const projectDir = await fs.promises.mkdtemp(path.join(os.tmpdir(), 'image-runtime-')) |
| 103 | temporaryDirs.push(projectDir) |
| 104 | state.generate.mockResolvedValue([ |
| 105 | { bytes: Buffer.from('image-bytes'), mimeType: 'image/png', extension: '.png' } |
| 106 | ]) |
| 107 | const { coordinator, events, generate, getState } = await register(projectDir) |
| 108 | |
| 109 | const result = await generate({}, imagePayload) |
| 110 | const runId = (result as { history: { id: string; assets: Array<{ relativePath: string }> } }).history |
| 111 | |
| 112 | expect(runId.id).toBe('history-1') |
| 113 | expect(runId.assets).toHaveLength(1) |
| 114 | expect(runId.assets[0]?.relativePath).toMatch(/^\.\/images\/generated-cover-.+\.png$/) |
| 115 | expect(state.generate).toHaveBeenCalledWith( |
| 116 | expect.objectContaining({ provider: 'openaiCompatible' }), |
| 117 | expect.objectContaining({ signal: expect.any(AbortSignal) }) |
| 118 | ) |
| 119 | expect(await getState({}, 'session-1')).toMatchObject({ status: 'completed', progress: 100 }) |
| 120 | expect(coordinator.getByOwner({ kind: 'image-history', id: 'session-1' })).toBeNull() |
| 121 | expect(events.map((event) => event.type)).toEqual([ |
| 122 | 'job.started', |
| 123 | 'image.progress', |
| 124 | 'image.progress', |
| 125 | 'image.progress', |
| 126 | 'job.completed' |
| 127 | ]) |
| 128 | }) |
| 129 | |
| 130 | it('rejects a concurrent session image request and cancels through the shared lease', async () => { |
| 131 | const projectDir = await fs.promises.mkdtemp(path.join(os.tmpdir(), 'image-runtime-')) |
| 132 | temporaryDirs.push(projectDir) |
| 133 | state.generate.mockImplementation( |
| 134 | (_config: unknown, input: { signal?: AbortSignal }) => |
| 135 | new Promise((_, reject) => { |
| 136 | input.signal?.addEventListener( |
| 137 | 'abort', |
| 138 | () => reject(new Error('Image generation cancelled')), |
| 139 | { once: true } |
| 140 | ) |
| 141 | }) |
| 142 | ) |
| 143 | const { coordinator, events, generate, cancel, getState } = await register(projectDir) |
| 144 | |
| 145 | const first = generate({}, imagePayload) |
| 146 | await vi.waitFor(() => expect(state.generate).toHaveBeenCalledTimes(1)) |
| 147 | await expect(generate({}, imagePayload)).rejects.toThrow('already running') |
| 148 | expect(await cancel({}, 'session-1')).toEqual({ success: true }) |
| 149 | await expect(first).rejects.toThrow('Image generation cancelled') |
| 150 | |
| 151 | expect(await getState({}, 'session-1')).toMatchObject({ status: 'cancelled', progress: 100 }) |
| 152 | expect(coordinator.getByOwner({ kind: 'image-history', id: 'session-1' })).toBeNull() |
| 153 | expect(events.map((event) => event.type)).toEqual([ |
| 154 | 'job.started', |
| 155 | 'image.progress', |
| 156 | 'image.progress', |
| 157 | 'job.cancelled' |
| 158 | ]) |
| 159 | }) |
| 160 | |
| 161 | it('publishes a stable failure event and releases the image-history lease', async () => { |
| 162 | const projectDir = await fs.promises.mkdtemp(path.join(os.tmpdir(), 'image-runtime-')) |
| 163 | temporaryDirs.push(projectDir) |
| 164 | state.generate.mockRejectedValue(new Error('Image provider unavailable')) |
| 165 | const { coordinator, events, generate, getState } = await register(projectDir) |
| 166 | |
| 167 | await expect(generate({}, imagePayload)).rejects.toThrow('Image provider unavailable') |
| 168 | |
| 169 | expect(await getState({}, 'session-1')).toMatchObject({ |
| 170 | status: 'failed', |
| 171 | error: 'Image provider unavailable' |
| 172 | }) |
| 173 | expect(coordinator.getByOwner({ kind: 'image-history', id: 'session-1' })).toBeNull() |
| 174 | expect(events).toContainEqual( |
| 175 | expect.objectContaining({ |
| 176 | type: 'job.failed', |
| 177 | payload: { |
| 178 | errorCode: 'image_generation_failed', |
| 179 | errorMessage: 'Image provider unavailable' |
| 180 | } |
| 181 | }) |
| 182 | ) |
| 183 | }) |
| 184 | |
| 185 | it('removes generated files when history persistence fails', async () => { |
| 186 | const projectDir = await fs.promises.mkdtemp(path.join(os.tmpdir(), 'image-runtime-')) |
| 187 | temporaryDirs.push(projectDir) |
| 188 | state.generate.mockResolvedValue([ |
| 189 | { bytes: Buffer.from('first-image'), mimeType: 'image/png', extension: '.png' }, |
| 190 | { bytes: Buffer.from('second-image'), mimeType: 'image/png', extension: '.png' } |
| 191 | ]) |
| 192 | const context = createContext(projectDir) |
| 193 | context.db.insertImageGenerationHistory.mockRejectedValue(new Error('history unavailable')) |
| 194 | const { registerImageGenerationHandlers } = await import( |
| 195 | '../../../src/main/image-generation/handlers' |
| 196 | ) |
| 197 | const coordinator = new JobCoordinator() |
| 198 | registerImageGenerationHandlers(context as never, coordinator, new TypedEventBus()) |
| 199 | const generate = state.handlers.get('images:generate') |
| 200 | if (!generate) throw new Error('Image IPC handler was not registered') |
| 201 | |
| 202 | await expect(generate({}, { ...imagePayload, count: 2 })).rejects.toThrow('history unavailable') |
| 203 | |
| 204 | const imagesDir = path.join(projectDir, 'images') |
| 205 | expect(await fs.promises.readdir(imagesDir)).toEqual([]) |
| 206 | expect(coordinator.getByOwner({ kind: 'image-history', id: 'session-1' })).toBeNull() |
| 207 | }) |
| 208 | }) |
| 209 |