refactor(worker): adopt operation-based API and opKey deduplication for compute flows

Migrate worker and backend contract to an operation-centric model, introducing opKey-based deduplication for Whisper alignment and PDF layout parsing. Add new operation request/state types and update worker endpoints to accept generic operation requests and expose operation status via SSE. Refactor client and server logic to support operation flows. Update deployment docs, environment examples, and Dockerfile to reflect new architecture and dependency management.
This commit is contained in:
Richard R 2026-05-21 09:05:26 -06:00
parent 8190b57ff2
commit a7b955e1ca
11 changed files with 1247 additions and 294 deletions

View file

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

View file

@ -67,3 +67,33 @@ export interface WorkerJobStatusResponse<Result> {
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<Result = unknown> {
opId: string;
opKey: string;
kind: WorkerOperationKind;
jobId: string;
status: WorkerJobState;
queuedAt: number;
updatedAt: number;
startedAt?: number;
result?: Result;
error?: WorkerJobErrorShape;
timing?: WorkerJobTiming;
}

View file

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

View file

@ -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<TPayload> {
jobId: string;
opId: string;
opKey: string;
kind: WorkerOperationKind;
queuedAt: number;
payload: TPayload;
}
interface StoredJobState<Result> extends WorkerJobStatusResponse<Result> {
interface StoredJobState<Result> {
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<T> = {
@ -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<T>(): JsonCodec<T> {
const encoder = new TextEncoder();
const decoder = new TextDecoder();
@ -222,6 +249,36 @@ function createJsonCodec<T>(): JsonCodec<T> {
};
}
function sleep(ms: number): Promise<void> {
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<Result>(
kv: KV,
codec: JsonCodec<StoredJobState<Result>>,
jobId: string,
state: StoredJobState<Result>,
): Promise<void> {
await kv.put(jobId, codec.encode(state));
}
async function getState<Result>(
kv: KV,
codec: JsonCodec<StoredJobState<Result>>,
jobId: string,
): Promise<StoredJobState<Result> | 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<TPayload, TResult>(input: {
consumer: Consumer;
kv: KV;
stateCodec: JsonCodec<StoredJobState<TResult>>;
jobCodec: JsonCodec<QueuedJob<TPayload>>;
run: (payload: TPayload, queueWaitMs: number) => Promise<TResult>;
maxAttempts: number;
logLabel: string;
shouldStop: () => boolean;
log: {
error: (obj: Record<string, unknown>, msg: string) => void;
info: (obj: Record<string, unknown>, msg: string) => void;
};
}): Promise<void> {
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<TPayload> | 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<TResult> = 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<void> {
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<void> {
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<void> {
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<void> {
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<void> {
const app = Fastify({ logger: buildLoggerConfig() });
const whisperStateCodec = createJsonCodec<StoredJobState<WhisperAlignJobResult>>();
const layoutStateCodec = createJsonCodec<StoredJobState<PdfLayoutJobResult>>();
const opIndexCodec = createJsonCodec<OpIndexEntry>();
const opStateCodec = createJsonCodec<WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult>>();
const whisperJobCodec = createJsonCodec<QueuedJob<WhisperAlignJobRequest>>();
const layoutJobCodec = createJsonCodec<QueuedJob<PdfLayoutJobRequest>>();
const jobStateCodec = createJsonCodec<StoredJobState<WhisperAlignJobResult | PdfLayoutJobResult>>();
const putOpState = async (state: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult>): Promise<void> => {
await kv.put(opStateKvKey(state.opId), opStateCodec.encode(state));
};
const getOpState = async (opId: string): Promise<WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult> | 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<WhisperAlignJobResult | PdfLayoutJobResult>): Promise<void> => {
await kv.put(jobStateKvKey(state.jobId), jobStateCodec.encode(state));
};
const publishQueuedJob = async (
op: WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult>,
payload: WhisperAlignJobRequest | PdfLayoutJobRequest,
): Promise<void> => {
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<WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult>> => {
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<WhisperAlignJobResult | PdfLayoutJobResult> = {
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<WhisperAlignJobResult | PdfLayoutJobResult> = {
...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<WhisperAlignJobResult | PdfLayoutJobResult> = {
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<WhisperAlignJobResult | PdfLayoutJobResult> = {
...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<void> {
}
});
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<void> {
};
}
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<WhisperAlignJobResult> = {
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<PdfLayoutJobResult> = {
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<WhisperAlignJobResult | PdfLayoutJobResult>): 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<void> {
};
};
async function processMessage<TPayload, TResult>(input: {
msg: JsMsg;
codec: JsonCodec<QueuedJob<TPayload>>;
run: (payload: TPayload, queueWaitMs: number) => Promise<TResult>;
workerLabel: string;
}): Promise<void> {
let decoded: QueuedJob<TPayload> | 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<WhisperAlignJobResult | PdfLayoutJobResult> = {
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<WhisperAlignJobResult | PdfLayoutJobResult> = {
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<WhisperAlignJobResult | PdfLayoutJobResult> = {
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<WhisperAlignJobResult | PdfLayoutJobResult> = {
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<TPayload, TResult>(input: {
consumer: Consumer;
codec: JsonCodec<QueuedJob<TPayload>>;
run: (payload: TPayload, queueWaitMs: number) => Promise<TResult>;
workerLabel: string;
}): Promise<void> {
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<void>[] = [];
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}`,
}));
}

View file

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

View file

@ -25,6 +25,7 @@ const isVercel = process.env.VERCEL === '1';
const serverExternalPackages = [
'@napi-rs/canvas',
'better-sqlite3',
'ffmpeg-static',
...(computeLocal ? ['onnxruntime-node', '@huggingface/tokenizers'] : []),
];

View file

@ -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: () => <DocumentSkeleton />
loading: () => null
}
);
@ -55,6 +54,7 @@ export default function PDFViewerPage() {
const { isAtLimit } = useAuthRateLimit();
const [error, setError] = useState<string | null>(null);
const [isLoading, setIsLoading] = useState(true);
const [isPdfViewerReady, setIsPdfViewerReady] = useState(false);
const [zoomLevel, setZoomLevel] = useState<number>(100);
const [activeSidebar, setActiveSidebar] = useState<null | 'settings' | 'audiobook' | 'segments'>(null);
const [containerHeight, setContainerHeight] = useState<string>('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() {
</div>
}
/>
<div className="overflow-hidden" style={{ height: containerHeight }}>
{isLoading || !isParseReady ? (
renderPdfStatusLoader()
) : (
<PDFViewer zoomLevel={zoomLevel} pdfState={pdfState} />
)}
<div className="relative overflow-hidden" style={{ height: containerHeight }}>
{isParseReady ? (
<div className={isPdfViewerReady ? 'h-full' : 'h-full opacity-0 pointer-events-none'}>
<PDFViewer
zoomLevel={zoomLevel}
onDocumentReady={() => setIsPdfViewerReady(true)}
pdfState={pdfState}
/>
</div>
) : null}
{isLoading || !isParseReady || !isPdfViewerReady ? (
<div className="absolute inset-0 z-10">
{renderPdfStatusLoader()}
</div>
) : null}
</div>
{canExportAudiobook && (
<AudiobookExportModal

View file

@ -1,4 +1,5 @@
import { NextRequest, NextResponse } from 'next/server';
import { randomUUID } from 'node:crypto';
import { and, eq, inArray } from 'drizzle-orm';
import { db } from '@/db';
import { ttsSegmentEntries, ttsSegmentVariants } from '@/db/schema';
@ -43,6 +44,7 @@ import type {
export const runtime = 'nodejs';
export const dynamic = 'force-dynamic';
const GENERATING_STALE_MS = 45_000;
function attachDeviceIdCookie(response: NextResponse, deviceId: string | null, didCreate: boolean) {
if (didCreate && deviceId) {
@ -150,6 +152,8 @@ async function deleteEntryIfUnused(userId: string, segmentEntryId: string): Prom
export async function POST(request: NextRequest) {
let didCreateDeviceIdCookie = false;
let deviceIdToSet: string | null = null;
const requestId = randomUUID();
const requestStartedAt = Date.now();
let computeBackendPromise: ReturnType<typeof getCompute> | 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<string, number> = {};
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,

View file

@ -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<HTMLDivElement>(null);
const [isPageRendering, setIsPageRendering] = useState(false);
const scaleRef = useRef<number>(1);
@ -495,11 +495,12 @@ export function PDFViewer({ zoomLevel, pdfState }: PDFViewerProps) {
>
<Document
key={currDocId || 'pdf'}
loading={<DocumentSkeleton />}
noData={<DocumentSkeleton />}
loading={null}
noData={null}
file={documentFile}
onLoadSuccess={(pdf) => {
onDocumentLoadSuccess(pdf);
onDocumentReady?.();
}}
onItemClick={(args: PDFOnLinkClickArgs) => {
if (args?.pageNumber) {

View file

@ -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,

View file

@ -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<string, unknown> {
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<string, unknown>): 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<string, unknown> {
if (!value || typeof value !== 'object') return {};
const record = value as Record<string, unknown>;
const summary: Record<string, unknown> = {};
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<T>(attempts: number, operation: () => Promise<T>): Promise<T> {
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<T>(
attempts: number,
operation: (attempt: number) => Promise<T>,
onAttemptError?: (meta: RetryMeta) => void,
): Promise<T> {
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<WhisperAlignJobResult>(`/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<WorkerOperationState<WhisperAlignJobResult>>('POST', '/ops', {
kind: 'whisper_align',
opKey,
payload,
}, {
traceId,
kind: 'whisper_align',
opKeyHash,
attempt,
});
const final = isTerminalStatus(op.status)
? op
: await this.waitForOperation<WhisperAlignJobResult>(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<PdfLayoutJobResult>(`/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<WorkerOperationState<PdfLayoutJobResult>>('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<PdfLayoutJobResult>(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<T>(method: 'GET' | 'POST', path: string, body?: unknown): Promise<T> {
private async requestJson<T>(
method: 'GET' | 'POST',
path: string,
body?: unknown,
context: Record<string, unknown> = {},
): Promise<T> {
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<T>;
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<Result>(path: string): Promise<WorkerJobStatusResponse<Result>> {
const started = Date.now();
let interval = POLL_INTERVAL_MS;
while ((Date.now() - started) < this.waitTimeoutMs) {
const status = await this.requestJson<WorkerJobStatusResponse<Result>>('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<Result>(
opId: string,
context: Record<string, unknown> = {},
): Promise<WorkerOperationState<Result>> {
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<Result> | 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<Result>;
try {
snapshot = JSON.parse(payload) as WorkerOperationState<Result>;
} 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`);
}
}