From 88809b826b05bb098433388148f9f8490899d023 Mon Sep 17 00:00:00 2001 From: Richard R Date: Wed, 20 May 2026 15:38:24 -0600 Subject: [PATCH] fix(compute): unify PDF layout result handling with object key support Update compute backend interfaces and job logic to consistently handle parsed PDF layout results as either direct data or S3 object key references. Refactor result types, local and worker backends, and job processing to support this pattern. Improves flexibility for large document parsing and object storage integration. --- compute/core/src/contracts.ts | 15 +++++++++---- compute/worker/src/server.ts | 36 ++++++++++++++++++++++++++---- src/lib/server/compute/local.ts | 2 +- src/lib/server/compute/types.ts | 6 ++++- src/lib/server/compute/worker.ts | 8 ++++++- src/lib/server/jobs/parsePdfJob.ts | 10 ++++++--- 6 files changed, 63 insertions(+), 14 deletions(-) diff --git a/compute/core/src/contracts.ts b/compute/core/src/contracts.ts index 04e1fb4..e97e672 100644 --- a/compute/core/src/contracts.ts +++ b/compute/core/src/contracts.ts @@ -36,10 +36,17 @@ export interface PdfLayoutJobRequest { documentObjectKey: string; } -export interface PdfLayoutJobResult { - parsed: ParsedPdfDocument; - timing?: WorkerJobTiming; -} +export type PdfLayoutJobResult = + | { + parsed: ParsedPdfDocument; + parsedObjectKey?: never; + timing?: WorkerJobTiming; + } + | { + parsed?: never; + parsedObjectKey: string; + timing?: WorkerJobTiming; + }; export type WorkerJobState = 'queued' | 'running' | 'succeeded' | 'failed'; diff --git a/compute/worker/src/server.ts b/compute/worker/src/server.ts index c59150e..1b44af5 100644 --- a/compute/worker/src/server.ts +++ b/compute/worker/src/server.ts @@ -33,7 +33,7 @@ import { type WorkerJobStatusResponse, type WorkerJobTiming, } from '@openreader/compute-core/contracts'; -import { GetObjectCommand, S3Client } from '@aws-sdk/client-s3'; +import { GetObjectCommand, PutObjectCommand, S3Client } from '@aws-sdk/client-s3'; const JOBS_STREAM_NAME = 'compute_jobs'; const WHISPER_JOBS_SUBJECT = 'jobs.whisper'; @@ -44,6 +44,8 @@ const JOB_STATES_BUCKET = 'job_states'; const JOB_STATES_TTL_MS = 24 * 60 * 60 * 1000; const PULL_EXPIRES_MS = 1000; const LOOP_ERROR_BACKOFF_MS = 500; +const DOCUMENT_ID_REGEX = /^[a-f0-9]{64}$/i; +const SAFE_NAMESPACE_REGEX = /^[a-zA-Z0-9._-]{1,128}$/; interface QueuedJob { jobId: string; @@ -104,6 +106,20 @@ function normalizeS3Prefix(prefix: string | undefined): string { return value ? value.replace(/^\/+|\/+$/g, '') : 'openreader'; } +function sanitizeNamespace(namespace: string | null): string | null { + if (!namespace) return null; + return SAFE_NAMESPACE_REGEX.test(namespace) ? namespace : null; +} + +function documentParsedKey(id: string, namespace: string | null, prefix: string): string { + if (!DOCUMENT_ID_REGEX.test(id)) { + throw new Error(`Invalid document id: ${id}`); + } + const ns = sanitizeNamespace(namespace); + const nsSegment = ns ? `ns/${ns}/` : ''; + return `${prefix}/documents_v1/parsed_v1/${nsSegment}${id}.json`; +} + function buildS3Client(): S3Client { const bucket = requireEnv('S3_BUCKET'); const region = requireEnv('S3_REGION'); @@ -468,6 +484,19 @@ async function main(): Promise { return toArrayBuffer(new Uint8Array(bytes)); }; + const putParsedObject = async (documentId: string, namespace: string | null, parsed: unknown): Promise => { + const key = documentParsedKey(documentId, namespace, s3Prefix); + const body = Buffer.from(JSON.stringify(parsed)); + await s3.send(new PutObjectCommand({ + Bucket: s3Bucket, + Key: key, + Body: body, + ContentType: 'application/json', + ServerSideEncryption: 'AES256', + })); + return key; + }; + if (prewarmModels) { await ensureComputeModels(); } @@ -638,7 +667,6 @@ async function main(): Promise { return { ...result, timing: { - ...(result.timing ?? {}), queueWaitMs, s3FetchMs, computeMs, @@ -667,10 +695,10 @@ async function main(): Promise { ); const computeMs = Date.now() - computeStartedAt; + const parsedObjectKey = await putParsedObject(parsed.documentId, parsed.namespace, result.parsed); return { - ...result, + parsedObjectKey, timing: { - ...(result.timing ?? {}), queueWaitMs, s3FetchMs, computeMs, diff --git a/src/lib/server/compute/local.ts b/src/lib/server/compute/local.ts index 41cea36..0315682 100644 --- a/src/lib/server/compute/local.ts +++ b/src/lib/server/compute/local.ts @@ -35,6 +35,6 @@ export class LocalComputeBackend implements ComputeBackend { if (!pdfBytes) { throw new Error('Local compute PDF layout requires pdfBytes or (documentId + namespace)'); } - return (await runPdfLayoutFromPdfBuffer({ documentId: input.documentId, pdfBytes })).parsed; + return { parsed: (await runPdfLayoutFromPdfBuffer({ documentId: input.documentId, pdfBytes })).parsed }; } } diff --git a/src/lib/server/compute/types.ts b/src/lib/server/compute/types.ts index 4a5ce41..aacbb10 100644 --- a/src/lib/server/compute/types.ts +++ b/src/lib/server/compute/types.ts @@ -21,8 +21,12 @@ export interface PdfLayoutInput { pdfBytes?: ArrayBuffer; } +export type PdfLayoutResult = + | { parsed: ParsedPdfDocument; parsedObjectKey?: never } + | { parsed?: never; parsedObjectKey: string }; + export interface ComputeBackend { mode: ComputeMode; alignWords(input: WhisperAlignInput): Promise; - parsePdfLayout(input: PdfLayoutInput): Promise; + parsePdfLayout(input: PdfLayoutInput): Promise; } diff --git a/src/lib/server/compute/worker.ts b/src/lib/server/compute/worker.ts index 79af3ba..5205e28 100644 --- a/src/lib/server/compute/worker.ts +++ b/src/lib/server/compute/worker.ts @@ -124,7 +124,13 @@ export class WorkerComputeBackend implements ComputeBackend { if (status.status !== 'succeeded' || !status.result) { throw new Error(status.error?.message || 'PDF layout worker job did not complete'); } - return status.result.parsed; + if (status.result.parsedObjectKey) { + return { parsedObjectKey: status.result.parsedObjectKey }; + } + if (status.result.parsed) { + return { parsed: status.result.parsed }; + } + throw new Error('PDF layout worker job completed without parsed output'); }); } diff --git a/src/lib/server/jobs/parsePdfJob.ts b/src/lib/server/jobs/parsePdfJob.ts index c64cea9..593d06d 100644 --- a/src/lib/server/jobs/parsePdfJob.ts +++ b/src/lib/server/jobs/parsePdfJob.ts @@ -29,14 +29,18 @@ export async function parsePdfJob(input: ParsePdfJobInput): Promise { .where(and(eq(documents.id, input.documentId), eq(documents.userId, input.userId))); const compute = await getCompute(); - const parsed = await compute.parsePdfLayout({ + const layout = await compute.parsePdfLayout({ documentId: input.documentId, namespace: input.namespace, documentObjectKey: documentKey(input.documentId, input.namespace), }); - const parsedJson = Buffer.from(JSON.stringify(parsed)); - const parsedJsonKey = await putParsedDocumentBlob(input.documentId, parsedJson, input.namespace); + let parsedJsonKey = layout.parsedObjectKey ?? null; + if (!parsedJsonKey) { + if (!layout.parsed) throw new Error('Compute backend did not return parsed result'); + const parsedJson = Buffer.from(JSON.stringify(layout.parsed)); + parsedJsonKey = await putParsedDocumentBlob(input.documentId, parsedJson, input.namespace); + } const cleared = await clearTtsSegmentCache({ userId: input.userId,