diff --git a/compute-worker/openapi.json b/compute-worker/openapi.json index c0973f1..5b1f63a 100644 --- a/compute-worker/openapi.json +++ b/compute-worker/openapi.json @@ -272,7 +272,7 @@ ], "additionalProperties": false }, - "PublicOperation": { + "ComputeOperation": { "type": "object", "properties": { "opId": { @@ -413,7 +413,7 @@ ], "additionalProperties": false }, - "PublicOperationEvent": { + "ComputeOperationEvent": { "type": "object", "properties": { "eventId": { @@ -1706,12 +1706,12 @@ ], "responses": { "200": { - "description": "Server-sent PublicOperationEvent stream", + "description": "Server-sent ComputeOperationEvent stream", "content": { "application/json": { "schema": { "type": "string", - "description": "Server-sent PublicOperationEvent stream" + "description": "Server-sent ComputeOperationEvent stream" } } } diff --git a/compute-worker/src/api/app.ts b/compute-worker/src/api/app.ts index 77b5a55..7ea327c 100644 --- a/compute-worker/src/api/app.ts +++ b/compute-worker/src/api/app.ts @@ -1,18 +1,9 @@ -import Fastify, { type FastifyInstance, type FastifyRequest } from 'fastify'; +import Fastify, { type FastifyInstance } from 'fastify'; import swagger from '@fastify/swagger'; import { - connect, credsAuthenticator, - type NatsConnection, + type ConnectionOptions, } from '@nats-io/transport-node'; -import { - jetstream, - jetstreamManager, - type Consumer, - type JetStreamClient, - type JetStreamManager, -} from '@nats-io/jetstream'; -import { Kvm } from '@nats-io/kv'; import { ensureComputeModels, } from '../inference/runtime'; @@ -21,30 +12,26 @@ import { getComputeOpStaleMs, getAvailableCpuCores, getOnnxThreadsPerJob, + buildLoggerConfig, + normalizeNatsReplicas, + readBoolEnv, + readPositiveIntEnv, + requireEnv, } from '../infrastructure/config'; -import { encodeSseFrame, OperationOrchestrator } from '../operations'; +import { OperationOrchestrator } from '../operations'; import type { PdfLayoutJobRequest, PdfLayoutJobResult, - WorkerOperationEvent, WhisperAlignJobRequest, WhisperAlignJobResult, - WorkerJobTiming, - WorkerOperationRequest, - WorkerOperationState, -} from '../api/contracts'; +} from '../operations/contracts'; import { JetStreamOperationEventStream, JetStreamOperationQueue, JetStreamOperationStateStore, - hashOpKey, } from '../infrastructure/nats-adapters'; import { createJsonCodec } from '../infrastructure/json-codec'; -import { - recoverOrphanedOperations, - type StreamedOperationState, -} from '../operations/recovery'; -import { parsedPdfArtifactKey } from '../storage/artifact-addressing'; +import { createOperationReconciler } from '../operations/reconciliation'; import { createArtifactStorage, createS3ClientFromEnv, @@ -53,44 +40,31 @@ import { } from '../infrastructure/storage'; import { createJobHandlers } from '../jobs/handlers'; import { createWorkerLoopController, type QueuedJob } from '../jobs/worker-loop'; +import { createNatsSessionManager } from '../infrastructure/nats-session'; import { - COMPUTE_STATE_BUCKET, - COMPUTE_STATE_TTL_MS, EVENTS_STREAM_NAME, - JOBS_STREAM_NAME, - LAYOUT_CONSUMER_NAME, LAYOUT_JOBS_SUBJECT, NATS_API_TIMEOUT_MS, - WHISPER_CONSUMER_NAME, WHISPER_JOBS_SUBJECT, - ensureJetStreamResources, } from '../infrastructure/nats'; -import { buildPdfOperationKey, buildWhisperOperationKey } from './operation-keys'; -import { toPublicOperation, type PublicOperationEvent } from './public-operation'; +import { registerHttpHooks } from './http-hooks'; +import { + registerComputeWorkerRoutes, + type ComputeWorkerRouteDeps, +} from './routes'; import { apiErrorResponseSchema, artifactReferenceSchema, jsonSchema, - operationEventsQuerySchema, - operationParamsSchema, operationErrorSchema, parsedPdfDocumentSchema, pdfLayoutProgressSchema, pdfLayoutResolutionSchema, - pdfOperationCreateSchema, - pdfResolveSchema, - publicOperationEventSchema, - publicOperationSchema, + computeOperationEventSchema, + computeOperationSchema, ttsSentenceAlignmentSchema, - whisperOperationCreateSchema, } from './schemas'; -const OP_EVENTS_KEEPALIVE_MS = 15_000; -// Reconnection delay handed to the browser EventSource via the SSE `retry:` -// directive. When a silent stream is torn down for idle sleep, this keeps the -// client from immediately reconnecting and re-waking the worker; instead it -// reconnects on a slow cadence so the container stays asleep most of the time. -const OP_EVENTS_RECONNECT_HINT_MS = 120_000; // Disconnect from NATS after this much continuous idle so the worker stops // generating outbound traffic (pull polling + keepalive PINGs) and Railway can // put it to sleep. Reconnect happens lazily on the next inbound request. @@ -100,76 +74,9 @@ const IDLE_STATUS_LOG_INTERVAL_MS = 60_000; const ORPHAN_SWEEP_INTERVAL_MS = 15_000; // Bounded pull window so consumer loops yield periodically and can be stopped // cleanly when going idle, instead of blocking on a long-lived pull. -const REQUEST_STARTED_AT_MS_KEY = Symbol('request-started-at-ms'); -const REQUEST_COUNTED_KEY = Symbol('request-activity-counted'); const WHISPER_MAX_DELIVER = 1; -interface NatsSession { - nc: NatsConnection; - js: JetStreamClient; - jsm: JetStreamManager; - kv: Awaited>; - whisperConsumer: Consumer; - layoutConsumer: Consumer; -} - -interface OperationEventStreamLike { - subscribe(input: { - opId: string; - sinceEventId?: number; - onEvent: (event: WorkerOperationEvent) => void | Promise; - onError?: (error: unknown) => void; - }): Promise<() => void>; -} - -interface OperationStateStoreLike { - getOpState(opId: string): Promise; - getOpStateRecord?(opId: string): Promise<{ state: StreamedOperationState; revision: number } | null>; - getOpIndex?(opKey: string): Promise<{ opId: string } | null>; - listOpStates?(): Promise; -} - -interface OrchestratorLike { - enqueueOrReuse(request: WorkerOperationRequest): Promise; - markRunning(input: { - opId: string; - startedAt?: number; - updatedAt?: number; - timing?: WorkerJobTiming; - }): Promise; - markProgress(input: { - opId: string; - progress: WorkerOperationState['progress']; - updatedAt?: number; - timing?: WorkerJobTiming; - }): Promise; - markSucceeded(input: { - opId: string; - result: unknown; - updatedAt?: number; - timing?: WorkerJobTiming; - }): Promise; - markFailed(input: { - opId: string; - error: { message: string; code?: string } | string; - updatedAt?: number; - timing?: WorkerJobTiming; - }): Promise; - markFailedIfUnchanged?(input: { - current: StreamedOperationState; - expectedRevision: number; - error: { message: string; code?: string } | string; - updatedAt?: number; - timing?: WorkerJobTiming; - }): Promise; -} - -export interface ComputeWorkerRouteDeps { - orchestrator: OrchestratorLike; - operationStateStore: OperationStateStoreLike; - operationEventStream: OperationEventStreamLike; - artifactExists?: (key: string) => Promise; -} +export type { ComputeWorkerRouteDeps } from './routes'; export interface CreateComputeWorkerAppOptions { host?: string; @@ -187,122 +94,27 @@ export interface ComputeWorkerApp { close(): Promise; } -function requireEnv(name: string): string { - const value = process.env[name]?.trim(); - if (!value) throw new Error(`${name} is required`); - return value; -} - -function readIntEnv(name: string, fallback: number): number { - const raw = process.env[name]?.trim(); - if (!raw) return fallback; - const parsed = Number(raw); - if (!Number.isFinite(parsed) || parsed <= 0) return fallback; - return Math.floor(parsed); -} - -function normalizeNatsReplicas(value: number): number { - if (value === 3 || value === 5) return value; - return 1; -} - -function parseBoolEnv(name: string, fallback: boolean): boolean { - const raw = process.env[name]?.trim(); - if (!raw) return fallback; - const normalized = raw.toLowerCase(); - return normalized === '1' || normalized === 'true' || normalized === 'yes' || normalized === 'on'; -} - -function buildLoggerConfig(): boolean | Record { - const format = (process.env.LOG_FORMAT?.trim().toLowerCase() || 'pretty'); - const level = process.env.COMPUTE_LOG_LEVEL?.trim() || 'info'; - if (format === 'json') { - return { - level, - base: null, - }; - } - return { - level, - base: null, - transport: { - target: 'pino-pretty', - options: { - colorize: true, - translateTime: 'SYS:standard', - ignore: 'pid,hostname', - }, - }, - }; -} - -function isAuthed(request: FastifyRequest, expectedToken: string): boolean { - const auth = request.headers.authorization; - if (!auth?.startsWith('Bearer ')) return false; - const token = auth.slice('Bearer '.length).trim(); - return token === expectedToken; -} - -function toErrorMessage(error: unknown): string { - if (error instanceof Error && error.message) return error.message; - return String(error); -} - -function requestPath(request: FastifyRequest): string { - return request.url.split('?')[0] ?? request.url; -} - -function isHealthPath(path: string): boolean { - return path === '/health/live' || path === '/health/ready'; -} - -function extractTraceId(request: FastifyRequest): string | null { - const header = request.headers['x-openreader-trace-id']; - if (Array.isArray(header)) return header[0] ?? null; - return typeof header === 'string' ? header : null; -} - -function extractOpId(request: FastifyRequest, path: string): string | null { - const params = request.params as { opId?: unknown } | undefined; - if (params && typeof params.opId === 'string' && params.opId.trim()) { - return params.opId.trim(); - } - const match = path.match(/^\/v1\/operations\/([^/]+)/); - if (!match?.[1]) return null; - try { - return decodeURIComponent(match[1]); - } catch { - return match[1]; - } -} - -function isTerminalStatus(status: import('../api/contracts').WorkerJobState): boolean { - return status === 'succeeded' || status === 'failed'; -} - -const errorResponseSchema = jsonSchema(apiErrorResponseSchema); - export async function createComputeWorkerApp(options: CreateComputeWorkerAppOptions = {}): Promise { - const port = options.port ?? readIntEnv('PORT', 8081); + const port = options.port ?? readPositiveIntEnv('PORT', 8081); const host = options.host ?? (process.env.COMPUTE_WORKER_HOST?.trim() || '0.0.0.0'); const workerToken = options.workerToken ?? requireEnv('COMPUTE_WORKER_TOKEN'); const disableWorkers = options.disableWorkers ?? false; const natsUrl = requireEnv('NATS_URL'); const timeoutConfig = getComputeTimeoutConfig(); - const jobConcurrency = readIntEnv('COMPUTE_JOB_CONCURRENCY', 1); + const jobConcurrency = readPositiveIntEnv('COMPUTE_JOB_CONCURRENCY', 1); const whisperTimeoutMs = timeoutConfig.whisperTimeoutMs; const pdfTimeoutMs = timeoutConfig.pdfTimeoutMs; const pdfHardCapMs = timeoutConfig.pdfHardCapMs; - const pdfAttempts = readIntEnv('COMPUTE_PDF_JOB_ATTEMPTS', 1); - const prewarmModels = parseBoolEnv('COMPUTE_PREWARM_MODELS', false); - const jobsStreamMaxBytes = readIntEnv('COMPUTE_JOBS_STREAM_MAX_BYTES', 256 * 1024 * 1024); - const eventsStreamMaxBytes = readIntEnv('COMPUTE_EVENTS_STREAM_MAX_BYTES', 128 * 1024 * 1024); - const jobStatesMaxBytes = readIntEnv('COMPUTE_JOB_STATES_MAX_BYTES', 64 * 1024 * 1024); - const natsReplicas = normalizeNatsReplicas(readIntEnv('COMPUTE_NATS_REPLICAS', 1)); + const pdfAttempts = readPositiveIntEnv('COMPUTE_PDF_JOB_ATTEMPTS', 1); + const prewarmModels = readBoolEnv('COMPUTE_PREWARM_MODELS', false); + const jobsStreamMaxBytes = readPositiveIntEnv('COMPUTE_JOBS_STREAM_MAX_BYTES', 256 * 1024 * 1024); + const eventsStreamMaxBytes = readPositiveIntEnv('COMPUTE_EVENTS_STREAM_MAX_BYTES', 128 * 1024 * 1024); + const jobStatesMaxBytes = readPositiveIntEnv('COMPUTE_JOB_STATES_MAX_BYTES', 64 * 1024 * 1024); + const natsReplicas = normalizeNatsReplicas(readPositiveIntEnv('COMPUTE_NATS_REPLICAS', 1)); const opStaleMs = getComputeOpStaleMs(); - const connectOpts: Parameters[0] = { servers: natsUrl }; + const connectOpts: ConnectionOptions = { servers: natsUrl }; const natsCreds = process.env.NATS_CREDS?.trim(); const natsCredsFile = process.env.NATS_CREDS_FILE?.trim(); @@ -316,160 +128,18 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti connectOpts.authenticator = credsAuthenticator(credsData); } - // Lazy NATS connection lifecycle. The worker connects on demand (first request - // that needs the queue/KV) and disconnects after IDLE_DISCONNECT_MS of full idle - // so it stops emitting outbound traffic and Railway can sleep it. Reconnect is - // transparent: any inbound operation request both wakes the container and re-establishes - // the session via ensureConnected(). - let session: NatsSession | null = null; - let connecting: Promise | null = null; - let idleTimer: NodeJS.Timeout | null = null; - let orphanSweepTimer: NodeJS.Timeout | null = null; let stopping = false; - - // Activity accounting feeding the idle detector. The worker is considered idle - // only when no HTTP request is in flight, no SSE stream is open, no job is - // processing, and nothing has happened for IDLE_DISCONNECT_MS. let inFlightHttp = 0; let activeSse = 0; let inFlightJobs = 0; let lastActivityAt = Date.now(); - let lastActivityReason = 'startup'; - let lastIdleStatusLogAt = 0; - + let lastActivityReason = "startup"; const markActivity = (reason: string): void => { lastActivityAt = Date.now(); lastActivityReason = reason; }; - - function startIdleTimer(): void { - if (idleTimer) return; - idleTimer = setInterval(() => { - if (!session || stopping) return; - const now = Date.now(); - const idleForMs = now - lastActivityAt; - if (now - lastIdleStatusLogAt >= IDLE_STATUS_LOG_INTERVAL_MS) { - lastIdleStatusLogAt = now; - app.log.info({ - activeSse, - idleForMs, - inFlightHttp, - inFlightJobs, - lastActivityReason, - disconnectEligible: inFlightHttp === 0 - && inFlightJobs === 0 - && idleForMs >= IDLE_DISCONNECT_MS, - }, 'nats idle status'); - } - // Hard work in flight always blocks idle. An open SSE no longer blocks just - // by existing — only by delivering events, which refresh lastActivityAt via - // markActivity() in the stream's onEvent. So a silent/stuck-op stream lets - // the idle window elapse and is torn down by disconnect() below. - if (inFlightHttp > 0 || inFlightJobs > 0) return; - if (idleForMs < IDLE_DISCONNECT_MS) return; - void disconnect('idle'); - }, IDLE_CHECK_INTERVAL_MS); - // Don't let the idle checker keep the process alive on its own. - idleTimer.unref?.(); - } - - function startOrphanSweepTimer(): void { - if (orphanSweepTimer) return; - orphanSweepTimer = setInterval(() => { - if (!session || stopping) return; - void runOrphanedOpRecovery({ force: true }).catch((error) => { - app.log.error({ - error: toErrorMessage(error), - }, 'periodic orphaned operation recovery failed'); - }); - }, ORPHAN_SWEEP_INTERVAL_MS); - orphanSweepTimer.unref?.(); - } - - async function disconnect(reason: string): Promise { - const current = session; - if (!current) return; - // Snapshot what's still attached so a dropped connection is visible in logs - // (e.g. SSE streams torn down because we went idle while they were open). - app.log.info({ - reason, - activeSse, - inFlightHttp, - inFlightJobs, - idleForMs: Date.now() - lastActivityAt, - }, 'nats dropping connection'); - // Clear synchronously (before any await) so concurrent requests reconnect a - // fresh session instead of using the connection we're about to close. - session = null; - if (idleTimer) { - clearInterval(idleTimer); - idleTimer = null; - } - if (orphanSweepTimer) { - clearInterval(orphanSweepTimer); - orphanSweepTimer = null; - } - try { - await current.nc.close(); - } catch { - // ignore close errors - } - await workerLoops.stop(); - app.log.info({ reason }, 'nats disconnected'); - } - - async function ensureConnected(): Promise { - if (session) return session; - if (connecting) return connecting; - connecting = (async () => { - const nc: NatsConnection = await connect(connectOpts); - const js: JetStreamClient = jetstream(nc, { timeout: NATS_API_TIMEOUT_MS }); - const jsm: JetStreamManager = await jetstreamManager(nc, { timeout: NATS_API_TIMEOUT_MS }); - await ensureJetStreamResources({ - jsm, - whisperTimeoutMs, - pdfTimeoutMs, - pdfAttempts, - jobsMaxBytes: jobsStreamMaxBytes, - eventsMaxBytes: eventsStreamMaxBytes, - natsReplicas, - }); - const kv = await new Kvm(js).create(COMPUTE_STATE_BUCKET, { - replicas: natsReplicas, - history: 1, - ttl: COMPUTE_STATE_TTL_MS, - max_bytes: jobStatesMaxBytes, - }); - const whisperConsumer = await js.consumers.get(JOBS_STREAM_NAME, WHISPER_CONSUMER_NAME); - const layoutConsumer = await js.consumers.get(JOBS_STREAM_NAME, LAYOUT_CONSUMER_NAME); - const next: NatsSession = { nc, js, jsm, kv, whisperConsumer, layoutConsumer }; - session = next; - sessionGeneration += 1; - orphanRecoveryDoneForGeneration = -1; - markActivity('nats_connected'); - if (!disableWorkers) { - workerLoops.start(next, { - whisper: next.whisperConsumer, - pdfLayout: next.layoutConsumer, - }); - } - startIdleTimer(); - startOrphanSweepTimer(); - // Safety net: if the connection closes for any reason (network drop after - // exhausting reconnects, or our own disconnect), drop the stale session so - // the next request reconnects cleanly. - void nc.closed().then(() => { - if (session?.nc === nc) session = null; - }); - app.log.info('nats connected'); - return next; - })(); - try { - return await connecting; - } finally { - connecting = null; - } - } + let sessionManager!: ReturnType; + const ensureConnected = () => sessionManager.ensureConnected(); const s3Prefix = normalizeS3Prefix(process.env.S3_PREFIX); const storageDisabled = async (): Promise => { @@ -517,8 +187,8 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti ErrorResponse: jsonSchema(apiErrorResponseSchema), OperationError: jsonSchema(operationErrorSchema), PdfLayoutProgress: jsonSchema(pdfLayoutProgressSchema), - PublicOperation: jsonSchema(publicOperationSchema), - PublicOperationEvent: jsonSchema(publicOperationEventSchema), + ComputeOperation: jsonSchema(computeOperationSchema), + ComputeOperationEvent: jsonSchema(computeOperationEventSchema), PdfLayoutResolution: jsonSchema(pdfLayoutResolutionSchema), }, }, @@ -571,392 +241,44 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti const operationStateStore = options.routeDeps?.operationStateStore ?? defaultOperationStateStore; const operationEventStream = options.routeDeps?.operationEventStream ?? defaultOperationEventStream; const orchestrator = options.routeDeps?.orchestrator ?? defaultOrchestrator; - let orphanRecoveryPromise: Promise | null = null; - let orphanRecoveryDoneForGeneration = -1; - let sessionGeneration = options.routeDeps ? 0 : -1; + const reconciler = createOperationReconciler({ + stateStore: operationStateStore, + orchestrator, + whisperTimeoutMs, + pdfTimeoutMs, + opStaleMs, + getGeneration: () => options.routeDeps ? 0 : sessionManager.getGeneration(), + logger: app.log, + }); + const ensureOrphanedOpRecovery = () => reconciler.run(); + const getOpState = (opId: string) => reconciler.getOpState(opId); - const runOrphanedOpRecovery = async (options?: { force?: boolean }): Promise => { - if (typeof operationStateStore.listOpStates !== 'function') return; - if (typeof operationStateStore.getOpStateRecord !== 'function') return; - if (typeof orchestrator.markFailedIfUnchanged !== 'function') return; - if (!options?.force && orphanRecoveryDoneForGeneration === sessionGeneration) return; - if (orphanRecoveryPromise) { - await orphanRecoveryPromise; - return; - } - - orphanRecoveryPromise = (async () => { - const recoveredStates = await recoverOrphanedOperations({ - operationStateStore: operationStateStore as typeof operationStateStore & { - getOpStateRecord(opId: string): Promise<{ state: StreamedOperationState; revision: number } | null>; - listOpStates(): Promise; - }, - orchestrator: orchestrator as typeof orchestrator & { - markFailedIfUnchanged(input: { - current: StreamedOperationState; - expectedRevision: number; - error: { message: string; code?: string } | string; - updatedAt?: number; - timing?: WorkerJobTiming; - }): Promise; - }, - whisperTimeoutMs, - pdfTimeoutMs, - opStaleMs, - }); - - if (recoveredStates.length > 0) { - app.log.warn({ - recoveredCount: recoveredStates.length, - ops: recoveredStates.map((state) => ({ - opId: state.opId, - kind: state.kind, - status: state.status, - })), - }, 'recovered stale in-flight operations during reconciliation'); - } - - orphanRecoveryDoneForGeneration = sessionGeneration; - })().finally(() => { - orphanRecoveryPromise = null; - }); - - await orphanRecoveryPromise; - }; - - const ensureOrphanedOpRecovery = async (): Promise => { - await runOrphanedOpRecovery(); - }; - - const getOpState = async (opId: string): Promise => { - await ensureOrphanedOpRecovery(); - return await operationStateStore.getOpState(opId); - }; - - const releaseHttp = (request: FastifyRequest): void => { - const counted = request as FastifyRequest & { [REQUEST_COUNTED_KEY]?: boolean }; - if (!counted[REQUEST_COUNTED_KEY]) return; - counted[REQUEST_COUNTED_KEY] = false; - inFlightHttp = Math.max(0, inFlightHttp - 1); - markActivity('http_completed'); - }; - - app.addHook('onRequest', async (request, reply) => { - const path = requestPath(request); - (request as FastifyRequest & { [REQUEST_STARTED_AT_MS_KEY]?: number })[REQUEST_STARTED_AT_MS_KEY] = Date.now(); - // Count every request as in-flight activity so the idle detector never - // disconnects mid-request. Released in onResponse, or manually after hijack - // for SSE streams (where onResponse does not fire). - (request as FastifyRequest & { [REQUEST_COUNTED_KEY]?: boolean })[REQUEST_COUNTED_KEY] = true; - inFlightHttp += 1; - markActivity(`http_started:${path}`); - if (isHealthPath(path)) return; - if (!isAuthed(request, workerToken)) { - return reply.code(401).send({ error: 'Unauthorized' }); - } - return; + const { releaseHttp } = registerHttpHooks({ + app, + workerToken, + markActivity, + onInFlightHttpChanged: (delta) => { + inFlightHttp = Math.max(0, inFlightHttp + delta); + }, }); - app.addHook('onResponse', async (request, reply) => { - releaseHttp(request); - const path = requestPath(request); - if (isHealthPath(path)) return; - if (reply.statusCode >= 500) { - const startedAt = (request as FastifyRequest & { [REQUEST_STARTED_AT_MS_KEY]?: number })[REQUEST_STARTED_AT_MS_KEY]; - const durationMs = Number.isFinite(startedAt) ? Math.max(0, Date.now() - (startedAt as number)) : -1; - app.log.error({ - reqId: request.id, - method: request.method, - path, - statusCode: reply.statusCode, - durationMs, - traceId: extractTraceId(request) ?? null, - opId: extractOpId(request, path), - }, 'http.error'); - } - }); - - app.get('/health/live', { - schema: { - security: [], - response: { 200: { type: 'object', properties: { ok: { type: 'boolean' } }, required: ['ok'] } }, + registerComputeWorkerRoutes({ + app, + deps: { + orchestrator, + operationStateStore, + operationEventStream, + artifactExists: options.routeDeps?.artifactExists ?? storage.objectExists, }, - }, async () => ({ ok: true })); - - // Reports readiness without forcing a NATS round-trip. Probing NATS here would - // reconnect (and keep) the connection open, defeating idle sleep, so we only - // report the current connection state. The worker reconnects lazily on the next - // operation request regardless of what this returns. - app.get('/health/ready', { - schema: { - security: [], - response: { - 200: { - type: 'object', - properties: { ok: { type: 'boolean' }, natsConnected: { type: 'boolean' } }, - required: ['ok', 'natsConnected'], - }, - }, + s3Prefix, + ensureOrphanedOpRecovery, + getOpState, + getNatsConnected: () => sessionManager.isConnected(), + releaseHttp, + markActivity, + onActiveSseChanged: (delta) => { + activeSse = Math.max(0, activeSse + delta); }, - }, async () => ({ ok: true, natsConnected: session !== null })); - - app.post('/v1/whisper-align/operations', { - schema: { - body: jsonSchema(whisperOperationCreateSchema), - response: { 202: jsonSchema(publicOperationSchema), 400: errorResponseSchema }, - }, - }, async (request, reply) => { - const parsed = whisperOperationCreateSchema.safeParse(request.body); - if (!parsed.success) { - reply.code(400); - return { - error: 'Invalid request body', - issues: parsed.error.issues, - }; - } - - const requestOp: WorkerOperationRequest = { - kind: 'whisper_align', - opKey: buildWhisperOperationKey(parsed.data), - payload: parsed.data, - }; - await ensureOrphanedOpRecovery(); - const op = await orchestrator.enqueueOrReuse(requestOp); - app.log.info({ - kind: requestOp.kind, - opId: op.opId, - jobId: op.jobId, - status: op.status, - opKeyHash: hashOpKey(requestOp.opKey.trim()).slice(0, 16), - }, 'op.accepted'); - reply.code(202); - return toPublicOperation(op); - }); - - app.post('/v1/pdf-layout/operations', { - schema: { - body: jsonSchema(pdfOperationCreateSchema), - response: { 202: jsonSchema(publicOperationSchema), 400: errorResponseSchema }, - }, - }, async (request, reply) => { - const parsed = pdfOperationCreateSchema.safeParse(request.body); - if (!parsed.success) { - reply.code(400); - return { - error: 'Invalid request body', - issues: parsed.error.issues, - }; - } - - const requestOp: WorkerOperationRequest = { - kind: 'pdf_layout', - opKey: buildPdfOperationKey(parsed.data), - payload: { - documentId: parsed.data.documentId, - namespace: parsed.data.namespace, - documentObjectKey: parsed.data.documentObjectKey, - }, - }; - await ensureOrphanedOpRecovery(); - const op = await orchestrator.enqueueOrReuse(requestOp); - reply.code(202); - return toPublicOperation(op); - }); - - app.post('/v1/pdf-layout/resolve', { - schema: { - body: jsonSchema(pdfResolveSchema), - response: { - 200: jsonSchema(pdfLayoutResolutionSchema), - 400: errorResponseSchema, - }, - }, - }, async (request, reply) => { - const parsed = pdfResolveSchema.safeParse(request.body); - if (!parsed.success) { - reply.code(400); - return { error: 'Invalid request body', issues: parsed.error.issues }; - } - await ensureOrphanedOpRecovery(); - const artifactKey = parsedPdfArtifactKey({ - documentId: parsed.data.documentId, - namespace: parsed.data.namespace, - prefix: s3Prefix, - }); - const hasArtifact = await (options.routeDeps?.artifactExists ?? storage.objectExists)(artifactKey); - const opKey = buildPdfOperationKey(parsed.data); - const index = await operationStateStore.getOpIndex?.(opKey); - const operation = index?.opId ? await operationStateStore.getOpState(index.opId) : null; - return { - artifact: hasArtifact ? { objectKey: artifactKey } : null, - operation: operation ? toPublicOperation(operation) : null, - }; - }); - - app.get('/v1/operations/:opId', { - schema: { - params: jsonSchema(operationParamsSchema), - response: { 200: jsonSchema(publicOperationSchema), 400: errorResponseSchema, 404: errorResponseSchema }, - }, - }, async (request, reply) => { - const params = operationParamsSchema.safeParse(request.params); - if (!params.success) { - reply.code(400); - return { error: 'Invalid op id' }; - } - - const state = await getOpState(params.data.opId); - if (!state) { - reply.code(404); - return { error: 'Operation not found' }; - } - - return toPublicOperation(state); - }); - - app.get('/v1/operations/:opId/events', { - schema: { - params: jsonSchema(operationParamsSchema), - querystring: jsonSchema(operationEventsQuerySchema), - response: { - 200: { type: 'string', description: 'Server-sent PublicOperationEvent stream' }, - 400: errorResponseSchema, - 404: errorResponseSchema, - }, - }, - }, async (request, reply) => { - const params = operationParamsSchema.safeParse(request.params); - if (!params.success) { - reply.code(400); - return { error: 'Invalid op id' }; - } - - const initial = await getOpState(params.data.opId); - if (!initial) { - reply.code(404); - return { error: 'Operation not found' }; - } - - const cursorQueryRaw = request.query as { sinceEventId?: string | number | null } | undefined; - const cursorFromQuery = Number(cursorQueryRaw?.sinceEventId ?? 0); - const lastEventIdHeader = request.headers['last-event-id']; - const cursorFromHeader = Number( - Array.isArray(lastEventIdHeader) ? (lastEventIdHeader[0] ?? 0) : (lastEventIdHeader ?? 0), - ); - const sinceEventId = Math.max( - 0, - Number.isFinite(cursorFromQuery) ? Math.floor(cursorFromQuery) : 0, - Number.isFinite(cursorFromHeader) ? Math.floor(cursorFromHeader) : 0, - ); - - reply.hijack(); - // onResponse will not fire for a hijacked reply, so release the HTTP in-flight - // count here and track the long-lived stream via activeSse instead. - releaseHttp(request); - activeSse += 1; - markActivity('sse_started'); - 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'); - // Tell the browser EventSource to back off before reconnecting. If this stream - // is torn down because the worker went idle (NATS dropped), we don't want the - // client to reconnect immediately and re-wake the container. - reply.raw.write(encodeSseFrame({ retry: OP_EVENTS_RECONNECT_HINT_MS })); - - let closed = false; - let unsubscribe: (() => void) | null = null; - let keepalive: NodeJS.Timeout | null = null; - - const writeSnapshot = (snapshot: StreamedOperationState, eventId: number): void => { - if (closed || reply.raw.writableEnded) return; - const frameEvent: PublicOperationEvent = { - eventId, - snapshot: toPublicOperation(snapshot), - }; - reply.raw.write(encodeSseFrame({ - id: eventId, - event: 'snapshot', - data: frameEvent, - })); - }; - - const closeStream = (): void => { - if (closed) return; - closed = true; - if (unsubscribe) { - unsubscribe(); - unsubscribe = null; - } - if (keepalive) { - clearInterval(keepalive); - keepalive = null; - } - activeSse = Math.max(0, activeSse - 1); - markActivity('sse_closed'); - if (!reply.raw.writableEnded) { - reply.raw.end(); - } - }; - - request.raw.on('close', () => { - closeStream(); - }); - - try { - let current = initial; - let signature = JSON.stringify(current); - writeSnapshot(current, sinceEventId > 0 ? sinceEventId : 0); - if (isTerminalStatus(current.status)) { - return reply; - } - - keepalive = setInterval(() => { - if (closed || reply.raw.writableEnded) return; - reply.raw.write(': keepalive\n\n'); - }, OP_EVENTS_KEEPALIVE_MS); - - unsubscribe = await operationEventStream.subscribe({ - opId: params.data.opId, - sinceEventId, - onEvent: (event) => { - if (closed) return; - if (event.snapshot.opId !== params.data.opId) return; - const nextSignature = JSON.stringify(event.snapshot); - if (nextSignature !== signature) { - current = event.snapshot; - signature = nextSignature; - // A real event is progress: refresh the idle window so an actively - // streaming op keeps the worker awake. A silent stream does not. - markActivity('sse_event'); - writeSnapshot(current, event.eventId); - } - if (isTerminalStatus(event.snapshot.status)) { - closeStream(); - } - }, - onError: (error) => { - app.log.warn({ - opId: params.data.opId, - error: toErrorMessage(error), - }, 'op events stream loop error'); - closeStream(); - }, - }); - - await new Promise((resolve) => { - request.raw.once('close', () => resolve()); - }); - } catch (error) { - app.log.warn({ - opId: params.data.opId, - error: toErrorMessage(error), - }, 'op events stream loop error'); - } finally { - closeStream(); - } - - return reply; }); const jobHandlers = createJobHandlers({ @@ -974,7 +296,7 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti pdfAttempts, whisperCodec: whisperJobCodec, pdfCodec: layoutJobCodec, - isOwnerActive: (owner) => session === owner, + isOwnerActive: (owner) => sessionManager.isOwnerActive(owner), isStopping: () => stopping, markActivity, onInFlightJobsChanged: (delta) => { @@ -982,32 +304,41 @@ export async function createComputeWorkerApp(options: CreateComputeWorkerAppOpti }, }); + sessionManager = createNatsSessionManager({ + connectOptions: connectOpts, + logger: app.log, + whisperTimeoutMs, + pdfTimeoutMs, + pdfAttempts, + jobsStreamMaxBytes, + eventsStreamMaxBytes, + jobStatesMaxBytes, + natsReplicas, + isStopping: () => stopping, + getActivity: () => ({ + activeSse, + inFlightHttp, + inFlightJobs, + lastActivityAt, + lastActivityReason, + }), + markActivity, + startWorkers: (session) => { + if (disableWorkers) return; + workerLoops.start(session, { + whisper: session.whisperConsumer, + pdfLayout: session.layoutConsumer, + }); + }, + stopWorkers: () => workerLoops.stop(), + runReconciliation: () => reconciler.run({ force: true }), + }); + const close = async (): Promise => { if (stopping) return; stopping = true; - if (idleTimer) { - clearInterval(idleTimer); - idleTimer = null; - } - if (orphanSweepTimer) { - clearInterval(orphanSweepTimer); - orphanSweepTimer = null; - } await app.close(); - await workerLoops.stop(); - const current = session; - session = null; - if (current) { - try { - await current.nc.drain(); - } catch { - try { - await current.nc.close(); - } catch { - // ignore close errors - } - } - } + await sessionManager.close(); }; const registerSignalHandlers = (): void => { diff --git a/compute-worker/src/api/public-operation.ts b/compute-worker/src/api/compute-operation.ts similarity index 72% rename from compute-worker/src/api/public-operation.ts rename to compute-worker/src/api/compute-operation.ts index ac9b43b..03c6360 100644 --- a/compute-worker/src/api/public-operation.ts +++ b/compute-worker/src/api/compute-operation.ts @@ -1,13 +1,13 @@ -import type { WorkerOperationState } from '../api/contracts'; -import { pdfSubjectFromOperationKey } from './operation-keys'; +import type { WorkerOperationState } from '../operations/contracts'; +import { pdfSubjectFromOperationKey } from '../operations/keys'; -export type PublicOperationSubject = +export type ComputeOperationSubject = | { kind: 'whisper_align' } | { kind: 'pdf_layout'; documentId: string; namespace: string | null }; -export interface PublicOperation { +export interface ComputeOperation { opId: string; - subject: PublicOperationSubject; + subject: ComputeOperationSubject; status: WorkerOperationState['status']; queuedAt: number; updatedAt: number; @@ -18,14 +18,14 @@ export interface PublicOperation { progress?: WorkerOperationState['progress']; } -export interface PublicOperationEvent { +export interface ComputeOperationEvent { eventId: number; - snapshot: PublicOperation; + snapshot: ComputeOperation; } -export function toPublicOperation( +export function toComputeOperation( state: WorkerOperationState, -): PublicOperation { +): ComputeOperation { const subject = state.kind === 'pdf_layout' ? (pdfSubjectFromOperationKey(state.opKey) ?? { kind: 'pdf_layout', documentId: '', namespace: null }) : { kind: 'whisper_align' as const }; diff --git a/compute-worker/src/api/http-hooks.ts b/compute-worker/src/api/http-hooks.ts new file mode 100644 index 0000000..d633da3 --- /dev/null +++ b/compute-worker/src/api/http-hooks.ts @@ -0,0 +1,80 @@ +import type { FastifyInstance, FastifyRequest } from 'fastify'; + +const REQUEST_STARTED_AT_MS_KEY = Symbol('request-started-at-ms'); +const REQUEST_COUNTED_KEY = Symbol('request-activity-counted'); + +function requestPath(request: FastifyRequest): string { + return request.url.split('?')[0] ?? request.url; +} + +function isHealthPath(path: string): boolean { + return path === '/health/live' || path === '/health/ready'; +} + +function isAuthed(request: FastifyRequest, expectedToken: string): boolean { + const auth = request.headers.authorization; + return auth?.startsWith('Bearer ') === true + && auth.slice('Bearer '.length).trim() === expectedToken; +} + +function extractTraceId(request: FastifyRequest): string | null { + const header = request.headers['x-openreader-trace-id']; + if (Array.isArray(header)) return header[0] ?? null; + return typeof header === 'string' ? header : null; +} + +function extractOpId(request: FastifyRequest, path: string): string | null { + const params = request.params as { opId?: unknown } | undefined; + if (typeof params?.opId === 'string' && params.opId.trim()) return params.opId.trim(); + const match = path.match(/^\/v1\/operations\/([^/]+)/); + if (!match?.[1]) return null; + try { + return decodeURIComponent(match[1]); + } catch { + return match[1]; + } +} + +export function registerHttpHooks(input: { + app: FastifyInstance; + workerToken: string; + markActivity: (reason: string) => void; + onInFlightHttpChanged: (delta: number) => void; +}) { + const releaseHttp = (request: FastifyRequest): void => { + const counted = request as FastifyRequest & { [REQUEST_COUNTED_KEY]?: boolean }; + if (!counted[REQUEST_COUNTED_KEY]) return; + counted[REQUEST_COUNTED_KEY] = false; + input.onInFlightHttpChanged(-1); + input.markActivity('http_completed'); + }; + + input.app.addHook('onRequest', async (request, reply) => { + const path = requestPath(request); + (request as FastifyRequest & { [REQUEST_STARTED_AT_MS_KEY]?: number })[REQUEST_STARTED_AT_MS_KEY] = Date.now(); + (request as FastifyRequest & { [REQUEST_COUNTED_KEY]?: boolean })[REQUEST_COUNTED_KEY] = true; + input.onInFlightHttpChanged(1); + input.markActivity(`http_started:${path}`); + if (!isHealthPath(path) && !isAuthed(request, input.workerToken)) { + return reply.code(401).send({ error: 'Unauthorized' }); + } + }); + + input.app.addHook('onResponse', async (request, reply) => { + releaseHttp(request); + const path = requestPath(request); + if (isHealthPath(path) || reply.statusCode < 500) return; + const startedAt = (request as FastifyRequest & { [REQUEST_STARTED_AT_MS_KEY]?: number })[REQUEST_STARTED_AT_MS_KEY]; + input.app.log.error({ + reqId: request.id, + method: request.method, + path, + statusCode: reply.statusCode, + durationMs: Number.isFinite(startedAt) ? Math.max(0, Date.now() - (startedAt as number)) : -1, + traceId: extractTraceId(request), + opId: extractOpId(request, path), + }, 'http.error'); + }); + + return { releaseHttp }; +} diff --git a/compute-worker/src/api/routes.ts b/compute-worker/src/api/routes.ts new file mode 100644 index 0000000..481e965 --- /dev/null +++ b/compute-worker/src/api/routes.ts @@ -0,0 +1,356 @@ +import type { FastifyInstance, FastifyRequest } from 'fastify'; +import { encodeSseFrame } from '../operations'; +import type { + PdfLayoutJobResult, + WorkerJobTiming, + WorkerOperationEvent, + WorkerOperationRequest, + WorkerOperationState, + WhisperAlignJobResult, +} from '../operations/contracts'; +import { hashOpKey } from '../infrastructure/nats-adapters'; +import type { StreamedOperationState } from '../operations/recovery'; +import type { ReconciliationStateStore } from '../operations/reconciliation'; +import { parsedPdfArtifactKey } from '../storage/artifact-addressing'; +import { buildPdfOperationKey, buildWhisperOperationKey } from '../operations/keys'; +import { toComputeOperation, type ComputeOperationEvent } from './compute-operation'; +import { + apiErrorResponseSchema, + jsonSchema, + operationEventsQuerySchema, + operationParamsSchema, + pdfLayoutResolutionSchema, + pdfOperationCreateSchema, + pdfResolveSchema, + computeOperationSchema, + whisperOperationCreateSchema, +} from './schemas'; + +const OP_EVENTS_KEEPALIVE_MS = 15_000; +const OP_EVENTS_RECONNECT_HINT_MS = 120_000; +const errorResponseSchema = jsonSchema(apiErrorResponseSchema); + +interface OperationEventStreamLike { + subscribe(input: { + opId: string; + sinceEventId?: number; + onEvent: (event: WorkerOperationEvent) => void | Promise; + onError?: (error: unknown) => void; + }): Promise<() => void>; +} + +interface OperationStateStoreLike extends ReconciliationStateStore {} + +interface OrchestratorLike { + enqueueOrReuse(request: WorkerOperationRequest): Promise; + markRunning(input: { + opId: string; + startedAt?: number; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; + markProgress(input: { + opId: string; + progress: WorkerOperationState['progress']; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; + markSucceeded(input: { + opId: string; + result: unknown; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; + markFailed(input: { + opId: string; + error: { message: string; code?: string } | string; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; + markFailedIfUnchanged?(input: { + current: StreamedOperationState; + expectedRevision: number; + error: { message: string; code?: string } | string; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; +} + +export interface ComputeWorkerRouteDeps { + orchestrator: OrchestratorLike; + operationStateStore: OperationStateStoreLike; + operationEventStream: OperationEventStreamLike; + artifactExists?: (key: string) => Promise; +} + +function toErrorMessage(error: unknown): string { + if (error instanceof Error && error.message) return error.message; + return String(error); +} + +function isTerminalStatus(status: import('../operations/contracts').WorkerJobState): boolean { + return status === 'succeeded' || status === 'failed'; +} + +export function registerComputeWorkerRoutes(input: { + app: FastifyInstance; + deps: ComputeWorkerRouteDeps; + s3Prefix: string; + ensureOrphanedOpRecovery: () => Promise; + getOpState: (opId: string) => Promise; + getNatsConnected: () => boolean; + releaseHttp: (request: FastifyRequest) => void; + markActivity: (reason: string) => void; + onActiveSseChanged: (delta: number) => void; +}) { + const { + app, + deps, + s3Prefix, + ensureOrphanedOpRecovery, + getOpState, + getNatsConnected, + releaseHttp, + markActivity, + onActiveSseChanged, + } = input; + + app.get('/health/live', { + schema: { + security: [], + response: { 200: { type: 'object', properties: { ok: { type: 'boolean' } }, required: ['ok'] } }, + }, + }, async () => ({ ok: true })); + + app.get('/health/ready', { + schema: { + security: [], + response: { + 200: { + type: 'object', + properties: { ok: { type: 'boolean' }, natsConnected: { type: 'boolean' } }, + required: ['ok', 'natsConnected'], + }, + }, + }, + }, async () => ({ ok: true, natsConnected: getNatsConnected() })); + + app.post('/v1/whisper-align/operations', { + schema: { + body: jsonSchema(whisperOperationCreateSchema), + response: { 202: jsonSchema(computeOperationSchema), 400: errorResponseSchema }, + }, + }, async (request, reply) => { + const parsed = whisperOperationCreateSchema.safeParse(request.body); + if (!parsed.success) { + reply.code(400); + return { error: 'Invalid request body', issues: parsed.error.issues }; + } + + const requestOp: WorkerOperationRequest = { + kind: 'whisper_align', + opKey: buildWhisperOperationKey(parsed.data), + payload: parsed.data, + }; + await ensureOrphanedOpRecovery(); + const op = await deps.orchestrator.enqueueOrReuse(requestOp); + app.log.info({ + kind: requestOp.kind, + opId: op.opId, + jobId: op.jobId, + status: op.status, + opKeyHash: hashOpKey(requestOp.opKey.trim()).slice(0, 16), + }, 'op.accepted'); + reply.code(202); + return toComputeOperation(op); + }); + + app.post('/v1/pdf-layout/operations', { + schema: { + body: jsonSchema(pdfOperationCreateSchema), + response: { 202: jsonSchema(computeOperationSchema), 400: errorResponseSchema }, + }, + }, async (request, reply) => { + const parsed = pdfOperationCreateSchema.safeParse(request.body); + if (!parsed.success) { + reply.code(400); + return { error: 'Invalid request body', issues: parsed.error.issues }; + } + + const requestOp: WorkerOperationRequest = { + kind: 'pdf_layout', + opKey: buildPdfOperationKey(parsed.data), + payload: { + documentId: parsed.data.documentId, + namespace: parsed.data.namespace, + documentObjectKey: parsed.data.documentObjectKey, + }, + }; + await ensureOrphanedOpRecovery(); + const op = await deps.orchestrator.enqueueOrReuse(requestOp); + reply.code(202); + return toComputeOperation(op); + }); + + app.post('/v1/pdf-layout/resolve', { + schema: { + body: jsonSchema(pdfResolveSchema), + response: { 200: jsonSchema(pdfLayoutResolutionSchema), 400: errorResponseSchema }, + }, + }, async (request, reply) => { + const parsed = pdfResolveSchema.safeParse(request.body); + if (!parsed.success) { + reply.code(400); + return { error: 'Invalid request body', issues: parsed.error.issues }; + } + await ensureOrphanedOpRecovery(); + const artifactKey = parsedPdfArtifactKey({ + documentId: parsed.data.documentId, + namespace: parsed.data.namespace, + prefix: s3Prefix, + }); + const hasArtifact = await deps.artifactExists?.(artifactKey) ?? false; + const opKey = buildPdfOperationKey(parsed.data); + const index = await deps.operationStateStore.getOpIndex?.(opKey); + const operation = index?.opId ? await deps.operationStateStore.getOpState(index.opId) : null; + return { + artifact: hasArtifact ? { objectKey: artifactKey } : null, + operation: operation ? toComputeOperation(operation) : null, + }; + }); + + app.get('/v1/operations/:opId', { + schema: { + params: jsonSchema(operationParamsSchema), + response: { 200: jsonSchema(computeOperationSchema), 400: errorResponseSchema, 404: errorResponseSchema }, + }, + }, async (request, reply) => { + const params = operationParamsSchema.safeParse(request.params); + if (!params.success) { + reply.code(400); + return { error: 'Invalid op id' }; + } + + const state = await getOpState(params.data.opId); + if (!state) { + reply.code(404); + return { error: 'Operation not found' }; + } + return toComputeOperation(state); + }); + + app.get('/v1/operations/:opId/events', { + schema: { + params: jsonSchema(operationParamsSchema), + querystring: jsonSchema(operationEventsQuerySchema), + response: { + 200: { type: 'string', description: 'Server-sent ComputeOperationEvent stream' }, + 400: errorResponseSchema, + 404: errorResponseSchema, + }, + }, + }, async (request, reply) => { + const params = operationParamsSchema.safeParse(request.params); + if (!params.success) { + reply.code(400); + return { error: 'Invalid op id' }; + } + + const initial = await getOpState(params.data.opId); + if (!initial) { + reply.code(404); + return { error: 'Operation not found' }; + } + + const cursorQueryRaw = request.query as { sinceEventId?: string | number | null } | undefined; + const cursorFromQuery = Number(cursorQueryRaw?.sinceEventId ?? 0); + const lastEventIdHeader = request.headers['last-event-id']; + const cursorFromHeader = Number( + Array.isArray(lastEventIdHeader) ? (lastEventIdHeader[0] ?? 0) : (lastEventIdHeader ?? 0), + ); + const sinceEventId = Math.max( + 0, + Number.isFinite(cursorFromQuery) ? Math.floor(cursorFromQuery) : 0, + Number.isFinite(cursorFromHeader) ? Math.floor(cursorFromHeader) : 0, + ); + + reply.hijack(); + releaseHttp(request); + onActiveSseChanged(1); + markActivity('sse_started'); + 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'); + reply.raw.write(encodeSseFrame({ retry: OP_EVENTS_RECONNECT_HINT_MS })); + + let closed = false; + let unsubscribe: (() => void) | null = null; + let keepalive: NodeJS.Timeout | null = null; + + const writeSnapshot = (snapshot: StreamedOperationState, eventId: number): void => { + if (closed || reply.raw.writableEnded) return; + const frameEvent: ComputeOperationEvent = { + eventId, + snapshot: toComputeOperation(snapshot), + }; + reply.raw.write(encodeSseFrame({ id: eventId, event: 'snapshot', data: frameEvent })); + }; + + const closeStream = (): void => { + if (closed) return; + closed = true; + if (unsubscribe) { + unsubscribe(); + unsubscribe = null; + } + if (keepalive) { + clearInterval(keepalive); + keepalive = null; + } + onActiveSseChanged(-1); + markActivity('sse_closed'); + if (!reply.raw.writableEnded) reply.raw.end(); + }; + + request.raw.on('close', closeStream); + + try { + let current = initial; + let signature = JSON.stringify(current); + writeSnapshot(current, sinceEventId > 0 ? sinceEventId : 0); + if (isTerminalStatus(current.status)) return reply; + + keepalive = setInterval(() => { + if (!closed && !reply.raw.writableEnded) reply.raw.write(': keepalive\n\n'); + }, OP_EVENTS_KEEPALIVE_MS); + + unsubscribe = await deps.operationEventStream.subscribe({ + opId: params.data.opId, + sinceEventId, + onEvent: (event) => { + if (closed || event.snapshot.opId !== params.data.opId) return; + const nextSignature = JSON.stringify(event.snapshot); + if (nextSignature !== signature) { + current = event.snapshot; + signature = nextSignature; + markActivity('sse_event'); + writeSnapshot(current, event.eventId); + } + if (isTerminalStatus(event.snapshot.status)) closeStream(); + }, + onError: (error) => { + app.log.warn({ opId: params.data.opId, error: toErrorMessage(error) }, 'op events stream loop error'); + closeStream(); + }, + }); + + await new Promise((resolve) => request.raw.once('close', resolve)); + } catch (error) { + app.log.warn({ opId: params.data.opId, error: toErrorMessage(error) }, 'op events stream loop error'); + } finally { + closeStream(); + } + return reply; + }); +} diff --git a/compute-worker/src/api/schemas.ts b/compute-worker/src/api/schemas.ts index 66df5c5..0175d86 100644 --- a/compute-worker/src/api/schemas.ts +++ b/compute-worker/src/api/schemas.ts @@ -95,7 +95,7 @@ export const pdfLayoutProgressSchema = z.object({ phase: z.enum(['infer', 'merge']), }); -export const publicOperationSchema = z.object({ +export const computeOperationSchema = z.object({ opId: z.string(), subject: z.discriminatedUnion('kind', [ z.object({ kind: z.literal('whisper_align') }), @@ -119,14 +119,14 @@ export const publicOperationSchema = z.object({ progress: pdfLayoutProgressSchema.optional(), }); -export const publicOperationEventSchema = z.object({ +export const computeOperationEventSchema = z.object({ eventId: z.number(), - snapshot: publicOperationSchema, + snapshot: computeOperationSchema, }); export const pdfLayoutResolutionSchema = z.object({ artifact: artifactReferenceSchema.nullable(), - operation: publicOperationSchema.nullable(), + operation: computeOperationSchema.nullable(), }); export function jsonSchema(schema: z.ZodType): Record { diff --git a/compute-worker/src/inference/pdf/parse.ts b/compute-worker/src/inference/pdf/parse.ts index b4491b2..4f70f53 100644 --- a/compute-worker/src/inference/pdf/parse.ts +++ b/compute-worker/src/inference/pdf/parse.ts @@ -4,7 +4,7 @@ import { ensureModel } from './model'; import { runLayoutModel } from './layout-model'; import { mergeTextWithRegions } from './document-layout'; import { configurePdfjsNodeRuntime, resolvePdfjsStandardFontDataUrl } from './pdfjs'; -import { PDF_PARSER_VERSION } from '../../api/contracts'; +import { PDF_PARSER_VERSION } from '../../operations/contracts'; import { stitchCrossPageBlocks } from './stitch'; import { renderPage } from './render'; import { normalizeTextItemsForLayout } from './normalize-text'; diff --git a/compute-worker/src/inference/whisper/align.ts b/compute-worker/src/inference/whisper/align.ts index 073afc4..a271623 100644 --- a/compute-worker/src/inference/whisper/align.ts +++ b/compute-worker/src/inference/whisper/align.ts @@ -29,6 +29,11 @@ import { WHISPER_DECODER_MERGED_MODEL_PATH, WHISPER_DECODER_WITH_PAST_MODEL_PATH, } from './model'; +import { + applyTokenSuppression, + applyWhisperTimestampLogitsRules, + argmax, +} from './decoder'; export function buildGoertzelCoefficients(freqBins: number, fftSize: number): Float64Array { const coeffs = new Float64Array(freqBins); @@ -448,118 +453,6 @@ function buildEmptyPastFeeds() { return state.emptyPastFeedsTemplate; } -function argmax(values: Float32Array): number | null { - let bestIdx = 0; - let bestScore = Number.NEGATIVE_INFINITY; - - for (let i = 0; i < values.length; i += 1) { - const score = values[i]; - if (score > bestScore) { - bestScore = score; - bestIdx = i; - } - } - - return Number.isFinite(bestScore) ? bestIdx : null; -} - -function applyTokenSuppression(logits: Float32Array, tokens: Set) { - for (const tokenId of tokens) { - if (tokenId >= 0 && tokenId < logits.length) { - logits[tokenId] = Number.NEGATIVE_INFINITY; - } - } -} - -function logSoftmax(input: Float32Array): Float32Array { - let max = Number.NEGATIVE_INFINITY; - for (let i = 0; i < input.length; i += 1) { - if (input[i] > max) max = input[i]; - } - if (!Number.isFinite(max)) { - return new Float32Array(input.length).fill(Number.NEGATIVE_INFINITY); - } - - let sum = 0; - for (let i = 0; i < input.length; i += 1) { - sum += Math.exp(input[i] - max); - } - const logSum = Math.log(sum); - - const out = new Float32Array(input.length); - for (let i = 0; i < input.length; i += 1) { - out[i] = input[i] - max - logSum; - } - return out; -} - -function applyWhisperTimestampLogitsRules(input: { - logits: Float32Array; - generated: number[]; - beginIndex: number; - eosTokenId: number; - noTimestampsTokenId: number; - timestampBeginTokenId: number; - maxInitialTimestampIndex: number; -}) { - const { - logits, - generated, - beginIndex, - eosTokenId, - noTimestampsTokenId, - timestampBeginTokenId, - maxInitialTimestampIndex, - } = input; - - if (noTimestampsTokenId >= 0 && noTimestampsTokenId < logits.length) { - logits[noTimestampsTokenId] = Number.NEGATIVE_INFINITY; - } - - if (generated.length === beginIndex) { - const upper = Math.min(timestampBeginTokenId, logits.length); - for (let i = 0; i < upper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } - - const seq = generated.slice(beginIndex); - const lastWasTimestamp = seq.length >= 1 && seq[seq.length - 1] >= timestampBeginTokenId; - const penultimateWasTimestamp = seq.length < 2 || seq[seq.length - 2] >= timestampBeginTokenId; - - if (lastWasTimestamp) { - if (penultimateWasTimestamp) { - for (let i = timestampBeginTokenId; i < logits.length; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } else { - const upper = Math.min(eosTokenId, logits.length); - for (let i = 0; i < upper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } - } - - if (generated.length === beginIndex && Number.isFinite(maxInitialTimestampIndex)) { - const lastAllowed = timestampBeginTokenId + maxInitialTimestampIndex; - for (let i = lastAllowed + 1; i < logits.length; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } - - const textUpper = Math.min(timestampBeginTokenId, logits.length); - if (textUpper <= 0 || textUpper >= logits.length) return; - - const logprobs = logSoftmax(logits); - - let maxTextTokenLogprob = Number.NEGATIVE_INFINITY; - for (let i = 0; i < textUpper; i += 1) { - if (logprobs[i] > maxTextTokenLogprob) maxTextTokenLogprob = logprobs[i]; - } - - let timestampProbMass = 0; - for (let i = textUpper; i < logprobs.length; i += 1) { - timestampProbMass += Math.exp(logprobs[i]); - } - const timestampLogprob = timestampProbMass > 0 ? Math.log(timestampProbMass) : Number.NEGATIVE_INFINITY; - - if (timestampLogprob > maxTextTokenLogprob) { - for (let i = 0; i < textUpper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; - } -} - async function getRuntime(): Promise { if (state.runtimePromise) return state.runtimePromise; diff --git a/compute-worker/src/inference/whisper/decoder.ts b/compute-worker/src/inference/whisper/decoder.ts new file mode 100644 index 0000000..777047e --- /dev/null +++ b/compute-worker/src/inference/whisper/decoder.ts @@ -0,0 +1,111 @@ +export function argmax(values: Float32Array): number | null { + let bestIdx = 0; + let bestScore = Number.NEGATIVE_INFINITY; + + for (let i = 0; i < values.length; i += 1) { + const score = values[i]; + if (score > bestScore) { + bestScore = score; + bestIdx = i; + } + } + + return Number.isFinite(bestScore) ? bestIdx : null; +} + +export function applyTokenSuppression(logits: Float32Array, tokens: Set) { + for (const tokenId of tokens) { + if (tokenId >= 0 && tokenId < logits.length) { + logits[tokenId] = Number.NEGATIVE_INFINITY; + } + } +} + +function logSoftmax(input: Float32Array): Float32Array { + let max = Number.NEGATIVE_INFINITY; + for (let i = 0; i < input.length; i += 1) { + if (input[i] > max) max = input[i]; + } + if (!Number.isFinite(max)) { + return new Float32Array(input.length).fill(Number.NEGATIVE_INFINITY); + } + + let sum = 0; + for (let i = 0; i < input.length; i += 1) { + sum += Math.exp(input[i] - max); + } + const logSum = Math.log(sum); + + const out = new Float32Array(input.length); + for (let i = 0; i < input.length; i += 1) { + out[i] = input[i] - max - logSum; + } + return out; +} + +export function applyWhisperTimestampLogitsRules(input: { + logits: Float32Array; + generated: number[]; + beginIndex: number; + eosTokenId: number; + noTimestampsTokenId: number; + timestampBeginTokenId: number; + maxInitialTimestampIndex: number; +}) { + const { + logits, + generated, + beginIndex, + eosTokenId, + noTimestampsTokenId, + timestampBeginTokenId, + maxInitialTimestampIndex, + } = input; + + if (noTimestampsTokenId >= 0 && noTimestampsTokenId < logits.length) { + logits[noTimestampsTokenId] = Number.NEGATIVE_INFINITY; + } + + if (generated.length === beginIndex) { + const upper = Math.min(timestampBeginTokenId, logits.length); + for (let i = 0; i < upper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } + + const seq = generated.slice(beginIndex); + const lastWasTimestamp = seq.length >= 1 && seq[seq.length - 1] >= timestampBeginTokenId; + const penultimateWasTimestamp = seq.length < 2 || seq[seq.length - 2] >= timestampBeginTokenId; + + if (lastWasTimestamp) { + if (penultimateWasTimestamp) { + for (let i = timestampBeginTokenId; i < logits.length; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } else { + const upper = Math.min(eosTokenId, logits.length); + for (let i = 0; i < upper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } + } + + if (generated.length === beginIndex && Number.isFinite(maxInitialTimestampIndex)) { + const lastAllowed = timestampBeginTokenId + maxInitialTimestampIndex; + for (let i = lastAllowed + 1; i < logits.length; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } + + const textUpper = Math.min(timestampBeginTokenId, logits.length); + if (textUpper <= 0 || textUpper >= logits.length) return; + + const logprobs = logSoftmax(logits); + + let maxTextTokenLogprob = Number.NEGATIVE_INFINITY; + for (let i = 0; i < textUpper; i += 1) { + if (logprobs[i] > maxTextTokenLogprob) maxTextTokenLogprob = logprobs[i]; + } + + let timestampProbMass = 0; + for (let i = textUpper; i < logprobs.length; i += 1) { + timestampProbMass += Math.exp(logprobs[i]); + } + const timestampLogprob = timestampProbMass > 0 ? Math.log(timestampProbMass) : Number.NEGATIVE_INFINITY; + + if (timestampLogprob > maxTextTokenLogprob) { + for (let i = 0; i < textUpper; i += 1) logits[i] = Number.NEGATIVE_INFINITY; + } +} diff --git a/compute-worker/src/infrastructure/config.ts b/compute-worker/src/infrastructure/config.ts index 2547f3b..791ac4e 100644 --- a/compute-worker/src/infrastructure/config.ts +++ b/compute-worker/src/infrastructure/config.ts @@ -19,7 +19,13 @@ export type IdleTimeoutAndHardCapInput = { label: string; }; -function readPositiveIntEnv(name: string, fallback: number): number { +export function requireEnv(name: string): string { + const value = process.env[name]?.trim(); + if (!value) throw new Error(`${name} is required`); + return value; +} + +export function readPositiveIntEnv(name: string, fallback: number): number { const raw = process.env[name]?.trim(); if (!raw) return fallback; const parsed = Number(raw); @@ -27,6 +33,34 @@ function readPositiveIntEnv(name: string, fallback: number): number { return Math.floor(parsed); } +export function readBoolEnv(name: string, fallback: boolean): boolean { + const raw = process.env[name]?.trim(); + if (!raw) return fallback; + return ['1', 'true', 'yes', 'on'].includes(raw.toLowerCase()); +} + +export function normalizeNatsReplicas(value: number): number { + return value === 3 || value === 5 ? value : 1; +} + +export function buildLoggerConfig(): boolean | Record { + const format = process.env.LOG_FORMAT?.trim().toLowerCase() || 'pretty'; + const level = process.env.COMPUTE_LOG_LEVEL?.trim() || 'info'; + if (format === 'json') return { level, base: null }; + return { + level, + base: null, + transport: { + target: 'pino-pretty', + options: { + colorize: true, + translateTime: 'SYS:standard', + ignore: 'pid,hostname', + }, + }, + }; +} + let timeoutConfigCache: ComputeTimeoutConfig | null = null; let opStaleMsCache: number | null = null; diff --git a/compute-worker/src/infrastructure/nats-adapters.ts b/compute-worker/src/infrastructure/nats-adapters.ts index 00d2526..d6e4ddb 100644 --- a/compute-worker/src/infrastructure/nats-adapters.ts +++ b/compute-worker/src/infrastructure/nats-adapters.ts @@ -13,7 +13,7 @@ import type { PdfLayoutJobRequest, WhisperAlignJobRequest, WorkerOperationKind, -} from '../api/contracts'; +} from '../operations/contracts'; import { createJsonCodec } from './json-codec'; export interface KvEntryLike { diff --git a/compute-worker/src/infrastructure/nats-session.ts b/compute-worker/src/infrastructure/nats-session.ts new file mode 100644 index 0000000..356a0de --- /dev/null +++ b/compute-worker/src/infrastructure/nats-session.ts @@ -0,0 +1,204 @@ +import type { Consumer, JetStreamClient, JetStreamManager } from '@nats-io/jetstream'; +import { jetstream, jetstreamManager } from '@nats-io/jetstream'; +import { Kvm } from '@nats-io/kv'; +import { connect, type ConnectionOptions, type NatsConnection } from '@nats-io/transport-node'; +import { + COMPUTE_STATE_BUCKET, + COMPUTE_STATE_TTL_MS, + JOBS_STREAM_NAME, + LAYOUT_CONSUMER_NAME, + NATS_API_TIMEOUT_MS, + WHISPER_CONSUMER_NAME, + ensureJetStreamResources, +} from './nats'; + +const IDLE_DISCONNECT_MS = 120_000; +const IDLE_CHECK_INTERVAL_MS = 5_000; +const IDLE_STATUS_LOG_INTERVAL_MS = 60_000; +const ORPHAN_SWEEP_INTERVAL_MS = 15_000; + +export interface NatsSession { + nc: NatsConnection; + js: JetStreamClient; + jsm: JetStreamManager; + kv: Awaited>; + whisperConsumer: Consumer; + layoutConsumer: Consumer; +} + +interface NatsSessionLogger { + info(data: unknown, message?: string): void; + info(message: string): void; + error(data: unknown, message?: string): void; +} + +export interface NatsActivitySnapshot { + activeSse: number; + inFlightHttp: number; + inFlightJobs: number; + lastActivityAt: number; + lastActivityReason: string; +} + +function toErrorMessage(error: unknown): string { + return error instanceof Error && error.message ? error.message : String(error); +} + +export function createNatsSessionManager(input: { + connectOptions: ConnectionOptions; + logger: NatsSessionLogger; + whisperTimeoutMs: number; + pdfTimeoutMs: number; + pdfAttempts: number; + jobsStreamMaxBytes: number; + eventsStreamMaxBytes: number; + jobStatesMaxBytes: number; + natsReplicas: number; + isStopping: () => boolean; + getActivity: () => NatsActivitySnapshot; + markActivity: (reason: string) => void; + startWorkers: (session: NatsSession) => void; + stopWorkers: () => Promise; + runReconciliation: () => Promise; +}) { + let session: NatsSession | null = null; + let connecting: Promise | null = null; + let idleTimer: NodeJS.Timeout | null = null; + let orphanSweepTimer: NodeJS.Timeout | null = null; + let lastIdleStatusLogAt = 0; + let generation = -1; + + const clearTimers = () => { + if (idleTimer) clearInterval(idleTimer); + if (orphanSweepTimer) clearInterval(orphanSweepTimer); + idleTimer = null; + orphanSweepTimer = null; + }; + + const disconnect = async (reason: string): Promise => { + const current = session; + if (!current) return; + const activity = input.getActivity(); + input.logger.info({ + reason, + activeSse: activity.activeSse, + inFlightHttp: activity.inFlightHttp, + inFlightJobs: activity.inFlightJobs, + idleForMs: Date.now() - activity.lastActivityAt, + }, 'nats dropping connection'); + session = null; + clearTimers(); + try { + await current.nc.close(); + } catch { + // Closing an already-closed session is harmless. + } + await input.stopWorkers(); + input.logger.info({ reason }, 'nats disconnected'); + }; + + const startTimers = () => { + if (!idleTimer) { + idleTimer = setInterval(() => { + if (!session || input.isStopping()) return; + const activity = input.getActivity(); + const now = Date.now(); + const idleForMs = now - activity.lastActivityAt; + if (now - lastIdleStatusLogAt >= IDLE_STATUS_LOG_INTERVAL_MS) { + lastIdleStatusLogAt = now; + input.logger.info({ + ...activity, + idleForMs, + disconnectEligible: activity.inFlightHttp === 0 + && activity.inFlightJobs === 0 + && idleForMs >= IDLE_DISCONNECT_MS, + }, 'nats idle status'); + } + if (activity.inFlightHttp > 0 || activity.inFlightJobs > 0 || idleForMs < IDLE_DISCONNECT_MS) return; + void disconnect('idle'); + }, IDLE_CHECK_INTERVAL_MS); + idleTimer.unref?.(); + } + if (!orphanSweepTimer) { + orphanSweepTimer = setInterval(() => { + if (!session || input.isStopping()) return; + void input.runReconciliation().catch((error) => { + input.logger.error({ error: toErrorMessage(error) }, 'periodic orphaned operation recovery failed'); + }); + }, ORPHAN_SWEEP_INTERVAL_MS); + orphanSweepTimer.unref?.(); + } + }; + + const ensureConnected = async (): Promise => { + if (session) return session; + if (connecting) return connecting; + connecting = (async () => { + const nc = await connect(input.connectOptions); + const js = jetstream(nc, { timeout: NATS_API_TIMEOUT_MS }); + const jsm = await jetstreamManager(nc, { timeout: NATS_API_TIMEOUT_MS }); + await ensureJetStreamResources({ + jsm, + whisperTimeoutMs: input.whisperTimeoutMs, + pdfTimeoutMs: input.pdfTimeoutMs, + pdfAttempts: input.pdfAttempts, + jobsMaxBytes: input.jobsStreamMaxBytes, + eventsMaxBytes: input.eventsStreamMaxBytes, + natsReplicas: input.natsReplicas, + }); + const kv = await new Kvm(js).create(COMPUTE_STATE_BUCKET, { + replicas: input.natsReplicas, + history: 1, + ttl: COMPUTE_STATE_TTL_MS, + max_bytes: input.jobStatesMaxBytes, + }); + const next: NatsSession = { + nc, + js, + jsm, + kv, + whisperConsumer: await js.consumers.get(JOBS_STREAM_NAME, WHISPER_CONSUMER_NAME), + layoutConsumer: await js.consumers.get(JOBS_STREAM_NAME, LAYOUT_CONSUMER_NAME), + }; + session = next; + generation += 1; + input.markActivity('nats_connected'); + input.startWorkers(next); + startTimers(); + void nc.closed().then(() => { + if (session?.nc === nc) session = null; + }); + input.logger.info('nats connected'); + return next; + })(); + try { + return await connecting; + } finally { + connecting = null; + } + }; + + return { + ensureConnected, + disconnect, + isConnected: () => session !== null, + isOwnerActive: (owner: object) => session === owner, + getGeneration: () => generation, + async close(): Promise { + clearTimers(); + await input.stopWorkers(); + const current = session; + session = null; + if (!current) return; + try { + await current.nc.drain(); + } catch { + try { + await current.nc.close(); + } catch { + // Closing an already-closed session is harmless. + } + } + }, + }; +} diff --git a/compute-worker/src/jobs/handlers.ts b/compute-worker/src/jobs/handlers.ts index a013602..0f64c24 100644 --- a/compute-worker/src/jobs/handlers.ts +++ b/compute-worker/src/jobs/handlers.ts @@ -10,7 +10,7 @@ import type { PdfLayoutProgress, WhisperAlignJobRequest, WhisperAlignJobResult, -} from '../api/contracts'; +} from '../operations/contracts'; import type { ArtifactStorage } from '../infrastructure/storage'; import { persistParsedPdfWhileSourceExists } from './pdf-artifact-persistence'; import { buildInferProgressForPageParsed, buildInferProgressForPageStart } from './pdf-progress'; diff --git a/compute-worker/src/jobs/pdf-progress.ts b/compute-worker/src/jobs/pdf-progress.ts index 4aefe08..acc3867 100644 --- a/compute-worker/src/jobs/pdf-progress.ts +++ b/compute-worker/src/jobs/pdf-progress.ts @@ -1,4 +1,4 @@ -import type { PdfLayoutProgress } from '../api/contracts'; +import type { PdfLayoutProgress } from '../operations/contracts'; export function buildInferProgressForPageStart(input: { pageNumber: number; diff --git a/compute-worker/src/jobs/worker-loop-policy.ts b/compute-worker/src/jobs/worker-loop-policy.ts index ce3547a..c5f9c4e 100644 --- a/compute-worker/src/jobs/worker-loop-policy.ts +++ b/compute-worker/src/jobs/worker-loop-policy.ts @@ -1,4 +1,4 @@ -import type { WorkerOperationKind } from '../api/contracts'; +import type { WorkerOperationKind } from '../operations/contracts'; export type RetryAction = 'nak_retry' | 'term_fail'; diff --git a/compute-worker/src/jobs/worker-loop.ts b/compute-worker/src/jobs/worker-loop.ts index 194bf62..27c5afe 100644 --- a/compute-worker/src/jobs/worker-loop.ts +++ b/compute-worker/src/jobs/worker-loop.ts @@ -7,7 +7,7 @@ import type { WhisperAlignJobResult, WorkerJobTiming, WorkerOperationKind, -} from '../api/contracts'; +} from '../operations/contracts'; import type { JsonCodec } from '../infrastructure/json-codec'; import type { JobHandlers } from './handlers'; import { buildQueueWaitTiming, decideRetryAction } from './worker-loop-policy'; @@ -30,7 +30,7 @@ export interface QueuedJob { payload: TPayload; } -interface WorkerLoopOrchestrator { +export interface WorkerLoopOrchestrator { markRunning(input: { opId: string; startedAt?: number; updatedAt?: number; timing?: WorkerJobTiming }): Promise; markProgress(input: { opId: string; @@ -47,7 +47,7 @@ interface WorkerLoopOrchestrator { }): Promise; } -interface WorkerLogger { +export interface WorkerLogger { info(data: unknown, message?: string): void; warn(data: unknown, message?: string): void; error(data: unknown, message?: string): void; @@ -126,6 +126,17 @@ export function createWorkerLoopController(input: { latestProgress?: PdfLayoutProgress; }; + type JobRunner = ( + payload: TPayload, + queueWaitMs: number, + hooks?: { onProgress?: (progress: PdfLayoutProgress) => Promise }, + ) => Promise; + + type WorkDefinition = { + codec: JsonCodec>; + run: JobRunner; + }; + const markRunning = async (context: Context, updatedAt: number): Promise => { if (context.latestProgress) { await input.orchestrator.markProgress({ @@ -144,10 +155,8 @@ export function createWorkerLoopController(input: { }); }; - const processMessage = async (work: { + const processMessage = async (work: WorkDefinition & { msg: JsMsg; - codec: JsonCodec>; - run: JobHandlers['runWhisper'] | JobHandlers['runPdfLayout']; workerLabel: string; }): Promise => { let context: Context | null = null; @@ -180,7 +189,7 @@ export function createWorkerLoopController(input: { }, 'failed to persist operation heartbeat state'); }); }, RUNNING_HEARTBEAT_MS); - const result = await work.run(decoded.payload as never, context.queueWaitTiming?.queueWaitMs ?? 0, { + const result = await work.run(decoded.payload, context.queueWaitTiming?.queueWaitMs ?? 0, { onProgress: async (progress) => { try { work.msg.working(); @@ -259,11 +268,9 @@ export function createWorkerLoopController(input: { } }; - const runLoop = async (work: { + const runLoop = async (work: WorkDefinition & { owner: object; consumer: Consumer; - codec: JsonCodec>; - run: JobHandlers['runWhisper'] | JobHandlers['runPdfLayout']; workerLabel: string; }): Promise => { const detached = () => input.isStopping() || stopRequested || !input.isOwnerActive(work.owner); @@ -298,9 +305,17 @@ export function createWorkerLoopController(input: { start(owner: object, consumers: { whisper: Consumer; pdfLayout: Consumer }): void { stopRequested = false; loops = []; + const whisperWork: WorkDefinition = { + codec: input.whisperCodec, + run: input.handlers.runWhisper, + }; + const pdfWork: WorkDefinition = { + codec: input.pdfCodec, + run: input.handlers.runPdfLayout, + }; for (let i = 0; i < input.jobConcurrency; i += 1) { - loops.push(runLoop({ owner, consumer: consumers.whisper, codec: input.whisperCodec, run: input.handlers.runWhisper, workerLabel: `whisper-${i + 1}` })); - loops.push(runLoop({ owner, consumer: consumers.pdfLayout, codec: input.pdfCodec, run: input.handlers.runPdfLayout, workerLabel: `layout-${i + 1}` })); + loops.push(runLoop({ owner, consumer: consumers.whisper, ...whisperWork, workerLabel: `whisper-${i + 1}` })); + loops.push(runLoop({ owner, consumer: consumers.pdfLayout, ...pdfWork, workerLabel: `layout-${i + 1}` })); } }, async stop(): Promise { diff --git a/compute-worker/src/api/contracts.ts b/compute-worker/src/operations/contracts.ts similarity index 95% rename from compute-worker/src/api/contracts.ts rename to compute-worker/src/operations/contracts.ts index 4f78782..5180972 100644 --- a/compute-worker/src/api/contracts.ts +++ b/compute-worker/src/operations/contracts.ts @@ -1,18 +1,18 @@ -import type { TTSSentenceAlignment, ParsedPdfDocument } from './types'; +import type { TTSSentenceAlignment, ParsedPdfDocument } from '../api/types'; export type { TTSAudioBuffer, TTSAudioBytes, TTSSentenceAlignment, TTSSentenceWord, -} from './types'; +} from '../api/types'; export type { ParsedPdfBlockKind, ParsedPdfBlockFragment, ParsedPdfBlock, ParsedPdfPage, ParsedPdfDocument, -} from './types'; +} from '../api/types'; export const ALIGN_QUEUE_NAME = 'whisper-align'; export const PDF_LAYOUT_QUEUE_NAME = 'pdf-layout'; diff --git a/compute-worker/src/api/operation-keys.ts b/compute-worker/src/operations/keys.ts similarity index 96% rename from compute-worker/src/api/operation-keys.ts rename to compute-worker/src/operations/keys.ts index d814fb2..02fc2c8 100644 --- a/compute-worker/src/api/operation-keys.ts +++ b/compute-worker/src/operations/keys.ts @@ -1,5 +1,5 @@ import { createHash } from 'node:crypto'; -import { PDF_PARSER_VERSION } from '../api/contracts'; +import { PDF_PARSER_VERSION } from './contracts'; function sha256Hex(input: string): string { return createHash('sha256').update(input).digest('hex'); diff --git a/compute-worker/src/operations/reconciliation.ts b/compute-worker/src/operations/reconciliation.ts new file mode 100644 index 0000000..df8dc9e --- /dev/null +++ b/compute-worker/src/operations/reconciliation.ts @@ -0,0 +1,95 @@ +import type { WorkerJobTiming } from './contracts'; +import { + recoverOrphanedOperations, + type StreamedOperationState, +} from './recovery'; + +export interface ReconciliationStateStore { + getOpState(opId: string): Promise; + getOpStateRecord?(opId: string): Promise<{ state: StreamedOperationState; revision: number } | null>; + getOpIndex?(opKey: string): Promise<{ opId: string } | null>; + listOpStates?(): Promise; +} + +export interface ReconciliationOrchestrator { + markFailedIfUnchanged?(input: { + current: StreamedOperationState; + expectedRevision: number; + error: { message: string; code?: string } | string; + updatedAt?: number; + timing?: WorkerJobTiming; + }): Promise; +} + +interface ReconciliationLogger { + warn(data: unknown, message?: string): void; +} + +function supportsRecovery( + stateStore: ReconciliationStateStore, + orchestrator: ReconciliationOrchestrator, +): boolean { + return typeof stateStore.listOpStates === 'function' + && typeof stateStore.getOpStateRecord === 'function' + && typeof orchestrator.markFailedIfUnchanged === 'function'; +} + +export function createOperationReconciler(input: { + stateStore: ReconciliationStateStore; + orchestrator: ReconciliationOrchestrator; + whisperTimeoutMs: number; + pdfTimeoutMs: number; + opStaleMs: number; + getGeneration: () => number; + logger: ReconciliationLogger; +}) { + let recoveryPromise: Promise | null = null; + let recoveredGeneration = -1; + + const run = async (options?: { force?: boolean }): Promise => { + if (!supportsRecovery(input.stateStore, input.orchestrator)) return; + const generation = input.getGeneration(); + if (!options?.force && recoveredGeneration === generation) return; + if (recoveryPromise) return await recoveryPromise; + + recoveryPromise = (async () => { + const recoveredStates = await recoverOrphanedOperations({ + operationStateStore: { + getOpStateRecord: input.stateStore.getOpStateRecord!, + listOpStates: input.stateStore.listOpStates!, + }, + orchestrator: { + markFailedIfUnchanged: async (request) => { + const result = await input.orchestrator.markFailedIfUnchanged!(request); + return result as StreamedOperationState | null; + }, + }, + whisperTimeoutMs: input.whisperTimeoutMs, + pdfTimeoutMs: input.pdfTimeoutMs, + opStaleMs: input.opStaleMs, + }); + if (recoveredStates.length > 0) { + input.logger.warn({ + recoveredCount: recoveredStates.length, + ops: recoveredStates.map((state) => ({ + opId: state.opId, + kind: state.kind, + status: state.status, + })), + }, 'recovered stale in-flight operations during reconciliation'); + } + recoveredGeneration = generation; + })().finally(() => { + recoveryPromise = null; + }); + await recoveryPromise; + }; + + return { + run, + async getOpState(opId: string): Promise { + await run(); + return await input.stateStore.getOpState(opId); + }, + }; +} diff --git a/compute-worker/src/operations/recovery.ts b/compute-worker/src/operations/recovery.ts index 5aaaa10..b453a9a 100644 --- a/compute-worker/src/operations/recovery.ts +++ b/compute-worker/src/operations/recovery.ts @@ -4,7 +4,7 @@ import type { WorkerJobTiming, WorkerJobState, WorkerOperationState, -} from '../api/contracts'; +} from '../operations/contracts'; export type StreamedOperationState = WorkerOperationState; diff --git a/compute-worker/src/operations/service.ts b/compute-worker/src/operations/service.ts index 1469277..67430e9 100644 --- a/compute-worker/src/operations/service.ts +++ b/compute-worker/src/operations/service.ts @@ -5,7 +5,7 @@ import type { WorkerOperationKind, WorkerOperationRequest, WorkerOperationState, -} from '../api/contracts'; +} from '../operations/contracts'; import { buildQueuedState, createErrorShape, diff --git a/compute-worker/src/operations/state-machine.ts b/compute-worker/src/operations/state-machine.ts index 3144e8c..cd4076f 100644 --- a/compute-worker/src/operations/state-machine.ts +++ b/compute-worker/src/operations/state-machine.ts @@ -4,7 +4,7 @@ import type { WorkerOperationKind, WorkerOperationRequest, WorkerOperationState, -} from '../api/contracts'; +} from '../operations/contracts'; export function isTerminalStatus(status: WorkerJobState): boolean { return status === 'succeeded' || status === 'failed'; diff --git a/compute-worker/src/operations/types.ts b/compute-worker/src/operations/types.ts index 1d60200..d3f24c2 100644 --- a/compute-worker/src/operations/types.ts +++ b/compute-worker/src/operations/types.ts @@ -2,7 +2,7 @@ import type { WorkerOperationEvent, WorkerOperationKind, WorkerOperationState, -} from '../api/contracts'; +} from '../operations/contracts'; export type OperationState = WorkerOperationState; export type OperationEvent = WorkerOperationEvent; diff --git a/compute-worker/src/storage/artifact-addressing.ts b/compute-worker/src/storage/artifact-addressing.ts index bf4e312..10bb820 100644 --- a/compute-worker/src/storage/artifact-addressing.ts +++ b/compute-worker/src/storage/artifact-addressing.ts @@ -1,5 +1,5 @@ -import { PDF_PARSER_VERSION } from '../api/contracts'; -import { encodeParserVersion } from '../api/contracts'; +import { PDF_PARSER_VERSION } from '../operations/contracts'; +import { encodeParserVersion } from '../operations/contracts'; const DOCUMENT_ID_REGEX = /^[a-f0-9]{64}$/i; const SAFE_NAMESPACE_REGEX = /^[a-zA-Z0-9._-]{1,128}$/; diff --git a/compute-worker/tests/compute/algorithms/whisper-decoder.test.ts b/compute-worker/tests/compute/algorithms/whisper-decoder.test.ts new file mode 100644 index 0000000..07f405e --- /dev/null +++ b/compute-worker/tests/compute/algorithms/whisper-decoder.test.ts @@ -0,0 +1,61 @@ +import { describe, expect, test } from 'vitest'; +import { + applyTokenSuppression, + applyWhisperTimestampLogitsRules, + argmax, +} from '../../../src/inference/whisper/decoder'; + +describe('whisper decoder logits', () => { + test('selects finite maximums and suppresses valid token ids', () => { + const logits = new Float32Array([1, 5, 3]); + + applyTokenSuppression(logits, new Set([-1, 1, 10])); + + expect(argmax(logits)).toBe(2); + expect(logits[1]).toBe(Number.NEGATIVE_INFINITY); + expect(argmax(new Float32Array([Number.NEGATIVE_INFINITY]))).toBeNull(); + }); + + test('forces an initial timestamp within the configured range', () => { + const logits = new Float32Array([4, 3, 2, 1, 8, 7, 6, 5]); + + applyWhisperTimestampLogitsRules({ + logits, + generated: [10, 11], + beginIndex: 2, + eosTokenId: 3, + noTimestampsTokenId: 2, + timestampBeginTokenId: 4, + maxInitialTimestampIndex: 1, + }); + + expect(Array.from(logits.slice(0, 4))).toEqual([ + Number.NEGATIVE_INFINITY, + Number.NEGATIVE_INFINITY, + Number.NEGATIVE_INFINITY, + Number.NEGATIVE_INFINITY, + ]); + expect(logits[4]).toBe(8); + expect(logits[5]).toBe(7); + expect(logits[6]).toBe(Number.NEGATIVE_INFINITY); + expect(logits[7]).toBe(Number.NEGATIVE_INFINITY); + }); + + test('requires text after a timestamp pair', () => { + const logits = new Float32Array([4, 3, 2, 1, 8, 7]); + + applyWhisperTimestampLogitsRules({ + logits, + generated: [10, 4, 5], + beginIndex: 1, + eosTokenId: 3, + noTimestampsTokenId: -1, + timestampBeginTokenId: 4, + maxInitialTimestampIndex: Number.POSITIVE_INFINITY, + }); + + expect(logits[4]).toBe(Number.NEGATIVE_INFINITY); + expect(logits[5]).toBe(Number.NEGATIVE_INFINITY); + expect(logits[0]).toBe(4); + }); +}); diff --git a/compute-worker/tests/compute/control-plane/orchestrator.test.ts b/compute-worker/tests/compute/control-plane/orchestrator.test.ts index 66964d3..092cfdb 100644 --- a/compute-worker/tests/compute/control-plane/orchestrator.test.ts +++ b/compute-worker/tests/compute/control-plane/orchestrator.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from 'vitest'; -import type { WorkerOperationRequest } from '../../../src/api/contracts'; +import type { WorkerOperationRequest } from '../../../src/operations/contracts'; import { OperationOrchestrator } from '../../../src/operations'; import { InMemoryOperationEventStream, diff --git a/compute-worker/tests/compute/control-plane/state-machine.test.ts b/compute-worker/tests/compute/control-plane/state-machine.test.ts index 8fba44a..667084f 100644 --- a/compute-worker/tests/compute/control-plane/state-machine.test.ts +++ b/compute-worker/tests/compute/control-plane/state-machine.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from 'vitest'; -import type { WorkerOperationState } from '../../../src/api/contracts'; +import type { WorkerOperationState } from '../../../src/operations/contracts'; import { explainReplacementReason, isInflightStatus, diff --git a/compute-worker/tests/compute/helpers/in-memory-control-plane.ts b/compute-worker/tests/compute/helpers/in-memory-control-plane.ts index 811521b..0348bff 100644 --- a/compute-worker/tests/compute/helpers/in-memory-control-plane.ts +++ b/compute-worker/tests/compute/helpers/in-memory-control-plane.ts @@ -1,5 +1,5 @@ import { EventEmitter } from 'node:events'; -import type { WorkerOperationKind } from '../../../src/api/contracts'; +import type { WorkerOperationKind } from '../../../src/operations/contracts'; import type { OperationEvent, OperationEventStream, diff --git a/compute-worker/tests/fixtures/fake-control-plane.ts b/compute-worker/tests/fixtures/fake-control-plane.ts index 23f2760..0cd1c86 100644 --- a/compute-worker/tests/fixtures/fake-control-plane.ts +++ b/compute-worker/tests/fixtures/fake-control-plane.ts @@ -4,7 +4,7 @@ import type { WorkerOperationEvent, WorkerOperationRequest, WorkerOperationState, -} from '../../src/api/contracts'; +} from '../../src/operations/contracts'; import type { ComputeWorkerRouteDeps } from '../../src/api/app'; type ComputeResult = WhisperAlignJobResult | PdfLayoutJobResult; diff --git a/compute-worker/tests/unit/jetstream-adapters.test.ts b/compute-worker/tests/unit/jetstream-adapters.test.ts index a238996..8425061 100644 --- a/compute-worker/tests/unit/jetstream-adapters.test.ts +++ b/compute-worker/tests/unit/jetstream-adapters.test.ts @@ -1,6 +1,6 @@ import { describe, expect, test } from 'vitest'; import { OperationOrchestrator } from '../../src/operations'; -import type { WorkerOperationRequest } from '../../src/api/contracts'; +import type { WorkerOperationRequest } from '../../src/operations/contracts'; import { JetStreamOperationEventStream, JetStreamOperationQueue, diff --git a/compute-worker/tests/unit/key-and-progress.test.ts b/compute-worker/tests/unit/key-and-progress.test.ts index d6274db..dc1656a 100644 --- a/compute-worker/tests/unit/key-and-progress.test.ts +++ b/compute-worker/tests/unit/key-and-progress.test.ts @@ -4,7 +4,7 @@ import { buildInferProgressForPageParsed, buildInferProgressForPageStart, } from '../../src/jobs/pdf-progress'; -import { buildPdfOperationKey } from '../../src/api/operation-keys'; +import { buildPdfOperationKey } from '../../src/operations/keys'; import { parsedPdfArtifactKey } from '../../src/storage/artifact-addressing'; describe('compute worker helpers', () => { diff --git a/compute-worker/tests/unit/worker-loop.test.ts b/compute-worker/tests/unit/worker-loop.test.ts new file mode 100644 index 0000000..3b9ba07 --- /dev/null +++ b/compute-worker/tests/unit/worker-loop.test.ts @@ -0,0 +1,170 @@ +import type { Consumer, JsMsg } from '@nats-io/jetstream'; +import { describe, expect, test, vi } from 'vitest'; +import type { PdfLayoutProgress } from '../../src/operations/contracts'; +import { createJsonCodec } from '../../src/infrastructure/json-codec'; +import { + createWorkerLoopController, + type QueuedJob, + type WorkerLoopOrchestrator, +} from '../../src/jobs/worker-loop'; + +function createMessage(job: QueuedJob, deliveryCount = 1) { + const codec = createJsonCodec>(); + const ack = vi.fn(); + const nak = vi.fn(); + const term = vi.fn(); + const working = vi.fn(); + return { + codec, + msg: { + data: codec.encode(job), + info: { deliveryCount }, + ack, + nak, + term, + working, + } as unknown as JsMsg, + ack, + nak, + term, + working, + }; +} + +function createConsumer(message?: JsMsg): Consumer { + let delivered = false; + return { + next: async () => { + if (!delivered && message) { + delivered = true; + return message; + } + await new Promise((resolve) => setTimeout(resolve, 1)); + return null; + }, + } as unknown as Consumer; +} + +function createOrchestrator() { + const calls: Array<{ method: string; input: unknown }> = []; + const record = (method: string) => async (input: unknown) => { + calls.push({ method, input }); + return input; + }; + const orchestrator: WorkerLoopOrchestrator = { + markRunning: record('running'), + markProgress: record('progress'), + markSucceeded: record('succeeded'), + markFailed: record('failed'), + }; + return { orchestrator, calls }; +} + +describe('worker loop controller', () => { + test('processes PDF progress, persists success, ACKs, and releases in-flight activity', async () => { + const owner = {}; + let active = true; + let inFlight = 0; + const { orchestrator, calls } = createOrchestrator(); + let complete!: () => void; + const completed = new Promise((resolve) => { complete = resolve; }); + const progress: PdfLayoutProgress = { totalPages: 2, pagesParsed: 1, currentPage: 1, phase: 'infer' }; + const pdf = createMessage({ + jobId: 'job-pdf', + opId: 'op-pdf', + opKey: 'pdf-key', + kind: 'pdf_layout', + queuedAt: Date.now() - 10, + payload: { + documentId: 'a'.repeat(64), + namespace: null, + documentObjectKey: 'openreader/doc.pdf', + }, + }); + const whisperCodec = createJsonCodec>(); + const controller = createWorkerLoopController({ + orchestrator, + handlers: { + runWhisper: async () => ({ alignments: [] }), + runPdfLayout: async (_payload, _queueWaitMs, hooks) => { + await hooks?.onProgress?.(progress); + active = false; + complete(); + return { parsedObjectKey: 'openreader/parsed.json' }; + }, + }, + logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn() }, + jobConcurrency: 1, + pdfAttempts: 2, + whisperCodec, + pdfCodec: pdf.codec, + isOwnerActive: () => active, + isStopping: () => false, + markActivity: vi.fn(), + onInFlightJobsChanged: (delta) => { inFlight += delta; }, + }); + + controller.start(owner, { whisper: createConsumer(), pdfLayout: createConsumer(pdf.msg) }); + await completed; + await controller.stop(); + + expect(pdf.working).toHaveBeenCalledOnce(); + expect(pdf.ack).toHaveBeenCalledOnce(); + expect(pdf.nak).not.toHaveBeenCalled(); + expect(inFlight).toBe(0); + expect(calls.map((call) => call.method)).toContain('progress'); + expect(calls.map((call) => call.method)).toContain('succeeded'); + }); + + test('NAKs a retryable PDF failure and does not mark it terminal', async () => { + const owner = {}; + let active = true; + const { orchestrator, calls } = createOrchestrator(); + let attempted!: () => void; + const attemptCompleted = new Promise((resolve) => { attempted = resolve; }); + const pdf = createMessage({ + jobId: 'job-pdf', + opId: 'op-pdf', + opKey: 'pdf-key', + kind: 'pdf_layout', + queuedAt: Date.now(), + payload: { + documentId: 'a'.repeat(64), + namespace: null, + documentObjectKey: 'openreader/doc.pdf', + }, + }); + const controller = createWorkerLoopController({ + orchestrator, + handlers: { + runWhisper: async () => ({ alignments: [] }), + runPdfLayout: async () => { + active = false; + attempted(); + throw new Error('retry me'); + }, + }, + logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn() }, + jobConcurrency: 1, + pdfAttempts: 2, + whisperCodec: createJsonCodec(), + pdfCodec: pdf.codec, + isOwnerActive: () => active, + isStopping: () => false, + markActivity: vi.fn(), + onInFlightJobsChanged: vi.fn(), + }); + + controller.start(owner, { whisper: createConsumer(), pdfLayout: createConsumer(pdf.msg) }); + await attemptCompleted; + await new Promise((resolve) => setTimeout(resolve, 0)); + await controller.stop(); + + expect(pdf.nak).toHaveBeenCalledOnce(); + expect(pdf.term).not.toHaveBeenCalled(); + expect(calls.map((call) => call.method)).not.toContain('failed'); + }); +}); diff --git a/src/lib/server/compute-worker/client.ts b/src/lib/server/compute-worker/client.ts index 2f8c0b6..a8f35bf 100644 --- a/src/lib/server/compute-worker/client.ts +++ b/src/lib/server/compute-worker/client.ts @@ -4,10 +4,10 @@ import type { PdfLayoutResult, WhisperAlignRequest, WhisperAlignResult, - WorkerOperation, - WorkerOperationEvent, + ComputeOperation, + ComputeOperationEvent, } from './protocol'; -import { isWorkerOperation } from './protocol'; +import { isComputeOperation } from './protocol'; class WorkerHttpError extends Error { constructor( @@ -88,7 +88,7 @@ function parseSseFrame(frame: string): { id: number | null; data: string | null return { id, data: data.length > 0 ? data.join('\n') : null }; } -function isTerminal(operation: WorkerOperation): boolean { +function isTerminal(operation: ComputeOperation): boolean { return operation.status === 'succeeded' || operation.status === 'failed'; } @@ -105,7 +105,7 @@ export class ComputeWorkerClient { let lastError: unknown; for (let attempt = 1; attempt <= 2; attempt += 1) { try { - const operation = await this.requestJson>( + const operation = await this.requestJson>( 'POST', '/v1/whisper-align/operations', input, @@ -129,7 +129,7 @@ export class ComputeWorkerClient { throw lastError instanceof Error ? lastError : new Error('Unknown compute worker failure'); } - createPdfLayoutOperation(input: PdfLayoutRequest): Promise> { + createPdfLayoutOperation(input: PdfLayoutRequest): Promise> { return this.requestJson('POST', '/v1/pdf-layout/operations', input); } @@ -141,7 +141,7 @@ export class ComputeWorkerClient { }); } - async getOperation(opId: string): Promise | null> { + async getOperation(opId: string): Promise | null> { try { return await this.requestJson('GET', `/v1/operations/${encodeURIComponent(opId)}`); } catch (error) { @@ -199,7 +199,7 @@ export class ComputeWorkerClient { return error instanceof Error && /network|timeout|fetch failed/i.test(error.message); } - private async waitForOperation(opId: string, timeoutMs: number): Promise> { + private async waitForOperation(opId: string, timeoutMs: number): Promise> { const controller = new AbortController(); const timeout = setTimeout(() => controller.abort(), timeoutMs); let lastEventId: number | null = null; @@ -227,8 +227,8 @@ export class ComputeWorkerClient { buffer = buffer.slice(end + 2); if (frame.id !== null) lastEventId = frame.id; if (!frame.data) continue; - const event = JSON.parse(frame.data) as WorkerOperationEvent; - if (!event?.snapshot || !isWorkerOperation(event.snapshot) || event.snapshot.opId !== opId) continue; + const event = JSON.parse(frame.data) as ComputeOperationEvent; + if (!event?.snapshot || !isComputeOperation(event.snapshot) || event.snapshot.opId !== opId) continue; if (isTerminal(event.snapshot)) return event.snapshot; } } diff --git a/src/lib/server/compute-worker/generated.ts b/src/lib/server/compute-worker/generated.ts index feb2060..cdac950 100644 --- a/src/lib/server/compute-worker/generated.ts +++ b/src/lib/server/compute-worker/generated.ts @@ -469,7 +469,7 @@ export interface paths { }; requestBody?: never; responses: { - /** @description Server-sent PublicOperationEvent stream */ + /** @description Server-sent ComputeOperationEvent stream */ 200: { headers: { [name: string]: unknown; @@ -579,7 +579,7 @@ export interface components { /** @enum {string} */ phase: "infer" | "merge"; }; - PublicOperation: { + ComputeOperation: { opId: string; subject: { /** @enum {string} */ @@ -613,7 +613,7 @@ export interface components { phase: "infer" | "merge"; }; }; - PublicOperationEvent: { + ComputeOperationEvent: { eventId: number; snapshot: { opId: string; diff --git a/src/lib/server/compute-worker/protocol.ts b/src/lib/server/compute-worker/protocol.ts index 73b196c..2fa0c83 100644 --- a/src/lib/server/compute-worker/protocol.ts +++ b/src/lib/server/compute-worker/protocol.ts @@ -9,17 +9,17 @@ export type ParsedPdfBlockFragment = ParsedPdfBlock['fragments'][number]; export type TTSSentenceAlignment = components['schemas']['TTSSentenceAlignment']; export type TTSSentenceWord = TTSSentenceAlignment['words'][number]; -export type PdfLayoutProgress = NonNullable; -export type WorkerOperationStatus = components['schemas']['PublicOperation']['status']; -export type WorkerOperationSubject = components['schemas']['PublicOperation']['subject']; +export type PdfLayoutProgress = NonNullable; +export type ComputeOperationStatus = components['schemas']['ComputeOperation']['status']; +export type ComputeOperationSubject = components['schemas']['ComputeOperation']['subject']; -export type WorkerOperation = - Omit +export type ComputeOperation = + Omit & { result?: Result }; -export type WorkerOperationEvent = { +export type ComputeOperationEvent = { eventId: number; - snapshot: WorkerOperation; + snapshot: ComputeOperation; }; export type WhisperAlignRequest = @@ -31,20 +31,20 @@ export type PdfLayoutResolveRequest = export type WhisperAlignResult = { alignments: TTSSentenceAlignment[]; - timing?: components['schemas']['PublicOperation']['timing']; + timing?: components['schemas']['ComputeOperation']['timing']; }; export type PdfLayoutResult = { parsedObjectKey: string; - timing?: components['schemas']['PublicOperation']['timing']; + timing?: components['schemas']['ComputeOperation']['timing']; }; export type PdfLayoutResolution = { artifact: { objectKey: string } | null; - operation: WorkerOperation | null; + operation: ComputeOperation | null; }; -export function isWorkerOperation(value: unknown): value is WorkerOperation { +export function isComputeOperation(value: unknown): value is ComputeOperation { if (!value || typeof value !== 'object') return false; const record = value as Record; return typeof record.opId === 'string' diff --git a/src/lib/server/pdf-parse/operation.ts b/src/lib/server/pdf-parse/operation.ts index 60a924a..a10136b 100644 --- a/src/lib/server/pdf-parse/operation.ts +++ b/src/lib/server/pdf-parse/operation.ts @@ -3,7 +3,7 @@ import { documentKey } from '@/lib/server/documents/blobstore'; import type { PdfLayoutResolution, PdfLayoutResult, - WorkerOperation, + ComputeOperation, } from '@/lib/server/compute-worker/protocol'; function currentPdfOperationInput(documentId: string, namespace: string | null, forceToken?: string): { @@ -33,7 +33,7 @@ export async function resolveCurrentPdfParse(input: { export async function lookupCurrentPdfParseOperation(input: { documentId: string; namespace: string | null; -}): Promise | null> { +}): Promise | null> { return (await resolveCurrentPdfParse(input)).operation; } @@ -41,7 +41,7 @@ export async function createOrReuseCurrentPdfParseOperation(input: { documentId: string; namespace: string | null; forceToken?: string; -}): Promise> { +}): Promise> { return getComputeWorkerClient().createPdfLayoutOperation(currentPdfOperationInput( input.documentId, input.namespace, @@ -49,12 +49,12 @@ export async function createOrReuseCurrentPdfParseOperation(input: { )); } -export async function fetchPdfParseOperation(opId: string): Promise | null> { +export async function fetchPdfParseOperation(opId: string): Promise | null> { return getComputeWorkerClient().getOperation(opId); } export function isPdfParseOperationForDocument( - state: WorkerOperation, + state: ComputeOperation, input: { documentId: string; namespace: string | null; diff --git a/src/lib/server/pdf-parse/snapshot.ts b/src/lib/server/pdf-parse/snapshot.ts index 2fba7f9..b4e8924 100644 --- a/src/lib/server/pdf-parse/snapshot.ts +++ b/src/lib/server/pdf-parse/snapshot.ts @@ -1,8 +1,8 @@ -import type { PdfLayoutResult, WorkerOperation } from '@/lib/server/compute-worker/protocol'; +import type { PdfLayoutResult, ComputeOperation } from '@/lib/server/compute-worker/protocol'; import type { PdfParseStatus } from '@/types/parsed-pdf'; import type { PdfParseSnapshot } from '@/lib/server/pdf-parse/types'; -function mapWorkerStatusToParseStatus(status: WorkerOperation['status']): PdfParseStatus { +function mapWorkerStatusToParseStatus(status: ComputeOperation['status']): PdfParseStatus { switch (status) { case 'queued': return 'pending'; @@ -19,7 +19,7 @@ function mapWorkerStatusToParseStatus(status: WorkerOperation['status']): PdfPar } export function parsedObjectKeyFromWorkerState( - state: WorkerOperation, + state: ComputeOperation, ): string | null { const result = state.result; if (!result || typeof result !== 'object' || !('parsedObjectKey' in result)) return null; @@ -30,7 +30,7 @@ export function parsedObjectKeyFromWorkerState( } export function pdfParseSnapshotFromWorkerState( - state: WorkerOperation, + state: ComputeOperation, ): PdfParseSnapshot { const parseStatus = mapWorkerStatusToParseStatus(state.status); return {