From 95348f7faf18775d68f6b6c7a53ce8c9f5280712 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Tue, 14 Jul 2026 12:58:44 -0400 Subject: [PATCH 1/7] feat(workspaces): move file storage and extraction to R2 --- config/r2/cors.production.json | 12 + config/r2/cors.staging.json | 12 + containers/liteparse/Dockerfile | 4 + containers/liteparse/server.mjs | 198 ++++++++- docs/configuration/deployments.mdx | 18 + package.json | 3 +- pnpm-lock.yaml | 35 +- .../conversion/container-file-conversion.ts | 49 ++- .../conversion/image-file-converter.ts | 52 ++- .../conversion/office-pdf-converter.ts | 37 +- .../extraction/liteparse-projection.ts | 48 ++- .../extraction/page-markdown-projection.ts | 19 +- .../extraction/providers/firecrawl.ts | 49 +-- .../providers/liteparse-response.ts | 26 +- .../extraction/providers/liteparse.test.ts | 35 +- .../extraction/providers/liteparse.ts | 93 +++-- .../extraction/providers/llama-parse.ts | 45 +- .../providers/workers-ai-to-markdown.ts | 13 +- .../request-workspace-file-extraction.ts | 2 +- src/features/workspaces/extraction/types.ts | 2 +- .../workspace-file-extraction-workflow.ts | 75 ++-- .../workspace-file-preview-projection.ts | 109 +++++ .../extraction/workspace-file-source.ts | 20 + .../workspace-page-projection.test.ts | 179 ++++++++ .../extraction/workspace-page-projection.ts | 290 +++++++++++++ .../workspaces/files/pdfium-server.ts | 33 -- .../files/workspace-file-object-keys.ts | 29 ++ .../files/workspace-file-preview.constants.ts | 1 - .../files/workspace-file-preview.ts | 160 ++----- .../files/workspace-file-processor.ts | 45 ++ .../workspaces/files/workspace-file-upload.ts | 60 ++- .../workspace-file-r2-migration-route.ts | 72 ++++ .../kernel/workspace-kernel-access.ts | 14 +- .../kernel/workspace-kernel-file-commands.ts | 380 +++++++++-------- .../workspace-kernel-file-migrations.ts | 163 ++++++++ .../kernel/workspace-kernel-files.ts | 4 - .../kernel/workspace-kernel-item-commands.ts | 11 +- .../kernel/workspace-kernel-rows.ts | 20 + .../kernel/workspace-kernel-schema.ts | 20 + .../kernel/workspace-kernel-types.ts | 51 ++- .../workspaces/kernel/workspace-kernel.ts | 48 +-- .../workspaces/model/workspace-file/limits.ts | 6 +- .../workspaces/operations/read-items.ts | 20 +- .../operations/read-page-selection.test.ts | 21 +- .../operations/read-page-selection.ts | 94 ----- .../operations/workspace-tool-definitions.ts | 2 +- .../operations/workspace-tool-schemas.ts | 2 +- .../workspaces/read-page-selection.ts | 47 +++ .../upload/pdf-upload-validation.test.ts | 48 ++- .../upload/pdf-upload-validation.ts | 65 +-- .../workspace-file-direct-upload.test.ts | 123 ++++++ .../upload/workspace-file-direct-upload.ts | 155 +++++++ ...orkspace-file-upload-normalization.test.ts | 52 --- .../workspace-file-upload-normalization.ts | 113 ----- .../upload/workspace-file-upload-protocol.ts | 8 + .../workspace-file-upload-storage.test.ts | 138 +++++++ .../upload/workspace-file-upload-storage.ts | 159 +++++++ .../upload/workspace-upload-intake.test.ts | 43 ++ .../upload/workspace-upload-intake.ts | 18 +- src/lib/http/streaming-multipart.ts | 63 +++ src/lib/r2.test.ts | 33 ++ src/lib/r2.ts | 15 + .../v1/workspaces.$workspaceId.file-upload.ts | 391 +++++++++++------- ...aces.$workspaceId.files.$itemId.content.ts | 19 +- ...aces.$workspaceId.files.$itemId.preview.ts | 14 +- src/server.ts | 9 +- worker-configuration.d.ts | 24 +- wrangler.jsonc | 43 +- 68 files changed, 3030 insertions(+), 1231 deletions(-) create mode 100644 config/r2/cors.production.json create mode 100644 config/r2/cors.staging.json create mode 100644 src/features/workspaces/extraction/workspace-file-preview-projection.ts create mode 100644 src/features/workspaces/extraction/workspace-file-source.ts create mode 100644 src/features/workspaces/extraction/workspace-page-projection.test.ts create mode 100644 src/features/workspaces/extraction/workspace-page-projection.ts delete mode 100644 src/features/workspaces/files/pdfium-server.ts create mode 100644 src/features/workspaces/files/workspace-file-object-keys.ts create mode 100644 src/features/workspaces/files/workspace-file-processor.ts create mode 100644 src/features/workspaces/kernel/workspace-file-r2-migration-route.ts create mode 100644 src/features/workspaces/kernel/workspace-kernel-file-migrations.ts delete mode 100644 src/features/workspaces/operations/read-page-selection.ts create mode 100644 src/features/workspaces/read-page-selection.ts create mode 100644 src/features/workspaces/upload/workspace-file-direct-upload.test.ts create mode 100644 src/features/workspaces/upload/workspace-file-direct-upload.ts delete mode 100644 src/features/workspaces/upload/workspace-file-upload-normalization.test.ts delete mode 100644 src/features/workspaces/upload/workspace-file-upload-normalization.ts create mode 100644 src/features/workspaces/upload/workspace-file-upload-protocol.ts create mode 100644 src/features/workspaces/upload/workspace-file-upload-storage.test.ts create mode 100644 src/features/workspaces/upload/workspace-file-upload-storage.ts create mode 100644 src/features/workspaces/upload/workspace-upload-intake.test.ts create mode 100644 src/lib/http/streaming-multipart.ts create mode 100644 src/lib/r2.test.ts create mode 100644 src/lib/r2.ts diff --git a/config/r2/cors.production.json b/config/r2/cors.production.json new file mode 100644 index 00000000..4eed2493 --- /dev/null +++ b/config/r2/cors.production.json @@ -0,0 +1,12 @@ +{ + "rules": [ + { + "allowed": { + "origins": ["https://thinkex.app"], + "methods": ["PUT"], + "headers": ["content-type"] + }, + "maxAgeSeconds": 3600 + } + ] +} diff --git a/config/r2/cors.staging.json b/config/r2/cors.staging.json new file mode 100644 index 00000000..5b56bbd3 --- /dev/null +++ b/config/r2/cors.staging.json @@ -0,0 +1,12 @@ +{ + "rules": [ + { + "allowed": { + "origins": ["https://staging.thinkex.app", "http://localhost:3000"], + "methods": ["PUT"], + "headers": ["content-type"] + }, + "maxAgeSeconds": 3600 + } + ] +} diff --git a/containers/liteparse/Dockerfile b/containers/liteparse/Dockerfile index 03188325..a2c123bc 100644 --- a/containers/liteparse/Dockerfile +++ b/containers/liteparse/Dockerfile @@ -1,5 +1,9 @@ FROM node:22-bookworm-slim +RUN apt-get update \ + && apt-get install -y --no-install-recommends libvips-tools poppler-utils \ + && rm -rf /var/lib/apt/lists/* + WORKDIR /app COPY package.json package-lock.json ./ diff --git a/containers/liteparse/server.mjs b/containers/liteparse/server.mjs index 7b361e6e..fba95aa2 100644 --- a/containers/liteparse/server.mjs +++ b/containers/liteparse/server.mjs @@ -1,5 +1,13 @@ import { LiteParse } from "@llamaindex/liteparse"; +import { execFile } from "node:child_process"; +import { createWriteStream } from "node:fs"; +import { mkdtemp, readFile, rm } from "node:fs/promises"; import { createServer } from "node:http"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { pipeline } from "node:stream/promises"; +import { Transform } from "node:stream"; +import { promisify } from "node:util"; const port = 8080; const parser = new LiteParse({ @@ -10,6 +18,8 @@ const parser = new LiteParse({ quiet: true, }); const parseTimeoutMs = 90_000; +const maxInputBytes = 200 * 1024 * 1024; +const execFileAsync = promisify(execFile); createServer(async (request, response) => { const startedAt = Date.now(); @@ -20,38 +30,52 @@ createServer(async (request, response) => { let errorMessage = null; try { - if (request.method !== "POST" || request.url !== "/parse") { - status = 404; - return sendJson(response, status, { error: "Not found." }); + if (request.method === "POST" && request.url === "/validate/pdf") { + inputBytes = await validatePdfRequest(request); + status = 204; + response.writeHead(status); + return response.end(); + } + + if ( + request.method === "POST" && + (request.url === "/preview/pdf" || request.url === "/preview/image") + ) { + const preview = await generatePreviewRequest( + request, + request.url === "/preview/pdf" ? "pdf" : "image", + ); + inputBytes = preview.inputBytes; + status = 200; + response.writeHead(status, { + "content-length": String(preview.bytes.byteLength), + "content-type": "image/webp", + }); + return response.end(preview.bytes); } - const webRequest = new Request(`http://localhost${request.url}`, { - body: request, - duplex: "half", - headers: request.headers, - method: request.method, - }); - const formData = await webRequest.formData(); - const file = formData.get("file"); - - if (!(file instanceof File)) { - status = 400; - return sendJson(response, status, { error: "A PDF file is required." }); + if (request.method !== "POST" || request.url !== "/parse/pdf") { + status = 404; + return sendJson(response, status, { error: "Not found." }); } - const bytes = new Uint8Array(await file.arrayBuffer()); + const bytes = await readPdfRequestBytes(request); inputBytes = bytes.byteLength; const result = await withTimeout(parser.parse(bytes), parseTimeoutMs); - const pages = result.pages.map((page) => ({ - markdown: page.markdown, - pageNumber: page.pageNum, - })); - pageCount = pages.length; + pageCount = result.pages.length; status = 200; - return sendJson(response, status, { pages }); + response.writeHead(status, { "content-type": "application/x-ndjson; charset=utf-8" }); + for (const page of result.pages) { + response.write(`${JSON.stringify({ markdown: page.markdown, pageNumber: page.pageNum })}\n`); + } + return response.end(); } catch (error) { errorType = error instanceof Error ? error.name : "UnknownError"; errorMessage = error instanceof Error ? error.message : String(error); + if (error instanceof PdfValidationError) { + status = error.status; + return sendJson(response, status, { code: error.code, error: error.message }); + } return sendJson(response, status, { error: "PDF parsing failed." }); } finally { console.info( @@ -71,6 +95,136 @@ createServer(async (request, response) => { } }).listen(port); +async function validatePdfRequest(request) { + return withRequestFile(request, "thinkex-pdf-", async ({ filePath, sizeBytes }) => { + try { + await execFileAsync("pdfinfo", [filePath], { timeout: 30_000 }); + } catch (error) { + const message = getProcessErrorMessage(error); + + if (/password|encrypted/i.test(message)) { + throw new PdfValidationError( + 422, + "PASSWORD_PROTECTED_PDF", + "Password-protected PDFs are not supported.", + ); + } + + throw new PdfValidationError(422, "INVALID_PDF", "PDF is damaged or invalid."); + } + + return sizeBytes; + }); +} + +async function readPdfRequestBytes(request) { + return withRequestFile( + request, + "thinkex-parse-", + async ({ filePath }) => new Uint8Array(await readFile(filePath)), + ); +} + +async function generatePreviewRequest(request, kind) { + return withRequestFile(request, "thinkex-preview-", async ({ filePath, sizeBytes, tempDir }) => { + const previewInputPath = + kind === "pdf" ? await renderPdfPreviewInput(filePath, tempDir) : filePath; + const outputPath = join(tempDir, "preview.webp"); + + await execFileAsync( + "vips", + [ + "thumbnail", + previewInputPath, + `${outputPath}[Q=80,strip]`, + "480", + "--height", + "1200", + "--size", + "down", + "--auto-rotate", + ], + { timeout: 60_000 }, + ); + + return { bytes: await readFile(outputPath), inputBytes: sizeBytes }; + }); +} + +async function renderPdfPreviewInput(filePath, tempDir) { + const outputPrefix = join(tempDir, "page"); + + await execFileAsync( + "pdftoppm", + [ + "-f", + "1", + "-l", + "1", + "-singlefile", + "-scale-to-x", + "480", + "-scale-to-y", + "-1", + "-png", + filePath, + outputPrefix, + ], + { timeout: 60_000 }, + ); + + return `${outputPrefix}.png`; +} + +async function withRequestFile(request, prefix, run) { + const tempDir = await mkdtemp(join(tmpdir(), prefix)); + const filePath = join(tempDir, "upload.bin"); + let sizeBytes = 0; + const limiter = new Transform({ + transform(chunk, _encoding, callback) { + sizeBytes += chunk.byteLength; + + if (sizeBytes > maxInputBytes) { + callback(new PdfValidationError(413, "UPLOAD_TOO_LARGE", "File is too large.")); + return; + } + + callback(null, chunk); + }, + }); + + try { + await pipeline(request, limiter, createWriteStream(filePath)); + + if (sizeBytes === 0) { + throw new PdfValidationError(422, "INVALID_FILE", "File is empty."); + } + + return await run({ filePath, sizeBytes, tempDir }); + } finally { + await rm(tempDir, { force: true, recursive: true }); + } +} + +function getProcessErrorMessage(error) { + if (typeof error !== "object" || error === null) { + return ""; + } + + const stderr = "stderr" in error ? String(error.stderr).trim() : ""; + const stdout = "stdout" in error ? String(error.stdout).trim() : ""; + const message = error instanceof Error ? error.message.trim() : ""; + return stderr || stdout || message; +} + +class PdfValidationError extends Error { + constructor(status, code, message) { + super(message); + this.status = status; + this.code = code; + } +} + function sendJson(response, status, body) { response.writeHead(status, { "content-type": "application/json" }); response.end(JSON.stringify(body)); diff --git a/docs/configuration/deployments.mdx b/docs/configuration/deployments.mdx index 4ed4c4eb..f540482c 100644 --- a/docs/configuration/deployments.mdx +++ b/docs/configuration/deployments.mdx @@ -53,3 +53,21 @@ vp run deploy:worker:staging ## Secrets Runtime and deploy secrets are synced into GitHub Actions secrets from Infisical. Production and staging workflows require Cloudflare credentials, PostHog variables, and the runtime secrets declared by the Worker configuration. + +Direct workspace uploads additionally require bucket-scoped R2 Object Read & Write credentials in `R2_ACCESS_KEY_ID` and `R2_SECRET_ACCESS_KEY`, plus `R2_ACCOUNT_ID`. Configure browser PUT access from the checked-in policies: + +```bash +wrangler r2 bucket cors set thinkex-workspace-kernel-files-staging --file config/r2/cors.staging.json +wrangler r2 bucket cors set thinkex-workspace-kernel-files --file config/r2/cors.production.json +``` + +After inspecting existing lifecycle rules, add a one-day expiry for abandoned direct-upload objects in each bucket: + +```bash +wrangler r2 bucket lifecycle add thinkex-workspace-kernel-files-staging expire-abandoned-workspace-file-uploads workspace_file_uploads/ --expire-days 1 +wrangler r2 bucket lifecycle add thinkex-workspace-kernel-files expire-abandoned-workspace-file-uploads workspace_file_uploads/ --expire-days 1 +``` + + + CORS `set` replaces the bucket's current CORS policy. Reconcile the checked-in policy with any existing rules before applying it. + diff --git a/package.json b/package.json index 50751877..c1801855 100644 --- a/package.json +++ b/package.json @@ -40,7 +40,6 @@ "dependencies": { "@ai-sdk/react": "^3.0.210", "@base-ui/react": "^1.6.0", - "@cf-wasm/photon": "^0.3.6", "@cloudflare/codemode": "^0.4.3", "@cloudflare/containers": "^0.3.7", "@cloudflare/sandbox": "0.12.3", @@ -108,6 +107,7 @@ "agents": "^0.17.4", "ai": "^6.0.208", "autumn-js": "^1.2.35", + "aws4fetch": "^1.0.20", "better-auth": "^1.6.23", "class-variance-authority": "^0.7.1", "clsx": "^2.1.1", @@ -142,7 +142,6 @@ "tw-animate-css": "^1.3.6", "unified": "^11.0.5", "unist-util-visit": "^5.1.0", - "unpdf": "1.6.2", "y-partyserver": "^2.2.0", "y-protocols": "^1.0.7", "yjs": "^13.6.31", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index caa08fa8..b38e5c00 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -23,9 +23,6 @@ importers: '@base-ui/react': specifier: ^1.6.0 version: 1.6.0(@date-fns/tz@1.5.0)(@types/react@19.2.17)(date-fns@4.3.0)(react-dom@19.2.7(react@19.2.7))(react@19.2.7) - '@cf-wasm/photon': - specifier: ^0.3.6 - version: 0.3.6 '@cloudflare/codemode': specifier: ^0.4.3 version: 0.4.3(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.4.3))(@tanstack/ai@0.22.0(@opentelemetry/api@1.9.1))(ai@6.0.208(zod@4.4.3))(zod@4.4.3) @@ -227,6 +224,9 @@ importers: autumn-js: specifier: ^1.2.35 version: 1.2.35(better-auth@1.6.23(@cloudflare/workers-types@4.20260702.1)(@opentelemetry/api@1.9.1)(@tanstack/react-start@1.168.27(@voidzero-dev/vite-plus-core@0.2.2(@types/node@26.1.0)(esbuild@0.28.1)(jiti@2.7.0)(tsx@4.23.0)(typescript@6.0.3)(yaml@2.9.0))(react-dom@19.2.7(react@19.2.7))(react@19.2.7))(drizzle-kit@0.31.10)(drizzle-orm@0.45.2(@cloudflare/workers-types@4.20260702.1)(@opentelemetry/api@1.9.1)(@types/pg@8.20.0)(kysely@0.28.17)(pg@8.22.0)(sql.js@1.14.1))(pg@8.22.0)(react-dom@19.2.7(react@19.2.7))(react@19.2.7)(solid-js@1.9.12)(svelte@5.56.3(@typescript-eslint/types@8.61.1))(vitest@4.1.9)(vue@3.5.38(typescript@6.0.3)))(better-call@1.3.7(zod@4.4.3))(express@5.2.1)(hono@4.12.27)(react@19.2.7) + aws4fetch: + specifier: ^1.0.20 + version: 1.0.20 better-auth: specifier: ^1.6.23 version: 1.6.23(@cloudflare/workers-types@4.20260702.1)(@opentelemetry/api@1.9.1)(@tanstack/react-start@1.168.27(@voidzero-dev/vite-plus-core@0.2.2(@types/node@26.1.0)(esbuild@0.28.1)(jiti@2.7.0)(tsx@4.23.0)(typescript@6.0.3)(yaml@2.9.0))(react-dom@19.2.7(react@19.2.7))(react@19.2.7))(drizzle-kit@0.31.10)(drizzle-orm@0.45.2(@cloudflare/workers-types@4.20260702.1)(@opentelemetry/api@1.9.1)(@types/pg@8.20.0)(kysely@0.28.17)(pg@8.22.0)(sql.js@1.14.1))(pg@8.22.0)(react-dom@19.2.7(react@19.2.7))(react@19.2.7)(solid-js@1.9.12)(svelte@5.56.3(@typescript-eslint/types@8.61.1))(vitest@4.1.9)(vue@3.5.38(typescript@6.0.3)) @@ -329,9 +329,6 @@ importers: unist-util-visit: specifier: ^5.1.0 version: 5.1.0 - unpdf: - specifier: 1.6.2 - version: 1.6.2 y-partyserver: specifier: ^2.2.0 version: 2.2.0(@cloudflare/workers-types@4.20260702.1)(partyserver@0.5.8(@cloudflare/workers-types@4.20260702.1))(yjs@13.6.31) @@ -353,7 +350,7 @@ importers: version: 8.0.2(@babel/core@7.29.7) '@cloudflare/vitest-pool-workers': specifier: ^0.18.0 - version: 0.18.0(@cloudflare/workers-types@4.20260702.1)(@vitest/runner@4.1.9)(@vitest/snapshot@4.1.9)(vitest@4.1.9) + version: 0.18.0(@cloudflare/workers-types@4.20260702.1)(@vitest/runner@4.1.9)(@vitest/snapshot@4.1.9)(vitest@4.1.9(@opentelemetry/api@1.9.1)(@types/node@26.1.0)(@vitest/browser-preview@4.1.9)(@voidzero-dev/vite-plus-core@0.2.2(@types/node@26.1.0)(esbuild@0.28.1)(jiti@2.7.0)(tsx@4.23.0)(typescript@6.0.3)(yaml@2.9.0))(jsdom@28.1.0(@noble/hashes@2.2.0))(msw@2.14.6(@types/node@26.1.0)(typescript@6.0.3))) '@cloudflare/workers-types': specifier: ^4.20260702.1 version: 4.20260702.1 @@ -855,12 +852,6 @@ packages: resolution: {integrity: sha512-ctxtJ/eA+t+6q2++vj5j7FYX3nRu311q1wfYH3xjlLOsczhlhxAg2FWNUXhpGvAw3BWo1xBcvOV6/YLc2r5FJw==} hasBin: true - '@cf-wasm/internals@0.1.2': - resolution: {integrity: sha512-9d/I3JFv1IpQFYOrIw5RQShQPyuZRw9DyeBylF39Uj/MH7my8+EzKDPpUCHiqZ5O7tZS/A6zwP08egpeI5NBhA==} - - '@cf-wasm/photon@0.3.6': - resolution: {integrity: sha512-LfLfJ10+Z+DrohjQTSBkgmqxp6d4gvlGuwFkb5c4RHJGAMyVJxrAQWGrLunWE6TvZssIkT5fhKAjQZYhFSYzqg==} - '@cfworker/json-schema@4.1.1': resolution: {integrity: sha512-gAmrUZSGtKc3AiBL71iNWxDsyUC5uMaKKGdvzYsBoTW/xi42JQHl7eKV2OYzCUqvc+D2RCcf7EXY2iCyFIk6og==} @@ -8297,14 +8288,6 @@ packages: resolution: {integrity: sha512-gptHNQghINnc/vTGIk0SOFGFNXw7JVrlRUtConJRlvaw6DuX0wO5Jeko9sWrMBhh+PsYAZ7oXAiOnf/UKogyiw==} engines: {node: '>= 10.0.0'} - unpdf@1.6.2: - resolution: {integrity: sha512-zQ80ySoPuPHOsvIoRp/nJyQt8TOUoTh1+WBCGcBvlddQNgKDLRwm0AY3x8Q35I7+kIiRSgqMx+Ma2pl9McIp7A==} - peerDependencies: - '@napi-rs/canvas': ^0.1.69 - peerDependenciesMeta: - '@napi-rs/canvas': - optional: true - unpipe@1.0.0: resolution: {integrity: sha512-pjy2bYhSsufwWlKwPc+l3cN7+wuJlK6uz0YdJEOlQDbl6jo/YlPi4mb8agUkVC8BF7V8NuzeyPNqRksA3hztKQ==} engines: {node: '>= 0.8'} @@ -9183,12 +9166,6 @@ snapshots: css-tree: 3.2.1 optional: true - '@cf-wasm/internals@0.1.2': {} - - '@cf-wasm/photon@0.3.6': - dependencies: - '@cf-wasm/internals': 0.1.2 - '@cfworker/json-schema@4.1.1': {} '@chevrotain/types@11.1.2': {} @@ -9280,7 +9257,7 @@ snapshots: - utf-8-validate - workerd - '@cloudflare/vitest-pool-workers@0.18.0(@cloudflare/workers-types@4.20260702.1)(@vitest/runner@4.1.9)(@vitest/snapshot@4.1.9)(vitest@4.1.9)': + '@cloudflare/vitest-pool-workers@0.18.0(@cloudflare/workers-types@4.20260702.1)(@vitest/runner@4.1.9)(@vitest/snapshot@4.1.9)(vitest@4.1.9(@opentelemetry/api@1.9.1)(@types/node@26.1.0)(@vitest/browser-preview@4.1.9)(@voidzero-dev/vite-plus-core@0.2.2(@types/node@26.1.0)(esbuild@0.28.1)(jiti@2.7.0)(tsx@4.23.0)(typescript@6.0.3)(yaml@2.9.0))(jsdom@28.1.0(@noble/hashes@2.2.0))(msw@2.14.6(@types/node@26.1.0)(typescript@6.0.3)))': dependencies: '@vitest/runner': 4.1.9 '@vitest/snapshot': 4.1.9 @@ -16609,8 +16586,6 @@ snapshots: universalify@2.0.1: {} - unpdf@1.6.2: {} - unpipe@1.0.0: {} unplugin@3.0.0: diff --git a/src/features/workspaces/conversion/container-file-conversion.ts b/src/features/workspaces/conversion/container-file-conversion.ts index e65d13d0..b83da391 100644 --- a/src/features/workspaces/conversion/container-file-conversion.ts +++ b/src/features/workspaces/conversion/container-file-conversion.ts @@ -14,33 +14,65 @@ export async function convertFileWithContainer(input: { formFieldName: string; url: string; }): Promise { + const response = await convertFileStreamWithContainer({ + ...input, + body: input.file.stream(), + contentType: input.file.type || "application/octet-stream", + sizeBytes: input.file.size, + }); + const bytes = await response.arrayBuffer(); + + if (bytes.byteLength === 0) { + throw input.error(input.emptyMessage); + } + + return bytes; +} + +export async function convertFileStreamWithContainer(input: { + container: FileConversionContainer; + body: ReadableStream; + contentType: string; + emptyMessage: string; + error: (message: string) => Error; + fileName: string; + formFieldName: string; + sizeBytes: number; + url: string; +}): Promise { await input.container.startAndWaitForPorts({ cancellationOptions: { portReadyTimeoutMS: 60_000, }, }); - const formData = new FormData(); - formData.set(input.formFieldName, input.file, input.fileName); + const multipart = createStreamingMultipartFile({ + body: input.body, + contentType: input.contentType, + fileName: input.fileName, + formFieldName: input.formFieldName, + sizeBytes: input.sizeBytes, + }); const response = await input.container.fetch( new Request(input.url, { - body: formData, + body: multipart.body, + duplex: "half", + headers: { "content-type": multipart.contentType }, method: "POST", - }), + } as RequestInit & { duplex: "half" }), ); + await multipart.done; if (!response.ok) { throw input.error(await getConversionErrorMessage(response)); } - const bytes = await response.arrayBuffer(); - - if (bytes.byteLength === 0) { + if (!response.body) { throw input.error(input.emptyMessage); } - return bytes; + return response; } async function getConversionErrorMessage(response: Response) { @@ -50,3 +82,4 @@ async function getConversionErrorMessage(response: Response) { return message ? `${fallback} ${message}` : fallback; } +import { createStreamingMultipartFile } from "#/lib/http/streaming-multipart"; diff --git a/src/features/workspaces/conversion/image-file-converter.ts b/src/features/workspaces/conversion/image-file-converter.ts index d540c975..8d3c4021 100644 --- a/src/features/workspaces/conversion/image-file-converter.ts +++ b/src/features/workspaces/conversion/image-file-converter.ts @@ -1,6 +1,9 @@ import { Container, getRandom } from "@cloudflare/containers"; -import { convertFileWithContainer } from "#/features/workspaces/conversion/container-file-conversion"; +import { + convertFileStreamWithContainer, + convertFileWithContainer, +} from "#/features/workspaces/conversion/container-file-conversion"; import { WorkspaceFileConversionError } from "#/features/workspaces/conversion/errors"; const imageConverterPort = 8080; @@ -34,11 +37,50 @@ export class ImageFileConversionError extends WorkspaceFileConversionError { } } -export async function convertImageFileToJpeg( +export async function convertImageStreamToJpeg( env: Cloudflare.Env, - input: ConvertImageFileToJpegInput, -): Promise { - return convertImageFileToJpegAtPath(env, input, imageConverterPath); + input: { + body: ReadableStream; + contentType: string; + fileName: string; + sizeBytes: number; + }, +) { + return convertImageStreamToJpegAtPath(env, input, imageConverterPath); +} + +export async function convertImageStreamToChatJpeg( + env: Cloudflare.Env, + input: { + body: ReadableStream; + contentType: string; + fileName: string; + sizeBytes: number; + }, +) { + return convertImageStreamToJpegAtPath(env, input, chatImageConverterPath); +} + +async function convertImageStreamToJpegAtPath( + env: Cloudflare.Env, + input: { + body: ReadableStream; + contentType: string; + fileName: string; + sizeBytes: number; + }, + path: string, +) { + const converter = await getRandom(env.IMAGE_FILE_CONVERTER, converterPoolSize); + + return convertFileStreamWithContainer({ + ...input, + container: converter, + emptyMessage: "Image conversion returned an empty JPEG.", + error: (message) => new ImageFileConversionError(message), + formFieldName: "file", + url: `http://image-file-converter${path}`, + }); } export async function convertImageFileToChatJpeg( diff --git a/src/features/workspaces/conversion/office-pdf-converter.ts b/src/features/workspaces/conversion/office-pdf-converter.ts index 9bc2c919..b16e1f7b 100644 --- a/src/features/workspaces/conversion/office-pdf-converter.ts +++ b/src/features/workspaces/conversion/office-pdf-converter.ts @@ -1,11 +1,10 @@ import { Container, getRandom } from "@cloudflare/containers"; -import { convertFileWithContainer } from "#/features/workspaces/conversion/container-file-conversion"; +import { convertFileStreamWithContainer } from "#/features/workspaces/conversion/container-file-conversion"; import { WorkspaceFileConversionError } from "#/features/workspaces/conversion/errors"; const gotenbergPort = 3000; const gotenbergLibreOfficeConvertPath = "/forms/libreoffice/convert"; -const pdfContentType = "application/pdf"; const converterPoolSize = 2; export class OfficePdfConverter extends Container { @@ -20,17 +19,6 @@ export class OfficePdfConverter extends Container { }; } -export interface ConvertOfficeFileToPdfInput { - file: File; - fileName: string; -} - -export interface ConvertOfficeFileToPdfResult { - bytes: ArrayBuffer; - contentType: typeof pdfContentType; - sizeBytes: number; -} - export class OfficePdfConversionError extends WorkspaceFileConversionError { constructor(message: string) { super(message, "Unable to convert this file to PDF right now."); @@ -38,24 +26,23 @@ export class OfficePdfConversionError extends WorkspaceFileConversionError { } } -export async function convertOfficeFileToPdf( +export async function convertOfficeStreamToPdf( env: Cloudflare.Env, - input: ConvertOfficeFileToPdfInput, -): Promise { + input: { + body: ReadableStream; + contentType: string; + fileName: string; + sizeBytes: number; + }, +) { const converter = await getRandom(env.OFFICE_PDF_CONVERTER, converterPoolSize); - const bytes = await convertFileWithContainer({ + + return convertFileStreamWithContainer({ + ...input, container: converter, emptyMessage: "Office file conversion returned an empty PDF.", error: (message) => new OfficePdfConversionError(message), - file: input.file, - fileName: input.fileName, formFieldName: "files", url: `http://office-pdf-converter${gotenbergLibreOfficeConvertPath}`, }); - - return { - bytes, - contentType: pdfContentType, - sizeBytes: bytes.byteLength, - }; } diff --git a/src/features/workspaces/extraction/liteparse-projection.ts b/src/features/workspaces/extraction/liteparse-projection.ts index 5eed35d2..605e6e0c 100644 --- a/src/features/workspaces/extraction/liteparse-projection.ts +++ b/src/features/workspaces/extraction/liteparse-projection.ts @@ -1,22 +1,19 @@ import type { WorkflowStep } from "cloudflare:workers"; -import { sha256Base64Url } from "#/features/workspaces/extraction/binary"; -import { - joinMarkdownProjectionPages, - parseMarkdownPagesProjection, - serializeMarkdownPagesProjection, -} from "#/features/workspaces/extraction/page-markdown-projection"; import { extractPdfWithLiteParse } from "#/features/workspaces/extraction/providers/liteparse"; import type { LiteParseStageOutcome, WorkspaceFileExtractionWorkflowParams, } from "#/features/workspaces/extraction/types"; +import { getWorkspaceFileSourceObject } from "#/features/workspaces/extraction/workspace-file-source"; +import { writeWorkspacePageProjection } from "#/features/workspaces/extraction/workspace-page-projection"; import { getWorkspaceKernelFromEnv } from "#/features/workspaces/kernel/workspace-kernel-access"; export async function publishLiteParseProjection( env: Cloudflare.Env, step: WorkflowStep, params: WorkspaceFileExtractionWorkflowParams, + runId: string, ): Promise { if (params.assetKind !== "pdf") { return { durationMs: 0, outcome: "skipped" }; @@ -33,27 +30,38 @@ export async function publishLiteParseProjection( }, async () => { const kernel = await getWorkspaceKernelFromEnv(env, params.workspaceId); - const source = await kernel.readFileContent({ itemId: params.itemId }); - const sourceHash = await sha256Base64Url(source.bytes); - const pages = await extractPdfWithLiteParse(env, { - bytes: source.bytes, - fileName: source.fileName, + const { object, source } = await getWorkspaceFileSourceObject({ + env, + itemId: params.itemId, + kernel, + }); + const projection = await writeWorkspacePageProjection({ + bucket: env.WORKSPACE_KERNEL_FILES, + itemId: params.itemId, + pages: extractPdfWithLiteParse(env, { + body: object.body, + fileName: source.fileName, + sizeBytes: source.sizeBytes, + }), + provider: "liteparse", + providerMode: "fast", + runId, + sourceHash: object.etag, + tier: "fast", + workspaceId: params.workspaceId, }); - const content = serializeMarkdownPagesProjection(pages); - const projectionPages = parseMarkdownPagesProjection(content); - const markdownLength = joinMarkdownProjectionPages(projectionPages).length; await kernel.upsertFileProjection({ itemId: params.itemId, format: "pages", status: "ready", - content, + objectKey: projection.manifestObjectKey, provider: "liteparse", providerMode: "fast", - sourceHash, + sourceHash: object.etag, metadataJson: { - markdownLength, - pageCount: projectionPages.length, + markdownLength: projection.manifest.markdownLength, + pageCount: projection.manifest.pageCount, provisional: true, }, actorUserId: params.actorUserId, @@ -61,9 +69,9 @@ export async function publishLiteParseProjection( return { durationMs: Date.now() - startedAt, - markdownLength, + markdownLength: projection.manifest.markdownLength, outcome: "success" as const, - pageCount: projectionPages.length, + pageCount: projection.manifest.pageCount, }; }, ); diff --git a/src/features/workspaces/extraction/page-markdown-projection.ts b/src/features/workspaces/extraction/page-markdown-projection.ts index 0b0d1bdd..b5aa37dd 100644 --- a/src/features/workspaces/extraction/page-markdown-projection.ts +++ b/src/features/workspaces/extraction/page-markdown-projection.ts @@ -8,17 +8,7 @@ export function createSingleMarkdownProjectionPage(markdown: string): MarkdownPr return trimmed ? [{ pageNumber: 1, markdown: trimmed }] : []; } -export function serializeMarkdownPagesProjection(pages: readonly MarkdownProjectionPage[]) { - return JSON.stringify( - pages - .map((page) => ({ - pageNumber: page.pageNumber, - markdown: page.markdown.trim(), - })) - .filter((page) => Number.isInteger(page.pageNumber) && page.pageNumber > 0 && page.markdown), - ); -} - +/** Parses the pre-R2 projection shape during lazy migration only. */ export function parseMarkdownPagesProjection(content: string | null) { if (!content?.trim()) { return []; @@ -39,13 +29,6 @@ export function parseMarkdownPagesProjection(content: string | null) { } } -export function joinMarkdownProjectionPages(pages: readonly MarkdownProjectionPage[]) { - return pages - .map((page) => `## Page ${page.pageNumber}\n\n${page.markdown}`) - .join("\n\n") - .trim(); -} - function parseMarkdownProjectionPage(value: unknown): MarkdownProjectionPage | null { if (!value || typeof value !== "object" || Array.isArray(value)) { return null; diff --git a/src/features/workspaces/extraction/providers/firecrawl.ts b/src/features/workspaces/extraction/providers/firecrawl.ts index 103e6d8c..e58d4de2 100644 --- a/src/features/workspaces/extraction/providers/firecrawl.ts +++ b/src/features/workspaces/extraction/providers/firecrawl.ts @@ -1,4 +1,3 @@ -import { toArrayBuffer } from "#/features/workspaces/extraction/binary"; import type { FirecrawlPdfMode, MarkdownExtractionInput, @@ -13,6 +12,7 @@ import { getRecordValue, getStringValue, } from "#/integrations/firecrawl/client"; +import { createStreamingMultipartFile } from "#/lib/http/streaming-multipart"; const firecrawlParseTimeoutMs = 300_000; @@ -23,30 +23,33 @@ export function createFirecrawlPdfExtractionProvider( id: "firecrawl", async extract(input) { const mode = normalizeFirecrawlMode(input.mode); - const formData = new FormData(); - formData.set( - "options", - JSON.stringify({ - formats: ["markdown"], - parsers: [{ type: "pdf", mode }], - timeout: firecrawlParseTimeoutMs, - }), - ); - formData.set( - "file", - new File([toArrayBuffer(input.bytes)], input.fileName, { - type: input.contentType || "application/pdf", - }), - ); - - const responseJson = await firecrawlJsonRequest({ - env, - path: "/v2/parse", - operation: "Firecrawl PDF parsing", - method: "POST", - body: formData, + const multipart = createStreamingMultipartFile({ + body: input.body, + contentType: input.contentType || "application/pdf", + fields: { + options: JSON.stringify({ + formats: ["markdown"], + parsers: [{ type: "pdf", mode }], + timeout: firecrawlParseTimeoutMs, + }), + }, + fileName: input.fileName, + formFieldName: "file", + sizeBytes: input.sizeBytes, }); + const [responseJson] = await Promise.all([ + firecrawlJsonRequest({ + env, + path: "/v2/parse", + operation: "Firecrawl PDF parsing", + method: "POST", + headers: { "content-type": multipart.contentType }, + body: multipart.body, + }), + multipart.done, + ]); + const markdown = getFirecrawlMarkdown(responseJson); if (!markdown) { diff --git a/src/features/workspaces/extraction/providers/liteparse-response.ts b/src/features/workspaces/extraction/providers/liteparse-response.ts index 11e3cb5d..aa5e0f95 100644 --- a/src/features/workspaces/extraction/providers/liteparse-response.ts +++ b/src/features/workspaces/extraction/providers/liteparse-response.ts @@ -1,23 +1,17 @@ import type { MarkdownProjectionPage } from "#/features/workspaces/extraction/page-markdown-projection"; -export function parseLiteParsePages(payload: unknown): MarkdownProjectionPage[] { - if (!isRecord(payload) || !Array.isArray(payload.pages)) { - throw new Error("LiteParse returned an invalid response."); +export function parseLiteParsePage(payload: unknown): MarkdownProjectionPage { + if ( + !isRecord(payload) || + typeof payload.pageNumber !== "number" || + !Number.isInteger(payload.pageNumber) || + payload.pageNumber < 1 || + typeof payload.markdown !== "string" + ) { + throw new Error("LiteParse returned an invalid page."); } - return payload.pages.map((page) => { - if ( - !isRecord(page) || - typeof page.pageNumber !== "number" || - !Number.isInteger(page.pageNumber) || - page.pageNumber < 1 || - typeof page.markdown !== "string" - ) { - throw new Error("LiteParse returned an invalid page."); - } - - return { pageNumber: page.pageNumber, markdown: page.markdown.trim() }; - }); + return { pageNumber: payload.pageNumber, markdown: payload.markdown.trim() }; } function isRecord(value: unknown): value is Record { diff --git a/src/features/workspaces/extraction/providers/liteparse.test.ts b/src/features/workspaces/extraction/providers/liteparse.test.ts index 4cbc0f26..9e50131a 100644 --- a/src/features/workspaces/extraction/providers/liteparse.test.ts +++ b/src/features/workspaces/extraction/providers/liteparse.test.ts @@ -1,36 +1,29 @@ import { describe, expect, it } from "vitest"; -import { parseLiteParsePages } from "#/features/workspaces/extraction/providers/liteparse-response"; +import { parseLiteParsePage } from "#/features/workspaces/extraction/providers/liteparse-response"; describe("LiteParse response parsing", () => { - it("accepts canonical per-page Markdown", () => { - expect( - parseLiteParsePages({ - pages: [ - { markdown: " # First ", pageNumber: 1 }, - { markdown: "\nSecond\n", pageNumber: 2 }, - ], - }), - ).toEqual([ - { markdown: "# First", pageNumber: 1 }, - { markdown: "Second", pageNumber: 2 }, - ]); + it("accepts one canonical page record", () => { + expect(parseLiteParsePage({ markdown: " # First ", pageNumber: 1 })).toEqual({ + markdown: "# First", + pageNumber: 1, + }); }); - it("keeps blank pages for the projection serializer to filter", () => { - expect(parseLiteParsePages({ pages: [{ markdown: " \n\t", pageNumber: 1 }] })).toEqual([ - { markdown: "", pageNumber: 1 }, - ]); + it("keeps blank pages so PDF page numbering remains stable", () => { + expect(parseLiteParsePage({ markdown: " \n\t", pageNumber: 1 })).toEqual({ + markdown: "", + pageNumber: 1, + }); }); it.each([ undefined, {}, { markdown: "document markdown without page data" }, - { pages: "invalid" }, - { pages: [{ markdown: "Text", pageNumber: 0 }] }, - { pages: [{ markdown: 123, pageNumber: 1 }] }, + { markdown: "Text", pageNumber: 0 }, + { markdown: 123, pageNumber: 1 }, ])("rejects malformed container responses", (payload) => { - expect(() => parseLiteParsePages(payload)).toThrow("LiteParse returned an invalid"); + expect(() => parseLiteParsePage(payload)).toThrow("LiteParse returned an invalid"); }); }); diff --git a/src/features/workspaces/extraction/providers/liteparse.ts b/src/features/workspaces/extraction/providers/liteparse.ts index 476783e4..6e79f3d3 100644 --- a/src/features/workspaces/extraction/providers/liteparse.ts +++ b/src/features/workspaces/extraction/providers/liteparse.ts @@ -1,50 +1,69 @@ -import { Container, getRandom } from "@cloudflare/containers"; - import type { MarkdownProjectionPage } from "#/features/workspaces/extraction/page-markdown-projection"; -import { parseLiteParsePages } from "#/features/workspaces/extraction/providers/liteparse-response"; - -const liteParsePort = 8080; -const liteParsePoolSize = 2; - -export class LiteParsePdfExtractor extends Container { - defaultPort = liteParsePort; - requiredPorts = [liteParsePort]; - sleepAfter = "5m"; - enableInternet = false; -} +import { parseLiteParsePage } from "#/features/workspaces/extraction/providers/liteparse-response"; +import { requestWorkspaceFileProcessor } from "#/features/workspaces/files/workspace-file-processor"; -export async function extractPdfWithLiteParse( +export async function* extractPdfWithLiteParse( env: Cloudflare.Env, - input: { bytes: Uint8Array; fileName: string }, -): Promise { - const extractor = await getRandom(env.LITEPARSE_PDF_EXTRACTOR, liteParsePoolSize); - await extractor.startAndWaitForPorts({ - cancellationOptions: { portReadyTimeoutMS: 60_000 }, + input: { + body: ReadableStream; + fileName: string; + sizeBytes: number; + }, +): AsyncGenerator { + const response = await requestWorkspaceFileProcessor(env, { + body: input.body, + contentType: "application/pdf", + fileName: input.fileName, + path: "/parse/pdf", + sizeBytes: input.sizeBytes, }); - const formData = new FormData(); - formData.set( - "file", - new File([input.bytes as BlobPart], input.fileName, { type: "application/pdf" }), - input.fileName, - ); - const response = await extractor.fetch( - new Request("http://liteparse-pdf-extractor/parse", { - body: formData, - method: "POST", - }), - ); - if (!response.ok) { throw new Error(`LiteParse failed with status ${response.status}.`); } - const payload: unknown = await response.json(); - const pages = parseLiteParsePages(payload); + if (!response.body) { + throw new Error("LiteParse completed without a response body."); + } - if (!pages.some((page) => page.markdown.trim().length > 0)) { - throw new Error("LiteParse did not extract usable page Markdown."); + for await (const line of readNdjsonLines(response.body)) { + let payload: unknown; + try { + payload = JSON.parse(line); + } catch { + throw new Error("LiteParse returned invalid NDJSON."); + } + yield parseLiteParsePage(payload); } +} + +async function* readNdjsonLines(body: ReadableStream): AsyncGenerator { + const reader = body.getReader(); + const decoder = new TextDecoder(); + let buffer = ""; + + try { + while (true) { + const { done, value } = await reader.read(); + buffer += decoder.decode(value, { stream: !done }); + const lines = buffer.split("\n"); + buffer = lines.pop() ?? ""; - return pages; + for (const line of lines) { + if (line.trim()) { + yield line; + } + } + + if (done) { + break; + } + } + + if (buffer.trim()) { + yield buffer; + } + } finally { + reader.releaseLock(); + } } diff --git a/src/features/workspaces/extraction/providers/llama-parse.ts b/src/features/workspaces/extraction/providers/llama-parse.ts index f3435008..0771a717 100644 --- a/src/features/workspaces/extraction/providers/llama-parse.ts +++ b/src/features/workspaces/extraction/providers/llama-parse.ts @@ -1,4 +1,3 @@ -import { toArrayBuffer } from "#/features/workspaces/extraction/binary"; import { createSingleMarkdownProjectionPage, type MarkdownProjectionPage, @@ -16,6 +15,7 @@ import { getStringValue, llamaCloudJsonRequest, } from "#/integrations/llamaparse/client"; +import { createStreamingMultipartFile } from "#/lib/http/streaming-multipart"; const llamaParsePollIntervalMs = 2_000; const llamaParseMaxPollMs = 300_000; @@ -63,22 +63,25 @@ function normalizeLlamaParseTier(mode: MarkdownExtractionInput["mode"]): LlamaPa } async function uploadLlamaParseFile(env: Env, input: MarkdownExtractionInput) { - const formData = new FormData(); - formData.set("purpose", "parse"); - formData.set( - "file", - new File([toArrayBuffer(input.bytes)], input.fileName, { - type: input.contentType || "application/pdf", - }), - ); - - const responseJson = await llamaCloudJsonRequest({ - env, - path: "/api/v1/beta/files", - operation: "LlamaParse file upload", - method: "POST", - body: formData, + const multipart = createStreamingMultipartFile({ + body: input.body, + contentType: input.contentType || "application/pdf", + fields: { purpose: "parse" }, + fileName: input.fileName, + formFieldName: "file", + sizeBytes: input.sizeBytes, }); + const [responseJson] = await Promise.all([ + llamaCloudJsonRequest({ + env, + path: "/api/v1/beta/files", + operation: "LlamaParse file upload", + method: "POST", + headers: { "content-type": multipart.contentType }, + body: multipart.body, + }), + multipart.done, + ]); const fileId = getStringValue(responseJson, "id"); if (!fileId) { @@ -136,14 +139,18 @@ async function pollLlamaParseJob(env: Env, jobId: string) { while (Date.now() - startedAt < llamaParseMaxPollMs) { const responseJson = await llamaCloudJsonRequest({ env, - path: `/api/v2/parse/${jobId}?expand=markdown,metadata,job_metadata`, - operation: "LlamaParse job result", + path: `/api/v2/parse/${jobId}`, + operation: "LlamaParse job status", }); const job = getRecordValue(responseJson, "job") ?? responseJson; const status = getStringValue(job, "status"); if (status === "COMPLETED") { - return responseJson; + return await llamaCloudJsonRequest({ + env, + path: `/api/v2/parse/${jobId}?expand=markdown,metadata,job_metadata`, + operation: "LlamaParse job result", + }); } if (status === "FAILED" || status === "CANCELLED") { diff --git a/src/features/workspaces/extraction/providers/workers-ai-to-markdown.ts b/src/features/workspaces/extraction/providers/workers-ai-to-markdown.ts index 24f8f6f0..8b699a83 100644 --- a/src/features/workspaces/extraction/providers/workers-ai-to-markdown.ts +++ b/src/features/workspaces/extraction/providers/workers-ai-to-markdown.ts @@ -1,4 +1,4 @@ -import { toArrayBuffer } from "#/features/workspaces/extraction/binary"; +import { convertImageStreamToChatJpeg } from "#/features/workspaces/conversion/image-file-converter"; import type { MarkdownExtractionProvider, MarkdownExtractionResult, @@ -9,11 +9,18 @@ export function createWorkersAiToMarkdownProvider(env: Env): MarkdownExtractionP return { id: "workers_ai_to_markdown", async extract(input) { + const conversion = await convertImageStreamToChatJpeg(env, { + body: input.body, + contentType: input.contentType, + fileName: input.fileName, + sizeBytes: input.sizeBytes, + }); + const bytes = await conversion.arrayBuffer(); const result = await env.AI.toMarkdown( { name: input.fileName, - blob: new Blob([toArrayBuffer(input.bytes)], { - type: input.contentType || "application/octet-stream", + blob: new Blob([bytes], { + type: "image/jpeg", }), }, { diff --git a/src/features/workspaces/extraction/request-workspace-file-extraction.ts b/src/features/workspaces/extraction/request-workspace-file-extraction.ts index 1be91d6a..2718a559 100644 --- a/src/features/workspaces/extraction/request-workspace-file-extraction.ts +++ b/src/features/workspaces/extraction/request-workspace-file-extraction.ts @@ -89,7 +89,7 @@ async function getWorkspaceFileExtractionWorkflowId(input: { assetKind: WorkspaceFileAssetKind; }) { const digest = await sha256Base64UrlText( - `${input.workspaceId}:${input.itemId}:${input.assetKind}-extraction:v1`, + `${input.workspaceId}:${input.itemId}:${input.assetKind}-extraction:v2`, ); return `${input.assetKind}-${digest.slice(0, 48)}`; diff --git a/src/features/workspaces/extraction/types.ts b/src/features/workspaces/extraction/types.ts index cf755bac..5dbc9b86 100644 --- a/src/features/workspaces/extraction/types.ts +++ b/src/features/workspaces/extraction/types.ts @@ -36,7 +36,7 @@ export type LiteParseStageOutcome = export interface MarkdownExtractionInput { workspaceId: string; itemId: string; - bytes: Uint8Array; + body: ReadableStream; fileName: string; contentType: string; sizeBytes: number; diff --git a/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts b/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts index 1c8b30f4..7f4c4035 100644 --- a/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts +++ b/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts @@ -1,19 +1,16 @@ import { WorkflowEntrypoint, type WorkflowEvent, type WorkflowStep } from "cloudflare:workers"; -import { sha256Base64Url } from "#/features/workspaces/extraction/binary"; import { publishLiteParseProjection } from "#/features/workspaces/extraction/liteparse-projection"; +import { publishWorkspaceFilePreview } from "#/features/workspaces/extraction/workspace-file-preview-projection"; import { recordWorkspaceFileExtractionOutcome } from "#/features/workspaces/extraction/workspace-file-extraction-observability"; -import { - joinMarkdownProjectionPages, - parseMarkdownPagesProjection, - serializeMarkdownPagesProjection, -} from "#/features/workspaces/extraction/page-markdown-projection"; import { createMarkdownExtractionProvider } from "#/features/workspaces/extraction/providers/index"; import type { MarkdownExtractionProviderId, MarkdownExtractionProviderMode, WorkspaceFileExtractionWorkflowParams, } from "#/features/workspaces/extraction/types"; +import { getWorkspaceFileSourceObject } from "#/features/workspaces/extraction/workspace-file-source"; +import { writeWorkspacePageProjection } from "#/features/workspaces/extraction/workspace-page-projection"; import { getWorkspaceKernelFromEnv } from "#/features/workspaces/kernel/workspace-kernel-access"; import { getWorkspaceUploadFamily } from "#/features/workspaces/model/workspace-file"; @@ -26,7 +23,6 @@ export class WorkspaceFileExtractionWorkflow extends WorkflowEntrypoint< step: WorkflowStep, ) { const params = assertWorkflowParams(event.payload); - const artifactKey = getExtractionArtifactKey(event.instanceId); const schedule = (task: Promise) => this.ctx.waitUntil(task); await step.do("mark extraction processing", async () => { @@ -41,7 +37,8 @@ export class WorkspaceFileExtractionWorkflow extends WorkflowEntrypoint< return { status: "processing" }; }); - const liteParse = await publishLiteParseProjection(this.env, step, params); + const liteParse = await publishLiteParseProjection(this.env, step, params, event.instanceId); + await publishWorkspaceFilePreview(this.env, step, params); const enhancementStartedAt = Date.now(); let extraction: StagedPageExtractionResult; let result: { @@ -64,37 +61,46 @@ export class WorkspaceFileExtractionWorkflow extends WorkflowEntrypoint< }, async (): Promise => { const kernel = await getWorkspaceKernelFromEnv(this.env, params.workspaceId); - const source = await kernel.readFileContent({ + const { object, source } = await getWorkspaceFileSourceObject({ + env: this.env, itemId: params.itemId, + kernel, }); - const sourceHash = await sha256Base64Url(source.bytes); const route = getWorkspaceUploadFamily(params.assetKind).extractionRoute; const provider = createMarkdownExtractionProvider(route.provider, this.env); const extraction = await provider.extract({ workspaceId: params.workspaceId, itemId: params.itemId, - bytes: source.bytes, + body: object.body, fileName: source.fileName, contentType: source.contentType, sizeBytes: source.sizeBytes, - sourceHash, + sourceHash: object.etag, mode: route.mode, }); - const pagesJson = serializeMarkdownPagesProjection(extraction.pages); - - await this.env.WORKSPACE_KERNEL_FILES.put(artifactKey, pagesJson, { - httpMetadata: { contentType: "application/json" }, + const projection = await writeWorkspacePageProjection({ + bucket: this.env.WORKSPACE_KERNEL_FILES, + itemId: params.itemId, + metadata: extraction.metadata, + pages: extraction.pages, + provider: extraction.provider, + providerMode: extraction.providerMode, + runId: event.instanceId, + sourceHash: object.etag, + tier: "enhanced", + workspaceId: params.workspaceId, }); return { - artifactKey, + manifestObjectKey: projection.manifestObjectKey, + markdownLength: projection.manifest.markdownLength, provider: extraction.provider, providerMode: extraction.providerMode, metadata: extraction.metadata, - pageCount: extraction.pages.length, + pageCount: projection.manifest.pageCount, routeReason: route.reason, - sourceHash, + sourceHash: object.etag, }; }, ); @@ -111,26 +117,18 @@ export class WorkspaceFileExtractionWorkflow extends WorkflowEntrypoint< }, async () => { const kernel = await getWorkspaceKernelFromEnv(this.env, params.workspaceId); - const artifact = await this.env.WORKSPACE_KERNEL_FILES.get(extraction.artifactKey); - - if (!artifact) { - throw new Error("Staged page extraction artifact was not found."); - } - - const pagesJson = await artifact.text(); - const pages = parseMarkdownPagesProjection(pagesJson); const metadataJson = { ...extraction.metadata, routeReason: extraction.routeReason, pageCount: extraction.pageCount, - markdownLength: joinMarkdownProjectionPages(pages).length, + markdownLength: extraction.markdownLength, }; await kernel.upsertFileProjection({ itemId: params.itemId, format: "pages", status: "ready", - content: pagesJson, + objectKey: extraction.manifestObjectKey, provider: extraction.provider, providerMode: extraction.providerMode, sourceHash: extraction.sourceHash, @@ -214,18 +212,6 @@ export class WorkspaceFileExtractionWorkflow extends WorkflowEntrypoint< throw error; } - await step.do("delete staged extraction artifact", async () => { - try { - await this.env.WORKSPACE_KERNEL_FILES.delete(extraction.artifactKey); - return { deleted: extraction.artifactKey }; - } catch (error) { - return { - deleted: null, - errorType: error instanceof Error ? error.name : "UnknownError", - }; - } - }); - await step.do("record extraction outcome", async () => { recordWorkspaceFileExtractionOutcome({ durationMs: Date.now() - event.timestamp.getTime(), @@ -252,7 +238,8 @@ export class WorkspaceFileExtractionWorkflow extends WorkflowEntrypoint< } interface StagedPageExtractionResult { - artifactKey: string; + manifestObjectKey: string; + markdownLength: number; provider: MarkdownExtractionProviderId; providerMode: MarkdownExtractionProviderMode; metadata: Record; @@ -277,10 +264,6 @@ function assertWorkflowParams( }; } -function getExtractionArtifactKey(instanceId: string) { - return `workflow-artifacts/page-extraction/${instanceId}/projection.json`; -} - function getErrorMessage(error: unknown) { return error instanceof Error ? error.message : String(error); } diff --git a/src/features/workspaces/extraction/workspace-file-preview-projection.ts b/src/features/workspaces/extraction/workspace-file-preview-projection.ts new file mode 100644 index 00000000..fdf24e23 --- /dev/null +++ b/src/features/workspaces/extraction/workspace-file-preview-projection.ts @@ -0,0 +1,109 @@ +import type { WorkflowStep } from "cloudflare:workers"; + +import { getWorkspaceFileSourceObject } from "#/features/workspaces/extraction/workspace-file-source"; +import type { WorkspaceFileExtractionWorkflowParams } from "#/features/workspaces/extraction/types"; +import { getWorkspaceFilePreviewObjectKey } from "#/features/workspaces/files/workspace-file-object-keys"; +import { + createWorkspaceFilePreview, + WORKSPACE_FILE_PREVIEW_CONTENT_TYPE, +} from "#/features/workspaces/files/workspace-file-preview"; +import { getWorkspaceKernelFromEnv } from "#/features/workspaces/kernel/workspace-kernel-access"; +import { getWorkspaceUploadFamily } from "#/features/workspaces/model/workspace-file"; +import { recordOperationalOutcome } from "#/integrations/observability/operational-events"; + +export async function publishWorkspaceFilePreview( + env: Cloudflare.Env, + step: WorkflowStep, + params: WorkspaceFileExtractionWorkflowParams, +) { + return step.do( + "publish workspace file preview", + { retries: { limit: 1, delay: "5 seconds", backoff: "constant" }, timeout: "2 minutes" }, + async () => generateWorkspaceFilePreview(env, params), + ); +} + +async function generateWorkspaceFilePreview( + env: Cloudflare.Env, + params: WorkspaceFileExtractionWorkflowParams, +) { + const startedAt = Date.now(); + const descriptor = getWorkspaceUploadFamily(params.assetKind); + + if (!descriptor.previewGenerator) { + return { outcome: "skipped" as const }; + } + + const kernel = await getWorkspaceKernelFromEnv(env, params.workspaceId); + const objectKey = getWorkspaceFilePreviewObjectKey({ + workspaceId: params.workspaceId, + itemId: params.itemId, + }); + let failure: unknown; + let inputBytes = 0; + let outputBytes = 0; + + try { + const { object: sourceObject, source } = await getWorkspaceFileSourceObject({ + env, + itemId: params.itemId, + kernel, + }); + inputBytes = sourceObject.size; + const preview = await createWorkspaceFilePreview(env, { + assetKind: params.assetKind, + body: sourceObject.body, + contentType: source.contentType, + sizeBytes: source.sizeBytes, + }); + const object = await env.WORKSPACE_KERNEL_FILES.put(objectKey, preview.body, { + httpMetadata: { contentType: WORKSPACE_FILE_PREVIEW_CONTENT_TYPE }, + }); + + if (!object) { + throw new Error("Workspace preview could not be stored."); + } + outputBytes = object.size; + + await kernel.upsertFileProjection({ + itemId: params.itemId, + format: "preview", + status: "ready", + objectKey, + sourceHash: sourceObject.etag, + metadataJson: { + contentType: WORKSPACE_FILE_PREVIEW_CONTENT_TYPE, + sizeBytes: object.size, + }, + actorUserId: params.actorUserId, + }); + + return { outcome: "success" as const }; + } catch (error) { + failure = error; + await env.WORKSPACE_KERNEL_FILES.delete(objectKey); + await kernel.upsertFileProjection({ + itemId: params.itemId, + format: "preview", + status: "failed", + errorMessage: error instanceof Error ? error.message : String(error), + actorUserId: params.actorUserId, + }); + return { outcome: "error" as const }; + } finally { + recordOperationalOutcome({ + distinctId: params.actorUserId ?? undefined, + error: failure, + event: "workspace_file_preview", + fields: { + asset_kind: params.assetKind, + duration_ms: Date.now() - startedAt, + input_bytes: inputBytes, + item_id: params.itemId, + output_bytes: outputBytes, + user_id: params.actorUserId, + workspace_id: params.workspaceId, + }, + }); + } +} diff --git a/src/features/workspaces/extraction/workspace-file-source.ts b/src/features/workspaces/extraction/workspace-file-source.ts new file mode 100644 index 00000000..d6bdd7c0 --- /dev/null +++ b/src/features/workspaces/extraction/workspace-file-source.ts @@ -0,0 +1,20 @@ +import type { WorkspaceKernelClient } from "#/features/workspaces/kernel/workspace-kernel-access"; + +export async function getWorkspaceFileSourceObject(input: { + env: Cloudflare.Env; + itemId: string; + kernel: WorkspaceKernelClient; +}) { + const source = await input.kernel.getFileSource({ itemId: input.itemId }); + const object = await input.env.WORKSPACE_KERNEL_FILES.get(source.objectKey); + + if (!object) { + throw new Error("Workspace file source object was not found."); + } + + if (object.size !== source.sizeBytes) { + throw new Error("Workspace file source size does not match its metadata."); + } + + return { object, source }; +} diff --git a/src/features/workspaces/extraction/workspace-page-projection.test.ts b/src/features/workspaces/extraction/workspace-page-projection.test.ts new file mode 100644 index 00000000..ea8c0eaa --- /dev/null +++ b/src/features/workspaces/extraction/workspace-page-projection.test.ts @@ -0,0 +1,179 @@ +import { describe, expect, it } from "vitest"; + +import { + getWorkspacePageObjectKey, + readWorkspacePageProjection, + writeWorkspacePageProjection, +} from "#/features/workspaces/extraction/workspace-page-projection"; + +describe("workspace page projections", () => { + it("publishes immutable page objects and reads only selected pages", async () => { + const storage = createObjectStorage(); + const reference = await writeWorkspacePageProjection({ + bucket: storage.bucket, + itemId: "item-1", + pages: [ + { pageNumber: 1, markdown: "First" }, + { pageNumber: 2, markdown: "Second" }, + { pageNumber: 3, markdown: "Third" }, + ], + provider: "liteparse", + providerMode: "fast", + runId: "run-1", + sourceHash: "etag-1", + tier: "fast", + workspaceId: "workspace-1", + }); + + storage.readKeys.length = 0; + const result = await readWorkspacePageProjection({ + bucket: storage.bucket, + manifestObjectKey: reference.manifestObjectKey, + pages: "2-3", + }); + + expect(result).toEqual({ + content: "## Page 2\n\nSecond\n\n## Page 3\n\nThird", + pages: { requested: "2-3", returned: [2, 3], total: 3 }, + }); + const prefix = reference.manifestObjectKey.slice(0, -"manifest.json".length); + expect(storage.readKeys).toEqual([ + reference.manifestObjectKey, + getWorkspacePageObjectKey(prefix, 2), + getWorkspacePageObjectKey(prefix, 3), + ]); + }); + + it("preserves missing page numbers as blank pages", async () => { + const storage = createObjectStorage(); + const reference = await writeWorkspacePageProjection({ + bucket: storage.bucket, + itemId: "item-1", + pages: [ + { pageNumber: 1, markdown: "First" }, + { pageNumber: 3, markdown: "Third" }, + ], + provider: "llama_parse", + providerMode: "agentic", + runId: "run-1", + sourceHash: "etag-1", + tier: "enhanced", + workspaceId: "workspace-1", + }); + + await expect( + readWorkspacePageProjection({ + bucket: storage.bucket, + manifestObjectKey: reference.manifestObjectKey, + pages: "2", + }), + ).resolves.toEqual({ + content: "## Page 2", + pages: { requested: "2", returned: [2], total: 3 }, + }); + }); + + it("bounds page reads", async () => { + const storage = createObjectStorage(); + const reference = await writeWorkspacePageProjection({ + bucket: storage.bucket, + itemId: "item-1", + pages: Array.from({ length: 21 }, (_, index) => ({ + pageNumber: index + 1, + markdown: `Page ${index + 1}`, + })), + provider: "liteparse", + providerMode: "fast", + runId: "run-1", + sourceHash: "etag-1", + tier: "fast", + workspaceId: "workspace-1", + }); + + await expect( + readWorkspacePageProjection({ + bucket: storage.bucket, + manifestObjectKey: reference.manifestObjectKey, + pages: "1-21", + }), + ).rejects.toMatchObject({ code: "page_selection_too_large" }); + }); + + it("rejects oversized page numbers before writing page gaps", async () => { + const storage = createObjectStorage(); + + await expect( + writeWorkspacePageProjection({ + bucket: storage.bucket, + itemId: "item-1", + pages: [{ pageNumber: 2_001, markdown: "Too far" }], + provider: "liteparse", + providerMode: "fast", + runId: "run-1", + sourceHash: "etag-1", + tier: "fast", + workspaceId: "workspace-1", + }), + ).rejects.toThrow("Extraction exceeds the 2000-page limit."); + expect(storage.values.size).toBe(0); + }); + + it("removes partial artifacts when publication fails", async () => { + const storage = createObjectStorage(); + + await expect( + writeWorkspacePageProjection({ + bucket: storage.bucket, + itemId: "item-1", + pages: [ + { pageNumber: 1, markdown: "First" }, + { pageNumber: 1, markdown: "Duplicate" }, + ], + provider: "liteparse", + providerMode: "fast", + runId: "run-1", + sourceHash: "etag-1", + tier: "fast", + workspaceId: "workspace-1", + }), + ).rejects.toThrow("Extracted pages must be ordered"); + expect(storage.values.size).toBe(0); + }); +}); + +function createObjectStorage() { + const values = new Map(); + const readKeys: string[] = []; + const bucket = { + async delete(keys: string | string[]) { + for (const key of Array.isArray(keys) ? keys : [keys]) { + values.delete(key); + } + }, + async get(key: string) { + readKeys.push(key); + const value = values.get(key); + if (value === undefined) { + return null; + } + return { + key, + size: new TextEncoder().encode(value).byteLength, + text: async () => value, + json: async () => JSON.parse(value) as unknown, + }; + }, + async put(key: string, value: string) { + values.set(key, value); + return { key, size: new TextEncoder().encode(value).byteLength }; + }, + async list(input: { prefix?: string }) { + const objects = Array.from(values.keys()) + .filter((key) => key.startsWith(input.prefix ?? "")) + .map((key) => ({ key })); + return { objects, truncated: false }; + }, + } as R2Bucket; + + return { bucket, readKeys, values }; +} diff --git a/src/features/workspaces/extraction/workspace-page-projection.ts b/src/features/workspaces/extraction/workspace-page-projection.ts new file mode 100644 index 00000000..245fc978 --- /dev/null +++ b/src/features/workspaces/extraction/workspace-page-projection.ts @@ -0,0 +1,290 @@ +import { jsonValueSchema, type JsonValue } from "#/features/workspaces/contracts"; +import type { MarkdownProjectionPage } from "#/features/workspaces/extraction/page-markdown-projection"; +import { getWorkspaceFileItemObjectPrefix } from "#/features/workspaces/files/workspace-file-object-keys"; +import { + parseWorkspacePageRange, + WorkspacePageSelectionError, + type WorkspaceReadPages, +} from "#/features/workspaces/read-page-selection"; +import { deleteR2Prefix } from "#/lib/r2"; + +const projectionSchemaVersion = 1; +const pageNumberWidth = 6; +const pageWriteConcurrency = 8; +const maxProjectionPages = 2_000; +const maxPageMarkdownBytes = 1024 * 1024; +const maxPageReadCount = 20; +const maxPageReadBytes = 2 * 1024 * 1024; + +export interface WorkspacePageProjectionManifest { + createdAt: string; + itemId: string; + markdownBytes: number; + markdownLength: number; + metadata: Record; + pageCount: number; + provider: string; + providerMode: string; + runId: string; + schemaVersion: typeof projectionSchemaVersion; + sourceHash: string; + workspaceId: string; +} + +export interface WorkspacePageProjectionReference { + manifestObjectKey: string; + manifest: WorkspacePageProjectionManifest; +} + +export async function writeWorkspacePageProjection(input: { + bucket: R2Bucket; + itemId: string; + metadata?: Record; + pages: AsyncIterable | Iterable; + provider: string; + providerMode: string; + runId: string; + sourceHash: string; + tier: "enhanced" | "fast" | "legacy"; + workspaceId: string; +}): Promise { + const prefix = getWorkspacePageProjectionPrefix(input); + const encoder = new TextEncoder(); + const writes: Promise[] = []; + let lastPageNumber = 0; + let markdownBytes = 0; + let markdownLength = 0; + let usablePageCount = 0; + + try { + for await (const rawPage of input.pages) { + const page = normalizeProjectionPage(rawPage); + if (page.pageNumber > maxProjectionPages) { + throw new Error(`Extraction exceeds the ${maxProjectionPages}-page limit.`); + } + if (page.pageNumber <= lastPageNumber) { + throw new Error("Extracted pages must be ordered by unique, increasing page number."); + } + + for (let pageNumber = lastPageNumber + 1; pageNumber < page.pageNumber; pageNumber += 1) { + await schedulePageWrite(input.bucket, writes, prefix, pageNumber, ""); + } + + const pageBytes = encoder.encode(page.markdown).byteLength; + if (pageBytes > maxPageMarkdownBytes) { + throw new Error(`Extracted page ${page.pageNumber} exceeds the page size limit.`); + } + + await schedulePageWrite(input.bucket, writes, prefix, page.pageNumber, page.markdown); + lastPageNumber = page.pageNumber; + markdownBytes += pageBytes; + markdownLength += page.markdown.length; + if (page.markdown.length > 0) { + usablePageCount += 1; + } + } + + await flushPageWrites(writes); + if (lastPageNumber === 0 || usablePageCount === 0) { + throw new Error("Extraction did not produce usable page Markdown."); + } + + const manifest: WorkspacePageProjectionManifest = { + createdAt: new Date().toISOString(), + itemId: input.itemId, + markdownBytes, + markdownLength, + metadata: input.metadata ?? {}, + pageCount: lastPageNumber, + provider: input.provider, + providerMode: input.providerMode, + runId: input.runId, + schemaVersion: projectionSchemaVersion, + sourceHash: input.sourceHash, + workspaceId: input.workspaceId, + }; + const manifestObjectKey = `${prefix}manifest.json`; + await input.bucket.put(manifestObjectKey, JSON.stringify(manifest), { + httpMetadata: { contentType: "application/json" }, + }); + + return { manifest, manifestObjectKey }; + } catch (error) { + await flushPageWrites(writes).catch(() => undefined); + await deleteR2Prefix(input.bucket, prefix).catch(() => undefined); + throw error; + } +} + +export async function readWorkspacePageProjection(input: { + bucket: R2Bucket; + manifestObjectKey: string; + pages?: string; +}): Promise<{ content: string; pages: WorkspaceReadPages }> { + const manifest = await readWorkspacePageProjectionManifest(input.bucket, input.manifestObjectKey); + const requested = input.pages?.trim() || "1"; + const selectedPageNumbers = parseWorkspacePageRange(requested, manifest.pageCount); + + if (selectedPageNumbers.length > maxPageReadCount) { + throw new WorkspacePageSelectionError("page_selection_too_large"); + } + + const prefix = getManifestPrefix(input.manifestObjectKey); + const objects = await Promise.all( + selectedPageNumbers.map(async (pageNumber) => { + const object = await input.bucket.get(getWorkspacePageObjectKey(prefix, pageNumber)); + if (!object) { + throw new Error(`Extracted page ${pageNumber} was not found.`); + } + return { object, pageNumber }; + }), + ); + const totalBytes = objects.reduce((total, entry) => total + entry.object.size, 0); + if (totalBytes > maxPageReadBytes) { + throw new WorkspacePageSelectionError("page_selection_too_large"); + } + + const pages = await Promise.all( + objects.map(async ({ object, pageNumber }) => ({ + markdown: (await object.text()).trim(), + pageNumber, + })), + ); + + return { + content: pages + .map((page) => `## Page ${page.pageNumber}\n\n${page.markdown}`.trimEnd()) + .join("\n\n"), + pages: { + requested, + returned: selectedPageNumbers, + total: manifest.pageCount, + }, + }; +} + +export async function readWorkspacePageProjectionManifest( + bucket: R2Bucket, + manifestObjectKey: string, +): Promise { + const object = await bucket.get(manifestObjectKey); + if (!object) { + throw new Error("Workspace page projection manifest was not found."); + } + + return parseWorkspacePageProjectionManifest(await object.json()); +} + +export function getWorkspacePageProjectionPrefix(input: { + itemId: string; + runId: string; + tier: "enhanced" | "fast" | "legacy"; + workspaceId: string; +}) { + return `${getWorkspaceFileItemObjectPrefix(input)}extractions/${encodePathPart(input.runId)}/${input.tier}/`; +} + +export function getWorkspacePageObjectKey(prefix: string, pageNumber: number) { + return `${prefix}pages/${String(pageNumber).padStart(pageNumberWidth, "0")}.md`; +} + +function parseWorkspacePageProjectionManifest(value: unknown): WorkspacePageProjectionManifest { + if (!isRecord(value)) { + throw new Error("Workspace page projection manifest is invalid."); + } + if ( + value.schemaVersion !== projectionSchemaVersion || + typeof value.workspaceId !== "string" || + typeof value.itemId !== "string" || + typeof value.runId !== "string" || + typeof value.sourceHash !== "string" || + typeof value.provider !== "string" || + typeof value.providerMode !== "string" || + typeof value.pageCount !== "number" || + !Number.isInteger(value.pageCount) || + value.pageCount < 1 || + typeof value.markdownLength !== "number" || + typeof value.markdownBytes !== "number" || + typeof value.createdAt !== "string" || + !isRecord(value.metadata) + ) { + throw new Error("Workspace page projection manifest is invalid."); + } + + const metadata = jsonValueSchema.parse(value.metadata); + if (!isJsonObject(metadata)) { + throw new Error("Workspace page projection manifest is invalid."); + } + + return { + createdAt: value.createdAt, + itemId: value.itemId, + markdownBytes: value.markdownBytes, + markdownLength: value.markdownLength, + metadata, + pageCount: value.pageCount, + provider: value.provider, + providerMode: value.providerMode, + runId: value.runId, + schemaVersion: projectionSchemaVersion, + sourceHash: value.sourceHash, + workspaceId: value.workspaceId, + }; +} + +async function schedulePageWrite( + bucket: R2Bucket, + writes: Promise[], + prefix: string, + pageNumber: number, + markdown: string, +) { + const write = bucket + .put(getWorkspacePageObjectKey(prefix, pageNumber), markdown, { + httpMetadata: { contentType: "text/markdown; charset=utf-8" }, + }) + .then(() => undefined); + writes.push(write); + + if (writes.length >= pageWriteConcurrency) { + await flushPageWrites(writes); + } +} + +async function flushPageWrites(writes: Promise[]) { + const results = await Promise.allSettled(writes.splice(0)); + const failure = results.find((result) => result.status === "rejected"); + + if (failure?.status === "rejected") { + throw failure.reason; + } +} + +function normalizeProjectionPage(page: MarkdownProjectionPage): MarkdownProjectionPage { + if (!Number.isInteger(page.pageNumber) || page.pageNumber < 1) { + throw new Error("Extracted page number is invalid."); + } + if (typeof page.markdown !== "string") { + throw new Error("Extracted page Markdown is invalid."); + } + return { pageNumber: page.pageNumber, markdown: page.markdown.trim() }; +} + +function getManifestPrefix(manifestObjectKey: string) { + if (!manifestObjectKey.endsWith("/manifest.json")) { + throw new Error("Workspace page projection manifest key is invalid."); + } + return manifestObjectKey.slice(0, -"manifest.json".length); +} + +function encodePathPart(value: string) { + return encodeURIComponent(value); +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function isJsonObject(value: JsonValue): value is { [key: string]: JsonValue } { + return typeof value === "object" && value !== null && !Array.isArray(value); +} diff --git a/src/features/workspaces/files/pdfium-server.ts b/src/features/workspaces/files/pdfium-server.ts deleted file mode 100644 index d6b2ed5d..00000000 --- a/src/features/workspaces/files/pdfium-server.ts +++ /dev/null @@ -1,33 +0,0 @@ -import { PdfiumNative } from "@embedpdf/engines/pdfium"; -import { init } from "@embedpdf/pdfium"; -import pdfiumWasmModule from "@embedpdf/pdfium/pdfium.wasm"; - -let pdfiumNativePromise: Promise | null = null; - -/** - * Server-side PDFium init for Workers / Durable Objects. - * - * Cloudflare's Vite plugin imports `.wasm` as a precompiled `WebAssembly.Module`. - * EmbedPDF's Emscripten glue accepts that via `instantiateWasm` — no fetch, no CDN, - * no broken static ArrayBuffer imports under TanStack Start SSR. - */ -export async function getPdfiumNative() { - if (!pdfiumNativePromise) { - pdfiumNativePromise = (async () => { - const pdfiumModule = await init({ - instantiateWasm: ( - imports: WebAssembly.Imports, - receiveInstance: (instance: WebAssembly.Instance, module: WebAssembly.Module) => void, - ) => { - const instance = new WebAssembly.Instance(pdfiumWasmModule, imports); - receiveInstance(instance, pdfiumWasmModule); - return instance.exports; - }, - }); - - return new PdfiumNative(pdfiumModule, { fontFallback: null }); - })(); - } - - return pdfiumNativePromise; -} diff --git a/src/features/workspaces/files/workspace-file-object-keys.ts b/src/features/workspaces/files/workspace-file-object-keys.ts new file mode 100644 index 00000000..1faf191d --- /dev/null +++ b/src/features/workspaces/files/workspace-file-object-keys.ts @@ -0,0 +1,29 @@ +const workspaceFileObjectPrefix = "workspace_file_objects"; +const workspaceFileUploadPrefix = "workspace_file_uploads"; + +function getWorkspaceFileObjectPrefix(workspaceId: string) { + return `${workspaceFileObjectPrefix}/${workspaceId}/`; +} + +export function getWorkspaceFileItemObjectPrefix(input: { workspaceId: string; itemId: string }) { + return `${getWorkspaceFileObjectPrefix(input.workspaceId)}${input.itemId}/`; +} + +export function getWorkspaceFileSourceObjectKey(input: { workspaceId: string; itemId: string }) { + return `${getWorkspaceFileItemObjectPrefix(input)}source`; +} + +export function getWorkspaceFilePreviewObjectKey(input: { workspaceId: string; itemId: string }) { + return `${getWorkspaceFileItemObjectPrefix(input)}preview.webp`; +} + +export function getWorkspaceFileUploadObjectKey(input: { itemId: string; workspaceId: string }) { + return `${workspaceFileUploadPrefix}/${input.workspaceId}/${input.itemId}/source`; +} + +export function getWorkspaceFileUploadCompletionKey(input: { + itemId: string; + workspaceId: string; +}) { + return `${workspaceFileUploadPrefix}/${input.workspaceId}/${input.itemId}/completion`; +} diff --git a/src/features/workspaces/files/workspace-file-preview.constants.ts b/src/features/workspaces/files/workspace-file-preview.constants.ts index bd7cb5b5..814e2763 100644 --- a/src/features/workspaces/files/workspace-file-preview.constants.ts +++ b/src/features/workspaces/files/workspace-file-preview.constants.ts @@ -1,2 +1 @@ -export const WORKSPACE_FILE_PREVIEW_MAX_WIDTH = 480; export const WORKSPACE_FILE_PREVIEW_CONTENT_TYPE = "image/webp"; diff --git a/src/features/workspaces/files/workspace-file-preview.ts b/src/features/workspaces/files/workspace-file-preview.ts index f99a3eb8..47337bcb 100644 --- a/src/features/workspaces/files/workspace-file-preview.ts +++ b/src/features/workspaces/files/workspace-file-preview.ts @@ -1,135 +1,27 @@ -import { PhotonImage, resize, SamplingFilter } from "@cf-wasm/photon/workerd"; -import type { ImageDataLike } from "@embedpdf/models"; - -import { getPdfiumNative } from "#/features/workspaces/files/pdfium-server"; -import { WORKSPACE_FILE_PREVIEW_MAX_WIDTH } from "#/features/workspaces/files/workspace-file-preview.constants"; -import type { - WorkspaceFilePreviewGeneratorId, - WorkspaceFileTypeDescriptor, -} from "#/features/workspaces/model/workspace-file"; - -export { - WORKSPACE_FILE_PREVIEW_CONTENT_TYPE, - WORKSPACE_FILE_PREVIEW_MAX_WIDTH, -} from "#/features/workspaces/files/workspace-file-preview.constants"; - -export interface WorkspaceFilePreviewResult { - bytes: Uint8Array; - width: number; - height: number; -} - -type PreviewGenerator = (bytes: Uint8Array) => Promise; - -const workspaceFilePreviewGenerators: Record = { - pdf_webp: generatePdfPreviewWebp, - image_webp: generateImagePreviewWebp, -}; - -export function resolveUploadPreviewGenerator( - descriptor: WorkspaceFileTypeDescriptor, -): PreviewGenerator | null { - if (!descriptor.previewGenerator) { - return null; - } - - return workspaceFilePreviewGenerators[descriptor.previewGenerator] ?? null; -} - -export async function generateImagePreviewWebp( - bytes: Uint8Array, -): Promise { - const inputImage = PhotonImage.new_from_byteslice(bytes); - - try { - return encodePreviewWebp(inputImage); - } finally { - inputImage.free(); - } -} - -export async function generatePdfPreviewWebp( - bytes: Uint8Array, -): Promise { - const native = await getPdfiumNative(); - const document = await native - .openDocumentBuffer({ id: "preview", content: toArrayBuffer(bytes) }) - .toPromise(); - const page = document.pages[0]; - - if (!page) { - await native.closeDocument(document).toPromise(); - throw new Error("PDF has no pages."); - } - - const scaleFactor = WORKSPACE_FILE_PREVIEW_MAX_WIDTH / page.size.width; - - try { - const rendered = await native - .renderPageRaw(document, page, { - scaleFactor, - withAnnotations: false, - }) - .toPromise(); - - return rgbaToPreviewWebp(rendered); - } finally { - await native.closeDocument(document).toPromise(); - } -} - -function rgbaToPreviewWebp(image: ImageDataLike): WorkspaceFilePreviewResult { - const inputImage = new PhotonImage( - new Uint8Array(image.data.buffer, image.data.byteOffset, image.data.byteLength), - image.width, - image.height, - ); - - try { - return encodePreviewWebp(inputImage); - } finally { - inputImage.free(); - } -} - -function encodePreviewWebp(inputImage: PhotonImage): WorkspaceFilePreviewResult { - const outputImage = resizePreviewImage(inputImage); - const previewBytes = outputImage.get_bytes_webp(); - const result = { - bytes: previewBytes, - width: outputImage.get_width(), - height: outputImage.get_height(), - }; - - if (outputImage !== inputImage) { - outputImage.free(); - } - - return result; -} - -function resizePreviewImage(image: PhotonImage): PhotonImage { - const width = image.get_width(); - const height = image.get_height(); - - if (width <= WORKSPACE_FILE_PREVIEW_MAX_WIDTH) { - return image; - } - - const newWidth = WORKSPACE_FILE_PREVIEW_MAX_WIDTH; - const newHeight = Math.max(1, Math.round(height * (newWidth / width))); - - return resize(image, newWidth, newHeight, SamplingFilter.Lanczos3); -} - -function toArrayBuffer(bytes: Uint8Array): ArrayBuffer { - if ( - bytes.byteOffset === 0 && - bytes.byteLength === bytes.buffer.byteLength && - bytes.buffer instanceof ArrayBuffer - ) { - return bytes.buffer; - } - - return bytes.slice().buffer; +import { requestWorkspaceFileProcessor } from "#/features/workspaces/files/workspace-file-processor"; +import type { WorkspaceFileAssetKind } from "#/features/workspaces/model/workspace-file"; + +export { WORKSPACE_FILE_PREVIEW_CONTENT_TYPE } from "#/features/workspaces/files/workspace-file-preview.constants"; + +export async function createWorkspaceFilePreview( + env: Cloudflare.Env, + input: { + assetKind: WorkspaceFileAssetKind; + body: ReadableStream; + contentType: string; + sizeBytes: number; + }, +) { + const response = await requestWorkspaceFileProcessor(env, { + body: input.body, + contentType: input.contentType, + path: input.assetKind === "pdf" ? "/preview/pdf" : "/preview/image", + sizeBytes: input.sizeBytes, + }); + + if (!response.ok || !response.body) { + throw new Error(`Workspace preview generation failed with status ${response.status}.`); + } + + return response; } diff --git a/src/features/workspaces/files/workspace-file-processor.ts b/src/features/workspaces/files/workspace-file-processor.ts new file mode 100644 index 00000000..b24128db --- /dev/null +++ b/src/features/workspaces/files/workspace-file-processor.ts @@ -0,0 +1,45 @@ +import { Container, getRandom } from "@cloudflare/containers"; + +const workspaceFileProcessorPort = 8080; +const workspaceFileProcessorPoolSize = 2; + +export class WorkspaceFileProcessor extends Container { + defaultPort = workspaceFileProcessorPort; + requiredPorts = [workspaceFileProcessorPort]; + sleepAfter = "5m"; + enableInternet = false; +} + +export async function requestWorkspaceFileProcessor( + env: Cloudflare.Env, + input: { + body: ReadableStream; + contentType: string; + fileName?: string; + path: "/parse/pdf" | "/preview/image" | "/preview/pdf" | "/validate/pdf"; + sizeBytes: number; + }, +) { + const processor = await getRandom(env.WORKSPACE_FILE_PROCESSOR, workspaceFileProcessorPoolSize); + await processor.startAndWaitForPorts({ + cancellationOptions: { portReadyTimeoutMS: 60_000 }, + }); + + const headers = new Headers({ + "content-type": input.contentType, + "x-file-size": String(input.sizeBytes), + }); + + if (input.fileName) { + headers.set("x-file-name", encodeURIComponent(input.fileName)); + } + + return processor.fetch( + new Request(`http://workspace-file-processor${input.path}`, { + body: input.body, + duplex: "half", + headers, + method: "POST", + } as RequestInit & { duplex: "half" }), + ); +} diff --git a/src/features/workspaces/files/workspace-file-upload.ts b/src/features/workspaces/files/workspace-file-upload.ts index 7c1faea4..02bcc7fb 100644 --- a/src/features/workspaces/files/workspace-file-upload.ts +++ b/src/features/workspaces/files/workspace-file-upload.ts @@ -11,6 +11,10 @@ import { partitionWorkspaceUploadSelection, uploadPlanCreatesDocument, } from "#/features/workspaces/upload/workspace-upload-intake"; +import { + type CompleteWorkspaceDirectUploadInput, + type WorkspaceDirectUploadSession, +} from "#/features/workspaces/upload/workspace-file-upload-protocol"; import { prepareWorkspaceClientMutationInput } from "#/features/workspaces/use-workspace-client-mutation-echo"; import { apiErrorSchema } from "#/lib/api/contracts"; import { getErrorMessage } from "#/lib/error-message"; @@ -76,25 +80,61 @@ export async function runWorkspaceFileUploadBatch( async function postWorkspaceFileUpload( job: WorkspaceFileUploadJob, ): Promise> { - const formData = new FormData(); + const uploadResponse = await postWorkspaceDirectUpload(job); + + if (!uploadResponse.ok) { + throw new Error(await getWorkspaceFileUploadErrorMessage(uploadResponse)); + } - formData.set("file", job.file); - formData.set("clientMutationId", job.clientMutationId); + return (await uploadResponse.json()) as WorkspaceCommandResult; +} + +async function postWorkspaceDirectUpload(job: WorkspaceFileUploadJob) { + const endpoint = `/api/v1/workspaces/${job.workspaceId}/file-upload`; + const contentType = job.file.type || "application/octet-stream"; + const session = await requestUploadJson( + `${endpoint}?action=initiate`, + { + body: JSON.stringify({ + clientMutationId: job.clientMutationId, + contentType, + fileName: job.file.name, + fileSize: job.file.size, + parentId: job.parentId, + }), + headers: { "content-type": "application/json" }, + method: "POST", + }, + ); - if (job.parentId) { - formData.set("parentId", job.parentId); + const uploadResponse = await fetch(session.uploadUrl, { + body: job.file, + headers: { "content-type": contentType }, + method: "PUT", + }); + + if (!uploadResponse.ok) { + throw new Error(`Direct file upload failed with status ${uploadResponse.status}.`); } - const uploadResponse = await fetch(`/api/v1/workspaces/${job.workspaceId}/file-upload`, { + const completeInput: CompleteWorkspaceDirectUploadInput = { + completionToken: session.completionToken, + }; + return fetch(`${endpoint}?action=complete`, { + body: JSON.stringify(completeInput), + headers: { "content-type": "application/json" }, method: "POST", - body: formData, }); +} - if (!uploadResponse.ok) { - throw new Error(await getWorkspaceFileUploadErrorMessage(uploadResponse)); +async function requestUploadJson(url: string, init: RequestInit): Promise { + const response = await fetch(url, init); + + if (!response.ok) { + throw new Error(await getWorkspaceFileUploadErrorMessage(response)); } - return (await uploadResponse.json()) as WorkspaceCommandResult; + return (await response.json()) as T; } function toUploadJob(input: { diff --git a/src/features/workspaces/kernel/workspace-file-r2-migration-route.ts b/src/features/workspaces/kernel/workspace-file-r2-migration-route.ts new file mode 100644 index 00000000..dc921dd0 --- /dev/null +++ b/src/features/workspaces/kernel/workspace-file-r2-migration-route.ts @@ -0,0 +1,72 @@ +import { getAgentByName } from "agents"; + +import { workspaceKernelAgentName } from "#/features/workspaces/agent-routes"; + +const migrationPath = "/api/internal/migrations/workspace-files-r2"; +const maximumBatchSize = 25; +const migrationConcurrency = 5; + +interface MigrationEnvironment extends Cloudflare.Env { + WORKSPACE_FILE_R2_MIGRATION_TOKEN: string; +} + +interface MigrationRequestBody { + after?: string; + limit?: number; +} + +export async function routeWorkspaceFileR2Migration(request: Request, environment: Cloudflare.Env) { + const url = new URL(request.url); + + if (url.pathname !== migrationPath) { + return null; + } + if (request.method !== "POST") { + return new Response("Method not allowed", { status: 405 }); + } + + const env = environment as MigrationEnvironment; + const authorization = request.headers.get("authorization"); + + if ( + !env.WORKSPACE_FILE_R2_MIGRATION_TOKEN || + authorization !== `Bearer ${env.WORKSPACE_FILE_R2_MIGRATION_TOKEN}` + ) { + return new Response("Unauthorized", { status: 401 }); + } + + const body = (await request.json()) as MigrationRequestBody; + const after = body.after ?? ""; + const limit = Math.min(Math.max(body.limit ?? maximumBatchSize, 1), maximumBatchSize); + const query = await env.DB.prepare("SELECT id FROM workspaces WHERE id > ? ORDER BY id LIMIT ?") + .bind(after, limit) + .all<{ id: string }>(); + const results = []; + + for (let index = 0; index < query.results.length; index += migrationConcurrency) { + const batch = query.results.slice(index, index + migrationConcurrency); + const batchResults = await Promise.all( + batch.map(async (workspace) => { + try { + const kernel = await getAgentByName(env[workspaceKernelAgentName], workspace.id); + return { + workspaceId: workspace.id, + result: await kernel.migrateLegacyFileStorage(), + }; + } catch (error) { + return { + workspaceId: workspace.id, + error: error instanceof Error ? error.message : String(error), + }; + } + }), + ); + results.push(...batchResults); + } + + return Response.json({ + done: query.results.length < limit, + nextAfter: query.results.at(-1)?.id ?? after, + results, + }); +} diff --git a/src/features/workspaces/kernel/workspace-kernel-access.ts b/src/features/workspaces/kernel/workspace-kernel-access.ts index bca43c4b..a92e2600 100644 --- a/src/features/workspaces/kernel/workspace-kernel-access.ts +++ b/src/features/workspaces/kernel/workspace-kernel-access.ts @@ -24,6 +24,7 @@ import type { ReadWorkspaceKernelFileProjectionArgs, ReadWorkspaceKernelFileProjectionResult, UpsertWorkspaceKernelFileProjectionArgs, + WorkspaceKernelFileSource, WorkspaceKernelItemRelation, WorkspaceKernelNameConflictPolicy, } from "#/features/workspaces/kernel/workspace-kernel-types"; @@ -97,12 +98,7 @@ export interface WorkspaceKernelClient { readItem(input: { itemId: string; }): Promise<{ item: WorkspaceItemSummary; content: string | null }>; - readFileContent(input: { itemId: string }): Promise<{ - bytes: Uint8Array; - contentType: string; - fileName: string; - sizeBytes: number; - }>; + getFileSource(input: { itemId: string }): Promise; readFilePreview(input: { itemId: string }): Promise; upsertFileProjection(input: UpsertWorkspaceKernelFileProjectionArgs): Promise; readFileProjection( @@ -117,7 +113,7 @@ export interface WorkspaceKernelClient { purgeForDeletion(): Promise; } -export async function readWorkspaceKernelFileContent(input: { +export async function readWorkspaceKernelFileSource(input: { workspaceId: string; userId: string; itemId: string; @@ -128,7 +124,7 @@ export async function readWorkspaceKernelFileContent(input: { await assertCanReadWorkspace(dbContext.db, input); const kernel = await getWorkspaceKernel(input.workspaceId); - return await kernel.readFileContent({ itemId: input.itemId }); + return await kernel.getFileSource({ itemId: input.itemId }); } finally { await dbContext.dispose(); } @@ -199,6 +195,7 @@ export async function createWorkspaceKernelItem( } export async function createWorkspaceFileFromUpload(input: { + id: string; workspaceId: string; userId: string; parentId?: string | null; @@ -217,6 +214,7 @@ export async function createWorkspaceFileFromUpload(input: { const kernel = await getWorkspaceKernel(input.workspaceId); return await kernel.createFileFromUpload({ + id: input.id, parentId: input.parentId ?? null, fileName: input.fileName, fileSize: input.fileSize, diff --git a/src/features/workspaces/kernel/workspace-kernel-file-commands.ts b/src/features/workspaces/kernel/workspace-kernel-file-commands.ts index e62df549..c2177de1 100644 --- a/src/features/workspaces/kernel/workspace-kernel-file-commands.ts +++ b/src/features/workspaces/kernel/workspace-kernel-file-commands.ts @@ -1,29 +1,26 @@ import type { Workspace as ShellWorkspace } from "@cloudflare/shell"; import type { JsonValue, WorkspaceItemSummary } from "#/features/workspaces/contracts"; -import { sha256Base64Url } from "#/features/workspaces/extraction/binary"; -import { - resolveUploadPreviewGenerator, - WORKSPACE_FILE_PREVIEW_CONTENT_TYPE, -} from "#/features/workspaces/files/workspace-file-preview"; +import { getWorkspaceFileItemObjectPrefix } from "#/features/workspaces/files/workspace-file-object-keys"; +import { WORKSPACE_FILE_PREVIEW_CONTENT_TYPE } from "#/features/workspaces/files/workspace-file-preview.constants"; import type { WorkspaceKernelEventBus } from "#/features/workspaces/kernel/workspace-kernel-events"; -import { - getWorkspaceKernelFilePreviewShellPath, - getWorkspaceKernelFileShellPath, -} from "#/features/workspaces/kernel/workspace-kernel-files"; +import { WorkspaceKernelFileMigrator } from "#/features/workspaces/kernel/workspace-kernel-file-migrations"; +import { getWorkspaceKernelFileShellPath } from "#/features/workspaces/kernel/workspace-kernel-files"; import { parseWorkspaceMetadataJson } from "#/features/workspaces/kernel/workspace-kernel-metadata"; import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; +import type { + KernelItemProjectionRow, + KernelItemRow, +} from "#/features/workspaces/kernel/workspace-kernel-rows"; import type { WorkspaceKernelStore } from "#/features/workspaces/kernel/workspace-kernel-store"; import type { CreateWorkspaceKernelFileFromUploadArgs, - ReadWorkspaceKernelFileContentArgs, - ReadWorkspaceKernelFileContentResult, + ReadWorkspaceKernelFileSourceArgs, ReadWorkspaceKernelFilePreviewResult, ReadWorkspaceKernelFileProjectionArgs, ReadWorkspaceKernelFileProjectionResult, UpsertWorkspaceKernelFileProjectionArgs, WorkspaceKernelFileProjectionFormat, - WorkspaceKernelFileProjectionStatus, } from "#/features/workspaces/kernel/workspace-kernel-types"; import { getMetadataNumber, @@ -38,9 +35,11 @@ import { } from "#/features/workspaces/model/workspace-file/policy"; import type { WorkspaceCommandResult } from "#/features/workspaces/realtime/messages"; import { recordOperationalOutcome } from "#/integrations/observability/operational-events"; +import { deleteR2Prefix } from "#/lib/r2"; export class WorkspaceKernelFileCommands { private readonly events: WorkspaceKernelEventBus; + private readonly migrator: WorkspaceKernelFileMigrator; private readonly r2: R2Bucket; private readonly sql: WorkspaceKernelSql; private readonly store: WorkspaceKernelStore; @@ -56,6 +55,12 @@ export class WorkspaceKernelFileCommands { workspaceId: () => string; }) { this.events = input.events; + this.migrator = new WorkspaceKernelFileMigrator({ + bucket: input.r2, + sql: input.sql, + workspace: input.workspace, + workspaceId: input.workspaceId, + }); this.r2 = input.r2; this.sql = input.sql; this.store = input.store; @@ -70,7 +75,11 @@ export class WorkspaceKernelFileCommands { this.store.assertParentIsValid(parentId); - const object = await this.r2.get(input.objectKey); + if (this.store.getItemRowIncludingDeleted(input.id)) { + throw new Error("Workspace item id already exists."); + } + + const object = await this.r2.head(input.objectKey); if (!object) { throw new Error("Uploaded file was not found."); @@ -80,7 +89,6 @@ export class WorkspaceKernelFileCommands { throw new Error("Uploaded file size did not match the upload request."); } - const bytes = new Uint8Array(await object.arrayBuffer()); const descriptor = getWorkspaceUploadFamily(input.assetKind); const contentType = resolveWorkspaceFileContentType({ contentType: input.contentType, @@ -89,7 +97,7 @@ export class WorkspaceKernelFileCommands { }); const now = Date.now(); - const itemId = crypto.randomUUID(); + const itemId = input.id; const requestedName = normalizeWorkspaceUploadFileName(input.fileName, descriptor); const name = this.store.resolveItemName({ itemId, @@ -109,13 +117,10 @@ export class WorkspaceKernelFileCommands { contentType, descriptor, originalName: requestedName, - sizeBytes: bytes.byteLength, + sizeBytes: object.size, source: input.source, }); - await this.workspace.writeFileBytes(shellPath, bytes, contentType); - await this.r2.delete(input.objectKey); - this.sql` INSERT INTO kernel_items ( id, @@ -126,6 +131,7 @@ export class WorkspaceKernelFileCommands { metadata_json, sort_order, shell_path, + object_key, created_at, updated_at, deleted_at @@ -139,6 +145,7 @@ export class WorkspaceKernelFileCommands { ${JSON.stringify(metadataJson)}, ${this.store.getNextSortOrder(parentId)}, ${shellPath}, + ${input.objectKey}, ${now}, ${now}, NULL @@ -153,52 +160,37 @@ export class WorkspaceKernelFileCommands { payload: { item }, }); - const previewGenerator = resolveUploadPreviewGenerator(descriptor); - - if (previewGenerator) { - await this.tryCreateUploadPreview({ - actorUserId: input.actorUserId ?? null, - bytes, - generate: previewGenerator, - itemId, - label: descriptor.assetKind, - now, - }); - } - return { result: item, event }; } - async readFileContent( - input: ReadWorkspaceKernelFileContentArgs, - ): Promise { + async getFileSource(input: ReadWorkspaceKernelFileSourceArgs) { const row = this.store.assertActiveItem(input.itemId); if (row.type !== "file") { throw new Error("Workspace item is not a file."); } - const bytes = await this.workspace.readFileBytes(row.shell_path); - - if (!bytes) { - throw new Error("Workspace file content was not found."); - } - const item = this.store.requireItem(input.itemId); const contentType = getMetadataString(item.metadataJson, "mimeType"); const originalName = getMetadataString(item.metadataJson, "originalName"); const sizeBytes = getMetadataNumber(item.metadataJson, "sizeBytes"); + const objectKey = + row.object_key ?? + (await this.migrator.migrateSource(row.shell_path, item.id, { + contentType: contentType ?? "application/octet-stream", + expectedSize: sizeBytes, + })); return { - bytes, + objectKey, contentType: contentType ?? "application/octet-stream", fileName: originalName ?? item.name, - sizeBytes: sizeBytes ?? bytes.byteLength, + sizeBytes: sizeBytes ?? (await this.requireObject(objectKey)).size, }; } async readFilePreview( - input: ReadWorkspaceKernelFileContentArgs, + input: ReadWorkspaceKernelFileSourceArgs, ): Promise { const row = this.store.assertActiveItem(input.itemId); @@ -215,18 +207,20 @@ export class WorkspaceKernelFileCommands { return null; } - const bytes = - projection.status === "ready" && projection.content_shell_path - ? await this.workspace.readFileBytes(projection.content_shell_path) + const objectKey = + projection.status === "ready" + ? (projection.object_key ?? (await this.migrator.migratePreview(projection, input.itemId))) : null; + const metadataJson = parseProjectionMetadataJson(projection.metadata_json); return { itemId: projection.item_id, status: projection.status, - bytes, + objectKey, contentType: WORKSPACE_FILE_PREVIEW_CONTENT_TYPE, + sizeBytes: getMetadataNumber(metadataJson, "sizeBytes"), sourceHash: projection.source_hash, - metadataJson: parseProjectionMetadataJson(projection.metadata_json), + metadataJson, updatedAt: new Date(projection.updated_at).toISOString(), }; } @@ -237,6 +231,21 @@ export class WorkspaceKernelFileCommands { if (row.type !== "file") { throw new Error("Workspace item is not a file."); } + if (input.status === "ready" && input.format === "pages") { + if (!row.object_key) { + throw new Error("Ready page projections require a source hash and manifest object."); + } + const [source, manifest] = await Promise.all([ + this.r2.head(row.object_key), + this.r2.head(input.objectKey), + ]); + if (!source || source.etag !== input.sourceHash) { + throw new Error("The file source changed before its extraction could be published."); + } + if (!manifest) { + throw new Error("The page projection manifest was not found."); + } + } const now = Date.now(); @@ -253,45 +262,10 @@ export class WorkspaceKernelFileCommands { projection: UpsertWorkspaceKernelFileProjectionArgs; now: number; }) { - const contentShellPath = - input.projection.content == null && input.projection.contentBytes == null - ? this.getExistingProjectionPath({ - itemId: input.itemId, - format: input.projection.format, - }) - : getWorkspaceKernelProjectionShellPath({ - itemId: input.itemId, - format: input.projection.format, - }); - - if (input.projection.content != null) { - const projectionShellPath = getWorkspaceKernelProjectionShellPath({ - itemId: input.itemId, - format: input.projection.format, - }); - - await this.workspace.mkdir(`/items/${input.itemId}/projections`, { - recursive: true, - }); - await this.workspace.writeFile( - projectionShellPath, - input.projection.content, - getProjectionContentType(input.projection.format), - ); - } - - if (input.projection.contentBytes != null) { - const previewShellPath = getWorkspaceKernelFilePreviewShellPath(input.itemId); - - await this.workspace.mkdir(`/items/${input.itemId}/derivatives`, { - recursive: true, - }); - await this.workspace.writeFileBytes( - previewShellPath, - input.projection.contentBytes, - WORKSPACE_FILE_PREVIEW_CONTENT_TYPE, - ); - } + const legacyContentPath = this.getProjectionRow({ + itemId: input.itemId, + format: input.projection.format, + })?.content_shell_path; this.sql` INSERT INTO kernel_item_projections ( @@ -301,6 +275,7 @@ export class WorkspaceKernelFileCommands { provider, provider_mode, content_shell_path, + object_key, error_message, source_hash, metadata_json, @@ -313,7 +288,8 @@ export class WorkspaceKernelFileCommands { ${input.projection.status}, ${input.projection.provider ?? null}, ${input.projection.providerMode ?? null}, - ${contentShellPath}, + NULL, + ${input.projection.objectKey ?? null}, ${input.projection.errorMessage ?? null}, ${input.projection.sourceHash ?? null}, ${JSON.stringify(input.projection.metadataJson ?? {})}, @@ -324,12 +300,17 @@ export class WorkspaceKernelFileCommands { status = excluded.status, provider = excluded.provider, provider_mode = excluded.provider_mode, - content_shell_path = COALESCE(excluded.content_shell_path, kernel_item_projections.content_shell_path), + content_shell_path = NULL, + object_key = excluded.object_key, error_message = excluded.error_message, source_hash = excluded.source_hash, metadata_json = excluded.metadata_json, updated_at = excluded.updated_at `; + + if (legacyContentPath) { + await this.workspace.rm(legacyContentPath, { force: true }); + } } async readFileProjection( @@ -341,19 +322,29 @@ export class WorkspaceKernelFileCommands { throw new Error("Workspace item is not a file."); } - const projection = this.getProjectionRow(input); + let projection = this.getProjectionRow(input); if (!projection) { return null; } + if ( + projection.status === "ready" && + projection.format === "pages" && + !projection.object_key && + projection.content_shell_path + ) { + await this.migrator.migratePageProjection(projection); + projection = this.getProjectionRow(input); + if (!projection) { + return null; + } + } return { itemId: projection.item_id, format: projection.format, status: projection.status, - content: projection.content_shell_path - ? await this.workspace.readFile(projection.content_shell_path) - : null, + objectKey: projection.object_key, provider: projection.provider, providerMode: projection.provider_mode, errorMessage: projection.error_message, @@ -363,13 +354,6 @@ export class WorkspaceKernelFileCommands { }; } - private getExistingProjectionPath(input: { - itemId: string; - format: WorkspaceKernelFileProjectionFormat; - }) { - return this.getProjectionRow(input)?.content_shell_path ?? null; - } - private getProjectionRow(input: { itemId: string; format: WorkspaceKernelFileProjectionFormat }) { return ( this.sql` @@ -381,90 +365,113 @@ export class WorkspaceKernelFileCommands { ); } - private async tryCreateUploadPreview(input: { - actorUserId: string | null; - bytes: Uint8Array; - generate: (bytes: Uint8Array) => Promise<{ bytes: Uint8Array; width: number; height: number }>; - itemId: string; - label: string; - now: number; - }) { - const startedAt = Date.now(); - let failure: unknown; - let previewMetrics: { height: number; outputBytes: number; width: number } | undefined; - - try { - const preview = await input.generate(input.bytes); - previewMetrics = { - height: preview.height, - outputBytes: preview.bytes.byteLength, - width: preview.width, - }; - const sourceHash = await sha256Base64Url(input.bytes); - - await this.writeProjectionRow({ - itemId: input.itemId, - now: input.now, - projection: { - itemId: input.itemId, - format: "preview", - status: "ready", - contentBytes: preview.bytes, - sourceHash, - metadataJson: { - contentType: WORKSPACE_FILE_PREVIEW_CONTENT_TYPE, - width: preview.width, - height: preview.height, - }, - }, - }); - } catch (error) { - failure = error; - - await this.writeProjectionRow({ - itemId: input.itemId, - now: input.now, - projection: { - itemId: input.itemId, - format: "preview", - status: "failed", - errorMessage: getErrorMessage(error), - }, - }); - } finally { + async deleteObjects(itemIds: string[]) { + const fileItemIds = itemIds.filter( + (itemId) => this.store.getItemRowIncludingDeleted(itemId)?.type === "file", + ); + const results = await Promise.allSettled( + fileItemIds.map((itemId) => + deleteR2Prefix( + this.r2, + getWorkspaceFileItemObjectPrefix({ workspaceId: this.workspaceId(), itemId }), + ), + ), + ); + const failure = results.find((result) => result.status === "rejected"); + + if (failure?.status === "rejected") { recordOperationalOutcome({ - distinctId: input.actorUserId ?? undefined, - error: failure, - event: "workspace_file_preview", + error: failure.reason, + event: "workspace_file_object_cleanup", fields: { - asset_kind: input.label, - duration_ms: Date.now() - startedAt, - height: previewMetrics?.height, - input_bytes: input.bytes.byteLength, - item_id: input.itemId, - output_bytes: previewMetrics?.outputBytes, - user_id: input.actorUserId, - width: previewMetrics?.width, + item_count: fileItemIds.length, workspace_id: this.workspaceId(), }, }); } } -} -type KernelItemProjectionRow = { - item_id: string; - format: WorkspaceKernelFileProjectionFormat; - status: WorkspaceKernelFileProjectionStatus; - provider: string | null; - provider_mode: string | null; - content_shell_path: string | null; - error_message: string | null; - source_hash: string | null; - metadata_json: string; - created_at: number; - updated_at: number; -}; + async migrateLegacyStorage() { + const sourceRows = this.sql` + SELECT * FROM kernel_items + WHERE type = 'file' AND object_key IS NULL + `; + const projectionRows = this.sql` + SELECT * FROM kernel_item_projections + WHERE content_shell_path IS NOT NULL + `; + const orphanedItemIds = new Set(); + let orphanedProjectionCount = 0; + + for (const row of sourceRows) { + if (!(await this.workspace.stat(row.shell_path))) { + orphanedItemIds.add(row.id); + const orphanedProjections = projectionRows.filter( + (candidate) => candidate.item_id === row.id && candidate.content_shell_path, + ); + orphanedProjectionCount += orphanedProjections.length; + for (const projection of orphanedProjections) { + await this.workspace.rm(projection.content_shell_path!, { force: true }); + } + this.sql`DELETE FROM kernel_item_projections WHERE item_id = ${row.id}`; + this + .sql`DELETE FROM kernel_relations WHERE from_item_id = ${row.id} OR to_item_id = ${row.id}`; + this.sql`DELETE FROM kernel_items WHERE id = ${row.id}`; + await deleteR2Prefix( + this.r2, + getWorkspaceFileItemObjectPrefix({ workspaceId: this.workspaceId(), itemId: row.id }), + ); + continue; + } + + const metadata = parseWorkspaceMetadataJson(row.metadata_json); + await this.migrator.migrateSource(row.shell_path, row.id, { + contentType: getMetadataString(metadata, "mimeType") ?? "application/octet-stream", + expectedSize: getMetadataNumber(metadata, "sizeBytes"), + }); + } + + for (const projection of projectionRows) { + if (orphanedItemIds.has(projection.item_id)) { + continue; + } + if (projection.format === "pages") { + await this.migrator.migratePageProjection(projection); + } else { + await this.migrator.migratePreview(projection, projection.item_id); + } + } + + const remainingSources = + this.sql<{ count: number }>` + SELECT COUNT(*) AS count FROM kernel_items + WHERE type = 'file' AND object_key IS NULL + `[0]?.count ?? 0; + const remainingProjections = + this.sql<{ count: number }>` + SELECT COUNT(*) AS count FROM kernel_item_projections + WHERE content_shell_path IS NOT NULL + `[0]?.count ?? 0; + + return { + migratedSources: sourceRows.length - orphanedItemIds.size, + migratedProjections: projectionRows.length - orphanedProjectionCount, + orphanedFilesRemoved: orphanedItemIds.size, + remainingSources, + remainingProjections, + }; + } + + private async requireObject(objectKey: string) { + const object = await this.r2.head(objectKey); + + if (!object) { + throw new Error("Workspace file object was not found."); + } + + return object; + } +} function createFileMetadata(input: { contentType: string; @@ -492,23 +499,4 @@ function createFileMetadata(input: { return metadata; } -function getWorkspaceKernelProjectionShellPath(input: { - itemId: string; - format: WorkspaceKernelFileProjectionFormat; -}) { - if (input.format === "preview") { - return getWorkspaceKernelFilePreviewShellPath(input.itemId); - } - - return `/items/${input.itemId}/projections/${input.format}.json`; -} - -function getProjectionContentType(format: WorkspaceKernelFileProjectionFormat) { - return format === "pages" ? "application/json" : "text/markdown"; -} - -function getErrorMessage(error: unknown) { - return error instanceof Error ? error.message : String(error); -} - const parseProjectionMetadataJson = parseWorkspaceMetadataJson; diff --git a/src/features/workspaces/kernel/workspace-kernel-file-migrations.ts b/src/features/workspaces/kernel/workspace-kernel-file-migrations.ts new file mode 100644 index 00000000..bfaba023 --- /dev/null +++ b/src/features/workspaces/kernel/workspace-kernel-file-migrations.ts @@ -0,0 +1,163 @@ +import type { Workspace as ShellWorkspace } from "@cloudflare/shell"; + +import { parseMarkdownPagesProjection } from "#/features/workspaces/extraction/page-markdown-projection"; +import { writeWorkspacePageProjection } from "#/features/workspaces/extraction/workspace-page-projection"; +import { + getWorkspaceFilePreviewObjectKey, + getWorkspaceFileSourceObjectKey, +} from "#/features/workspaces/files/workspace-file-object-keys"; +import { WORKSPACE_FILE_PREVIEW_CONTENT_TYPE } from "#/features/workspaces/files/workspace-file-preview.constants"; +import { parseWorkspaceMetadataJson } from "#/features/workspaces/kernel/workspace-kernel-metadata"; +import type { KernelItemProjectionRow } from "#/features/workspaces/kernel/workspace-kernel-rows"; +import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; + +export class WorkspaceKernelFileMigrator { + private readonly bucket: R2Bucket; + private readonly sql: WorkspaceKernelSql; + private readonly workspace: ShellWorkspace; + private readonly workspaceId: () => string; + + constructor(input: { + bucket: R2Bucket; + sql: WorkspaceKernelSql; + workspace: ShellWorkspace; + workspaceId: () => string; + }) { + this.bucket = input.bucket; + this.sql = input.sql; + this.workspace = input.workspace; + this.workspaceId = input.workspaceId; + } + + async migratePageProjection(projection: KernelItemProjectionRow) { + if (!projection.content_shell_path) { + return; + } + + const content = await this.workspace.readFile(projection.content_shell_path); + if (content === null) { + throw new Error("Legacy page projection content was not found."); + } + + const sourceHash = projection.source_hash ?? `legacy-${projection.updated_at}`; + const reference = await writeWorkspacePageProjection({ + bucket: this.bucket, + itemId: projection.item_id, + metadata: parseWorkspaceMetadataJson(projection.metadata_json), + pages: parseMarkdownPagesProjection(content), + provider: projection.provider ?? "legacy", + providerMode: projection.provider_mode ?? "legacy", + runId: `legacy-${sourceHash}`, + sourceHash, + tier: "legacy", + workspaceId: this.workspaceId(), + }); + + this.sql` + UPDATE kernel_item_projections + SET object_key = ${reference.manifestObjectKey}, content_shell_path = NULL + WHERE item_id = ${projection.item_id} AND format = ${projection.format} + `; + await this.workspace.rm(projection.content_shell_path, { force: true }); + } + + async migrateSource( + shellPath: string, + itemId: string, + input: { contentType: string; expectedSize: number | null }, + ) { + const objectKey = getWorkspaceFileSourceObjectKey({ + workspaceId: this.workspaceId(), + itemId, + }); + const object = await this.copyLegacyFile({ + contentType: input.contentType, + expectedSize: input.expectedSize, + objectKey, + shellPath, + }); + + if (input.expectedSize !== null && object.size !== input.expectedSize) { + await this.bucket.delete(objectKey); + throw new Error("Legacy workspace file migration produced an invalid object."); + } + + this.sql`UPDATE kernel_items SET object_key = ${objectKey} WHERE id = ${itemId}`; + await this.workspace.rm(shellPath, { force: true }); + return objectKey; + } + + async migratePreview(projection: KernelItemProjectionRow, itemId: string) { + if (!projection.content_shell_path) { + return null; + } + if (projection.object_key && (await this.bucket.head(projection.object_key))) { + this.sql` + UPDATE kernel_item_projections + SET content_shell_path = NULL + WHERE item_id = ${itemId} AND format = 'preview' + `; + await this.workspace.rm(projection.content_shell_path, { force: true }); + return projection.object_key; + } + + const objectKey = getWorkspaceFilePreviewObjectKey({ + workspaceId: this.workspaceId(), + itemId, + }); + const object = await this.copyLegacyFile({ + contentType: WORKSPACE_FILE_PREVIEW_CONTENT_TYPE, + expectedSize: null, + objectKey, + shellPath: projection.content_shell_path, + }); + + const metadataJson = { + ...parseWorkspaceMetadataJson(projection.metadata_json), + sizeBytes: object.size, + }; + this.sql` + UPDATE kernel_item_projections + SET object_key = ${objectKey}, + content_shell_path = NULL, + metadata_json = ${JSON.stringify(metadataJson)} + WHERE item_id = ${itemId} AND format = 'preview' + `; + await this.workspace.rm(projection.content_shell_path, { force: true }); + return objectKey; + } + + private async copyLegacyFile(input: { + contentType: string; + expectedSize: number | null; + objectKey: string; + shellPath: string; + }) { + const [stat, stream] = await Promise.all([ + this.workspace.stat(input.shellPath), + this.workspace.readFileStream(input.shellPath), + ]); + + if (!stat || !stream) { + throw new Error("Legacy workspace file content was not found."); + } + if (input.expectedSize !== null && stat.size !== input.expectedSize) { + throw new Error("Legacy workspace file size did not match its metadata."); + } + + const fixedLengthStream = new FixedLengthStream(stat.size); + const [object] = await Promise.all([ + this.bucket.put(input.objectKey, fixedLengthStream.readable, { + httpMetadata: { contentType: input.contentType }, + }), + stream.pipeTo(fixedLengthStream.writable), + ]); + + if (!object || object.size !== stat.size) { + await this.bucket.delete(input.objectKey); + throw new Error("Legacy workspace file migration produced an invalid object."); + } + + return object; + } +} diff --git a/src/features/workspaces/kernel/workspace-kernel-files.ts b/src/features/workspaces/kernel/workspace-kernel-files.ts index 0c1e4e47..4ac20a43 100644 --- a/src/features/workspaces/kernel/workspace-kernel-files.ts +++ b/src/features/workspaces/kernel/workspace-kernel-files.ts @@ -16,10 +16,6 @@ export function getWorkspaceKernelFileShellPath(input: { itemId: string; extensi return `/items/${input.itemId}/content.${input.extension}`; } -export function getWorkspaceKernelFilePreviewShellPath(itemId: string) { - return `/items/${itemId}/derivatives/preview.webp`; -} - export function getWorkspaceKernelContentMimeType(type: WorkspaceItemType) { switch (type) { case "document": diff --git a/src/features/workspaces/kernel/workspace-kernel-item-commands.ts b/src/features/workspaces/kernel/workspace-kernel-item-commands.ts index dd2575d9..0a52b721 100644 --- a/src/features/workspaces/kernel/workspace-kernel-item-commands.ts +++ b/src/features/workspaces/kernel/workspace-kernel-item-commands.ts @@ -65,6 +65,11 @@ export class WorkspaceKernelItemCommands { input: CreateWorkspaceKernelItemArgs, ): Promise> { const type = workspaceItemTypeSchema.parse(input.type); + + if (type === "file") { + throw new Error("Binary workspace files must be created through the upload flow."); + } + const id = input.id ?? crypto.randomUUID(); const parentId = input.parentId ?? null; const color = resolveWorkspaceItemColorForCreate({ @@ -274,7 +279,7 @@ export class WorkspaceKernelItemCommands { const item = this.store.assertActiveItem(input.itemId); const itemSummary = mapKernelItemRow(item, this.workspaceId()); - return item.type === "folder" + return item.type === "folder" || item.type === "file" ? { item: itemSummary, content: null } : { item: itemSummary, @@ -288,8 +293,8 @@ export class WorkspaceKernelItemCommands { const item = this.store.assertActiveItem(input.itemId); const type = workspaceItemTypeSchema.parse(item.type); - if (type === "folder") { - throw new Error("Folders do not have writable content."); + if (type === "folder" || type === "file") { + throw new Error("This workspace item does not have writable text content."); } await this.workspace.writeFile( diff --git a/src/features/workspaces/kernel/workspace-kernel-rows.ts b/src/features/workspaces/kernel/workspace-kernel-rows.ts index ac6e7f0c..2ef9d030 100644 --- a/src/features/workspaces/kernel/workspace-kernel-rows.ts +++ b/src/features/workspaces/kernel/workspace-kernel-rows.ts @@ -3,6 +3,10 @@ import { workspaceItemTypeSchema } from "#/features/workspaces/contracts"; import { getWorkspaceItemTypeMeta } from "#/features/workspaces/defaults"; import { parseWorkspaceMetadataJson } from "#/features/workspaces/kernel/workspace-kernel-metadata"; import type { WorkspaceRealtimeEvent } from "#/features/workspaces/realtime/messages"; +import type { + WorkspaceKernelFileProjectionFormat, + WorkspaceKernelFileProjectionStatus, +} from "#/features/workspaces/kernel/workspace-kernel-types"; export type KernelItemRow = { id: string; @@ -13,6 +17,7 @@ export type KernelItemRow = { metadata_json: string; sort_order: number; shell_path: string; + object_key: string | null; created_at: number; updated_at: number; deleted_at: number | null; @@ -28,6 +33,21 @@ export type KernelEventRow = { created_at: number; }; +export type KernelItemProjectionRow = { + item_id: string; + format: WorkspaceKernelFileProjectionFormat; + status: WorkspaceKernelFileProjectionStatus; + provider: string | null; + provider_mode: string | null; + content_shell_path: string | null; + object_key: string | null; + error_message: string | null; + source_hash: string | null; + metadata_json: string; + created_at: number; + updated_at: number; +}; + export function mapKernelItemRow(row: KernelItemRow, workspaceId: string): WorkspaceItemSummary { const type = workspaceItemTypeSchema.parse(row.type); diff --git a/src/features/workspaces/kernel/workspace-kernel-schema.ts b/src/features/workspaces/kernel/workspace-kernel-schema.ts index 4b0d8ca1..f4c5d290 100644 --- a/src/features/workspaces/kernel/workspace-kernel-schema.ts +++ b/src/features/workspaces/kernel/workspace-kernel-schema.ts @@ -20,11 +20,13 @@ export function initializeWorkspaceKernelStorage(sql: WorkspaceKernelSql) { metadata_json TEXT NOT NULL DEFAULT '{}', sort_order INTEGER NOT NULL, shell_path TEXT NOT NULL UNIQUE, + object_key TEXT, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, deleted_at INTEGER ) `; + ensureKernelItemObjectKeyColumn(sql); sql`CREATE INDEX IF NOT EXISTS kernel_items_parent_idx ON kernel_items (parent_id, deleted_at, sort_order)`; sql`CREATE INDEX IF NOT EXISTS kernel_items_type_idx @@ -38,6 +40,7 @@ export function initializeWorkspaceKernelStorage(sql: WorkspaceKernelSql) { provider TEXT, provider_mode TEXT, content_shell_path TEXT, + object_key TEXT, error_message TEXT, source_hash TEXT, metadata_json TEXT NOT NULL DEFAULT '{}', @@ -46,6 +49,7 @@ export function initializeWorkspaceKernelStorage(sql: WorkspaceKernelSql) { PRIMARY KEY (item_id, format) ) `; + ensureProjectionObjectKeyColumn(sql); sql`CREATE INDEX IF NOT EXISTS kernel_item_projections_status_idx ON kernel_item_projections (status, updated_at)`; sql` @@ -86,6 +90,22 @@ export function initializeWorkspaceKernelStorage(sql: WorkspaceKernelSql) { ON kernel_events (revision)`; } +function ensureKernelItemObjectKeyColumn(sql: WorkspaceKernelSql) { + const columns = sql<{ name: string }>`PRAGMA table_info(kernel_items)`; + + if (!columns.some((column) => column.name === "object_key")) { + sql`ALTER TABLE kernel_items ADD COLUMN object_key TEXT`; + } +} + +function ensureProjectionObjectKeyColumn(sql: WorkspaceKernelSql) { + const columns = sql<{ name: string }>`PRAGMA table_info(kernel_item_projections)`; + + if (!columns.some((column) => column.name === "object_key")) { + sql`ALTER TABLE kernel_item_projections ADD COLUMN object_key TEXT`; + } +} + function createSiblingNameIndexes(sql: WorkspaceKernelSql) { try { sql` diff --git a/src/features/workspaces/kernel/workspace-kernel-types.ts b/src/features/workspaces/kernel/workspace-kernel-types.ts index 02669fbc..99f72eee 100644 --- a/src/features/workspaces/kernel/workspace-kernel-types.ts +++ b/src/features/workspaces/kernel/workspace-kernel-types.ts @@ -96,12 +96,12 @@ export interface ReadWorkspaceKernelItemArgs { itemId: string; } -export interface ReadWorkspaceKernelFileContentArgs { +export interface ReadWorkspaceKernelFileSourceArgs { itemId: string; } -export interface ReadWorkspaceKernelFileContentResult { - bytes: Uint8Array; +export interface WorkspaceKernelFileSource { + objectKey: string; contentType: string; fileName: string; sizeBytes: number; @@ -116,26 +116,48 @@ export type WorkspaceKernelFileProjectionStatus = | "ready" | "failed"; -export interface UpsertWorkspaceKernelFileProjectionArgs { +interface WorkspaceKernelFileProjectionMutationBase { itemId: string; format: WorkspaceKernelFileProjectionFormat; - status: WorkspaceKernelFileProjectionStatus; - content?: string | null; - contentBytes?: Uint8Array | null; - provider?: string | null; - providerMode?: string | null; - errorMessage?: string | null; - sourceHash?: string | null; - metadataJson?: Record; actorUserId?: string | null; clientMutationId?: string | null; } +export type UpsertWorkspaceKernelFileProjectionArgs = + | (WorkspaceKernelFileProjectionMutationBase & { + status: "not_started" | "queued" | "processing"; + errorMessage?: never; + metadataJson?: never; + objectKey?: never; + provider?: never; + providerMode?: never; + sourceHash?: never; + }) + | (WorkspaceKernelFileProjectionMutationBase & { + status: "failed"; + errorMessage: string; + metadataJson?: never; + objectKey?: never; + provider?: never; + providerMode?: never; + sourceHash?: never; + }) + | (WorkspaceKernelFileProjectionMutationBase & { + status: "ready"; + errorMessage?: never; + metadataJson?: Record; + objectKey: string; + provider?: string | null; + providerMode?: string | null; + sourceHash: string; + }); + export interface ReadWorkspaceKernelFilePreviewResult { itemId: string; status: WorkspaceKernelFileProjectionStatus; - bytes: Uint8Array | null; + objectKey: string | null; contentType: string; + sizeBytes: number | null; sourceHash: string | null; metadataJson: Record; updatedAt: string; @@ -150,7 +172,7 @@ export interface ReadWorkspaceKernelFileProjectionResult { itemId: string; format: WorkspaceKernelFileProjectionFormat; status: WorkspaceKernelFileProjectionStatus; - content: string | null; + objectKey: string | null; provider: string | null; providerMode: string | null; errorMessage: string | null; @@ -167,6 +189,7 @@ export interface WriteWorkspaceKernelItemArgs { } export interface CreateWorkspaceKernelFileFromUploadArgs { + id: string; parentId?: string | null; fileName: string; fileSize: number; diff --git a/src/features/workspaces/kernel/workspace-kernel.ts b/src/features/workspaces/kernel/workspace-kernel.ts index aa8e25e6..4e02bcdf 100644 --- a/src/features/workspaces/kernel/workspace-kernel.ts +++ b/src/features/workspaces/kernel/workspace-kernel.ts @@ -29,7 +29,7 @@ import type { ListWorkspaceKernelItemsArgs, MoveWorkspaceKernelItemsArgs, MoveWorkspaceKernelItemsResult, - ReadWorkspaceKernelFileContentArgs, + ReadWorkspaceKernelFileSourceArgs, ReadWorkspaceKernelFileProjectionArgs, ReadWorkspaceKernelItemArgs, RenameWorkspaceKernelItemArgs, @@ -46,6 +46,7 @@ import type { WorkspaceRealtimeServerMessage, } from "#/features/workspaces/realtime/messages"; import { recordOperationalOutcome } from "#/integrations/observability/operational-events"; +import { deleteR2Prefix } from "#/lib/r2"; const workspaceKernelInlineThresholdBytes = 1_500_000; @@ -152,11 +153,11 @@ export class WorkspaceKernel extends Agent { return await this.fileCommands.createFileFromUpload(input); } - async readFileContent(input: ReadWorkspaceKernelFileContentArgs) { - return await this.fileCommands.readFileContent(input); + async getFileSource(input: ReadWorkspaceKernelFileSourceArgs) { + return await this.fileCommands.getFileSource(input); } - async readFilePreview(input: ReadWorkspaceKernelFileContentArgs) { + async readFilePreview(input: ReadWorkspaceKernelFileSourceArgs) { return await this.fileCommands.readFilePreview(input); } @@ -168,6 +169,10 @@ export class WorkspaceKernel extends Agent { return await this.fileCommands.readFileProjection(input); } + async migrateLegacyFileStorage() { + return await this.fileCommands.migrateLegacyStorage(); + } + async renameItem( input: RenameWorkspaceKernelItemArgs, ): Promise> { @@ -193,9 +198,11 @@ export class WorkspaceKernel extends Agent { async deleteItems( input: DeleteWorkspaceKernelItemsArgs, ): Promise> { - return this.runMutation("delete_items", input, input.itemIds.length, () => - this.itemCommands.deleteItems(input), - ); + return this.runMutation("delete_items", input, input.itemIds.length, async () => { + const command = await this.itemCommands.deleteItems(input); + await this.fileCommands.deleteObjects(command.result.deletedItemIds); + return command; + }); } async readItem(input: ReadWorkspaceKernelItemArgs) { @@ -263,31 +270,20 @@ export class WorkspaceKernel extends Agent { } await Promise.all([ - this.deleteR2Prefix(getChatAttachmentWorkspacePrefix(workspaceId)), - this.deleteR2Prefix(`uploads/workspaces/${workspaceId}/`), - this.deleteR2Prefix(`workspace_kernel_files/${workspaceId}/`), + deleteR2Prefix( + this.env.WORKSPACE_KERNEL_FILES, + getChatAttachmentWorkspacePrefix(workspaceId), + ), + deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `uploads/workspaces/${workspaceId}/`), + deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_kernel_files/${workspaceId}/`), + deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_file_objects/${workspaceId}/`), + deleteR2Prefix(this.env.WORKSPACE_KERNEL_FILES, `workspace_file_uploads/${workspaceId}/`), ]); await this.ctx.storage.deleteAll(); return { attempted: documentItemIds.length + 1, failed }; } - private async deleteR2Prefix(prefix: string) { - const bucket = this.env.WORKSPACE_KERNEL_FILES; - - if (!bucket) { - return; - } - - let cursor: string | undefined; - - do { - const listed = await bucket.list({ prefix, cursor }); - await Promise.all(listed.objects.map((object) => bucket.delete(object.key))); - cursor = listed.truncated ? listed.cursor : undefined; - } while (cursor); - } - private broadcastPresenceSnapshot() { this.broadcastRealtimeMessage({ type: "presence.snapshot", diff --git a/src/features/workspaces/model/workspace-file/limits.ts b/src/features/workspaces/model/workspace-file/limits.ts index 2089881b..1d808868 100644 --- a/src/features/workspaces/model/workspace-file/limits.ts +++ b/src/features/workspaces/model/workspace-file/limits.ts @@ -1,5 +1,7 @@ export const workspaceFileUploadLimits = { maxFilesPerSelection: 50, - maxBytesPerSelection: 200 * 1024 * 1024, - concurrency: 5, + maxFileBytes: 200 * 1024 * 1024, + maxSelectionBytes: 200 * 1024 * 1024, + maxDocumentImportBytes: 10 * 1024 * 1024, + concurrency: 3, } as const; diff --git a/src/features/workspaces/operations/read-items.ts b/src/features/workspaces/operations/read-items.ts index a42c4925..a79314cd 100644 --- a/src/features/workspaces/operations/read-items.ts +++ b/src/features/workspaces/operations/read-items.ts @@ -8,13 +8,12 @@ import { } from "#/features/workspaces/operations/relations"; import { parseWorkspacePageRange, - readWorkspaceProjectionPages, WorkspacePageSelectionError, type WorkspaceReadPages, -} from "#/features/workspaces/operations/read-page-selection"; +} from "#/features/workspaces/read-page-selection"; import type { WorkspaceItemSummary } from "#/features/workspaces/contracts"; import { serializeTiptapDocumentToMarkdown } from "#/features/workspaces/documents/document-markdown"; -import { parseMarkdownPagesProjection } from "#/features/workspaces/extraction/page-markdown-projection"; +import { readWorkspacePageProjection } from "#/features/workspaces/extraction/workspace-page-projection"; import { parseTiptapDocumentJson } from "#/features/workspaces/documents/tiptap-document"; import type { WorkspaceKernelClient } from "#/features/workspaces/kernel/workspace-kernel-access"; import { buildWorkspaceKernelItemPathIndex } from "#/features/workspaces/kernel/workspace-kernel-paths"; @@ -46,6 +45,7 @@ const TRUNCATED_LINE_SUFFIX = `... (line truncated to ${MAX_WORKSPACE_READ_LINE_ export const readWorkspaceItemsFailureCodes = [ "page_range_out_of_range", + "page_selection_too_large", "path_is_folder", "path_not_absolute", "path_not_found", @@ -222,16 +222,15 @@ async function readWorkspaceFileItem(input: { return createWorkspaceFileStatusItem(input.path, "pending"); } - if (pagesProjection.status !== "ready" || pagesProjection.content === null) { + if (pagesProjection.status !== "ready" || pagesProjection.objectKey === null) { return createWorkspaceFileStatusItem(input.path, "failed"); } - const pageRead = readWorkspaceProjectionPages( - parseMarkdownPagesProjection(pagesProjection.content), - { - pages: input.pages, - }, - ); + const pageRead = await readWorkspacePageProjection({ + bucket: env.WORKSPACE_KERNEL_FILES, + manifestObjectKey: pagesProjection.objectKey, + pages: input.pages, + }); return { content: pageRead.content, @@ -294,3 +293,4 @@ function truncateWorkspaceMarkdownLine(line: string) { value: line.slice(0, MAX_WORKSPACE_READ_LINE_LENGTH) + TRUNCATED_LINE_SUFFIX, }; } +import { env } from "cloudflare:workers"; diff --git a/src/features/workspaces/operations/read-page-selection.test.ts b/src/features/workspaces/operations/read-page-selection.test.ts index 9210d99a..1c476538 100644 --- a/src/features/workspaces/operations/read-page-selection.test.ts +++ b/src/features/workspaces/operations/read-page-selection.test.ts @@ -1,25 +1,16 @@ import { describe, expect, it } from "vitest"; import { - readWorkspaceProjectionPages, + parseWorkspacePageRange, WorkspacePageSelectionError, -} from "#/features/workspaces/operations/read-page-selection"; +} from "#/features/workspaces/read-page-selection"; describe("workspace read page selection", () => { - it("returns an empty first page for empty ready projections", () => { - expect(readWorkspaceProjectionPages([], {})).toEqual({ - content: "", - pages: { - requested: "1", - returned: [1], - total: 1, - }, - }); + it("returns sorted unique page numbers", () => { + expect(parseWorkspacePageRange("3, 1-2, 2", 3)).toEqual([1, 2, 3]); }); - it("rejects out-of-range empty projection page requests", () => { - expect(() => readWorkspaceProjectionPages([], { pages: "2" })).toThrow( - WorkspacePageSelectionError, - ); + it("rejects out-of-range page requests", () => { + expect(() => parseWorkspacePageRange("4", 3)).toThrow(WorkspacePageSelectionError); }); }); diff --git a/src/features/workspaces/operations/read-page-selection.ts b/src/features/workspaces/operations/read-page-selection.ts deleted file mode 100644 index dfaf5b65..00000000 --- a/src/features/workspaces/operations/read-page-selection.ts +++ /dev/null @@ -1,94 +0,0 @@ -import { - joinMarkdownProjectionPages, - type MarkdownProjectionPage, -} from "#/features/workspaces/extraction/page-markdown-projection"; - -export interface WorkspaceReadPages { - requested: string; - returned: number[]; - total: number; -} - -export class WorkspacePageSelectionError extends Error { - constructor(readonly code: "page_range_out_of_range") { - super(code); - } -} - -export function readWorkspaceProjectionPages( - pages: readonly MarkdownProjectionPage[], - input: { - pages?: string; - }, -): { content: string; pages: WorkspaceReadPages } { - const requested = input.pages?.trim() || "1"; - if (pages.length === 0) { - const selectedPageNumbers = parseWorkspacePageRange(requested, 1); - - return { - content: "", - pages: { - requested, - returned: selectedPageNumbers, - total: 1, - }, - }; - } - - const maxPageNumber = pages.reduce((max, page) => Math.max(max, page.pageNumber), 0); - const selectedPageNumbers = parseWorkspacePageRange(requested, maxPageNumber); - const selectedPages = selectedPageNumbers.map((pageNumber) => { - const page = pages.find((candidate) => candidate.pageNumber === pageNumber); - - if (!page) { - throw new WorkspacePageSelectionError("page_range_out_of_range"); - } - - return page; - }); - - return { - content: joinMarkdownProjectionPages(selectedPages), - pages: { - requested, - returned: selectedPages.map((page) => page.pageNumber), - total: maxPageNumber, - }, - }; -} - -export function parseWorkspacePageRange(value: string, totalPages: number) { - const selected = new Set(); - const parts = value.split(","); - - for (const rawPart of parts) { - const part = rawPart.trim(); - - if (!part) { - continue; - } - - const rangeMatch = /^(\d+)(?:\s*-\s*(\d+))?$/.exec(part); - - if (!rangeMatch) { - throw new WorkspacePageSelectionError("page_range_out_of_range"); - } - - const start = Number(rangeMatch[1]); - const end = Number(rangeMatch[2] ?? rangeMatch[1]); - - if (start < 1 || end < start || end > totalPages) { - throw new WorkspacePageSelectionError("page_range_out_of_range"); - } - - for (let pageNumber = start; pageNumber <= end; pageNumber += 1) { - selected.add(pageNumber); - } - } - - if (selected.size === 0) { - throw new WorkspacePageSelectionError("page_range_out_of_range"); - } - - return Array.from(selected).sort((left, right) => left - right); -} diff --git a/src/features/workspaces/operations/workspace-tool-definitions.ts b/src/features/workspaces/operations/workspace-tool-definitions.ts index f24d52ed..920aa2f8 100644 --- a/src/features/workspaces/operations/workspace-tool-definitions.ts +++ b/src/features/workspaces/operations/workspace-tool-definitions.ts @@ -128,7 +128,7 @@ export const workspaceToolDefinitions = [ defineWorkspaceTool({ name: "workspace_read_items", description: - "Read ThinkEx documents and files by absolute path. Use pages for continuation: PDF pages for PDFs, 1000-line Markdown pages for documents and extracted files. Defaults to page 1. Check pages.total before reading more.", + "Read ThinkEx documents and files by absolute path. Use pages for continuation: PDF pages for PDFs, 1000-line Markdown pages for documents and extracted files. Defaults to page 1; read at most 20 pages per call and check pages.total before continuing.", inputSchema: workspaceReadItemsInputSchema, inputExamples: workspaceReadItemsInputExamples, outputSchema: workspaceReadItemsOutputSchema, diff --git a/src/features/workspaces/operations/workspace-tool-schemas.ts b/src/features/workspaces/operations/workspace-tool-schemas.ts index 6da95797..5a60e034 100644 --- a/src/features/workspaces/operations/workspace-tool-schemas.ts +++ b/src/features/workspaces/operations/workspace-tool-schemas.ts @@ -90,7 +90,7 @@ export const workspacePageRangeSchema = z .min(1) .regex(/^\d+(?:\s*-\s*\d+)?(?:\s*,\s*\d+(?:\s*-\s*\d+)?)*$/) .describe( - "1-based pages to read, like 1, 3, 5-7, or 1,4-6. For PDFs, pages are PDF pages. For Markdown-backed items, each page is 1000 Markdown lines. Defaults to 1.", + "Up to 20 1-based pages to read, like 1, 3, 5-7, or 1,4-6. For PDFs, pages are PDF pages. For Markdown-backed items, each page is 1000 Markdown lines. Defaults to 1.", ); export const workspaceListItemsInputSchema = z.object({ diff --git a/src/features/workspaces/read-page-selection.ts b/src/features/workspaces/read-page-selection.ts new file mode 100644 index 00000000..277ddeca --- /dev/null +++ b/src/features/workspaces/read-page-selection.ts @@ -0,0 +1,47 @@ +export interface WorkspaceReadPages { + requested: string; + returned: number[]; + total: number; +} + +export class WorkspacePageSelectionError extends Error { + constructor(readonly code: "page_range_out_of_range" | "page_selection_too_large") { + super(code); + } +} + +export function parseWorkspacePageRange(value: string, totalPages: number) { + const selected = new Set(); + const parts = value.split(","); + + for (const rawPart of parts) { + const part = rawPart.trim(); + + if (!part) { + continue; + } + + const rangeMatch = /^(\d+)(?:\s*-\s*(\d+))?$/.exec(part); + + if (!rangeMatch) { + throw new WorkspacePageSelectionError("page_range_out_of_range"); + } + + const start = Number(rangeMatch[1]); + const end = Number(rangeMatch[2] ?? rangeMatch[1]); + + if (start < 1 || end < start || end > totalPages) { + throw new WorkspacePageSelectionError("page_range_out_of_range"); + } + + for (let pageNumber = start; pageNumber <= end; pageNumber += 1) { + selected.add(pageNumber); + } + } + + if (selected.size === 0) { + throw new WorkspacePageSelectionError("page_range_out_of_range"); + } + + return Array.from(selected).sort((left, right) => left - right); +} diff --git a/src/features/workspaces/upload/pdf-upload-validation.test.ts b/src/features/workspaces/upload/pdf-upload-validation.test.ts index 0ba57655..c3f53114 100644 --- a/src/features/workspaces/upload/pdf-upload-validation.test.ts +++ b/src/features/workspaces/upload/pdf-upload-validation.test.ts @@ -1,42 +1,43 @@ -import { afterEach, describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import { WorkspaceFileUploadError } from "#/features/workspaces/model/workspace-file"; import { assertReadablePdfUpload } from "#/features/workspaces/upload/pdf-upload-validation"; -const getDocumentProxy = vi.hoisted(() => vi.fn()); +const requestWorkspaceFileProcessor = vi.hoisted(() => vi.fn()); -vi.mock("unpdf", () => ({ getDocumentProxy })); +vi.mock("#/features/workspaces/files/workspace-file-processor", () => ({ + requestWorkspaceFileProcessor, +})); describe("PDF upload validation", () => { - afterEach(() => { - getDocumentProxy.mockReset(); + beforeEach(() => { + requestWorkspaceFileProcessor.mockReset(); }); - it("accepts a readable PDF and releases the parser document", async () => { - const destroy = vi.fn().mockResolvedValue(undefined); - getDocumentProxy.mockResolvedValue({ destroy }); + it("accepts a PDF approved by the isolated validator", async () => { + requestWorkspaceFileProcessor.mockResolvedValue(new Response(null, { status: 204 })); - await expect(assertReadablePdfUpload(new Uint8Array([1, 2, 3]))).resolves.toBeUndefined(); - expect(destroy).toHaveBeenCalledOnce(); + await expect(validate(new Uint8Array([1, 2, 3]))).resolves.toBeUndefined(); + expect(requestWorkspaceFileProcessor).toHaveBeenCalledOnce(); }); it("rejects a PDF that requires a password", async () => { - getDocumentProxy.mockRejectedValue( - Object.assign(new Error("No password given"), { name: "PasswordException", code: 1 }), + requestWorkspaceFileProcessor.mockResolvedValue( + Response.json({ code: "PASSWORD_PROTECTED_PDF" }, { status: 422 }), ); - await expect(assertReadablePdfUpload(new Uint8Array([1, 2, 3]))).rejects.toMatchObject({ + await expect(validate(new Uint8Array([1, 2, 3]))).rejects.toMatchObject({ code: "PASSWORD_PROTECTED_PDF", status: 422, }); }); - it("rejects malformed PDFs without misclassifying them as password protected", async () => { - getDocumentProxy.mockRejectedValue( - Object.assign(new Error("Invalid PDF structure"), { name: "InvalidPDFException" }), + it("rejects malformed PDFs without misclassifying them", async () => { + requestWorkspaceFileProcessor.mockResolvedValue( + Response.json({ code: "INVALID_PDF" }, { status: 422 }), ); - await expect(assertReadablePdfUpload(new Uint8Array([1, 2, 3]))).rejects.toEqual( + await expect(validate(new Uint8Array([1, 2, 3]))).rejects.toEqual( expect.objectContaining>({ code: "INVALID_PDF", status: 422, @@ -44,3 +45,16 @@ describe("PDF upload validation", () => { ); }); }); + +function validate(bytes: Uint8Array) { + const body = new Response(bytes.slice().buffer).body; + + if (!body) { + throw new Error("Test stream was not created."); + } + + return assertReadablePdfUpload({ + env: {} as Cloudflare.Env, + object: { body } as R2ObjectBody, + }); +} diff --git a/src/features/workspaces/upload/pdf-upload-validation.ts b/src/features/workspaces/upload/pdf-upload-validation.ts index 1257d5cf..a30f183d 100644 --- a/src/features/workspaces/upload/pdf-upload-validation.ts +++ b/src/features/workspaces/upload/pdf-upload-validation.ts @@ -1,41 +1,46 @@ -import { getDocumentProxy } from "unpdf"; - +import { requestWorkspaceFileProcessor } from "#/features/workspaces/files/workspace-file-processor"; import { WorkspaceFileUploadError } from "#/features/workspaces/model/workspace-file"; -export async function assertReadablePdfUpload(bytes: ArrayBuffer | Uint8Array): Promise { - let document: Awaited> | null = null; - - try { - document = await getDocumentProxy(toUint8Array(bytes)); - } catch (error) { - if (isPdfPasswordError(error)) { - throw new WorkspaceFileUploadError({ - code: "PASSWORD_PROTECTED_PDF", - message: - "Password-protected PDFs aren’t supported. Remove the password and upload the PDF again.", - status: 422, - }); - } +export async function assertReadablePdfUpload(input: { + env: Cloudflare.Env; + object: R2ObjectBody; +}): Promise { + const response = await requestWorkspaceFileProcessor(input.env, { + body: input.object.body, + contentType: "application/pdf", + path: "/validate/pdf", + sizeBytes: input.object.size, + }); + + if (response.ok) { + return; + } + + const payload = await readValidationFailure(response); + if (payload.code === "PASSWORD_PROTECTED_PDF") { throw new WorkspaceFileUploadError({ - code: "INVALID_PDF", - message: "This PDF is damaged or invalid and could not be opened.", + code: "PASSWORD_PROTECTED_PDF", + message: + "Password-protected PDFs aren’t supported. Remove the password and upload the PDF again.", status: 422, }); - } finally { - await document?.destroy(); } -} -function isPdfPasswordError(error: unknown) { - return ( - typeof error === "object" && - error !== null && - "name" in error && - error.name === "PasswordException" - ); + throw new WorkspaceFileUploadError({ + code: "INVALID_PDF", + message: "This PDF is damaged or invalid and could not be opened.", + status: 422, + }); } -function toUint8Array(bytes: ArrayBuffer | Uint8Array) { - return bytes instanceof Uint8Array ? bytes : new Uint8Array(bytes); +async function readValidationFailure(response: Response) { + const payload: unknown = await response.json().catch(() => null); + + if (typeof payload !== "object" || payload === null) { + return { code: null }; + } + + const code = "code" in payload && typeof payload.code === "string" ? payload.code : null; + return { code }; } diff --git a/src/features/workspaces/upload/workspace-file-direct-upload.test.ts b/src/features/workspaces/upload/workspace-file-direct-upload.test.ts new file mode 100644 index 00000000..3647443a --- /dev/null +++ b/src/features/workspaces/upload/workspace-file-direct-upload.test.ts @@ -0,0 +1,123 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; + +import { + claimWorkspaceDirectUploadCompletion, + createWorkspaceDirectUploadSession, + verifyWorkspaceDirectUploadToken, +} from "#/features/workspaces/upload/workspace-file-direct-upload"; + +const sign = vi.hoisted(() => vi.fn()); + +vi.mock("aws4fetch", () => ({ + AwsClient: class { + sign = sign; + }, +})); + +describe("workspace direct upload sessions", () => { + afterEach(() => { + vi.useRealTimers(); + sign.mockReset(); + }); + + it("binds a presigned R2 PUT to signed completion claims", async () => { + sign.mockImplementation(async (request: Request) => request); + const env = createEnv(); + const session = await createWorkspaceDirectUploadSession(env, { + clientMutationId: "mutation-1", + contentType: "application/pdf", + fileName: "report.pdf", + fileSize: 42, + parentId: null, + userId: "user-1", + workspaceId: "workspace-1", + }); + const signedRequest = sign.mock.calls[0]?.[0] as Request; + + expect(signedRequest.method).toBe("PUT"); + expect(signedRequest.headers.get("content-type")).toBe("application/pdf"); + expect(signedRequest.url).toContain( + "/thinkex-workspace-kernel-files/workspace_file_uploads/workspace-1/", + ); + + await expect( + verifyWorkspaceDirectUploadToken(env, session.completionToken), + ).resolves.toMatchObject({ + clientMutationId: "mutation-1", + fileName: "report.pdf", + fileSize: 42, + userId: "user-1", + workspaceId: "workspace-1", + }); + }); + + it("rejects tampered and expired completion tokens", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-07-14T12:00:00Z")); + sign.mockImplementation(async (request: Request) => request); + const env = createEnv(); + const session = await createWorkspaceDirectUploadSession(env, { + clientMutationId: "mutation-1", + contentType: "image/png", + fileName: "image.png", + fileSize: 10, + parentId: null, + userId: "user-1", + workspaceId: "workspace-1", + }); + const lastCharacter = session.completionToken.at(-1); + const tampered = `${session.completionToken.slice(0, -1)}${lastCharacter === "x" ? "y" : "x"}`; + + await expect(verifyWorkspaceDirectUploadToken(env, tampered)).rejects.toThrow("invalid"); + + vi.setSystemTime(new Date("2026-07-14T12:31:00Z")); + await expect(verifyWorkspaceDirectUploadToken(env, session.completionToken)).rejects.toThrow( + "expired", + ); + }); + + it("allows only one upload completion owner", async () => { + const values = new Set(); + const env = { + ...createEnv(), + WORKSPACE_KERNEL_FILES: { + async put(key: string) { + if (values.has(key)) { + return null; + } + values.add(key); + return { key }; + }, + } as unknown as R2Bucket, + }; + const claims = { + clientMutationId: "mutation-1", + contentType: "application/pdf", + expiresAt: Math.floor(Date.now() / 1_000) + 60, + fileName: "report.pdf", + fileSize: 42, + itemId: crypto.randomUUID(), + parentId: null, + userId: "user-1", + version: 1 as const, + workspaceId: "workspace-1", + }; + + const [first, second] = await Promise.all([ + claimWorkspaceDirectUploadCompletion(env, claims), + claimWorkspaceDirectUploadCompletion(env, claims), + ]); + + expect(first).toContain("/completion"); + expect(second).toBeNull(); + }); +}); + +function createEnv() { + return { + R2_ACCESS_KEY_ID: "access-key", + R2_ACCOUNT_ID: "account-id", + R2_BUCKET_NAME: "thinkex-workspace-kernel-files", + R2_SECRET_ACCESS_KEY: "secret-key", + } as Cloudflare.Env; +} diff --git a/src/features/workspaces/upload/workspace-file-direct-upload.ts b/src/features/workspaces/upload/workspace-file-direct-upload.ts new file mode 100644 index 00000000..8c720cea --- /dev/null +++ b/src/features/workspaces/upload/workspace-file-direct-upload.ts @@ -0,0 +1,155 @@ +import { AwsClient } from "aws4fetch"; +import { z } from "zod"; + +import { + getWorkspaceFileUploadCompletionKey, + getWorkspaceFileUploadObjectKey, +} from "#/features/workspaces/files/workspace-file-object-keys"; + +const uploadUrlLifetimeSeconds = 30 * 60; +const uploadTokenVersion = 1; +const encoder = new TextEncoder(); +const decoder = new TextDecoder(); +const uploadClaimsSchema = z.object({ + clientMutationId: z.string().min(1), + contentType: z.string().min(1), + expiresAt: z.number().int().positive(), + fileName: z.string().min(1), + fileSize: z.number().int().positive(), + itemId: z.uuid(), + parentId: z.string().min(1).nullable(), + userId: z.string().min(1), + version: z.literal(uploadTokenVersion), + workspaceId: z.string().min(1), +}); + +export type WorkspaceDirectUploadClaims = z.infer; + +export async function createWorkspaceDirectUploadSession( + env: Cloudflare.Env, + input: Omit, +) { + const claims: WorkspaceDirectUploadClaims = { + ...input, + expiresAt: Math.floor(Date.now() / 1_000) + uploadUrlLifetimeSeconds, + itemId: crypto.randomUUID(), + version: uploadTokenVersion, + }; + const uploadUrl = await createPresignedUploadUrl(env, { + contentType: claims.contentType, + objectKey: getWorkspaceFileUploadObjectKey(claims), + }); + + return { + completionToken: await signUploadClaims(env.R2_SECRET_ACCESS_KEY, claims), + uploadUrl, + }; +} + +export async function verifyWorkspaceDirectUploadToken( + env: Cloudflare.Env, + token: string, +): Promise { + const [encodedPayload, encodedSignature, extra] = token.split("."); + + if (!encodedPayload || !encodedSignature || extra) { + throw new Error("Upload completion token is invalid."); + } + + const key = await createSigningKey(env.R2_SECRET_ACCESS_KEY, ["verify"]); + const valid = await crypto.subtle.verify( + "HMAC", + key, + decodeBase64Url(encodedSignature), + encoder.encode(encodedPayload), + ); + + if (!valid) { + throw new Error("Upload completion token is invalid."); + } + + const claims = uploadClaimsSchema.parse(JSON.parse(decodeBase64UrlText(encodedPayload))); + + if (claims.expiresAt < Math.floor(Date.now() / 1_000)) { + throw new Error("Upload completion token has expired."); + } + + return claims; +} + +export async function claimWorkspaceDirectUploadCompletion( + env: Cloudflare.Env, + claims: WorkspaceDirectUploadClaims, +): Promise { + const objectKey = getWorkspaceFileUploadCompletionKey(claims); + const claim = await env.WORKSPACE_KERNEL_FILES.put(objectKey, "", { + onlyIf: { etagDoesNotMatch: "*" }, + }); + + return claim ? objectKey : null; +} + +async function createPresignedUploadUrl( + env: Cloudflare.Env, + input: { contentType: string; objectKey: string }, +) { + const client = new AwsClient({ + accessKeyId: env.R2_ACCESS_KEY_ID, + region: "auto", + secretAccessKey: env.R2_SECRET_ACCESS_KEY, + service: "s3", + }); + const objectPath = input.objectKey.split("/").map(encodeURIComponent).join("/"); + const bucketName = encodeURIComponent(env.R2_BUCKET_NAME); + const url = new URL( + `https://${env.R2_ACCOUNT_ID}.r2.cloudflarestorage.com/${bucketName}/${objectPath}`, + ); + url.searchParams.set("X-Amz-Expires", String(uploadUrlLifetimeSeconds)); + const signed = await client.sign( + new Request(url, { + headers: { "content-type": input.contentType }, + method: "PUT", + }), + { aws: { signQuery: true } }, + ); + + return signed.url; +} + +async function signUploadClaims(secret: string, claims: WorkspaceDirectUploadClaims) { + const payload = encodeBase64Url(encoder.encode(JSON.stringify(claims))); + const key = await createSigningKey(secret, ["sign"]); + const signature = await crypto.subtle.sign("HMAC", key, encoder.encode(payload)); + return `${payload}.${encodeBase64Url(new Uint8Array(signature))}`; +} + +function createSigningKey(secret: string, usages: KeyUsage[]) { + return crypto.subtle.importKey( + "raw", + encoder.encode(`thinkex-workspace-upload:${secret}`), + { hash: "SHA-256", name: "HMAC" }, + false, + usages, + ); +} + +function encodeBase64Url(bytes: Uint8Array) { + let binary = ""; + + for (const byte of bytes) { + binary += String.fromCharCode(byte); + } + + return btoa(binary).replaceAll("+", "-").replaceAll("/", "_").replace(/=+$/, ""); +} + +function decodeBase64Url(value: string) { + const base64 = value.replaceAll("-", "+").replaceAll("_", "/"); + const padded = base64.padEnd(Math.ceil(base64.length / 4) * 4, "="); + const binary = atob(padded); + return Uint8Array.from(binary, (character) => character.charCodeAt(0)); +} + +function decodeBase64UrlText(value: string) { + return decoder.decode(decodeBase64Url(value)); +} diff --git a/src/features/workspaces/upload/workspace-file-upload-normalization.test.ts b/src/features/workspaces/upload/workspace-file-upload-normalization.test.ts deleted file mode 100644 index a09da094..00000000 --- a/src/features/workspaces/upload/workspace-file-upload-normalization.test.ts +++ /dev/null @@ -1,52 +0,0 @@ -import { describe, expect, it, vi } from "vitest"; - -import { requireWorkspaceFileTypeFromHint } from "#/features/workspaces/model/workspace-file"; -import { prepareWorkspaceFileUpload } from "#/features/workspaces/upload/workspace-file-upload-normalization"; - -const assertReadablePdfUpload = vi.hoisted(() => vi.fn()); - -vi.mock("#/features/workspaces/conversion/image-file-converter", () => ({ - convertImageFileToJpeg: vi.fn(), -})); -vi.mock("#/features/workspaces/conversion/office-pdf-converter", () => ({ - convertOfficeFileToPdf: vi.fn(), -})); -vi.mock("#/features/workspaces/upload/pdf-upload-validation", () => ({ - assertReadablePdfUpload, -})); - -describe("workspace file upload normalization", () => { - it("validates converted PDF bytes without consuming the upload body", async () => { - const convertedBytes = new Uint8Array([1, 2, 3]).buffer; - assertReadablePdfUpload.mockImplementation((bytes: ArrayBuffer) => { - new Uint8Array(bytes)[0] = 9; - }); - - const file = new File(["source"], "notes.docx", { - type: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", - }); - const descriptor = requireWorkspaceFileTypeFromHint({ - fileName: file.name, - contentType: file.type, - }); - - const prepared = await prepareWorkspaceFileUpload({ - converters: { - heic_to_jpeg: vi.fn(), - office_to_pdf: vi.fn().mockResolvedValue({ - bytes: convertedBytes, - contentType: "application/pdf", - sizeBytes: 3, - }), - }, - descriptor, - env: {} as Cloudflare.Env, - file, - }); - - expect(prepared.body).toBe(convertedBytes); - expect(new Uint8Array(prepared.body as ArrayBuffer)).toEqual(new Uint8Array([1, 2, 3])); - expect(assertReadablePdfUpload).toHaveBeenCalledWith(expect.any(ArrayBuffer)); - expect(assertReadablePdfUpload.mock.calls[0]?.[0]).not.toBe(convertedBytes); - }); -}); diff --git a/src/features/workspaces/upload/workspace-file-upload-normalization.ts b/src/features/workspaces/upload/workspace-file-upload-normalization.ts deleted file mode 100644 index 7945015d..00000000 --- a/src/features/workspaces/upload/workspace-file-upload-normalization.ts +++ /dev/null @@ -1,113 +0,0 @@ -import { convertImageFileToJpeg } from "#/features/workspaces/conversion/image-file-converter"; -import { convertOfficeFileToPdf } from "#/features/workspaces/conversion/office-pdf-converter"; -import { - getWorkspaceConvertedFileName, - requireWorkspaceFileTypeFromHint, - resolveWorkspaceUploadConversion, - type WorkspaceFileTypeDescriptor, - WorkspaceFileUploadError, - type WorkspaceUploadConversion, - workspaceFileUploadLimits, -} from "#/features/workspaces/model/workspace-file"; -import { assertReadablePdfUpload } from "#/features/workspaces/upload/pdf-upload-validation"; - -export interface PreparedWorkspaceFileUpload { - body: ArrayBuffer | File; - contentType: string; - descriptor: WorkspaceFileTypeDescriptor; - fileName: string; - fileSize: number; - source?: { - conversion: WorkspaceUploadConversion; - fileName: string; - mimeType: string | null; - sizeBytes: number; - }; -} - -type WorkspaceUploadConverter = ( - env: Cloudflare.Env, - input: { file: File; fileName: string }, -) => Promise<{ bytes: ArrayBuffer; contentType: string; sizeBytes: number }>; - -const defaultConverters = { - heic_to_jpeg: convertImageFileToJpeg, - office_to_pdf: convertOfficeFileToPdf, -} satisfies Record; - -export async function prepareWorkspaceFileUpload(input: { - converters?: Record; - descriptor: WorkspaceFileTypeDescriptor; - env: Cloudflare.Env; - file: File; -}): Promise { - const conversion = resolveWorkspaceUploadConversion({ - fileName: input.file.name, - contentType: input.file.type, - }); - const prepared = conversion - ? await prepareConvertedWorkspaceFileUpload(input, conversion) - : { - body: input.file, - contentType: input.file.type || "application/octet-stream", - descriptor: input.descriptor, - fileName: input.file.name, - fileSize: input.file.size, - }; - - await validatePreparedWorkspaceFileUpload(prepared); - return prepared; -} - -async function prepareConvertedWorkspaceFileUpload( - input: { - converters?: Record; - descriptor: WorkspaceFileTypeDescriptor; - env: Cloudflare.Env; - file: File; - }, - conversion: WorkspaceUploadConversion, -): Promise { - const converters = input.converters ?? defaultConverters; - const converted = await converters[conversion](input.env, { - file: input.file, - fileName: input.file.name, - }); - - if (converted.sizeBytes > workspaceFileUploadLimits.maxBytesPerSelection) { - throw new WorkspaceFileUploadError({ - code: "SELECTION_TOO_LARGE", - message: "Converted file is outside the supported upload limit.", - status: 413, - }); - } - - const fileName = getWorkspaceConvertedFileName(input.file.name, conversion); - const descriptor = requireWorkspaceFileTypeFromHint({ - fileName, - contentType: converted.contentType, - }); - - return { - body: converted.bytes, - contentType: converted.contentType, - descriptor, - fileName, - fileSize: converted.sizeBytes, - source: { - conversion, - fileName: input.file.name, - mimeType: input.file.type || null, - sizeBytes: input.file.size, - }, - }; -} - -async function validatePreparedWorkspaceFileUpload(input: PreparedWorkspaceFileUpload) { - if (input.descriptor.assetKind !== "pdf") { - return; - } - - const bytes = input.body instanceof File ? await input.body.arrayBuffer() : input.body.slice(0); - await assertReadablePdfUpload(bytes); -} diff --git a/src/features/workspaces/upload/workspace-file-upload-protocol.ts b/src/features/workspaces/upload/workspace-file-upload-protocol.ts new file mode 100644 index 00000000..400048b3 --- /dev/null +++ b/src/features/workspaces/upload/workspace-file-upload-protocol.ts @@ -0,0 +1,8 @@ +export interface WorkspaceDirectUploadSession { + completionToken: string; + uploadUrl: string; +} + +export interface CompleteWorkspaceDirectUploadInput { + completionToken: string; +} diff --git a/src/features/workspaces/upload/workspace-file-upload-storage.test.ts b/src/features/workspaces/upload/workspace-file-upload-storage.test.ts new file mode 100644 index 00000000..774e9c79 --- /dev/null +++ b/src/features/workspaces/upload/workspace-file-upload-storage.test.ts @@ -0,0 +1,138 @@ +import { describe, expect, it, vi } from "vitest"; + +import { + requireWorkspaceFileTypeFromHint, + type WorkspaceUploadConversion, +} from "#/features/workspaces/model/workspace-file"; +import { + storeWorkspaceFileUpload, + type WorkspaceUploadStreamConverter, +} from "#/features/workspaces/upload/workspace-file-upload-storage"; + +vi.mock("#/features/workspaces/conversion/image-file-converter", () => ({ + convertImageStreamToJpeg: vi.fn(), +})); +vi.mock("#/features/workspaces/conversion/office-pdf-converter", () => ({ + convertOfficeStreamToPdf: vi.fn(), +})); +vi.mock("#/features/workspaces/upload/pdf-upload-validation", () => ({ + assertReadablePdfUpload: vi.fn(), +})); + +describe("workspace file upload storage", () => { + it("streams an unchanged binary upload into its permanent object", async () => { + const bucket = createR2Bucket(); + const bytes = new Uint8Array([1, 2, 3, 4]); + + const result = await storeWorkspaceFileUpload({ + body: stream(bytes), + contentType: "image/png", + descriptor: requireWorkspaceFileTypeFromHint({ + fileName: "diagram.png", + contentType: "image/png", + }), + env: createEnv(bucket), + fileName: "diagram.png", + fileSize: bytes.byteLength, + objectKey: "workspace_file_objects/workspace/item/source", + }); + + expect(result).toMatchObject({ + fileName: "diagram.png", + fileSize: 4, + objectKey: "workspace_file_objects/workspace/item/source", + }); + expect(bucket.bytes()).toEqual(bytes); + }); + + it("streams conversion output to R2 and records source provenance", async () => { + const bucket = createR2Bucket(); + const converted = new Uint8Array([9, 8, 7]); + const converter: WorkspaceUploadStreamConverter = vi + .fn() + .mockResolvedValue(new Response(converted.slice().buffer)); + + const result = await storeWorkspaceFileUpload({ + body: stream(new Uint8Array([1, 2, 3, 4, 5])), + contentType: "image/heic", + converters: createConverters(converter), + descriptor: requireWorkspaceFileTypeFromHint({ + fileName: "photo.heic", + contentType: "image/heic", + }), + env: createEnv(bucket), + fileName: "photo.heic", + fileSize: 5, + objectKey: "workspace_file_objects/workspace/item/source", + }); + + expect(result).toMatchObject({ + contentType: "image/jpeg", + fileName: "photo.jpg", + fileSize: 3, + source: { + conversion: "heic_to_jpeg", + fileName: "photo.heic", + sizeBytes: 5, + }, + }); + expect(bucket.bytes()).toEqual(converted); + }); + + it("removes a partial object when the uploaded size is inconsistent", async () => { + const bucket = createR2Bucket(); + + await expect( + storeWorkspaceFileUpload({ + body: stream(new Uint8Array([1, 2, 3])), + contentType: "image/png", + descriptor: requireWorkspaceFileTypeFromHint({ + fileName: "diagram.png", + contentType: "image/png", + }), + env: createEnv(bucket), + fileName: "diagram.png", + fileSize: 4, + objectKey: "workspace_file_objects/workspace/item/source", + }), + ).rejects.toThrow("did not match"); + expect(bucket.bytes()).toBeNull(); + }); +}); + +function createConverters(converter: WorkspaceUploadStreamConverter) { + return { + heic_to_jpeg: converter, + office_to_pdf: converter, + } satisfies Record; +} + +function createEnv(bucket: R2Bucket) { + return { WORKSPACE_KERNEL_FILES: bucket } as Cloudflare.Env; +} + +function stream(bytes: Uint8Array) { + const body = new Response(bytes.slice().buffer).body; + + if (!body) { + throw new Error("Test stream was not created."); + } + + return body; +} + +function createR2Bucket() { + let value: Uint8Array | null = null; + const bucket = { + async put(key: string, body: ReadableStream) { + value = new Uint8Array(await new Response(body).arrayBuffer()); + return { key, size: value.byteLength } as R2Object; + }, + async delete() { + value = null; + }, + bytes: () => value, + }; + + return bucket as typeof bucket & R2Bucket; +} diff --git a/src/features/workspaces/upload/workspace-file-upload-storage.ts b/src/features/workspaces/upload/workspace-file-upload-storage.ts new file mode 100644 index 00000000..7faf93b6 --- /dev/null +++ b/src/features/workspaces/upload/workspace-file-upload-storage.ts @@ -0,0 +1,159 @@ +import { convertImageStreamToJpeg } from "#/features/workspaces/conversion/image-file-converter"; +import { convertOfficeStreamToPdf } from "#/features/workspaces/conversion/office-pdf-converter"; +import { + getWorkspaceConvertedFileName, + requireWorkspaceFileTypeFromHint, + resolveWorkspaceUploadConversion, + type WorkspaceFileTypeDescriptor, + WorkspaceFileUploadError, + type WorkspaceUploadConversion, + workspaceFileUploadLimits, +} from "#/features/workspaces/model/workspace-file"; +import { assertReadablePdfUpload } from "#/features/workspaces/upload/pdf-upload-validation"; + +export interface StoredWorkspaceFileUpload { + contentType: string; + descriptor: WorkspaceFileTypeDescriptor; + fileName: string; + fileSize: number; + objectKey: string; + source?: { + conversion: WorkspaceUploadConversion; + fileName: string; + mimeType: string | null; + sizeBytes: number; + }; +} + +export type WorkspaceUploadStreamConverter = ( + env: Cloudflare.Env, + input: { + body: ReadableStream; + contentType: string; + fileName: string; + sizeBytes: number; + }, +) => Promise; + +const defaultConverters = { + heic_to_jpeg: convertImageStreamToJpeg, + office_to_pdf: convertOfficeStreamToPdf, +} satisfies Record; + +export async function storeWorkspaceFileUpload(input: { + body: ReadableStream; + contentType: string; + converters?: Record; + descriptor: WorkspaceFileTypeDescriptor; + env: Cloudflare.Env; + fileName: string; + fileSize: number; + objectKey: string; +}): Promise { + const conversion = resolveWorkspaceUploadConversion({ + fileName: input.fileName, + contentType: input.contentType, + }); + const prepared = conversion + ? await convertWorkspaceFileUpload(input, conversion) + : { + body: input.body, + contentType: input.contentType || "application/octet-stream", + descriptor: input.descriptor, + fileName: input.fileName, + source: undefined, + }; + + try { + const stored = await input.env.WORKSPACE_KERNEL_FILES.put(input.objectKey, prepared.body, { + httpMetadata: { contentType: prepared.contentType }, + }); + + if (!stored) { + throw new Error("Workspace file could not be stored."); + } + + if (!conversion && stored.size !== input.fileSize) { + throw new Error("Stored workspace file size did not match the upload request."); + } + + if (stored.size > workspaceFileUploadLimits.maxFileBytes) { + throw createConvertedFileSizeError(); + } + + if (prepared.descriptor.assetKind === "pdf") { + const object = await input.env.WORKSPACE_KERNEL_FILES.get(input.objectKey); + + if (!object) { + throw new Error("Stored workspace PDF could not be read for validation."); + } + + await assertReadablePdfUpload({ env: input.env, object }); + } + + return { + contentType: prepared.contentType, + descriptor: prepared.descriptor, + fileName: prepared.fileName, + fileSize: stored.size, + objectKey: input.objectKey, + source: prepared.source, + }; + } catch (error) { + await input.env.WORKSPACE_KERNEL_FILES.delete(input.objectKey); + throw error; + } +} + +async function convertWorkspaceFileUpload( + input: { + body: ReadableStream; + contentType: string; + converters?: Record; + descriptor: WorkspaceFileTypeDescriptor; + env: Cloudflare.Env; + fileName: string; + fileSize: number; + }, + conversion: WorkspaceUploadConversion, +) { + const converters = input.converters ?? defaultConverters; + const response = await converters[conversion](input.env, { + body: input.body, + contentType: input.contentType, + fileName: input.fileName, + sizeBytes: input.fileSize, + }); + const contentType = getConvertedContentType(conversion); + const fileName = getWorkspaceConvertedFileName(input.fileName, conversion); + const descriptor = requireWorkspaceFileTypeFromHint({ fileName, contentType }); + + if (!response.body) { + throw new Error("Workspace file conversion returned an empty response."); + } + + return { + body: response.body, + contentType, + descriptor, + fileName, + source: { + conversion, + fileName: input.fileName, + mimeType: input.contentType || null, + sizeBytes: input.fileSize, + }, + }; +} + +function getConvertedContentType(conversion: WorkspaceUploadConversion) { + return conversion === "office_to_pdf" ? "application/pdf" : "image/jpeg"; +} + +function createConvertedFileSizeError(): WorkspaceFileUploadError { + return new WorkspaceFileUploadError({ + code: "SELECTION_TOO_LARGE", + message: "Converted file is outside the supported upload limit.", + status: 413, + }); +} diff --git a/src/features/workspaces/upload/workspace-upload-intake.test.ts b/src/features/workspaces/upload/workspace-upload-intake.test.ts new file mode 100644 index 00000000..36ef3d93 --- /dev/null +++ b/src/features/workspaces/upload/workspace-upload-intake.test.ts @@ -0,0 +1,43 @@ +import { describe, expect, it } from "vitest"; + +import { workspaceFileUploadLimits } from "#/features/workspaces/model/workspace-file"; +import { + getWorkspaceUploadSelectionValidationError, + validateWorkspaceUpload, +} from "#/features/workspaces/upload/workspace-upload-intake"; + +describe("workspace upload intake", () => { + it("accepts a binary file at the 200 MB limit", () => { + expect( + validateWorkspaceUpload({ + contentType: "application/pdf", + fileName: "research.pdf", + sizeBytes: workspaceFileUploadLimits.maxFileBytes, + }), + ).toMatchObject({ ok: true, plan: { kind: "file" } }); + }); + + it("bounds document imports before they are materialized in Worker memory", () => { + expect( + validateWorkspaceUpload({ + contentType: "text/csv", + fileName: "large.csv", + sizeBytes: workspaceFileUploadLimits.maxDocumentImportBytes + 1, + }), + ).toMatchObject({ + error: { code: "SELECTION_TOO_LARGE", status: 413 }, + ok: false, + }); + }); + + it("enforces the total selection limit separately from the per-file limit", () => { + const file = new File(["x"], "second.pdf", { type: "application/pdf" }); + const error = getWorkspaceUploadSelectionValidationError({ + acceptedCount: 1, + file, + selectionBytes: workspaceFileUploadLimits.maxSelectionBytes, + }); + + expect(error).toMatchObject({ code: "SELECTION_TOO_LARGE", status: 413 }); + }); +}); diff --git a/src/features/workspaces/upload/workspace-upload-intake.ts b/src/features/workspaces/upload/workspace-upload-intake.ts index ab4138d6..d893fc76 100644 --- a/src/features/workspaces/upload/workspace-upload-intake.ts +++ b/src/features/workspaces/upload/workspace-upload-intake.ts @@ -124,7 +124,21 @@ export function validateWorkspaceUpload(input: { }; } - if (input.sizeBytes > workspaceFileUploadLimits.maxBytesPerSelection) { + if ( + plan.kind === "document" && + input.sizeBytes > workspaceFileUploadLimits.maxDocumentImportBytes + ) { + return { + error: { + code: "SELECTION_TOO_LARGE", + message: "Import text, code, Markdown, CSV, or TSV files up to 10 MB.", + status: 413, + }, + ok: false, + }; + } + + if (input.sizeBytes > workspaceFileUploadLimits.maxFileBytes) { return { error: { code: "SELECTION_TOO_LARGE", @@ -161,7 +175,7 @@ export function getWorkspaceUploadSelectionValidationError(input: { return validationError; } - if (input.selectionBytes + input.file.size > workspaceFileUploadLimits.maxBytesPerSelection) { + if (input.selectionBytes + input.file.size > workspaceFileUploadLimits.maxSelectionBytes) { return { code: "SELECTION_TOO_LARGE", message: "Upload up to 200 MB at once.", diff --git a/src/lib/http/streaming-multipart.ts b/src/lib/http/streaming-multipart.ts new file mode 100644 index 00000000..b820e7e3 --- /dev/null +++ b/src/lib/http/streaming-multipart.ts @@ -0,0 +1,63 @@ +export function createStreamingMultipartFile(input: { + body: ReadableStream; + contentType: string; + fields?: Record; + fileName: string; + formFieldName: string; + sizeBytes: number; +}) { + const boundary = `thinkex-${crypto.randomUUID()}`; + const encoder = new TextEncoder(); + const fieldParts = Object.entries(input.fields ?? {}).map( + ([name, value]) => + `--${boundary}\r\nContent-Disposition: form-data; name="${sanitizePartValue(name)}"\r\n\r\n${value}\r\n`, + ); + const filePart = `--${boundary}\r\nContent-Disposition: form-data; name="${sanitizePartValue(input.formFieldName)}"; filename="${sanitizePartValue(input.fileName)}"\r\nContent-Type: ${input.contentType}\r\n\r\n`; + const prefix = encoder.encode(fieldParts.join("") + filePart); + const suffix = encoder.encode(`\r\n--${boundary}--\r\n`); + const stream = new FixedLengthStream(prefix.byteLength + input.sizeBytes + suffix.byteLength); + const done = pipeMultipartBody(stream.writable, input.body, prefix, suffix); + + return { + body: stream.readable, + contentType: `multipart/form-data; boundary=${boundary}`, + done, + }; +} + +async function pipeMultipartBody( + destination: WritableStream, + source: ReadableStream, + prefix: Uint8Array, + suffix: Uint8Array, +) { + const writer = destination.getWriter(); + const reader = source.getReader(); + + try { + await writer.write(prefix); + + while (true) { + const { done, value } = await reader.read(); + + if (done) { + break; + } + + await writer.write(value); + } + + await writer.write(suffix); + await writer.close(); + } catch (error) { + await writer.abort(error).catch(() => undefined); + throw error; + } finally { + reader.releaseLock(); + writer.releaseLock(); + } +} + +function sanitizePartValue(value: string) { + return value.replace(/["\r\n\\]/g, "_"); +} diff --git a/src/lib/r2.test.ts b/src/lib/r2.test.ts new file mode 100644 index 00000000..87d13b1e --- /dev/null +++ b/src/lib/r2.test.ts @@ -0,0 +1,33 @@ +import { describe, expect, it } from "vitest"; + +import { deleteR2Prefix } from "#/lib/r2"; + +describe("R2 helpers", () => { + it("deletes every page of an object prefix", async () => { + const keys = new Set([ + ...Array.from({ length: 1_500 }, (_, index) => `target/${index}`), + "other/keep", + ]); + const bucket = { + async delete(input: string | string[]) { + for (const key of Array.isArray(input) ? input : [input]) { + keys.delete(key); + } + }, + async list(input: { limit?: number; prefix?: string }) { + const objects = Array.from(keys) + .filter((key) => key.startsWith(input.prefix ?? "")) + .slice(0, input.limit) + .map((key) => ({ key })); + return { + objects, + truncated: objects.length === input.limit, + }; + }, + } as unknown as R2Bucket; + + await deleteR2Prefix(bucket, "target/"); + + expect(keys).toEqual(new Set(["other/keep"])); + }); +}); diff --git a/src/lib/r2.ts b/src/lib/r2.ts new file mode 100644 index 00000000..4563899e --- /dev/null +++ b/src/lib/r2.ts @@ -0,0 +1,15 @@ +export async function deleteR2Prefix(bucket: R2Bucket, prefix: string): Promise { + while (true) { + const page = await bucket.list({ prefix, limit: 1_000 }); + + if (page.objects.length === 0) { + return; + } + + await bucket.delete(page.objects.map((object) => object.key)); + + if (!page.truncated) { + return; + } + } +} diff --git a/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts b/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts index 8254d84d..1a611014 100644 --- a/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts +++ b/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts @@ -1,9 +1,14 @@ import { env } from "cloudflare:workers"; import { createFileRoute } from "@tanstack/react-router"; +import { z } from "zod"; import { createDbContext } from "#/db/server"; import { WorkspaceFileConversionError } from "#/features/workspaces/conversion/errors"; import { requestWorkspaceFileExtraction } from "#/features/workspaces/extraction/request-workspace-file-extraction"; +import { + getWorkspaceFileSourceObjectKey, + getWorkspaceFileUploadObjectKey, +} from "#/features/workspaces/files/workspace-file-object-keys"; import { createWorkspaceFileFromUpload, getWorkspaceKernel, @@ -17,224 +22,330 @@ import { WorkspaceForbiddenError, } from "#/features/workspaces/server/permissions"; import { - createDocumentContentFromWorkspaceUpload, - validateWorkspaceUpload, - type WorkspaceUploadPlan, -} from "#/features/workspaces/upload/workspace-upload-intake"; -import { prepareWorkspaceFileUpload } from "#/features/workspaces/upload/workspace-file-upload-normalization"; + claimWorkspaceDirectUploadCompletion, + createWorkspaceDirectUploadSession, + verifyWorkspaceDirectUploadToken, + type WorkspaceDirectUploadClaims, +} from "#/features/workspaces/upload/workspace-file-direct-upload"; import { observeWorkspaceFileIntake, type WorkspaceFileIntakeObservation, } from "#/features/workspaces/upload/workspace-file-intake-observability"; +import type { CompleteWorkspaceDirectUploadInput } from "#/features/workspaces/upload/workspace-file-upload-protocol"; +import { storeWorkspaceFileUpload } from "#/features/workspaces/upload/workspace-file-upload-storage"; +import { + createDocumentContentFromWorkspaceUpload, + validateWorkspaceUpload, + type WorkspaceUploadPlan, +} from "#/features/workspaces/upload/workspace-upload-intake"; import { apiError, apiJson, getRequestId } from "#/lib/api/http"; import { getSessionFromRequest } from "#/lib/auth-queries.server"; -const fileFormKey = "file"; -const parentIdFormKey = "parentId"; -const clientMutationIdFormKey = "clientMutationId"; +const uploadIntentSchema = z.object({ + clientMutationId: z.string().min(1), + contentType: z.string().min(1), + fileName: z.string().min(1), + fileSize: z.number().int().positive(), + parentId: z.string().min(1).nullable(), +}); +const uploadCompletionSchema = z.object({ completionToken: z.string().min(1) }); + +async function handleWorkspaceUploadPost(request: Request, workspaceId: string) { + const action = new URL(request.url).searchParams.get("action"); + + if (action === "initiate") { + return initiateWorkspaceFileUpload(request, workspaceId); + } + + if (action === "complete") { + return completeWorkspaceFileUpload(request, workspaceId); + } -async function handleWorkspaceFileUpload(request: Request, workspaceId: string) { + return apiError(getRequestId(request), 400, "INVALID_UPLOAD", "Unknown upload action."); +} + +async function initiateWorkspaceFileUpload(request: Request, workspaceId: string) { + const requestId = getRequestId(request); + + try { + const userId = await authorizeWorkspaceUpload(request, workspaceId); + const input = await readUploadIntent(request); + const validation = validateWorkspaceUpload({ + contentType: input.contentType, + fileName: input.fileName, + sizeBytes: input.fileSize, + }); + + if (!validation.ok) { + return apiError( + requestId, + validation.error.status, + validation.error.code, + validation.error.message, + ); + } + + const session = await createWorkspaceDirectUploadSession(env, { + ...input, + userId, + workspaceId, + }); + return apiJson(session, requestId, 201); + } catch (error) { + return workspaceUploadErrorResponse(requestId, error); + } +} + +async function completeWorkspaceFileUpload(request: Request, workspaceId: string) { const requestId = getRequestId(request); return observeWorkspaceFileIntake({ kind: "workspace_file", request, requestId, - run: (observation) => executeWorkspaceFileUpload(request, workspaceId, requestId, observation), + run: (observation) => finalizeWorkspaceFileUpload(request, workspaceId, requestId, observation), workspaceId, }); } -async function executeWorkspaceFileUpload( +async function finalizeWorkspaceFileUpload( request: Request, workspaceId: string, requestId: string, observation: WorkspaceFileIntakeObservation, ) { - let objectKey: string | null = null; + let stagingObjectKey: string | null = null; + let finalObjectKey: string | null = null; + let completionClaimKey: string | null = null; try { - const session = await getSessionFromRequest(request); + const userId = await authorizeWorkspaceUpload(request, workspaceId); + observation.userId = userId; + const claims = await readUploadClaims(request); - if (!session) { - return apiError( - requestId, - 401, - "UNAUTHORIZED", - "You must be signed in to upload workspace files.", - ); + if (claims.workspaceId !== workspaceId || claims.userId !== userId) { + throw invalidUpload("Upload completion token does not belong to this workspace."); } - observation.userId = session.user.id; - const dbContext = await createDbContext(); + const validation = validateWorkspaceUpload({ + contentType: claims.contentType, + fileName: claims.fileName, + sizeBytes: claims.fileSize, + }); - try { - await assertCanMutateWorkspace(dbContext.db, { - workspaceId, - userId: session.user.id, - }); - } finally { - await dbContext.dispose(); + if (!validation.ok) { + throw invalidUpload("Upload completion metadata is invalid."); } - - const formData = await request.formData(); - const file = formData.get(fileFormKey); - - if (!(file instanceof File)) { - return apiError(requestId, 400, "INVALID_UPLOAD", "File upload is missing a file."); + completionClaimKey = await claimWorkspaceDirectUploadCompletion(env, claims); + if (!completionClaimKey) { + throw invalidUpload("Upload is already being completed."); } - observation.inputBytes = file.size; - const uploadValidation = validateWorkspaceUpload({ - fileName: file.name, - sizeBytes: file.size, - contentType: file.type, - }); + observation.inputBytes = claims.fileSize; + observation.plan = validation.plan.kind; + stagingObjectKey = getWorkspaceFileUploadObjectKey(claims); + const stagingObject = await env.WORKSPACE_KERNEL_FILES.get(stagingObjectKey); - if (!uploadValidation.ok) { - return apiError( - requestId, - uploadValidation.error.status, - uploadValidation.error.code, - uploadValidation.error.message, - ); + if (!stagingObject || stagingObject.size !== claims.fileSize) { + throw invalidUpload("Uploaded file size does not match the selected file."); } - const parentId = getNullableString(formData.get(parentIdFormKey)); - const clientMutationId = getNullableString(formData.get(clientMutationIdFormKey)); - const uploadPlan = uploadValidation.plan; - observation.plan = uploadPlan.kind; - - if (uploadPlan.kind === "document") { + if (validation.plan.kind === "document") { const command = await createWorkspaceDocumentFromUpload({ - clientMutationId, - file, - parentId, - plan: uploadPlan, - userId: session.user.id, - workspaceId, + claims, + file: new File([await stagingObject.arrayBuffer()], claims.fileName, { + type: claims.contentType, + }), + plan: validation.plan, }); - observation.itemId = command.result.id; - observation.outputBytes = file.size; + observation.outputBytes = claims.fileSize; return apiJson(command, requestId); } - const upload = await prepareWorkspaceFileUpload({ - descriptor: uploadPlan.descriptor, + finalObjectKey = getWorkspaceFileSourceObjectKey(claims); + const upload = await storeWorkspaceFileUpload({ + body: stagingObject.body, + contentType: claims.contentType, + descriptor: validation.plan.descriptor, env, - file, + fileName: claims.fileName, + fileSize: claims.fileSize, + objectKey: finalObjectKey, }); observation.assetKind = upload.descriptor.assetKind; observation.conversion = upload.source?.conversion; observation.outputBytes = upload.fileSize; - objectKey = getWorkspaceFileUploadObjectKey(workspaceId); - await env.WORKSPACE_KERNEL_FILES.put(objectKey, upload.body, { - httpMetadata: { - contentType: upload.contentType, - }, - }); - const command = await createWorkspaceFileFromUpload({ - workspaceId, - userId: session.user.id, - parentId, + assetKind: upload.descriptor.assetKind, + clientMutationId: claims.clientMutationId, + contentType: upload.contentType, fileName: upload.fileName, fileSize: upload.fileSize, - objectKey, - contentType: upload.contentType, - assetKind: upload.descriptor.assetKind, + id: claims.itemId, + objectKey: upload.objectKey, + parentId: claims.parentId, source: upload.source, - clientMutationId, + userId, + workspaceId, }); - objectKey = null; + finalObjectKey = null; observation.itemId = command.result.id; - if ( - resolveWorkspaceFileAiReadStrategy({ - fileName: upload.fileName, - contentType: upload.contentType, - descriptor: upload.descriptor, - }) === "markdown_extraction" - ) { - await requestWorkspaceFileExtraction({ - workspaceId, - itemId: command.result.id, - actorUserId: session.user.id, - assetKind: upload.descriptor.assetKind, - requestId, - }); - } + await queueWorkspaceFileExtraction(upload, { + itemId: command.result.id, + requestId, + userId, + workspaceId, + }); return apiJson(command, requestId); } catch (error) { observation.error = error; - if (error instanceof WorkspaceForbiddenError) { - observation.error = undefined; - return apiError( - requestId, - 403, - "FORBIDDEN", - "You do not have permission to upload files to this workspace.", - ); - } - - if (error instanceof WorkspaceFileUploadError) { - observation.error = undefined; - return apiError(requestId, error.status, error.code, error.message); - } - - if (error instanceof WorkspaceFileConversionError) { - return apiError(requestId, 422, "CONVERSION_FAILED", error.userMessage); - } - - return apiError(requestId, 500, "UPLOAD_FAILED", "Unable to upload file right now."); + return workspaceUploadErrorResponse(requestId, error); } finally { - if (objectKey) { - await env.WORKSPACE_KERNEL_FILES.delete(objectKey); + if (stagingObjectKey) { + await env.WORKSPACE_KERNEL_FILES.delete(stagingObjectKey).catch(() => undefined); } + await Promise.allSettled([ + finalObjectKey ? env.WORKSPACE_KERNEL_FILES.delete(finalObjectKey) : Promise.resolve(), + completionClaimKey + ? env.WORKSPACE_KERNEL_FILES.delete(completionClaimKey) + : Promise.resolve(), + ]); } } -export const Route = createFileRoute("/api/v1/workspaces/$workspaceId/file-upload")({ - server: { - handlers: { - POST: ({ params, request }) => handleWorkspaceFileUpload(request, params.workspaceId), - }, - }, -}); +async function queueWorkspaceFileExtraction( + upload: Awaited>, + input: { itemId: string; requestId: string; userId: string; workspaceId: string }, +) { + if ( + resolveWorkspaceFileAiReadStrategy({ + contentType: upload.contentType, + descriptor: upload.descriptor, + fileName: upload.fileName, + }) !== "markdown_extraction" + ) { + return; + } -function getWorkspaceFileUploadObjectKey(workspaceId: string) { - return `uploads/workspaces/${workspaceId}/${crypto.randomUUID()}/source`; + await requestWorkspaceFileExtraction({ + actorUserId: input.userId, + assetKind: upload.descriptor.assetKind, + itemId: input.itemId, + requestId: input.requestId, + workspaceId: input.workspaceId, + }); } async function createWorkspaceDocumentFromUpload(input: { - clientMutationId: string | null; + claims: WorkspaceDirectUploadClaims; file: File; - parentId: string | null; plan: Extract; - userId: string; - workspaceId: string; }) { const documentContent = await createDocumentContentFromWorkspaceUpload({ file: input.file, plan: input.plan, }); - const kernel = await getWorkspaceKernel(input.workspaceId); + const kernel = await getWorkspaceKernel(input.claims.workspaceId); return kernel.createItem({ - parentId: input.parentId, - type: "document", - name: documentContent.name, - metadataJson: documentContent.metadataJson, + id: input.claims.itemId, + actorUserId: input.claims.userId, + clientMutationId: input.claims.clientMutationId, initialContent: documentContent.initialContent, - actorUserId: input.userId, - clientMutationId: input.clientMutationId, + metadataJson: documentContent.metadataJson, + name: documentContent.name, + parentId: input.claims.parentId, + type: "document", }); } -function getNullableString(value: FormDataEntryValue | null) { - return typeof value === "string" && value.trim() ? value : null; +async function authorizeWorkspaceUpload(request: Request, workspaceId: string) { + const session = await getSessionFromRequest(request); + + if (!session) { + throw new WorkspaceUploadRequestError(401, "UNAUTHORIZED", "You must be signed in."); + } + + const dbContext = await createDbContext(); + + try { + await assertCanMutateWorkspace(dbContext.db, { userId: session.user.id, workspaceId }); + } finally { + await dbContext.dispose(); + } + + return session.user.id; +} + +async function readUploadIntent(request: Request) { + return parseJsonRequest(request, uploadIntentSchema); } -export { handleWorkspaceFileUpload }; +async function readUploadClaims(request: Request): Promise { + const input: CompleteWorkspaceDirectUploadInput = await parseJsonRequest( + request, + uploadCompletionSchema, + ); + + try { + return await verifyWorkspaceDirectUploadToken(env, input.completionToken); + } catch { + throw invalidUpload("Upload completion token is invalid or expired."); + } +} + +async function parseJsonRequest(request: Request, schema: z.ZodType): Promise { + const value: unknown = await request.json().catch(() => null); + const result = schema.safeParse(value); + + if (!result.success) { + throw invalidUpload("Upload request is invalid."); + } + + return result.data; +} + +function invalidUpload(message: string) { + return new WorkspaceUploadRequestError(400, "INVALID_UPLOAD", message); +} + +function workspaceUploadErrorResponse(requestId: string, error: unknown) { + if (error instanceof WorkspaceUploadRequestError) { + return apiError(requestId, error.status, error.code, error.message); + } + + if (error instanceof WorkspaceForbiddenError) { + return apiError(requestId, 403, "FORBIDDEN", "You cannot upload to this workspace."); + } + + if (error instanceof WorkspaceFileUploadError) { + return apiError(requestId, error.status, error.code, error.message); + } + + if (error instanceof WorkspaceFileConversionError) { + return apiError(requestId, 422, "CONVERSION_FAILED", error.userMessage); + } + + return apiError(requestId, 500, "UPLOAD_FAILED", "Unable to upload file right now."); +} + +class WorkspaceUploadRequestError extends Error { + constructor( + readonly status: number, + readonly code: string, + message: string, + ) { + super(message); + } +} + +export const Route = createFileRoute("/api/v1/workspaces/$workspaceId/file-upload")({ + server: { + handlers: { + POST: ({ params, request }) => handleWorkspaceUploadPost(request, params.workspaceId), + }, + }, +}); diff --git a/src/routes/api/v1/workspaces.$workspaceId.files.$itemId.content.ts b/src/routes/api/v1/workspaces.$workspaceId.files.$itemId.content.ts index 310b9ae8..07bd0b2b 100644 --- a/src/routes/api/v1/workspaces.$workspaceId.files.$itemId.content.ts +++ b/src/routes/api/v1/workspaces.$workspaceId.files.$itemId.content.ts @@ -1,6 +1,7 @@ +import { env } from "cloudflare:workers"; import { createFileRoute } from "@tanstack/react-router"; -import { readWorkspaceKernelFileContent } from "#/features/workspaces/kernel/workspace-kernel-access"; +import { readWorkspaceKernelFileSource } from "#/features/workspaces/kernel/workspace-kernel-access"; import { WorkspaceForbiddenError } from "#/features/workspaces/server/permissions"; import { apiError, getRequestId } from "#/lib/api/http"; import { getSessionFromRequest } from "#/lib/auth-queries.server"; @@ -20,19 +21,23 @@ async function handleWorkspaceFileContent(request: Request, workspaceId: string, ); } - const content = await readWorkspaceKernelFileContent({ + const source = await readWorkspaceKernelFileSource({ workspaceId, userId: session.user.id, itemId, }); - const body = new Uint8Array(content.bytes).buffer; + const object = await env.WORKSPACE_KERNEL_FILES.get(source.objectKey); - return new Response(body, { + if (!object) { + throw new Error("Workspace file object was not found."); + } + + return new Response(object.body, { headers: { "cache-control": "private, max-age=60", - "content-disposition": `inline; filename="${sanitizeHeaderFileName(content.fileName)}"`, - "content-length": String(body.byteLength), - "content-type": content.contentType, + "content-disposition": `inline; filename="${sanitizeHeaderFileName(source.fileName)}"`, + "content-length": String(object.size), + "content-type": source.contentType, "x-request-id": requestId, }, }); diff --git a/src/routes/api/v1/workspaces.$workspaceId.files.$itemId.preview.ts b/src/routes/api/v1/workspaces.$workspaceId.files.$itemId.preview.ts index 3694fdbc..ac4586be 100644 --- a/src/routes/api/v1/workspaces.$workspaceId.files.$itemId.preview.ts +++ b/src/routes/api/v1/workspaces.$workspaceId.files.$itemId.preview.ts @@ -1,3 +1,4 @@ +import { env } from "cloudflare:workers"; import { createFileRoute } from "@tanstack/react-router"; import { WORKSPACE_FILE_PREVIEW_CONTENT_TYPE } from "#/features/workspaces/files/workspace-file-preview.constants"; import { readWorkspaceKernelFilePreview } from "#/features/workspaces/kernel/workspace-kernel-access"; @@ -26,7 +27,7 @@ async function handleWorkspaceFilePreview(request: Request, workspaceId: string, itemId, }); - if (!preview || preview.status !== "ready" || !preview.bytes) { + if (!preview || preview.status !== "ready" || !preview.objectKey) { return apiError( requestId, 404, @@ -35,17 +36,22 @@ async function handleWorkspaceFilePreview(request: Request, workspaceId: string, ); } - const body = new Uint8Array(preview.bytes).buffer; + const object = await env.WORKSPACE_KERNEL_FILES.get(preview.objectKey); + + if (!object) { + throw new Error("Workspace file preview object was not found."); + } + const cacheKey = preview.sourceHash ?? preview.updatedAt; const headers = { "cache-control": "private, max-age=86400, immutable", - "content-length": String(body.byteLength), + "content-length": String(object.size), "content-type": preview.contentType || WORKSPACE_FILE_PREVIEW_CONTENT_TYPE, etag: `"${cacheKey}"`, "x-request-id": requestId, }; - return new Response(body, { headers }); + return new Response(object.body, { headers }); } catch (error) { if (error instanceof WorkspaceForbiddenError) { return apiError( diff --git a/src/server.ts b/src/server.ts index 61a06f3e..7340f180 100644 --- a/src/server.ts +++ b/src/server.ts @@ -3,6 +3,7 @@ import handler from "@tanstack/react-start/server-entry"; import { routeUserAIRequest } from "#/features/workspaces/ai/auth"; import { routeDocumentSessionRequest } from "#/features/workspaces/documents/document-session-auth"; import { routeWorkspaceKernelRequest } from "#/features/workspaces/kernel/workspace-kernel-auth"; +import { routeWorkspaceFileR2Migration } from "#/features/workspaces/kernel/workspace-file-r2-migration-route"; import { recordOperationalFailure } from "#/integrations/observability/operational-events"; import { posthogHost, posthogHostOrigin, posthogProjectToken } from "#/integrations/posthog/config"; import { getTelemetryRequestDetails } from "#/integrations/posthog/server-context"; @@ -14,7 +15,7 @@ export { ImageFileConverter } from "#/features/workspaces/conversion/image-file- export { OfficePdfConverter } from "#/features/workspaces/conversion/office-pdf-converter"; export { DocumentSession } from "#/features/workspaces/documents/document-session"; export { WorkspaceFileExtractionWorkflow } from "#/features/workspaces/extraction/workspace-file-extraction-workflow"; -export { LiteParsePdfExtractor } from "#/features/workspaces/extraction/providers/liteparse"; +export { WorkspaceFileProcessor } from "#/features/workspaces/files/workspace-file-processor"; export { WorkspaceKernel } from "#/features/workspaces/kernel/workspace-kernel"; const isProduction = import.meta.env.PROD; @@ -112,6 +113,12 @@ function withSecurityHeaders(response: Response) { export default { async fetch(request, env) { try { + const migrationResponse = await routeWorkspaceFileR2Migration(request, env); + + if (migrationResponse) { + return migrationResponse; + } + const chatResponse = await routeUserAIRequest(request, env); if (chatResponse) { diff --git a/worker-configuration.d.ts b/worker-configuration.d.ts index 93639739..6e577d4b 100644 --- a/worker-configuration.d.ts +++ b/worker-configuration.d.ts @@ -1,5 +1,5 @@ /* eslint-disable */ -// Generated by Wrangler by running `wrangler types` (hash: 2096aa769661065b6446723bfa40549d) +// Generated by Wrangler by running `wrangler types` (hash: 66ce1a1950974042862fea1ba4c75c2b) // Runtime types generated with workerd@1.20260701.1 2026-03-24 nodejs_compat interface __BaseEnv_Env { WORKSPACE_KERNEL_FILES: R2Bucket; @@ -10,6 +10,7 @@ interface __BaseEnv_Env { AI: Ai; CF_VERSION_METADATA: WorkerVersionMetadata; FIRECRAWL_API_URL: "https://api.firecrawl.dev"; + R2_BUCKET_NAME: "thinkex-workspace-kernel-files-staging" | "thinkex-workspace-kernel-files"; SANDBOX_TRANSPORT: "rpc"; BETTER_AUTH_URL: string; BETTER_AUTH_SECRET: string; @@ -20,19 +21,22 @@ interface __BaseEnv_Env { AI_GATEWAY_API_KEY: string; AUTUMN_SECRET_KEY: string; TCC_API_KEY: string; + R2_ACCOUNT_ID: string; + R2_ACCESS_KEY_ID: string; + R2_SECRET_ACCESS_KEY: string; UserAIStore: DurableObjectNamespace; WorkspaceKernel: DurableObjectNamespace; DocumentSession: DurableObjectNamespace; CODE_SANDBOX: DurableObjectNamespace; OFFICE_PDF_CONVERTER: DurableObjectNamespace; IMAGE_FILE_CONVERTER: DurableObjectNamespace; - LITEPARSE_PDF_EXTRACTOR: DurableObjectNamespace; + WORKSPACE_FILE_PROCESSOR: DurableObjectNamespace; WORKSPACE_FILE_EXTRACTION_WORKFLOW: Workflow[0]['payload']>; } declare namespace Cloudflare { interface GlobalProps { mainModule: typeof import("./src/server"); - durableNamespaces: "UserAIStore" | "WorkspaceKernel" | "DocumentSession" | "OfficePdfConverter" | "Sandbox" | "ImageFileConverter" | "LiteParsePdfExtractor"; + durableNamespaces: "UserAIStore" | "WorkspaceKernel" | "DocumentSession" | "OfficePdfConverter" | "Sandbox" | "ImageFileConverter" | "WorkspaceFileProcessor"; } interface StagingEnv { WORKSPACE_KERNEL_FILES: R2Bucket; @@ -43,6 +47,7 @@ declare namespace Cloudflare { AI: Ai; CF_VERSION_METADATA: WorkerVersionMetadata; FIRECRAWL_API_URL: "https://api.firecrawl.dev"; + R2_BUCKET_NAME: "thinkex-workspace-kernel-files-staging"; SANDBOX_TRANSPORT: "rpc"; BETTER_AUTH_URL: string; BETTER_AUTH_SECRET: string; @@ -53,13 +58,16 @@ declare namespace Cloudflare { AI_GATEWAY_API_KEY: string; AUTUMN_SECRET_KEY: string; TCC_API_KEY: string; + R2_ACCOUNT_ID: string; + R2_ACCESS_KEY_ID: string; + R2_SECRET_ACCESS_KEY: string; UserAIStore: DurableObjectNamespace; WorkspaceKernel: DurableObjectNamespace; DocumentSession: DurableObjectNamespace; CODE_SANDBOX: DurableObjectNamespace; OFFICE_PDF_CONVERTER: DurableObjectNamespace; IMAGE_FILE_CONVERTER: DurableObjectNamespace; - LITEPARSE_PDF_EXTRACTOR: DurableObjectNamespace; + WORKSPACE_FILE_PROCESSOR: DurableObjectNamespace; WORKSPACE_FILE_EXTRACTION_WORKFLOW: Workflow[0]['payload']>; } interface ProductionEnv { @@ -71,6 +79,7 @@ declare namespace Cloudflare { AI: Ai; CF_VERSION_METADATA: WorkerVersionMetadata; FIRECRAWL_API_URL: "https://api.firecrawl.dev"; + R2_BUCKET_NAME: "thinkex-workspace-kernel-files"; SANDBOX_TRANSPORT: "rpc"; BETTER_AUTH_URL: string; BETTER_AUTH_SECRET: string; @@ -81,13 +90,16 @@ declare namespace Cloudflare { AI_GATEWAY_API_KEY: string; AUTUMN_SECRET_KEY: string; TCC_API_KEY: string; + R2_ACCOUNT_ID: string; + R2_ACCESS_KEY_ID: string; + R2_SECRET_ACCESS_KEY: string; UserAIStore: DurableObjectNamespace; WorkspaceKernel: DurableObjectNamespace; DocumentSession: DurableObjectNamespace; CODE_SANDBOX: DurableObjectNamespace; OFFICE_PDF_CONVERTER: DurableObjectNamespace; IMAGE_FILE_CONVERTER: DurableObjectNamespace; - LITEPARSE_PDF_EXTRACTOR: DurableObjectNamespace; + WORKSPACE_FILE_PROCESSOR: DurableObjectNamespace; WORKSPACE_FILE_EXTRACTION_WORKFLOW: Workflow[0]['payload']>; } interface Env extends __BaseEnv_Env {} @@ -97,7 +109,7 @@ type StringifyValues> = { [Binding in keyof EnvType]: EnvType[Binding] extends string ? EnvType[Binding] : string; }; declare namespace NodeJS { - interface ProcessEnv extends StringifyValues> {} + interface ProcessEnv extends StringifyValues> {} } // Begin runtime types diff --git a/wrangler.jsonc b/wrangler.jsonc index d2ee162b..41ed18d7 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -40,6 +40,7 @@ }, "vars": { "FIRECRAWL_API_URL": "https://api.firecrawl.dev", + "R2_BUCKET_NAME": "thinkex-workspace-kernel-files-staging", "SANDBOX_TRANSPORT": "rpc", }, "send_email": [ @@ -51,7 +52,8 @@ "r2_buckets": [ { "binding": "WORKSPACE_KERNEL_FILES", - "bucket_name": "thinkex-workspace-kernel-files", + "bucket_name": "thinkex-workspace-kernel-files-staging", + "remote": true, }, ], "secrets": { @@ -65,6 +67,9 @@ "AI_GATEWAY_API_KEY", "AUTUMN_SECRET_KEY", "TCC_API_KEY", + "R2_ACCOUNT_ID", + "R2_ACCESS_KEY_ID", + "R2_SECRET_ACCESS_KEY", ], }, "workflows": [ @@ -93,7 +98,7 @@ "max_instances": 1, }, { - "class_name": "LiteParsePdfExtractor", + "class_name": "WorkspaceFileProcessor", "image": "./containers/liteparse/Dockerfile", "instance_type": "standard-2", "max_instances": 2, @@ -127,8 +132,8 @@ "class_name": "ImageFileConverter", }, { - "name": "LITEPARSE_PDF_EXTRACTOR", - "class_name": "LiteParsePdfExtractor", + "name": "WORKSPACE_FILE_PROCESSOR", + "class_name": "WorkspaceFileProcessor", }, ], }, @@ -161,6 +166,15 @@ "tag": "v8_liteparse_pdf_extractor", "new_sqlite_classes": ["LiteParsePdfExtractor"], }, + { + "tag": "v9_workspace_file_processor", + "renamed_classes": [ + { + "from": "LiteParsePdfExtractor", + "to": "WorkspaceFileProcessor", + }, + ], + }, ], "d1_databases": [ { @@ -185,6 +199,7 @@ // Provision dedicated staging resources before the first deploy. "vars": { "FIRECRAWL_API_URL": "https://api.firecrawl.dev", + "R2_BUCKET_NAME": "thinkex-workspace-kernel-files-staging", "SANDBOX_TRANSPORT": "rpc", }, "secrets": { @@ -198,6 +213,9 @@ "AI_GATEWAY_API_KEY", "AUTUMN_SECRET_KEY", "TCC_API_KEY", + "R2_ACCOUNT_ID", + "R2_ACCESS_KEY_ID", + "R2_SECRET_ACCESS_KEY", ], }, "ai": { @@ -251,7 +269,7 @@ "max_instances": 1, }, { - "class_name": "LiteParsePdfExtractor", + "class_name": "WorkspaceFileProcessor", "image": "./containers/liteparse/Dockerfile", "instance_type": "standard-2", "max_instances": 2, @@ -284,8 +302,8 @@ "class_name": "ImageFileConverter", }, { - "name": "LITEPARSE_PDF_EXTRACTOR", - "class_name": "LiteParsePdfExtractor", + "name": "WORKSPACE_FILE_PROCESSOR", + "class_name": "WorkspaceFileProcessor", }, ], }, @@ -311,6 +329,7 @@ ], "vars": { "FIRECRAWL_API_URL": "https://api.firecrawl.dev", + "R2_BUCKET_NAME": "thinkex-workspace-kernel-files", "SANDBOX_TRANSPORT": "rpc", }, "secrets": { @@ -324,6 +343,9 @@ "AI_GATEWAY_API_KEY", "AUTUMN_SECRET_KEY", "TCC_API_KEY", + "R2_ACCOUNT_ID", + "R2_ACCESS_KEY_ID", + "R2_SECRET_ACCESS_KEY", ], }, "ai": { @@ -377,10 +399,11 @@ "max_instances": 1, }, { - "class_name": "LiteParsePdfExtractor", + "class_name": "WorkspaceFileProcessor", "image": "./containers/liteparse/Dockerfile", "instance_type": "standard-2", "max_instances": 2, + "name": "thinkex-liteparsepdfextractor-production", }, ], "durable_objects": { @@ -410,8 +433,8 @@ "class_name": "ImageFileConverter", }, { - "name": "LITEPARSE_PDF_EXTRACTOR", - "class_name": "LiteParsePdfExtractor", + "name": "WORKSPACE_FILE_PROCESSOR", + "class_name": "WorkspaceFileProcessor", }, ], }, From 1e54e0e9ab637b7cc5535b1f1a44f03a6ed4f147 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Tue, 14 Jul 2026 15:29:31 -0400 Subject: [PATCH 2/7] refactor(workspaces): remove legacy file storage paths --- .../extraction/page-markdown-projection.ts | 44 ----- .../extraction/workspace-page-projection.ts | 4 +- .../workspace-file-r2-migration-route.ts | 72 -------- .../kernel/workspace-kernel-file-commands.ts | 138 ++------------- .../workspace-kernel-file-migrations.ts | 163 ------------------ .../kernel/workspace-kernel-rows.ts | 1 - .../kernel/workspace-kernel-schema.ts | 19 -- .../workspaces/kernel/workspace-kernel.ts | 5 - src/server.ts | 7 - 9 files changed, 16 insertions(+), 437 deletions(-) delete mode 100644 src/features/workspaces/kernel/workspace-file-r2-migration-route.ts delete mode 100644 src/features/workspaces/kernel/workspace-kernel-file-migrations.ts diff --git a/src/features/workspaces/extraction/page-markdown-projection.ts b/src/features/workspaces/extraction/page-markdown-projection.ts index b5aa37dd..99f50db2 100644 --- a/src/features/workspaces/extraction/page-markdown-projection.ts +++ b/src/features/workspaces/extraction/page-markdown-projection.ts @@ -7,47 +7,3 @@ export function createSingleMarkdownProjectionPage(markdown: string): MarkdownPr const trimmed = markdown.trim(); return trimmed ? [{ pageNumber: 1, markdown: trimmed }] : []; } - -/** Parses the pre-R2 projection shape during lazy migration only. */ -export function parseMarkdownPagesProjection(content: string | null) { - if (!content?.trim()) { - return []; - } - - try { - const parsed = JSON.parse(content) as unknown; - - if (!Array.isArray(parsed)) { - return []; - } - - return parsed - .map(parseMarkdownProjectionPage) - .filter((page): page is MarkdownProjectionPage => page !== null); - } catch { - return []; - } -} - -function parseMarkdownProjectionPage(value: unknown): MarkdownProjectionPage | null { - if (!value || typeof value !== "object" || Array.isArray(value)) { - return null; - } - - const record = value as Record; - const pageNumber = - getPositiveInteger(record.pageNumber) ?? - getPositiveInteger(record.page) ?? - getPositiveInteger(record.index); - const markdown = typeof record.markdown === "string" ? record.markdown.trim() : ""; - - if (!pageNumber || !markdown) { - return null; - } - - return { pageNumber, markdown }; -} - -function getPositiveInteger(value: unknown) { - return typeof value === "number" && Number.isInteger(value) && value > 0 ? value : null; -} diff --git a/src/features/workspaces/extraction/workspace-page-projection.ts b/src/features/workspaces/extraction/workspace-page-projection.ts index 245fc978..8148ef0c 100644 --- a/src/features/workspaces/extraction/workspace-page-projection.ts +++ b/src/features/workspaces/extraction/workspace-page-projection.ts @@ -45,7 +45,7 @@ export async function writeWorkspacePageProjection(input: { providerMode: string; runId: string; sourceHash: string; - tier: "enhanced" | "fast" | "legacy"; + tier: "enhanced" | "fast"; workspaceId: string; }): Promise { const prefix = getWorkspacePageProjectionPrefix(input); @@ -178,7 +178,7 @@ export async function readWorkspacePageProjectionManifest( export function getWorkspacePageProjectionPrefix(input: { itemId: string; runId: string; - tier: "enhanced" | "fast" | "legacy"; + tier: "enhanced" | "fast"; workspaceId: string; }) { return `${getWorkspaceFileItemObjectPrefix(input)}extractions/${encodePathPart(input.runId)}/${input.tier}/`; diff --git a/src/features/workspaces/kernel/workspace-file-r2-migration-route.ts b/src/features/workspaces/kernel/workspace-file-r2-migration-route.ts deleted file mode 100644 index dc921dd0..00000000 --- a/src/features/workspaces/kernel/workspace-file-r2-migration-route.ts +++ /dev/null @@ -1,72 +0,0 @@ -import { getAgentByName } from "agents"; - -import { workspaceKernelAgentName } from "#/features/workspaces/agent-routes"; - -const migrationPath = "/api/internal/migrations/workspace-files-r2"; -const maximumBatchSize = 25; -const migrationConcurrency = 5; - -interface MigrationEnvironment extends Cloudflare.Env { - WORKSPACE_FILE_R2_MIGRATION_TOKEN: string; -} - -interface MigrationRequestBody { - after?: string; - limit?: number; -} - -export async function routeWorkspaceFileR2Migration(request: Request, environment: Cloudflare.Env) { - const url = new URL(request.url); - - if (url.pathname !== migrationPath) { - return null; - } - if (request.method !== "POST") { - return new Response("Method not allowed", { status: 405 }); - } - - const env = environment as MigrationEnvironment; - const authorization = request.headers.get("authorization"); - - if ( - !env.WORKSPACE_FILE_R2_MIGRATION_TOKEN || - authorization !== `Bearer ${env.WORKSPACE_FILE_R2_MIGRATION_TOKEN}` - ) { - return new Response("Unauthorized", { status: 401 }); - } - - const body = (await request.json()) as MigrationRequestBody; - const after = body.after ?? ""; - const limit = Math.min(Math.max(body.limit ?? maximumBatchSize, 1), maximumBatchSize); - const query = await env.DB.prepare("SELECT id FROM workspaces WHERE id > ? ORDER BY id LIMIT ?") - .bind(after, limit) - .all<{ id: string }>(); - const results = []; - - for (let index = 0; index < query.results.length; index += migrationConcurrency) { - const batch = query.results.slice(index, index + migrationConcurrency); - const batchResults = await Promise.all( - batch.map(async (workspace) => { - try { - const kernel = await getAgentByName(env[workspaceKernelAgentName], workspace.id); - return { - workspaceId: workspace.id, - result: await kernel.migrateLegacyFileStorage(), - }; - } catch (error) { - return { - workspaceId: workspace.id, - error: error instanceof Error ? error.message : String(error), - }; - } - }), - ); - results.push(...batchResults); - } - - return Response.json({ - done: query.results.length < limit, - nextAfter: query.results.at(-1)?.id ?? after, - results, - }); -} diff --git a/src/features/workspaces/kernel/workspace-kernel-file-commands.ts b/src/features/workspaces/kernel/workspace-kernel-file-commands.ts index c2177de1..18ff034f 100644 --- a/src/features/workspaces/kernel/workspace-kernel-file-commands.ts +++ b/src/features/workspaces/kernel/workspace-kernel-file-commands.ts @@ -1,17 +1,11 @@ -import type { Workspace as ShellWorkspace } from "@cloudflare/shell"; - import type { JsonValue, WorkspaceItemSummary } from "#/features/workspaces/contracts"; import { getWorkspaceFileItemObjectPrefix } from "#/features/workspaces/files/workspace-file-object-keys"; import { WORKSPACE_FILE_PREVIEW_CONTENT_TYPE } from "#/features/workspaces/files/workspace-file-preview.constants"; import type { WorkspaceKernelEventBus } from "#/features/workspaces/kernel/workspace-kernel-events"; -import { WorkspaceKernelFileMigrator } from "#/features/workspaces/kernel/workspace-kernel-file-migrations"; import { getWorkspaceKernelFileShellPath } from "#/features/workspaces/kernel/workspace-kernel-files"; import { parseWorkspaceMetadataJson } from "#/features/workspaces/kernel/workspace-kernel-metadata"; import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; -import type { - KernelItemProjectionRow, - KernelItemRow, -} from "#/features/workspaces/kernel/workspace-kernel-rows"; +import type { KernelItemProjectionRow } from "#/features/workspaces/kernel/workspace-kernel-rows"; import type { WorkspaceKernelStore } from "#/features/workspaces/kernel/workspace-kernel-store"; import type { CreateWorkspaceKernelFileFromUploadArgs, @@ -39,11 +33,9 @@ import { deleteR2Prefix } from "#/lib/r2"; export class WorkspaceKernelFileCommands { private readonly events: WorkspaceKernelEventBus; - private readonly migrator: WorkspaceKernelFileMigrator; private readonly r2: R2Bucket; private readonly sql: WorkspaceKernelSql; private readonly store: WorkspaceKernelStore; - private readonly workspace: ShellWorkspace; private readonly workspaceId: () => string; constructor(input: { @@ -51,20 +43,12 @@ export class WorkspaceKernelFileCommands { r2: R2Bucket; sql: WorkspaceKernelSql; store: WorkspaceKernelStore; - workspace: ShellWorkspace; workspaceId: () => string; }) { this.events = input.events; - this.migrator = new WorkspaceKernelFileMigrator({ - bucket: input.r2, - sql: input.sql, - workspace: input.workspace, - workspaceId: input.workspaceId, - }); this.r2 = input.r2; this.sql = input.sql; this.store = input.store; - this.workspace = input.workspace; this.workspaceId = input.workspaceId; } @@ -174,12 +158,10 @@ export class WorkspaceKernelFileCommands { const contentType = getMetadataString(item.metadataJson, "mimeType"); const originalName = getMetadataString(item.metadataJson, "originalName"); const sizeBytes = getMetadataNumber(item.metadataJson, "sizeBytes"); - const objectKey = - row.object_key ?? - (await this.migrator.migrateSource(row.shell_path, item.id, { - contentType: contentType ?? "application/octet-stream", - expectedSize: sizeBytes, - })); + const objectKey = row.object_key; + if (!objectKey) { + throw new Error("Workspace file source object is missing."); + } return { objectKey, @@ -207,10 +189,10 @@ export class WorkspaceKernelFileCommands { return null; } - const objectKey = - projection.status === "ready" - ? (projection.object_key ?? (await this.migrator.migratePreview(projection, input.itemId))) - : null; + if (projection.status === "ready" && !projection.object_key) { + throw new Error("Ready workspace file preview object is missing."); + } + const objectKey = projection.status === "ready" ? projection.object_key : null; const metadataJson = parseProjectionMetadataJson(projection.metadata_json); return { @@ -249,24 +231,19 @@ export class WorkspaceKernelFileCommands { const now = Date.now(); - await this.writeProjectionRow({ + this.writeProjectionRow({ itemId: input.itemId, projection: input, now, }); } - private async writeProjectionRow(input: { + private writeProjectionRow(input: { createdAt?: number; itemId: string; projection: UpsertWorkspaceKernelFileProjectionArgs; now: number; }) { - const legacyContentPath = this.getProjectionRow({ - itemId: input.itemId, - format: input.projection.format, - })?.content_shell_path; - this.sql` INSERT INTO kernel_item_projections ( item_id, @@ -274,7 +251,6 @@ export class WorkspaceKernelFileCommands { status, provider, provider_mode, - content_shell_path, object_key, error_message, source_hash, @@ -288,7 +264,6 @@ export class WorkspaceKernelFileCommands { ${input.projection.status}, ${input.projection.provider ?? null}, ${input.projection.providerMode ?? null}, - NULL, ${input.projection.objectKey ?? null}, ${input.projection.errorMessage ?? null}, ${input.projection.sourceHash ?? null}, @@ -300,17 +275,12 @@ export class WorkspaceKernelFileCommands { status = excluded.status, provider = excluded.provider, provider_mode = excluded.provider_mode, - content_shell_path = NULL, object_key = excluded.object_key, error_message = excluded.error_message, source_hash = excluded.source_hash, metadata_json = excluded.metadata_json, updated_at = excluded.updated_at `; - - if (legacyContentPath) { - await this.workspace.rm(legacyContentPath, { force: true }); - } } async readFileProjection( @@ -322,22 +292,13 @@ export class WorkspaceKernelFileCommands { throw new Error("Workspace item is not a file."); } - let projection = this.getProjectionRow(input); + const projection = this.getProjectionRow(input); if (!projection) { return null; } - if ( - projection.status === "ready" && - projection.format === "pages" && - !projection.object_key && - projection.content_shell_path - ) { - await this.migrator.migratePageProjection(projection); - projection = this.getProjectionRow(input); - if (!projection) { - return null; - } + if (projection.status === "ready" && !projection.object_key) { + throw new Error("Ready workspace file projection object is missing."); } return { @@ -391,77 +352,6 @@ export class WorkspaceKernelFileCommands { } } - async migrateLegacyStorage() { - const sourceRows = this.sql` - SELECT * FROM kernel_items - WHERE type = 'file' AND object_key IS NULL - `; - const projectionRows = this.sql` - SELECT * FROM kernel_item_projections - WHERE content_shell_path IS NOT NULL - `; - const orphanedItemIds = new Set(); - let orphanedProjectionCount = 0; - - for (const row of sourceRows) { - if (!(await this.workspace.stat(row.shell_path))) { - orphanedItemIds.add(row.id); - const orphanedProjections = projectionRows.filter( - (candidate) => candidate.item_id === row.id && candidate.content_shell_path, - ); - orphanedProjectionCount += orphanedProjections.length; - for (const projection of orphanedProjections) { - await this.workspace.rm(projection.content_shell_path!, { force: true }); - } - this.sql`DELETE FROM kernel_item_projections WHERE item_id = ${row.id}`; - this - .sql`DELETE FROM kernel_relations WHERE from_item_id = ${row.id} OR to_item_id = ${row.id}`; - this.sql`DELETE FROM kernel_items WHERE id = ${row.id}`; - await deleteR2Prefix( - this.r2, - getWorkspaceFileItemObjectPrefix({ workspaceId: this.workspaceId(), itemId: row.id }), - ); - continue; - } - - const metadata = parseWorkspaceMetadataJson(row.metadata_json); - await this.migrator.migrateSource(row.shell_path, row.id, { - contentType: getMetadataString(metadata, "mimeType") ?? "application/octet-stream", - expectedSize: getMetadataNumber(metadata, "sizeBytes"), - }); - } - - for (const projection of projectionRows) { - if (orphanedItemIds.has(projection.item_id)) { - continue; - } - if (projection.format === "pages") { - await this.migrator.migratePageProjection(projection); - } else { - await this.migrator.migratePreview(projection, projection.item_id); - } - } - - const remainingSources = - this.sql<{ count: number }>` - SELECT COUNT(*) AS count FROM kernel_items - WHERE type = 'file' AND object_key IS NULL - `[0]?.count ?? 0; - const remainingProjections = - this.sql<{ count: number }>` - SELECT COUNT(*) AS count FROM kernel_item_projections - WHERE content_shell_path IS NOT NULL - `[0]?.count ?? 0; - - return { - migratedSources: sourceRows.length - orphanedItemIds.size, - migratedProjections: projectionRows.length - orphanedProjectionCount, - orphanedFilesRemoved: orphanedItemIds.size, - remainingSources, - remainingProjections, - }; - } - private async requireObject(objectKey: string) { const object = await this.r2.head(objectKey); diff --git a/src/features/workspaces/kernel/workspace-kernel-file-migrations.ts b/src/features/workspaces/kernel/workspace-kernel-file-migrations.ts deleted file mode 100644 index bfaba023..00000000 --- a/src/features/workspaces/kernel/workspace-kernel-file-migrations.ts +++ /dev/null @@ -1,163 +0,0 @@ -import type { Workspace as ShellWorkspace } from "@cloudflare/shell"; - -import { parseMarkdownPagesProjection } from "#/features/workspaces/extraction/page-markdown-projection"; -import { writeWorkspacePageProjection } from "#/features/workspaces/extraction/workspace-page-projection"; -import { - getWorkspaceFilePreviewObjectKey, - getWorkspaceFileSourceObjectKey, -} from "#/features/workspaces/files/workspace-file-object-keys"; -import { WORKSPACE_FILE_PREVIEW_CONTENT_TYPE } from "#/features/workspaces/files/workspace-file-preview.constants"; -import { parseWorkspaceMetadataJson } from "#/features/workspaces/kernel/workspace-kernel-metadata"; -import type { KernelItemProjectionRow } from "#/features/workspaces/kernel/workspace-kernel-rows"; -import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; - -export class WorkspaceKernelFileMigrator { - private readonly bucket: R2Bucket; - private readonly sql: WorkspaceKernelSql; - private readonly workspace: ShellWorkspace; - private readonly workspaceId: () => string; - - constructor(input: { - bucket: R2Bucket; - sql: WorkspaceKernelSql; - workspace: ShellWorkspace; - workspaceId: () => string; - }) { - this.bucket = input.bucket; - this.sql = input.sql; - this.workspace = input.workspace; - this.workspaceId = input.workspaceId; - } - - async migratePageProjection(projection: KernelItemProjectionRow) { - if (!projection.content_shell_path) { - return; - } - - const content = await this.workspace.readFile(projection.content_shell_path); - if (content === null) { - throw new Error("Legacy page projection content was not found."); - } - - const sourceHash = projection.source_hash ?? `legacy-${projection.updated_at}`; - const reference = await writeWorkspacePageProjection({ - bucket: this.bucket, - itemId: projection.item_id, - metadata: parseWorkspaceMetadataJson(projection.metadata_json), - pages: parseMarkdownPagesProjection(content), - provider: projection.provider ?? "legacy", - providerMode: projection.provider_mode ?? "legacy", - runId: `legacy-${sourceHash}`, - sourceHash, - tier: "legacy", - workspaceId: this.workspaceId(), - }); - - this.sql` - UPDATE kernel_item_projections - SET object_key = ${reference.manifestObjectKey}, content_shell_path = NULL - WHERE item_id = ${projection.item_id} AND format = ${projection.format} - `; - await this.workspace.rm(projection.content_shell_path, { force: true }); - } - - async migrateSource( - shellPath: string, - itemId: string, - input: { contentType: string; expectedSize: number | null }, - ) { - const objectKey = getWorkspaceFileSourceObjectKey({ - workspaceId: this.workspaceId(), - itemId, - }); - const object = await this.copyLegacyFile({ - contentType: input.contentType, - expectedSize: input.expectedSize, - objectKey, - shellPath, - }); - - if (input.expectedSize !== null && object.size !== input.expectedSize) { - await this.bucket.delete(objectKey); - throw new Error("Legacy workspace file migration produced an invalid object."); - } - - this.sql`UPDATE kernel_items SET object_key = ${objectKey} WHERE id = ${itemId}`; - await this.workspace.rm(shellPath, { force: true }); - return objectKey; - } - - async migratePreview(projection: KernelItemProjectionRow, itemId: string) { - if (!projection.content_shell_path) { - return null; - } - if (projection.object_key && (await this.bucket.head(projection.object_key))) { - this.sql` - UPDATE kernel_item_projections - SET content_shell_path = NULL - WHERE item_id = ${itemId} AND format = 'preview' - `; - await this.workspace.rm(projection.content_shell_path, { force: true }); - return projection.object_key; - } - - const objectKey = getWorkspaceFilePreviewObjectKey({ - workspaceId: this.workspaceId(), - itemId, - }); - const object = await this.copyLegacyFile({ - contentType: WORKSPACE_FILE_PREVIEW_CONTENT_TYPE, - expectedSize: null, - objectKey, - shellPath: projection.content_shell_path, - }); - - const metadataJson = { - ...parseWorkspaceMetadataJson(projection.metadata_json), - sizeBytes: object.size, - }; - this.sql` - UPDATE kernel_item_projections - SET object_key = ${objectKey}, - content_shell_path = NULL, - metadata_json = ${JSON.stringify(metadataJson)} - WHERE item_id = ${itemId} AND format = 'preview' - `; - await this.workspace.rm(projection.content_shell_path, { force: true }); - return objectKey; - } - - private async copyLegacyFile(input: { - contentType: string; - expectedSize: number | null; - objectKey: string; - shellPath: string; - }) { - const [stat, stream] = await Promise.all([ - this.workspace.stat(input.shellPath), - this.workspace.readFileStream(input.shellPath), - ]); - - if (!stat || !stream) { - throw new Error("Legacy workspace file content was not found."); - } - if (input.expectedSize !== null && stat.size !== input.expectedSize) { - throw new Error("Legacy workspace file size did not match its metadata."); - } - - const fixedLengthStream = new FixedLengthStream(stat.size); - const [object] = await Promise.all([ - this.bucket.put(input.objectKey, fixedLengthStream.readable, { - httpMetadata: { contentType: input.contentType }, - }), - stream.pipeTo(fixedLengthStream.writable), - ]); - - if (!object || object.size !== stat.size) { - await this.bucket.delete(input.objectKey); - throw new Error("Legacy workspace file migration produced an invalid object."); - } - - return object; - } -} diff --git a/src/features/workspaces/kernel/workspace-kernel-rows.ts b/src/features/workspaces/kernel/workspace-kernel-rows.ts index 2ef9d030..0a64d201 100644 --- a/src/features/workspaces/kernel/workspace-kernel-rows.ts +++ b/src/features/workspaces/kernel/workspace-kernel-rows.ts @@ -39,7 +39,6 @@ export type KernelItemProjectionRow = { status: WorkspaceKernelFileProjectionStatus; provider: string | null; provider_mode: string | null; - content_shell_path: string | null; object_key: string | null; error_message: string | null; source_hash: string | null; diff --git a/src/features/workspaces/kernel/workspace-kernel-schema.ts b/src/features/workspaces/kernel/workspace-kernel-schema.ts index f4c5d290..1fb0c590 100644 --- a/src/features/workspaces/kernel/workspace-kernel-schema.ts +++ b/src/features/workspaces/kernel/workspace-kernel-schema.ts @@ -26,7 +26,6 @@ export function initializeWorkspaceKernelStorage(sql: WorkspaceKernelSql) { deleted_at INTEGER ) `; - ensureKernelItemObjectKeyColumn(sql); sql`CREATE INDEX IF NOT EXISTS kernel_items_parent_idx ON kernel_items (parent_id, deleted_at, sort_order)`; sql`CREATE INDEX IF NOT EXISTS kernel_items_type_idx @@ -39,7 +38,6 @@ export function initializeWorkspaceKernelStorage(sql: WorkspaceKernelSql) { status TEXT NOT NULL, provider TEXT, provider_mode TEXT, - content_shell_path TEXT, object_key TEXT, error_message TEXT, source_hash TEXT, @@ -49,7 +47,6 @@ export function initializeWorkspaceKernelStorage(sql: WorkspaceKernelSql) { PRIMARY KEY (item_id, format) ) `; - ensureProjectionObjectKeyColumn(sql); sql`CREATE INDEX IF NOT EXISTS kernel_item_projections_status_idx ON kernel_item_projections (status, updated_at)`; sql` @@ -90,22 +87,6 @@ export function initializeWorkspaceKernelStorage(sql: WorkspaceKernelSql) { ON kernel_events (revision)`; } -function ensureKernelItemObjectKeyColumn(sql: WorkspaceKernelSql) { - const columns = sql<{ name: string }>`PRAGMA table_info(kernel_items)`; - - if (!columns.some((column) => column.name === "object_key")) { - sql`ALTER TABLE kernel_items ADD COLUMN object_key TEXT`; - } -} - -function ensureProjectionObjectKeyColumn(sql: WorkspaceKernelSql) { - const columns = sql<{ name: string }>`PRAGMA table_info(kernel_item_projections)`; - - if (!columns.some((column) => column.name === "object_key")) { - sql`ALTER TABLE kernel_item_projections ADD COLUMN object_key TEXT`; - } -} - function createSiblingNameIndexes(sql: WorkspaceKernelSql) { try { sql` diff --git a/src/features/workspaces/kernel/workspace-kernel.ts b/src/features/workspaces/kernel/workspace-kernel.ts index 4e02bcdf..9c02b1da 100644 --- a/src/features/workspaces/kernel/workspace-kernel.ts +++ b/src/features/workspaces/kernel/workspace-kernel.ts @@ -86,7 +86,6 @@ export class WorkspaceKernel extends Agent { r2: this.env.WORKSPACE_KERNEL_FILES, sql: this.kernelSql, store: this.store, - workspace: this.workspace, workspaceId: () => this.name, }); @@ -169,10 +168,6 @@ export class WorkspaceKernel extends Agent { return await this.fileCommands.readFileProjection(input); } - async migrateLegacyFileStorage() { - return await this.fileCommands.migrateLegacyStorage(); - } - async renameItem( input: RenameWorkspaceKernelItemArgs, ): Promise> { diff --git a/src/server.ts b/src/server.ts index 7340f180..3f1e7adc 100644 --- a/src/server.ts +++ b/src/server.ts @@ -3,7 +3,6 @@ import handler from "@tanstack/react-start/server-entry"; import { routeUserAIRequest } from "#/features/workspaces/ai/auth"; import { routeDocumentSessionRequest } from "#/features/workspaces/documents/document-session-auth"; import { routeWorkspaceKernelRequest } from "#/features/workspaces/kernel/workspace-kernel-auth"; -import { routeWorkspaceFileR2Migration } from "#/features/workspaces/kernel/workspace-file-r2-migration-route"; import { recordOperationalFailure } from "#/integrations/observability/operational-events"; import { posthogHost, posthogHostOrigin, posthogProjectToken } from "#/integrations/posthog/config"; import { getTelemetryRequestDetails } from "#/integrations/posthog/server-context"; @@ -113,12 +112,6 @@ function withSecurityHeaders(response: Response) { export default { async fetch(request, env) { try { - const migrationResponse = await routeWorkspaceFileR2Migration(request, env); - - if (migrationResponse) { - return migrationResponse; - } - const chatResponse = await routeUserAIRequest(request, env); if (chatResponse) { From f6afa33d3f0dee6560c36294a3ed2f9aba99fe5d Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Tue, 14 Jul 2026 16:08:03 -0400 Subject: [PATCH 3/7] fix(workspaces): harden upload and extraction boundaries Close review-found failures in upload replay, empty output, and processor errors. Make projection publication atomic and bound streaming reads. Tests: pnpm check; pnpm test; pnpm exec wrangler types --check --- .dev.vars.example | 3 + containers/liteparse/server.mjs | 9 ++- docs/configuration/deployments.mdx | 2 +- .../container-file-conversion.test.ts | 69 +++++++++++++++++++ .../conversion/container-file-conversion.ts | 39 ++++++++--- .../extraction/liteparse-projection.ts | 34 +++++---- .../extraction/providers/liteparse.test.ts | 30 +++++++- .../extraction/providers/liteparse.ts | 22 ++++++ .../workspace-file-extraction-workflow.ts | 28 +++++--- .../workspace-file-preview-projection.ts | 35 ++++++---- .../workspace-page-projection.test.ts | 35 +++++++++- .../extraction/workspace-page-projection.ts | 32 +++++++-- .../files/workspace-file-processor.ts | 7 ++ .../kernel/workspace-kernel-file-commands.ts | 10 +-- .../workspaces/operations/read-items.ts | 3 +- .../operations/read-page-selection.test.ts | 15 ++-- .../workspaces/read-page-selection.ts | 6 ++ .../upload/pdf-upload-validation.test.ts | 21 ++++++ .../upload/pdf-upload-validation.ts | 15 ++++ .../workspace-file-direct-upload.test.ts | 6 +- .../upload/workspace-file-direct-upload.ts | 4 +- .../upload/workspace-file-upload-storage.ts | 3 + src/lib/http/streaming-multipart.ts | 2 +- .../v1/workspaces.$workspaceId.file-upload.ts | 5 +- worker-configuration.d.ts | 7 +- wrangler.jsonc | 3 + 26 files changed, 367 insertions(+), 78 deletions(-) create mode 100644 src/features/workspaces/conversion/container-file-conversion.test.ts diff --git a/.dev.vars.example b/.dev.vars.example index 1ecdaa76..47c05003 100644 --- a/.dev.vars.example +++ b/.dev.vars.example @@ -38,3 +38,6 @@ BETTER_AUTH_URL="http://localhost:3000" # Document extraction (LlamaParse). # LLAMA_CLOUD_API_KEY="" + +# Direct workspace-upload completion claims. Use a separate random 32+ byte value. +# WORKSPACE_UPLOAD_TOKEN_SECRET="" diff --git a/containers/liteparse/server.mjs b/containers/liteparse/server.mjs index fba95aa2..2eb3e712 100644 --- a/containers/liteparse/server.mjs +++ b/containers/liteparse/server.mjs @@ -1,5 +1,6 @@ import { LiteParse } from "@llamaindex/liteparse"; import { execFile } from "node:child_process"; +import { once } from "node:events"; import { createWriteStream } from "node:fs"; import { mkdtemp, readFile, rm } from "node:fs/promises"; import { createServer } from "node:http"; @@ -66,7 +67,13 @@ createServer(async (request, response) => { status = 200; response.writeHead(status, { "content-type": "application/x-ndjson; charset=utf-8" }); for (const page of result.pages) { - response.write(`${JSON.stringify({ markdown: page.markdown, pageNumber: page.pageNum })}\n`); + if ( + !response.write( + `${JSON.stringify({ markdown: page.markdown, pageNumber: page.pageNum })}\n`, + ) + ) { + await once(response, "drain"); + } } return response.end(); } catch (error) { diff --git a/docs/configuration/deployments.mdx b/docs/configuration/deployments.mdx index f540482c..b22bfe65 100644 --- a/docs/configuration/deployments.mdx +++ b/docs/configuration/deployments.mdx @@ -54,7 +54,7 @@ vp run deploy:worker:staging Runtime and deploy secrets are synced into GitHub Actions secrets from Infisical. Production and staging workflows require Cloudflare credentials, PostHog variables, and the runtime secrets declared by the Worker configuration. -Direct workspace uploads additionally require bucket-scoped R2 Object Read & Write credentials in `R2_ACCESS_KEY_ID` and `R2_SECRET_ACCESS_KEY`, plus `R2_ACCOUNT_ID`. Configure browser PUT access from the checked-in policies: +Direct workspace uploads additionally require bucket-scoped R2 Object Read & Write credentials in `R2_ACCESS_KEY_ID` and `R2_SECRET_ACCESS_KEY`, `R2_ACCOUNT_ID`, and an independent random `WORKSPACE_UPLOAD_TOKEN_SECRET` for completion claims. Configure browser PUT access from the checked-in policies: ```bash wrangler r2 bucket cors set thinkex-workspace-kernel-files-staging --file config/r2/cors.staging.json diff --git a/src/features/workspaces/conversion/container-file-conversion.test.ts b/src/features/workspaces/conversion/container-file-conversion.test.ts new file mode 100644 index 00000000..2867d0e2 --- /dev/null +++ b/src/features/workspaces/conversion/container-file-conversion.test.ts @@ -0,0 +1,69 @@ +import { beforeAll, describe, expect, it, vi } from "vitest"; + +import { convertFileStreamWithContainer } from "#/features/workspaces/conversion/container-file-conversion"; + +beforeAll(() => { + vi.stubGlobal( + "FixedLengthStream", + class { + readonly readable: ReadableStream; + readonly writable: WritableStream; + + constructor() { + const stream = new TransformStream(); + this.readable = stream.readable; + this.writable = stream.writable; + } + }, + ); +}); + +describe("container file conversion", () => { + it("rejects an empty successful conversion response", async () => { + const container = createContainer(new Uint8Array()); + + await expect(convert(container)).rejects.toThrow("Conversion returned no content"); + }); + + it("preserves a non-empty conversion response", async () => { + const output = new Uint8Array([4, 5, 6]); + const container = createContainer(output); + + const response = await convert(container); + + expect(new Uint8Array(await response.arrayBuffer())).toEqual(output); + expect(container.startAndWaitForPorts).toHaveBeenCalledOnce(); + }); +}); + +function convert(container: ReturnType) { + return convertFileStreamWithContainer({ + body: stream(new Uint8Array([1, 2, 3])), + container, + contentType: "application/octet-stream", + emptyMessage: "Conversion returned no content", + error: (message) => new Error(message), + fileName: "input.bin", + formFieldName: "file", + sizeBytes: 3, + url: "http://container/convert", + }); +} + +function createContainer(output: Uint8Array) { + return { + fetch: vi.fn(async (request: Request) => { + await request.arrayBuffer(); + return new Response(output.slice().buffer); + }), + startAndWaitForPorts: vi.fn(async () => undefined), + }; +} + +function stream(bytes: Uint8Array) { + const body = new Response(bytes.slice().buffer).body; + if (!body) { + throw new Error("Test stream was not created."); + } + return body; +} diff --git a/src/features/workspaces/conversion/container-file-conversion.ts b/src/features/workspaces/conversion/container-file-conversion.ts index b83da391..78438aa1 100644 --- a/src/features/workspaces/conversion/container-file-conversion.ts +++ b/src/features/workspaces/conversion/container-file-conversion.ts @@ -1,3 +1,5 @@ +import { createStreamingMultipartFile } from "#/lib/http/streaming-multipart"; + type FileConversionContainer = { fetch(request: Request): Promise; startAndWaitForPorts(input: { @@ -54,15 +56,17 @@ export async function convertFileStreamWithContainer(input: { sizeBytes: input.sizeBytes, }); - const response = await input.container.fetch( - new Request(input.url, { - body: multipart.body, - duplex: "half", - headers: { "content-type": multipart.contentType }, - method: "POST", - } as RequestInit & { duplex: "half" }), - ); - await multipart.done; + const [response] = await Promise.all([ + input.container.fetch( + new Request(input.url, { + body: multipart.body, + duplex: "half", + headers: { "content-type": multipart.contentType }, + method: "POST", + } as RequestInit & { duplex: "half" }), + ), + multipart.done, + ]); if (!response.ok) { throw input.error(await getConversionErrorMessage(response)); @@ -72,7 +76,21 @@ export async function convertFileStreamWithContainer(input: { throw input.error(input.emptyMessage); } - return response; + return requireNonEmptyResponse(response, () => input.error(input.emptyMessage)); +} + +async function requireNonEmptyResponse(response: Response, createError: () => Error) { + const [probe, body] = response.body!.tee(); + const reader = probe.getReader(); + const first = await reader.read(); + void reader.cancel().catch(() => undefined); + + if (first.done || first.value.byteLength === 0) { + await body.cancel().catch(() => undefined); + throw createError(); + } + + return new Response(body, response); } async function getConversionErrorMessage(response: Response) { @@ -82,4 +100,3 @@ async function getConversionErrorMessage(response: Response) { return message ? `${fallback} ${message}` : fallback; } -import { createStreamingMultipartFile } from "#/lib/http/streaming-multipart"; diff --git a/src/features/workspaces/extraction/liteparse-projection.ts b/src/features/workspaces/extraction/liteparse-projection.ts index 605e6e0c..07bc046b 100644 --- a/src/features/workspaces/extraction/liteparse-projection.ts +++ b/src/features/workspaces/extraction/liteparse-projection.ts @@ -6,7 +6,10 @@ import type { WorkspaceFileExtractionWorkflowParams, } from "#/features/workspaces/extraction/types"; import { getWorkspaceFileSourceObject } from "#/features/workspaces/extraction/workspace-file-source"; -import { writeWorkspacePageProjection } from "#/features/workspaces/extraction/workspace-page-projection"; +import { + commitWorkspacePageProjection, + writeWorkspacePageProjection, +} from "#/features/workspaces/extraction/workspace-page-projection"; import { getWorkspaceKernelFromEnv } from "#/features/workspaces/kernel/workspace-kernel-access"; export async function publishLiteParseProjection( @@ -51,20 +54,23 @@ export async function publishLiteParseProjection( workspaceId: params.workspaceId, }); - await kernel.upsertFileProjection({ - itemId: params.itemId, - format: "pages", - status: "ready", - objectKey: projection.manifestObjectKey, - provider: "liteparse", - providerMode: "fast", - sourceHash: object.etag, - metadataJson: { - markdownLength: projection.manifest.markdownLength, - pageCount: projection.manifest.pageCount, - provisional: true, + await commitWorkspacePageProjection({ + bucket: env.WORKSPACE_KERNEL_FILES, + kernel, + manifestObjectKey: projection.manifestObjectKey, + mutation: { + itemId: params.itemId, + format: "pages", + provider: "liteparse", + providerMode: "fast", + sourceHash: object.etag, + metadataJson: { + markdownLength: projection.manifest.markdownLength, + pageCount: projection.manifest.pageCount, + provisional: true, + }, + actorUserId: params.actorUserId, }, - actorUserId: params.actorUserId, }); return { diff --git a/src/features/workspaces/extraction/providers/liteparse.test.ts b/src/features/workspaces/extraction/providers/liteparse.test.ts index 9e50131a..7e4ee1e0 100644 --- a/src/features/workspaces/extraction/providers/liteparse.test.ts +++ b/src/features/workspaces/extraction/providers/liteparse.test.ts @@ -1,6 +1,12 @@ -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; +import { extractPdfWithLiteParse } from "#/features/workspaces/extraction/providers/liteparse"; import { parseLiteParsePage } from "#/features/workspaces/extraction/providers/liteparse-response"; +import { requestWorkspaceFileProcessor } from "#/features/workspaces/files/workspace-file-processor"; + +vi.mock("#/features/workspaces/files/workspace-file-processor", () => ({ + requestWorkspaceFileProcessor: vi.fn(), +})); describe("LiteParse response parsing", () => { it("accepts one canonical page record", () => { @@ -26,4 +32,26 @@ describe("LiteParse response parsing", () => { ])("rejects malformed container responses", (payload) => { expect(() => parseLiteParsePage(payload)).toThrow("LiteParse returned an invalid"); }); + + it("rejects an oversized NDJSON record and cancels the processor response", async () => { + const cancel = vi.fn(); + const body = new ReadableStream({ + cancel, + start(controller) { + controller.enqueue(new Uint8Array(8 * 1024 * 1024 + 1)); + }, + }); + vi.mocked(requestWorkspaceFileProcessor).mockResolvedValue({ + body, + ok: true, + } as Response); + const pages = extractPdfWithLiteParse({} as Cloudflare.Env, { + body: new ReadableStream(), + fileName: "document.pdf", + sizeBytes: 1, + }); + + await expect(pages.next()).rejects.toThrow("oversized NDJSON page"); + expect(cancel).toHaveBeenCalledOnce(); + }); }); diff --git a/src/features/workspaces/extraction/providers/liteparse.ts b/src/features/workspaces/extraction/providers/liteparse.ts index 6e79f3d3..88b1e91f 100644 --- a/src/features/workspaces/extraction/providers/liteparse.ts +++ b/src/features/workspaces/extraction/providers/liteparse.ts @@ -2,6 +2,8 @@ import type { MarkdownProjectionPage } from "#/features/workspaces/extraction/pa import { parseLiteParsePage } from "#/features/workspaces/extraction/providers/liteparse-response"; import { requestWorkspaceFileProcessor } from "#/features/workspaces/files/workspace-file-processor"; +const maxNdjsonLineBytes = 8 * 1024 * 1024; + export async function* extractPdfWithLiteParse( env: Cloudflare.Env, input: { @@ -41,10 +43,29 @@ async function* readNdjsonLines(body: ReadableStream): AsyncGenerato const reader = body.getReader(); const decoder = new TextDecoder(); let buffer = ""; + let pendingLineBytes = 0; try { while (true) { const { done, value } = await reader.read(); + + if (value) { + let offset = 0; + while (offset < value.byteLength) { + const newlineOffset = value.indexOf(0x0a, offset); + const segmentEnd = newlineOffset === -1 ? value.byteLength : newlineOffset; + pendingLineBytes += segmentEnd - offset; + if (pendingLineBytes > maxNdjsonLineBytes) { + throw new Error("LiteParse returned an oversized NDJSON page."); + } + if (newlineOffset === -1) { + break; + } + pendingLineBytes = 0; + offset = newlineOffset + 1; + } + } + buffer += decoder.decode(value, { stream: !done }); const lines = buffer.split("\n"); buffer = lines.pop() ?? ""; @@ -64,6 +85,7 @@ async function* readNdjsonLines(body: ReadableStream): AsyncGenerato yield buffer; } } finally { + await reader.cancel().catch(() => undefined); reader.releaseLock(); } } diff --git a/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts b/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts index 7f4c4035..3ea2cacf 100644 --- a/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts +++ b/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts @@ -10,7 +10,10 @@ import type { WorkspaceFileExtractionWorkflowParams, } from "#/features/workspaces/extraction/types"; import { getWorkspaceFileSourceObject } from "#/features/workspaces/extraction/workspace-file-source"; -import { writeWorkspacePageProjection } from "#/features/workspaces/extraction/workspace-page-projection"; +import { + commitWorkspacePageProjection, + writeWorkspacePageProjection, +} from "#/features/workspaces/extraction/workspace-page-projection"; import { getWorkspaceKernelFromEnv } from "#/features/workspaces/kernel/workspace-kernel-access"; import { getWorkspaceUploadFamily } from "#/features/workspaces/model/workspace-file"; @@ -124,16 +127,19 @@ export class WorkspaceFileExtractionWorkflow extends WorkflowEntrypoint< markdownLength: extraction.markdownLength, }; - await kernel.upsertFileProjection({ - itemId: params.itemId, - format: "pages", - status: "ready", - objectKey: extraction.manifestObjectKey, - provider: extraction.provider, - providerMode: extraction.providerMode, - sourceHash: extraction.sourceHash, - metadataJson, - actorUserId: params.actorUserId, + await commitWorkspacePageProjection({ + bucket: this.env.WORKSPACE_KERNEL_FILES, + kernel, + manifestObjectKey: extraction.manifestObjectKey, + mutation: { + itemId: params.itemId, + format: "pages", + provider: extraction.provider, + providerMode: extraction.providerMode, + sourceHash: extraction.sourceHash, + metadataJson, + actorUserId: params.actorUserId, + }, }); return { diff --git a/src/features/workspaces/extraction/workspace-file-preview-projection.ts b/src/features/workspaces/extraction/workspace-file-preview-projection.ts index fdf24e23..b4a1f48c 100644 --- a/src/features/workspaces/extraction/workspace-file-preview-projection.ts +++ b/src/features/workspaces/extraction/workspace-file-preview-projection.ts @@ -16,11 +16,25 @@ export async function publishWorkspaceFilePreview( step: WorkflowStep, params: WorkspaceFileExtractionWorkflowParams, ) { - return step.do( - "publish workspace file preview", - { retries: { limit: 1, delay: "5 seconds", backoff: "constant" }, timeout: "2 minutes" }, - async () => generateWorkspaceFilePreview(env, params), - ); + try { + return await step.do( + "publish workspace file preview", + { retries: { limit: 1, delay: "5 seconds", backoff: "constant" }, timeout: "2 minutes" }, + async () => generateWorkspaceFilePreview(env, params), + ); + } catch (error) { + return step.do("record workspace file preview failure", async () => { + const kernel = await getWorkspaceKernelFromEnv(env, params.workspaceId); + await kernel.upsertFileProjection({ + itemId: params.itemId, + format: "preview", + status: "failed", + errorMessage: error instanceof Error ? error.message : String(error), + actorUserId: params.actorUserId, + }); + return { outcome: "error" as const }; + }); + } } async function generateWorkspaceFilePreview( @@ -81,15 +95,8 @@ async function generateWorkspaceFilePreview( return { outcome: "success" as const }; } catch (error) { failure = error; - await env.WORKSPACE_KERNEL_FILES.delete(objectKey); - await kernel.upsertFileProjection({ - itemId: params.itemId, - format: "preview", - status: "failed", - errorMessage: error instanceof Error ? error.message : String(error), - actorUserId: params.actorUserId, - }); - return { outcome: "error" as const }; + await env.WORKSPACE_KERNEL_FILES.delete(objectKey).catch(() => undefined); + throw error; } finally { recordOperationalOutcome({ distinctId: params.actorUserId ?? undefined, diff --git a/src/features/workspaces/extraction/workspace-page-projection.test.ts b/src/features/workspaces/extraction/workspace-page-projection.test.ts index ea8c0eaa..d4f920f4 100644 --- a/src/features/workspaces/extraction/workspace-page-projection.test.ts +++ b/src/features/workspaces/extraction/workspace-page-projection.test.ts @@ -1,6 +1,7 @@ -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { + commitWorkspacePageProjection, getWorkspacePageObjectKey, readWorkspacePageProjection, writeWorkspacePageProjection, @@ -139,6 +140,38 @@ describe("workspace page projections", () => { ).rejects.toThrow("Extracted pages must be ordered"); expect(storage.values.size).toBe(0); }); + + it("removes a completed projection when its kernel publication fails", async () => { + const storage = createObjectStorage(); + const reference = await writeWorkspacePageProjection({ + bucket: storage.bucket, + itemId: "item-1", + pages: [{ pageNumber: 1, markdown: "First" }], + provider: "liteparse", + providerMode: "fast", + runId: "run-1", + sourceHash: "etag-1", + tier: "fast", + workspaceId: "workspace-1", + }); + const publicationError = new Error("Kernel publication failed"); + + await expect( + commitWorkspacePageProjection({ + bucket: storage.bucket, + kernel: { upsertFileProjection: vi.fn().mockRejectedValue(publicationError) }, + manifestObjectKey: reference.manifestObjectKey, + mutation: { + format: "pages", + itemId: "item-1", + provider: "liteparse", + providerMode: "fast", + sourceHash: "etag-1", + }, + }), + ).rejects.toBe(publicationError); + expect(storage.values.size).toBe(0); + }); }); function createObjectStorage() { diff --git a/src/features/workspaces/extraction/workspace-page-projection.ts b/src/features/workspaces/extraction/workspace-page-projection.ts index 8148ef0c..28ff4b47 100644 --- a/src/features/workspaces/extraction/workspace-page-projection.ts +++ b/src/features/workspaces/extraction/workspace-page-projection.ts @@ -1,6 +1,8 @@ import { jsonValueSchema, type JsonValue } from "#/features/workspaces/contracts"; import type { MarkdownProjectionPage } from "#/features/workspaces/extraction/page-markdown-projection"; import { getWorkspaceFileItemObjectPrefix } from "#/features/workspaces/files/workspace-file-object-keys"; +import type { WorkspaceKernelClient } from "#/features/workspaces/kernel/workspace-kernel-access"; +import type { UpsertWorkspaceKernelFileProjectionArgs } from "#/features/workspaces/kernel/workspace-kernel-types"; import { parseWorkspacePageRange, WorkspacePageSelectionError, @@ -13,7 +15,6 @@ const pageNumberWidth = 6; const pageWriteConcurrency = 8; const maxProjectionPages = 2_000; const maxPageMarkdownBytes = 1024 * 1024; -const maxPageReadCount = 20; const maxPageReadBytes = 2 * 1024 * 1024; export interface WorkspacePageProjectionManifest { @@ -36,6 +37,31 @@ export interface WorkspacePageProjectionReference { manifest: WorkspacePageProjectionManifest; } +type ReadyProjectionMutation = Omit< + Extract, + "objectKey" | "status" +>; + +export async function commitWorkspacePageProjection(input: { + bucket: R2Bucket; + kernel: Pick; + manifestObjectKey: string; + mutation: ReadyProjectionMutation; +}) { + try { + await input.kernel.upsertFileProjection({ + ...input.mutation, + objectKey: input.manifestObjectKey, + status: "ready", + }); + } catch (error) { + await deleteR2Prefix(input.bucket, getManifestPrefix(input.manifestObjectKey)).catch( + () => undefined, + ); + throw error; + } +} + export async function writeWorkspacePageProjection(input: { bucket: R2Bucket; itemId: string; @@ -125,10 +151,6 @@ export async function readWorkspacePageProjection(input: { const requested = input.pages?.trim() || "1"; const selectedPageNumbers = parseWorkspacePageRange(requested, manifest.pageCount); - if (selectedPageNumbers.length > maxPageReadCount) { - throw new WorkspacePageSelectionError("page_selection_too_large"); - } - const prefix = getManifestPrefix(input.manifestObjectKey); const objects = await Promise.all( selectedPageNumbers.map(async (pageNumber) => { diff --git a/src/features/workspaces/files/workspace-file-processor.ts b/src/features/workspaces/files/workspace-file-processor.ts index b24128db..f738a183 100644 --- a/src/features/workspaces/files/workspace-file-processor.ts +++ b/src/features/workspaces/files/workspace-file-processor.ts @@ -2,6 +2,12 @@ import { Container, getRandom } from "@cloudflare/containers"; const workspaceFileProcessorPort = 8080; const workspaceFileProcessorPoolSize = 2; +const processorRequestTimeoutMs = { + "/parse/pdf": 10 * 60_000, + "/preview/image": 2 * 60_000, + "/preview/pdf": 2 * 60_000, + "/validate/pdf": 2 * 60_000, +} as const; export class WorkspaceFileProcessor extends Container { defaultPort = workspaceFileProcessorPort; @@ -40,6 +46,7 @@ export async function requestWorkspaceFileProcessor( duplex: "half", headers, method: "POST", + signal: AbortSignal.timeout(processorRequestTimeoutMs[input.path]), } as RequestInit & { duplex: "half" }), ); } diff --git a/src/features/workspaces/kernel/workspace-kernel-file-commands.ts b/src/features/workspaces/kernel/workspace-kernel-file-commands.ts index 18ff034f..04a36ea1 100644 --- a/src/features/workspaces/kernel/workspace-kernel-file-commands.ts +++ b/src/features/workspaces/kernel/workspace-kernel-file-commands.ts @@ -213,19 +213,19 @@ export class WorkspaceKernelFileCommands { if (row.type !== "file") { throw new Error("Workspace item is not a file."); } - if (input.status === "ready" && input.format === "pages") { + if (input.status === "ready") { if (!row.object_key) { - throw new Error("Ready page projections require a source hash and manifest object."); + throw new Error("Ready file projections require a current source object."); } - const [source, manifest] = await Promise.all([ + const [source, projectionObject] = await Promise.all([ this.r2.head(row.object_key), this.r2.head(input.objectKey), ]); if (!source || source.etag !== input.sourceHash) { throw new Error("The file source changed before its extraction could be published."); } - if (!manifest) { - throw new Error("The page projection manifest was not found."); + if (!projectionObject) { + throw new Error("The file projection object was not found."); } } diff --git a/src/features/workspaces/operations/read-items.ts b/src/features/workspaces/operations/read-items.ts index a79314cd..1a615715 100644 --- a/src/features/workspaces/operations/read-items.ts +++ b/src/features/workspaces/operations/read-items.ts @@ -1,3 +1,5 @@ +import { env } from "cloudflare:workers"; + import { getWorkspaceOperationContext, resolveWorkspaceOperationPath, @@ -293,4 +295,3 @@ function truncateWorkspaceMarkdownLine(line: string) { value: line.slice(0, MAX_WORKSPACE_READ_LINE_LENGTH) + TRUNCATED_LINE_SUFFIX, }; } -import { env } from "cloudflare:workers"; diff --git a/src/features/workspaces/operations/read-page-selection.test.ts b/src/features/workspaces/operations/read-page-selection.test.ts index 1c476538..4e6d3487 100644 --- a/src/features/workspaces/operations/read-page-selection.test.ts +++ b/src/features/workspaces/operations/read-page-selection.test.ts @@ -1,9 +1,6 @@ import { describe, expect, it } from "vitest"; -import { - parseWorkspacePageRange, - WorkspacePageSelectionError, -} from "#/features/workspaces/read-page-selection"; +import { parseWorkspacePageRange } from "#/features/workspaces/read-page-selection"; describe("workspace read page selection", () => { it("returns sorted unique page numbers", () => { @@ -11,6 +8,14 @@ describe("workspace read page selection", () => { }); it("rejects out-of-range page requests", () => { - expect(() => parseWorkspacePageRange("4", 3)).toThrow(WorkspacePageSelectionError); + expect(() => parseWorkspacePageRange("4", 3)).toThrow( + expect.objectContaining({ code: "page_range_out_of_range" }), + ); + }); + + it("rejects selections larger than the shared read limit", () => { + expect(() => parseWorkspacePageRange("1-21", 21)).toThrow( + expect.objectContaining({ code: "page_selection_too_large" }), + ); }); }); diff --git a/src/features/workspaces/read-page-selection.ts b/src/features/workspaces/read-page-selection.ts index 277ddeca..7ab492a0 100644 --- a/src/features/workspaces/read-page-selection.ts +++ b/src/features/workspaces/read-page-selection.ts @@ -4,6 +4,8 @@ export interface WorkspaceReadPages { total: number; } +export const maxWorkspacePageReadCount = 20; + export class WorkspacePageSelectionError extends Error { constructor(readonly code: "page_range_out_of_range" | "page_selection_too_large") { super(code); @@ -36,6 +38,10 @@ export function parseWorkspacePageRange(value: string, totalPages: number) { for (let pageNumber = start; pageNumber <= end; pageNumber += 1) { selected.add(pageNumber); + + if (selected.size > maxWorkspacePageReadCount) { + throw new WorkspacePageSelectionError("page_selection_too_large"); + } } } diff --git a/src/features/workspaces/upload/pdf-upload-validation.test.ts b/src/features/workspaces/upload/pdf-upload-validation.test.ts index c3f53114..9770b8ad 100644 --- a/src/features/workspaces/upload/pdf-upload-validation.test.ts +++ b/src/features/workspaces/upload/pdf-upload-validation.test.ts @@ -44,6 +44,27 @@ describe("PDF upload validation", () => { }), ); }); + + it("preserves processor failures as retryable server errors", async () => { + requestWorkspaceFileProcessor.mockResolvedValue( + Response.json({ error: "temporary failure" }, { status: 500 }), + ); + + await expect(validate(new Uint8Array([1, 2, 3]))).rejects.toThrow( + "Workspace file processor failed with status 500", + ); + }); + + it("preserves the processor upload limit response", async () => { + requestWorkspaceFileProcessor.mockResolvedValue( + Response.json({ code: "UPLOAD_TOO_LARGE" }, { status: 413 }), + ); + + await expect(validate(new Uint8Array([1, 2, 3]))).rejects.toMatchObject({ + code: "UPLOAD_TOO_LARGE", + status: 413, + }); + }); }); function validate(bytes: Uint8Array) { diff --git a/src/features/workspaces/upload/pdf-upload-validation.ts b/src/features/workspaces/upload/pdf-upload-validation.ts index a30f183d..27d6270c 100644 --- a/src/features/workspaces/upload/pdf-upload-validation.ts +++ b/src/features/workspaces/upload/pdf-upload-validation.ts @@ -18,6 +18,18 @@ export async function assertReadablePdfUpload(input: { const payload = await readValidationFailure(response); + if (response.status >= 500) { + throw new Error(`Workspace file processor failed with status ${response.status}.`); + } + + if (response.status === 413 || payload.code === "UPLOAD_TOO_LARGE") { + throw new WorkspaceFileUploadError({ + code: "UPLOAD_TOO_LARGE", + message: "This PDF exceeds the supported upload limit.", + status: 413, + }); + } + if (payload.code === "PASSWORD_PROTECTED_PDF") { throw new WorkspaceFileUploadError({ code: "PASSWORD_PROTECTED_PDF", @@ -26,6 +38,9 @@ export async function assertReadablePdfUpload(input: { status: 422, }); } + if (payload.code !== "INVALID_PDF" && payload.code !== "INVALID_FILE") { + throw new Error(`Workspace file processor returned status ${response.status}.`); + } throw new WorkspaceFileUploadError({ code: "INVALID_PDF", diff --git a/src/features/workspaces/upload/workspace-file-direct-upload.test.ts b/src/features/workspaces/upload/workspace-file-direct-upload.test.ts index 3647443a..5efb8e7e 100644 --- a/src/features/workspaces/upload/workspace-file-direct-upload.test.ts +++ b/src/features/workspaces/upload/workspace-file-direct-upload.test.ts @@ -65,8 +65,9 @@ describe("workspace direct upload sessions", () => { userId: "user-1", workspaceId: "workspace-1", }); - const lastCharacter = session.completionToken.at(-1); - const tampered = `${session.completionToken.slice(0, -1)}${lastCharacter === "x" ? "y" : "x"}`; + const [payload, signature] = session.completionToken.split(".") as [string, string]; + const tamperedPayload = `${payload.startsWith("x") ? "y" : "x"}${payload.slice(1)}`; + const tampered = `${tamperedPayload}.${signature}`; await expect(verifyWorkspaceDirectUploadToken(env, tampered)).rejects.toThrow("invalid"); @@ -119,5 +120,6 @@ function createEnv() { R2_ACCOUNT_ID: "account-id", R2_BUCKET_NAME: "thinkex-workspace-kernel-files", R2_SECRET_ACCESS_KEY: "secret-key", + WORKSPACE_UPLOAD_TOKEN_SECRET: "upload-token-secret", } as Cloudflare.Env; } diff --git a/src/features/workspaces/upload/workspace-file-direct-upload.ts b/src/features/workspaces/upload/workspace-file-direct-upload.ts index 8c720cea..60a8bc25 100644 --- a/src/features/workspaces/upload/workspace-file-direct-upload.ts +++ b/src/features/workspaces/upload/workspace-file-direct-upload.ts @@ -41,7 +41,7 @@ export async function createWorkspaceDirectUploadSession( }); return { - completionToken: await signUploadClaims(env.R2_SECRET_ACCESS_KEY, claims), + completionToken: await signUploadClaims(env.WORKSPACE_UPLOAD_TOKEN_SECRET, claims), uploadUrl, }; } @@ -56,7 +56,7 @@ export async function verifyWorkspaceDirectUploadToken( throw new Error("Upload completion token is invalid."); } - const key = await createSigningKey(env.R2_SECRET_ACCESS_KEY, ["verify"]); + const key = await createSigningKey(env.WORKSPACE_UPLOAD_TOKEN_SECRET, ["verify"]); const valid = await crypto.subtle.verify( "HMAC", key, diff --git a/src/features/workspaces/upload/workspace-file-upload-storage.ts b/src/features/workspaces/upload/workspace-file-upload-storage.ts index 7faf93b6..da2bdb2c 100644 --- a/src/features/workspaces/upload/workspace-file-upload-storage.ts +++ b/src/features/workspaces/upload/workspace-file-upload-storage.ts @@ -72,6 +72,9 @@ export async function storeWorkspaceFileUpload(input: { if (!stored) { throw new Error("Workspace file could not be stored."); } + if (stored.size === 0) { + throw new Error("Workspace file conversion produced an empty file."); + } if (!conversion && stored.size !== input.fileSize) { throw new Error("Stored workspace file size did not match the upload request."); diff --git a/src/lib/http/streaming-multipart.ts b/src/lib/http/streaming-multipart.ts index b820e7e3..e9cce862 100644 --- a/src/lib/http/streaming-multipart.ts +++ b/src/lib/http/streaming-multipart.ts @@ -50,7 +50,7 @@ async function pipeMultipartBody( await writer.write(suffix); await writer.close(); } catch (error) { - await writer.abort(error).catch(() => undefined); + await Promise.allSettled([reader.cancel(error), writer.abort(error)]); throw error; } finally { reader.releaseLock(); diff --git a/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts b/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts index 1a611014..58cfcf8b 100644 --- a/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts +++ b/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts @@ -116,6 +116,7 @@ async function finalizeWorkspaceFileUpload( let stagingObjectKey: string | null = null; let finalObjectKey: string | null = null; let completionClaimKey: string | null = null; + let uploadCompleted = false; try { const userId = await authorizeWorkspaceUpload(request, workspaceId); @@ -159,6 +160,7 @@ async function finalizeWorkspaceFileUpload( }); observation.itemId = command.result.id; observation.outputBytes = claims.fileSize; + uploadCompleted = true; return apiJson(command, requestId); } @@ -192,6 +194,7 @@ async function finalizeWorkspaceFileUpload( finalObjectKey = null; observation.itemId = command.result.id; + uploadCompleted = true; await queueWorkspaceFileExtraction(upload, { itemId: command.result.id, requestId, @@ -209,7 +212,7 @@ async function finalizeWorkspaceFileUpload( } await Promise.allSettled([ finalObjectKey ? env.WORKSPACE_KERNEL_FILES.delete(finalObjectKey) : Promise.resolve(), - completionClaimKey + completionClaimKey && !uploadCompleted ? env.WORKSPACE_KERNEL_FILES.delete(completionClaimKey) : Promise.resolve(), ]); diff --git a/worker-configuration.d.ts b/worker-configuration.d.ts index 6e577d4b..a9e7da70 100644 --- a/worker-configuration.d.ts +++ b/worker-configuration.d.ts @@ -1,5 +1,5 @@ /* eslint-disable */ -// Generated by Wrangler by running `wrangler types` (hash: 66ce1a1950974042862fea1ba4c75c2b) +// Generated by Wrangler by running `wrangler types` (hash: 58de4e88b517f3692dd84708e9339a59) // Runtime types generated with workerd@1.20260701.1 2026-03-24 nodejs_compat interface __BaseEnv_Env { WORKSPACE_KERNEL_FILES: R2Bucket; @@ -24,6 +24,7 @@ interface __BaseEnv_Env { R2_ACCOUNT_ID: string; R2_ACCESS_KEY_ID: string; R2_SECRET_ACCESS_KEY: string; + WORKSPACE_UPLOAD_TOKEN_SECRET: string; UserAIStore: DurableObjectNamespace; WorkspaceKernel: DurableObjectNamespace; DocumentSession: DurableObjectNamespace; @@ -61,6 +62,7 @@ declare namespace Cloudflare { R2_ACCOUNT_ID: string; R2_ACCESS_KEY_ID: string; R2_SECRET_ACCESS_KEY: string; + WORKSPACE_UPLOAD_TOKEN_SECRET: string; UserAIStore: DurableObjectNamespace; WorkspaceKernel: DurableObjectNamespace; DocumentSession: DurableObjectNamespace; @@ -93,6 +95,7 @@ declare namespace Cloudflare { R2_ACCOUNT_ID: string; R2_ACCESS_KEY_ID: string; R2_SECRET_ACCESS_KEY: string; + WORKSPACE_UPLOAD_TOKEN_SECRET: string; UserAIStore: DurableObjectNamespace; WorkspaceKernel: DurableObjectNamespace; DocumentSession: DurableObjectNamespace; @@ -109,7 +112,7 @@ type StringifyValues> = { [Binding in keyof EnvType]: EnvType[Binding] extends string ? EnvType[Binding] : string; }; declare namespace NodeJS { - interface ProcessEnv extends StringifyValues> {} + interface ProcessEnv extends StringifyValues> {} } // Begin runtime types diff --git a/wrangler.jsonc b/wrangler.jsonc index 41ed18d7..da356b80 100644 --- a/wrangler.jsonc +++ b/wrangler.jsonc @@ -70,6 +70,7 @@ "R2_ACCOUNT_ID", "R2_ACCESS_KEY_ID", "R2_SECRET_ACCESS_KEY", + "WORKSPACE_UPLOAD_TOKEN_SECRET", ], }, "workflows": [ @@ -216,6 +217,7 @@ "R2_ACCOUNT_ID", "R2_ACCESS_KEY_ID", "R2_SECRET_ACCESS_KEY", + "WORKSPACE_UPLOAD_TOKEN_SECRET", ], }, "ai": { @@ -346,6 +348,7 @@ "R2_ACCOUNT_ID", "R2_ACCESS_KEY_ID", "R2_SECRET_ACCESS_KEY", + "WORKSPACE_UPLOAD_TOKEN_SECRET", ], }, "ai": { From a1b83de5b61e58629a353b0f989651e6cb24ea63 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Tue, 14 Jul 2026 16:10:08 -0400 Subject: [PATCH 4/7] perf(workspaces): remove avoidable upload waits Skip processor error-body parsing for 5xx responses. Acquire the workspace kernel while document conversion runs. --- .../workspaces/upload/pdf-upload-validation.ts | 4 ++-- .../api/v1/workspaces.$workspaceId.file-upload.ts | 12 +++++++----- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/src/features/workspaces/upload/pdf-upload-validation.ts b/src/features/workspaces/upload/pdf-upload-validation.ts index 27d6270c..91219744 100644 --- a/src/features/workspaces/upload/pdf-upload-validation.ts +++ b/src/features/workspaces/upload/pdf-upload-validation.ts @@ -16,12 +16,12 @@ export async function assertReadablePdfUpload(input: { return; } - const payload = await readValidationFailure(response); - if (response.status >= 500) { throw new Error(`Workspace file processor failed with status ${response.status}.`); } + const payload = await readValidationFailure(response); + if (response.status === 413 || payload.code === "UPLOAD_TOO_LARGE") { throw new WorkspaceFileUploadError({ code: "UPLOAD_TOO_LARGE", diff --git a/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts b/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts index 58cfcf8b..9f33bda9 100644 --- a/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts +++ b/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts @@ -247,11 +247,13 @@ async function createWorkspaceDocumentFromUpload(input: { file: File; plan: Extract; }) { - const documentContent = await createDocumentContentFromWorkspaceUpload({ - file: input.file, - plan: input.plan, - }); - const kernel = await getWorkspaceKernel(input.claims.workspaceId); + const [documentContent, kernel] = await Promise.all([ + createDocumentContentFromWorkspaceUpload({ + file: input.file, + plan: input.plan, + }), + getWorkspaceKernel(input.claims.workspaceId), + ]); return kernel.createItem({ id: input.claims.itemId, From b42820c52461e806794f6fc9026e4c939f57f9a0 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Wed, 15 Jul 2026 10:30:07 -0400 Subject: [PATCH 5/7] fix(workspaces): make R2 publication retry-safe Preserve immutable R2 artifacts across ambiguous Durable Object failures. Recover retried item creation from the kernel event log and avoid duplicate extraction queues. Tests: pnpm check; pnpm test; pnpm knip --- .../container-file-conversion.test.ts | 12 +++- .../extraction/liteparse-projection.ts | 34 +++++------- .../request-workspace-file-extraction.ts | 9 +++ .../workspace-file-extraction-workflow.ts | 28 ++++------ .../workspace-file-preview-projection.ts | 10 +++- .../workspace-page-projection.test.ts | 35 +----------- .../extraction/workspace-page-projection.ts | 27 --------- .../kernel/workspace-kernel-events.test.ts | 55 +++++++++++++++++++ .../kernel/workspace-kernel-events.ts | 22 ++++++++ .../kernel/workspace-kernel-file-commands.ts | 10 ++++ .../kernel/workspace-kernel-item-commands.ts | 12 ++++ .../kernel/workspace-kernel-schema.ts | 3 + .../v1/workspaces.$workspaceId.file-upload.ts | 15 ++--- 13 files changed, 162 insertions(+), 110 deletions(-) create mode 100644 src/features/workspaces/kernel/workspace-kernel-events.test.ts diff --git a/src/features/workspaces/conversion/container-file-conversion.test.ts b/src/features/workspaces/conversion/container-file-conversion.test.ts index 2867d0e2..744671d8 100644 --- a/src/features/workspaces/conversion/container-file-conversion.test.ts +++ b/src/features/workspaces/conversion/container-file-conversion.test.ts @@ -34,6 +34,14 @@ describe("container file conversion", () => { expect(new Uint8Array(await response.arrayBuffer())).toEqual(output); expect(container.startAndWaitForPorts).toHaveBeenCalledOnce(); }); + + it("propagates a failed conversion response", async () => { + const container = createContainer(new Response("converter unavailable", { status: 503 })); + + await expect(convert(container)).rejects.toThrow( + "File conversion failed with status 503. converter unavailable", + ); + }); }); function convert(container: ReturnType) { @@ -50,11 +58,11 @@ function convert(container: ReturnType) { }); } -function createContainer(output: Uint8Array) { +function createContainer(output: Response | Uint8Array) { return { fetch: vi.fn(async (request: Request) => { await request.arrayBuffer(); - return new Response(output.slice().buffer); + return output instanceof Response ? output : new Response(output.slice().buffer); }), startAndWaitForPorts: vi.fn(async () => undefined), }; diff --git a/src/features/workspaces/extraction/liteparse-projection.ts b/src/features/workspaces/extraction/liteparse-projection.ts index 07bc046b..605e6e0c 100644 --- a/src/features/workspaces/extraction/liteparse-projection.ts +++ b/src/features/workspaces/extraction/liteparse-projection.ts @@ -6,10 +6,7 @@ import type { WorkspaceFileExtractionWorkflowParams, } from "#/features/workspaces/extraction/types"; import { getWorkspaceFileSourceObject } from "#/features/workspaces/extraction/workspace-file-source"; -import { - commitWorkspacePageProjection, - writeWorkspacePageProjection, -} from "#/features/workspaces/extraction/workspace-page-projection"; +import { writeWorkspacePageProjection } from "#/features/workspaces/extraction/workspace-page-projection"; import { getWorkspaceKernelFromEnv } from "#/features/workspaces/kernel/workspace-kernel-access"; export async function publishLiteParseProjection( @@ -54,23 +51,20 @@ export async function publishLiteParseProjection( workspaceId: params.workspaceId, }); - await commitWorkspacePageProjection({ - bucket: env.WORKSPACE_KERNEL_FILES, - kernel, - manifestObjectKey: projection.manifestObjectKey, - mutation: { - itemId: params.itemId, - format: "pages", - provider: "liteparse", - providerMode: "fast", - sourceHash: object.etag, - metadataJson: { - markdownLength: projection.manifest.markdownLength, - pageCount: projection.manifest.pageCount, - provisional: true, - }, - actorUserId: params.actorUserId, + await kernel.upsertFileProjection({ + itemId: params.itemId, + format: "pages", + status: "ready", + objectKey: projection.manifestObjectKey, + provider: "liteparse", + providerMode: "fast", + sourceHash: object.etag, + metadataJson: { + markdownLength: projection.manifest.markdownLength, + pageCount: projection.manifest.pageCount, + provisional: true, }, + actorUserId: params.actorUserId, }); return { diff --git a/src/features/workspaces/extraction/request-workspace-file-extraction.ts b/src/features/workspaces/extraction/request-workspace-file-extraction.ts index 2718a559..e3c7f7f3 100644 --- a/src/features/workspaces/extraction/request-workspace-file-extraction.ts +++ b/src/features/workspaces/extraction/request-workspace-file-extraction.ts @@ -25,6 +25,15 @@ export async function requestWorkspaceFileExtraction(input: { requestId: input.requestId, } satisfies WorkspaceFileExtractionWorkflowParams; const kernel = await getWorkspaceKernel(input.workspaceId); + const existing = await kernel.readFileProjection({ + itemId: input.itemId, + format: "pages", + }); + + if (existing) { + return; + } + await kernel.upsertFileProjection({ itemId: input.itemId, format: "pages", diff --git a/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts b/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts index 3ea2cacf..7f4c4035 100644 --- a/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts +++ b/src/features/workspaces/extraction/workspace-file-extraction-workflow.ts @@ -10,10 +10,7 @@ import type { WorkspaceFileExtractionWorkflowParams, } from "#/features/workspaces/extraction/types"; import { getWorkspaceFileSourceObject } from "#/features/workspaces/extraction/workspace-file-source"; -import { - commitWorkspacePageProjection, - writeWorkspacePageProjection, -} from "#/features/workspaces/extraction/workspace-page-projection"; +import { writeWorkspacePageProjection } from "#/features/workspaces/extraction/workspace-page-projection"; import { getWorkspaceKernelFromEnv } from "#/features/workspaces/kernel/workspace-kernel-access"; import { getWorkspaceUploadFamily } from "#/features/workspaces/model/workspace-file"; @@ -127,19 +124,16 @@ export class WorkspaceFileExtractionWorkflow extends WorkflowEntrypoint< markdownLength: extraction.markdownLength, }; - await commitWorkspacePageProjection({ - bucket: this.env.WORKSPACE_KERNEL_FILES, - kernel, - manifestObjectKey: extraction.manifestObjectKey, - mutation: { - itemId: params.itemId, - format: "pages", - provider: extraction.provider, - providerMode: extraction.providerMode, - sourceHash: extraction.sourceHash, - metadataJson, - actorUserId: params.actorUserId, - }, + await kernel.upsertFileProjection({ + itemId: params.itemId, + format: "pages", + status: "ready", + objectKey: extraction.manifestObjectKey, + provider: extraction.provider, + providerMode: extraction.providerMode, + sourceHash: extraction.sourceHash, + metadataJson, + actorUserId: params.actorUserId, }); return { diff --git a/src/features/workspaces/extraction/workspace-file-preview-projection.ts b/src/features/workspaces/extraction/workspace-file-preview-projection.ts index b4a1f48c..63529709 100644 --- a/src/features/workspaces/extraction/workspace-file-preview-projection.ts +++ b/src/features/workspaces/extraction/workspace-file-preview-projection.ts @@ -25,6 +25,15 @@ export async function publishWorkspaceFilePreview( } catch (error) { return step.do("record workspace file preview failure", async () => { const kernel = await getWorkspaceKernelFromEnv(env, params.workspaceId); + const existing = await kernel.readFileProjection({ + itemId: params.itemId, + format: "preview", + }); + + if (existing?.status === "ready") { + return { outcome: "success" as const }; + } + await kernel.upsertFileProjection({ itemId: params.itemId, format: "preview", @@ -95,7 +104,6 @@ async function generateWorkspaceFilePreview( return { outcome: "success" as const }; } catch (error) { failure = error; - await env.WORKSPACE_KERNEL_FILES.delete(objectKey).catch(() => undefined); throw error; } finally { recordOperationalOutcome({ diff --git a/src/features/workspaces/extraction/workspace-page-projection.test.ts b/src/features/workspaces/extraction/workspace-page-projection.test.ts index d4f920f4..ea8c0eaa 100644 --- a/src/features/workspaces/extraction/workspace-page-projection.test.ts +++ b/src/features/workspaces/extraction/workspace-page-projection.test.ts @@ -1,7 +1,6 @@ -import { describe, expect, it, vi } from "vitest"; +import { describe, expect, it } from "vitest"; import { - commitWorkspacePageProjection, getWorkspacePageObjectKey, readWorkspacePageProjection, writeWorkspacePageProjection, @@ -140,38 +139,6 @@ describe("workspace page projections", () => { ).rejects.toThrow("Extracted pages must be ordered"); expect(storage.values.size).toBe(0); }); - - it("removes a completed projection when its kernel publication fails", async () => { - const storage = createObjectStorage(); - const reference = await writeWorkspacePageProjection({ - bucket: storage.bucket, - itemId: "item-1", - pages: [{ pageNumber: 1, markdown: "First" }], - provider: "liteparse", - providerMode: "fast", - runId: "run-1", - sourceHash: "etag-1", - tier: "fast", - workspaceId: "workspace-1", - }); - const publicationError = new Error("Kernel publication failed"); - - await expect( - commitWorkspacePageProjection({ - bucket: storage.bucket, - kernel: { upsertFileProjection: vi.fn().mockRejectedValue(publicationError) }, - manifestObjectKey: reference.manifestObjectKey, - mutation: { - format: "pages", - itemId: "item-1", - provider: "liteparse", - providerMode: "fast", - sourceHash: "etag-1", - }, - }), - ).rejects.toBe(publicationError); - expect(storage.values.size).toBe(0); - }); }); function createObjectStorage() { diff --git a/src/features/workspaces/extraction/workspace-page-projection.ts b/src/features/workspaces/extraction/workspace-page-projection.ts index 28ff4b47..3b5ed245 100644 --- a/src/features/workspaces/extraction/workspace-page-projection.ts +++ b/src/features/workspaces/extraction/workspace-page-projection.ts @@ -1,8 +1,6 @@ import { jsonValueSchema, type JsonValue } from "#/features/workspaces/contracts"; import type { MarkdownProjectionPage } from "#/features/workspaces/extraction/page-markdown-projection"; import { getWorkspaceFileItemObjectPrefix } from "#/features/workspaces/files/workspace-file-object-keys"; -import type { WorkspaceKernelClient } from "#/features/workspaces/kernel/workspace-kernel-access"; -import type { UpsertWorkspaceKernelFileProjectionArgs } from "#/features/workspaces/kernel/workspace-kernel-types"; import { parseWorkspacePageRange, WorkspacePageSelectionError, @@ -37,31 +35,6 @@ export interface WorkspacePageProjectionReference { manifest: WorkspacePageProjectionManifest; } -type ReadyProjectionMutation = Omit< - Extract, - "objectKey" | "status" ->; - -export async function commitWorkspacePageProjection(input: { - bucket: R2Bucket; - kernel: Pick; - manifestObjectKey: string; - mutation: ReadyProjectionMutation; -}) { - try { - await input.kernel.upsertFileProjection({ - ...input.mutation, - objectKey: input.manifestObjectKey, - status: "ready", - }); - } catch (error) { - await deleteR2Prefix(input.bucket, getManifestPrefix(input.manifestObjectKey)).catch( - () => undefined, - ); - throw error; - } -} - export async function writeWorkspacePageProjection(input: { bucket: R2Bucket; itemId: string; diff --git a/src/features/workspaces/kernel/workspace-kernel-events.test.ts b/src/features/workspaces/kernel/workspace-kernel-events.test.ts new file mode 100644 index 00000000..d86a442f --- /dev/null +++ b/src/features/workspaces/kernel/workspace-kernel-events.test.ts @@ -0,0 +1,55 @@ +import { describe, expect, it, vi } from "vitest"; + +import { WorkspaceKernelEventBus } from "#/features/workspaces/kernel/workspace-kernel-events"; +import type { KernelEventRow } from "#/features/workspaces/kernel/workspace-kernel-rows"; +import type { WorkspaceKernelSql } from "#/features/workspaces/kernel/workspace-kernel-schema"; + +describe("workspace kernel event idempotency", () => { + it("recovers a committed item creation by client mutation id", () => { + const events = createEventBus(createItemEventRow("item-1")); + + expect( + events.getCreatedItemEvent({ + clientMutationId: "mutation-1", + itemId: "item-1", + }), + ).toMatchObject({ + clientMutationId: "mutation-1", + payload: { item: { id: "item-1" } }, + type: "workspace.item.created", + }); + }); + + it("rejects reuse of a client mutation id for another item", () => { + const events = createEventBus(createItemEventRow("item-1")); + + expect(() => + events.getCreatedItemEvent({ + clientMutationId: "mutation-1", + itemId: "item-2", + }), + ).toThrow("client mutation id was already used"); + }); +}); + +function createEventBus(row: KernelEventRow) { + const sql = vi.fn(() => [row]) as unknown as WorkspaceKernelSql; + return new WorkspaceKernelEventBus({ + broadcast: vi.fn(), + getNextRevision: () => 2, + sql, + workspaceId: () => "workspace-1", + }); +} + +function createItemEventRow(itemId: string): KernelEventRow { + return { + actor_user_id: "user-1", + client_mutation_id: "mutation-1", + created_at: Date.parse("2026-07-15T00:00:00Z"), + id: "event-1", + payload_json: JSON.stringify({ item: { id: itemId } }), + revision: 1, + type: "workspace.item.created", + }; +} diff --git a/src/features/workspaces/kernel/workspace-kernel-events.ts b/src/features/workspaces/kernel/workspace-kernel-events.ts index a4d68bb3..3aa3c0de 100644 --- a/src/features/workspaces/kernel/workspace-kernel-events.ts +++ b/src/features/workspaces/kernel/workspace-kernel-events.ts @@ -40,6 +40,28 @@ export class WorkspaceKernelEventBus { return rows.map((row) => mapKernelEventRow(row, this.workspaceId())); } + getCreatedItemEvent(input: { clientMutationId: string; itemId: string }) { + const [row] = this.sql` + SELECT * + FROM kernel_events + WHERE type = 'workspace.item.created' + AND client_mutation_id = ${input.clientMutationId} + ORDER BY revision ASC + LIMIT 1 + `; + + if (!row) { + return null; + } + + const event = mapKernelEventRow(row, this.workspaceId()); + if (event.type !== "workspace.item.created" || event.payload.item.id !== input.itemId) { + throw new Error("Workspace client mutation id was already used."); + } + + return event; + } + commit(input: Omit) { const createdAt = Date.now(); const event = { diff --git a/src/features/workspaces/kernel/workspace-kernel-file-commands.ts b/src/features/workspaces/kernel/workspace-kernel-file-commands.ts index 04a36ea1..dc913cf0 100644 --- a/src/features/workspaces/kernel/workspace-kernel-file-commands.ts +++ b/src/features/workspaces/kernel/workspace-kernel-file-commands.ts @@ -56,6 +56,16 @@ export class WorkspaceKernelFileCommands { input: CreateWorkspaceKernelFileFromUploadArgs, ): Promise> { const parentId = input.parentId ?? null; + const priorEvent = input.clientMutationId + ? this.events.getCreatedItemEvent({ + clientMutationId: input.clientMutationId, + itemId: input.id, + }) + : null; + + if (priorEvent) { + return { event: priorEvent, result: this.store.requireItem(input.id) }; + } this.store.assertParentIsValid(parentId); diff --git a/src/features/workspaces/kernel/workspace-kernel-item-commands.ts b/src/features/workspaces/kernel/workspace-kernel-item-commands.ts index 0a52b721..5e9980e4 100644 --- a/src/features/workspaces/kernel/workspace-kernel-item-commands.ts +++ b/src/features/workspaces/kernel/workspace-kernel-item-commands.ts @@ -72,6 +72,18 @@ export class WorkspaceKernelItemCommands { const id = input.id ?? crypto.randomUUID(); const parentId = input.parentId ?? null; + const priorEvent = + input.id && input.clientMutationId + ? this.events.getCreatedItemEvent({ + clientMutationId: input.clientMutationId, + itemId: input.id, + }) + : null; + + if (priorEvent) { + return { event: priorEvent, result: this.store.requireItem(id) }; + } + const color = resolveWorkspaceItemColorForCreate({ type, color: input.color, diff --git a/src/features/workspaces/kernel/workspace-kernel-schema.ts b/src/features/workspaces/kernel/workspace-kernel-schema.ts index 1fb0c590..88eb7686 100644 --- a/src/features/workspaces/kernel/workspace-kernel-schema.ts +++ b/src/features/workspaces/kernel/workspace-kernel-schema.ts @@ -85,6 +85,9 @@ export function initializeWorkspaceKernelStorage(sql: WorkspaceKernelSql) { `; sql`CREATE INDEX IF NOT EXISTS kernel_events_revision_idx ON kernel_events (revision)`; + sql`CREATE INDEX IF NOT EXISTS kernel_events_client_mutation_idx + ON kernel_events (client_mutation_id, type, revision) + WHERE client_mutation_id IS NOT NULL`; } function createSiblingNameIndexes(sql: WorkspaceKernelSql) { diff --git a/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts b/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts index 9f33bda9..2deb238d 100644 --- a/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts +++ b/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts @@ -114,7 +114,6 @@ async function finalizeWorkspaceFileUpload( observation: WorkspaceFileIntakeObservation, ) { let stagingObjectKey: string | null = null; - let finalObjectKey: string | null = null; let completionClaimKey: string | null = null; let uploadCompleted = false; @@ -136,13 +135,13 @@ async function finalizeWorkspaceFileUpload( if (!validation.ok) { throw invalidUpload("Upload completion metadata is invalid."); } + observation.inputBytes = claims.fileSize; + observation.plan = validation.plan.kind; completionClaimKey = await claimWorkspaceDirectUploadCompletion(env, claims); if (!completionClaimKey) { throw invalidUpload("Upload is already being completed."); } - observation.inputBytes = claims.fileSize; - observation.plan = validation.plan.kind; stagingObjectKey = getWorkspaceFileUploadObjectKey(claims); const stagingObject = await env.WORKSPACE_KERNEL_FILES.get(stagingObjectKey); @@ -164,7 +163,7 @@ async function finalizeWorkspaceFileUpload( return apiJson(command, requestId); } - finalObjectKey = getWorkspaceFileSourceObjectKey(claims); + const finalObjectKey = getWorkspaceFileSourceObjectKey(claims); const upload = await storeWorkspaceFileUpload({ body: stagingObject.body, contentType: claims.contentType, @@ -192,7 +191,6 @@ async function finalizeWorkspaceFileUpload( workspaceId, }); - finalObjectKey = null; observation.itemId = command.result.id; uploadCompleted = true; await queueWorkspaceFileExtraction(upload, { @@ -207,11 +205,10 @@ async function finalizeWorkspaceFileUpload( observation.error = error; return workspaceUploadErrorResponse(requestId, error); } finally { - if (stagingObjectKey) { - await env.WORKSPACE_KERNEL_FILES.delete(stagingObjectKey).catch(() => undefined); - } await Promise.allSettled([ - finalObjectKey ? env.WORKSPACE_KERNEL_FILES.delete(finalObjectKey) : Promise.resolve(), + stagingObjectKey && uploadCompleted + ? env.WORKSPACE_KERNEL_FILES.delete(stagingObjectKey) + : Promise.resolve(), completionClaimKey && !uploadCompleted ? env.WORKSPACE_KERNEL_FILES.delete(completionClaimKey) : Promise.resolve(), From 1b5d4f2723c6f08668d1c14666d988ce6e2ff96b Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Wed, 15 Jul 2026 10:40:45 -0400 Subject: [PATCH 6/7] refactor(workspaces): simplify upload completion coordination --- .../request-workspace-file-extraction.ts | 23 +---- .../files/workspace-file-object-keys.ts | 7 -- .../kernel/workspace-kernel-types.ts | 9 +- .../workspaces/operations/read-items.ts | 10 +- .../workspace-file-direct-upload.test.ts | 37 ------- .../upload/workspace-file-direct-upload.ts | 17 +--- .../v1/workspaces.$workspaceId.file-upload.ts | 99 ++++++++----------- 7 files changed, 46 insertions(+), 156 deletions(-) diff --git a/src/features/workspaces/extraction/request-workspace-file-extraction.ts b/src/features/workspaces/extraction/request-workspace-file-extraction.ts index e3c7f7f3..18f5b5ce 100644 --- a/src/features/workspaces/extraction/request-workspace-file-extraction.ts +++ b/src/features/workspaces/extraction/request-workspace-file-extraction.ts @@ -24,33 +24,12 @@ export async function requestWorkspaceFileExtraction(input: { assetKind: input.assetKind, requestId: input.requestId, } satisfies WorkspaceFileExtractionWorkflowParams; - const kernel = await getWorkspaceKernel(input.workspaceId); - const existing = await kernel.readFileProjection({ - itemId: input.itemId, - format: "pages", - }); - - if (existing) { - return; - } - - await kernel.upsertFileProjection({ - itemId: input.itemId, - format: "pages", - status: "queued", - actorUserId: input.actorUserId, - }); - - const [instance] = await env.WORKSPACE_FILE_EXTRACTION_WORKFLOW.createBatch([ + await env.WORKSPACE_FILE_EXTRACTION_WORKFLOW.createBatch([ { id: workflowId, params, }, ]); - - if (!instance) { - throw new Error("Workspace file extraction workflow was not created."); - } } catch (error) { recordOperationalFailure({ distinctId: input.actorUserId ?? undefined, diff --git a/src/features/workspaces/files/workspace-file-object-keys.ts b/src/features/workspaces/files/workspace-file-object-keys.ts index 1faf191d..0863eb91 100644 --- a/src/features/workspaces/files/workspace-file-object-keys.ts +++ b/src/features/workspaces/files/workspace-file-object-keys.ts @@ -20,10 +20,3 @@ export function getWorkspaceFilePreviewObjectKey(input: { workspaceId: string; i export function getWorkspaceFileUploadObjectKey(input: { itemId: string; workspaceId: string }) { return `${workspaceFileUploadPrefix}/${input.workspaceId}/${input.itemId}/source`; } - -export function getWorkspaceFileUploadCompletionKey(input: { - itemId: string; - workspaceId: string; -}) { - return `${workspaceFileUploadPrefix}/${input.workspaceId}/${input.itemId}/completion`; -} diff --git a/src/features/workspaces/kernel/workspace-kernel-types.ts b/src/features/workspaces/kernel/workspace-kernel-types.ts index 99f72eee..088ed82f 100644 --- a/src/features/workspaces/kernel/workspace-kernel-types.ts +++ b/src/features/workspaces/kernel/workspace-kernel-types.ts @@ -109,12 +109,7 @@ export interface WorkspaceKernelFileSource { export type WorkspaceKernelFileProjectionFormat = "pages" | "preview"; -export type WorkspaceKernelFileProjectionStatus = - | "not_started" - | "queued" - | "processing" - | "ready" - | "failed"; +export type WorkspaceKernelFileProjectionStatus = "processing" | "ready" | "failed"; interface WorkspaceKernelFileProjectionMutationBase { itemId: string; @@ -125,7 +120,7 @@ interface WorkspaceKernelFileProjectionMutationBase { export type UpsertWorkspaceKernelFileProjectionArgs = | (WorkspaceKernelFileProjectionMutationBase & { - status: "not_started" | "queued" | "processing"; + status: "processing"; errorMessage?: never; metadataJson?: never; objectKey?: never; diff --git a/src/features/workspaces/operations/read-items.ts b/src/features/workspaces/operations/read-items.ts index 1a615715..729fb456 100644 --- a/src/features/workspaces/operations/read-items.ts +++ b/src/features/workspaces/operations/read-items.ts @@ -212,15 +212,7 @@ async function readWorkspaceFileItem(input: { format: "pages", }); - if ( - pagesProjection?.status === "queued" || - pagesProjection?.status === "processing" || - pagesProjection?.status === "not_started" - ) { - return createWorkspaceFileStatusItem(input.path, "pending"); - } - - if (!pagesProjection) { + if (!pagesProjection || pagesProjection.status === "processing") { return createWorkspaceFileStatusItem(input.path, "pending"); } diff --git a/src/features/workspaces/upload/workspace-file-direct-upload.test.ts b/src/features/workspaces/upload/workspace-file-direct-upload.test.ts index 5efb8e7e..49a57260 100644 --- a/src/features/workspaces/upload/workspace-file-direct-upload.test.ts +++ b/src/features/workspaces/upload/workspace-file-direct-upload.test.ts @@ -1,7 +1,6 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { - claimWorkspaceDirectUploadCompletion, createWorkspaceDirectUploadSession, verifyWorkspaceDirectUploadToken, } from "#/features/workspaces/upload/workspace-file-direct-upload"; @@ -76,42 +75,6 @@ describe("workspace direct upload sessions", () => { "expired", ); }); - - it("allows only one upload completion owner", async () => { - const values = new Set(); - const env = { - ...createEnv(), - WORKSPACE_KERNEL_FILES: { - async put(key: string) { - if (values.has(key)) { - return null; - } - values.add(key); - return { key }; - }, - } as unknown as R2Bucket, - }; - const claims = { - clientMutationId: "mutation-1", - contentType: "application/pdf", - expiresAt: Math.floor(Date.now() / 1_000) + 60, - fileName: "report.pdf", - fileSize: 42, - itemId: crypto.randomUUID(), - parentId: null, - userId: "user-1", - version: 1 as const, - workspaceId: "workspace-1", - }; - - const [first, second] = await Promise.all([ - claimWorkspaceDirectUploadCompletion(env, claims), - claimWorkspaceDirectUploadCompletion(env, claims), - ]); - - expect(first).toContain("/completion"); - expect(second).toBeNull(); - }); }); function createEnv() { diff --git a/src/features/workspaces/upload/workspace-file-direct-upload.ts b/src/features/workspaces/upload/workspace-file-direct-upload.ts index 60a8bc25..0aec5e03 100644 --- a/src/features/workspaces/upload/workspace-file-direct-upload.ts +++ b/src/features/workspaces/upload/workspace-file-direct-upload.ts @@ -1,10 +1,7 @@ import { AwsClient } from "aws4fetch"; import { z } from "zod"; -import { - getWorkspaceFileUploadCompletionKey, - getWorkspaceFileUploadObjectKey, -} from "#/features/workspaces/files/workspace-file-object-keys"; +import { getWorkspaceFileUploadObjectKey } from "#/features/workspaces/files/workspace-file-object-keys"; const uploadUrlLifetimeSeconds = 30 * 60; const uploadTokenVersion = 1; @@ -77,18 +74,6 @@ export async function verifyWorkspaceDirectUploadToken( return claims; } -export async function claimWorkspaceDirectUploadCompletion( - env: Cloudflare.Env, - claims: WorkspaceDirectUploadClaims, -): Promise { - const objectKey = getWorkspaceFileUploadCompletionKey(claims); - const claim = await env.WORKSPACE_KERNEL_FILES.put(objectKey, "", { - onlyIf: { etagDoesNotMatch: "*" }, - }); - - return claim ? objectKey : null; -} - async function createPresignedUploadUrl( env: Cloudflare.Env, input: { contentType: string; objectKey: string }, diff --git a/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts b/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts index 2deb238d..482630b5 100644 --- a/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts +++ b/src/routes/api/v1/workspaces.$workspaceId.file-upload.ts @@ -22,7 +22,6 @@ import { WorkspaceForbiddenError, } from "#/features/workspaces/server/permissions"; import { - claimWorkspaceDirectUploadCompletion, createWorkspaceDirectUploadSession, verifyWorkspaceDirectUploadToken, type WorkspaceDirectUploadClaims, @@ -113,10 +112,6 @@ async function finalizeWorkspaceFileUpload( requestId: string, observation: WorkspaceFileIntakeObservation, ) { - let stagingObjectKey: string | null = null; - let completionClaimKey: string | null = null; - let uploadCompleted = false; - try { const userId = await authorizeWorkspaceUpload(request, workspaceId); observation.userId = userId; @@ -137,20 +132,18 @@ async function finalizeWorkspaceFileUpload( } observation.inputBytes = claims.fileSize; observation.plan = validation.plan.kind; - completionClaimKey = await claimWorkspaceDirectUploadCompletion(env, claims); - if (!completionClaimKey) { - throw invalidUpload("Upload is already being completed."); - } - stagingObjectKey = getWorkspaceFileUploadObjectKey(claims); + const stagingObjectKey = getWorkspaceFileUploadObjectKey(claims); const stagingObject = await env.WORKSPACE_KERNEL_FILES.get(stagingObjectKey); if (!stagingObject || stagingObject.size !== claims.fileSize) { throw invalidUpload("Uploaded file size does not match the selected file."); } + let command: Awaited>; + if (validation.plan.kind === "document") { - const command = await createWorkspaceDocumentFromUpload({ + command = await createWorkspaceDocumentFromUpload({ claims, file: new File([await stagingObject.arrayBuffer()], claims.fileName, { type: claims.contentType, @@ -159,60 +152,50 @@ async function finalizeWorkspaceFileUpload( }); observation.itemId = command.result.id; observation.outputBytes = claims.fileSize; - uploadCompleted = true; - return apiJson(command, requestId); - } - - const finalObjectKey = getWorkspaceFileSourceObjectKey(claims); - const upload = await storeWorkspaceFileUpload({ - body: stagingObject.body, - contentType: claims.contentType, - descriptor: validation.plan.descriptor, - env, - fileName: claims.fileName, - fileSize: claims.fileSize, - objectKey: finalObjectKey, - }); - observation.assetKind = upload.descriptor.assetKind; - observation.conversion = upload.source?.conversion; - observation.outputBytes = upload.fileSize; + } else { + const finalObjectKey = getWorkspaceFileSourceObjectKey(claims); + const upload = await storeWorkspaceFileUpload({ + body: stagingObject.body, + contentType: claims.contentType, + descriptor: validation.plan.descriptor, + env, + fileName: claims.fileName, + fileSize: claims.fileSize, + objectKey: finalObjectKey, + }); + observation.assetKind = upload.descriptor.assetKind; + observation.conversion = upload.source?.conversion; + observation.outputBytes = upload.fileSize; + + command = await createWorkspaceFileFromUpload({ + assetKind: upload.descriptor.assetKind, + clientMutationId: claims.clientMutationId, + contentType: upload.contentType, + fileName: upload.fileName, + fileSize: upload.fileSize, + id: claims.itemId, + objectKey: upload.objectKey, + parentId: claims.parentId, + source: upload.source, + userId, + workspaceId, + }); - const command = await createWorkspaceFileFromUpload({ - assetKind: upload.descriptor.assetKind, - clientMutationId: claims.clientMutationId, - contentType: upload.contentType, - fileName: upload.fileName, - fileSize: upload.fileSize, - id: claims.itemId, - objectKey: upload.objectKey, - parentId: claims.parentId, - source: upload.source, - userId, - workspaceId, - }); + observation.itemId = command.result.id; + await queueWorkspaceFileExtraction(upload, { + itemId: command.result.id, + requestId, + userId, + workspaceId, + }); + } - observation.itemId = command.result.id; - uploadCompleted = true; - await queueWorkspaceFileExtraction(upload, { - itemId: command.result.id, - requestId, - userId, - workspaceId, - }); + await env.WORKSPACE_KERNEL_FILES.delete(stagingObjectKey).catch(() => undefined); return apiJson(command, requestId); } catch (error) { observation.error = error; return workspaceUploadErrorResponse(requestId, error); - } finally { - await Promise.allSettled([ - stagingObjectKey && uploadCompleted - ? env.WORKSPACE_KERNEL_FILES.delete(stagingObjectKey) - : Promise.resolve(), - completionClaimKey && !uploadCompleted - ? env.WORKSPACE_KERNEL_FILES.delete(completionClaimKey) - : Promise.resolve(), - ]); } } From ebe8bc756e32fa3536c7b952e3f8ec4ac8245ec6 Mon Sep 17 00:00:00 2001 From: Urjit Chakraborty <135136842+urjitc@users.noreply.github.com> Date: Wed, 15 Jul 2026 10:51:24 -0400 Subject: [PATCH 7/7] fix(workspaces): close upload publication races --- .../workspace-file-preview-projection.ts | 6 +++- .../kernel/workspace-kernel-file-commands.ts | 28 +++++++++++------ .../kernel/workspace-kernel-item-commands.ts | 30 ++++++++++++------- .../kernel/workspace-kernel-types.ts | 7 ++++- .../workspaces/operations/read-items.ts | 7 ++++- 5 files changed, 56 insertions(+), 22 deletions(-) diff --git a/src/features/workspaces/extraction/workspace-file-preview-projection.ts b/src/features/workspaces/extraction/workspace-file-preview-projection.ts index 63529709..8e537246 100644 --- a/src/features/workspaces/extraction/workspace-file-preview-projection.ts +++ b/src/features/workspaces/extraction/workspace-file-preview-projection.ts @@ -30,7 +30,11 @@ export async function publishWorkspaceFilePreview( format: "preview", }); - if (existing?.status === "ready") { + if ( + existing?.status === "ready" && + existing.objectKey && + (await env.WORKSPACE_KERNEL_FILES.head(existing.objectKey)) + ) { return { outcome: "success" as const }; } diff --git a/src/features/workspaces/kernel/workspace-kernel-file-commands.ts b/src/features/workspaces/kernel/workspace-kernel-file-commands.ts index dc913cf0..cd3c9f75 100644 --- a/src/features/workspaces/kernel/workspace-kernel-file-commands.ts +++ b/src/features/workspaces/kernel/workspace-kernel-file-commands.ts @@ -56,15 +56,20 @@ export class WorkspaceKernelFileCommands { input: CreateWorkspaceKernelFileFromUploadArgs, ): Promise> { const parentId = input.parentId ?? null; - const priorEvent = input.clientMutationId - ? this.events.getCreatedItemEvent({ - clientMutationId: input.clientMutationId, - itemId: input.id, - }) - : null; - - if (priorEvent) { - return { event: priorEvent, result: this.store.requireItem(input.id) }; + const getPriorResult = () => { + const event = input.clientMutationId + ? this.events.getCreatedItemEvent({ + clientMutationId: input.clientMutationId, + itemId: input.id, + }) + : null; + + return event ? { event, result: this.store.requireItem(input.id) } : null; + }; + const priorResult = getPriorResult(); + + if (priorResult) { + return priorResult; } this.store.assertParentIsValid(parentId); @@ -83,6 +88,11 @@ export class WorkspaceKernelFileCommands { throw new Error("Uploaded file size did not match the upload request."); } + const concurrentResult = getPriorResult(); + if (concurrentResult) { + return concurrentResult; + } + const descriptor = getWorkspaceUploadFamily(input.assetKind); const contentType = resolveWorkspaceFileContentType({ contentType: input.contentType, diff --git a/src/features/workspaces/kernel/workspace-kernel-item-commands.ts b/src/features/workspaces/kernel/workspace-kernel-item-commands.ts index 5e9980e4..96e39e18 100644 --- a/src/features/workspaces/kernel/workspace-kernel-item-commands.ts +++ b/src/features/workspaces/kernel/workspace-kernel-item-commands.ts @@ -72,16 +72,21 @@ export class WorkspaceKernelItemCommands { const id = input.id ?? crypto.randomUUID(); const parentId = input.parentId ?? null; - const priorEvent = - input.id && input.clientMutationId - ? this.events.getCreatedItemEvent({ - clientMutationId: input.clientMutationId, - itemId: input.id, - }) - : null; - - if (priorEvent) { - return { event: priorEvent, result: this.store.requireItem(id) }; + const getPriorResult = () => { + const event = + input.id && input.clientMutationId + ? this.events.getCreatedItemEvent({ + clientMutationId: input.clientMutationId, + itemId: input.id, + }) + : null; + + return event ? { event, result: this.store.requireItem(id) } : null; + }; + const priorResult = getPriorResult(); + + if (priorResult) { + return priorResult; } const color = resolveWorkspaceItemColorForCreate({ @@ -117,6 +122,11 @@ export class WorkspaceKernelItemCommands { initialContent, }); + const concurrentResult = getPriorResult(); + if (concurrentResult) { + return concurrentResult; + } + this.sql` INSERT INTO kernel_items ( id, diff --git a/src/features/workspaces/kernel/workspace-kernel-types.ts b/src/features/workspaces/kernel/workspace-kernel-types.ts index 088ed82f..c96a2812 100644 --- a/src/features/workspaces/kernel/workspace-kernel-types.ts +++ b/src/features/workspaces/kernel/workspace-kernel-types.ts @@ -109,7 +109,12 @@ export interface WorkspaceKernelFileSource { export type WorkspaceKernelFileProjectionFormat = "pages" | "preview"; -export type WorkspaceKernelFileProjectionStatus = "processing" | "ready" | "failed"; +export type WorkspaceKernelFileProjectionStatus = + | "not_started" + | "queued" + | "processing" + | "ready" + | "failed"; interface WorkspaceKernelFileProjectionMutationBase { itemId: string; diff --git a/src/features/workspaces/operations/read-items.ts b/src/features/workspaces/operations/read-items.ts index 729fb456..c473fe17 100644 --- a/src/features/workspaces/operations/read-items.ts +++ b/src/features/workspaces/operations/read-items.ts @@ -212,7 +212,12 @@ async function readWorkspaceFileItem(input: { format: "pages", }); - if (!pagesProjection || pagesProjection.status === "processing") { + if ( + !pagesProjection || + pagesProjection.status === "not_started" || + pagesProjection.status === "queued" || + pagesProjection.status === "processing" + ) { return createWorkspaceFileStatusItem(input.path, "pending"); }