feat(parse-progress): add local in-memory realtime bus

This commit is contained in:
Richard R 2026-05-25 21:06:18 -06:00
parent 52083c6a09
commit 2df49b54de
8 changed files with 425 additions and 79 deletions

View file

@ -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<Uint8Array>({
start(controller) {
let closed = false;
let unsubscribe: (() => void) | null = null;
let keepaliveTimer: ReturnType<typeof setInterval> | null = null;
let resyncTimer: ReturnType<typeof setInterval> | 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() {

View file

@ -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<void> {
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);
};
}
}

View file

@ -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<ParseProgressBus> | null = null;
export async function getParseProgressBus(): Promise<ParseProgressBus> {
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';

View file

@ -0,0 +1,14 @@
import type { ParseProgressEvent } from '@/lib/server/documents/parse-progress-events';
export interface ParseProgressBusSubscribeInput {
documentId: string;
userIdHashes: Set<string>;
onEvent: (event: ParseProgressEvent) => void | Promise<void>;
onError?: (error: unknown) => void;
}
export interface ParseProgressBus {
readonly kind: 'memory';
publish(event: ParseProgressEvent): Promise<void>;
subscribe(input: ParseProgressBusSubscribeInput): Promise<() => void>;
}

View file

@ -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);
}

View file

@ -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;
}

View file

@ -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<void> {
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<void> {
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<void>
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<void>
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<void>
userIds: cohortUserIds,
parseState: runningState,
});
await emitParseStateEvents({
documentId: input.documentId,
userIds: cohortUserIds,
state: runningStateData,
});
const compute = await getCompute();
const writeProgress = async (progress: PdfLayoutProgress): Promise<void> => {
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<void>
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<void>
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', {

View file

@ -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');
});
});