From be7ec2c254aec29127464a5b4caf70ef170507a9 Mon Sep 17 00:00:00 2001 From: Richard R Date: Mon, 25 May 2026 21:37:27 -0600 Subject: [PATCH] feat(worker-events): stream op events via jetstream and proxy in app --- compute/core/src/api-contracts/index.ts | 5 + compute/worker/src/server.ts | 136 ++++++- .../api/documents/[id]/parsed/events/route.ts | 334 ++++++++++++++++-- src/lib/server/compute/types.ts | 2 + src/lib/server/compute/worker.ts | 256 ++++++++------ src/lib/server/jobs/user-pdf-layout-job.ts | 88 ++++- 6 files changed, 650 insertions(+), 171 deletions(-) diff --git a/compute/core/src/api-contracts/index.ts b/compute/core/src/api-contracts/index.ts index ff4b2a0..18faa08 100644 --- a/compute/core/src/api-contracts/index.ts +++ b/compute/core/src/api-contracts/index.ts @@ -113,3 +113,8 @@ export interface WorkerOperationState { timing?: WorkerJobTiming; progress?: PdfLayoutProgress; } + +export interface WorkerOperationEvent { + eventId: number; + snapshot: WorkerOperationState; +} diff --git a/compute/worker/src/server.ts b/compute/worker/src/server.ts index 0b3a381..ea30c8b 100644 --- a/compute/worker/src/server.ts +++ b/compute/worker/src/server.ts @@ -37,6 +37,7 @@ import { import type { PdfLayoutJobRequest, PdfLayoutJobResult, + WorkerOperationEvent, WhisperAlignJobRequest, WhisperAlignJobResult, WorkerJobErrorShape, @@ -54,11 +55,12 @@ const WHISPER_JOBS_SUBJECT = 'jobs.whisper'; const LAYOUT_JOBS_SUBJECT = 'jobs.layout'; const WHISPER_CONSUMER_NAME = 'compute_whisper'; const LAYOUT_CONSUMER_NAME = 'compute_layout'; +const EVENTS_STREAM_NAME = 'compute_events'; const COMPUTE_STATE_BUCKET = 'compute_state'; const COMPUTE_STATE_TTL_MS = 24 * 60 * 60 * 1000; const LOOP_ERROR_BACKOFF_MS = 500; -const SSE_POLL_INTERVAL_MS = 400; const RUNNING_HEARTBEAT_MS = 5000; +const OP_EVENTS_SUBJECT_PREFIX = 'ops.events'; const DOCUMENT_ID_REGEX = /^[a-f0-9]{64}$/i; const SAFE_NAMESPACE_REGEX = /^[a-zA-Z0-9._-]{1,128}$/; const WHISPER_MAX_DELIVER = 1; @@ -115,6 +117,8 @@ interface NatsSession { layoutConsumer: Consumer; } +type StreamedOperationState = WorkerOperationState; + type JsonCodec = { encode(value: T): Uint8Array; decode(data: Uint8Array): T; @@ -363,6 +367,10 @@ function jobStateKvKey(jobId: string): string { return `job_state.${jobId}`; } +function opEventsSubject(opId: string): string { + return `${OP_EVENTS_SUBJECT_PREFIX}.${opId}`; +} + function extractResultRef(kind: WorkerOperationKind, result: unknown): string | undefined { if (kind !== 'pdf_layout' || !result || typeof result !== 'object') return undefined; const maybe = result as { parsedObjectKey?: unknown }; @@ -400,7 +408,8 @@ async function ensureJetStreamResources( whisperTimeoutMs: number, pdfTimeoutMs: number, pdfAttempts: number, - maxBytes: number, + jobsMaxBytes: number, + eventsMaxBytes: number, natsReplicas: number, ): Promise { const streamConfig = { @@ -408,7 +417,7 @@ async function ensureJetStreamResources( subjects: [WHISPER_JOBS_SUBJECT, LAYOUT_JOBS_SUBJECT], retention: RetentionPolicy.Workqueue, storage: StorageType.File, - max_bytes: maxBytes, + max_bytes: jobsMaxBytes, num_replicas: natsReplicas, }; @@ -418,7 +427,29 @@ async function ensureJetStreamResources( if (!isAlreadyExistsError(error)) throw error; await jsm.streams.update(JOBS_STREAM_NAME, { subjects: [WHISPER_JOBS_SUBJECT, LAYOUT_JOBS_SUBJECT], - max_bytes: maxBytes, + max_bytes: jobsMaxBytes, + num_replicas: natsReplicas, + }); + } + + const eventsStreamConfig = { + name: EVENTS_STREAM_NAME, + subjects: [`${OP_EVENTS_SUBJECT_PREFIX}.*`], + retention: RetentionPolicy.Limits, + storage: StorageType.File, + max_bytes: eventsMaxBytes, + max_age: nanos(COMPUTE_STATE_TTL_MS), + num_replicas: natsReplicas, + }; + + try { + await jsm.streams.add(eventsStreamConfig); + } catch (error) { + if (!isAlreadyExistsError(error)) throw error; + await jsm.streams.update(EVENTS_STREAM_NAME, { + subjects: [`${OP_EVENTS_SUBJECT_PREFIX}.*`], + max_bytes: eventsMaxBytes, + max_age: nanos(COMPUTE_STATE_TTL_MS), num_replicas: natsReplicas, }); } @@ -471,6 +502,7 @@ async function main(): Promise { const pdfAttempts = readIntEnv('COMPUTE_PDF_JOB_ATTEMPTS', 1); const prewarmModels = parseBoolEnv('COMPUTE_PREWARM_MODELS', true); const jobsStreamMaxBytes = readIntEnv('COMPUTE_JOBS_STREAM_MAX_BYTES', 256 * 1024 * 1024); + const eventsStreamMaxBytes = readIntEnv('COMPUTE_EVENTS_STREAM_MAX_BYTES', 128 * 1024 * 1024); const jobStatesMaxBytes = readIntEnv('COMPUTE_JOB_STATES_MAX_BYTES', 64 * 1024 * 1024); const natsReplicas = normalizeNatsReplicas(readIntEnv('COMPUTE_NATS_REPLICAS', 1)); const opStaleMs = getComputeOpStaleMs(); @@ -555,7 +587,15 @@ async function main(): Promise { const nc: NatsConnection = await connect(connectOpts); const js: JetStreamClient = jetstream(nc, { timeout: NATS_API_TIMEOUT_MS }); const jsm: JetStreamManager = await jetstreamManager(nc, { timeout: NATS_API_TIMEOUT_MS }); - await ensureJetStreamResources(jsm, whisperTimeoutMs, pdfTimeoutMs, pdfAttempts, jobsStreamMaxBytes, natsReplicas); + await ensureJetStreamResources( + jsm, + whisperTimeoutMs, + pdfTimeoutMs, + pdfAttempts, + jobsStreamMaxBytes, + eventsStreamMaxBytes, + natsReplicas, + ); const kv = await new Kvm(js).create(COMPUTE_STATE_BUCKET, { replicas: natsReplicas, history: 1, @@ -638,21 +678,32 @@ async function main(): Promise { onnxThreadsPerJob: getOnnxThreadsPerJob(), natsApiTimeoutMs: NATS_API_TIMEOUT_MS, natsReplicas, + eventsStreamMaxBytes, pdfLayoutHardCapMs: pdfHardCapMs, }, 'compute runtime config'); const opIndexCodec = createJsonCodec(); - const opStateCodec = createJsonCodec>(); + const opStateCodec = createJsonCodec(); const whisperJobCodec = createJsonCodec>(); const layoutJobCodec = createJsonCodec>(); const jobStateCodec = createJsonCodec>(); + const opEventCodec = createJsonCodec(); - const putOpState = async (state: WorkerOperationState): Promise => { - const { kv } = await ensureConnected(); + const putOpState = async (state: StreamedOperationState): Promise => { + const { kv, js } = await ensureConnected(); await kv.put(opStateKvKey(state.opId), opStateCodec.encode(state)); + try { + await js.publish(opEventsSubject(state.opId), opEventCodec.encode(state)); + } catch (error) { + app.log.warn({ + opId: state.opId, + status: state.status, + error: toErrorMessage(error), + }, 'failed to publish op event'); + } }; - const getOpState = async (opId: string): Promise | null> => { + const getOpState = async (opId: string): Promise => { const { kv } = await ensureConnected(); const entry = await kv.get(opStateKvKey(opId)); if (!entry || entry.operation !== 'PUT') return null; @@ -1007,6 +1058,18 @@ async function main(): Promise { return { error: 'Operation not found' }; } + const cursorQueryRaw = request.query as { sinceEventId?: string | number | null } | undefined; + const cursorFromQuery = Number(cursorQueryRaw?.sinceEventId ?? 0); + const lastEventIdHeader = request.headers['last-event-id']; + const cursorFromHeader = Number( + Array.isArray(lastEventIdHeader) ? (lastEventIdHeader[0] ?? 0) : (lastEventIdHeader ?? 0), + ); + const sinceEventId = Math.max( + 0, + Number.isFinite(cursorFromQuery) ? Math.floor(cursorFromQuery) : 0, + Number.isFinite(cursorFromHeader) ? Math.floor(cursorFromHeader) : 0, + ); + reply.hijack(); // onResponse will not fire for a hijacked reply, so release the HTTP in-flight // count here and track the long-lived stream via activeSse instead. @@ -1018,12 +1081,18 @@ async function main(): Promise { reply.raw.setHeader('Connection', 'keep-alive'); reply.raw.setHeader('X-Accel-Buffering', 'no'); - const writeSnapshot = (snapshot: WorkerOperationState): void => { + const writeSnapshot = (snapshot: StreamedOperationState, eventId: number): void => { if (closed || reply.raw.writableEnded) return; - reply.raw.write(`event: snapshot\ndata: ${JSON.stringify(snapshot)}\n\n`); + const frameEvent: WorkerOperationEvent = { + eventId, + snapshot, + }; + reply.raw.write(`id: ${eventId}\nevent: snapshot\ndata: ${JSON.stringify(frameEvent)}\n\n`); }; let closed = false; + let consumerName: string | null = null; + let messages: Awaited> | null = null; request.raw.on('close', () => { closed = true; activeSse = Math.max(0, activeSse - 1); @@ -1032,25 +1101,54 @@ async function main(): Promise { try { let current = initial; - writeSnapshot(current); let signature = JSON.stringify(current); + writeSnapshot(current, sinceEventId > 0 ? sinceEventId : 0); + if (isTerminalStatus(current.status)) { + return reply; + } - while (!closed && !isTerminalStatus(current.status)) { - await sleep(SSE_POLL_INTERVAL_MS); - const next = await getOpState(params.data.opId); - if (!next) break; - const nextSignature = JSON.stringify(next); + const { jsm, js } = await ensureConnected(); + const consumerConfig = { + name: `op_events_${params.data.opId.slice(0, 12)}_${crypto.randomUUID().replaceAll('-', '').slice(0, 12)}`, + ack_policy: AckPolicy.None, + deliver_policy: sinceEventId > 0 ? DeliverPolicy.StartSequence : DeliverPolicy.New, + replay_policy: ReplayPolicy.Instant, + filter_subject: opEventsSubject(params.data.opId), + max_deliver: 1, + inactive_threshold: nanos(60_000_000_000), + ...(sinceEventId > 0 ? { opt_start_seq: sinceEventId + 1 } : {}), + }; + const info = await jsm.consumers.add(EVENTS_STREAM_NAME, consumerConfig); + consumerName = info.name; + const consumer = await js.consumers.get(EVENTS_STREAM_NAME, info.name); + messages = await consumer.consume(); + + for await (const msg of messages) { + if (closed) break; + const state = opEventCodec.decode(msg.data); + if (state.opId !== params.data.opId) continue; + + const nextSignature = JSON.stringify(state); if (nextSignature !== signature) { - current = next; + current = state; signature = nextSignature; - writeSnapshot(current); + writeSnapshot(current, msg.seq); } + if (isTerminalStatus(state.status)) break; } } catch (error) { app.log.warn({ opId: params.data.opId, error: toErrorMessage(error), }, 'op events stream loop error'); + } finally { + if (messages) { + await messages.close().catch(() => undefined); + } + if (consumerName) { + const { jsm } = await ensureConnected(); + await jsm.consumers.delete(EVENTS_STREAM_NAME, consumerName).catch(() => undefined); + } } if (!reply.raw.writableEnded) { diff --git a/src/app/api/documents/[id]/parsed/events/route.ts b/src/app/api/documents/[id]/parsed/events/route.ts index 44df139..536b509 100644 --- a/src/app/api/documents/[id]/parsed/events/route.ts +++ b/src/app/api/documents/[id]/parsed/events/route.ts @@ -4,6 +4,7 @@ import { db } from '@/db'; import { documents } from '@/db/schema'; import { requireAuthContext } from '@/lib/server/auth/auth'; import { readComputeMode } from '@/lib/server/compute/mode'; +import { getWorkerClientConfigFromEnv } from '@/lib/server/compute/worker'; 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'; @@ -12,6 +13,7 @@ import { healStaleDocumentParseState } from '@/lib/server/documents/parse-state- 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 type { PdfLayoutJobResult, WorkerOperationEvent, WorkerOperationState } from '@openreader/compute-core/api-contracts'; export const dynamic = 'force-dynamic'; export const runtime = 'nodejs'; @@ -31,6 +33,11 @@ type ParsedSnapshot = { parseProgress: PdfParseProgress | null; }; +type SnapshotState = { + snapshot: ParsedSnapshot; + opId: string | null; +}; + function s3NotConfiguredResponse(): NextResponse { return NextResponse.json( { error: 'Documents storage is not configured. Set S3_* environment variables.' }, @@ -42,7 +49,53 @@ function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } -async function toSnapshot(row: ParseRow): Promise { +function parseSsePayload(frame: string): string | null { + const lines = frame.replace(/\r\n/g, '\n').split('\n'); + const dataLines: string[] = []; + for (const line of lines) { + if (!line.startsWith('data:')) continue; + dataLines.push(line.slice('data:'.length).trimStart()); + } + if (dataLines.length === 0) return null; + return dataLines.join('\n'); +} + +function parseSseEventId(frame: string): number | null { + const lines = frame.replace(/\r\n/g, '\n').split('\n'); + for (const line of lines) { + if (!line.startsWith('id:')) continue; + const value = Number(line.slice('id:'.length).trim()); + if (Number.isFinite(value) && value > 0) { + return Math.floor(value); + } + } + return null; +} + +function mapWorkerStatusToParseStatus(status: WorkerOperationState['status']): PdfParseStatus { + switch (status) { + case 'queued': + return 'pending'; + case 'running': + return 'running'; + case 'succeeded': + return 'ready'; + case 'failed': + return 'failed'; + default: + return 'pending'; + } +} + +function snapshotFromWorkerState(state: WorkerOperationState): ParsedSnapshot { + const parseStatus = mapWorkerStatusToParseStatus(state.status); + return { + parseStatus, + parseProgress: parseStatus === 'running' ? (state.progress ?? null) : null, + }; +} + +async function toSnapshotState(row: ParseRow): Promise { const state = await healStaleDocumentParseState({ documentId: row.id, userId: row.userId, @@ -50,8 +103,11 @@ async function toSnapshot(row: ParseRow): Promise { }); const parseStatus = normalizeParseStatus(state.status); return { - parseStatus, - parseProgress: state.progress ?? null, + snapshot: { + parseStatus, + parseProgress: state.progress ?? null, + }, + opId: typeof state.opId === 'string' && state.opId.trim() ? state.opId.trim() : null, }; } @@ -72,6 +128,33 @@ async function loadPreferredRow(input: { return rows.find((candidate) => candidate.userId === input.storageUserId) ?? rows[0] ?? null; } +async function syncFromDb(input: { + documentId: string; + storageUserId: string; + allowedUserIds: string[]; + signature: string; + writeSnapshot: (snapshot: ParsedSnapshot) => void; +}): Promise<{ snapshot: ParsedSnapshot; opId: string | null; signature: string } | null> { + const nextRow = await loadPreferredRow({ + documentId: input.documentId, + storageUserId: input.storageUserId, + allowedUserIds: input.allowedUserIds, + }); + if (!nextRow) return null; + + const nextState = await toSnapshotState(nextRow); + const nextSignature = JSON.stringify(nextState.snapshot); + if (nextSignature !== input.signature) { + input.writeSnapshot(nextState.snapshot); + } + + return { + snapshot: nextState.snapshot, + opId: nextState.opId, + signature: nextSignature, + }; +} + export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string }> }) { try { if (!isS3Configured()) return s3NotConfiguredResponse(); @@ -100,10 +183,11 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string return NextResponse.json({ error: 'Not found' }, { status: 404 }); } - const initial = await toSnapshot(row); + const initialState = await toSnapshotState(row); const computeMode = readComputeMode(); const bus = computeMode === 'local' ? await getParseProgressBus() : null; + const workerCfg = computeMode === 'worker' ? getWorkerClientConfigFromEnv() : null; const encoder = new TextEncoder(); const stream = new ReadableStream({ @@ -112,6 +196,7 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string let unsubscribe: (() => void) | null = null; let keepaliveTimer: ReturnType | null = null; let resyncTimer: ReturnType | null = null; + let workerAbort: AbortController | null = null; const allowedUserHashes = new Set(allowedUserIds.map((candidate) => hashUserId(candidate))); @@ -134,6 +219,10 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string unsubscribe(); unsubscribe = null; } + if (workerAbort) { + workerAbort.abort(); + workerAbort = null; + } try { controller.close(); } catch { @@ -141,29 +230,8 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string } }; - 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 current = initialState.snapshot; let signature = JSON.stringify(current); writeSnapshot(current); @@ -179,7 +247,13 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string resyncTimer = setInterval(() => { if (closed) return; - void syncFromDb({ signature }).then((next) => { + void syncFromDb({ + documentId: id, + storageUserId, + allowedUserIds, + signature, + writeSnapshot, + }).then((next) => { if (closed) return; if (!next) { closeStream(); @@ -202,7 +276,13 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string documentId: id, userIdHashes: allowedUserHashes, onEvent: async () => { - const next = await syncFromDb({ signature }); + const next = await syncFromDb({ + documentId: id, + storageUserId, + allowedUserIds, + signature, + writeSnapshot, + }); if (!next) { closeStream(); return; @@ -226,24 +306,204 @@ export async function GET(req: NextRequest, ctx: { params: Promise<{ id: string unsubscribe = nextUnsubscribe; }; - const runWorkerPolling = async () => { - let current = initial; - writeSnapshot(current); + const runWorkerProxy = async () => { + if (!workerCfg) throw new Error('Worker mode selected but worker config is unavailable'); + + let current = initialState.snapshot; let signature = JSON.stringify(current); + let currentOpId = initialState.opId; + let lastEventId: number | null = null; + + 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({ + documentId: id, + storageUserId, + allowedUserIds, + signature, + writeSnapshot, + }).then((next) => { + if (closed) return; + if (!next) { + closeStream(); + return; + } + current = next.snapshot; + signature = next.signature; + if (next.opId !== currentOpId) { + currentOpId = next.opId; + lastEventId = null; + if (workerAbort) { + workerAbort.abort(); + workerAbort = null; + } + } + 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`)); + }); + }, SSE_RESYNC_INTERVAL_MS); while (!closed) { - if (current.parseStatus === 'ready' || current.parseStatus === 'failed') break; - await sleep(SSE_POLL_INTERVAL_MS); - if (closed) break; + if (!currentOpId) { + await sleep(SSE_POLL_INTERVAL_MS); + const next = await syncFromDb({ + documentId: id, + storageUserId, + allowedUserIds, + signature, + writeSnapshot, + }); + if (!next) { + closeStream(); + return; + } + current = next.snapshot; + signature = next.signature; + currentOpId = next.opId; + if (current.parseStatus === 'ready' || current.parseStatus === 'failed') { + closeStream(); + return; + } + continue; + } - const next = await syncFromDb({ signature }); - if (!next) break; + workerAbort = new AbortController(); + const query = lastEventId && lastEventId > 0 + ? `?sinceEventId=${encodeURIComponent(String(lastEventId))}` + : ''; + const response = await fetch( + `${workerCfg.baseUrl}/ops/${encodeURIComponent(currentOpId)}/events${query}`, + { + method: 'GET', + headers: { + Authorization: `Bearer ${workerCfg.token}`, + Accept: 'text/event-stream', + ...(lastEventId && lastEventId > 0 ? { 'Last-Event-ID': String(lastEventId) } : {}), + }, + cache: 'no-store', + signal: workerAbort.signal, + }, + ); + + if (closed) return; + + if (!response.ok) { + const detail = await response.text().catch(() => ''); + console.warn('[parsed/events] worker stream request failed', { + documentId: id, + opId: currentOpId, + status: response.status, + detail, + }); + await sleep(500); + continue; + } + if (!response.body) { + await sleep(500); + continue; + } + + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ''; + let streamEnded = false; + + while (!closed && !streamEnded) { + const read = await reader.read(); + if (read.done) { + streamEnded = true; + break; + } + + buffer += decoder.decode(read.value, { stream: true }); + + while (true) { + const frameEnd = buffer.indexOf('\n\n'); + if (frameEnd < 0) break; + const frame = buffer.slice(0, frameEnd); + buffer = buffer.slice(frameEnd + 2); + + const eventId = parseSseEventId(frame); + if (eventId && eventId > 0) { + lastEventId = eventId; + } + + const payload = parseSsePayload(frame); + if (!payload) continue; + + let parsed: WorkerOperationEvent | WorkerOperationState; + try { + parsed = JSON.parse(payload) as WorkerOperationEvent | WorkerOperationState; + } catch { + continue; + } + + const workerSnapshot: WorkerOperationState = ( + parsed && typeof parsed === 'object' && 'snapshot' in parsed + ? parsed.snapshot + : parsed as WorkerOperationState + ); + if (!workerSnapshot || workerSnapshot.opId !== currentOpId) continue; + + const nextSnapshot = snapshotFromWorkerState(workerSnapshot); + const nextSignature = JSON.stringify(nextSnapshot); + if (nextSignature !== signature) { + current = nextSnapshot; + signature = nextSignature; + writeSnapshot(current); + } + + if (current.parseStatus === 'ready' || current.parseStatus === 'failed') { + closeStream(); + return; + } + } + } + + if (closed) return; + + const next = await syncFromDb({ + documentId: id, + storageUserId, + allowedUserIds, + signature, + writeSnapshot, + }); + if (!next) { + closeStream(); + return; + } current = next.snapshot; signature = next.signature; + if (next.opId !== currentOpId) { + currentOpId = next.opId; + lastEventId = null; + } + if (current.parseStatus === 'ready' || current.parseStatus === 'failed') { + closeStream(); + return; + } + await sleep(250); } }; - const run = computeMode === 'local' ? runLocalRealtime : runWorkerPolling; + const run = computeMode === 'local' ? runLocalRealtime : runWorkerProxy; void run() .catch((error) => { diff --git a/src/lib/server/compute/types.ts b/src/lib/server/compute/types.ts index ab178ac..f93a8f3 100644 --- a/src/lib/server/compute/types.ts +++ b/src/lib/server/compute/types.ts @@ -4,6 +4,7 @@ import type { ParsedPdfDocument, } from '@openreader/compute-core/types'; import type { PdfLayoutProgress, WhisperAlignJobBase } from '@openreader/compute-core/api-contracts'; +import type { PdfLayoutJobResult, WorkerOperationState } from '@openreader/compute-core/api-contracts'; export type ComputeMode = 'local' | 'worker'; @@ -23,6 +24,7 @@ export interface PdfLayoutInput { pdfBytes?: ArrayBuffer; forceToken?: string; onProgress?: (progress: PdfLayoutProgress) => void | Promise; + onWorkerSnapshot?: (snapshot: WorkerOperationState) => void | Promise; } export type PdfLayoutResult = diff --git a/src/lib/server/compute/worker.ts b/src/lib/server/compute/worker.ts index 0113fe1..10d4dc2 100644 --- a/src/lib/server/compute/worker.ts +++ b/src/lib/server/compute/worker.ts @@ -4,6 +4,7 @@ import { getWorkerClientWaitTimeoutMs } from '@openreader/compute-core'; import type { PdfLayoutJobRequest, PdfLayoutJobResult, + WorkerOperationEvent, WhisperAlignJobRequest, WhisperAlignJobResult, WorkerOperationState, @@ -118,6 +119,13 @@ function normalizeWorkerBaseUrl(raw: string): string { return parsed.toString().replace(/\/+$/, ''); } +export function getWorkerClientConfigFromEnv(): { baseUrl: string; token: string } { + return { + baseUrl: normalizeWorkerBaseUrl(readRequiredEnv('COMPUTE_WORKER_URL')), + token: readRequiredEnv('COMPUTE_WORKER_TOKEN'), + }; +} + function parseRetryAfterMs(value: string | null): number | null { if (!value) return null; const asNum = Number(value); @@ -189,6 +197,18 @@ function extractSsePayload(frame: string): string | null { return dataLines.join('\n'); } +function extractSseId(frame: string): number | null { + const normalized = frame.replace(/\r\n/g, '\n'); + for (const line of normalized.split('\n')) { + if (!line.startsWith('id:')) continue; + const value = Number(line.slice('id:'.length).trim()); + if (Number.isFinite(value) && value > 0) { + return Math.floor(value); + } + } + return null; +} + type RetryMeta = { attempt: number, maxAttempts: number, @@ -240,8 +260,9 @@ export class WorkerComputeBackend implements ComputeBackend { private readonly retries: number; constructor() { - this.baseUrl = normalizeWorkerBaseUrl(readRequiredEnv('COMPUTE_WORKER_URL')); - this.token = readRequiredEnv('COMPUTE_WORKER_TOKEN'); + const cfg = getWorkerClientConfigFromEnv(); + this.baseUrl = cfg.baseUrl; + this.token = cfg.token; this.waitTimeoutMsByKind = { whisper_align: getWorkerClientWaitTimeoutMs('whisper_align'), pdf_layout: getWorkerClientWaitTimeoutMs('pdf_layout'), @@ -371,6 +392,7 @@ export class WorkerComputeBackend implements ComputeBackend { documentId: input.documentId, attempt, }); + await input.onWorkerSnapshot?.(op); const final = isTerminalStatus(op.status) ? op @@ -382,6 +404,7 @@ export class WorkerComputeBackend implements ComputeBackend { attempt, waitTimeoutMs: this.waitTimeoutMsByKind.pdf_layout, onSnapshot: (snapshot) => { + void input.onWorkerSnapshot?.(snapshot); if (snapshot.progress) { void input.onProgress?.(snapshot.progress); } @@ -513,123 +536,150 @@ export class WorkerComputeBackend implements ComputeBackend { }); try { - const res = await fetch(`${this.baseUrl}/ops/${encodeURIComponent(opId)}/events`, { - method: 'GET', - headers: { - Authorization: `Bearer ${this.token}`, - Accept: 'text/event-stream', - 'x-openreader-trace-id': traceId, - }, - signal: controller.signal, - }); + let latest: WorkerOperationState | null = null; + let lastEventId: number | null = null; + let eventCount = 0; + let lastStatus: string | null = null; - if (!res.ok) { - const retryAfterMs = parseRetryAfterMs(res.headers.get('retry-after')); - const detail = await res.text().catch(() => ''); - logWorker(res.status >= 500 ? 'warn' : 'error', 'sse.wait.http_failed', { + while (!controller.signal.aborted) { + const query = lastEventId && lastEventId > 0 + ? `?sinceEventId=${encodeURIComponent(String(lastEventId))}` + : ''; + const res = await fetch(`${this.baseUrl}/ops/${encodeURIComponent(opId)}/events${query}`, { + method: 'GET', + headers: { + Authorization: `Bearer ${this.token}`, + Accept: 'text/event-stream', + 'x-openreader-trace-id': traceId, + ...(lastEventId && lastEventId > 0 ? { 'Last-Event-ID': String(lastEventId) } : {}), + }, + signal: controller.signal, + }); + + if (!res.ok) { + const retryAfterMs = parseRetryAfterMs(res.headers.get('retry-after')); + const detail = await res.text().catch(() => ''); + logWorker(res.status >= 500 ? 'warn' : 'error', 'sse.wait.http_failed', { + ...context, + traceId, + opId, + status: res.status, + retryAfterMs, + durationMs: Date.now() - startedAt, + detail: truncateForLog(detail), + }); + throw new WorkerHttpError( + `Worker request failed (GET /ops/${encodeURIComponent(opId)}/events): ${res.status}${detail ? ` ${detail}` : ''}`, + res.status, + retryAfterMs, + ); + } + + if (!res.body) { + logWorker('error', 'sse.wait.no_body', { + ...context, + traceId, + opId, + durationMs: Date.now() - startedAt, + }); + throw new Error('Worker operation stream response has no body'); + } + + logWorker('info', 'sse.wait.connected', { ...context, traceId, opId, status: res.status, - retryAfterMs, - durationMs: Date.now() - startedAt, - detail: truncateForLog(detail), }); - throw new WorkerHttpError( - `Worker request failed (GET /ops/${encodeURIComponent(opId)}/events): ${res.status}${detail ? ` ${detail}` : ''}`, - res.status, - retryAfterMs, - ); - } - - if (!res.body) { - logWorker('error', 'sse.wait.no_body', { - ...context, - traceId, - opId, - durationMs: Date.now() - startedAt, - }); - throw new Error('Worker operation stream response has no body'); - } - - logWorker('info', 'sse.wait.connected', { - ...context, - traceId, - opId, - status: res.status, - }); - const reader = res.body.getReader(); - const decoder = new TextDecoder(); - let buffer = ''; - let latest: WorkerOperationState | null = null; - let eventCount = 0; - let lastStatus: string | null = null; - - while (true) { - const { done, value } = await reader.read(); - if (done) break; - buffer += decoder.decode(value, { stream: true }); + const reader = res.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ''; while (true) { - const frameEnd = buffer.indexOf('\n\n'); - if (frameEnd < 0) break; - const frame = buffer.slice(0, frameEnd); - buffer = buffer.slice(frameEnd + 2); + const { done, value } = await reader.read(); + if (done) break; + buffer += decoder.decode(value, { stream: true }); - const payload = extractSsePayload(frame); - if (!payload) continue; + while (true) { + const frameEnd = buffer.indexOf('\n\n'); + if (frameEnd < 0) break; + const frame = buffer.slice(0, frameEnd); + buffer = buffer.slice(frameEnd + 2); - let snapshot: WorkerOperationState; - try { - snapshot = JSON.parse(payload) as WorkerOperationState; - } catch { - logWorker('warn', 'sse.wait.json_parse_skipped', { - ...context, - traceId, - opId, - sample: truncateForLog(payload), - }); - continue; - } + const eventId = extractSseId(frame); + if (eventId && eventId > 0) { + lastEventId = eventId; + } + const payload = extractSsePayload(frame); + if (!payload) continue; - eventCount += 1; - latest = snapshot; - context.onSnapshot?.(snapshot); - if (snapshot.status !== lastStatus) { - lastStatus = snapshot.status; - logWorker('info', 'sse.wait.status', { - ...context, - traceId, - opId, - eventCount, - status: snapshot.status, - jobId: snapshot.jobId ?? null, - updatedAt: snapshot.updatedAt ?? null, - }); - } - if (isTerminalStatus(snapshot.status)) { - logWorker('info', 'sse.wait.terminal', { - ...context, - traceId, - opId, - eventCount, - status: snapshot.status, - durationMs: Date.now() - startedAt, - }); - return snapshot; + let event: WorkerOperationEvent | WorkerOperationState; + try { + event = JSON.parse(payload) as WorkerOperationEvent | WorkerOperationState; + } catch { + logWorker('warn', 'sse.wait.json_parse_skipped', { + ...context, + traceId, + opId, + sample: truncateForLog(payload), + }); + continue; + } + + const snapshot: WorkerOperationState = ( + event && typeof event === 'object' && 'snapshot' in event + ? (event as WorkerOperationEvent).snapshot + : (event as WorkerOperationState) + ); + if (!snapshot || typeof snapshot !== 'object') continue; + if (snapshot.opId !== opId) continue; + + eventCount += 1; + latest = snapshot; + context.onSnapshot?.(snapshot); + if (snapshot.status !== lastStatus) { + lastStatus = snapshot.status; + logWorker('info', 'sse.wait.status', { + ...context, + traceId, + opId, + eventCount, + status: snapshot.status, + jobId: snapshot.jobId ?? null, + updatedAt: snapshot.updatedAt ?? null, + }); + } + if (isTerminalStatus(snapshot.status)) { + logWorker('info', 'sse.wait.terminal', { + ...context, + traceId, + opId, + eventCount, + status: snapshot.status, + durationMs: Date.now() - startedAt, + }); + return snapshot; + } } } + + if (latest && isTerminalStatus(latest.status)) { + logWorker('info', 'sse.wait.terminal_after_close', { + ...context, + traceId, + opId, + eventCount, + status: latest.status, + durationMs: Date.now() - startedAt, + }); + return latest; + } + + if (controller.signal.aborted) break; + await sleep(250); } if (latest && isTerminalStatus(latest.status)) { - logWorker('info', 'sse.wait.terminal_after_close', { - ...context, - traceId, - opId, - eventCount, - status: latest.status, - durationMs: Date.now() - startedAt, - }); return latest; } diff --git a/src/lib/server/jobs/user-pdf-layout-job.ts b/src/lib/server/jobs/user-pdf-layout-job.ts index 6a40d94..1b75556 100644 --- a/src/lib/server/jobs/user-pdf-layout-job.ts +++ b/src/lib/server/jobs/user-pdf-layout-job.ts @@ -13,7 +13,12 @@ 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'; -import type { PdfLayoutJobBase, PdfLayoutProgress } from '@openreader/compute-core/api-contracts'; +import type { + PdfLayoutJobBase, + PdfLayoutJobResult, + PdfLayoutProgress, + WorkerOperationState, +} from '@openreader/compute-core/api-contracts'; type UserPdfLayoutJobRequest = PdfLayoutJobBase & { userId: string; @@ -24,6 +29,7 @@ const running = new Set(); const FOLLOWER_WAIT_TIMEOUT_MS = 180_000; const FOLLOWER_POLL_MS = 1_200; +const WORKER_PROGRESS_DB_THROTTLE_MS = 10_000; type ParseRow = { userId: string; @@ -290,7 +296,72 @@ export async function parsePdfJob(input: UserPdfLayoutJobRequest): Promise }); const compute = await getCompute(); + const isWorkerCompute = compute.mode === 'worker'; + let activeOpId: string = claimOpId; + let activeJobId: string | undefined; + let lastWorkerProgressWriteAt = 0; + let lastWorkerSnapshotWriteAt = 0; + let lastWorkerSnapshotStatus: 'pending' | 'running' | null = null; + let lastWorkerSnapshotOpId: string = claimOpId; + let lastWorkerSnapshotJobId: string | undefined; + + const persistRunningState = async (state: DocumentParseState): Promise => { + await updateParseStateForUsers({ + documentId: input.documentId, + userIds: cohortUserIds, + parseState: stringifyDocumentParseState(state), + }); + await emitParseStateEvents({ + documentId: input.documentId, + userIds: cohortUserIds, + state, + }); + }; + + const onWorkerSnapshot = async (snapshot: WorkerOperationState): Promise => { + if (!isWorkerCompute) return; + if (snapshot.opId) activeOpId = snapshot.opId; + if (snapshot.jobId) activeJobId = snapshot.jobId; + + const mappedStatus = snapshot.status === 'queued' + ? 'pending' + : snapshot.status === 'running' + ? 'running' + : null; + if (!mappedStatus) return; + + const now = Date.now(); + const forceWrite = ( + mappedStatus !== lastWorkerSnapshotStatus + || activeOpId !== lastWorkerSnapshotOpId + || activeJobId !== lastWorkerSnapshotJobId + ); + if (!forceWrite && (now - lastWorkerSnapshotWriteAt) < WORKER_PROGRESS_DB_THROTTLE_MS) { + return; + } + + const nextState: DocumentParseState = { + status: mappedStatus, + progress: snapshot.progress ?? null, + updatedAt: now, + opId: activeOpId, + ...(activeJobId ? { jobId: activeJobId } : {}), + }; + await persistRunningState(nextState); + lastWorkerSnapshotWriteAt = now; + lastWorkerSnapshotStatus = mappedStatus; + lastWorkerSnapshotOpId = activeOpId; + lastWorkerSnapshotJobId = activeJobId; + }; + const writeProgress = async (progress: PdfLayoutProgress): Promise => { + if (isWorkerCompute) { + const now = Date.now(); + if ((now - lastWorkerProgressWriteAt) < WORKER_PROGRESS_DB_THROTTLE_MS && progress.pagesParsed < progress.totalPages) { + return; + } + lastWorkerProgressWriteAt = now; + } const runningProgressState: DocumentParseState = { status: 'running', progress: { @@ -300,18 +371,10 @@ export async function parsePdfJob(input: UserPdfLayoutJobRequest): Promise phase: progress.phase, }, updatedAt: Date.now(), - opId: claimOpId, + opId: activeOpId, + ...(activeJobId ? { jobId: activeJobId } : {}), }; - await updateParseStateForUsers({ - documentId: input.documentId, - userIds: cohortUserIds, - parseState: stringifyDocumentParseState(runningProgressState), - }); - await emitParseStateEvents({ - documentId: input.documentId, - userIds: cohortUserIds, - state: runningProgressState, - }); + await persistRunningState(runningProgressState); }; const layout = await compute.parsePdfLayout({ documentId: input.documentId, @@ -319,6 +382,7 @@ export async function parsePdfJob(input: UserPdfLayoutJobRequest): Promise documentObjectKey: documentKey(input.documentId, input.namespace), forceToken: input.forceToken, onProgress: writeProgress, + onWorkerSnapshot, }); let parsedJsonKey = layout.parsedObjectKey ?? null;