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