danvics/openreader
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor(compute-worker): modularize operation contracts and API, update protocol naming

Browse source 
- Move operation contracts from api/contracts to operations/contracts for clearer separation
- Delete legacy API modules: contracts.ts, operation-keys.ts, public-operation.ts
- Add new modular files: compute-operation.ts, http-hooks.ts, routes.ts, operations/keys.ts, operations/reconciliation.ts
- Refactor all imports to use new operations/contracts and keys modules
- Rename protocol types from PublicOperation to ComputeOperation for clarity
- Update OpenAPI schema and all references to use ComputeOperation/Event naming
- Extract whisper decoder helpers to inference/whisper/decoder.ts
- Add new infra: nats-session.ts for NATS session management
- Update tests and storage logic to use new contracts and keys modules
- Refactor worker loop and orchestrator to use modular job definitions
- Update client and protocol in src/lib/server/compute-worker to use new ComputeOperation types

This refactor modularizes compute-worker operation logic, clarifies type naming, and improves maintainability by separating contracts, keys, and reconciliation logic. No breaking changes to external API, but all internal references and tests are updated for the new structure.
This commit is contained in:
Richard R 2026-06-13 05:04:24 -06:00
parent 6e69b48103
commit 5403aa00cc
37 changed files with 1313 additions and 963 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
compute-worker/openapi.json
View file

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

865
compute-worker/src/api/app.ts
View file

