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,