From 2df49b54dea2c2e27f0bfcef3871bfea0303231c Mon Sep 17 00:00:00 2001 From: Richard R Date: Mon, 25 May 2026 21:06:18 -0600 Subject: [PATCH] feat(parse-progress): add local in-memory realtime bus --- .../api/documents/[id]/parsed/events/route.ts | 166 +++++++++++++---- .../documents/parse-progress-bus.memory.ts | 36 ++++ .../server/documents/parse-progress-bus.ts | 16 ++ .../documents/parse-progress-bus.types.ts | 14 ++ .../server/documents/parse-progress-events.ts | 18 ++ .../server/documents/parse-state-healing.ts | 24 ++- src/lib/server/jobs/user-pdf-layout-job.ts | 173 +++++++++++++----- tests/unit/parse-progress-bus.spec.ts | 57 ++++++ 8 files changed, 425 insertions(+), 79 deletions(-) create mode 100644 src/lib/server/documents/parse-progress-bus.memory.ts create mode 100644 src/lib/server/documents/parse-progress-bus.ts create mode 100644 src/lib/server/documents/parse-progress-bus.types.ts create mode 100644 src/lib/server/documents/parse-progress-events.ts create mode 100644 tests/unit/parse-progress-bus.spec.ts diff --git a/src/app/api/documents/[id]/parsed/events/route.ts b/src/app/api/documents/[id]/parsed/events/route.ts index d48771c..44df139 100644 --- a/src/app/api/documents/[id]/parsed/events/route.ts +++ b/src/app/api/documents/[id]/parsed/events/route.ts @@ -3,16 +3,22 @@ import { and, eq, inArray } from 'drizzle-orm'; import { db } from '@/db'; import { documents } from '@/db/schema'; import { requireAuthContext } from '@/lib/server/auth/auth'; -import { getOpenReaderTestNamespace, getUnclaimedUserIdForNamespace } from '@/lib/server/testing/test-namespace'; -import { isS3Configured } from '@/lib/server/storage/s3'; -import type { PdfParseProgress, PdfParseStatus } from '@/types/parsed-pdf'; +import { readComputeMode } from '@/lib/server/compute/mode'; +import { getParseProgressBus } from '@/lib/server/documents/parse-progress-bus'; +import { hashUserId } from '@/lib/server/documents/parse-progress-events'; import { isValidDocumentId } from '@/lib/server/documents/blobstore'; import { normalizeParseStatus, parseDocumentParseState } from '@/lib/server/documents/parse-state'; import { healStaleDocumentParseState } from '@/lib/server/documents/parse-state-healing'; +import { getOpenReaderTestNamespace, getUnclaimedUserIdForNamespace } from '@/lib/server/testing/test-namespace'; +import { isS3Configured } from '@/lib/server/storage/s3'; +import type { PdfParseProgress, PdfParseStatus } from '@/types/parsed-pdf'; export const dynamic = 'force-dynamic'; +export const runtime = 'nodejs'; const SSE_POLL_INTERVAL_MS = 1200; +const SSE_KEEPALIVE_MS = 15_000; +const SSE_RESYNC_INTERVAL_MS = 30_000; type ParseRow = { id: string; @@ -95,18 +101,132 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string } const initial = await toSnapshot(row); + const computeMode = readComputeMode(); + const bus = computeMode === 'local' ? await getParseProgressBus() : null; const encoder = new TextEncoder(); const stream = new ReadableStream({ start(controller) { let closed = false; + let unsubscribe: (() => void) | null = null; + let keepaliveTimer: ReturnType | null = null; + let resyncTimer: ReturnType | null = null; + + const allowedUserHashes = new Set(allowedUserIds.map((candidate) => hashUserId(candidate))); const writeSnapshot = (snapshot: ParsedSnapshot): void => { controller.enqueue(encoder.encode(`event: snapshot\ndata: ${JSON.stringify(snapshot)}\n\n`)); }; - const run = async () => { + const closeStream = (): void => { + if (closed) return; + closed = true; + if (keepaliveTimer) { + clearInterval(keepaliveTimer); + keepaliveTimer = null; + } + if (resyncTimer) { + clearInterval(resyncTimer); + resyncTimer = null; + } + if (unsubscribe) { + unsubscribe(); + unsubscribe = null; + } + try { + controller.close(); + } catch { + // no-op + } + }; + + const syncFromDb = async (input: { + signature: string; + }): Promise<{ snapshot: ParsedSnapshot; signature: string } | null> => { + const nextRow = await loadPreferredRow({ + documentId: id, + storageUserId, + allowedUserIds, + }); + if (!nextRow) return null; + + const nextSnapshot = await toSnapshot(nextRow); + const nextSignature = JSON.stringify(nextSnapshot); + if (nextSignature !== input.signature) { + writeSnapshot(nextSnapshot); + } + return { + snapshot: nextSnapshot, + signature: nextSignature, + }; + }; + + const runLocalRealtime = async () => { + let current = initial; + let signature = JSON.stringify(current); + writeSnapshot(current); + + if (current.parseStatus === 'ready' || current.parseStatus === 'failed') { + closeStream(); + return; + } + + keepaliveTimer = setInterval(() => { + if (closed) return; + controller.enqueue(encoder.encode(': keepalive\n\n')); + }, SSE_KEEPALIVE_MS); + + resyncTimer = setInterval(() => { + if (closed) return; + void syncFromDb({ signature }).then((next) => { + if (closed) return; + if (!next) { + closeStream(); + return; + } + current = next.snapshot; + signature = next.signature; + if (current.parseStatus === 'ready' || current.parseStatus === 'failed') { + closeStream(); + } + }).catch((error) => { + if (closed) return; + controller.enqueue(encoder.encode(`event: error\ndata: ${JSON.stringify({ error: String(error) })}\n\n`)); + closeStream(); + }); + }, SSE_RESYNC_INTERVAL_MS); + + if (!bus) throw new Error('Local parse progress bus unavailable'); + const nextUnsubscribe = await bus.subscribe({ + documentId: id, + userIdHashes: allowedUserHashes, + onEvent: async () => { + const next = await syncFromDb({ signature }); + if (!next) { + closeStream(); + return; + } + current = next.snapshot; + signature = next.signature; + if (current.parseStatus === 'ready' || current.parseStatus === 'failed') { + closeStream(); + } + }, + onError: (error) => { + if (closed) return; + controller.enqueue(encoder.encode(`event: error\ndata: ${JSON.stringify({ error: String(error) })}\n\n`)); + closeStream(); + }, + }); + if (closed) { + nextUnsubscribe(); + return; + } + unsubscribe = nextUnsubscribe; + }; + + const runWorkerPolling = async () => { let current = initial; writeSnapshot(current); let signature = JSON.stringify(current); @@ -116,24 +236,15 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string await sleep(SSE_POLL_INTERVAL_MS); if (closed) break; - const nextRow = await loadPreferredRow({ - documentId: id, - storageUserId, - allowedUserIds, - }); - if (!nextRow) break; - - const next = await toSnapshot(nextRow); - - const nextSignature = JSON.stringify(next); - if (nextSignature !== signature) { - current = next; - signature = nextSignature; - writeSnapshot(current); - } + const next = await syncFromDb({ signature }); + if (!next) break; + current = next.snapshot; + signature = next.signature; } }; + const run = computeMode === 'local' ? runLocalRealtime : runWorkerPolling; + void run() .catch((error) => { if (!closed) { @@ -141,24 +252,11 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string } }) .finally(() => { - if (!closed) { - closed = true; - try { - controller.close(); - } catch { - // no-op - } - } + closeStream(); }); req.signal.addEventListener('abort', () => { - if (closed) return; - closed = true; - try { - controller.close(); - } catch { - // no-op - } + closeStream(); }, { once: true }); }, cancel() { diff --git a/src/lib/server/documents/parse-progress-bus.memory.ts b/src/lib/server/documents/parse-progress-bus.memory.ts new file mode 100644 index 0000000..9a707f1 --- /dev/null +++ b/src/lib/server/documents/parse-progress-bus.memory.ts @@ -0,0 +1,36 @@ +import { EventEmitter } from 'node:events'; +import type { ParseProgressBus, ParseProgressBusSubscribeInput } from '@/lib/server/documents/parse-progress-bus.types'; +import type { ParseProgressEvent } from '@/lib/server/documents/parse-progress-events'; + +type Listener = (event: ParseProgressEvent) => void; + +function topicForDocument(documentId: string): string { + return `parse-progress.${documentId}`; +} + +export class InMemoryParseProgressBus implements ParseProgressBus { + readonly kind = 'memory' as const; + private readonly emitter = new EventEmitter(); + + constructor() { + this.emitter.setMaxListeners(0); + } + + async publish(event: ParseProgressEvent): Promise { + this.emitter.emit(topicForDocument(event.documentId), event); + } + + async subscribe(input: ParseProgressBusSubscribeInput): Promise<() => void> { + const listener: Listener = (event) => { + if (!input.userIdHashes.has(event.userIdHash)) return; + Promise.resolve(input.onEvent(event)).catch((error) => { + input.onError?.(error); + }); + }; + const topic = topicForDocument(input.documentId); + this.emitter.on(topic, listener); + return () => { + this.emitter.off(topic, listener); + }; + } +} diff --git a/src/lib/server/documents/parse-progress-bus.ts b/src/lib/server/documents/parse-progress-bus.ts new file mode 100644 index 0000000..df6e27b --- /dev/null +++ b/src/lib/server/documents/parse-progress-bus.ts @@ -0,0 +1,16 @@ +import { InMemoryParseProgressBus } from '@/lib/server/documents/parse-progress-bus.memory'; +import type { ParseProgressBus } from '@/lib/server/documents/parse-progress-bus.types'; + +let busPromise: Promise | null = null; + +export async function getParseProgressBus(): Promise { + if (busPromise) return busPromise; + busPromise = Promise.resolve(new InMemoryParseProgressBus()); + return busPromise; +} + +export function __resetParseProgressBusForTests(): void { + busPromise = null; +} + +export type { ParseProgressBus, ParseProgressBusSubscribeInput } from '@/lib/server/documents/parse-progress-bus.types'; diff --git a/src/lib/server/documents/parse-progress-bus.types.ts b/src/lib/server/documents/parse-progress-bus.types.ts new file mode 100644 index 0000000..423d257 --- /dev/null +++ b/src/lib/server/documents/parse-progress-bus.types.ts @@ -0,0 +1,14 @@ +import type { ParseProgressEvent } from '@/lib/server/documents/parse-progress-events'; + +export interface ParseProgressBusSubscribeInput { + documentId: string; + userIdHashes: Set; + onEvent: (event: ParseProgressEvent) => void | Promise; + onError?: (error: unknown) => void; +} + +export interface ParseProgressBus { + readonly kind: 'memory'; + publish(event: ParseProgressEvent): Promise; + subscribe(input: ParseProgressBusSubscribeInput): Promise<() => void>; +} diff --git a/src/lib/server/documents/parse-progress-events.ts b/src/lib/server/documents/parse-progress-events.ts new file mode 100644 index 0000000..f43fc91 --- /dev/null +++ b/src/lib/server/documents/parse-progress-events.ts @@ -0,0 +1,18 @@ +import { createHash } from 'node:crypto'; +import type { PdfParseProgress, PdfParseStatus } from '@/types/parsed-pdf'; + +export interface ParseProgressEvent { + version: 1; + documentId: string; + userIdHash: string; + parseStatus: PdfParseStatus; + parseProgress: PdfParseProgress | null; + updatedAt: number; + opId?: string; + jobId?: string; + error?: string | null; +} + +export function hashUserId(userId: string): string { + return createHash('sha256').update(userId).digest('hex').slice(0, 24); +} diff --git a/src/lib/server/documents/parse-state-healing.ts b/src/lib/server/documents/parse-state-healing.ts index b7a946c..909bbe3 100644 --- a/src/lib/server/documents/parse-state-healing.ts +++ b/src/lib/server/documents/parse-state-healing.ts @@ -8,6 +8,8 @@ import { stringifyDocumentParseState, type DocumentParseState, } from '@/lib/server/documents/parse-state'; +import { getParseProgressBus } from '@/lib/server/documents/parse-progress-bus'; +import { hashUserId } from '@/lib/server/documents/parse-progress-events'; export async function healStaleDocumentParseState(input: { documentId: string; @@ -29,6 +31,26 @@ export async function healStaleDocumentParseState(input: { .set({ parseState: stringifyDocumentParseState(nextState) }) .where(and(eq(documents.id, input.documentId), eq(documents.userId, input.userId))); + try { + const bus = await getParseProgressBus(); + await bus.publish({ + version: 1, + documentId: input.documentId, + userIdHash: hashUserId(input.userId), + parseStatus: nextState.status, + parseProgress: nextState.progress ?? null, + updatedAt: nextState.updatedAt ?? Date.now(), + ...(nextState.opId ? { opId: nextState.opId } : {}), + ...(nextState.jobId ? { jobId: nextState.jobId } : {}), + ...(nextState.error ? { error: nextState.error } : {}), + }); + } catch (error) { + console.warn('[parse-state-healing] failed to publish stale state event', { + documentId: input.documentId, + userId: input.userId, + error: error instanceof Error ? error.message : String(error), + }); + } + return nextState; } - diff --git a/src/lib/server/jobs/user-pdf-layout-job.ts b/src/lib/server/jobs/user-pdf-layout-job.ts index deec00e..6a40d94 100644 --- a/src/lib/server/jobs/user-pdf-layout-job.ts +++ b/src/lib/server/jobs/user-pdf-layout-job.ts @@ -3,7 +3,13 @@ import { and, eq, inArray, isNull } from 'drizzle-orm'; import { db } from '@/db'; import { documents } from '@/db/schema'; import { documentKey, putParsedDocumentBlob } from '@/lib/server/documents/blobstore'; -import { parseDocumentParseState, stringifyDocumentParseState } from '@/lib/server/documents/parse-state'; +import { + parseDocumentParseState, + stringifyDocumentParseState, + type DocumentParseState, +} from '@/lib/server/documents/parse-state'; +import { getParseProgressBus } from '@/lib/server/documents/parse-progress-bus'; +import { hashUserId } from '@/lib/server/documents/parse-progress-events'; import { getCompute } from '@/lib/server/compute'; import { clearTtsSegmentCache } from '@/lib/server/tts/segments-cache'; import { UNCLAIMED_USER_ID } from '@/lib/server/storage/docstore-legacy'; @@ -110,6 +116,41 @@ async function updateParseStateForUsers(input: { ); } +function normalizeEventState(state: DocumentParseState): DocumentParseState { + return parseDocumentParseState(stringifyDocumentParseState(state)); +} + +async function emitParseStateEvents(input: { + documentId: string; + userIds: string[]; + state: DocumentParseState; +}): Promise { + if (input.userIds.length === 0) return; + const normalized = normalizeEventState(input.state); + const bus = await getParseProgressBus(); + + const publishResults = await Promise.allSettled(input.userIds.map(async (userId) => bus.publish({ + version: 1, + documentId: input.documentId, + userIdHash: hashUserId(userId), + parseStatus: normalized.status, + parseProgress: normalized.progress ?? null, + updatedAt: normalized.updatedAt ?? Date.now(), + ...(normalized.opId ? { opId: normalized.opId } : {}), + ...(normalized.jobId ? { jobId: normalized.jobId } : {}), + ...(normalized.error ? { error: normalized.error } : {}), + }))); + + publishResults.forEach((result, index) => { + if (result.status === 'fulfilled') return; + console.warn('[parsePdfJob] failed to publish parse progress event', { + documentId: input.documentId, + userId: input.userIds[index], + error: result.reason instanceof Error ? result.reason.message : String(result.reason), + }); + }); +} + async function syncCallerToSharedResult(input: UserPdfLayoutJobRequest): Promise { const deadline = Date.now() + FOLLOWER_WAIT_TIMEOUT_MS; @@ -117,16 +158,22 @@ async function syncCallerToSharedResult(input: UserPdfLayoutJobRequest): Promise const rows = await loadScopedRows(input); const ready = rows.find(isReadyRow); if (ready) { + const readyState: DocumentParseState = { + status: 'ready', + progress: null, + updatedAt: Date.now(), + }; await updateParseStateForUsers({ documentId: input.documentId, userIds: [input.userId], - parseState: stringifyDocumentParseState({ - status: 'ready', - progress: null, - updatedAt: Date.now(), - }), + parseState: stringifyDocumentParseState(readyState), parsedJsonKey: ready.parsedJsonKey, }); + await emitParseStateEvents({ + documentId: input.documentId, + userIds: [input.userId], + state: readyState, + }); return; } @@ -134,15 +181,21 @@ async function syncCallerToSharedResult(input: UserPdfLayoutJobRequest): Promise const hasInFlight = statuses.some((state) => state.status === 'pending' || state.status === 'running'); const failedState = statuses.find((state) => state.status === 'failed'); if (failedState && !hasInFlight) { + const nextFailedState: DocumentParseState = { + status: 'failed', + progress: null, + updatedAt: Date.now(), + ...(failedState.error ? { error: failedState.error } : {}), + }; await updateParseStateForUsers({ documentId: input.documentId, userIds: [input.userId], - parseState: stringifyDocumentParseState({ - status: 'failed', - progress: null, - updatedAt: Date.now(), - ...(failedState.error ? { error: failedState.error } : {}), - }), + parseState: stringifyDocumentParseState(nextFailedState), + }); + await emitParseStateEvents({ + documentId: input.documentId, + userIds: [input.userId], + state: nextFailedState, }); return; } @@ -170,16 +223,22 @@ export async function parsePdfJob(input: UserPdfLayoutJobRequest): Promise if (!input.forceToken?.trim()) { const ready = scopedRows.find(isReadyRow); if (ready) { + const readyState: DocumentParseState = { + status: 'ready', + progress: null, + updatedAt: Date.now(), + }; await updateParseStateForUsers({ documentId: input.documentId, userIds: [input.userId], - parseState: stringifyDocumentParseState({ - status: 'ready', - progress: null, - updatedAt: Date.now(), - }), + parseState: stringifyDocumentParseState(readyState), parsedJsonKey: ready.parsedJsonKey, }); + await emitParseStateEvents({ + documentId: input.documentId, + userIds: [input.userId], + state: readyState, + }); return; } } @@ -188,12 +247,13 @@ export async function parsePdfJob(input: UserPdfLayoutJobRequest): Promise if (!coordinator) return; const claimOpId = randomUUID(); - const runningState = stringifyDocumentParseState({ + const runningStateData: DocumentParseState = { status: 'running', progress: null, updatedAt: Date.now(), opId: claimOpId, - }); + }; + const runningState = stringifyDocumentParseState(runningStateData); const claimRows = (await db .update(documents) @@ -223,23 +283,34 @@ export async function parsePdfJob(input: UserPdfLayoutJobRequest): Promise userIds: cohortUserIds, parseState: runningState, }); + await emitParseStateEvents({ + documentId: input.documentId, + userIds: cohortUserIds, + state: runningStateData, + }); const compute = await getCompute(); const writeProgress = async (progress: PdfLayoutProgress): Promise => { + const runningProgressState: DocumentParseState = { + status: 'running', + progress: { + totalPages: progress.totalPages, + pagesParsed: progress.pagesParsed, + currentPage: progress.currentPage, + phase: progress.phase, + }, + updatedAt: Date.now(), + opId: claimOpId, + }; await updateParseStateForUsers({ documentId: input.documentId, userIds: cohortUserIds, - parseState: stringifyDocumentParseState({ - status: 'running', - progress: { - totalPages: progress.totalPages, - pagesParsed: progress.pagesParsed, - currentPage: progress.currentPage, - phase: progress.phase, - }, - updatedAt: Date.now(), - opId: claimOpId, - }), + parseState: stringifyDocumentParseState(runningProgressState), + }); + await emitParseStateEvents({ + documentId: input.documentId, + userIds: cohortUserIds, + state: runningProgressState, }); }; const layout = await compute.parsePdfLayout({ @@ -259,19 +330,26 @@ export async function parsePdfJob(input: UserPdfLayoutJobRequest): Promise const finalScopedRows = await loadScopedRows(input); const finalUserIds = userIdsFromRows(finalScopedRows); + const readyState: DocumentParseState = { + status: 'ready', + progress: null, + updatedAt: Date.now(), + }; + const readyUserIds = finalUserIds.length > 0 ? finalUserIds : cohortUserIds; await updateParseStateForUsers({ documentId: input.documentId, - userIds: finalUserIds.length > 0 ? finalUserIds : cohortUserIds, - parseState: stringifyDocumentParseState({ - status: 'ready', - progress: null, - updatedAt: Date.now(), - }), + userIds: readyUserIds, + parseState: stringifyDocumentParseState(readyState), parsedJsonKey, }); + await emitParseStateEvents({ + documentId: input.documentId, + userIds: readyUserIds, + state: readyState, + }); // Best-effort cache invalidation should not block parse readiness. - for (const userId of (finalUserIds.length > 0 ? finalUserIds : cohortUserIds)) { + for (const userId of readyUserIds) { void clearTtsSegmentCache({ userId, documentId: input.documentId, @@ -300,15 +378,22 @@ export async function parsePdfJob(input: UserPdfLayoutJobRequest): Promise try { const scopedRows = await loadScopedRows(input); const scopedUserIds = userIdsFromRows(scopedRows); + const failedState: DocumentParseState = { + status: 'failed', + progress: null, + updatedAt: Date.now(), + error: message, + }; + const failedUserIds = scopedUserIds.length > 0 ? scopedUserIds : [input.userId]; await updateParseStateForUsers({ documentId: input.documentId, - userIds: scopedUserIds.length > 0 ? scopedUserIds : [input.userId], - parseState: stringifyDocumentParseState({ - status: 'failed', - progress: null, - updatedAt: Date.now(), - error: message, - }), + userIds: failedUserIds, + parseState: stringifyDocumentParseState(failedState), + }); + await emitParseStateEvents({ + documentId: input.documentId, + userIds: failedUserIds, + state: failedState, }); } catch (statusError) { console.error('[parsePdfJob] failed to write parse status', { diff --git a/tests/unit/parse-progress-bus.spec.ts b/tests/unit/parse-progress-bus.spec.ts new file mode 100644 index 0000000..754c71d --- /dev/null +++ b/tests/unit/parse-progress-bus.spec.ts @@ -0,0 +1,57 @@ +import { expect, test } from '@playwright/test'; +import { __resetParseProgressBusForTests, getParseProgressBus } from '../../src/lib/server/documents/parse-progress-bus'; +import { InMemoryParseProgressBus } from '../../src/lib/server/documents/parse-progress-bus.memory'; +import { hashUserId, type ParseProgressEvent } from '../../src/lib/server/documents/parse-progress-events'; + +test.describe('parse progress bus', () => { + test.afterEach(() => { + __resetParseProgressBusForTests(); + }); + + test('uses in-memory singleton bus', async () => { + const busA = await getParseProgressBus(); + const busB = await getParseProgressBus(); + expect(busA.kind).toBe('memory'); + expect(busA).toBe(busB); + }); + + test('in-memory bus publishes only to subscribed user hashes', async () => { + const bus = new InMemoryParseProgressBus(); + const documentId = 'd'.repeat(64); + const wantedHash = hashUserId('user-a'); + const otherHash = hashUserId('user-b'); + const received: ParseProgressEvent[] = []; + + const close = await bus.subscribe({ + documentId, + userIdHashes: new Set([wantedHash]), + onEvent: (event) => { + received.push(event); + }, + }); + + await bus.publish({ + version: 1, + documentId, + userIdHash: otherHash, + parseStatus: 'running', + parseProgress: { totalPages: 4, pagesParsed: 1, currentPage: 2, phase: 'infer' }, + updatedAt: Date.now(), + }); + + await bus.publish({ + version: 1, + documentId, + userIdHash: wantedHash, + parseStatus: 'running', + parseProgress: { totalPages: 4, pagesParsed: 2, currentPage: 3, phase: 'infer' }, + updatedAt: Date.now(), + opId: 'op-1', + }); + + close(); + expect(received).toHaveLength(1); + expect(received[0].userIdHash).toBe(wantedHash); + expect(received[0].opId).toBe('op-1'); + }); +});