@ -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<ReturnType<Kvm['create']>>;
whisperConsumer: Consumer;
layoutConsumer: Consumer;
}
interface OperationEventStreamLike {
subscribe(input: {
opId: string;
sinceEventId?: number;
onEvent: (event: WorkerOperationEvent<WhisperAlignJobResult | PdfLayoutJobResult>) => void | Promise<void>;
onError?: (error: unknown) => void;
}): Promise<() => void>;
}
interface OperationStateStoreLike {
getOpState(opId: string): Promise<StreamedOperationState | null>;
getOpStateRecord?(opId: string): Promise<{ state: StreamedOperationState; revision: number } | null>;
getOpIndex?(opKey: string): Promise<{ opId: string } | null>;
listOpStates?(): Promise<StreamedOperationState[]>;
}
interface OrchestratorLike {
enqueueOrReuse(request: WorkerOperationRequest): Promise<StreamedOperationState>;
markRunning(input: {
opId: string;
startedAt?: number;
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<StreamedOperationState>;
markProgress(input: {
opId: string;
progress: WorkerOperationState['progress'];
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<StreamedOperationState>;
markSucceeded(input: {
opId: string;
result: unknown;
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<StreamedOperationState>;
markFailed(input: {
opId: string;
error: { message: string; code?: string } | string;
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<StreamedOperationState>;
markFailedIfUnchanged?(input: {
current: StreamedOperationState;
expectedRevision: number;
error: { message: string; code?: string } | string;
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<StreamedOperationState | null>;
}
export interface ComputeWorkerRouteDeps {
orchestrator: OrchestratorLike;
operationStateStore: OperationStateStoreLike;
operationEventStream: OperationEventStreamLike;
artifactExists?: (key: string) => Promise<boolean>;
}
export type { ComputeWorkerRouteDeps } from './routes';
export interface CreateComputeWorkerAppOptions {
host?: string;
@ -187,122 +94,27 @@ export interface ComputeWorkerApp {
close(): Promise<void>;
}
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<string, unknown> {
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<ComputeWorkerApp> {
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<typeof connect>[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<NatsSession> | 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<void> {
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<NatsSession> {
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<typeof createNatsSessionManager>;
const ensureConnected = () => sessionManager.ensureConnected();
const s3Prefix = normalizeS3Prefix(process.env.S3_PREFIX);
const storageDisabled = async (): Promise<never> => {
@ -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<void> | 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<void> => {
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<StreamedOperationState[]>;
},
orchestrator: orchestrator as typeof orchestrator & {
markFailedIfUnchanged(input: {
current: StreamedOperationState;
expectedRevision: number;
error: { message: string; code?: string } | string;
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<StreamedOperationState | null>;
},
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<void> => {
await runOrphanedOpRecovery();
};
const getOpState = async (opId: string): Promise<StreamedOperationState | null> => {
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<WhisperAlignJobResult | PdfLayoutJobResult> = {
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<void>((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<void> => {
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 => {

18
compute-worker/src/api/public-operation.ts → compute-worker/src/api/compute-operation.ts
View file

@ -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<Result = unknown> {
export interface ComputeOperation<Result = unknown> {
opId: string;
subject: PublicOperationSubject;
subject: ComputeOperationSubject;
status: WorkerOperationState['status'];
queuedAt: number;
updatedAt: number;
@ -18,14 +18,14 @@ export interface PublicOperation<Result = unknown> {
progress?: WorkerOperationState['progress'];
}
export interface PublicOperationEvent<Result = unknown> {
export interface ComputeOperationEvent<Result = unknown> {
eventId: number;
snapshot: PublicOperation<Result>;
snapshot: ComputeOperation<Result>;
}
export function toPublicOperation<Result>(
export function toComputeOperation<Result>(
state: WorkerOperationState<Result>,
): PublicOperation<Result> {
): ComputeOperation<Result> {
const subject = state.kind === 'pdf_layout'
? (pdfSubjectFromOperationKey(state.opKey) ?? { kind: 'pdf_layout', documentId: '', namespace: null })
: { kind: 'whisper_align' as const };

80
compute-worker/src/api/http-hooks.ts Normal file
View file

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

356
compute-worker/src/api/routes.ts Normal file
View file

@ -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<WhisperAlignJobResult | PdfLayoutJobResult>) => void | Promise<void>;
onError?: (error: unknown) => void;
}): Promise<() => void>;
}
interface OperationStateStoreLike extends ReconciliationStateStore {}
interface OrchestratorLike {
enqueueOrReuse(request: WorkerOperationRequest): Promise<WorkerOperationState>;
markRunning(input: {
opId: string;
startedAt?: number;
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<unknown>;
markProgress(input: {
opId: string;
progress: WorkerOperationState['progress'];
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<unknown>;
markSucceeded(input: {
opId: string;
result: unknown;
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<unknown>;
markFailed(input: {
opId: string;
error: { message: string; code?: string } | string;
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<unknown>;
markFailedIfUnchanged?(input: {
current: StreamedOperationState;
expectedRevision: number;
error: { message: string; code?: string } | string;
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<unknown>;
}
export interface ComputeWorkerRouteDeps {
orchestrator: OrchestratorLike;
operationStateStore: OperationStateStoreLike;
operationEventStream: OperationEventStreamLike;
artifactExists?: (key: string) => Promise<boolean>;
}
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<void>;
getOpState: (opId: string) => Promise<StreamedOperationState | null>;
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<WhisperAlignJobResult | PdfLayoutJobResult> = {
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<void>((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;
});
}

8
compute-worker/src/api/schemas.ts
View file

@ -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<string, unknown> {

2
compute-worker/src/inference/pdf/parse.ts
View file

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

117
compute-worker/src/inference/whisper/align.ts
View file

@ -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<number>) {
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<WhisperRuntime> {
if (state.runtimePromise) return state.runtimePromise;

111
compute-worker/src/inference/whisper/decoder.ts Normal file
View file

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

36
compute-worker/src/infrastructure/config.ts
View file

@ -19,7 +19,13 @@ export type IdleTimeoutAndHardCapInput<T> = {
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<string, unknown> {
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;

2
compute-worker/src/infrastructure/nats-adapters.ts
View file

@ -13,7 +13,7 @@ import type {
PdfLayoutJobRequest,
WhisperAlignJobRequest,
WorkerOperationKind,
} from '../api/contracts';
} from '../operations/contracts';
import { createJsonCodec } from './json-codec';
export interface KvEntryLike {

204
compute-worker/src/infrastructure/nats-session.ts Normal file
View file

@ -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<ReturnType<Kvm['create']>>;
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<void>;
runReconciliation: () => Promise<void>;
}) {
let session: NatsSession | null = null;
let connecting: Promise<NatsSession> | 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<void> => {
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<NatsSession> => {
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<void> {
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.
}
}
},
};
}

2
compute-worker/src/jobs/handlers.ts
View file

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

2
compute-worker/src/jobs/pdf-progress.ts
View file

@ -1,4 +1,4 @@
import type { PdfLayoutProgress } from '../api/contracts';
import type { PdfLayoutProgress } from '../operations/contracts';
export function buildInferProgressForPageStart(input: {
pageNumber: number;

2
compute-worker/src/jobs/worker-loop-policy.ts
View file

@ -1,4 +1,4 @@
import type { WorkerOperationKind } from '../api/contracts';
import type { WorkerOperationKind } from '../operations/contracts';
export type RetryAction = 'nak_retry' | 'term_fail';

39
compute-worker/src/jobs/worker-loop.ts
View file

@ -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<TPayload> {
payload: TPayload;
}
interface WorkerLoopOrchestrator {
export interface WorkerLoopOrchestrator {
markRunning(input: { opId: string; startedAt?: number; updatedAt?: number; timing?: WorkerJobTiming }): Promise<unknown>;
markProgress(input: {
opId: string;
@ -47,7 +47,7 @@ interface WorkerLoopOrchestrator {
}): Promise<unknown>;
}
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<TPayload, TResult> = (
payload: TPayload,
queueWaitMs: number,
hooks?: { onProgress?: (progress: PdfLayoutProgress) => Promise<void> },
) => Promise<TResult>;
type WorkDefinition<TPayload, TResult> = {
codec: JsonCodec<QueuedJob<TPayload>>;
run: JobRunner<TPayload, TResult>;
};
const markRunning = async <TPayload>(context: Context<TPayload>, updatedAt: number): Promise<void> => {
if (context.latestProgress) {
await input.orchestrator.markProgress({
@ -144,10 +155,8 @@ export function createWorkerLoopController(input: {
});
};
const processMessage = async <TPayload, TResult>(work: {
const processMessage = async <TPayload, TResult>(work: WorkDefinition<TPayload, TResult> & {
msg: JsMsg;
codec: JsonCodec<QueuedJob<TPayload>>;
run: JobHandlers['runWhisper'] | JobHandlers['runPdfLayout'];
workerLabel: string;
}): Promise<void> => {
let context: Context<TPayload> | 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 <TPayload>(work: {
const runLoop = async <TPayload, TResult>(work: WorkDefinition<TPayload, TResult> & {
owner: object;
consumer: Consumer;
codec: JsonCodec<QueuedJob<TPayload>>;
run: JobHandlers['runWhisper'] | JobHandlers['runPdfLayout'];
workerLabel: string;
}): Promise<void> => {
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<WhisperAlignJobRequest, WhisperAlignJobResult> = {
codec: input.whisperCodec,
run: input.handlers.runWhisper,
};
const pdfWork: WorkDefinition<PdfLayoutJobRequest, PdfLayoutJobResult> = {
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<void> {

6
compute-worker/src/api/contracts.ts → compute-worker/src/operations/contracts.ts
View file

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

2
compute-worker/src/api/operation-keys.ts → compute-worker/src/operations/keys.ts
View file

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

95
compute-worker/src/operations/reconciliation.ts Normal file
View file

@ -0,0 +1,95 @@
import type { WorkerJobTiming } from './contracts';
import {
recoverOrphanedOperations,
type StreamedOperationState,
} from './recovery';
export interface ReconciliationStateStore {
getOpState(opId: string): Promise<StreamedOperationState | null>;
getOpStateRecord?(opId: string): Promise<{ state: StreamedOperationState; revision: number } | null>;
getOpIndex?(opKey: string): Promise<{ opId: string } | null>;
listOpStates?(): Promise<StreamedOperationState[]>;
}
export interface ReconciliationOrchestrator {
markFailedIfUnchanged?(input: {
current: StreamedOperationState;
expectedRevision: number;
error: { message: string; code?: string } | string;
updatedAt?: number;
timing?: WorkerJobTiming;
}): Promise<unknown>;
}
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<void> | null = null;
let recoveredGeneration = -1;
const run = async (options?: { force?: boolean }): Promise<void> => {
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<StreamedOperationState | null> {
await run();
return await input.stateStore.getOpState(opId);
},
};
}

2
compute-worker/src/operations/recovery.ts
View file

@ -4,7 +4,7 @@ import type {
WorkerJobTiming,
WorkerJobState,
WorkerOperationState,
} from '../api/contracts';
} from '../operations/contracts';
export type StreamedOperationState = WorkerOperationState<WhisperAlignJobResult | PdfLayoutJobResult>;

2
compute-worker/src/operations/service.ts
View file

@ -5,7 +5,7 @@ import type {
WorkerOperationKind,
WorkerOperationRequest,
WorkerOperationState,
} from '../api/contracts';
} from '../operations/contracts';
import {
buildQueuedState,
createErrorShape,

2
compute-worker/src/operations/state-machine.ts
View file

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

2
compute-worker/src/operations/types.ts
View file

@ -2,7 +2,7 @@ import type {
WorkerOperationEvent,
WorkerOperationKind,
WorkerOperationState,
} from '../api/contracts';
} from '../operations/contracts';
export type OperationState<Result = unknown> = WorkerOperationState<Result>;
export type OperationEvent<Result = unknown> = WorkerOperationEvent<Result>;

4
compute-worker/src/storage/artifact-addressing.ts
View file

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

61
compute-worker/tests/compute/algorithms/whisper-decoder.test.ts Normal file
View file

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

2
compute-worker/tests/compute/control-plane/orchestrator.test.ts
View file

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

2
compute-worker/tests/compute/control-plane/state-machine.test.ts
View file

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

2
compute-worker/tests/compute/helpers/in-memory-control-plane.ts
View file

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

2
compute-worker/tests/fixtures/fake-control-plane.ts vendored
View file

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

2
compute-worker/tests/unit/jetstream-adapters.test.ts
View file

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

2
compute-worker/tests/unit/key-and-progress.test.ts
View file

@ -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', () => {

170
compute-worker/tests/unit/worker-loop.test.ts Normal file
View file

@ -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<T>(job: QueuedJob<T>, deliveryCount = 1) {
const codec = createJsonCodec<QueuedJob<T>>();
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<void>((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<QueuedJob<{
text: string;
audioObjectKey: string;
}>>();
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<void>((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');
});
});

20
src/lib/server/compute-worker/client.ts
View file

@ -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<WorkerOperation<WhisperAlignResult>>(
const operation = await this.requestJson<ComputeOperation<WhisperAlignResult>>(
'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<WorkerOperation<PdfLayoutResult>> {
createPdfLayoutOperation(input: PdfLayoutRequest): Promise<ComputeOperation<PdfLayoutResult>> {
return this.requestJson('POST', '/v1/pdf-layout/operations', input);
}
@ -141,7 +141,7 @@ export class ComputeWorkerClient {
});
}
async getOperation<Result>(opId: string): Promise<WorkerOperation<Result> | null> {
async getOperation<Result>(opId: string): Promise<ComputeOperation<Result> | 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<Result>(opId: string, timeoutMs: number): Promise<WorkerOperation<Result>> {
private async waitForOperation<Result>(opId: string, timeoutMs: number): Promise<ComputeOperation<Result>> {
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<Result>;
if (!event?.snapshot || !isWorkerOperation(event.snapshot) || event.snapshot.opId !== opId) continue;
const event = JSON.parse(frame.data) as ComputeOperationEvent<Result>;
if (!event?.snapshot || !isComputeOperation(event.snapshot) || event.snapshot.opId !== opId) continue;
if (isTerminal(event.snapshot)) return event.snapshot;
}
}

6
src/lib/server/compute-worker/generated.ts
View file

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

22
src/lib/server/compute-worker/protocol.ts
View file

@ -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<components['schemas']['PublicOperation']['progress']>;
export type WorkerOperationStatus = components['schemas']['PublicOperation']['status'];
export type WorkerOperationSubject = components['schemas']['PublicOperation']['subject'];
export type PdfLayoutProgress = NonNullable<components['schemas']['ComputeOperation']['progress']>;
export type ComputeOperationStatus = components['schemas']['ComputeOperation']['status'];
export type ComputeOperationSubject = components['schemas']['ComputeOperation']['subject'];
export type WorkerOperation<Result = unknown> =
Omit<components['schemas']['PublicOperation'], 'result'>
export type ComputeOperation<Result = unknown> =
Omit<components['schemas']['ComputeOperation'], 'result'>
& { result?: Result };
export type WorkerOperationEvent<Result = unknown> = {
export type ComputeOperationEvent<Result = unknown> = {
eventId: number;
snapshot: WorkerOperation<Result>;
snapshot: ComputeOperation<Result>;
};
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<PdfLayoutResult> | null;
operation: ComputeOperation<PdfLayoutResult> | 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<string, unknown>;
return typeof record.opId === 'string'

10
src/lib/server/pdf-parse/operation.ts
View file

@ -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<WorkerOperation<PdfLayoutResult> | null> {
}): Promise<ComputeOperation<PdfLayoutResult> | null> {
return (await resolveCurrentPdfParse(input)).operation;
}
@ -41,7 +41,7 @@ export async function createOrReuseCurrentPdfParseOperation(input: {
documentId: string;
namespace: string | null;
forceToken?: string;
}): Promise<WorkerOperation<PdfLayoutResult>> {
}): Promise<ComputeOperation<PdfLayoutResult>> {
return getComputeWorkerClient().createPdfLayoutOperation(currentPdfOperationInput(
input.documentId,
input.namespace,
@ -49,12 +49,12 @@ export async function createOrReuseCurrentPdfParseOperation(input: {
));
}
export async function fetchPdfParseOperation(opId: string): Promise<WorkerOperation<PdfLayoutResult> | null> {
export async function fetchPdfParseOperation(opId: string): Promise<ComputeOperation<PdfLayoutResult> | null> {
return getComputeWorkerClient().getOperation<PdfLayoutResult>(opId);
}
export function isPdfParseOperationForDocument(
state: WorkerOperation<PdfLayoutResult>,
state: ComputeOperation<PdfLayoutResult>,
input: {
documentId: string;
namespace: string | null;

8
src/lib/server/pdf-parse/snapshot.ts
View file

@ -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<PdfLayoutResult>,
state: ComputeOperation<PdfLayoutResult>,
): 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<PdfLayoutResult>,
state: ComputeOperation<PdfLayoutResult>,
): PdfParseSnapshot {
const parseStatus = mapWorkerStatusToParseStatus(state.status);
return {