diff --git a/Dockerfile b/Dockerfile index 262bae0..ec8a7bb 100644 --- a/Dockerfile +++ b/Dockerfile @@ -45,19 +45,21 @@ FROM node:lts-alpine AS runner # ffmpeg is provided by ffmpeg-static from node_modules. RUN apk add --no-cache ca-certificates libreoffice-writer -# drizzle-kit is used by scripts/openreader-entrypoint.mjs for startup migrations. -RUN npm install -g drizzle-kit@0.31.10 +# Install pnpm and production deps from the repository lockfile. +RUN npm install -g pnpm@10.33.4 # App runtime directory WORKDIR /app -# Entry-point and migration scripts import dotenv directly. -RUN npm install --no-save dotenv@17.4.2 - # Copy standalone Next.js server and required static assets. COPY --from=app-builder /app/.next/standalone ./ COPY --from=app-builder /app/.next/static ./.next/static COPY --from=app-builder /app/public ./public +# Install runtime dependencies from lockfile (no Dockerfile version pin drift). +COPY --from=app-builder /app/package.json ./package.json +COPY --from=app-builder /app/pnpm-lock.yaml ./pnpm-lock.yaml +COPY --from=app-builder /app/pnpm-workspace.yaml ./pnpm-workspace.yaml +RUN CI=true pnpm install --prod --frozen-lockfile --config.confirmModulesPurge=false # Copy startup/migration scripts and migration files used by openreader-entrypoint. COPY --from=app-builder /app/scripts/openreader-entrypoint.mjs ./scripts/openreader-entrypoint.mjs COPY --from=app-builder /app/scripts/migrate-fs-v2.mjs ./scripts/migrate-fs-v2.mjs diff --git a/compute/core/src/contracts.ts b/compute/core/src/contracts.ts index e97e672..0abe295 100644 --- a/compute/core/src/contracts.ts +++ b/compute/core/src/contracts.ts @@ -67,3 +67,33 @@ export interface WorkerJobStatusResponse { error?: WorkerJobErrorShape; timing?: WorkerJobTiming; } + +export type WorkerOperationKind = 'whisper_align' | 'pdf_layout'; + +export interface WhisperAlignOperationRequest { + kind: 'whisper_align'; + opKey: string; + payload: WhisperAlignJobRequest; +} + +export interface PdfLayoutOperationRequest { + kind: 'pdf_layout'; + opKey: string; + payload: PdfLayoutJobRequest; +} + +export type WorkerOperationRequest = WhisperAlignOperationRequest | PdfLayoutOperationRequest; + +export interface WorkerOperationState { + opId: string; + opKey: string; + kind: WorkerOperationKind; + jobId: string; + status: WorkerJobState; + queuedAt: number; + updatedAt: number; + startedAt?: number; + result?: Result; + error?: WorkerJobErrorShape; + timing?: WorkerJobTiming; +} diff --git a/compute/worker/.env.example b/compute/worker/.env.example index 766f0ce..fa18018 100644 --- a/compute/worker/.env.example +++ b/compute/worker/.env.example @@ -32,3 +32,6 @@ S3_FORCE_PATH_STYLE=true COMPUTE_PREWARM_MODELS=true COMPUTE_JOBS_STREAM_MAX_BYTES=268435456 COMPUTE_JOB_STATES_MAX_BYTES=67108864 +# Optional stale window for reusing in-flight opKey entries before forcing a new attempt +# Default is max(30m, 4x max compute timeout); running jobs also refresh heartbeat every 5s +# COMPUTE_OP_STALE_MS=180000 diff --git a/compute/worker/src/server.ts b/compute/worker/src/server.ts index 4d3cbdd..6712192 100644 --- a/compute/worker/src/server.ts +++ b/compute/worker/src/server.ts @@ -1,3 +1,4 @@ +import { createHash } from 'node:crypto'; import Fastify, { type FastifyRequest } from 'fastify'; import { z } from 'zod'; import { @@ -30,8 +31,12 @@ import { type PdfLayoutJobResult, type WhisperAlignJobRequest, type WhisperAlignJobResult, - type WorkerJobStatusResponse, + type WorkerJobErrorShape, + type WorkerJobState, type WorkerJobTiming, + type WorkerOperationKind, + type WorkerOperationRequest, + type WorkerOperationState, } from '@openreader/compute-core/contracts'; import { GetObjectCommand, PutObjectCommand, S3Client } from '@aws-sdk/client-s3'; @@ -40,23 +45,40 @@ 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 JOB_STATES_BUCKET = 'job_states'; -const JOB_STATES_TTL_MS = 24 * 60 * 60 * 1000; +const COMPUTE_STATE_BUCKET = 'compute_state'; +const COMPUTE_STATE_TTL_MS = 24 * 60 * 60 * 1000; const PULL_EXPIRES_MS = 1000; const LOOP_ERROR_BACKOFF_MS = 500; +const SSE_POLL_INTERVAL_MS = 400; +const RUNNING_HEARTBEAT_MS = 5000; const DOCUMENT_ID_REGEX = /^[a-f0-9]{64}$/i; const SAFE_NAMESPACE_REGEX = /^[a-zA-Z0-9._-]{1,128}$/; interface QueuedJob { jobId: string; + opId: string; + opKey: string; + kind: WorkerOperationKind; queuedAt: number; payload: TPayload; } -interface StoredJobState extends WorkerJobStatusResponse { +interface StoredJobState { + jobId: string; + opId: string; + opKey: string; + kind: WorkerOperationKind; + status: WorkerJobState; timestamp: number; startedAt?: number; updatedAt: number; + result?: Result; + error?: WorkerJobErrorShape; + timing?: WorkerJobTiming; +} + +interface OpIndexEntry { + opId: string; } type JsonCodec = { @@ -209,6 +231,11 @@ function isAlreadyExistsError(error: unknown): boolean { return message.includes('already in use') || message.includes('already exists'); } +function isCasConflictError(error: unknown): boolean { + const message = toErrorMessage(error).toLowerCase(); + return message.includes('wrong last sequence') || message.includes('key exists') || message.includes('wrong last'); +} + function createJsonCodec(): JsonCodec { const encoder = new TextEncoder(); const decoder = new TextDecoder(); @@ -222,6 +249,36 @@ function createJsonCodec(): JsonCodec { }; } +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +function isTerminalStatus(status: WorkerJobState): boolean { + return status === 'succeeded' || status === 'failed'; +} + +function hashOpKey(opKey: string): string { + return createHash('sha256').update(opKey).digest('hex'); +} + +function opIndexKvKey(opKey: string): string { + return `op_index.${hashOpKey(opKey)}`; +} + +function opStateKvKey(opId: string): string { + return `op_state.${opId}`; +} + +function jobStateKvKey(jobId: string): string { + return `job_state.${jobId}`; +} + +function extractResultRef(kind: WorkerOperationKind, result: unknown): string | undefined { + if (kind !== 'pdf_layout' || !result || typeof result !== 'object') return undefined; + const maybe = result as { parsedObjectKey?: unknown }; + return typeof maybe.parsedObjectKey === 'string' ? maybe.parsedObjectKey : undefined; +} + const alignSchema = z.object({ text: z.string().trim().min(1), lang: z.string().trim().min(1).max(16).optional(), @@ -235,24 +292,18 @@ const layoutSchema = z.object({ documentObjectKey: z.string().trim().min(1).max(2048), }); -async function putState( - kv: KV, - codec: JsonCodec>, - jobId: string, - state: StoredJobState, -): Promise { - await kv.put(jobId, codec.encode(state)); -} - -async function getState( - kv: KV, - codec: JsonCodec>, - jobId: string, -): Promise | null> { - const entry = await kv.get(jobId); - if (!entry || entry.operation !== 'PUT') return null; - return codec.decode(entry.value); -} +const operationCreateSchema = z.discriminatedUnion('kind', [ + z.object({ + kind: z.literal('whisper_align'), + opKey: z.string().trim().min(1).max(1024), + payload: alignSchema, + }), + z.object({ + kind: z.literal('pdf_layout'), + opKey: z.string().trim().min(1).max(1024), + payload: layoutSchema, + }), +]); async function ensureJetStreamResources( jsm: JetStreamManager, @@ -308,119 +359,6 @@ async function ensureJetStreamResources( ]); } -async function createWorkerLoop(input: { - consumer: Consumer; - kv: KV; - stateCodec: JsonCodec>; - jobCodec: JsonCodec>; - run: (payload: TPayload, queueWaitMs: number) => Promise; - maxAttempts: number; - logLabel: string; - shouldStop: () => boolean; - log: { - error: (obj: Record, msg: string) => void; - info: (obj: Record, msg: string) => void; - }; -}): Promise { - while (!input.shouldStop()) { - let msg: JsMsg | null = null; - try { - msg = await input.consumer.next({ expires: PULL_EXPIRES_MS }); - } catch (error) { - if (input.shouldStop()) return; - input.log.error({ error: toErrorMessage(error), worker: input.logLabel }, 'worker pull failed'); - await new Promise((resolve) => setTimeout(resolve, LOOP_ERROR_BACKOFF_MS)); - continue; - } - - if (!msg) continue; - - let decoded: QueuedJob | null = null; - try { - const job = input.jobCodec.decode(msg.data); - decoded = job; - const startedAt = Date.now(); - const queueWaitMs = safeDurationMs(job.queuedAt, startedAt); - - await putState(input.kv, input.stateCodec, job.jobId, { - status: 'running', - timestamp: job.queuedAt, - startedAt, - updatedAt: startedAt, - ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), - }); - - const result = await input.run(job.payload, queueWaitMs ?? 0); - const resultTiming = result && typeof result === 'object' && 'timing' in result - ? (result as { timing?: WorkerJobTiming }).timing - : undefined; - - await putState(input.kv, input.stateCodec, job.jobId, { - status: 'succeeded', - timestamp: job.queuedAt, - startedAt, - updatedAt: Date.now(), - result, - ...(resultTiming ? { timing: resultTiming } : {}), - }); - - msg.ack(); - input.log.info({ worker: input.logLabel, jobId: job.jobId, timing: resultTiming }, 'job succeeded'); - } catch (error) { - const message = toErrorMessage(error); - const deliveryCount = msg.info.deliveryCount; - const hasRetriesLeft = deliveryCount < input.maxAttempts; - - if (decoded?.jobId) { - const now = Date.now(); - const queueWaitMs = safeDurationMs(decoded.queuedAt, now); - const state: StoredJobState = hasRetriesLeft - ? { - status: 'running', - timestamp: decoded.queuedAt, - updatedAt: now, - ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), - } - : { - status: 'failed', - timestamp: decoded.queuedAt, - updatedAt: now, - error: { message }, - ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), - }; - - await putState(input.kv, input.stateCodec, decoded.jobId, state).catch((stateError) => { - input.log.error({ - worker: input.logLabel, - jobId: decoded?.jobId, - error: toErrorMessage(stateError), - }, 'failed to persist failed state'); - }); - } - - if (hasRetriesLeft) { - msg.nak(); - input.log.error({ - worker: input.logLabel, - jobId: decoded?.jobId, - error: message, - deliveryCount, - maxAttempts: input.maxAttempts, - }, 'job failed, nacked for retry'); - } else { - msg.term(message); - input.log.error({ - worker: input.logLabel, - jobId: decoded?.jobId, - error: message, - deliveryCount, - maxAttempts: input.maxAttempts, - }, 'job failed, max attempts reached'); - } - } - } -} - async function main(): Promise { const port = readIntEnv('PORT', 8081); const host = process.env.COMPUTE_WORKER_HOST?.trim() || '0.0.0.0'; @@ -435,6 +373,10 @@ async function main(): Promise { const prewarmModels = parseBoolEnv('COMPUTE_PREWARM_MODELS', true); const jobsStreamMaxBytes = readIntEnv('COMPUTE_JOBS_STREAM_MAX_BYTES', 256 * 1024 * 1024); const jobStatesMaxBytes = readIntEnv('COMPUTE_JOB_STATES_MAX_BYTES', 64 * 1024 * 1024); + const opStaleMs = readIntEnv( + 'COMPUTE_OP_STALE_MS', + Math.max(30 * 60_000, Math.max(whisperTimeoutMs, pdfTimeoutMs) * 4), + ); const connectOpts: any = { servers: natsUrl }; const natsCreds = process.env.NATS_CREDS?.trim(); @@ -445,7 +387,7 @@ async function main(): Promise { connectOpts.authenticator = credsAuthenticator(new TextEncoder().encode(natsCreds)); } else if (natsCredsFile) { console.log(`[compute-worker] Connecting to NATS using credentials file: ${natsCredsFile}`); - const { readFileSync } = require('fs'); + const { readFileSync } = await import('node:fs'); const credsData = readFileSync(natsCredsFile); connectOpts.authenticator = credsAuthenticator(credsData); } @@ -456,9 +398,9 @@ async function main(): Promise { await ensureJetStreamResources(jsm, whisperTimeoutMs, pdfTimeoutMs, attempts, jobsStreamMaxBytes); - const kv = await new Kvm(js).create(JOB_STATES_BUCKET, { + const kv = await new Kvm(js).create(COMPUTE_STATE_BUCKET, { history: 1, - ttl: JOB_STATES_TTL_MS, + ttl: COMPUTE_STATE_TTL_MS, max_bytes: jobStatesMaxBytes, }); @@ -503,10 +445,187 @@ async function main(): Promise { const app = Fastify({ logger: buildLoggerConfig() }); - const whisperStateCodec = createJsonCodec>(); - const layoutStateCodec = createJsonCodec>(); + const opIndexCodec = createJsonCodec(); + const opStateCodec = createJsonCodec>(); const whisperJobCodec = createJsonCodec>(); const layoutJobCodec = createJsonCodec>(); + const jobStateCodec = createJsonCodec>(); + + const putOpState = async (state: WorkerOperationState): Promise => { + await kv.put(opStateKvKey(state.opId), opStateCodec.encode(state)); + }; + + const getOpState = async (opId: string): Promise | null> => { + const entry = await kv.get(opStateKvKey(opId)); + if (!entry || entry.operation !== 'PUT') return null; + return opStateCodec.decode(entry.value); + }; + + const putJobState = async (state: StoredJobState): Promise => { + await kv.put(jobStateKvKey(state.jobId), jobStateCodec.encode(state)); + }; + + const publishQueuedJob = async ( + op: WorkerOperationState, + payload: WhisperAlignJobRequest | PdfLayoutJobRequest, + ): Promise => { + if (op.kind === 'whisper_align') { + await js.publish(WHISPER_JOBS_SUBJECT, whisperJobCodec.encode({ + jobId: op.jobId, + opId: op.opId, + opKey: op.opKey, + kind: 'whisper_align', + queuedAt: op.queuedAt, + payload: payload as WhisperAlignJobRequest, + })); + return; + } + + await js.publish(LAYOUT_JOBS_SUBJECT, layoutJobCodec.encode({ + jobId: op.jobId, + opId: op.opId, + opKey: op.opKey, + kind: 'pdf_layout', + queuedAt: op.queuedAt, + payload: payload as PdfLayoutJobRequest, + })); + }; + + const enqueueOrReuseOperation = async ( + req: WorkerOperationRequest, + ): Promise> => { + const opKey = req.opKey.trim(); + const indexKey = opIndexKvKey(opKey); + + for (let attemptNo = 0; attemptNo < 10; attemptNo += 1) { + const indexEntry = await kv.get(indexKey); + if (indexEntry && indexEntry.operation === 'PUT') { + const pointer = opIndexCodec.decode(indexEntry.value); + const current = await getOpState(pointer.opId); + + if (!current) { + await sleep(25); + continue; + } + + if (current && current.kind === req.kind) { + const ageMs = Date.now() - current.updatedAt; + if (current.status === 'succeeded') return current; + if ((current.status === 'queued' || current.status === 'running') && ageMs <= opStaleMs) { + return current; + } + } + + const now = Date.now(); + const replacement: WorkerOperationState = { + opId: crypto.randomUUID(), + opKey, + kind: req.kind, + jobId: crypto.randomUUID(), + status: 'queued', + queuedAt: now, + updatedAt: now, + }; + + try { + await kv.update(indexKey, opIndexCodec.encode({ opId: replacement.opId }), indexEntry.revision); + } catch (error) { + if (isCasConflictError(error)) continue; + throw error; + } + + await putOpState(replacement); + await putJobState({ + jobId: replacement.jobId, + opId: replacement.opId, + opKey: replacement.opKey, + kind: replacement.kind, + status: 'queued', + timestamp: now, + updatedAt: now, + }); + + try { + await publishQueuedJob(replacement, req.payload); + return replacement; + } catch (error) { + const failed: WorkerOperationState = { + ...replacement, + status: 'failed', + updatedAt: Date.now(), + error: { message: toErrorMessage(error) }, + }; + await putOpState(failed); + await putJobState({ + jobId: replacement.jobId, + opId: replacement.opId, + opKey: replacement.opKey, + kind: replacement.kind, + status: 'failed', + timestamp: replacement.queuedAt, + updatedAt: failed.updatedAt, + error: failed.error, + }); + return failed; + } + } + + const now = Date.now(); + const created: WorkerOperationState = { + opId: crypto.randomUUID(), + opKey, + kind: req.kind, + jobId: crypto.randomUUID(), + status: 'queued', + queuedAt: now, + updatedAt: now, + }; + + try { + await kv.create(indexKey, opIndexCodec.encode({ opId: created.opId })); + } catch (error) { + if (isCasConflictError(error)) continue; + throw error; + } + + await putOpState(created); + await putJobState({ + jobId: created.jobId, + opId: created.opId, + opKey: created.opKey, + kind: created.kind, + status: 'queued', + timestamp: now, + updatedAt: now, + }); + + try { + await publishQueuedJob(created, req.payload); + return created; + } catch (error) { + const failed: WorkerOperationState = { + ...created, + status: 'failed', + updatedAt: Date.now(), + error: { message: toErrorMessage(error) }, + }; + await putOpState(failed); + await putJobState({ + jobId: created.jobId, + opId: created.opId, + opKey: created.opKey, + kind: created.kind, + status: 'failed', + timestamp: created.queuedAt, + updatedAt: failed.updatedAt, + error: failed.error, + }); + return failed; + } + } + + throw new Error('Unable to reserve operation after repeated CAS conflicts'); + }; app.addHook('onRequest', async (request, reply) => { const path = request.url.split('?')[0] ?? request.url; @@ -532,8 +651,8 @@ async function main(): Promise { } }); - app.post('/align/whisper/jobs', async (request, reply) => { - const parsed = alignSchema.safeParse(request.body); + app.post('/ops', async (request, reply) => { + const parsed = operationCreateSchema.safeParse(request.body); if (!parsed.success) { reply.code(400); return { @@ -542,98 +661,74 @@ async function main(): Promise { }; } - const jobId = crypto.randomUUID(); - const queuedAt = Date.now(); - - await putState(kv, whisperStateCodec, jobId, { - status: 'queued', - timestamp: queuedAt, - updatedAt: queuedAt, - }); - - await js.publish(WHISPER_JOBS_SUBJECT, whisperJobCodec.encode({ - jobId, - queuedAt, - payload: parsed.data, - })); - + const op = await enqueueOrReuseOperation(parsed.data as WorkerOperationRequest); reply.code(202); - return { jobId }; + return op; }); - app.get('/align/whisper/jobs/:jobId', async (request, reply) => { - const params = z.object({ jobId: z.string().trim().min(1) }).safeParse(request.params); + app.get('/ops/:opId', async (request, reply) => { + const params = z.object({ opId: z.string().trim().min(1) }).safeParse(request.params); if (!params.success) { reply.code(400); - return { error: 'Invalid job id' }; + return { error: 'Invalid op id' }; } - const state = await getState(kv, whisperStateCodec, params.data.jobId); + const state = await getOpState(params.data.opId); if (!state) { reply.code(404); - return { error: 'Job not found' }; + return { error: 'Operation not found' }; } - const response: WorkerJobStatusResponse = { - status: state.status, - ...(state.result ? { result: state.result } : {}), - ...(state.error ? { error: state.error } : {}), - ...(state.timing ? { timing: state.timing } : {}), - }; - - return response; + return state; }); - app.post('/layout/pdf/jobs', async (request, reply) => { - const parsed = layoutSchema.safeParse(request.body); - if (!parsed.success) { - reply.code(400); - return { - error: 'Invalid request body', - issues: parsed.error.issues, - }; - } - - const jobId = crypto.randomUUID(); - const queuedAt = Date.now(); - - await putState(kv, layoutStateCodec, jobId, { - status: 'queued', - timestamp: queuedAt, - updatedAt: queuedAt, - }); - - await js.publish(LAYOUT_JOBS_SUBJECT, layoutJobCodec.encode({ - jobId, - queuedAt, - payload: parsed.data, - })); - - reply.code(202); - return { jobId }; - }); - - app.get('/layout/pdf/jobs/:jobId', async (request, reply) => { - const params = z.object({ jobId: z.string().trim().min(1) }).safeParse(request.params); + app.get('/ops/:opId/events', async (request, reply) => { + const params = z.object({ opId: z.string().trim().min(1) }).safeParse(request.params); if (!params.success) { reply.code(400); - return { error: 'Invalid job id' }; + return { error: 'Invalid op id' }; } - const state = await getState(kv, layoutStateCodec, params.data.jobId); - if (!state) { + const initial = await getOpState(params.data.opId); + if (!initial) { reply.code(404); - return { error: 'Job not found' }; + return { error: 'Operation not found' }; } - const response: WorkerJobStatusResponse = { - status: state.status, - ...(state.result ? { result: state.result } : {}), - ...(state.error ? { error: state.error } : {}), - ...(state.timing ? { timing: state.timing } : {}), + reply.hijack(); + reply.raw.setHeader('Content-Type', 'text/event-stream; charset=utf-8'); + reply.raw.setHeader('Cache-Control', 'no-cache, no-transform'); + reply.raw.setHeader('Connection', 'keep-alive'); + reply.raw.setHeader('X-Accel-Buffering', 'no'); + + const writeSnapshot = (snapshot: WorkerOperationState): void => { + reply.raw.write(`event: snapshot\ndata: ${JSON.stringify(snapshot)}\n\n`); }; - return response; + let closed = false; + request.raw.on('close', () => { + closed = true; + }); + + let current = initial; + writeSnapshot(current); + let signature = JSON.stringify(current); + + 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); + if (nextSignature !== signature) { + current = next; + signature = nextSignature; + writeSnapshot(current); + } + } + + if (!reply.raw.writableEnded) { + reply.raw.end(); + } }); const whisperConsumer = await js.consumers.get(JOBS_STREAM_NAME, WHISPER_CONSUMER_NAME); @@ -706,33 +801,247 @@ async function main(): Promise { }; }; + async function processMessage(input: { + msg: JsMsg; + codec: JsonCodec>; + run: (payload: TPayload, queueWaitMs: number) => Promise; + workerLabel: string; + }): Promise { + let decoded: QueuedJob | null = null; + let heartbeat: NodeJS.Timeout | null = null; + try { + decoded = input.codec.decode(input.msg.data); + const startedAt = Date.now(); + const queueWaitMs = safeDurationMs(decoded.queuedAt, startedAt); + + const runningState: WorkerOperationState = { + opId: decoded.opId, + opKey: decoded.opKey, + kind: decoded.kind, + jobId: decoded.jobId, + status: 'running', + queuedAt: decoded.queuedAt, + startedAt, + updatedAt: startedAt, + ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), + }; + + await putOpState(runningState); + await putJobState({ + jobId: decoded.jobId, + opId: decoded.opId, + opKey: decoded.opKey, + kind: decoded.kind, + status: 'running', + timestamp: decoded.queuedAt, + startedAt, + updatedAt: startedAt, + ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), + }); + + heartbeat = setInterval(() => { + const now = Date.now(); + const heartbeatState: WorkerOperationState = { + opId: decoded!.opId, + opKey: decoded!.opKey, + kind: decoded!.kind, + jobId: decoded!.jobId, + status: 'running', + queuedAt: decoded!.queuedAt, + startedAt, + updatedAt: now, + ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), + }; + void putOpState(heartbeatState).catch((stateError) => { + app.log.error({ + worker: input.workerLabel, + opId: decoded?.opId, + jobId: decoded?.jobId, + error: toErrorMessage(stateError), + }, 'failed to persist running heartbeat op state'); + }); + void putJobState({ + jobId: decoded!.jobId, + opId: decoded!.opId, + opKey: decoded!.opKey, + kind: decoded!.kind, + status: 'running', + timestamp: decoded!.queuedAt, + startedAt, + updatedAt: now, + ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), + }).catch((stateError) => { + app.log.error({ + worker: input.workerLabel, + opId: decoded?.opId, + jobId: decoded?.jobId, + error: toErrorMessage(stateError), + }, 'failed to persist running heartbeat job state'); + }); + }, RUNNING_HEARTBEAT_MS); + + const result = await input.run(decoded.payload, queueWaitMs ?? 0); + const resultTiming = result && typeof result === 'object' && 'timing' in result + ? (result as { timing?: WorkerJobTiming }).timing + : undefined; + const now = Date.now(); + + const succeededState: WorkerOperationState = { + opId: decoded.opId, + opKey: decoded.opKey, + kind: decoded.kind, + jobId: decoded.jobId, + status: 'succeeded', + queuedAt: decoded.queuedAt, + startedAt, + updatedAt: now, + result: result as WhisperAlignJobResult | PdfLayoutJobResult, + ...(resultTiming ? { timing: resultTiming } : {}), + }; + + await putOpState(succeededState); + await putJobState({ + jobId: decoded.jobId, + opId: decoded.opId, + opKey: decoded.opKey, + kind: decoded.kind, + status: 'succeeded', + timestamp: decoded.queuedAt, + startedAt, + updatedAt: now, + result: result as WhisperAlignJobResult | PdfLayoutJobResult, + ...(resultTiming ? { timing: resultTiming } : {}), + }); + + input.msg.ack(); + app.log.info({ + worker: input.workerLabel, + opId: decoded.opId, + jobId: decoded.jobId, + resultRef: extractResultRef(decoded.kind, result), + timing: resultTiming, + }, 'job succeeded'); + } catch (error) { + const message = toErrorMessage(error); + const deliveryCount = input.msg.info.deliveryCount; + const hasRetriesLeft = deliveryCount < attempts; + + if (decoded) { + const now = Date.now(); + const queueWaitMs = safeDurationMs(decoded.queuedAt, now); + + const status: WorkerJobState = hasRetriesLeft ? 'running' : 'failed'; + const opState: WorkerOperationState = { + opId: decoded.opId, + opKey: decoded.opKey, + kind: decoded.kind, + jobId: decoded.jobId, + status, + queuedAt: decoded.queuedAt, + updatedAt: now, + ...(status === 'failed' ? { error: { message } } : {}), + ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), + }; + + await putOpState(opState).catch((stateError) => { + app.log.error({ + worker: input.workerLabel, + opId: decoded?.opId, + jobId: decoded?.jobId, + error: toErrorMessage(stateError), + }, 'failed to persist operation state'); + }); + + await putJobState({ + jobId: decoded.jobId, + opId: decoded.opId, + opKey: decoded.opKey, + kind: decoded.kind, + status, + timestamp: decoded.queuedAt, + updatedAt: now, + ...(status === 'failed' ? { error: { message } } : {}), + ...(typeof queueWaitMs === 'number' ? { timing: { queueWaitMs } } : {}), + }).catch((stateError) => { + app.log.error({ + worker: input.workerLabel, + opId: decoded?.opId, + jobId: decoded?.jobId, + error: toErrorMessage(stateError), + }, 'failed to persist job state'); + }); + } + + if (hasRetriesLeft) { + input.msg.nak(); + app.log.error({ + worker: input.workerLabel, + opId: decoded?.opId, + jobId: decoded?.jobId, + error: message, + deliveryCount, + maxAttempts: attempts, + }, 'job failed, nacked for retry'); + } else { + input.msg.term(message); + app.log.error({ + worker: input.workerLabel, + opId: decoded?.opId, + jobId: decoded?.jobId, + error: message, + deliveryCount, + maxAttempts: attempts, + }, 'job failed, max attempts reached'); + } + } finally { + if (heartbeat) clearInterval(heartbeat); + } + } + + async function createWorkerLoop(input: { + consumer: Consumer; + codec: JsonCodec>; + run: (payload: TPayload, queueWaitMs: number) => Promise; + workerLabel: string; + }): Promise { + while (!stopping) { + let msg: JsMsg | null = null; + try { + msg = await input.consumer.next({ expires: PULL_EXPIRES_MS }); + } catch (error) { + if (stopping) return; + app.log.error({ error: toErrorMessage(error), worker: input.workerLabel }, 'worker pull failed'); + await sleep(LOOP_ERROR_BACKOFF_MS); + continue; + } + + if (!msg) continue; + await processMessage({ + msg, + codec: input.codec, + run: input.run, + workerLabel: input.workerLabel, + }); + } + } + const workerLoops: Promise[] = []; for (let i = 0; i < whisperConcurrency; i += 1) { workerLoops.push(createWorkerLoop({ consumer: whisperConsumer, - kv, - stateCodec: whisperStateCodec, - jobCodec: whisperJobCodec, + codec: whisperJobCodec, run: runWhisper, - maxAttempts: attempts, - logLabel: `whisper-${i + 1}`, - shouldStop: () => stopping, - log: app.log, + workerLabel: `whisper-${i + 1}`, })); } for (let i = 0; i < pdfConcurrency; i += 1) { workerLoops.push(createWorkerLoop({ consumer: layoutConsumer, - kv, - stateCodec: layoutStateCodec, - jobCodec: layoutJobCodec, + codec: layoutJobCodec, run: runLayout, - maxAttempts: attempts, - logLabel: `layout-${i + 1}`, - shouldStop: () => stopping, - log: app.log, + workerLabel: `layout-${i + 1}`, })); } diff --git a/docs-site/docs/deploy/compute-worker.md b/docs-site/docs/deploy/compute-worker.md index 50bdfd4..454f704 100644 --- a/docs-site/docs/deploy/compute-worker.md +++ b/docs-site/docs/deploy/compute-worker.md @@ -7,10 +7,10 @@ Use this guide for `COMPUTE_MODE=worker` deployments where heavy compute runs ou The compute worker handles: -- Whisper word alignment (`/align/whisper/jobs`) -- PDF layout parsing (`/layout/pdf/jobs`) +- Whisper word alignment operations +- PDF layout parsing operations -The app server enqueues jobs and polls status. Queue durability and retries are backed by NATS JetStream WorkQueue consumers and NATS KV. +The app server submits operations to `POST /ops`, reuses in-flight work via required `opKey`, and consumes status updates via `GET /ops/:opId/events` (SSE). Queue durability and retries are backed by NATS JetStream WorkQueue consumers and NATS KV. ## Published image diff --git a/next.config.ts b/next.config.ts index 18fc239..0cbfb73 100644 --- a/next.config.ts +++ b/next.config.ts @@ -25,6 +25,7 @@ const isVercel = process.env.VERCEL === '1'; const serverExternalPackages = [ '@napi-rs/canvas', 'better-sqlite3', + 'ffmpeg-static', ...(computeLocal ? ['onnxruntime-node', '@huggingface/tokenizers'] : []), ]; diff --git a/src/app/(app)/pdf/[id]/page.tsx b/src/app/(app)/pdf/[id]/page.tsx index bad2b1e..9646da7 100644 --- a/src/app/(app)/pdf/[id]/page.tsx +++ b/src/app/(app)/pdf/[id]/page.tsx @@ -4,7 +4,6 @@ import dynamic from 'next/dynamic'; import { useParams, useRouter } from 'next/navigation'; import Link from 'next/link'; import { useCallback, useEffect, useRef, useState, type MouseEvent } from 'react'; -import { DocumentSkeleton } from '@/components/documents/DocumentSkeleton'; import { useTTS } from '@/contexts/TTSContext'; import { DocumentSettings } from '@/components/documents/DocumentSettings'; import { DocumentHeaderMenu } from '@/components/documents/DocumentHeaderMenu'; @@ -27,7 +26,7 @@ const PDFViewer = dynamic( () => import('@/components/views/PDFViewer').then((module) => module.PDFViewer), { ssr: false, - loading: () => + loading: () => null } ); @@ -55,6 +54,7 @@ export default function PDFViewerPage() { const { isAtLimit } = useAuthRateLimit(); const [error, setError] = useState(null); const [isLoading, setIsLoading] = useState(true); + const [isPdfViewerReady, setIsPdfViewerReady] = useState(false); const [zoomLevel, setZoomLevel] = useState(100); const [activeSidebar, setActiveSidebar] = useState(null); const [containerHeight, setContainerHeight] = useState('auto'); @@ -68,6 +68,7 @@ export default function PDFViewerPage() { useEffect(() => { setIsLoading(true); + setIsPdfViewerReady(false); setError(null); setActiveSidebar(null); inFlightDocIdRef.current = null; @@ -303,12 +304,21 @@ export default function PDFViewerPage() { } /> -
- {isLoading || !isParseReady ? ( - renderPdfStatusLoader() - ) : ( - - )} +
+ {isParseReady ? ( +
+ setIsPdfViewerReady(true)} + pdfState={pdfState} + /> +
+ ) : null} + {isLoading || !isParseReady || !isPdfViewerReady ? ( +
+ {renderPdfStatusLoader()} +
+ ) : null}
{canExportAudiobook && ( | null = null; const getComputeBackend = async () => { if (!computeBackendPromise) computeBackendPromise = getCompute(); @@ -325,6 +329,9 @@ export async function POST(request: NextRequest) { }; for (const segment of normalized) { + const segmentStartedAt = Date.now(); + const stageTimings: Record = {}; + let failedStage = 'unknown'; const locatorProjection = projectSegmentLocator(segment.locator); const segmentKeyForRow = typeof segment.original.segmentKey === 'string' && segment.original.segmentKey.trim() ? segment.original.segmentKey.trim() @@ -377,11 +384,13 @@ export async function POST(request: NextRequest) { // previously unavailable, retry alignment using the current segment text. if (!alignment) { try { + const alignStartedAt = Date.now(); const computeBackend = await getComputeBackend(); const aligned = (await computeBackend.alignWords({ audioObjectKey: existing.audioKey, text: segment.text, })).alignments; + stageTimings.selfHealAlignMs = Date.now() - alignStartedAt; alignment = aligned[0] ? { ...aligned[0], sentenceIndex: segment.original.segmentIndex } : null; if (alignment) { @@ -398,6 +407,7 @@ export async function POST(request: NextRequest) { } } catch (alignError) { console.warn('Whisper alignment still unavailable for completed segment; continuing without word highlights.', { + requestId, segmentId: segment.segmentId, error: alignError instanceof Error ? alignError.message : String(alignError), }); @@ -504,7 +514,92 @@ export async function POST(request: NextRequest) { await deleteEntryIfUnused(scope.storageUserId, movedFromEntryId); } + const [currentVariant] = await db + .select({ + status: ttsSegmentVariants.status, + updatedAt: ttsSegmentVariants.updatedAt, + error: ttsSegmentVariants.error, + audioKey: ttsSegmentVariants.audioKey, + }) + .from(ttsSegmentVariants) + .where(and( + eq(ttsSegmentVariants.segmentId, segment.segmentId), + eq(ttsSegmentVariants.userId, scope.storageUserId), + )) + .limit(1); + + if (!currentVariant) { + manifest.push({ + segmentId: segment.segmentId, + segmentIndex: segment.original.segmentIndex, + segmentKey: segmentKeyForRow, + audioPresignUrl: null, + audioFallbackUrl: null, + durationMs: 0, + alignment: null, + locator: segment.locator, + status: 'pending', + }); + continue; + } + + if (currentVariant.status === 'generating') { + const lastUpdatedAt = Number(currentVariant.updatedAt ?? 0); + const isFresh = lastUpdatedAt > 0 && (Date.now() - lastUpdatedAt) < GENERATING_STALE_MS; + if (isFresh) { + manifest.push({ + segmentId: segment.segmentId, + segmentIndex: segment.original.segmentIndex, + segmentKey: segmentKeyForRow, + audioPresignUrl: null, + audioFallbackUrl: null, + durationMs: 0, + alignment: null, + locator: segment.locator, + status: 'pending', + }); + continue; + } + } + + if (currentVariant.status === 'pending' || currentVariant.status === 'error' || currentVariant.status === 'generating') { + const expectedUpdatedAt = Number(currentVariant.updatedAt ?? 0); + const [claim] = await db + .update(ttsSegmentVariants) + .set({ + status: 'generating', + error: null, + updatedAt: Date.now(), + }) + .where(and( + eq(ttsSegmentVariants.segmentId, segment.segmentId), + eq(ttsSegmentVariants.userId, scope.storageUserId), + eq(ttsSegmentVariants.status, currentVariant.status), + eq(ttsSegmentVariants.updatedAt, expectedUpdatedAt), + )) + .returning({ + status: ttsSegmentVariants.status, + }); + + if (!claim) { + manifest.push({ + segmentId: segment.segmentId, + segmentIndex: segment.original.segmentIndex, + segmentKey: segmentKeyForRow, + audioPresignUrl: null, + audioFallbackUrl: null, + durationMs: 0, + alignment: null, + locator: segment.locator, + status: 'pending', + }); + continue; + } + } + try { + failedStage = 'tts.generate'; + const ttsStartedAt = Date.now(); const ttsBuffer = await generateTTSBuffer({ text: segment.text, voice: effectiveSettings.voice, @@ -517,33 +612,47 @@ export async function POST(request: NextRequest) { baseUrl: requestCreds.baseUrl, testNamespace: scope.testNamespace, }, request.signal); + stageTimings.generateTtsMs = Date.now() - ttsStartedAt; + failedStage = 's3.put_audio'; + const putStartedAt = Date.now(); await putTtsSegmentAudioObject(audioKey, ttsBuffer); + stageTimings.putAudioMs = Date.now() - putStartedAt; let persistedBuffer = ttsBuffer; if (persistedBuffer.byteLength === 0) { + failedStage = 's3.get_audio_after_empty_put'; + const getStartedAt = Date.now(); persistedBuffer = await getTtsSegmentAudioObject(audioKey); + stageTimings.getAudioAfterEmptyPutMs = Date.now() - getStartedAt; } + failedStage = 'audio.probe_duration'; + const probeStartedAt = Date.now(); const durationMs = await probeAudioDurationMsFromBuffer(persistedBuffer, request.signal); + stageTimings.probeDurationMs = Date.now() - probeStartedAt; let alignment: TTSSegmentManifestItem['alignment'] = null; try { - const whisperBytes = Uint8Array.from(persistedBuffer); + failedStage = 'whisper.align'; + const alignStartedAt = Date.now(); const computeBackend = await getComputeBackend(); const aligned = (await computeBackend.alignWords({ - audioBuffer: whisperBytes.buffer, audioObjectKey: audioKey, text: segment.text, })).alignments; + stageTimings.whisperAlignMs = Date.now() - alignStartedAt; alignment = aligned[0] ? { ...aligned[0], sentenceIndex: segment.original.segmentIndex } : null; } catch (alignError) { console.warn('Whisper alignment unavailable for segment; continuing without word highlights.', { + requestId, segmentId: segment.segmentId, error: alignError instanceof Error ? alignError.message : String(alignError), }); alignment = null; } + failedStage = 'db.mark_completed'; + const markCompletedStartedAt = Date.now(); await db .update(ttsSegmentVariants) .set({ @@ -557,12 +666,16 @@ export async function POST(request: NextRequest) { eq(ttsSegmentVariants.segmentId, segment.segmentId), eq(ttsSegmentVariants.userId, scope.storageUserId), )); + stageTimings.markCompletedMs = Date.now() - markCompletedStartedAt; + failedStage = 'resolve.audio_urls'; + const resolveUrlsStartedAt = Date.now(); const audioUrls = await resolveSegmentAudioUrls({ documentId: parsed.documentId, segmentId: segment.segmentId, audioKey, }); + stageTimings.resolveAudioUrlsMs = Date.now() - resolveUrlsStartedAt; manifest.push({ segmentId: segment.segmentId, @@ -586,6 +699,19 @@ export async function POST(request: NextRequest) { : upstreamStatus && upstreamStatus >= 500 ? 'UPSTREAM_TTS_ERROR' : 'TTS_SEGMENT_GENERATION_FAILED'; + console.error('[tts-segments/ensure] segment failed', { + requestId, + documentId: parsed.documentId, + segmentId: segment.segmentId, + failedStage, + elapsedMs: Date.now() - segmentStartedAt, + stageTimings, + aborted, + upstreamStatus, + retryAfterSeconds, + errorCode, + message, + }); await db .update(ttsSegmentVariants) .set({ @@ -619,11 +745,33 @@ export async function POST(request: NextRequest) { detail: message, ...(typeof upstreamStatus === 'number' ? { upstreamStatus } : {}), ...(typeof retryAfterSeconds === 'number' ? { retryAfterSeconds } : {}), - }, + }, }); } } + const completedCount = manifest.filter((s) => s.status === 'completed').length; + const pendingCount = manifest.filter((s) => s.status === 'pending').length; + const errorItems = manifest.filter((s) => s.status === 'error'); + if (errorItems.length > 0) { + console.error('[tts-segments/ensure] partial result', { + requestId, + documentId: parsed.documentId, + total: manifest.length, + completedCount, + pendingCount, + errorCount: errorItems.length, + elapsedMs: Date.now() - requestStartedAt, + errors: errorItems.slice(0, 5).map((item) => ({ + segmentId: item.segmentId, + code: item.error?.code ?? null, + detail: item.error?.detail ?? null, + upstreamStatus: item.error?.upstreamStatus ?? null, + retryAfterSeconds: item.error?.retryAfterSeconds ?? null, + })), + }); + } + const response = NextResponse.json({ documentId: parsed.documentId, segments: manifest, diff --git a/src/components/views/PDFViewer.tsx b/src/components/views/PDFViewer.tsx index 8bbecec..947f9c7 100644 --- a/src/components/views/PDFViewer.tsx +++ b/src/components/views/PDFViewer.tsx @@ -5,7 +5,6 @@ import { Document, Page } from 'react-pdf'; import type { Dest } from 'react-pdf/src/shared/types.js'; import 'react-pdf/dist/Page/AnnotationLayer.css'; import 'react-pdf/dist/Page/TextLayer.css'; -import { DocumentSkeleton } from '@/components/documents/DocumentSkeleton'; import { useTTS } from '@/contexts/TTSContext'; import { useConfig } from '@/contexts/ConfigContext'; import { usePDFResize } from '@/hooks/pdf/usePDFResize'; @@ -14,6 +13,7 @@ import type { ParsedPdfBlock, ParsedPdfPage } from '@/types/parsed-pdf'; interface PDFViewerProps { zoomLevel: number; + onDocumentReady?: () => void; pdfState: Pick< PdfDocumentState, | 'highlightPattern' @@ -36,7 +36,7 @@ interface PDFOnLinkClickArgs { dest?: Dest; } -export function PDFViewer({ zoomLevel, pdfState }: PDFViewerProps) { +export function PDFViewer({ zoomLevel, onDocumentReady, pdfState }: PDFViewerProps) { const containerRef = useRef(null); const [isPageRendering, setIsPageRendering] = useState(false); const scaleRef = useRef(1); @@ -495,11 +495,12 @@ export function PDFViewer({ zoomLevel, pdfState }: PDFViewerProps) { > } - noData={} + loading={null} + noData={null} file={documentFile} onLoadSuccess={(pdf) => { onDocumentLoadSuccess(pdf); + onDocumentReady?.(); }} onItemClick={(args: PDFOnLinkClickArgs) => { if (args?.pageNumber) { diff --git a/src/lib/server/compute/worker-contract.ts b/src/lib/server/compute/worker-contract.ts index a93c984..19e6bc5 100644 --- a/src/lib/server/compute/worker-contract.ts +++ b/src/lib/server/compute/worker-contract.ts @@ -1,11 +1,14 @@ export { - ALIGN_QUEUE_NAME, - PDF_LAYOUT_QUEUE_NAME, type PdfLayoutJobRequest, type PdfLayoutJobResult, + type PdfLayoutOperationRequest, type WhisperAlignJobRequest, type WhisperAlignJobResult, + type WhisperAlignOperationRequest, type WorkerJobErrorShape, + type WorkerOperationKind, + type WorkerOperationRequest, + type WorkerOperationState, type WorkerJobState, type WorkerJobTiming, type WorkerJobStatusResponse, diff --git a/src/lib/server/compute/worker.ts b/src/lib/server/compute/worker.ts index 500ad6f..55c9197 100644 --- a/src/lib/server/compute/worker.ts +++ b/src/lib/server/compute/worker.ts @@ -1,10 +1,11 @@ +import { createHash, randomUUID } from 'node:crypto'; import type { ComputeBackend, PdfLayoutInput, WhisperAlignInput, WhisperAlignResult } from '@/lib/server/compute/types'; import type { PdfLayoutJobRequest, PdfLayoutJobResult, WhisperAlignJobRequest, WhisperAlignJobResult, - WorkerJobStatusResponse, + WorkerOperationState, } from '@/lib/server/compute/worker-contract'; class WorkerHttpError extends Error { @@ -21,8 +22,72 @@ class WorkerHttpError extends Error { const DEFAULT_WAIT_TIMEOUT_MS = 120_000; const DEFAULT_RETRIES = 2; -const POLL_INTERVAL_MS = 400; -const POLL_MAX_INTERVAL_MS = 1_500; +const LOG_PREFIX = '[compute-worker-client]'; +const MAX_LOG_DETAIL_CHARS = 600; +const LOG_EVENTS = new Set([ + 'align.request.failed', + 'align.request.attempt_error', + 'pdf_layout.request.failed', + 'pdf_layout.request.attempt_error', + 'http.request.failed', + 'sse.wait.http_failed', + 'sse.wait.ended_without_terminal', + 'sse.wait.failed', +]); + +type WorkerLogLevel = 'info' | 'warn' | 'error'; + +function truncateForLog(value: string, maxChars = MAX_LOG_DETAIL_CHARS): string { + if (value.length <= maxChars) return value; + return `${value.slice(0, maxChars)}...`; +} + +function errorToLog(error: unknown): Record { + if (error instanceof WorkerHttpError) { + return { + name: error.name, + message: error.message, + status: error.status, + retryAfterMs: error.retryAfterMs, + stack: error.stack, + }; + } + if (error instanceof Error) { + return { name: error.name, message: error.message, stack: error.stack }; + } + return { message: String(error) }; +} + +function logWorker(level: WorkerLogLevel, event: string, fields: Record): void { + if (!LOG_EVENTS.has(event)) return; + + const payload = { + ts: new Date().toISOString(), + event, + ...fields, + }; + const line = `${LOG_PREFIX} ${JSON.stringify(payload)}`; + if (level === 'error') { + console.error(line); + return; + } + if (level === 'warn') { + console.warn(line); + return; + } + console.info(line); +} + +function opSummary(value: unknown): Record { + if (!value || typeof value !== 'object') return {}; + const record = value as Record; + const summary: Record = {}; + if (typeof record.opId === 'string') summary.opId = record.opId; + if (typeof record.status === 'string') summary.status = record.status; + if (typeof record.jobId === 'string') summary.jobId = record.jobId; + if (typeof record.updatedAt === 'string') summary.updatedAt = record.updatedAt; + return summary; +} function readRequiredEnv(name: string): string { const value = process.env[name]?.trim(); @@ -79,19 +144,88 @@ function shouldRetry(error: unknown): boolean { return false; } -async function withRetries(attempts: number, operation: () => Promise): Promise { +function isTerminalStatus(status: string): boolean { + return status === 'succeeded' || status === 'failed'; +} + +function sha256Hex(input: string): string { + return createHash('sha256').update(input).digest('hex'); +} + +function buildWhisperOpKey(input: WhisperAlignInput): string { + const cacheKey = input.cacheKey?.trim(); + if (cacheKey) { + return `whisper_align|v1|cache|${cacheKey}|${input.audioObjectKey}`; + } + return [ + 'whisper_align', + 'v1', + input.audioObjectKey, + input.lang ?? '', + sha256Hex(input.text), + ].join('|'); +} + +function buildPdfOpKey(input: PdfLayoutInput): string { + return [ + 'pdf_layout', + 'v1', + input.documentId, + input.namespace ?? '', + input.documentObjectKey ?? '', + ].join('|'); +} + +function extractSsePayload(frame: string): string | null { + const dataLines: string[] = []; + const normalized = frame.replace(/\r\n/g, '\n'); + for (const line of normalized.split('\n')) { + if (line.startsWith('data:')) { + dataLines.push(line.slice('data:'.length).trimStart()); + } + } + if (dataLines.length === 0) return null; + return dataLines.join('\n'); +} + +type RetryMeta = { + attempt: number, + maxAttempts: number, + willRetry: boolean, + delayMs: number | null, + error: unknown, +}; + +async function withRetries( + attempts: number, + operation: (attempt: number) => Promise, + onAttemptError?: (meta: RetryMeta) => void, +): Promise { let lastError: unknown = null; - for (let attempt = 0; attempt < attempts; attempt += 1) { + for (let attemptIndex = 0; attemptIndex < attempts; attemptIndex += 1) { + const attempt = attemptIndex + 1; try { - return await operation(); + return await operation(attempt); } catch (error) { lastError = error; - if (attempt === attempts - 1 || !shouldRetry(error)) break; - if (error instanceof WorkerHttpError && typeof error.retryAfterMs === 'number') { - await sleep(error.retryAfterMs); - } else { - await sleep((attempt + 1) * 250); + const willRetry = attemptIndex < attempts - 1 && shouldRetry(error); + let delayMs: number | null = null; + if (willRetry) { + if (error instanceof WorkerHttpError && typeof error.retryAfterMs === 'number') { + delayMs = error.retryAfterMs; + } else { + delayMs = attempt * 250; + } } + onAttemptError?.({ + attempt, + maxAttempts: attempts, + willRetry, + delayMs, + error, + }); + if (!willRetry) break; + await sleep(delayMs ?? 0); } } throw lastError instanceof Error ? lastError : new Error('Unknown worker compute failure'); @@ -121,15 +255,78 @@ export class WorkerComputeBackend implements ComputeBackend { cacheKey: input.cacheKey, audioObjectKey: input.audioObjectKey, }; - - return withRetries(this.retries, async () => { - const { jobId } = await this.requestJson<{ jobId: string }>('POST', '/align/whisper/jobs', payload); - const status = await this.waitForJob(`/align/whisper/jobs/${encodeURIComponent(jobId)}`); - if (status.status !== 'succeeded' || !status.result) { - throw new Error(status.error?.message || 'Whisper worker job did not complete'); - } - return { alignments: status.result.alignments }; + const traceId = randomUUID(); + const opKey = buildWhisperOpKey(input); + const opKeyHash = sha256Hex(opKey).slice(0, 16); + const startedAt = Date.now(); + logWorker('info', 'align.request.start', { + traceId, + kind: 'whisper_align', + opKeyHash, + audioObjectKey: input.audioObjectKey, + cacheKey: input.cacheKey ?? null, + lang: input.lang ?? null, + textLength: input.text.length, + waitTimeoutMs: this.waitTimeoutMs, + maxRetries: this.retries, }); + + try { + const result = await withRetries(this.retries, async (attempt) => { + const op = await this.requestJson>('POST', '/ops', { + kind: 'whisper_align', + opKey, + payload, + }, { + traceId, + kind: 'whisper_align', + opKeyHash, + attempt, + }); + + const final = isTerminalStatus(op.status) + ? op + : await this.waitForOperation(op.opId, { + traceId, + kind: 'whisper_align', + opKeyHash, + attempt, + }); + + if (final.status !== 'succeeded' || !final.result) { + throw new Error(final.error?.message || 'Whisper worker operation did not complete'); + } + return { alignments: final.result.alignments }; + }, ({ attempt, maxAttempts, willRetry, delayMs, error }) => { + logWorker(willRetry ? 'warn' : 'error', 'align.request.attempt_error', { + traceId, + kind: 'whisper_align', + opKeyHash, + attempt, + maxAttempts, + willRetry, + delayMs, + error: errorToLog(error), + }); + }); + + logWorker('info', 'align.request.succeeded', { + traceId, + kind: 'whisper_align', + opKeyHash, + durationMs: Date.now() - startedAt, + }); + return result; + } catch (error) { + logWorker('error', 'align.request.failed', { + traceId, + kind: 'whisper_align', + opKeyHash, + durationMs: Date.now() - startedAt, + error: errorToLog(error), + }); + throw error; + } } async parsePdfLayout(input: PdfLayoutInput) { @@ -141,27 +338,109 @@ export class WorkerComputeBackend implements ComputeBackend { namespace: input.namespace ?? null, documentObjectKey: input.documentObjectKey, }; - return withRetries(this.retries, async () => { - const { jobId } = await this.requestJson<{ jobId: string }>('POST', '/layout/pdf/jobs', payload); - const status = await this.waitForJob(`/layout/pdf/jobs/${encodeURIComponent(jobId)}`); - if (status.status !== 'succeeded' || !status.result) { - throw new Error(status.error?.message || 'PDF layout worker job did not complete'); - } - if (status.result.parsedObjectKey) { - return { parsedObjectKey: status.result.parsedObjectKey }; - } - if (status.result.parsed) { - return { parsed: status.result.parsed }; - } - throw new Error('PDF layout worker job completed without parsed output'); + const traceId = randomUUID(); + const opKey = buildPdfOpKey(input); + const opKeyHash = sha256Hex(opKey).slice(0, 16); + const startedAt = Date.now(); + logWorker('info', 'pdf_layout.request.start', { + traceId, + kind: 'pdf_layout', + opKeyHash, + documentId: input.documentId, + namespace: input.namespace ?? null, + documentObjectKey: input.documentObjectKey, + waitTimeoutMs: this.waitTimeoutMs, + maxRetries: this.retries, }); + + try { + const result = await withRetries(this.retries, async (attempt) => { + const op = await this.requestJson>('POST', '/ops', { + kind: 'pdf_layout', + opKey, + payload, + }, { + traceId, + kind: 'pdf_layout', + opKeyHash, + documentId: input.documentId, + attempt, + }); + + const final = isTerminalStatus(op.status) + ? op + : await this.waitForOperation(op.opId, { + traceId, + kind: 'pdf_layout', + opKeyHash, + documentId: input.documentId, + attempt, + }); + + if (final.status !== 'succeeded' || !final.result) { + throw new Error(final.error?.message || 'PDF layout worker operation did not complete'); + } + if (final.result.parsedObjectKey) { + return { parsedObjectKey: final.result.parsedObjectKey }; + } + if (final.result.parsed) { + return { parsed: final.result.parsed }; + } + throw new Error('PDF layout worker operation completed without parsed output'); + }, ({ attempt, maxAttempts, willRetry, delayMs, error }) => { + logWorker(willRetry ? 'warn' : 'error', 'pdf_layout.request.attempt_error', { + traceId, + kind: 'pdf_layout', + opKeyHash, + documentId: input.documentId, + attempt, + maxAttempts, + willRetry, + delayMs, + error: errorToLog(error), + }); + }); + + logWorker('info', 'pdf_layout.request.succeeded', { + traceId, + kind: 'pdf_layout', + opKeyHash, + documentId: input.documentId, + durationMs: Date.now() - startedAt, + }); + return result; + } catch (error) { + logWorker('error', 'pdf_layout.request.failed', { + traceId, + kind: 'pdf_layout', + opKeyHash, + documentId: input.documentId, + durationMs: Date.now() - startedAt, + error: errorToLog(error), + }); + throw error; + } } - private async requestJson(method: 'GET' | 'POST', path: string, body?: unknown): Promise { + private async requestJson( + method: 'GET' | 'POST', + path: string, + body?: unknown, + context: Record = {}, + ): Promise { + const startedAt = Date.now(); + const traceId = typeof context.traceId === 'string' ? context.traceId : randomUUID(); + logWorker('info', 'http.request.start', { + ...context, + traceId, + method, + path, + }); const res = await fetch(`${this.baseUrl}${path}`, { method, headers: { Authorization: `Bearer ${this.token}`, + 'x-openreader-trace-id': traceId, ...(method === 'POST' ? { 'Content-Type': 'application/json' } : {}), }, ...(method === 'POST' ? { body: JSON.stringify(body ?? {}) } : {}), @@ -170,6 +449,16 @@ export class WorkerComputeBackend implements ComputeBackend { if (!res.ok) { const retryAfterMs = parseRetryAfterMs(res.headers.get('retry-after')); const detail = await res.text().catch(() => ''); + logWorker(res.status >= 500 ? 'warn' : 'error', 'http.request.failed', { + ...context, + traceId, + method, + path, + status: res.status, + retryAfterMs, + durationMs: Date.now() - startedAt, + detail: truncateForLog(detail), + }); throw new WorkerHttpError( `Worker request failed (${method} ${path}): ${res.status}${detail ? ` ${detail}` : ''}`, res.status, @@ -177,18 +466,175 @@ export class WorkerComputeBackend implements ComputeBackend { ); } - return res.json() as Promise; + const parsed = await res.json() as T; + const operationSummary = opSummary(parsed); + logWorker('info', 'http.request.succeeded', { + ...context, + traceId, + method, + path, + httpStatus: res.status, + durationMs: Date.now() - startedAt, + ...operationSummary, + }); + return parsed; } - private async waitForJob(path: string): Promise> { - const started = Date.now(); - let interval = POLL_INTERVAL_MS; - while ((Date.now() - started) < this.waitTimeoutMs) { - const status = await this.requestJson>('GET', path); - if (status.status === 'succeeded' || status.status === 'failed') return status; - await sleep(interval); - interval = Math.min(POLL_MAX_INTERVAL_MS, Math.floor(interval * 1.5)); + private async waitForOperation( + opId: string, + context: Record = {}, + ): Promise> { + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), this.waitTimeoutMs); + const startedAt = Date.now(); + const traceId = typeof context.traceId === 'string' ? context.traceId : randomUUID(); + logWorker('info', 'sse.wait.start', { + ...context, + traceId, + opId, + waitTimeoutMs: this.waitTimeoutMs, + }); + + 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, + }); + + 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, + }); + 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 }); + + while (true) { + const frameEnd = buffer.indexOf('\n\n'); + if (frameEnd < 0) break; + const frame = buffer.slice(0, frameEnd); + buffer = buffer.slice(frameEnd + 2); + + const payload = extractSsePayload(frame); + if (!payload) continue; + + 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; + } + + eventCount += 1; + latest = 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; + } + + logWorker('error', 'sse.wait.ended_without_terminal', { + ...context, + traceId, + opId, + eventCount, + latestStatus: latest?.status ?? null, + durationMs: Date.now() - startedAt, + }); + throw new Error(`Operation stream ended before terminal state for op ${opId}`); + } catch (error) { + logWorker('error', 'sse.wait.failed', { + ...context, + traceId, + opId, + durationMs: Date.now() - startedAt, + error: errorToLog(error), + }); + throw error; + } finally { + clearTimeout(timeout); } - throw new Error(`Timed out waiting for worker job after ${this.waitTimeoutMs}ms`); } }