| 1 | import type { BrowserWindow } from 'electron' |
| 2 | import log from 'electron-log/main.js' |
| 3 | import type { GenerateChunkEvent } from '@shared/generation' |
| 4 | import { progressDisplayLabel, type AppLocale } from '@shared/progress' |
| 5 | import { TypedEventBus, type RuntimeDomain } from '../../agent-runtime' |
| 6 | import { isCancellationMessage } from '../../generation/status-utils' |
| 7 | import { |
| 8 | revealGenerationWindow, |
| 9 | shouldRevealGenerationWindow |
| 10 | } from '../../generation/generation-window-policy' |
| 11 | import { |
| 12 | getSessionRunPageCounts, |
| 13 | runtimeDomainForSessionRun, |
| 14 | type SessionRunState, |
| 15 | type SessionRunStateStore |
| 16 | } from './session-run-state' |
| 17 | |
| 18 | export type RuntimeJobStartedArgs = { |
| 19 | sessionId: string |
| 20 | jobId: string |
| 21 | domain?: RuntimeDomain |
| 22 | } |
| 23 | |
| 24 | export type RuntimeJobTerminalArgs = { |
| 25 | sessionId: string |
| 26 | jobId: string |
| 27 | domain?: RuntimeDomain |
| 28 | status: 'completed' | 'failed' | 'cancelled' |
| 29 | errorCode?: string |
| 30 | errorMessage?: string |
| 31 | cancellationReason?: 'user' | 'timeout' | 'shutdown' |
| 32 | } |
| 33 | |
| 34 | export type LlmStatusEmissionSnapshot = { |
| 35 | stage: string |
| 36 | label: string |
| 37 | detail: string |
| 38 | progress: number | null |
| 39 | emittedAt: number |
| 40 | } |
| 41 | |
| 42 | export type RuntimeEmitters = { |
| 43 | emitSessionRunLifecycle(state: SessionRunState): void |
| 44 | emitGenerateChunk(sessionId: string, chunk: GenerateChunkEvent): void |
| 45 | emitRuntimeJobStarted(args: RuntimeJobStartedArgs): void |
| 46 | emitRuntimeJobTerminal(args: RuntimeJobTerminalArgs): void |
| 47 | createDeckProgressEmitter( |
| 48 | sessionId: string, |
| 49 | appLocale?: AppLocale |
| 50 | ): (chunk: GenerateChunkEvent) => void |
| 51 | } |
| 52 | |
| 53 | const LLM_STATUS_MIN_PROGRESS_DELTA = 5 |
| 54 | |
| 55 | export function getDeckProgressStageBounds(stage: string): { min: number; max: number } { |
| 56 | if (stage === 'preflight' || stage === 'planning') return { min: 0, max: 10 } |
| 57 | if (stage === 'rendering') return { min: 10, max: 90 } |
| 58 | if (stage === 'finalizing') return { min: 80, max: 100 } |
| 59 | return { min: 0, max: 90 } |
| 60 | } |
| 61 | |
| 62 | export function shouldEmitLlmStatusUpdate( |
| 63 | previous: LlmStatusEmissionSnapshot | null, |
| 64 | next: Omit<LlmStatusEmissionSnapshot, 'emittedAt'>, |
| 65 | now: number |
| 66 | ): boolean { |
| 67 | if (!previous) return true |
| 68 | if ( |
| 69 | previous.stage !== next.stage || |
| 70 | previous.label !== next.label || |
| 71 | previous.detail !== next.detail |
| 72 | ) { |
| 73 | return true |
| 74 | } |
| 75 | if ( |
| 76 | next.progress !== null && |
| 77 | (previous.progress === null || next.progress - previous.progress >= LLM_STATUS_MIN_PROGRESS_DELTA) |
| 78 | ) { |
| 79 | return true |
| 80 | } |
| 81 | void now |
| 82 | return false |
| 83 | } |
| 84 | |
| 85 | const summarizeGenerateChunk = (chunk: GenerateChunkEvent): Record<string, unknown> => { |
| 86 | switch (chunk.type) { |
| 87 | case 'stage_started': |
| 88 | case 'stage_progress': |
| 89 | return { |
| 90 | type: chunk.type, |
| 91 | stage: chunk.payload.stage, |
| 92 | label: chunk.payload.label, |
| 93 | progress: chunk.payload.progress ?? null, |
| 94 | totalPages: chunk.payload.totalPages ?? null |
| 95 | } |
| 96 | case 'llm_status': |
| 97 | return { |
| 98 | type: chunk.type, |
| 99 | stage: chunk.payload.stage, |
| 100 | label: chunk.payload.label, |
| 101 | detail: chunk.payload.detail ?? null, |
| 102 | progress: chunk.payload.progress ?? null, |
| 103 | totalPages: chunk.payload.totalPages ?? null, |
| 104 | provider: chunk.payload.provider ?? null, |
| 105 | model: chunk.payload.model ?? null |
| 106 | } |
| 107 | case 'page_generated': |
| 108 | case 'page_updated': |
| 109 | return { |
| 110 | type: chunk.type, |
| 111 | stage: chunk.payload.stage, |
| 112 | pageNumber: chunk.payload.pageNumber, |
| 113 | pageId: chunk.payload.pageId, |
| 114 | title: chunk.payload.title, |
| 115 | progress: chunk.payload.progress ?? null, |
| 116 | htmlPath: chunk.payload.htmlPath ?? null |
| 117 | } |
| 118 | case 'page_planned': |
| 119 | case 'page_started': |
| 120 | case 'page_failed': |
| 121 | return { |
| 122 | type: chunk.type, |
| 123 | stage: chunk.payload.stage, |
| 124 | pageNumber: chunk.payload.pageNumber, |
| 125 | pageId: chunk.payload.pageId, |
| 126 | title: chunk.payload.title, |
| 127 | progress: chunk.payload.progress ?? null, |
| 128 | error: chunk.payload.error ?? null |
| 129 | } |
| 130 | case 'run_completed': |
| 131 | return { |
| 132 | type: chunk.type, |
| 133 | totalPages: chunk.payload.totalPages, |
| 134 | completedPageCount: chunk.payload.completedPageCount ?? null, |
| 135 | failedPageCount: chunk.payload.failedPageCount ?? null, |
| 136 | activityKind: chunk.payload.activityKind ?? null |
| 137 | } |
| 138 | case 'run_error': |
| 139 | return { |
| 140 | type: chunk.type, |
| 141 | message: chunk.payload.message, |
| 142 | activityKind: chunk.payload.activityKind ?? null |
| 143 | } |
| 144 | default: |
| 145 | return { type: chunk.type } |
| 146 | } |
| 147 | } |
| 148 | |
| 149 | export function createRuntimeEmitters(args: { |
| 150 | mainWindow: BrowserWindow |
| 151 | runtimeEvents: TypedEventBus |
| 152 | sessionRuns: SessionRunStateStore |
| 153 | }): RuntimeEmitters { |
| 154 | const { mainWindow, runtimeEvents, sessionRuns } = args |
| 155 | |
| 156 | const emitSessionRunLifecycle = (state: SessionRunState): void => { |
| 157 | runtimeEvents.emit({ |
| 158 | type: state.status === 'queued' ? 'job.queued' : 'job.started', |
| 159 | payload: {}, |
| 160 | jobId: state.runId, |
| 161 | domain: runtimeDomainForSessionRun(state), |
| 162 | owner: { sessionId: state.sessionId }, |
| 163 | audience: { kind: 'broadcast' }, |
| 164 | occurredAt: state.startedAt |
| 165 | }) |
| 166 | } |
| 167 | |
| 168 | const emitGenerateChunk = (sessionId: string, chunk: GenerateChunkEvent): void => { |
| 169 | let enrichedChunk = { |
| 170 | ...chunk, |
| 171 | payload: { |
| 172 | ...chunk.payload, |
| 173 | sessionId, |
| 174 | timestamp: new Date().toISOString() |
| 175 | } |
| 176 | } as GenerateChunkEvent |
| 177 | if (enrichedChunk.type === 'run_error') { |
| 178 | enrichedChunk = { |
| 179 | ...enrichedChunk, |
| 180 | payload: { |
| 181 | ...enrichedChunk.payload, |
| 182 | cancelled: |
| 183 | enrichedChunk.payload.cancelled ?? |
| 184 | isCancellationMessage(enrichedChunk.payload.message || '') |
| 185 | } |
| 186 | } |
| 187 | } |
| 188 | |
| 189 | sessionRuns.trackSessionRunChunk(sessionId, enrichedChunk) |
| 190 | const state = sessionRuns.sessionRunStates.get(sessionId) |
| 191 | if (state?.runId === enrichedChunk.payload.runId) { |
| 192 | const pageCounts = getSessionRunPageCounts(state) |
| 193 | enrichedChunk = { |
| 194 | ...enrichedChunk, |
| 195 | payload: { |
| 196 | ...enrichedChunk.payload, |
| 197 | activityKind: state.activityKind, |
| 198 | completedPageCount: pageCounts.completedPageCount, |
| 199 | failedPageCount: pageCounts.failedPageCount |
| 200 | } |
| 201 | } as GenerateChunkEvent |
| 202 | } |
| 203 | |
| 204 | if ( |
| 205 | enrichedChunk.type === 'stage_started' || |
| 206 | enrichedChunk.type === 'stage_progress' || |
| 207 | enrichedChunk.type === 'llm_status' || |
| 208 | enrichedChunk.type === 'page_planned' || |
| 209 | enrichedChunk.type === 'page_started' || |
| 210 | enrichedChunk.type === 'page_generated' || |
| 211 | enrichedChunk.type === 'page_updated' || |
| 212 | enrichedChunk.type === 'page_failed' || |
| 213 | enrichedChunk.type === 'run_completed' || |
| 214 | enrichedChunk.type === 'run_error' |
| 215 | ) { |
| 216 | log.info('[generate:chunk] emit', summarizeGenerateChunk(enrichedChunk)) |
| 217 | } |
| 218 | |
| 219 | if (shouldRevealGenerationWindow(enrichedChunk, state)) { |
| 220 | revealGenerationWindow(mainWindow) |
| 221 | } |
| 222 | |
| 223 | runtimeEvents.emit({ |
| 224 | type: 'generation.chunk', |
| 225 | payload: enrichedChunk, |
| 226 | jobId: enrichedChunk.payload.runId, |
| 227 | domain: runtimeDomainForSessionRun(state), |
| 228 | owner: { sessionId }, |
| 229 | audience: { kind: 'broadcast' }, |
| 230 | occurredAt: Date.now() |
| 231 | }) |
| 232 | } |
| 233 | |
| 234 | const emitRuntimeJobTerminal = (event: RuntimeJobTerminalArgs): void => { |
| 235 | const domain = |
| 236 | event.domain || runtimeDomainForSessionRun(sessionRuns.sessionRunStates.get(event.sessionId)) |
| 237 | runtimeEvents.emit({ |
| 238 | type: |
| 239 | event.status === 'completed' |
| 240 | ? 'job.completed' |
| 241 | : event.status === 'cancelled' |
| 242 | ? 'job.cancelled' |
| 243 | : 'job.failed', |
| 244 | payload: |
| 245 | event.status === 'completed' |
| 246 | ? {} |
| 247 | : event.status === 'cancelled' |
| 248 | ? { reason: event.cancellationReason || 'user' } |
| 249 | : { |
| 250 | errorCode: event.errorCode || 'generation_failed', |
| 251 | errorMessage: event.errorMessage || 'Generation failed' |
| 252 | }, |
| 253 | jobId: event.jobId, |
| 254 | domain, |
| 255 | owner: { sessionId: event.sessionId }, |
| 256 | audience: { kind: 'broadcast' }, |
| 257 | occurredAt: Date.now() |
| 258 | }) |
| 259 | } |
| 260 | |
| 261 | const emitRuntimeJobStarted = (event: RuntimeJobStartedArgs): void => { |
| 262 | runtimeEvents.emit({ |
| 263 | type: 'job.started', |
| 264 | payload: {}, |
| 265 | jobId: event.jobId, |
| 266 | domain: |
| 267 | event.domain || runtimeDomainForSessionRun(sessionRuns.sessionRunStates.get(event.sessionId)), |
| 268 | owner: { sessionId: event.sessionId }, |
| 269 | audience: { kind: 'broadcast' }, |
| 270 | occurredAt: Date.now() |
| 271 | }) |
| 272 | } |
| 273 | |
| 274 | const createDeckProgressEmitter = ( |
| 275 | sessionId: string, |
| 276 | appLocale?: AppLocale |
| 277 | ): ((chunk: GenerateChunkEvent) => void) => { |
| 278 | let normalizedProgress = 0 |
| 279 | let lastLlmStatusEmission: LlmStatusEmissionSnapshot | null = null |
| 280 | |
| 281 | const clamp = (value: number, min: number, max: number): number => |
| 282 | Math.max(min, Math.min(max, Math.round(value))) |
| 283 | |
| 284 | return (chunk: GenerateChunkEvent): void => { |
| 285 | if (chunk.type === 'run_completed') { |
| 286 | normalizedProgress = 100 |
| 287 | emitGenerateChunk(sessionId, chunk) |
| 288 | return |
| 289 | } |
| 290 | |
| 291 | if ( |
| 292 | chunk.type !== 'stage_started' && |
| 293 | chunk.type !== 'stage_progress' && |
| 294 | chunk.type !== 'llm_status' && |
| 295 | chunk.type !== 'page_started' && |
| 296 | chunk.type !== 'page_generated' && |
| 297 | chunk.type !== 'page_updated' && |
| 298 | chunk.type !== 'page_failed' |
| 299 | ) { |
| 300 | emitGenerateChunk(sessionId, chunk) |
| 301 | return |
| 302 | } |
| 303 | |
| 304 | const { min, max } = getDeckProgressStageBounds(chunk.payload.stage) |
| 305 | const rawProgress = |
| 306 | typeof chunk.payload.progress === 'number' && Number.isFinite(chunk.payload.progress) |
| 307 | ? chunk.payload.progress |
| 308 | : normalizedProgress |
| 309 | const bounded = clamp(rawProgress, min, max) |
| 310 | normalizedProgress = Math.max(normalizedProgress, bounded) |
| 311 | |
| 312 | const normalizedChunk = { |
| 313 | ...chunk, |
| 314 | payload: { |
| 315 | ...chunk.payload, |
| 316 | label: progressDisplayLabel(appLocale, chunk.payload.label), |
| 317 | progress: normalizedProgress |
| 318 | } |
| 319 | } as GenerateChunkEvent |
| 320 | |
| 321 | if (normalizedChunk.type === 'llm_status') { |
| 322 | const now = Date.now() |
| 323 | const next = { |
| 324 | stage: normalizedChunk.payload.stage, |
| 325 | label: normalizedChunk.payload.label, |
| 326 | detail: normalizedChunk.payload.detail || '', |
| 327 | progress: |
| 328 | typeof normalizedChunk.payload.progress === 'number' |
| 329 | ? normalizedChunk.payload.progress |
| 330 | : null |
| 331 | } |
| 332 | if (!shouldEmitLlmStatusUpdate(lastLlmStatusEmission, next, now)) return |
| 333 | lastLlmStatusEmission = { ...next, emittedAt: now } |
| 334 | } |
| 335 | |
| 336 | emitGenerateChunk(sessionId, normalizedChunk) |
| 337 | } |
| 338 | } |
| 339 | |
| 340 | return { |
| 341 | emitSessionRunLifecycle, |
| 342 | emitGenerateChunk, |
| 343 | emitRuntimeJobStarted, |
| 344 | emitRuntimeJobTerminal, |
| 345 | createDeckProgressEmitter |
| 346 | } |
| 347 | } |
| 348 |