| 1 | import type { GenerateChunkEvent } from '@shared/generation' |
| 2 | import type { RuntimeDomain } from '../../agent-runtime' |
| 3 | |
| 4 | export type SessionRunMode = |
| 5 | | 'generate' |
| 6 | | 'edit' |
| 7 | | 'retry' |
| 8 | | 'addPage' |
| 9 | | 'retrySinglePage' |
| 10 | | 'style-switch' |
| 11 | |
| 12 | export type SessionRunKind = |
| 13 | | 'standard' |
| 14 | | 'template' |
| 15 | | 'retry' |
| 16 | | 'add-page' |
| 17 | | 'single-page-retry' |
| 18 | | 'edit' |
| 19 | | 'page-edit' |
| 20 | | 'deck-edit' |
| 21 | | 'style-switch' |
| 22 | |
| 23 | export type SessionRunActivityKind = |
| 24 | | 'page-edit' |
| 25 | | 'deck-edit' |
| 26 | | 'edit' |
| 27 | | 'style-switch' |
| 28 | | 'single-page-retry' |
| 29 | | 'addPage' |
| 30 | |
| 31 | export type SessionRunState = { |
| 32 | sessionId: string |
| 33 | runId: string |
| 34 | mode: SessionRunMode |
| 35 | kind?: SessionRunKind |
| 36 | activityKind?: SessionRunActivityKind |
| 37 | targetPageId?: string |
| 38 | targetPageNumber?: number |
| 39 | previousSessionStatus?: string |
| 40 | status: 'queued' | 'running' | 'completed' | 'failed' | 'cancelled' |
| 41 | progress: number |
| 42 | totalPages: number |
| 43 | completedPageBaseCount: number |
| 44 | failedPageBaseKeys: string[] |
| 45 | completedPageKeys: string[] |
| 46 | failedPageKeys: string[] |
| 47 | events: GenerateChunkEvent[] |
| 48 | error: string | null |
| 49 | startedAt: number |
| 50 | updatedAt: number |
| 51 | } |
| 52 | |
| 53 | export type BeginSessionRunStateArgs = { |
| 54 | sessionId: string |
| 55 | runId: string |
| 56 | mode: SessionRunMode |
| 57 | kind?: SessionRunKind |
| 58 | activityKind?: SessionRunActivityKind |
| 59 | targetPageId?: string |
| 60 | targetPageNumber?: number |
| 61 | totalPages: number |
| 62 | previousSessionStatus?: string |
| 63 | status?: 'queued' | 'running' |
| 64 | completedPageBaseCount?: number |
| 65 | failedPageBaseKeys?: string[] |
| 66 | } |
| 67 | |
| 68 | export type SessionRunStateStore = { |
| 69 | sessionRunStates: Map<string, SessionRunState> |
| 70 | pruneFinishedSessionRunStates(now?: number): void |
| 71 | beginSessionRunState(args: BeginSessionRunStateArgs): SessionRunState |
| 72 | trackSessionRunChunk(sessionId: string, chunk: GenerateChunkEvent): void |
| 73 | } |
| 74 | |
| 75 | const MAX_SESSION_RUN_EVENTS = 500 |
| 76 | const FINISHED_SESSION_RUN_STATE_TTL_MS = 30 * 60 * 1000 |
| 77 | |
| 78 | export function getSessionRunPageCounts(state: { |
| 79 | completedPageBaseCount: number |
| 80 | failedPageBaseKeys: string[] |
| 81 | completedPageKeys: string[] |
| 82 | failedPageKeys: string[] |
| 83 | }): { completedPageCount: number; failedPageCount: number } { |
| 84 | const completedPageCount = state.completedPageBaseCount + state.completedPageKeys.length |
| 85 | const completed = new Set(state.completedPageKeys) |
| 86 | const failed = new Set(state.failedPageKeys) |
| 87 | for (const pageKey of state.failedPageBaseKeys) { |
| 88 | if (!completed.has(pageKey)) failed.add(pageKey) |
| 89 | } |
| 90 | return { completedPageCount, failedPageCount: failed.size } |
| 91 | } |
| 92 | |
| 93 | export function runtimeDomainForSessionRun(state: SessionRunState | undefined): RuntimeDomain { |
| 94 | const activityKind = state?.activityKind |
| 95 | if (activityKind === 'style-switch') return 'style' |
| 96 | if (activityKind === 'page-edit' || activityKind === 'deck-edit' || activityKind === 'edit') { |
| 97 | return 'edit' |
| 98 | } |
| 99 | return 'generation' |
| 100 | } |
| 101 | |
| 102 | export function createSessionRunStateStore(): SessionRunStateStore { |
| 103 | const sessionRunStates = new Map<string, SessionRunState>() |
| 104 | |
| 105 | const pruneFinishedSessionRunStates = (now = Date.now()): void => { |
| 106 | for (const [sessionId, state] of sessionRunStates) { |
| 107 | if (state.status === 'queued' || state.status === 'running') continue |
| 108 | if (now - state.updatedAt > FINISHED_SESSION_RUN_STATE_TTL_MS) { |
| 109 | sessionRunStates.delete(sessionId) |
| 110 | } |
| 111 | } |
| 112 | } |
| 113 | |
| 114 | const beginSessionRunState = (args: BeginSessionRunStateArgs): SessionRunState => { |
| 115 | const now = Date.now() |
| 116 | pruneFinishedSessionRunStates(now) |
| 117 | const state: SessionRunState = { |
| 118 | sessionId: args.sessionId, |
| 119 | runId: args.runId, |
| 120 | mode: args.mode, |
| 121 | kind: args.kind, |
| 122 | activityKind: args.activityKind, |
| 123 | targetPageId: args.targetPageId, |
| 124 | targetPageNumber: args.targetPageNumber, |
| 125 | previousSessionStatus: args.previousSessionStatus, |
| 126 | status: args.status || 'running', |
| 127 | progress: 0, |
| 128 | totalPages: Math.max(1, Math.floor(args.totalPages || 1)), |
| 129 | completedPageBaseCount: Math.max(0, Math.floor(args.completedPageBaseCount || 0)), |
| 130 | failedPageBaseKeys: args.failedPageBaseKeys || [], |
| 131 | completedPageKeys: [], |
| 132 | failedPageKeys: [], |
| 133 | events: [], |
| 134 | error: null, |
| 135 | startedAt: now, |
| 136 | updatedAt: now |
| 137 | } |
| 138 | sessionRunStates.set(args.sessionId, state) |
| 139 | return state |
| 140 | } |
| 141 | |
| 142 | const trackSessionRunChunk = (sessionId: string, chunk: GenerateChunkEvent): void => { |
| 143 | const state = sessionRunStates.get(sessionId) |
| 144 | if (!state || state.runId !== chunk.payload.runId) return |
| 145 | |
| 146 | if ( |
| 147 | chunk.type === 'page_generated' || |
| 148 | chunk.type === 'page_updated' || |
| 149 | chunk.type === 'page_failed' |
| 150 | ) { |
| 151 | const payload = chunk.payload as { pageId?: unknown; id?: unknown; pageNumber?: unknown } |
| 152 | const pageId = |
| 153 | typeof payload.pageId === 'string' && payload.pageId.trim().length > 0 |
| 154 | ? payload.pageId.trim() |
| 155 | : typeof payload.id === 'string' && payload.id.trim().length > 0 |
| 156 | ? payload.id.trim() |
| 157 | : '' |
| 158 | const pageKey = |
| 159 | pageId || |
| 160 | (typeof payload.pageNumber === 'number' && Number.isFinite(payload.pageNumber) |
| 161 | ? `page-number:${Math.floor(payload.pageNumber)}` |
| 162 | : '') |
| 163 | if (pageKey) { |
| 164 | const completed = new Set(state.completedPageKeys) |
| 165 | const failed = new Set(state.failedPageKeys) |
| 166 | if (chunk.type === 'page_failed') { |
| 167 | completed.delete(pageKey) |
| 168 | failed.add(pageKey) |
| 169 | } else { |
| 170 | failed.delete(pageKey) |
| 171 | completed.add(pageKey) |
| 172 | } |
| 173 | state.completedPageKeys = Array.from(completed) |
| 174 | state.failedPageKeys = Array.from(failed) |
| 175 | } |
| 176 | } |
| 177 | |
| 178 | const compactChunk = |
| 179 | chunk.type === 'page_generated' || chunk.type === 'page_updated' |
| 180 | ? ({ |
| 181 | ...chunk, |
| 182 | payload: { |
| 183 | ...chunk.payload, |
| 184 | html: '' |
| 185 | } |
| 186 | } as GenerateChunkEvent) |
| 187 | : chunk |
| 188 | |
| 189 | state.updatedAt = Date.now() |
| 190 | state.events.push(compactChunk) |
| 191 | if (state.events.length > MAX_SESSION_RUN_EVENTS) { |
| 192 | state.events.splice(0, state.events.length - MAX_SESSION_RUN_EVENTS) |
| 193 | } |
| 194 | |
| 195 | if (chunk.type === 'run_completed') { |
| 196 | state.status = 'completed' |
| 197 | state.progress = 100 |
| 198 | state.totalPages = Math.max( |
| 199 | state.totalPages, |
| 200 | Math.floor(chunk.payload.totalPages || state.totalPages) |
| 201 | ) |
| 202 | state.error = null |
| 203 | return |
| 204 | } |
| 205 | |
| 206 | if (chunk.type === 'run_error') { |
| 207 | state.status = /^(生成已取消|Generation cancelled|Generation canceled)$/i.test( |
| 208 | chunk.payload.message || '' |
| 209 | ) |
| 210 | ? 'cancelled' |
| 211 | : 'failed' |
| 212 | state.error = chunk.payload.message || 'Generation failed' |
| 213 | return |
| 214 | } |
| 215 | |
| 216 | if ( |
| 217 | 'totalPages' in chunk.payload && |
| 218 | typeof chunk.payload.totalPages === 'number' && |
| 219 | Number.isFinite(chunk.payload.totalPages) |
| 220 | ) { |
| 221 | state.totalPages = Math.max(1, Math.floor(chunk.payload.totalPages)) |
| 222 | } |
| 223 | if ( |
| 224 | 'progress' in chunk.payload && |
| 225 | typeof chunk.payload.progress === 'number' && |
| 226 | Number.isFinite(chunk.payload.progress) |
| 227 | ) { |
| 228 | const boundedProgress = Math.max(0, Math.min(100, Math.round(chunk.payload.progress))) |
| 229 | state.progress = Math.max(state.progress, boundedProgress) |
| 230 | } |
| 231 | } |
| 232 | |
| 233 | return { |
| 234 | sessionRunStates, |
| 235 | pruneFinishedSessionRunStates, |
| 236 | beginSessionRunState, |
| 237 | trackSessionRunChunk |
| 238 | } |
| 239 | } |
| 240 |