From f8a2d530d6361a2c962d33efd0b8ded40dee9614 Mon Sep 17 00:00:00 2001 From: Jyong Date: Tue, 18 Aug 2026 03:03:18 -0400 Subject: [PATCH] feat(knowledge-fs): harden PDF image materialization --- api/knowledge-fs-contract.lock.json | 2 +- docker/README.md | 9 + docker/docker-compose-template.yaml | 8 + docker/docker-compose.yaml | 8 + .../core-services/knowledge-fs.env.example | 17 + .../2026-08-18-api-image-poppler-runtime.md | 98 ++ knowledge-fs/README.md | 1 + knowledge-fs/apps/api/Dockerfile | 17 +- .../document-compilation-runtime-options.ts | 12 +- .../apps/api/src/multimodal-options.test.ts | 100 +- .../apps/api/src/multimodal-options.ts | 78 +- .../apps/api/src/parser-options.test.ts | 10 + knowledge-fs/apps/api/src/parser-options.ts | 30 + knowledge-fs/docs/production-deployment.md | 35 + knowledge-fs/infra/kubernetes/README.md | 6 + .../kubernetes/dify-integration-baseline.yaml | 10 + knowledge-fs/infra/local/.env.example | 13 +- knowledge-fs/infra/local/README.md | 6 + knowledge-fs/infra/local/compose.yaml | 14 +- .../adapters/src/dify-object-storage.test.ts | 73 +- .../adapters/src/dify-object-storage.ts | 245 ++- knowledge-fs/packages/adapters/src/node.ts | 5 + .../src/document-compilation-runtime.test.ts | 26 + .../api/src/document-compilation-runtime.ts | 22 + .../src/document-compilation-worker.test.ts | 597 ++++++- .../api/src/document-compilation-worker.ts | 414 +++-- .../document-image-variant-generator.test.ts | 42 + .../src/document-image-variant-generator.ts | 22 +- .../src/document-multimodal-artifact.test.ts | 88 ++ .../api/src/document-multimodal-artifact.ts | 151 ++ ...ocument-multimodal-asset-extractor.test.ts | 52 + .../document-multimodal-asset-extractor.ts | 16 +- .../document-pdf-rasterizer-coverage.test.ts | 105 +- .../api/src/document-pdf-rasterizer.test.ts | 1078 +++++++++++++ .../api/src/document-pdf-rasterizer.ts | 1372 ++++++++++++++--- knowledge-fs/packages/api/src/gateway.test.ts | 25 +- .../packages/api/src/index-reindexer.test.ts | 2 +- .../packages/api/src/index-reindexer.ts | 19 +- .../packages/api/src/knowledge-fs-errors.ts | 1 + .../api/src/parse-artifact-repository.test.ts | 206 ++- .../api/src/parse-artifact-repository.ts | 283 ++-- .../api/src/storage-path-utils.test.ts | 32 + .../packages/api/src/storage-path-utils.ts | 20 +- knowledge-fs/packages/parsers/src/index.ts | 536 ++++++- .../packages/parsers/src/parser.test.ts | 461 +++++- .../scripts/api-image-bundle-smoke.mjs | 25 + .../scripts/api-image-bundle-smoke.test.mjs | 16 + knowledge-fs/scripts/compose-apps.test.mjs | 138 +- 48 files changed, 5886 insertions(+), 660 deletions(-) create mode 100644 knowledge-fs/.harness/changes/2026-08-18-api-image-poppler-runtime.md create mode 100644 knowledge-fs/packages/api/src/document-multimodal-artifact.test.ts create mode 100644 knowledge-fs/packages/api/src/document-multimodal-artifact.ts diff --git a/api/knowledge-fs-contract.lock.json b/api/knowledge-fs-contract.lock.json index 75e22bb79c5..74eacab96b2 100644 --- a/api/knowledge-fs-contract.lock.json +++ b/api/knowledge-fs-contract.lock.json @@ -1,6 +1,6 @@ { "schemaVersion": 5, - "subtreeTree": "222470dcafc7a985cfb78896c31180f21cf88d75", + "subtreeTree": "5b0580d14b446eff3d6e5e6554137aec0dab065d", "openapiSha256": "3a712231fa850c4f5151bc283205da9062086a0b693f0d1ab01c2c526323f018", "capabilityV2AuthManifestSha256": "fc0a47e23cce12544882f0298522b4933002e892b84ce1815df7e81d36a7a0c7", "capabilityV2AuthTestVectorSha256": "ae0de37b1ff05c40f905cf17a7b410d8971acacf64db07d5ee3d6fecfa559ce3", diff --git a/docker/README.md b/docker/README.md index 126ffe89e3d..9e38bc41a2b 100644 --- a/docker/README.md +++ b/docker/README.md @@ -66,6 +66,15 @@ bucket, endpoint, or provider credentials in `knowledge-fs.env`. Feature-specifi and capacity tunables should be added only when deliberately overriding their safe runtime defaults. +The KnowledgeFS API image includes Poppler and enables its `pdftoppm` PDF image rasterizer with +bounded defaults (144 DPI, 48 DPI thumbnails, a 30-second timeout, at most 500 assets per document, +and two concurrent Poppler page batches per replica). Canonical settings live in the dedicated +`knowledge-fs.env`; whitelisted `DIFY_ROOT_*_OVERRIDE` proxies let explicitly set values in +`docker/.env` take precedence without injecting the complete root environment. When a root value +is unset, the service env or image default remains authoritative. Set +`KNOWLEDGE_PDF_RASTERIZER=off` in either operator env as an emergency or low-resource kill switch. Existing +documents whose parse artifacts lack image asset references must be ingested again after rollout. + Use a dedicated KnowledgeFS database. Do not point `DATABASE_URL` at Dify's application database, reuse Dataset/Document tables, or run a data migration as part of this service. KnowledgeFS migrations remain a separate controlled operator step. The selected Dify storage backend must diff --git a/docker/docker-compose-template.yaml b/docker/docker-compose-template.yaml index 1f57408500c..deea2cca5f9 100644 --- a/docker/docker-compose-template.yaml +++ b/docker/docker-compose-template.yaml @@ -663,6 +663,14 @@ services: - path: ./envs/core-services/knowledge-fs.env required: false environment: + # Canonical PDF rasterizer values stay in the service env/image. These whitelisted proxies + # let an explicitly set docker/.env value override them without injecting the root env. + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_OVERRIDE: ${KNOWLEDGE_PDF_RASTERIZER-} + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_DPI_OVERRIDE: ${KNOWLEDGE_PDF_RASTERIZER_DPI-} + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI_OVERRIDE: ${KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI-} + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS_OVERRIDE: ${KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS-} + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS_OVERRIDE: ${KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS-} + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY_OVERRIDE: ${KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY-} DIFY_INNER_API_URL: ${PLUGIN_DIFY_INNER_API_URL:-http://api:5001} DIFY_INNER_API_KEY: ${PLUGIN_DIFY_INNER_API_KEY:-QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1} KNOWLEDGE_INTEGRATED_MODE_ENABLED: ${KNOWLEDGE_INTEGRATED_MODE_ENABLED:-true} diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index ef4ea3f2c3a..8ddc084a752 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -669,6 +669,14 @@ services: - path: ./envs/core-services/knowledge-fs.env required: false environment: + # Canonical PDF rasterizer values stay in the service env/image. These whitelisted proxies + # let an explicitly set docker/.env value override them without injecting the root env. + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_OVERRIDE: ${KNOWLEDGE_PDF_RASTERIZER-} + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_DPI_OVERRIDE: ${KNOWLEDGE_PDF_RASTERIZER_DPI-} + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI_OVERRIDE: ${KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI-} + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS_OVERRIDE: ${KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS-} + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS_OVERRIDE: ${KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS-} + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY_OVERRIDE: ${KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY-} DIFY_INNER_API_URL: ${PLUGIN_DIFY_INNER_API_URL:-http://api:5001} DIFY_INNER_API_KEY: ${PLUGIN_DIFY_INNER_API_KEY:-QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1} KNOWLEDGE_INTEGRATED_MODE_ENABLED: ${KNOWLEDGE_INTEGRATED_MODE_ENABLED:-true} diff --git a/docker/envs/core-services/knowledge-fs.env.example b/docker/envs/core-services/knowledge-fs.env.example index 9dc9e8ab175..f1a13a4a013 100644 --- a/docker/envs/core-services/knowledge-fs.env.example +++ b/docker/envs/core-services/knowledge-fs.env.example @@ -14,6 +14,16 @@ DATABASE_URL= KNOWLEDGE_DOCUMENT_COMPILATION_RUNTIME=on +# PDF image elements are rasterized by the Poppler executable bundled in the API image. The +# limits are per document/asset operation. Set the mode to off as an emergency or low-resource +# kill switch; do not set a host-specific command path when using the published image. +KNOWLEDGE_PDF_RASTERIZER=poppler +KNOWLEDGE_PDF_RASTERIZER_DPI=144 +KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI=48 +KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS=30000 +KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS=500 +KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY=2 + # Ingestion-time LLM concurrency. The per-document outline limit improves one document's latency; # the process-wide fair limit prevents simultaneous imports from multiplying provider pressure. # The global value is per API replica and also covers embedding transport requests. @@ -50,3 +60,10 @@ KNOWLEDGE_QUERY_IMAGE_EXPANSION_TIMEOUT_MS=8000 # Leave both values blank only when PDF/Office parsing is intentionally unavailable. UNSTRUCTURED_API_URL= UNSTRUCTURED_API_KEY= +# Bound parser fan-out independently from the document-compilation batch size. +UNSTRUCTURED_MAX_CONCURRENCY=2 +UNSTRUCTURED_REQUEST_TIMEOUT_MS=120000 +UNSTRUCTURED_MAX_RESPONSE_BYTES=33554432 + +# Bound authenticated Dify object-storage calls so cleanup cannot hang a compilation lease forever. +DIFY_OBJECT_STORAGE_REQUEST_TIMEOUT_MS=60000 diff --git a/knowledge-fs/.harness/changes/2026-08-18-api-image-poppler-runtime.md b/knowledge-fs/.harness/changes/2026-08-18-api-image-poppler-runtime.md new file mode 100644 index 00000000000..8d5f91a9b03 --- /dev/null +++ b/knowledge-fs/.harness/changes/2026-08-18-api-image-poppler-runtime.md @@ -0,0 +1,98 @@ +# Hybrid PDF Image Materialization + +## Summary + +- Kept Unstructured as the layout parser, including image/table coordinates and OCR text, while + making Poppler the primary byte-materialization path for PDF figures and tables. +- Added a provider fallback that requests both `Image` and `Table` payloads when local PDF + rasterization is unavailable or cannot completely resolve the document. +- Bumped the Unstructured parser policy to `unstructured@5`; its artifact hash now includes the + normalized filename/MIME, parser hints, and the effective provider request policy so a parser + strategy change cannot silently reuse an incompatible artifact. +- Installed `poppler-utils` in the API image, enabled it by default with an `off` kill switch, and + aligned local Compose, Dify Compose, and Kubernetes configuration. +- Made parse-artifact materialization atomic and idempotent. The repository reports + `created`/`replaced`/`unchanged`, preserves canonical object references on identical replays, and + reconciles commit-acknowledgement ambiguity before cleanup. +- Added execution-owner object namespaces and a final multimodal materialization digest. Runtime + object keys are excluded from that digest, while image/variant bytes, crop metadata, dimensions, + renderer contract, and DPI remain part of publication identity. + +## Runtime Bounds And Failure Recovery + +- Poppler renders each distinct `(page, DPI)` once and reuses the decoded page for all crops. + Main and thumbnail output therefore require two Poppler calls per page at the default distinct + DPIs, rather than two calls per image/table element. +- PDF materialization uses an independent process-wide admission gate (default `2`, configurable + `1..8`) around source loading, rendering, object writes, and canonical materialization. A caller + abort waits for any uncancellable Sharp/storage/commit operation to settle before cleanup. +- Rasterization is bounded by a ten-minute document deadline, 30-second Poppler child timeout, + 4096-pixel page edge, 20-million decoded page pixels, 32 MiB encoded page/image caps, 64 MiB + per-page crop output, 128 MiB per-document output, and 500 candidate elements. +- Relative, PDF-point, and Unstructured `PixelSpace` coordinates are scaled against actual rendered + page dimensions. Fractional crops use `ceil(x + width)`/`ceil(y + height)` to avoid truncation. +- Partial local results are never published as complete: an unresolved figure/table triggers + compensation and provider fallback for the whole unresolved set. +- Object cleanup is execution-owner scoped, pre-registers writes before PUT, runs with concurrency + `4` and a 60-second scheduling budget, and preserves retryability when storage or cleanup is + incomplete. This prevents transient cleanup failures from becoming terminal compilation errors. +- Object-storage HTTP calls have a 60-second per-request deadline and classify transient transport, + 408/409/425/429, and 5xx failures as retryable. +- Thumbnail generation rejects compressed-image expansion beyond 20 million pixels and encoded + variants larger than 8 MiB. + +## Parser Resource Controls + +- Unstructured calls use an abort-aware FIFO gate (default `2`), a 120-second request/body deadline, + a streamed 32 MiB response cap, and bounded retry delay. Oversized chunked responses are canceled + before additional chunks are buffered. +- 408/409/425/429/5xx, transport, and timeout failures are retryable; invalid input, invalid schema, + oversized response, and ordinary 4xx failures remain terminal. +- `image_base64` is removed before cloning provider metadata, avoiding an extra in-memory copy of + large inline image payloads. +- Parser-specific failure codes are retained by the durable compilation runtime. + +## Deployment And Precedence + +- The final API image installs `poppler-utils` and executes a build-time `pdftoppm -v` check before + switching to the non-root `node` user. The runtime bundle smoke repeats that check as `node`. +- Dify Compose uses six whitelisted `DIFY_ROOT_*_OVERRIDE` proxies. Precedence is explicit root + `docker/.env` value, then service `knowledge-fs.env`, then image/code default; no unrelated root + secrets are injected into the KnowledgeFS container. +- `KNOWLEDGE_PDF_RASTERIZER=off` remains the emergency kill switch. If Poppler is disabled or its + result is incomplete, provider payload fallback remains available. +- Existing PDFs published without durable `assetRef.objectKey` values must be re-indexed after + rollout; a missing object reference cannot be repaired by refreshing the viewer URL. + +## Regression Coverage + +- `@knowledge/api`: 411 test files passed, 4536 tests passed, 3 skipped; typecheck passed. +- `@knowledge/parsers`: 50 tests passed; 95.82% statements/lines, 90.10% branches, 97.45% + functions; typecheck passed. +- `@knowledge/adapters`: 7 test files and 108 tests passed; typecheck passed. +- `@knowledge/api-app`: 45 test files and 257 tests passed; typecheck passed. +- API image and Compose/Kubernetes deployment checks: 15 tests passed. +- Full KnowledgeFS backend Biome gate checked 1061 files with no errors. +- `git diff --check` passed. + +## Measured/Deterministic Performance Result + +- Same-page `N` image/table crops now issue exactly two Poppler page renders at the default distinct + main/thumbnail DPIs instead of `2N`; two pages issue four. This is asserted with real Poppler + subprocess tests. +- The supplied invoice PDF was materialized locally into a 119,090-byte main image and a + 24,599-byte thumbnail at 144/48 DPI, proving the configured executable and crop pipeline work on + the reported input. +- No end-to-end percentage improvement is claimed: deployment-level queue time, parser latency, + object-storage latency, and container CPU/RSS have not been measured on a before/after workload. + +## Local Limitation + +- A full Docker image build/runtime smoke was not run locally because the Docker daemon is + unavailable at the configured socket. The Dockerfile build itself now fails if `pdftoppm` is + absent, and CI's image build plus runtime smoke exercise that gate. +- Execution-owned object cleanup is bounded and retryable, but a hard process exit after object + storage commits and before in-memory ownership is durably recorded can leave an unreachable + owner prefix. A durable owner-receipt GC is a separate storage-lifecycle follow-up; shared + deterministic keys are intentionally avoided because overlapping at-least-once workers could + delete a newer successful writer's object. diff --git a/knowledge-fs/README.md b/knowledge-fs/README.md index c2c5edbecc5..b67e7a4b3e8 100644 --- a/knowledge-fs/README.md +++ b/knowledge-fs/README.md @@ -57,6 +57,7 @@ The canonical Dify Compose service loads - `DATABASE_URL` - `DIFY_INNER_API_URL` and `DIFY_INNER_API_KEY`, injected by Dify Compose - `UNSTRUCTURED_API_URL` and optional `UNSTRUCTURED_API_KEY` +- bounded Poppler PDF image rasterization settings (the API image includes `pdftoppm`) - KnowledgeFS capability/JWKS and document-compilation rollout settings Do not configure storage-provider credentials, model-provider keys, datasource credentials, or a diff --git a/knowledge-fs/apps/api/Dockerfile b/knowledge-fs/apps/api/Dockerfile index 9ae43640776..c489e933844 100644 --- a/knowledge-fs/apps/api/Dockerfile +++ b/knowledge-fs/apps/api/Dockerfile @@ -45,10 +45,25 @@ RUN sharp_runtime_source="$(dirname "$(realpath apps/api/node_modules/sharp)")" FROM node:22-bookworm-slim -ENV NODE_ENV=production +ENV NODE_ENV=production \ + KNOWLEDGE_PDF_RASTERIZER=poppler \ + KNOWLEDGE_PDF_RASTERIZER_DPI=144 \ + KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI=48 \ + KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS=30000 \ + KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS=500 \ + KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY=2 WORKDIR /workspace +# PDF image elements are rasterized outside Unstructured. Keep Poppler in the +# final image and fail the build if the executable is not usable on the target +# platform. Deployments can set KNOWLEDGE_PDF_RASTERIZER=off as a kill switch. +RUN apt-get update \ + && apt-get install --yes --no-install-recommends poppler-utils \ + && rm -rf /var/lib/apt/lists/* \ + && command -v pdftoppm >/dev/null \ + && pdftoppm -v >/dev/null 2>&1 + COPY --from=builder /workspace/apps/api/dist/server.mjs ./server.mjs COPY --from=builder /workspace/apps/api/dist/migrate.mjs ./migrate.mjs COPY --from=builder /runtime/node_modules ./node_modules diff --git a/knowledge-fs/apps/api/src/document-compilation-runtime-options.ts b/knowledge-fs/apps/api/src/document-compilation-runtime-options.ts index 89bf4b9fd90..d6e14cafd3a 100644 --- a/knowledge-fs/apps/api/src/document-compilation-runtime-options.ts +++ b/knowledge-fs/apps/api/src/document-compilation-runtime-options.ts @@ -51,6 +51,7 @@ import { type ProjectionSetPublicationMemberRepository, type ProjectionSetPublicationRepository, type SemanticChunker, + createConcurrencyGate, createDatabaseDocumentCompilationCandidateValidator, createDatabaseDocumentCompilationIndexOverrideResolver, createDatabaseDocumentLogicalMutationReconciler, @@ -157,7 +158,7 @@ export interface CreateApiDocumentCompilationRuntimeOptions { } | undefined; readonly multimodal?: - | Pick< + | (Pick< KnowledgeGatewayOptions, | "documentMultimodalImageVariantGenerator" | "documentMultimodalLocalAssetAllowlist" @@ -165,7 +166,7 @@ export interface CreateApiDocumentCompilationRuntimeOptions { | "documentMultimodalMaxLocalAssetBytes" | "documentMultimodalMaxPdfRasterizedAssets" | "documentPdfRasterizer" - > + > & { readonly documentMultimodalMaxConcurrency?: number | undefined }) | undefined; readonly repositories: Partial; readonly semantic?: @@ -298,6 +299,9 @@ export function createApiDocumentCompilationRuntime({ "Document compilation runtime requires bounded projection getMany and status updates", ); } + const multimodalMaterializationGate = createConcurrencyGate( + multimodal?.documentMultimodalMaxConcurrency ?? 2, + ); const compilationJobs = createDurableDocumentCompilationJobStateMachine({ assertCompilationAdmission: (input) => @@ -576,6 +580,7 @@ export function createApiDocumentCompilationRuntime({ multimodalImageVariantGenerator: multimodal.documentMultimodalImageVariantGenerator, } : {}), + multimodalMaterializationGate, ...(multimodal?.documentMultimodalLocalAssetAllowlist ? { multimodalLocalAssetAllowlist: multimodal.documentMultimodalLocalAssetAllowlist, @@ -586,6 +591,9 @@ export function createApiDocumentCompilationRuntime({ multimodalMaxExtractedAssets: multimodal.documentMultimodalMaxExtractedAssets, } : {}), + ...(multimodal + ? { multimodalMaxConcurrency: multimodal.documentMultimodalMaxConcurrency } + : {}), ...(multimodal?.documentMultimodalMaxLocalAssetBytes ? { multimodalMaxLocalAssetBytes: multimodal.documentMultimodalMaxLocalAssetBytes, diff --git a/knowledge-fs/apps/api/src/multimodal-options.test.ts b/knowledge-fs/apps/api/src/multimodal-options.test.ts index 68d658679b8..7c0ab087f54 100644 --- a/knowledge-fs/apps/api/src/multimodal-options.test.ts +++ b/knowledge-fs/apps/api/src/multimodal-options.test.ts @@ -1,11 +1,27 @@ -import { describe, expect, it } from "vitest"; +import { createPopplerPdfRasterizer } from "@knowledge/api"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("@knowledge/api", async (importOriginal) => { + const actual = await importOriginal(); + + return { + ...actual, + createPopplerPdfRasterizer: vi.fn(actual.createPopplerPdfRasterizer), + }; +}); import { createApiMultimodalOptions } from "./multimodal-options"; describe("createApiMultimodalOptions", () => { + beforeEach(() => { + vi.mocked(createPopplerPdfRasterizer).mockClear(); + }); + it("leaves PDF rasterization disabled by default or when explicitly off", () => { - expect(createApiMultimodalOptions({}).documentPdfRasterizer).toBeUndefined(); - expect(createApiMultimodalOptions({}).documentMultimodalImageVariantGenerator).toBeDefined(); + const defaults = createApiMultimodalOptions({}); + expect(defaults.documentPdfRasterizer).toBeUndefined(); + expect(defaults.documentMultimodalImageVariantGenerator).toBeDefined(); + expect(defaults.documentMultimodalMaxConcurrency).toBe(2); expect(createApiMultimodalOptions({ KNOWLEDGE_PDF_RASTERIZER: "off" })).toMatchObject({ documentMultimodalImageVariantGenerator: expect.any(Object), }); @@ -15,7 +31,9 @@ describe("createApiMultimodalOptions", () => { }); it("can disable or configure non-PDF image thumbnails", () => { - expect(createApiMultimodalOptions({ KNOWLEDGE_IMAGE_THUMBNAILS: "off" })).toEqual({}); + expect(createApiMultimodalOptions({ KNOWLEDGE_IMAGE_THUMBNAILS: "off" })).toEqual({ + documentMultimodalMaxConcurrency: 2, + }); expect( createApiMultimodalOptions({ KNOWLEDGE_IMAGE_THUMBNAIL_MAX_DIMENSION: "96", @@ -30,13 +48,70 @@ describe("createApiMultimodalOptions", () => { KNOWLEDGE_PDF_RASTERIZER: "poppler", KNOWLEDGE_PDF_RASTERIZER_DPI: "200", KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS: "25", + KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY: "3", KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI: "64", KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_VARIANT: "small", KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS: "5000", }); expect(options.documentPdfRasterizer).toBeDefined(); + expect(options.documentMultimodalMaxConcurrency).toBe(3); expect(options.documentMultimodalMaxPdfRasterizedAssets).toBe(25); + expect(createPopplerPdfRasterizer).toHaveBeenCalledWith( + expect.objectContaining({ maxConcurrency: 3 }), + ); + }); + + it("uses service-specific PDF values when root override proxies are empty", () => { + const options = createApiMultimodalOptions({ + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY_OVERRIDE: " ", + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_OVERRIDE: "", + KNOWLEDGE_PDF_RASTERIZER: "poppler", + KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY: "4", + }); + + expect(options.documentPdfRasterizer).toBeDefined(); + expect(options.documentMultimodalMaxConcurrency).toBe(4); + expect(createPopplerPdfRasterizer).toHaveBeenCalledWith( + expect.objectContaining({ maxConcurrency: 4 }), + ); + }); + + it("lets whitelisted root proxies override service-specific PDF values", () => { + const options = createApiMultimodalOptions({ + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_DPI_OVERRIDE: "180", + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS_OVERRIDE: "11", + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY_OVERRIDE: "3", + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_OVERRIDE: "poppler", + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI_OVERRIDE: "60", + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS_OVERRIDE: "7000", + KNOWLEDGE_PDF_RASTERIZER: "off", + KNOWLEDGE_PDF_RASTERIZER_DPI: "100", + KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS: "10", + KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY: "7", + KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI: "40", + KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS: "6000", + }); + + expect(options.documentMultimodalMaxPdfRasterizedAssets).toBe(11); + expect(options.documentMultimodalMaxConcurrency).toBe(3); + expect(createPopplerPdfRasterizer).toHaveBeenCalledWith( + expect.objectContaining({ + dpi: 180, + maxConcurrency: 3, + thumbnailDpi: 60, + timeoutMs: 7000, + }), + ); + + vi.mocked(createPopplerPdfRasterizer).mockClear(); + expect( + createApiMultimodalOptions({ + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_OVERRIDE: "off", + KNOWLEDGE_PDF_RASTERIZER: "poppler", + }).documentPdfRasterizer, + ).toBeUndefined(); + expect(createPopplerPdfRasterizer).not.toHaveBeenCalled(); }); it("creates a Poppler rasterizer when a command path is configured", () => { @@ -45,6 +120,9 @@ describe("createApiMultimodalOptions", () => { }); expect(options.documentPdfRasterizer).toBeDefined(); + expect(createPopplerPdfRasterizer).toHaveBeenCalledWith( + expect.objectContaining({ maxConcurrency: 2 }), + ); }); it("rejects invalid rasterizer environment values", () => { @@ -78,5 +156,19 @@ describe("createApiMultimodalOptions", () => { KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS: "many", }), ).toThrow("KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS must be a positive integer"); + for (const value of ["0", "9", "1.5", "many"]) { + expect(() => + createApiMultimodalOptions({ + KNOWLEDGE_PDF_RASTERIZER: "poppler", + KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY: value, + }), + ).toThrow(/KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY must be/u); + } + expect(() => + createApiMultimodalOptions({ + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY_OVERRIDE: "9", + DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_OVERRIDE: "poppler", + }), + ).toThrow("KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY must be an integer between 1 and 8"); }); }); diff --git a/knowledge-fs/apps/api/src/multimodal-options.ts b/knowledge-fs/apps/api/src/multimodal-options.ts index 2bb7c29077b..83187e6cf5a 100644 --- a/knowledge-fs/apps/api/src/multimodal-options.ts +++ b/knowledge-fs/apps/api/src/multimodal-options.ts @@ -6,6 +6,12 @@ import { } from "@knowledge/api"; export interface ApiMultimodalEnv { + readonly DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_DPI_OVERRIDE?: string | undefined; + readonly DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS_OVERRIDE?: string | undefined; + readonly DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY_OVERRIDE?: string | undefined; + readonly DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_OVERRIDE?: string | undefined; + readonly DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI_OVERRIDE?: string | undefined; + readonly DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS_OVERRIDE?: string | undefined; readonly KNOWLEDGE_IMAGE_THUMBNAILS?: string | undefined; readonly KNOWLEDGE_IMAGE_THUMBNAIL_MAX_DIMENSION?: string | undefined; readonly KNOWLEDGE_IMAGE_THUMBNAIL_VARIANT?: string | undefined; @@ -16,10 +22,12 @@ export interface ApiMultimodalEnv { readonly KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_VARIANT?: string | undefined; readonly KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS?: string | undefined; readonly KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS?: string | undefined; + readonly KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY?: string | undefined; } export interface ApiMultimodalOptions { readonly documentMultimodalImageVariantGenerator?: DocumentImageVariantGenerator; + readonly documentMultimodalMaxConcurrency: number; readonly documentMultimodalMaxPdfRasterizedAssets?: number; readonly documentPdfRasterizer?: DocumentPdfRasterizer; } @@ -27,14 +35,43 @@ export interface ApiMultimodalOptions { export function createApiMultimodalOptions( env: ApiMultimodalEnv = process.env, ): ApiMultimodalOptions { - const rasterizerName = normalizedRasterizer(env.KNOWLEDGE_PDF_RASTERIZER); + const rasterizerMode = rootOverride( + env.DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_OVERRIDE, + env.KNOWLEDGE_PDF_RASTERIZER, + ); + const rasterizerDpi = rootOverride( + env.DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_DPI_OVERRIDE, + env.KNOWLEDGE_PDF_RASTERIZER_DPI, + ); + const rasterizerThumbnailDpi = rootOverride( + env.DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI_OVERRIDE, + env.KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI, + ); + const rasterizerTimeoutMs = rootOverride( + env.DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS_OVERRIDE, + env.KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS, + ); + const rasterizerMaxAssets = rootOverride( + env.DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS_OVERRIDE, + env.KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS, + ); + const rasterizerMaxConcurrency = rootOverride( + env.DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY_OVERRIDE, + env.KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY, + ); + const rasterizerName = normalizedRasterizer(rasterizerMode); const command = trimmed(env.KNOWLEDGE_PDF_RASTERIZER_COMMAND); const imageVariantOptions = imageThumbnailOptions(env); + const maxConcurrency = boundedPositiveIntegerEnv( + rasterizerMaxConcurrency ?? "2", + "KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY", + 8, + ); const maxAssets = - env.KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS !== undefined + rasterizerMaxAssets !== undefined ? { documentMultimodalMaxPdfRasterizedAssets: positiveIntegerEnv( - env.KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS, + rasterizerMaxAssets, "KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS", ), } @@ -44,26 +81,25 @@ export function createApiMultimodalOptions( return { ...imageVariantOptions, ...maxAssets, + documentMultimodalMaxConcurrency: maxConcurrency, }; } return { ...imageVariantOptions, ...maxAssets, + documentMultimodalMaxConcurrency: maxConcurrency, documentPdfRasterizer: createPopplerPdfRasterizer({ ...(command ? { command } : {}), - ...(env.KNOWLEDGE_PDF_RASTERIZER_DPI !== undefined + ...(rasterizerDpi !== undefined ? { - dpi: positiveIntegerEnv( - env.KNOWLEDGE_PDF_RASTERIZER_DPI, - "KNOWLEDGE_PDF_RASTERIZER_DPI", - ), + dpi: positiveIntegerEnv(rasterizerDpi, "KNOWLEDGE_PDF_RASTERIZER_DPI"), } : {}), - ...(env.KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI !== undefined + ...(rasterizerThumbnailDpi !== undefined ? { thumbnailDpi: positiveIntegerEnv( - env.KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI, + rasterizerThumbnailDpi, "KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI", ), } @@ -71,14 +107,15 @@ export function createApiMultimodalOptions( ...(trimmed(env.KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_VARIANT) ? { thumbnailVariantName: trimmed(env.KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_VARIANT) } : {}), - ...(env.KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS !== undefined + ...(rasterizerTimeoutMs !== undefined ? { timeoutMs: positiveIntegerEnv( - env.KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS, + rasterizerTimeoutMs, "KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS", ), } : {}), + maxConcurrency, }), }; } @@ -161,8 +198,25 @@ function positiveIntegerEnv(value: string | undefined, name: string): number { return parsed; } +function boundedPositiveIntegerEnv(value: string | undefined, name: string, max: number): number { + const parsed = positiveIntegerEnv(value, name); + + if (parsed > max) { + throw new Error(`${name} must be an integer between 1 and ${max}`); + } + + return parsed; +} + function trimmed(value: string | undefined): string | undefined { const text = value?.trim(); return text ? text : undefined; } + +function rootOverride( + proxyValue: string | undefined, + serviceValue: string | undefined, +): string | undefined { + return trimmed(proxyValue) ?? serviceValue; +} diff --git a/knowledge-fs/apps/api/src/parser-options.test.ts b/knowledge-fs/apps/api/src/parser-options.test.ts index 70781748af7..f53007e0982 100644 --- a/knowledge-fs/apps/api/src/parser-options.test.ts +++ b/knowledge-fs/apps/api/src/parser-options.test.ts @@ -130,5 +130,15 @@ describe("createApiDocumentParser", () => { env: { UNSTRUCTURED_API_URL: "http://parser", UNSTRUCTURED_MAX_RESPONSE_BYTES: "0" }, }), ).toThrow("UNSTRUCTURED_MAX_RESPONSE_BYTES must be at least 1"); + expect(() => + createApiDocumentParser({ + env: { UNSTRUCTURED_API_URL: "http://parser", UNSTRUCTURED_MAX_CONCURRENCY: "33" }, + }), + ).toThrow("UNSTRUCTURED_MAX_CONCURRENCY must be between 1 and 32"); + expect(() => + createApiDocumentParser({ + env: { UNSTRUCTURED_API_URL: "http://parser", UNSTRUCTURED_REQUEST_TIMEOUT_MS: "600001" }, + }), + ).toThrow("UNSTRUCTURED_REQUEST_TIMEOUT_MS must be between 1 and 600000"); }); }); diff --git a/knowledge-fs/apps/api/src/parser-options.ts b/knowledge-fs/apps/api/src/parser-options.ts index 47c280f92a4..f5340ff99f1 100644 --- a/knowledge-fs/apps/api/src/parser-options.ts +++ b/knowledge-fs/apps/api/src/parser-options.ts @@ -11,10 +11,12 @@ export interface ApiParserEnv { readonly NODE_ENV?: string | undefined; readonly UNSTRUCTURED_API_KEY?: string | undefined; readonly UNSTRUCTURED_API_URL?: string | undefined; + readonly UNSTRUCTURED_MAX_CONCURRENCY?: string | undefined; readonly UNSTRUCTURED_MAX_RESPONSE_BYTES?: string | undefined; readonly UNSTRUCTURED_MAX_RETRIES?: string | undefined; readonly UNSTRUCTURED_PORT?: string | undefined; readonly UNSTRUCTURED_RETRY_DELAY_MS?: string | undefined; + readonly UNSTRUCTURED_REQUEST_TIMEOUT_MS?: string | undefined; } export interface CreateApiDocumentParserOptions { @@ -61,6 +63,15 @@ function createApiUnstructuredParser({ endpoint, ...(env.UNSTRUCTURED_API_KEY?.trim() ? { apiKey: env.UNSTRUCTURED_API_KEY.trim() } : {}), ...(fetchImpl ? { fetch: fetchImpl } : {}), + ...(env.UNSTRUCTURED_MAX_CONCURRENCY !== undefined + ? { + maxConcurrency: parseBoundedPositiveInteger( + env.UNSTRUCTURED_MAX_CONCURRENCY, + "UNSTRUCTURED_MAX_CONCURRENCY", + 32, + ), + } + : {}), ...(env.UNSTRUCTURED_MAX_RESPONSE_BYTES !== undefined ? { maxResponseBytes: parsePositiveInteger( @@ -85,6 +96,15 @@ function createApiUnstructuredParser({ ), } : {}), + ...(env.UNSTRUCTURED_REQUEST_TIMEOUT_MS !== undefined + ? { + requestTimeoutMs: parseBoundedPositiveInteger( + env.UNSTRUCTURED_REQUEST_TIMEOUT_MS, + "UNSTRUCTURED_REQUEST_TIMEOUT_MS", + 600_000, + ), + } + : {}), }); } @@ -116,6 +136,16 @@ function parsePort(value: string, name: string): number { return port; } +function parseBoundedPositiveInteger(value: string, name: string, max: number): number { + const parsed = parsePositiveInteger(value, name); + + if (parsed > max) { + throw new Error(`${name} must be between 1 and ${max}`); + } + + return parsed; +} + function parsePositiveInteger(value: string, name: string): number { const parsed = parseInteger(value, name); diff --git a/knowledge-fs/docs/production-deployment.md b/knowledge-fs/docs/production-deployment.md index bd5fa7efdea..349c336b7a0 100644 --- a/knowledge-fs/docs/production-deployment.md +++ b/knowledge-fs/docs/production-deployment.md @@ -15,6 +15,7 @@ This deployment does not migrate, replace, or delete existing Dataset/Document d | Physical object storage | Dify `STORAGE_TYPE` implementation | Dify inner storage API | | KnowledgeFS relational state | KnowledgeFS database | `DATABASE_URL` | | Complex document parsing | Unstructured-compatible service | `UNSTRUCTURED_API_URL` | +| PDF image rasterization | KnowledgeFS API image | Bundled Poppler `pdftoppm` executable | | Capability signing | Dify | Public JWKS only in KnowledgeFS | KnowledgeFS must never receive model-provider keys, datasource secrets, direct Plugin Daemon @@ -50,17 +51,51 @@ the service: |---|---| | `DATABASE_URL` | KnowledgeFS PostgreSQL connection string. | | `KNOWLEDGE_DOCUMENT_COMPILATION_RUNTIME` | Durable document worker rollout. | +| `KNOWLEDGE_PDF_RASTERIZER` | PDF image rasterizer. The production image defaults to `poppler`; set `off` as a kill switch. | +| `KNOWLEDGE_PDF_RASTERIZER_DPI` | Main PDF image resolution; the bounded deployment default is `144`. | +| `KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI` | Thumbnail resolution; the bounded deployment default is `48`. | +| `KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS` | Poppler subprocess timeout; the deployment default is `30000`. | +| `KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS` | Maximum PDF assets rasterized for one document; the deployment default is `500`. | +| `KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY` | Maximum concurrent Poppler page batches per API replica; defaults to `2` and accepts `1..8`. | | `KNOWLEDGE_FS_CAPABILITY_V2_ENABLED` | Capability-v2 verifier rollout. | | `KNOWLEDGE_FS_CAPABILITY_V2_PUBLIC_JWKS` | Public verification key set issued by Dify. | | `KNOWLEDGE_QUERY_IMAGE_RETRIEVAL_ENABLED` | Opt in to query-image visual retrieval; requires an enabled visual-embedding provider/index and a query mode other than `off`. | | `KNOWLEDGE_QUERY_IMAGE_EXPANSION_TIMEOUT_MS` | Timeout for the single Deep/Research vision expansion call; defaults to 8000 ms. | | `UNSTRUCTURED_API_URL` | Parser endpoint for complex formats. | | `UNSTRUCTURED_API_KEY` | Optional parser authentication. | +| `UNSTRUCTURED_MAX_CONCURRENCY` | Process-wide parser request limit; defaults to `2`. | +| `UNSTRUCTURED_REQUEST_TIMEOUT_MS` | Total timeout for one parser request and response body; defaults to `120000`. | +| `UNSTRUCTURED_MAX_RESPONSE_BYTES` | Maximum parser response body; defaults to `33554432` (32 MiB). | Compose injects `DIFY_INNER_API_URL` and `DIFY_INNER_API_KEY`; do not duplicate them in the operator-owned env file. Do not add `MINIO_*`, cloud object-storage credentials, provider API keys, `PLUGIN_DAEMON_*`, datasource tokens, or OAuth client secrets. +`DIFY_OBJECT_STORAGE_REQUEST_TIMEOUT_MS` bounds each authenticated inner object-storage request, +including response consumption. It defaults to `60000`; transport failures and +408/409/425/429/5xx responses remain retryable at the durable compilation layer. + +## PDF image rasterization + +The production API image installs Poppler and verifies `pdftoppm` during the image build. Its image +defaults enable rasterization at 144 DPI, generate 48 DPI thumbnails, stop an individual Poppler +operation after 30 seconds, cap one document at 500 rasterized assets, and run at most two Poppler +page batches concurrently per API replica. The process still runs as the unprivileged `node` user. + +The Dify Compose service keeps canonical values in +`docker/envs/core-services/knowledge-fs.env`, where they override the image defaults. It maps only +whitelisted `DIFY_ROOT_*_OVERRIDE` proxies, so an explicitly set PDF rasterizer value in +`docker/.env` takes precedence without exposing the rest of the root environment. An unset or +empty root value leaves the service env (or image default when that file is absent) in control. +Set `KNOWLEDGE_PDF_RASTERIZER=off` in either operator env during an incident or on a deliberately +constrained deployment. Do not set +`KNOWLEDGE_PDF_RASTERIZER_COMMAND` for the published image; its bundled command is on `PATH`. + +Rasterization supplies durable image objects for PDF image elements when the parser returns layout +coordinates without image bytes. It does not repair already-published parse artifacts. Re-run the +document ingestion after deploying the corrected image to repopulate images that were previously +stored without an asset reference. + ## Image-query rollout The public Dify API accepts query images as actor-owned Dify `UploadFile` references. KnowledgeFS diff --git a/knowledge-fs/infra/kubernetes/README.md b/knowledge-fs/infra/kubernetes/README.md index 9c270153c83..f36d20977da 100644 --- a/knowledge-fs/infra/kubernetes/README.md +++ b/knowledge-fs/infra/kubernetes/README.md @@ -16,6 +16,12 @@ production verifier. The baseline keeps Capability v2, integrated mode, direct u streaming explicitly disabled, so any unchanged pod reports `503` from `/ready` even if `/health` is `200`. +The API image bundles Poppler. The baseline ConfigMap enables `pdftoppm` with 144 DPI images, 48 +DPI thumbnails, a 30-second subprocess timeout, a 500-asset per-document cap, and at most two +concurrent Poppler page batches per API replica. Set +`KNOWLEDGE_PDF_RASTERIZER=off` in the downstream ConfigMap as an explicit incident or low-resource +kill switch; do not add a host-specific command path to the container configuration. + Before a later controlled scale-up, replace the image with an immutable digest and create the `knowledge-fs-runtime` Secret with dedicated `DATABASE_URL` and `DIFY_INNER_API_KEY`. Model, datasource, and object-storage credentials stay in Dify and are never copied into KnowledgeFS. diff --git a/knowledge-fs/infra/kubernetes/dify-integration-baseline.yaml b/knowledge-fs/infra/kubernetes/dify-integration-baseline.yaml index dc172a1151e..fe642b2d3de 100644 --- a/knowledge-fs/infra/kubernetes/dify-integration-baseline.yaml +++ b/knowledge-fs/infra/kubernetes/dify-integration-baseline.yaml @@ -10,6 +10,12 @@ data: PORT: "8787" KNOWLEDGE_DATABASE_REPOSITORIES: "on" KNOWLEDGE_DOCUMENT_COMPILATION_RUNTIME: "on" + KNOWLEDGE_PDF_RASTERIZER: poppler + KNOWLEDGE_PDF_RASTERIZER_DPI: "144" + KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI: "48" + KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS: "30000" + KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS: "500" + KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY: "2" KNOWLEDGE_FS_CAPABILITY_V2_ENABLED: "false" KNOWLEDGE_INTEGRATED_MODE_ENABLED: "false" KNOWLEDGE_LEGACY_ACL_READ_ONLY: "false" @@ -18,11 +24,15 @@ data: KNOWLEDGE_DIRECT_STREAM_ENABLED: "off" DURABLE_DELETION_ENABLED: "off" DIFY_INNER_API_URL: http://api:5001 + DIFY_OBJECT_STORAGE_REQUEST_TIMEOUT_MS: "60000" DIFY_DATASOURCE_RUNTIME_MAX_RESPONSE_BYTES: "8388608" DIFY_DATASOURCE_RUNTIME_REQUEST_TIMEOUT_MS: "60000" DIFY_MODEL_RUNTIME_MAX_RESPONSE_BYTES: "8388608" DIFY_MODEL_RUNTIME_REQUEST_TIMEOUT_MS: "60000" UNSTRUCTURED_API_URL: http://unstructured:8000 + UNSTRUCTURED_MAX_CONCURRENCY: "2" + UNSTRUCTURED_REQUEST_TIMEOUT_MS: "120000" + UNSTRUCTURED_MAX_RESPONSE_BYTES: "33554432" --- apiVersion: apps/v1 kind: Deployment diff --git a/knowledge-fs/infra/local/.env.example b/knowledge-fs/infra/local/.env.example index 492dc9700f4..d4a41a66222 100644 --- a/knowledge-fs/infra/local/.env.example +++ b/knowledge-fs/infra/local/.env.example @@ -12,13 +12,16 @@ DURABLE_DELETION_HMAC_KEY_BASE64= UNSTRUCTURED_PORT=8000 UNSTRUCTURED_API_URL=http://127.0.0.1:8000 UNSTRUCTURED_API_KEY= -UNSTRUCTURED_MAX_RESPONSE_BYTES= +UNSTRUCTURED_MAX_CONCURRENCY=2 +UNSTRUCTURED_REQUEST_TIMEOUT_MS=120000 +UNSTRUCTURED_MAX_RESPONSE_BYTES=33554432 UNSTRUCTURED_MAX_RETRIES= UNSTRUCTURED_RETRY_DELAY_MS= # Dify is required. It owns object storage, model/plugin credentials, and datasource invocation. DIFY_INNER_API_URL=http://host.docker.internal:5001 DIFY_INNER_API_KEY= +DIFY_OBJECT_STORAGE_REQUEST_TIMEOUT_MS=60000 DIFY_DATASOURCE_RUNTIME_MAX_RESPONSE_BYTES=8388608 DIFY_DATASOURCE_RUNTIME_REQUEST_TIMEOUT_MS=60000 DIFY_MODEL_RUNTIME_MAX_RESPONSE_BYTES=8388608 @@ -26,6 +29,14 @@ DIFY_MODEL_RUNTIME_REQUEST_TIMEOUT_MS=60000 # Rollout/cutover gate only; it never enables a standalone runtime. KNOWLEDGE_INTEGRATED_MODE_ENABLED=false +# The API image includes Poppler. Use off as an emergency/low-resource kill switch. +KNOWLEDGE_PDF_RASTERIZER=poppler +KNOWLEDGE_PDF_RASTERIZER_DPI=144 +KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI=48 +KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS=30000 +KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS=500 +KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY=2 + API_PORT=8788 ADMIN_PORT=3000 KNOWLEDGE_API_BASE_URL=http://localhost:8788 diff --git a/knowledge-fs/infra/local/README.md b/knowledge-fs/infra/local/README.md index ea9b8d74323..121b3ba2943 100644 --- a/knowledge-fs/infra/local/README.md +++ b/knowledge-fs/infra/local/README.md @@ -9,6 +9,7 @@ model/plugin credentials, and datasource invocation. - A running Dify API reachable through `DIFY_INNER_API_URL`. - The matching Dify inner API key in `DIFY_INNER_API_KEY`. - Docker for the local PostgreSQL and Unstructured dependencies. +- Poppler (`pdftoppm`) when processing PDFs with `pnpm dev:api`; the API container already includes it. Copy `infra/local/.env.example` to the ignored `infra/local/.env` and set the Dify URL and key. Do not add MinIO, cloud-storage credentials, model-provider keys, or datasource credentials to the @@ -86,6 +87,11 @@ Default local endpoints: - KnowledgeFS API readiness: `http://localhost:8788/ready` - Unstructured API: `http://localhost:8000` +The app-profile API enables its bundled Poppler rasterizer with the same bounded defaults as the +production image, including at most two concurrent Poppler page batches. Set +`KNOWLEDGE_PDF_RASTERIZER=off` in the ignored `infra/local/.env` to disable +it, or install Poppler on the host before using the source-run API for PDFs. + This Compose file is not a supported KnowledgeFS deployment topology. Production KnowledgeFS must be started by Dify's Compose or Kubernetes deployment so the internal URL, authentication, storage, models, datasources, and lifecycle are wired together. diff --git a/knowledge-fs/infra/local/compose.yaml b/knowledge-fs/infra/local/compose.yaml index a39cc4eea48..042f204c1a5 100644 --- a/knowledge-fs/infra/local/compose.yaml +++ b/knowledge-fs/infra/local/compose.yaml @@ -38,17 +38,29 @@ services: DATABASE_URL: postgresql://${POSTGRES_USER:-knowledge_fs}:${POSTGRES_PASSWORD:-knowledge_fs}@postgres:5432/${POSTGRES_DB:-knowledge_fs} DIFY_INNER_API_KEY: ${DIFY_INNER_API_KEY:-} DIFY_INNER_API_URL: ${DIFY_INNER_API_URL:-http://host.docker.internal:5001} + DIFY_OBJECT_STORAGE_REQUEST_TIMEOUT_MS: ${DIFY_OBJECT_STORAGE_REQUEST_TIMEOUT_MS:-60000} DIFY_DATASOURCE_RUNTIME_MAX_RESPONSE_BYTES: ${DIFY_DATASOURCE_RUNTIME_MAX_RESPONSE_BYTES:-8388608} DIFY_DATASOURCE_RUNTIME_REQUEST_TIMEOUT_MS: ${DIFY_DATASOURCE_RUNTIME_REQUEST_TIMEOUT_MS:-60000} DIFY_MODEL_RUNTIME_MAX_RESPONSE_BYTES: ${DIFY_MODEL_RUNTIME_MAX_RESPONSE_BYTES:-8388608} DIFY_MODEL_RUNTIME_REQUEST_TIMEOUT_MS: ${DIFY_MODEL_RUNTIME_REQUEST_TIMEOUT_MS:-60000} KNOWLEDGE_INTEGRATED_MODE_ENABLED: ${KNOWLEDGE_INTEGRATED_MODE_ENABLED:-false} + KNOWLEDGE_PDF_RASTERIZER: ${KNOWLEDGE_PDF_RASTERIZER:-poppler} + KNOWLEDGE_PDF_RASTERIZER_DPI: ${KNOWLEDGE_PDF_RASTERIZER_DPI:-144} + KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI: ${KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI:-48} + KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS: ${KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS:-30000} + KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS: ${KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS:-500} + KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY: ${KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY:-2} NODE_ENV: development PORT: 8787 KNOWLEDGE_DEV_AUTH_TOKEN: ${KNOWLEDGE_DEV_AUTH_TOKEN:-dev-token} KNOWLEDGE_DEV_SUBJECT_ID: ${KNOWLEDGE_DEV_SUBJECT_ID:-dev-user} KNOWLEDGE_DEV_TENANT_ID: ${KNOWLEDGE_DEV_TENANT_ID:-tenant-dev} - UNSTRUCTURED_API_URL: http://unstructured:8000 + UNSTRUCTURED_API_URL: ${UNSTRUCTURED_API_URL:-http://unstructured:8000} + UNSTRUCTURED_MAX_CONCURRENCY: ${UNSTRUCTURED_MAX_CONCURRENCY:-2} + UNSTRUCTURED_REQUEST_TIMEOUT_MS: ${UNSTRUCTURED_REQUEST_TIMEOUT_MS:-120000} + UNSTRUCTURED_MAX_RESPONSE_BYTES: ${UNSTRUCTURED_MAX_RESPONSE_BYTES:-33554432} + UNSTRUCTURED_MAX_RETRIES: ${UNSTRUCTURED_MAX_RETRIES:-0} + UNSTRUCTURED_RETRY_DELAY_MS: ${UNSTRUCTURED_RETRY_DELAY_MS:-100} ports: - "${API_PORT:-8788}:8787" healthcheck: diff --git a/knowledge-fs/packages/adapters/src/dify-object-storage.test.ts b/knowledge-fs/packages/adapters/src/dify-object-storage.test.ts index db5e980f3cc..85770f0f4a1 100644 --- a/knowledge-fs/packages/adapters/src/dify-object-storage.test.ts +++ b/knowledge-fs/packages/adapters/src/dify-object-storage.test.ts @@ -1,6 +1,9 @@ import { describe, expect, it, vi } from "vitest"; -import { createDifyObjectStorageAdapter } from "./dify-object-storage"; +import { + DifyObjectStorageRequestError, + createDifyObjectStorageAdapter, +} from "./dify-object-storage"; const metadata = { checksumSha256Base64: "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=", @@ -81,4 +84,72 @@ describe("Dify object storage adapter", () => { await expect(adapter.health()).resolves.toBe(false); }); + + it("classifies transient transport and HTTP failures as retryable", async () => { + const offline = createDifyObjectStorageAdapter({ + apiKey: "inner-key", + baseUrl: "http://api:5001", + fetch: vi.fn().mockRejectedValue(new Error("offline")), + }); + const offlineRequest = offline.headObject("tenant-1/offline"); + await expect(offlineRequest).rejects.toBeInstanceOf(DifyObjectStorageRequestError); + await expect(offlineRequest).rejects.toMatchObject({ retryable: true }); + + const statusFetch = vi + .fn() + .mockResolvedValueOnce(new Response(null, { status: 503 })) + .mockResolvedValueOnce(new Response(null, { status: 400 })); + const statuses = createDifyObjectStorageAdapter({ + apiKey: "inner-key", + baseUrl: "http://api:5001", + fetch: statusFetch, + }); + await expect(statuses.deleteObject("tenant-1/transient")).rejects.toMatchObject({ + retryable: true, + status: 503, + }); + await expect(statuses.deleteObject("tenant-1/invalid")).rejects.toMatchObject({ + retryable: false, + status: 400, + }); + }); + + it("bounds stalled Dify object-storage headers and response bodies", async () => { + const stalledHeaders = createDifyObjectStorageAdapter({ + apiKey: "inner-key", + baseUrl: "http://api:5001", + fetch: async (input, init) => { + const signal = init?.signal ?? (input instanceof Request ? input.signal : undefined); + return await new Promise((_resolve, reject) => { + signal?.addEventListener("abort", () => reject(signal.reason), { once: true }); + }); + }, + requestTimeoutMs: 10, + }); + await expect(stalledHeaders.headObject("tenant-1/headers")).rejects.toMatchObject({ + retryable: true, + }); + + const stalledBody = createDifyObjectStorageAdapter({ + apiKey: "inner-key", + baseUrl: "http://api:5001", + fetch: async (input, init) => { + const signal = init?.signal ?? (input instanceof Request ? input.signal : undefined); + return new Response( + new ReadableStream({ + start: (controller) => { + signal?.addEventListener("abort", () => controller.error(signal.reason), { + once: true, + }); + }, + }), + { status: 200 }, + ); + }, + requestTimeoutMs: 10, + }); + await expect(stalledBody.headObject("tenant-1/body")).rejects.toMatchObject({ + retryable: true, + }); + }); }); diff --git a/knowledge-fs/packages/adapters/src/dify-object-storage.ts b/knowledge-fs/packages/adapters/src/dify-object-storage.ts index 3663596a67a..23d575f9424 100644 --- a/knowledge-fs/packages/adapters/src/dify-object-storage.ts +++ b/knowledge-fs/packages/adapters/src/dify-object-storage.ts @@ -10,13 +10,35 @@ export interface DifyObjectStorageOptions { readonly baseUrl: string; readonly fetch?: typeof globalThis.fetch; readonly maxObjectBytes?: number; + readonly requestTimeoutMs?: number; } const defaultMaxObjectBytes = 64 * 1024 * 1024; +const defaultRequestTimeoutMs = 60_000; const metadataHeader = "X-Knowledge-FS-Metadata"; const checksumHeader = "X-Knowledge-FS-Checksum-Sha256"; const contentTypeHeader = "X-Knowledge-FS-Content-Type"; +export class DifyObjectStorageRequestError extends Error { + readonly code = "dify_object_storage_request_failed"; + readonly retryable: boolean; + readonly status?: number; + + constructor( + message: string, + options: { + readonly cause?: unknown; + readonly retryable: boolean; + readonly status?: number; + }, + ) { + super(message, options.cause === undefined ? undefined : { cause: options.cause }); + this.name = "DifyObjectStorageRequestError"; + this.retryable = options.retryable; + if (options.status !== undefined) this.status = options.status; + } +} + /** * Uses Dify's authenticated inner API as the only physical object-storage owner. The adapter * deliberately omits direct-upload capabilities because Dify's portable @@ -27,77 +49,134 @@ export function createDifyObjectStorageAdapter({ baseUrl, fetch = globalThis.fetch, maxObjectBytes = defaultMaxObjectBytes, + requestTimeoutMs = defaultRequestTimeoutMs, }: DifyObjectStorageOptions): ObjectStorageAdapter { const normalizedBaseUrl = requiredBaseUrl(baseUrl); const normalizedApiKey = requiredString(apiKey, "Dify inner API key"); positiveSafeInteger(maxObjectBytes, "maxObjectBytes"); + positiveSafeInteger(requestTimeoutMs, "requestTimeoutMs"); - const request = (path: string, init: RequestInit = {}) => - fetch(new URL(path, normalizedBaseUrl), { - ...init, - headers: { - ...headersRecord(init.headers), - "X-Inner-Api-Key": normalizedApiKey, - }, - }); + const request = async (path: string, init: RequestInit = {}) => { + const deadline = createRequestDeadline(requestTimeoutMs); + try { + const response = await fetch(new URL(path, normalizedBaseUrl), { + ...init, + headers: { + ...headersRecord(init.headers), + "X-Inner-Api-Key": normalizedApiKey, + }, + signal: deadline.signal, + }); + return { deadline, response }; + } catch (error) { + deadline.dispose(); + throw requestTransportError(error, deadline.expired(), requestTimeoutMs); + } + }; + + const withResponse = async ( + path: string, + init: RequestInit, + consume: (response: Response) => Promise | T, + ): Promise => { + const { deadline, response } = await request(path, init); + try { + const result = await consume(response); + deadline.throwIfExpired(); + return result; + } catch (error) { + if (deadline.expired()) { + throw requestTransportError(error, true, requestTimeoutMs); + } + throw error; + } finally { + deadline.dispose(); + } + }; return { kind: "dify", deleteObject: async (key) => { - const response = await request( + await withResponse( objectPath("/inner/api/knowledge-fs/storage/object", { key }), { method: "DELETE", }, + (response) => assertStatus(response, [204]), ); - assertStatus(response, [204]); }, getObject: async (key) => { - const response = await request(objectPath("/inner/api/knowledge-fs/storage/object", { key })); - if (response.status === 404) return null; - assertStatus(response, [200]); - return readBoundedBody(response, maxObjectBytes); + return withResponse( + objectPath("/inner/api/knowledge-fs/storage/object", { key }), + {}, + async (response) => { + if (response.status === 404) return null; + assertStatus(response, [200]); + return readBoundedBody(response, maxObjectBytes); + }, + ); }, getObjectStream: async (key) => { - const response = await request(objectPath("/inner/api/knowledge-fs/storage/object", { key })); - if (response.status === 404) return null; - assertStatus(response, [200]); - return boundedResponseStream(response, maxObjectBytes); + const { deadline, response } = await request( + objectPath("/inner/api/knowledge-fs/storage/object", { key }), + ); + if (response.status === 404) { + deadline.dispose(); + return null; + } + try { + assertStatus(response, [200]); + return boundedResponseStream(response, maxObjectBytes, deadline, requestTimeoutMs); + } catch (error) { + deadline.dispose(); + throw error; + } }, health: async () => { try { - const response = await request("/inner/api/knowledge-fs/storage/health"); - if (!response.ok) return false; - const payload = asRecord(await response.json()); - return payload?.ok === true; + return await withResponse( + "/inner/api/knowledge-fs/storage/health", + {}, + async (response) => { + if (!response.ok) return false; + const payload = asRecord(await response.json()); + return payload?.ok === true; + }, + ); } catch { return false; } }, headObject: async (key) => { - const response = await request( + return withResponse( objectPath("/inner/api/knowledge-fs/storage/object/metadata", { key }), + {}, + async (response) => { + if (response.status === 404) return null; + assertStatus(response, [200]); + return parseObjectMetadata(await response.json()); + }, ); - if (response.status === 404) return null; - assertStatus(response, [200]); - return parseObjectMetadata(await response.json()); }, listObjects: async ({ cursor, limit, prefix }) => { - const response = await request( + return withResponse( objectPath("/inner/api/knowledge-fs/storage/objects", { ...(cursor ? { cursor } : {}), limit: String(limit), prefix, }), + {}, + async (response) => { + assertStatus(response, [200]); + return parseObjectList(await response.json()); + }, ); - assertStatus(response, [200]); - return parseObjectList(await response.json()); }, putObject: async (input) => { if (input.body.byteLength > maxObjectBytes) { throw new Error(`Object ${input.key} exceeds maxObjectBytes=${maxObjectBytes}`); } - const response = await request( + return withResponse( objectPath("/inner/api/knowledge-fs/storage/object", { key: input.key }), { body: requestBody(input.body), @@ -110,9 +189,11 @@ export function createDifyObjectStorageAdapter({ }, method: "PUT", }, + async (response) => { + assertStatus(response, [200]); + return parseObjectMetadata(await response.json()); + }, ); - assertStatus(response, [200]); - return parseObjectMetadata(await response.json()); }, }; } @@ -164,7 +245,10 @@ function headersRecord(headers: HeadersInit | undefined): Record function assertStatus(response: Response, expected: readonly number[]): void { if (!expected.includes(response.status)) { - throw new Error(`Dify object storage request failed with status ${response.status}`); + throw new DifyObjectStorageRequestError( + `Dify object storage request failed with status ${response.status}`, + { retryable: isRetryableStatus(response.status), status: response.status }, + ); } } @@ -246,32 +330,97 @@ async function readBoundedBody(response: Response, maxObjectBytes: number): Prom function boundedResponseStream( response: Response, maxObjectBytes: number, + deadline?: RequestDeadline, + requestTimeoutMs = defaultRequestTimeoutMs, ): ReadableStream { const declaredLength = Number(response.headers.get("content-length")); if (Number.isFinite(declaredLength) && declaredLength > maxObjectBytes) { throw new Error(`Dify object storage response exceeds maxObjectBytes=${maxObjectBytes}`); } const source = response.body; - if (!source) return new ReadableStream({ start: (controller) => controller.close() }); + if (!source) { + deadline?.dispose(); + return new ReadableStream({ start: (controller) => controller.close() }); + } const reader = source.getReader(); let totalBytes = 0; return new ReadableStream({ - cancel: (reason) => reader.cancel(reason), + cancel: async (reason) => { + try { + await reader.cancel(reason); + } finally { + deadline?.dispose(); + } + }, async pull(controller) { - const chunk = await reader.read(); - if (chunk.done) { - controller.close(); - return; + try { + const chunk = await reader.read(); + if (chunk.done) { + deadline?.throwIfExpired(); + deadline?.dispose(); + controller.close(); + return; + } + totalBytes += chunk.value.byteLength; + if (totalBytes > maxObjectBytes) { + await reader.cancel(); + deadline?.dispose(); + controller.error( + new Error(`Dify object storage response exceeds maxObjectBytes=${maxObjectBytes}`), + ); + return; + } + controller.enqueue(chunk.value); + } catch (error) { + const timedOut = deadline?.expired() === true; + deadline?.dispose(); + controller.error(timedOut ? requestTransportError(error, true, requestTimeoutMs) : error); } - totalBytes += chunk.value.byteLength; - if (totalBytes > maxObjectBytes) { - await reader.cancel(); - controller.error( - new Error(`Dify object storage response exceeds maxObjectBytes=${maxObjectBytes}`), - ); - return; - } - controller.enqueue(chunk.value); }, }); } + +interface RequestDeadline { + readonly signal: AbortSignal; + dispose(): void; + expired(): boolean; + throwIfExpired(): void; +} + +function createRequestDeadline(requestTimeoutMs: number): RequestDeadline { + const controller = new AbortController(); + const timeoutReason = new Error("Dify object storage request deadline exceeded"); + let timedOut = false; + const timer = setTimeout(() => { + timedOut = true; + controller.abort(timeoutReason); + }, requestTimeoutMs); + (timer as { unref?: () => void }).unref?.(); + + return { + signal: controller.signal, + dispose: () => clearTimeout(timer), + expired: () => timedOut, + throwIfExpired: () => { + if (timedOut) throw timeoutReason; + }, + }; +} + +function requestTransportError( + cause: unknown, + timedOut: boolean, + requestTimeoutMs: number, +): DifyObjectStorageRequestError { + if (cause instanceof DifyObjectStorageRequestError) return cause; + return new DifyObjectStorageRequestError( + timedOut + ? `Dify object storage request timed out after requestTimeoutMs=${requestTimeoutMs}` + : "Dify object storage request failed", + { cause, retryable: true }, + ); +} + +function isRetryableStatus(status: number): boolean { + return status === 408 || status === 409 || status === 425 || status === 429 || status >= 500; +} diff --git a/knowledge-fs/packages/adapters/src/node.ts b/knowledge-fs/packages/adapters/src/node.ts index 14b9320166b..4931087d125 100644 --- a/knowledge-fs/packages/adapters/src/node.ts +++ b/knowledge-fs/packages/adapters/src/node.ts @@ -22,6 +22,7 @@ export interface NodePlatformAdapterOptions { } const maxObjectBytes = 64 * 1024 * 1024; +const defaultObjectStorageRequestTimeoutMs = 60_000; const defaultDifyInnerApiUrl = "http://localhost:5001"; const defaultDifyInnerApiKey = "QaHbTe77CtuXmsfyhR7+vRjI/+XbV1AaFy691iy+kGDv2Jvy0/eAh8Y1"; @@ -87,6 +88,10 @@ function createNodeObjectStorageAdapter( baseUrl: env.DIFY_INNER_API_URL?.trim() || defaultDifyInnerApiUrl, ...(difyStorageFetch ? { fetch: difyStorageFetch } : {}), maxObjectBytes, + requestTimeoutMs: parsePositiveInteger( + env.DIFY_OBJECT_STORAGE_REQUEST_TIMEOUT_MS, + defaultObjectStorageRequestTimeoutMs, + ), }); } diff --git a/knowledge-fs/packages/api/src/document-compilation-runtime.test.ts b/knowledge-fs/packages/api/src/document-compilation-runtime.test.ts index e158b5d01ed..67934ca316d 100644 --- a/knowledge-fs/packages/api/src/document-compilation-runtime.test.ts +++ b/knowledge-fs/packages/api/src/document-compilation-runtime.test.ts @@ -119,6 +119,32 @@ describe("createDocumentCompilationRuntime", () => { }); }); + it.each([ + ["provider_input", false, "DOCUMENT_PARSER_INPUT_INVALID"], + ["provider_rate_limited", true, "DOCUMENT_PARSER_RATE_LIMITED"], + ["provider_request_failed", true, "DOCUMENT_PARSER_UNAVAILABLE"], + ["provider_response_invalid", false, "DOCUMENT_PARSER_RESPONSE_INVALID"], + ] as const)("maps %s to the parser-specific durable failure", (code, retryable, expectedCode) => { + const error = Object.assign(new Error(code), { code, retryable }); + + expect(defaultDocumentCompilationErrorClassifier(error)).toMatchObject({ + code: expectedCode, + retryable, + }); + }); + + it("keeps transient object-storage failures retryable", () => { + const error = Object.assign(new Error("storage unavailable"), { + code: "dify_object_storage_request_failed", + retryable: true, + }); + + expect(defaultDocumentCompilationErrorClassifier(error)).toMatchObject({ + code: "DOCUMENT_COMPILATION_RETRYABLE", + retryable: true, + }); + }); + it.each([null, [], {}, { attemptId: 42 }, { attemptId: "not-a-uuid" }])( "rejects malformed attempt locator payload %j", async (payload) => { diff --git a/knowledge-fs/packages/api/src/document-compilation-runtime.ts b/knowledge-fs/packages/api/src/document-compilation-runtime.ts index f3e3f889940..c5877b8b64f 100644 --- a/knowledge-fs/packages/api/src/document-compilation-runtime.ts +++ b/knowledge-fs/packages/api/src/document-compilation-runtime.ts @@ -694,6 +694,12 @@ export function defaultDocumentCompilationErrorClassifier( }); return { code: failure.code, message: failure.message, retryable: error.retryable }; } + if (isParserProviderError(error)) { + const failure = knowledgeFsFailureForCode(error.code, { + ...(stage ? { stage } : {}), + }); + return { code: failure.code, message: failure.message, retryable: error.retryable === true }; + } if (isRetryableProviderError(error)) { const failure = knowledgeFsFailureForCode( typeof error.code === "string" && error.code.trim() @@ -717,6 +723,22 @@ export function defaultDocumentCompilationErrorClassifier( }; } +function isParserProviderError( + error: unknown, +): error is Error & { readonly code: string; readonly retryable?: boolean } { + return ( + error instanceof Error && + "code" in error && + typeof (error as { readonly code?: unknown }).code === "string" && + [ + "provider_input", + "provider_rate_limited", + "provider_request_failed", + "provider_response_invalid", + ].includes((error as { readonly code: string }).code) + ); +} + function isRetryableProviderError( error: unknown, ): error is Error & { readonly code?: unknown; readonly retryable: true } { diff --git a/knowledge-fs/packages/api/src/document-compilation-worker.test.ts b/knowledge-fs/packages/api/src/document-compilation-worker.test.ts index 3819dba5d0d..6f054645589 100644 --- a/knowledge-fs/packages/api/src/document-compilation-worker.test.ts +++ b/knowledge-fs/packages/api/src/document-compilation-worker.test.ts @@ -1,7 +1,7 @@ import { createMemoryObjectStorageAdapter } from "@knowledge/adapters"; import { createNodePlatformAdapter } from "@knowledge/adapters/node"; import { ParseArtifactSchema, type PlatformAdapter } from "@knowledge/core"; -import type { ParserAdapter } from "@knowledge/parsers"; +import type { ParserAdapter, ParserRouteHints } from "@knowledge/parsers"; import { describe, expect, it } from "vitest"; import { @@ -11,6 +11,7 @@ import { } from "./deletion-lifecycle-fence"; import { + createConcurrencyGate, createDocumentCompilationJobStateMachine, createDocumentCompilationWorker, createDocumentMultimodalManifestBuilder, @@ -995,7 +996,7 @@ describe("createDocumentCompilationWorker lease integration", () => { }), }, reindexer: { - canonicalizeArtifact: async (input) => artifacts.create(input), + canonicalizeArtifact: async (input) => artifacts.materialize(input), reindex: async (input) => { reindexInputArtifactIds.push(input.parseArtifact.id); const canonicalArtifact = await artifacts.create(input.parseArtifact); @@ -1039,7 +1040,7 @@ describe("createDocumentCompilationWorker lease integration", () => { expect(semanticArtifactIds).toEqual([canonicalArtifactId]); await expect( artifacts.getByDocumentVersion({ documentAssetId: asset.id, version: asset.version }), - ).resolves.toMatchObject({ id: canonicalArtifactId }); + ).resolves.toMatchObject({ artifactHash: "2".repeat(64), id: canonicalArtifactId }); await expect(artifacts.getById({ id: retryArtifactId })).resolves.toBeNull(); await expect( outlines.getByDocumentVersion({ documentAssetId: asset.id, version: asset.version }), @@ -1100,7 +1101,7 @@ describe("createDocumentCompilationWorker lease integration", () => { outlines: createInMemoryDocumentOutlineRepository({ maxOutlines: 1 }), parser: parser(), reindexer: { - canonicalizeArtifact: async (input) => input, + canonicalizeArtifact: async (input) => ({ artifact: input, disposition: "created" }), reindex: async (input) => { reindexCalls += 1; return { @@ -1367,15 +1368,16 @@ describe("createDocumentCompilationWorker lease integration", () => { expect.objectContaining({ metadata: expect.objectContaining({ contentKind: "document-multimodal-asset", - itemId: "018f0d60-7a49-7cc2-9c1b-5b36f18f6a02:1:figure-1", + itemId: expect.stringMatching(/^[0-9a-f-]{36}:1:figure-1$/u), modality: "image", objectKey: expect.stringMatching( - /^tenant-1\/spaces\/018f0d60-7a49-7cc2-9c1b-5b36f18f2c42\/documents\/018f0d60-7a49-7cc2-9c1b-5b36f18f6a01\/assets\/figure-1-[a-f0-9]{12}\.png$/u, + /^tenant-1\/spaces\/018f0d60-7a49-7cc2-9c1b-5b36f18f2c42\/documents\/018f0d60-7a49-7cc2-9c1b-5b36f18f6a01\/assets\/[a-f0-9-]+\/figure-1-[a-f0-9]{12}\.png$/u, ), }), targetId: asset.id, - virtualPath: - "/knowledge/docs/Worker.md--018f0d60/assets/image-Worker-diagram--018f0d60.json", + virtualPath: expect.stringMatching( + /^\/knowledge\/docs\/Worker\.md--018f0d60\/assets\/image-Worker-diagram--[a-f0-9]{8}\.json$/u, + ), }), ]); await expect( @@ -1636,7 +1638,7 @@ describe("createDocumentCompilationWorker lease integration", () => { filename: "Paper.pdf", id: "018f0d60-7a49-7cc2-9c1b-5b36f18f6b01", knowledgeSpaceId: "018f0d60-7a49-7cc2-9c1b-5b36f18f2c42", - mimeType: "application/pdf", + mimeType: "application/pdf; charset=binary", objectKey: "tenant-1/spaces/space/documents/asset/Paper.pdf", sha256: "c".repeat(64), sizeBytes: 12, @@ -1662,6 +1664,11 @@ describe("createDocumentCompilationWorker lease integration", () => { maxListLimit: 10, maxPaths: 10, }); + const multimodalManifests = createInMemoryDocumentMultimodalManifestRepository({ + maxManifests: 4, + }); + const parserHints: (ParserRouteHints | undefined)[] = []; + const basePdfParser = pdfParser(); const worker = createDocumentCompilationWorker({ assets, generateKnowledgePathId: sequenceIds([ @@ -1671,13 +1678,18 @@ describe("createDocumentCompilationWorker lease integration", () => { "018f0d60-7a49-7cc2-9c1b-5b36f18f6b08", "018f0d60-7a49-7cc2-9c1b-5b36f18f6b09", ]), + generateMultimodalWriteOwnerId: () => "pdf-write-owner", jobs: compilationJobs, knowledgePaths, - multimodalManifests: createInMemoryDocumentMultimodalManifestRepository({ - maxManifests: 4, - }), + multimodalManifests, objectStorage: adapter.objectStorage, - parser: pdfParser(), + parser: { + ...basePdfParser, + parse: async (input) => { + parserHints.push(input.parserHints); + return basePdfParser.parse(input); + }, + }, pdfRasterizer: { render: async (input) => { expect(input).toMatchObject({ @@ -1722,6 +1734,20 @@ describe("createDocumentCompilationWorker lease integration", () => { version: asset.version, }), ).resolves.toMatchObject({ stage: "published" }); + expect(parserHints).toEqual([ + expect.objectContaining({ imagesHandledExternally: true, requiresImages: true }), + ]); + const manifest = await multimodalManifests.getByDocumentVersion({ + documentAssetId: asset.id, + version: asset.version, + }); + expect(manifest).toMatchObject({ + items: [ + expect.objectContaining({ id: expect.stringContaining(String(manifest?.parseArtifactId)) }), + ], + parseArtifactId: expect.stringMatching(/^[a-f0-9-]{36}$/u), + }); + expect(manifest?.artifactHash).not.toBe("d".repeat(64)); const assetPaths = await knowledgePaths.listPhysicalDescendants({ knowledgeSpaceId: asset.knowledgeSpaceId, @@ -1734,7 +1760,7 @@ describe("createDocumentCompilationWorker lease integration", () => { metadata: expect.objectContaining({ contentKind: "document-multimodal-asset", modality: "image", - objectKey: expect.stringMatching(/figure-1-[a-f0-9]{12}\.png$/u), + objectKey: expect.stringMatching(/assets\/pdf-write-owner\/figure-1-[a-f0-9]{12}\.png$/u), }), }), ]); @@ -1742,6 +1768,549 @@ describe("createDocumentCompilationWorker lease integration", () => { adapter.objectStorage.getObject(String(assetPaths.items[0]?.metadata.objectKey)), ).resolves.toEqual(new Uint8Array([9, 8, 7, 6])); }); + + it("reuses canonical multimodal objects when the same PDF materialization is retried", async () => { + const adapter = createTestPlatformAdapter(); + const assets = createInMemoryDocumentAssetRepository({ maxAssets: 1 }); + const asset = await assets.create({ + filename: "Stable.pdf", + id: "018f0d60-7a49-7cc2-9c1b-5b36f18f8b01", + knowledgeSpaceId: "018f0d60-7a49-7cc2-9c1b-5b36f18f2c42", + mimeType: "application/pdf", + objectKey: "tenant-1/spaces/space/documents/asset/Stable.pdf", + sha256: "b".repeat(64), + sizeBytes: 8, + }); + await adapter.objectStorage.putObject({ + body: new TextEncoder().encode("%PDF-1.7"), + contentType: asset.mimeType, + key: asset.objectKey, + metadata: {}, + }); + const compilationJobs = createDocumentCompilationJobStateMachine({ + generateId: sequenceIds(["stable-job-1", "stable-job-2"]), + jobs: adapter.jobs, + repository: createInMemoryDocumentCompilationJobRepository({ maxJobs: 2 }), + }); + const artifacts = createInMemoryParseArtifactRepository({ maxArtifacts: 1 }); + const multimodalManifests = createInMemoryDocumentMultimodalManifestRepository({ + maxManifests: 1, + }); + let simulateCommitAcknowledgementFailure = true; + const worker = createDocumentCompilationWorker({ + assets, + failureManagement: "caller", + generateMultimodalWriteOwnerId: sequenceIds(["owner-a", "owner-b"]), + jobs: compilationJobs, + multimodalManifests, + objectStorage: adapter.objectStorage, + parser: pdfParser(), + pdfRasterizer: { + render: async () => ({ + body: new Uint8Array([9, 8, 7, 6]), + contentType: "image/png", + }), + }, + reindexer: { + canonicalizeArtifact: async (input) => { + const materialized = await artifacts.materialize(input); + if (simulateCommitAcknowledgementFailure) { + simulateCommitAcknowledgementFailure = false; + throw Object.assign(new Error("artifact commit acknowledgement was lost"), { + retryable: true, + }); + } + return materialized; + }, + getCanonicalArtifact: (input) => artifacts.getByDocumentVersion(input), + reindex: async (input) => ({ + artifact: input.parseArtifact, + nodesCreated: 1, + projectionsCreated: 1, + status: "rebuilt", + }), + }, + }); + const process = async () => { + const job = await compilationJobs.start({ + documentAssetId: asset.id, + knowledgeSpaceId: asset.knowledgeSpaceId, + tenantId: "tenant-1", + version: asset.version, + }); + return worker.process({ + documentAssetId: asset.id, + documentCompilationJobId: job.id, + knowledgeSpaceId: asset.knowledgeSpaceId, + tenantId: "tenant-1", + version: asset.version, + }); + }; + + await expect(process()).resolves.toMatchObject({ stage: "published" }); + await expect(process()).resolves.toMatchObject({ stage: "published" }); + + const stored = await adapter.objectStorage.listObjects({ limit: 20, prefix: "tenant-1/" }); + const multimodalObjects = stored.objects.filter(({ key }) => key.includes("/assets/")); + expect(multimodalObjects).toHaveLength(1); + expect(multimodalObjects[0]?.key).toContain("/assets/owner-a/"); + const manifest = await multimodalManifests.getByDocumentVersion({ + documentAssetId: asset.id, + version: asset.version, + }); + expect(manifest?.items[0]?.assetRef?.objectKey).toBe(multimodalObjects[0]?.key); + }); + + it.each([ + [ + "a renderer error", + async () => { + throw new Error("pdftoppm is unavailable"); + }, + ], + ["an unresolved renderer result", async () => null], + ])("falls back to provider image payloads after %s", async (_scenario, render) => { + const adapter = createTestPlatformAdapter(); + const assets = createInMemoryDocumentAssetRepository({ maxAssets: 1 }); + const asset = await assets.create({ + filename: "Fallback.pdf", + id: "018f0d60-7a49-7cc2-9c1b-5b36f18f7b01", + knowledgeSpaceId: "018f0d60-7a49-7cc2-9c1b-5b36f18f2c42", + mimeType: "application/pdf; charset=binary", + objectKey: "tenant-1/spaces/space/documents/asset/Fallback.pdf", + sha256: "e".repeat(64), + sizeBytes: 12, + }); + await adapter.objectStorage.putObject({ + body: new TextEncoder().encode("%PDF-1.7"), + contentType: asset.mimeType, + key: asset.objectKey, + metadata: {}, + }); + const compilationJobs = createDocumentCompilationJobStateMachine({ + generateId: () => "document-compilation-job-pdf-fallback-1", + jobs: adapter.jobs, + repository: createInMemoryDocumentCompilationJobRepository({ maxJobs: 1 }), + }); + const compilationJob = await compilationJobs.start({ + documentAssetId: asset.id, + knowledgeSpaceId: asset.knowledgeSpaceId, + tenantId: "tenant-1", + version: asset.version, + }); + const multimodalManifests = createInMemoryDocumentMultimodalManifestRepository({ + maxManifests: 1, + }); + const parserHints: (ParserRouteHints | undefined)[] = []; + const basePdfParser = pdfParser(); + const worker = createDocumentCompilationWorker({ + assets, + jobs: compilationJobs, + multimodalManifests, + objectStorage: adapter.objectStorage, + parser: { + ...basePdfParser, + parse: async (input) => { + parserHints.push(input.parserHints); + const parsed = await basePdfParser.parse(input); + + if (input.parserHints?.imagesHandledExternally) { + return parsed; + } + + return ParseArtifactSchema.parse({ + ...parsed, + elements: [ + ...parsed.elements.map((element) => + element.id === "figure-1" + ? { + ...element, + metadata: { + ...element.metadata, + assetRef: { + contentType: "image/png", + uri: "data:image/png;base64,AQIDBA==", + }, + }, + } + : element, + ), + { + id: "table-1", + metadata: { + assetRef: { + contentType: "image/png", + uri: "data:image/png;base64,BQYHCA==", + }, + table: { html: "
42
" }, + }, + pageNumber: 2, + sectionPath: ["Paper"], + text: "42", + type: "table", + }, + ], + }); + }, + }, + pdfRasterizer: { + render, + }, + reindexer: { + reindex: async (input) => ({ + artifact: input.parseArtifact, + nodesCreated: 1, + projectionsCreated: 1, + status: "rebuilt", + }), + }, + }); + + await expect( + worker.process({ + documentAssetId: asset.id, + documentCompilationJobId: compilationJob.id, + knowledgeSpaceId: asset.knowledgeSpaceId, + tenantId: "tenant-1", + version: asset.version, + }), + ).resolves.toMatchObject({ stage: "published" }); + expect(parserHints).toEqual([ + expect.objectContaining({ imagesHandledExternally: true, requiresImages: true }), + expect.objectContaining({ imagesHandledExternally: false, requiresImages: true }), + ]); + const manifest = await multimodalManifests.getByDocumentVersion({ + documentAssetId: asset.id, + version: asset.version, + }); + expect(manifest?.items).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + assetRef: expect.objectContaining({ + contentType: "image/png", + objectKey: expect.stringMatching(/figure-1-[a-f0-9]{12}\.png$/u), + }), + enrichment: expect.objectContaining({ asset: "provided" }), + parseElementId: "figure-1", + }), + expect.objectContaining({ + assetRef: expect.objectContaining({ + contentType: "image/png", + objectKey: expect.stringMatching(/table-1-[a-f0-9]{12}\.png$/u), + }), + enrichment: expect.objectContaining({ asset: "provided" }), + parseElementId: "table-1", + }), + ]), + ); + const imageObjectKey = manifest?.items.find((item) => item.parseElementId === "figure-1") + ?.assetRef?.objectKey; + const tableObjectKey = manifest?.items.find((item) => item.parseElementId === "table-1") + ?.assetRef?.objectKey; + expect(imageObjectKey).toBeDefined(); + expect(tableObjectKey).toBeDefined(); + await expect(adapter.objectStorage.getObject(String(imageObjectKey))).resolves.toEqual( + new Uint8Array([1, 2, 3, 4]), + ); + await expect(adapter.objectStorage.getObject(String(tableObjectKey))).resolves.toEqual( + new Uint8Array([5, 6, 7, 8]), + ); + }); + + it("compensates an execution-owned asset when object storage commits and then throws", async () => { + const adapter = createTestPlatformAdapter(); + const assets = createInMemoryDocumentAssetRepository({ maxAssets: 1 }); + const asset = await assets.create({ + filename: "Ambiguous.md", + id: "018f0d60-7a49-7cc2-9c1b-5b36f18f7c01", + knowledgeSpaceId: "018f0d60-7a49-7cc2-9c1b-5b36f18f2c42", + mimeType: "text/markdown", + objectKey: "tenant-1/spaces/space/documents/asset/Ambiguous.md", + sha256: "f".repeat(64), + sizeBytes: 12, + }); + await adapter.objectStorage.putObject({ + body: new TextEncoder().encode("# Ambiguous"), + contentType: asset.mimeType, + key: asset.objectKey, + metadata: {}, + }); + const compilationJobs = createDocumentCompilationJobStateMachine({ + generateId: () => "document-compilation-job-object-ambiguity-1", + jobs: adapter.jobs, + repository: createInMemoryDocumentCompilationJobRepository({ maxJobs: 1 }), + }); + const compilationJob = await compilationJobs.start({ + documentAssetId: asset.id, + knowledgeSpaceId: asset.knowledgeSpaceId, + tenantId: "tenant-1", + version: asset.version, + }); + const ambiguousStorage: PlatformAdapter["objectStorage"] = { + ...adapter.objectStorage, + putObject: async (input) => { + const result = await adapter.objectStorage.putObject(input); + if (input.key.includes("/assets/failed-write-owner/")) { + throw new Error("object committed before transport failure"); + } + return result; + }, + }; + const worker = createDocumentCompilationWorker({ + assets, + failureManagement: "caller", + generateMultimodalWriteOwnerId: () => "failed-write-owner", + jobs: compilationJobs, + multimodalManifests: createInMemoryDocumentMultimodalManifestRepository({ maxManifests: 1 }), + objectStorage: ambiguousStorage, + parser: parser(), + reindexer: { + reindex: async (input) => ({ + artifact: input.parseArtifact, + nodesCreated: 1, + projectionsCreated: 1, + status: "rebuilt", + }), + }, + }); + + await expect( + worker.process({ + documentAssetId: asset.id, + documentCompilationJobId: compilationJob.id, + knowledgeSpaceId: asset.knowledgeSpaceId, + tenantId: "tenant-1", + version: asset.version, + }), + ).rejects.toThrow("object committed before transport failure"); + await expect( + adapter.objectStorage.listObjects({ + limit: 10, + prefix: `tenant-1/spaces/${asset.knowledgeSpaceId}/documents/${asset.id}/assets/failed-write-owner/`, + }), + ).resolves.toMatchObject({ objects: [] }); + }); + + it("keeps incomplete multimodal compensation retryable", async () => { + const adapter = createTestPlatformAdapter(); + const assets = createInMemoryDocumentAssetRepository({ maxAssets: 1 }); + const asset = await assets.create({ + filename: "Retry-cleanup.md", + id: "018f0d60-7a49-7cc2-9c1b-5b36f18f7c02", + knowledgeSpaceId: "018f0d60-7a49-7cc2-9c1b-5b36f18f2c42", + mimeType: "text/markdown", + objectKey: "tenant-1/spaces/space/documents/asset/Retry-cleanup.md", + sha256: "e".repeat(64), + sizeBytes: 15, + }); + await adapter.objectStorage.putObject({ + body: new TextEncoder().encode("# Retry cleanup"), + contentType: asset.mimeType, + key: asset.objectKey, + metadata: {}, + }); + const compilationJobs = createDocumentCompilationJobStateMachine({ + generateId: () => "document-compilation-job-cleanup-retry-1", + jobs: adapter.jobs, + repository: createInMemoryDocumentCompilationJobRepository({ maxJobs: 1 }), + }); + const compilationJob = await compilationJobs.start({ + documentAssetId: asset.id, + knowledgeSpaceId: asset.knowledgeSpaceId, + tenantId: "tenant-1", + version: asset.version, + }); + const unavailableStorage: PlatformAdapter["objectStorage"] = { + ...adapter.objectStorage, + deleteObject: async (key) => { + if (key.includes("/assets/cleanup-retry-owner/")) { + throw Object.assign(new Error("cleanup storage unavailable"), { retryable: true }); + } + return adapter.objectStorage.deleteObject(key); + }, + putObject: async (input) => { + const result = await adapter.objectStorage.putObject(input); + if (input.key.includes("/assets/cleanup-retry-owner/")) { + throw Object.assign(new Error("object write acknowledgement was lost"), { + retryable: true, + }); + } + return result; + }, + }; + const worker = createDocumentCompilationWorker({ + assets, + failureManagement: "caller", + generateMultimodalWriteOwnerId: () => "cleanup-retry-owner", + jobs: compilationJobs, + multimodalManifests: createInMemoryDocumentMultimodalManifestRepository({ maxManifests: 1 }), + objectStorage: unavailableStorage, + parser: parser(), + reindexer: { + reindex: async (input) => ({ + artifact: input.parseArtifact, + nodesCreated: 1, + projectionsCreated: 1, + status: "rebuilt", + }), + }, + }); + + const error = await worker + .process({ + documentAssetId: asset.id, + documentCompilationJobId: compilationJob.id, + knowledgeSpaceId: asset.knowledgeSpaceId, + tenantId: "tenant-1", + version: asset.version, + }) + .catch((caught: unknown) => caught); + + expect(error).toBeInstanceOf(AggregateError); + expect(error).toMatchObject({ + code: "DOCUMENT_COMPILATION_RETRYABLE", + retryable: true, + }); + await expect( + adapter.objectStorage.listObjects({ + limit: 10, + prefix: `tenant-1/spaces/${asset.knowledgeSpaceId}/documents/${asset.id}/assets/cleanup-retry-owner/`, + }), + ).resolves.toMatchObject({ objects: [expect.any(Object)] }); + }); + + it("limits PDF provider materialization even when the local rasterizer is disabled", async () => { + const adapter = createTestPlatformAdapter(); + const assets = createInMemoryDocumentAssetRepository({ maxAssets: 3 }); + const assetIds = [ + "018f0d60-7a49-7cc2-9c1b-5b36f18f7d01", + "018f0d60-7a49-7cc2-9c1b-5b36f18f7d02", + "018f0d60-7a49-7cc2-9c1b-5b36f18f7d03", + ] as const; + const artifactIds = [ + "018f0d60-7a49-7cc2-9c1b-5b36f18f7e01", + "018f0d60-7a49-7cc2-9c1b-5b36f18f7e02", + "018f0d60-7a49-7cc2-9c1b-5b36f18f7e03", + ] as const; + const createdAssets = []; + for (const [index, id] of assetIds.entries()) { + const asset = await assets.create({ + filename: `Queued-${index + 1}.pdf`, + id, + knowledgeSpaceId: "018f0d60-7a49-7cc2-9c1b-5b36f18f2c42", + mimeType: "application/pdf", + objectKey: `tenant-1/spaces/space/documents/${id}/source.pdf`, + sha256: String(index + 1).repeat(64), + sizeBytes: 8, + }); + createdAssets.push(asset); + await adapter.objectStorage.putObject({ + body: new TextEncoder().encode("%PDF-1.7"), + contentType: asset.mimeType, + key: asset.objectKey, + metadata: {}, + }); + } + const compilationJobs = createDocumentCompilationJobStateMachine({ + generateId: sequenceIds(["pdf-provider-job-1", "pdf-provider-job-2", "pdf-provider-job-3"]), + jobs: adapter.jobs, + repository: createInMemoryDocumentCompilationJobRepository({ maxJobs: 3 }), + }); + const jobs: Awaited>[] = []; + for (const asset of createdAssets) { + jobs.push( + await compilationJobs.start({ + documentAssetId: asset.id, + knowledgeSpaceId: asset.knowledgeSpaceId, + tenantId: "tenant-1", + version: asset.version, + }), + ); + } + let activeParses = 0; + let enteredParses = 0; + let maxActiveParses = 0; + let releaseParses!: () => void; + let resolveTwoEntered!: () => void; + const parseBarrier = new Promise((resolve) => { + releaseParses = resolve; + }); + const twoEntered = new Promise((resolve) => { + resolveTwoEntered = resolve; + }); + const multimodalMaterializationGate = createConcurrencyGate(2); + const multimodalManifests = createInMemoryDocumentMultimodalManifestRepository({ + maxManifests: 3, + }); + const createWorker = () => + createDocumentCompilationWorker({ + assets, + jobs: compilationJobs, + multimodalImageVariantGenerator: { generate: async () => [] }, + multimodalManifests, + multimodalMaterializationGate, + objectStorage: adapter.objectStorage, + parser: { + kind: "unstructured", + parse: async (input) => { + expect(input.parserHints).toMatchObject({ + imagesHandledExternally: false, + requiresImages: true, + }); + activeParses += 1; + enteredParses += 1; + maxActiveParses = Math.max(maxActiveParses, activeParses); + if (enteredParses === 2) resolveTwoEntered(); + await parseBarrier; + activeParses -= 1; + const artifactId = + artifactIds[assetIds.indexOf(input.documentAssetId as (typeof assetIds)[number])]; + if (!artifactId) throw new Error("Missing test parse artifact id"); + return ParseArtifactSchema.parse({ + artifactHash: input.documentAssetId.replaceAll("-", "").padEnd(64, "0").slice(0, 64), + contentType: "text", + createdAt: "2026-08-18T12:00:00.000Z", + documentAssetId: input.documentAssetId, + elements: [], + id: artifactId, + metadata: {}, + parser: "unstructured", + version: input.version, + }); + }, + }, + reindexer: { + reindex: async (input) => ({ + artifact: input.parseArtifact, + nodesCreated: 0, + projectionIds: [], + projectionsCreated: 0, + status: "rebuilt", + }), + }, + }); + + const processes = createdAssets.map((asset, index) => + createWorker().process({ + documentAssetId: asset.id, + documentCompilationJobId: jobs[index]?.id ?? "missing-job", + knowledgeSpaceId: asset.knowledgeSpaceId, + tenantId: "tenant-1", + version: asset.version, + }), + ); + await twoEntered; + await Promise.resolve(); + expect(enteredParses).toBe(2); + expect(maxActiveParses).toBe(2); + releaseParses(); + await expect(Promise.all(processes)).resolves.toEqual([ + expect.objectContaining({ stage: "published" }), + expect.objectContaining({ stage: "published" }), + expect.objectContaining({ stage: "published" }), + ]); + expect(maxActiveParses).toBe(2); + }); }); function createTestPlatformAdapter(): PlatformAdapter { diff --git a/knowledge-fs/packages/api/src/document-compilation-worker.ts b/knowledge-fs/packages/api/src/document-compilation-worker.ts index 7998c9653ae..4824578ca10 100644 --- a/knowledge-fs/packages/api/src/document-compilation-worker.ts +++ b/knowledge-fs/packages/api/src/document-compilation-worker.ts @@ -1,3 +1,5 @@ +import { createHash, randomUUID } from "node:crypto"; + import { z } from "@hono/zod-openapi"; import type { ChunkConfig } from "@knowledge/compute"; import { @@ -9,12 +11,18 @@ import { type KnowledgeSpaceEmbeddingProfile, type KnowledgeSpaceRetrievalProfile, type ParseArtifact, + ParseArtifactSchema, type PlatformAdapter, PublicationGenerationIdSchema, TenantIdSchema, } from "@knowledge/core"; import type { ParserAdapter, ParserRouteHints } from "@knowledge/parsers"; +import { + type ConcurrencyGate, + createConcurrencyGate, + mapWithConcurrency, +} from "./bounded-concurrency"; import { DeletionLifecycleFenceActiveError, type DeletionLifecycleFenceGuard, @@ -40,6 +48,7 @@ import { buildDocumentSectionKnowledgePaths, } from "./document-knowledge-paths"; import type { DocumentModelBudget } from "./document-model-budget"; +import { finalizeDocumentMultimodalArtifact } from "./document-multimodal-artifact"; import { extractDocumentMultimodalAssets } from "./document-multimodal-asset-extractor"; import { createDocumentMultimodalManifestBuilder } from "./document-multimodal-manifest-builder"; import type { DocumentMultimodalManifestRepository } from "./document-multimodal-manifest-repository"; @@ -48,6 +57,7 @@ import type { DocumentOutlineRepository } from "./document-outline-repository"; import type { DocumentOutlineSummaryEnhancer } from "./document-outline-summary-enhancer"; import { type DocumentPdfRasterizer, + DocumentPdfRenderError, rasterizeDocumentPdfMultimodalAssets, } from "./document-pdf-rasterizer"; import type { JointSemanticGraphMaterializer } from "./document-semantic-enrichment-processor"; @@ -90,12 +100,15 @@ export interface DocumentCompilationWorkerOptions { */ readonly failureManagement?: "caller" | "worker" | undefined; readonly generateKnowledgePathId?: (() => string) | undefined; + readonly generateMultimodalWriteOwnerId?: (() => string) | undefined; readonly jobs: DocumentCompilationJobStateMachine; readonly jointSemanticGraph?: JointSemanticGraphMaterializer | undefined; readonly knowledgePaths?: KnowledgePathRepository | undefined; readonly multimodalImageVariantGenerator?: DocumentImageVariantGenerator | undefined; + readonly multimodalMaterializationGate?: ConcurrencyGate | undefined; readonly multimodalLocalAssetAllowlist?: readonly string[] | undefined; readonly multimodalMaxExtractedAssets?: number | undefined; + readonly multimodalMaxConcurrency?: number | undefined; readonly multimodalMaxLocalAssetBytes?: number | undefined; readonly multimodalMaxPdfRasterizedAssets?: number | undefined; readonly multimodalManifests: DocumentMultimodalManifestRepository; @@ -223,12 +236,15 @@ export function createDocumentCompilationWorker({ indexOverrides, failureManagement = "worker", generateKnowledgePathId, + generateMultimodalWriteOwnerId = randomUUID, jobs, jointSemanticGraph, knowledgePaths, multimodalImageVariantGenerator, + multimodalMaterializationGate, multimodalLocalAssetAllowlist, multimodalMaxExtractedAssets, + multimodalMaxConcurrency = 2, multimodalMaxLocalAssetBytes, multimodalMaxPdfRasterizedAssets, multimodalManifests, @@ -248,6 +264,8 @@ export function createDocumentCompilationWorker({ smokeEvaluation, visualEmbeddingModel, }: DocumentCompilationWorkerOptions): DocumentCompilationWorker { + const effectiveMultimodalMaterializationGate = + multimodalMaterializationGate ?? createConcurrencyGate(multimodalMaxConcurrency); const stagedProjectionPublication = reindexer.publishProjections && reindexer.failProjections ? { @@ -271,6 +289,7 @@ export function createDocumentCompilationWorker({ signal?.throwIfAborted(); }; let cleanupStaleObjectWrites = async (): Promise => undefined; + let multimodalWritesDurable = false; try { await assertDocumentAvailable?.({ @@ -307,6 +326,7 @@ export function createDocumentCompilationWorker({ } const activeAsset = asset; + const multimodalWriteOwnerId = generateMultimodalWriteOwnerId(); const deletionToken = await deletionFence?.captureDeletionFence({ documentAssetId: activeAsset.id, knowledgeSpaceId: input.knowledgeSpaceId, @@ -320,18 +340,15 @@ export function createDocumentCompilationWorker({ } signal?.throwIfAborted(); }; - const multimodalObjectStorage = - deletionToken || objectWriteAdmission - ? createDeletionFencedCompilationObjectStorage({ - assertWritable, - objectWriteAdmission, - objectStorage, - onCleanupReady: (cleanup) => { - cleanupStaleObjectWrites = cleanup; - }, - scope: { knowledgeSpaceId: input.knowledgeSpaceId, tenantId: input.tenantId }, - }) - : objectStorage; + const multimodalObjectStorage = createDeletionFencedCompilationObjectStorage({ + assertWritable, + objectWriteAdmission, + objectStorage, + onCleanupReady: (cleanup) => { + cleanupStaleObjectWrites = cleanup; + }, + scope: { knowledgeSpaceId: input.knowledgeSpaceId, tenantId: input.tenantId }, + }); const compile = async () => { const initialJob = await jobs.get(input.documentCompilationJobId); if (!initialJob) { @@ -361,62 +378,146 @@ export function createDocumentCompilationWorker({ } canonicalArtifact = persistedArtifact; } else { - const body = await objectStorage.getObject(activeAsset.objectKey); + const requiresImages = Boolean( + visualEmbeddingModel || multimodalImageVariantGenerator || pdfRasterizer, + ); + const materializeSource = async (): Promise => { + signal?.throwIfAborted(); + const body = await objectStorage.getObject(activeAsset.objectKey); - if (!body) { - throw new Error("Document compilation object not found"); - } + if (!body) { + throw new Error("Document compilation object not found"); + } - const parsedArtifact = await parser.parse({ - body, - documentAssetId: activeAsset.id, - filename: activeAsset.filename, - mimeType: activeAsset.mimeType, - parserHints: documentParserHints({ - assetMetadata: activeAsset.metadata, - requiresImages: Boolean( - visualEmbeddingModel || multimodalImageVariantGenerator || pdfRasterizer, - ), - }), - ...(signal ? { signal } : {}), - version: activeAsset.version, - }); - await assertWritable(); - const rasterized = await rasterizeDocumentPdfMultimodalAssets({ - artifact: parsedArtifact, - documentBody: body, - documentMimeType: activeAsset.mimeType, - knowledgeSpaceId: input.knowledgeSpaceId, - ...(multimodalMaxPdfRasterizedAssets - ? { maxRasterizedAssets: multimodalMaxPdfRasterizedAssets } - : {}), - objectStorage: multimodalObjectStorage, - ...(pdfRasterizer ? { rasterizer: pdfRasterizer } : {}), - tenantId: input.tenantId, - }); - await assertWritable(); - const { artifact } = await extractDocumentMultimodalAssets({ - ...(multimodalLocalAssetAllowlist - ? { allowLocalAssetPaths: multimodalLocalAssetAllowlist } - : {}), - artifact: rasterized.artifact, - knowledgeSpaceId: input.knowledgeSpaceId, - ...(multimodalMaxExtractedAssets - ? { maxExtractedAssets: multimodalMaxExtractedAssets } - : {}), - ...(multimodalMaxLocalAssetBytes - ? { maxLocalAssetBytes: multimodalMaxLocalAssetBytes } - : {}), - ...(multimodalImageVariantGenerator - ? { imageVariantGenerator: multimodalImageVariantGenerator } - : {}), - objectStorage: multimodalObjectStorage, - tenantId: input.tenantId, - }); - await assertWritable(); - canonicalArtifact = reindexer.canonicalizeArtifact - ? await reindexer.canonicalizeArtifact(artifact) - : artifact; + const parseDocument = (imagesHandledExternally: boolean) => + parser.parse({ + body, + documentAssetId: activeAsset.id, + filename: activeAsset.filename, + mimeType: activeAsset.mimeType, + parserHints: documentParserHints({ + assetMetadata: activeAsset.metadata, + imagesHandledExternally, + requiresImages, + }), + ...(signal ? { signal } : {}), + version: activeAsset.version, + }); + const materializationArtifactId = await resolveMaterializationArtifactId({ + documentAssetId: activeAsset.id, + reindexer, + version: activeAsset.version, + }); + const bindMaterializationIdentity = (artifact: ParseArtifact) => + bindParseArtifactIdentity(artifact, materializationArtifactId); + const parsedArtifact = bindMaterializationIdentity( + await parseDocument(Boolean(pdfRasterizer) && isPdfDocument(activeAsset.mimeType)), + ); + await assertWritable(); + let multimodalArtifact: ParseArtifact; + + try { + const rasterized = await rasterizeDocumentPdfMultimodalAssets({ + artifact: parsedArtifact, + documentBody: body, + documentMimeType: activeAsset.mimeType, + knowledgeSpaceId: input.knowledgeSpaceId, + ...(multimodalMaxPdfRasterizedAssets + ? { maxRasterizedAssets: multimodalMaxPdfRasterizedAssets } + : {}), + objectStorage: multimodalObjectStorage, + ...(pdfRasterizer ? { rasterizer: pdfRasterizer } : {}), + ...(signal ? { signal } : {}), + tenantId: input.tenantId, + writeOwnerId: multimodalWriteOwnerId, + }); + const providerFallbackRequired = + Boolean(pdfRasterizer) && + requiresImages && + isPdfDocument(activeAsset.mimeType) && + rasterized.rasterizedCount === 0 && + rasterized.unresolvedCount > 0; + multimodalArtifact = providerFallbackRequired + ? bindMaterializationIdentity(await parseDocument(false)) + : rasterized.artifact; + } catch (error) { + if ( + !(error instanceof DocumentPdfRenderError) || + !pdfRasterizer || + !requiresImages || + !isPdfDocument(activeAsset.mimeType) + ) { + throw error; + } + + multimodalArtifact = bindMaterializationIdentity(await parseDocument(false)); + } + await assertWritable(); + const { artifact } = await extractDocumentMultimodalAssets({ + ...(multimodalLocalAssetAllowlist + ? { allowLocalAssetPaths: multimodalLocalAssetAllowlist } + : {}), + artifact: multimodalArtifact, + knowledgeSpaceId: input.knowledgeSpaceId, + ...(multimodalMaxExtractedAssets + ? { maxExtractedAssets: multimodalMaxExtractedAssets } + : {}), + ...(multimodalMaxLocalAssetBytes + ? { maxLocalAssetBytes: multimodalMaxLocalAssetBytes } + : {}), + ...(multimodalImageVariantGenerator + ? { imageVariantGenerator: multimodalImageVariantGenerator } + : {}), + objectStorage: multimodalObjectStorage, + tenantId: input.tenantId, + writeOwnerId: multimodalWriteOwnerId, + }); + await assertWritable(); + const finalizedArtifact = finalizeDocumentMultimodalArtifact(artifact); + if (reindexer.canonicalizeArtifact) { + let materialized: Awaited< + ReturnType> + >; + try { + materialized = await reindexer.canonicalizeArtifact(finalizedArtifact); + } catch (error) { + if (!reindexer.getCanonicalArtifact) { + multimodalWritesDurable = true; + throw ambiguousArtifactMaterializationError(error); + } + let reconciled: ParseArtifact | null; + try { + reconciled = await reindexer.getCanonicalArtifact({ + documentAssetId: activeAsset.id, + version: activeAsset.version, + }); + } catch (reconciliationError) { + multimodalWritesDurable = true; + throw ambiguousArtifactMaterializationError(error, reconciliationError); + } + if (reconciled?.artifactHash !== finalizedArtifact.artifactHash) { + throw error; + } + if (sameArtifactObjectReferences(reconciled, finalizedArtifact)) { + multimodalWritesDurable = true; + } else { + await cleanupStaleObjectWrites(); + } + return reconciled; + } + if (materialized.disposition === "unchanged") { + await cleanupStaleObjectWrites(); + } else { + multimodalWritesDurable = true; + } + return materialized.artifact; + } + return finalizedArtifact; + }; + canonicalArtifact = + requiresImages && isPdfDocument(activeAsset.mimeType) + ? await effectiveMultimodalMaterializationGate.run(materializeSource) + : await materializeSource(); } const documentIndexOverrides = indexOverrides ? await indexOverrides.resolve({ @@ -523,6 +624,9 @@ export function createDocumentCompilationWorker({ } await assertWritable(); persistedManifest = await multimodalManifests.upsert(multimodalManifest); + if (!reindexer.canonicalizeArtifact) { + multimodalWritesDurable = true; + } if (!deferOutlineUntilSemanticNodes) { await assertWritable(); await jobs.advance(input.documentCompilationJobId, "outline_built"); @@ -827,8 +931,15 @@ export function createDocumentCompilationWorker({ } } } - if (isDeletionWriteBlocked(effectiveError)) { - await cleanupStaleObjectWrites(); + const deletionWriteBlocked = isDeletionWriteBlocked(effectiveError); + if (deletionWriteBlocked || !multimodalWritesDurable) { + try { + await cleanupStaleObjectWrites(); + } catch (cleanupError) { + effectiveError = retryableMultimodalCleanupError(effectiveError, cleanupError); + } + } + if (deletionWriteBlocked) { throw effectiveError; } if (legacyStagedProjectionPublication && stagedProjectionIds.length > 0) { @@ -863,8 +974,13 @@ export function createDocumentCompilationWorker({ }; } +function isPdfDocument(mimeType: string): boolean { + return mimeType.split(";", 1)[0]?.trim().toLowerCase() === "application/pdf"; +} + function documentParserHints(input: { readonly assetMetadata: Readonly>; + readonly imagesHandledExternally: boolean; readonly requiresImages: boolean; }): ParserRouteHints { const language = @@ -877,6 +993,7 @@ function documentParserHints(input: { ? input.assetMetadata.layoutComplexity : undefined; return { + imagesHandledExternally: input.imagesHandledExternally, ...(language ? { language } : {}), ...(layoutComplexity ? { layoutComplexity } : {}), requiresImages: input.requiresImages, @@ -885,6 +1002,107 @@ function documentParserHints(input: { }; } +async function resolveMaterializationArtifactId({ + documentAssetId, + reindexer, + version, +}: { + readonly documentAssetId: string; + readonly reindexer: IncrementalReindexer; + readonly version: number; +}): Promise { + const existing = await reindexer.getCanonicalArtifact?.({ documentAssetId, version }); + + return existing?.id ?? deterministicParseArtifactId(documentAssetId, version); +} + +function deterministicParseArtifactId(documentAssetId: string, version: number): string { + const bytes = createHash("sha256") + .update(`knowledge-fs:parse-artifact:${documentAssetId}:${version}`) + .digest() + .subarray(0, 16); + bytes[6] = ((bytes[6] ?? 0) & 0x0f) | 0x50; + bytes[8] = ((bytes[8] ?? 0) & 0x3f) | 0x80; + const hex = bytes.toString("hex"); + + return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-${hex.slice(12, 16)}-${hex.slice(16, 20)}-${hex.slice(20)}`; +} + +function bindParseArtifactIdentity(artifact: ParseArtifact, artifactId: string): ParseArtifact { + const generatedElementIds = artifact.elements.every((element, index) => { + const match = element.id.match( + /^[0-9a-f]{8}-[0-9a-f]{4}-[1-8][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}:element-(\d+)$/iu, + ); + + return match?.[1] === String(index + 1); + }); + + return ParseArtifactSchema.parse({ + ...artifact, + elements: generatedElementIds + ? artifact.elements.map((element, index) => ({ + ...element, + id: `${artifactId}:element-${index + 1}`, + })) + : artifact.elements, + id: artifactId, + }); +} + +function sameArtifactObjectReferences(left: ParseArtifact, right: ParseArtifact): boolean { + return ( + JSON.stringify(artifactObjectReferences(left)) === + JSON.stringify(artifactObjectReferences(right)) + ); +} + +function artifactObjectReferences(artifact: ParseArtifact): readonly string[] { + const keys = new Set(); + const visit = (value: unknown): void => { + if (Array.isArray(value)) { + for (const item of value) visit(item); + return; + } + if (!value || typeof value !== "object") return; + for (const [key, nested] of Object.entries(value)) { + if (key === "objectKey" && typeof nested === "string") keys.add(nested); + else visit(nested); + } + }; + for (const element of artifact.elements) visit(element.metadata); + return [...keys].sort(); +} + +function ambiguousArtifactMaterializationError( + materializationError: unknown, + reconciliationError?: unknown, +): Error & { readonly code: string; readonly retryable: true } { + return Object.assign( + new AggregateError( + reconciliationError === undefined + ? [materializationError] + : [materializationError, reconciliationError], + "Parse artifact materialization outcome is ambiguous; execution-owned objects were retained", + { cause: materializationError }, + ), + { code: "DOCUMENT_COMPILATION_RETRYABLE", retryable: true as const }, + ); +} + +function retryableMultimodalCleanupError( + originalError: unknown, + cleanupError: unknown, +): Error & { readonly code: string; readonly retryable: true } { + return Object.assign( + new AggregateError( + [originalError, cleanupError], + "Document compilation failed and could not compensate multimodal object writes", + { cause: originalError }, + ), + { code: "DOCUMENT_COMPILATION_RETRYABLE", retryable: true as const }, + ); +} + function isDeletionWriteBlocked(error: unknown): boolean { return ( error instanceof DeletionLifecycleFenceActiveError || @@ -914,27 +1132,46 @@ function createDeletionFencedCompilationObjectStorage({ const createdKeys = new Set(); const cleanup = async (): Promise => { const failures: unknown[] = []; - for (const key of [...createdKeys]) { - const keyFailures: unknown[] = []; - for (let attempt = 1; attempt <= 3; attempt += 1) { - try { - await objectStorage.deleteObject(key); + const deadlineAt = Date.now() + 60_000; + let skippedForDeadline = 0; + await mapWithConcurrency([...createdKeys], 4, async (key) => { + if (Date.now() >= deadlineAt) { + skippedForDeadline += 1; + return; + } + try { + if ((await objectStorage.headObject(key)) === null) { createdKeys.delete(key); - break; - } catch (error) { - keyFailures.push(error); + return; + } + if (Date.now() >= deadlineAt) { + skippedForDeadline += 1; + return; + } + await objectStorage.deleteObject(key); + createdKeys.delete(key); + } catch (error) { + if (failures.length < 16) { + failures.push( + new AggregateError([error], `Failed to compensate late object write key=${key}`), + ); } } - if (createdKeys.has(key)) { - failures.push( - new AggregateError(keyFailures, `Failed to compensate late object write key=${key}`), - ); - } + }); + if (skippedForDeadline > 0) { + failures.push( + new Error( + `Document compilation object cleanup exceeded 60000ms with ${skippedForDeadline} key(s) unattempted`, + ), + ); } if (failures.length > 0) { - throw new AggregateError( - failures, - `Document compilation could not compensate ${failures.length} late object write(s)`, + throw Object.assign( + new AggregateError( + failures, + `Document compilation could not compensate ${failures.length} late object write(s)`, + ), + { code: "DOCUMENT_COMPILATION_RETRYABLE", retryable: true as const }, ); } }; @@ -946,7 +1183,11 @@ function createDeletionFencedCompilationObjectStorage({ close: () => objectStorage.close?.() ?? Promise.resolve(), } : {}), - deleteObject: (key) => objectStorage.deleteObject(key), + deleteObject: async (key) => { + const result = await objectStorage.deleteObject(key); + createdKeys.delete(key); + return result; + }, getObject: (key) => objectStorage.getObject(key), getObjectStream: (key) => objectStorage.getObjectStream(key), headObject: (key) => objectStorage.headObject(key), @@ -956,11 +1197,12 @@ function createDeletionFencedCompilationObjectStorage({ putObject: async (input) => { await assertWritable(); const existedBefore = (await objectStorage.headObject(input.key)) !== null; + // Register ownership before PUT so a committed write with a lost response is still cleaned. + if (!existedBefore) createdKeys.add(input.key); await assertWritable(); const stored = await withDeletionObjectWriteAdmission(objectWriteAdmission, scope, () => objectStorage.putObject(input), ); - if (!existedBefore) createdKeys.add(input.key); try { await assertWritable(); } catch (error) { diff --git a/knowledge-fs/packages/api/src/document-image-variant-generator.test.ts b/knowledge-fs/packages/api/src/document-image-variant-generator.test.ts index 385a7dd5b34..18a5eba16b3 100644 --- a/knowledge-fs/packages/api/src/document-image-variant-generator.test.ts +++ b/knowledge-fs/packages/api/src/document-image-variant-generator.test.ts @@ -42,8 +42,50 @@ describe("createSharpImageThumbnailVariantGenerator", () => { expect(() => createSharpImageThumbnailVariantGenerator({ maxDimension: 0 })).toThrow( "Sharp image thumbnail maxDimension must be at least 1", ); + expect(() => createSharpImageThumbnailVariantGenerator({ maxInputPixels: 0 })).toThrow( + "Sharp image thumbnail maxInputPixels must be at least 1", + ); + expect(() => createSharpImageThumbnailVariantGenerator({ maxOutputBytes: 0 })).toThrow( + "Sharp image thumbnail maxOutputBytes must be at least 1", + ); expect(() => createSharpImageThumbnailVariantGenerator({ variantName: "" })).toThrow( "Sharp image thumbnail variantName must be non-empty", ); }); + + it("rejects compressed images whose decoded dimensions exceed the pixel budget", async () => { + const generator = createSharpImageThumbnailVariantGenerator({ maxInputPixels: 100 }); + const oversizedSvg = new TextEncoder().encode( + '', + ); + + await expect( + generator.generate({ + body: oversizedSvg, + contentType: "image/svg+xml", + elementId: "figure-oversized", + }), + ).rejects.toThrow(/pixel limit/iu); + }); + + it("rejects thumbnail variants that exceed the encoded output budget", async () => { + const sharp = (await import("sharp")).default; + const generator = createSharpImageThumbnailVariantGenerator({ maxOutputBytes: 1 }); + const png = new Uint8Array( + await sharp({ + create: { + background: { alpha: 1, b: 255, g: 0, r: 0 }, + channels: 4, + height: 2, + width: 2, + }, + }) + .png() + .toBuffer(), + ); + + await expect( + generator.generate({ body: png, contentType: "image/png", elementId: "figure-output" }), + ).rejects.toThrow(/output exceeds maxOutputBytes/iu); + }); }); diff --git a/knowledge-fs/packages/api/src/document-image-variant-generator.ts b/knowledge-fs/packages/api/src/document-image-variant-generator.ts index bfd4ee7f6b4..0f02050516d 100644 --- a/knowledge-fs/packages/api/src/document-image-variant-generator.ts +++ b/knowledge-fs/packages/api/src/document-image-variant-generator.ts @@ -20,20 +20,34 @@ export interface DocumentImageVariantGenerator { export interface SharpImageThumbnailVariantGeneratorOptions { readonly maxDimension?: number | undefined; + readonly maxInputPixels?: number | undefined; + readonly maxOutputBytes?: number | undefined; readonly variantName?: string | undefined; } const defaultThumbnailMaxDimension = 320; +const defaultThumbnailMaxInputPixels = 20_000_000; +const defaultThumbnailMaxOutputBytes = 8 * 1024 * 1024; const defaultThumbnailVariantName = "thumbnail"; export function createSharpImageThumbnailVariantGenerator({ maxDimension = defaultThumbnailMaxDimension, + maxInputPixels = defaultThumbnailMaxInputPixels, + maxOutputBytes = defaultThumbnailMaxOutputBytes, variantName = defaultThumbnailVariantName, }: SharpImageThumbnailVariantGeneratorOptions = {}): DocumentImageVariantGenerator { if (!Number.isSafeInteger(maxDimension) || maxDimension < 1) { throw new Error("Sharp image thumbnail maxDimension must be at least 1"); } + if (!Number.isSafeInteger(maxInputPixels) || maxInputPixels < 1) { + throw new Error("Sharp image thumbnail maxInputPixels must be at least 1"); + } + + if (!Number.isSafeInteger(maxOutputBytes) || maxOutputBytes < 1) { + throw new Error("Sharp image thumbnail maxOutputBytes must be at least 1"); + } + if (!variantName.trim()) { throw new Error("Sharp image thumbnail variantName must be non-empty"); } @@ -45,7 +59,7 @@ export function createSharpImageThumbnailVariantGenerator({ } const sharp = (await import("sharp")).default; - const { data, info } = await sharp(body) + const { data, info } = await sharp(body, { limitInputPixels: maxInputPixels }) .rotate() .resize({ fit: "inside", @@ -56,6 +70,12 @@ export function createSharpImageThumbnailVariantGenerator({ .png() .toBuffer({ resolveWithObject: true }); + if (data.byteLength > maxOutputBytes) { + throw new Error( + `Sharp image thumbnail output exceeds maxOutputBytes (${data.byteLength} > ${maxOutputBytes})`, + ); + } + return [ { body: new Uint8Array(data), diff --git a/knowledge-fs/packages/api/src/document-multimodal-artifact.test.ts b/knowledge-fs/packages/api/src/document-multimodal-artifact.test.ts new file mode 100644 index 00000000000..1ed4a47bc66 --- /dev/null +++ b/knowledge-fs/packages/api/src/document-multimodal-artifact.test.ts @@ -0,0 +1,88 @@ +import { type ParseArtifact, ParseArtifactSchema } from "@knowledge/core"; +import { describe, expect, it } from "vitest"; + +import { finalizeDocumentMultimodalArtifact } from "./document-multimodal-artifact"; + +describe("finalizeDocumentMultimodalArtifact", () => { + it("changes lineage for visual bytes and renderer semantics but ignores owner object keys", () => { + const first = finalizeDocumentMultimodalArtifact(artifact({ owner: "owner-a" })); + const sameMaterial = finalizeDocumentMultimodalArtifact(artifact({ owner: "owner-b" })); + const changedBytes = finalizeDocumentMultimodalArtifact( + artifact({ owner: "owner-c", sha256: "c".repeat(64) }), + ); + const changedDpi = finalizeDocumentMultimodalArtifact(artifact({ dpi: 200, owner: "owner-d" })); + + expect(first.artifactHash).toBe(sameMaterial.artifactHash); + expect(changedBytes.artifactHash).not.toBe(first.artifactHash); + expect(changedDpi.artifactHash).not.toBe(first.artifactHash); + expect(first.metadata.multimodalMaterialization).toMatchObject({ + assetCount: 1, + contractVersion: 1, + digest: first.artifactHash, + sourceArtifactHash: "a".repeat(64), + }); + }); + + it("is idempotent and preserves parser lineage when no durable visual exists", () => { + const materialized = finalizeDocumentMultimodalArtifact(artifact({ owner: "owner-a" })); + + expect(finalizeDocumentMultimodalArtifact(materialized).artifactHash).toBe( + materialized.artifactHash, + ); + const withoutVisual = ParseArtifactSchema.parse({ + ...artifact({ owner: "owner-a" }), + elements: [], + }); + expect(finalizeDocumentMultimodalArtifact(withoutVisual)).toBe(withoutVisual); + }); +}); + +function artifact({ + dpi = 144, + owner, + sha256 = "b".repeat(64), +}: { + readonly dpi?: number; + readonly owner: string; + readonly sha256?: string; +}): ParseArtifact { + return ParseArtifactSchema.parse({ + artifactHash: "a".repeat(64), + contentType: "mixed", + createdAt: "2026-08-18T00:00:00.000Z", + documentAssetId: "018f0d60-7a49-7cc2-9c1b-5b36f18f2c43", + elements: [ + { + id: "018f0d60-7a49-7cc2-9c1b-5b36f18f2c44:element-1", + metadata: { + assetRef: { + contentType: "image/png", + objectKey: `tenant/spaces/space/documents/asset/assets/${owner}/figure.png`, + sha256, + source: "pdf-raster", + variants: { + thumbnail: { + contentType: "image/png", + objectKey: `tenant/spaces/space/documents/asset/assets/${owner}/thumb.png`, + sha256: "d".repeat(64), + }, + }, + }, + pdfRaster: { + boundingBox: { height: 40, width: 30, x: 10, y: 20 }, + cropKind: "figure", + pageNumber: 1, + renderer: { command: "/usr/bin/pdftoppm", dpi, thumbnailDpi: 48 }, + }, + }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + ], + id: "018f0d60-7a49-7cc2-9c1b-5b36f18f2c44", + metadata: {}, + parser: "unstructured", + version: 1, + }); +} diff --git a/knowledge-fs/packages/api/src/document-multimodal-artifact.ts b/knowledge-fs/packages/api/src/document-multimodal-artifact.ts new file mode 100644 index 00000000000..8b1e4648636 --- /dev/null +++ b/knowledge-fs/packages/api/src/document-multimodal-artifact.ts @@ -0,0 +1,151 @@ +import { createHash } from "node:crypto"; + +import { type ParseArtifact, ParseArtifactSchema, stableJson } from "@knowledge/core"; + +import { cloneJsonObject, isPlainObject } from "./json-utils"; + +const materializationContractVersion = 1; + +/** + * Makes durable visual bytes part of parse-artifact lineage without hashing attempt-scoped object + * keys. Re-running the same parser output with the same visual material is therefore idempotent, + * while a renderer/model/configuration change that changes bytes or relevant crop semantics + * produces a new publication fingerprint. + */ +export function finalizeDocumentMultimodalArtifact(artifact: ParseArtifact): ParseArtifact { + const existingMaterialization = isPlainObject(artifact.metadata.multimodalMaterialization) + ? artifact.metadata.multimodalMaterialization + : {}; + const sourceArtifactHash = + typeof existingMaterialization.sourceArtifactHash === "string" && + /^[a-f0-9]{64}$/u.test(existingMaterialization.sourceArtifactHash) + ? existingMaterialization.sourceArtifactHash + : artifact.artifactHash; + const visualMaterial = artifact.elements.flatMap((element, ordinal) => { + const assetRef = isPlainObject(element.metadata.assetRef) ? element.metadata.assetRef : {}; + const sha256 = sha256Value(assetRef.sha256); + + if (!sha256) { + return []; + } + + return [ + { + asset: { + contentType: stringValue(assetRef.contentType), + cropKind: stringValue(assetRef.cropKind), + height: numberValue(assetRef.height), + sha256, + source: stringValue(assetRef.source), + variants: materializedVariants(assetRef.variants), + width: numberValue(assetRef.width), + }, + ordinal, + pdfRaster: materializedPdfRaster(element.metadata.pdfRaster), + type: element.type, + }, + ]; + }); + + if (visualMaterial.length === 0) { + return artifact; + } + + const digest = createHash("sha256") + .update( + stableJson({ + contractVersion: materializationContractVersion, + sourceArtifactHash, + visualMaterial, + }), + ) + .digest("hex"); + + return ParseArtifactSchema.parse({ + ...artifact, + artifactHash: digest, + metadata: { + ...cloneJsonObject(artifact.metadata), + multimodalMaterialization: { + assetCount: visualMaterial.length, + contractVersion: materializationContractVersion, + digest, + sourceArtifactHash, + }, + }, + }); +} + +function materializedVariants(value: unknown): Readonly> { + if (!isPlainObject(value)) { + return {}; + } + + return Object.fromEntries( + Object.entries(value).flatMap(([name, candidate]) => { + if (!isPlainObject(candidate)) { + return []; + } + const sha256 = sha256Value(candidate.sha256); + + return sha256 + ? [ + [ + name, + { + contentType: stringValue(candidate.contentType), + height: numberValue(candidate.height), + sha256, + width: numberValue(candidate.width), + }, + ], + ] + : []; + }), + ); +} + +function materializedPdfRaster(value: unknown): Readonly> | null { + if (!isPlainObject(value)) { + return null; + } + + const renderer = isPlainObject(value.renderer) ? value.renderer : {}; + + return { + boundingBox: jsonValue(value.boundingBox), + cropKind: stringValue(value.cropKind), + geometry: jsonValue(value.geometry), + pageNumber: numberValue(value.pageNumber), + renderer: { + dimensionCapped: booleanValue(renderer.dimensionCapped), + dpi: numberValue(renderer.dpi), + thumbnailDpi: numberValue(renderer.thumbnailDpi), + variant: stringValue(renderer.variant), + }, + }; +} + +function sha256Value(value: unknown): string | null { + return typeof value === "string" && /^[a-f0-9]{64}$/u.test(value) ? value : null; +} + +function stringValue(value: unknown): string | null { + return typeof value === "string" ? value : null; +} + +function numberValue(value: unknown): number | null { + return typeof value === "number" && Number.isFinite(value) ? value : null; +} + +function booleanValue(value: unknown): boolean | null { + return typeof value === "boolean" ? value : null; +} + +function jsonValue(value: unknown): unknown { + return Array.isArray(value) + ? value.map((item) => jsonValue(item)) + : isPlainObject(value) + ? cloneJsonObject(value) + : (value ?? null); +} diff --git a/knowledge-fs/packages/api/src/document-multimodal-asset-extractor.test.ts b/knowledge-fs/packages/api/src/document-multimodal-asset-extractor.test.ts index 2513dbd0afe..da61de892d4 100644 --- a/knowledge-fs/packages/api/src/document-multimodal-asset-extractor.test.ts +++ b/knowledge-fs/packages/api/src/document-multimodal-asset-extractor.test.ts @@ -72,6 +72,58 @@ describe("extractDocumentMultimodalAssets", () => { ).resolves.toEqual(new Uint8Array([1, 2, 3, 4])); }); + it("stores provider table image payloads without discarding table semantics", async () => { + const adapter = createNodePlatformAdapter({ env: {} }); + + const result = await extractDocumentMultimodalAssets({ + artifact: { + artifactHash: "a".repeat(64), + contentType: "mixed", + createdAt: "2026-06-23T00:00:00.000Z", + documentAssetId, + elements: [ + { + id: "table-1", + metadata: { + assetRef: { + contentType: "image/png", + uri: "data:image/png;base64,AQIDBA==", + }, + table: { html: "
42
" }, + textAsHtml: "
42
", + }, + sectionPath: ["Results"], + text: "42", + type: "table", + }, + ], + id: parseArtifactId, + metadata: {}, + parser: "unstructured", + version: 1, + }, + knowledgeSpaceId, + objectStorage: adapter.objectStorage, + tenantId: "tenant-1", + }); + + expect(result.extractedCount).toBe(1); + expect(result.artifact.elements[0]).toMatchObject({ + metadata: { + assetRef: { + contentType: "image/png", + objectKey: expect.stringMatching(/table-1-[a-f0-9]{12}\.png$/u), + sha256: "9f64a747e1b97f131fabb6b447296c9b6f0201e79fb3c5356e6c77e89b6a806a", + }, + table: { html: "
42
" }, + textAsHtml: "
42
", + }, + text: "42", + type: "table", + }); + expect(result.artifact.elements[0]?.metadata).not.toHaveProperty("assetRef.uri"); + }); + it("records dimensions for extracted image asset refs when headers expose them", async () => { const adapter = createNodePlatformAdapter({ env: {} }); const pngBytes = new Uint8Array([ diff --git a/knowledge-fs/packages/api/src/document-multimodal-asset-extractor.ts b/knowledge-fs/packages/api/src/document-multimodal-asset-extractor.ts index 90ded5fdd9b..63572868468 100644 --- a/knowledge-fs/packages/api/src/document-multimodal-asset-extractor.ts +++ b/knowledge-fs/packages/api/src/document-multimodal-asset-extractor.ts @@ -25,12 +25,13 @@ export interface ExtractDocumentMultimodalAssetsInput { readonly imageVariantGenerator?: DocumentImageVariantGenerator | undefined; readonly objectStorage: PlatformAdapter["objectStorage"]; readonly tenantId: string; + readonly writeOwnerId?: string | undefined; } export interface ExtractDocumentMultimodalAssetsResult { readonly artifact: ParseArtifact; readonly extractedCount: number; - /** Number of extractable images left inline because the per-document cap was reached. */ + /** Number of extractable visual assets left inline because the per-document cap was reached. */ readonly skippedForCapCount: number; } @@ -61,6 +62,7 @@ export async function extractDocumentMultimodalAssets({ imageVariantGenerator, objectStorage, tenantId, + writeOwnerId, }: ExtractDocumentMultimodalAssetsInput): Promise { if (!Number.isSafeInteger(maxEmbeddedAssetBytes) || maxEmbeddedAssetBytes < 1) { throw new Error("Document multimodal embedded asset max bytes must be at least 1"); @@ -81,7 +83,7 @@ export async function extractDocumentMultimodalAssets({ const allowedLocalRoots = normalizeAllowedLocalAssetPaths(allowLocalAssetPaths); for (const element of artifact.elements) { - if (element.type !== "image") { + if (element.type !== "image" && element.type !== "table") { elements.push(element); continue; } @@ -118,6 +120,7 @@ export async function extractDocumentMultimodalAssets({ knowledgeSpaceId, sha256, tenantId, + ...(writeOwnerId ? { writeOwnerId } : {}), }); await objectStorage.putObject({ @@ -130,6 +133,7 @@ export async function extractDocumentMultimodalAssets({ parseElementId: element.id, sha256, tenantId, + ...(writeOwnerId ? { writeOwnerId } : {}), }, }); const variants = imageVariantGenerator @@ -141,6 +145,7 @@ export async function extractDocumentMultimodalAssets({ knowledgeSpaceId, objectStorage, tenantId, + ...(writeOwnerId ? { writeOwnerId } : {}), }) : {}; @@ -207,6 +212,7 @@ async function storeGeneratedImageVariants({ knowledgeSpaceId, objectStorage, tenantId, + writeOwnerId, }: { readonly assetId: string; readonly elementId: string; @@ -215,6 +221,7 @@ async function storeGeneratedImageVariants({ readonly knowledgeSpaceId: string; readonly objectStorage: PlatformAdapter["objectStorage"]; readonly tenantId: string; + readonly writeOwnerId?: string | undefined; }): Promise>> { const variants: Record> = {}; const generated = await generator.generate({ @@ -231,6 +238,7 @@ async function storeGeneratedImageVariants({ objectStorage, tenantId, variant, + ...(writeOwnerId ? { writeOwnerId } : {}), }); variants[variant.name] = stored; } @@ -245,6 +253,7 @@ async function storeGeneratedImageVariant({ objectStorage, tenantId, variant, + writeOwnerId, }: { readonly assetId: string; readonly elementId: string; @@ -252,6 +261,7 @@ async function storeGeneratedImageVariant({ readonly objectStorage: PlatformAdapter["objectStorage"]; readonly tenantId: string; readonly variant: GeneratedDocumentImageVariant; + readonly writeOwnerId?: string | undefined; }): Promise> { const sha256 = sha256Hex(variant.body); const objectKey = createDocumentMultimodalAssetVariantObjectKey({ @@ -262,6 +272,7 @@ async function storeGeneratedImageVariant({ sha256, tenantId, variant: variant.name, + ...(writeOwnerId ? { writeOwnerId } : {}), }); await objectStorage.putObject({ @@ -275,6 +286,7 @@ async function storeGeneratedImageVariant({ sha256, tenantId, variant: variant.name, + ...(writeOwnerId ? { writeOwnerId } : {}), }, }); diff --git a/knowledge-fs/packages/api/src/document-pdf-rasterizer-coverage.test.ts b/knowledge-fs/packages/api/src/document-pdf-rasterizer-coverage.test.ts index df2407e8ba8..dcf88c03a87 100644 --- a/knowledge-fs/packages/api/src/document-pdf-rasterizer-coverage.test.ts +++ b/knowledge-fs/packages/api/src/document-pdf-rasterizer-coverage.test.ts @@ -54,6 +54,32 @@ describe("document pdf rasterizer coverage", () => { tenantId: "tenant-1", }), ).rejects.toThrow("Document PDF rasterized asset max count must be at least 1"); + + await expect( + rasterizeDocumentPdfMultimodalAssets({ + artifact: artifact([]), + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + maxDurationMs: 0, + objectStorage: adapter.objectStorage, + rasterizer: recordingRasterizer([]), + tenantId: "tenant-1", + }), + ).rejects.toThrow("Document PDF rasterization maxDurationMs must be at least 1"); + + await expect( + rasterizeDocumentPdfMultimodalAssets({ + artifact: artifact([]), + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + maxRasterizedBytes: 0, + objectStorage: adapter.objectStorage, + rasterizer: recordingRasterizer([]), + tenantId: "tenant-1", + }), + ).rejects.toThrow("Document PDF rasterized byte max must be at least 1"); }); it("keeps elements unchanged when the renderer yields no image", async () => { @@ -85,7 +111,7 @@ describe("document pdf rasterizer coverage", () => { expect(result.artifact.metadata).not.toHaveProperty("pdfRasterAssets"); }); - it("skips elements without pages or boxes and elements that already have assets", async () => { + it("reports unresolved images without pages or boxes before rendering partial assets", async () => { const adapter = createNodePlatformAdapter({ env: {} }); const calls: RenderDocumentPdfPageInput[] = []; const boundingBox = { height: 10, width: 10, x: 0, y: 0 }; @@ -124,11 +150,46 @@ describe("document pdf rasterizer coverage", () => { tenantId: "tenant-1", }); - expect(calls.map((call) => call.elementId)).toEqual(["empty-asset-ref"]); - expect(result.rasterizedCount).toBe(1); + expect(calls).toEqual([]); + expect(result).toMatchObject({ + candidateCount: 1, + rasterizedCount: 0, + unresolvedCount: 3, + }); expect(result.artifact.elements[0]?.metadata).not.toHaveProperty("assetRef"); - expect(result.artifact.elements[4]?.metadata.assetRef).toMatchObject({ - source: "pdf-raster", + expect(result.artifact.elements[4]?.metadata.assetRef).toEqual({ + note: "not a stored asset yet", + }); + }); + + it("forces provider fallback for a table without a bounding box", async () => { + const adapter = createNodePlatformAdapter({ env: {} }); + const calls: RenderDocumentPdfPageInput[] = []; + + const result = await rasterizeDocumentPdfMultimodalAssets({ + artifact: artifact([ + { id: "table-no-box", metadata: {}, pageNumber: 1, sectionPath: [], type: "table" }, + { + id: "figure-1", + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + ]), + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + objectStorage: adapter.objectStorage, + rasterizer: recordingRasterizer(calls), + tenantId: "tenant-1", + }); + + expect(calls).toEqual([]); + expect(result).toMatchObject({ + candidateCount: 1, + rasterizedCount: 0, + unresolvedCount: 2, }); }); @@ -173,6 +234,27 @@ describe("document pdf rasterizer coverage", () => { expect(() => createPopplerPdfRasterizer({ timeoutMs: 0 })).toThrow( "Poppler PDF rasterizer timeoutMs must be at least 1", ); + expect(() => createPopplerPdfRasterizer({ maxConcurrency: 0 })).toThrow( + "Poppler PDF rasterizer maxConcurrency must be at least 1", + ); + expect(() => createPopplerPdfRasterizer({ maxPageDimension: 0 })).toThrow( + "Poppler PDF rasterizer maxPageDimension must be at least 1", + ); + expect(() => createPopplerPdfRasterizer({ maxPagePixels: 0 })).toThrow( + "Poppler PDF rasterizer maxPagePixels must be at least 1", + ); + expect(() => createPopplerPdfRasterizer({ maxEncodedPageBytes: 0 })).toThrow( + "Poppler PDF rasterizer maxEncodedPageBytes must be at least 1", + ); + expect(() => createPopplerPdfRasterizer({ maxEncodedImageBytes: 0 })).toThrow( + "Poppler PDF rasterizer maxEncodedImageBytes must be at least 1", + ); + expect(() => createPopplerPdfRasterizer({ maxEncodedCropBytes: 0 })).toThrow( + "Poppler PDF rasterizer maxEncodedCropBytes must be at least 1", + ); + expect(() => createPopplerPdfRasterizer({ pdfInfoCommand: "" })).toThrow( + "Poppler PDF rasterizer pdfInfoCommand must be non-empty", + ); }); it("infers crop kinds from explicit metadata kind hints", async () => { @@ -306,7 +388,7 @@ describe("document pdf rasterizer coverage", () => { }); }); - it("rejects malformed bounding boxes so their elements are skipped", async () => { + it("rejects malformed image boxes and avoids publishing a partial raster set", async () => { const adapter = createNodePlatformAdapter({ env: {} }); const calls: RenderDocumentPdfPageInput[] = []; @@ -370,13 +452,12 @@ describe("document pdf rasterizer coverage", () => { tenantId: "tenant-1", }); - // Only the element with a valid box (unknown units default to pixel) rasterizes. - expect(calls.map((call) => call.elementId)).toEqual(["unknown-unit"]); - expect(calls[0]).toMatchObject({ - boundingBox: { height: 4, width: 3, x: 1, y: 2 }, - boundingBoxGeometry: { coordinateSystem: "pixel" }, + expect(calls).toEqual([]); + expect(result).toMatchObject({ + candidateCount: 1, + rasterizedCount: 0, + unresolvedCount: 7, }); - expect(result.rasterizedCount).toBe(1); }); it("leaves bounding boxes unchanged for identity geometries", () => { diff --git a/knowledge-fs/packages/api/src/document-pdf-rasterizer.test.ts b/knowledge-fs/packages/api/src/document-pdf-rasterizer.test.ts index 66bdabe1c37..f16fa04f556 100644 --- a/knowledge-fs/packages/api/src/document-pdf-rasterizer.test.ts +++ b/knowledge-fs/packages/api/src/document-pdf-rasterizer.test.ts @@ -1,9 +1,16 @@ +import { chmod, mkdtemp, readFile, rm, stat, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + import { createNodePlatformAdapter } from "@knowledge/adapters/node"; import type { ParseArtifact } from "@knowledge/core"; import { describe, expect, it } from "vitest"; import { type DocumentPdfRasterizer, + DocumentPdfRenderError, + type RenderDocumentPdfPageRequest, + awaitUncancellablePdfRasterOperation, createPopplerPdfRasterizer, normalizePdfRasterBoundingBoxForDpi, rasterizeDocumentPdfMultimodalAssets, @@ -14,6 +21,88 @@ const documentAssetId = "018f0d60-7a49-7cc2-9c1b-5b36f18f2c43"; const parseArtifactId = "018f0d60-7a49-7cc2-9c1b-5b36f18f2c44"; const documentBody = new TextEncoder().encode("%PDF-1.7\n"); +interface FakePopplerCommand { + readonly cleanup: () => Promise; + readonly command: string; + readonly pdfInfoCommand: string; + readonly readInvocations: () => Promise; + readonly readWorkDirs: () => Promise; +} + +async function createFakePopplerCommand( + mode: "failure" | "success" | "timeout" = "success", +): Promise { + const root = await mkdtemp(join(tmpdir(), "knowledge-fs-fake-poppler-")); + const command = join(root, "pdftoppm.cjs"); + const invocationLog = join(root, "invocations.ndjson"); + const workDirLog = join(root, "workdirs.txt"); + const sharp = (await import("sharp")).default; + const pagePng = await sharp({ + create: { + background: { alpha: 1, b: 255, g: 127, r: 63 }, + channels: 4, + height: 80, + width: 100, + }, + }) + .png() + .toBuffer(); + const script = `#!/usr/bin/env node +const fs = require("node:fs"); +const path = require("node:path"); +const args = process.argv.slice(2); +if (args.includes("-box")) { + const pageNumber = args[args.indexOf("-f") + 1]; + process.stdout.write("Page " + pageNumber + " size: 612 x 792 pts\\n"); + process.exit(0); +} +const inputPath = args.at(-2); +const outputPrefix = args.at(-1); +const pageNumber = args[args.indexOf("-f") + 1]; +fs.appendFileSync(${JSON.stringify(invocationLog)}, JSON.stringify(args) + "\\n"); +fs.appendFileSync(${JSON.stringify(workDirLog)}, path.dirname(inputPath) + "\\n"); +if (${JSON.stringify(mode)} === "timeout") { + setTimeout(() => {}, 10_000); +} else if (${JSON.stringify(mode)} === "failure") { + process.stderr.write("fake pdftoppm failure"); + process.exitCode = 7; +} else { + fs.writeFileSync(outputPrefix + "-" + pageNumber + ".png", Buffer.from(${JSON.stringify(pagePng.toString("base64"))}, "base64")); +} +`; + await writeFile(command, script); + await chmod(command, 0o755); + + const readLines = async (path: string): Promise => { + try { + return (await readFile(path, "utf8")).trim().split("\n").filter(Boolean); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === "ENOENT") { + return []; + } + + throw error; + } + }; + + return { + cleanup: async () => rm(root, { force: true, recursive: true }), + command, + pdfInfoCommand: command, + readInvocations: async () => + (await readLines(invocationLog)).map((line) => JSON.parse(line) as string[]), + readWorkDirs: async () => readLines(workDirLog), + }; +} + +async function expectTemporaryDirectoriesRemoved(paths: readonly string[]): Promise { + expect(paths.length).toBeGreaterThan(0); + + for (const path of new Set(paths)) { + await expect(stat(path)).rejects.toMatchObject({ code: "ENOENT" }); + } +} + describe("rasterizeDocumentPdfMultimodalAssets", () => { it("stores rasterized PDF image crops and rewrites asset refs", async () => { const adapter = createNodePlatformAdapter({ env: {} }); @@ -240,6 +329,57 @@ describe("rasterizeDocumentPdfMultimodalAssets", () => { }); }); + it("recognizes Unstructured PixelSpace page geometry", async () => { + const adapter = createNodePlatformAdapter({ env: {} }); + const calls: unknown[] = []; + + await rasterizeDocumentPdfMultimodalAssets({ + artifact: artifact({ + elements: [ + { + id: "figure-pixel-space", + metadata: { + coordinates: { + layout_height: 800, + layout_width: 1_000, + system: "PixelSpace", + x1: 100, + x2: 400, + y1: 200, + y2: 600, + }, + }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + ], + }), + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + objectStorage: adapter.objectStorage, + rasterizer: { + render: async (input) => { + calls.push(input); + return { body: new Uint8Array([1, 2, 3]), contentType: "image/png" }; + }, + }, + tenantId: "tenant-1", + }); + + expect(calls).toEqual([ + expect.objectContaining({ + boundingBox: { height: 400, width: 300, x: 100, y: 200 }, + boundingBoxGeometry: { + coordinateSystem: "pixel", + pageHeight: 800, + pageWidth: 1_000, + }, + }), + ]); + }); + it("rasterizes table elements as table-specific visual crops", async () => { const adapter = createNodePlatformAdapter({ env: {} }); const calls: unknown[] = []; @@ -382,6 +522,385 @@ describe("rasterizeDocumentPdfMultimodalAssets", () => { ).rejects.toThrow("Document PDF rasterized asset count exceeds maxRasterizedAssets=1"); }); + it("reports unresolved PDF candidates when no rasterizer is configured", async () => { + const adapter = createNodePlatformAdapter({ env: {} }); + const input = artifact({ + elements: [ + { + id: "figure-1", + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + ], + }); + + const result = await rasterizeDocumentPdfMultimodalAssets({ + artifact: input, + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + objectStorage: adapter.objectStorage, + tenantId: "tenant-1", + }); + + expect(result).toMatchObject({ + artifact: input, + candidateCount: 1, + rasterizedCount: 0, + unresolvedCount: 1, + }); + }); + + it("counts null render results as unresolved candidates", async () => { + const adapter = createNodePlatformAdapter({ env: {} }); + let renderCount = 0; + + const result = await rasterizeDocumentPdfMultimodalAssets({ + artifact: artifact({ + elements: [ + { + id: "figure-1", + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + { + id: "figure-2", + metadata: { boundingBox: { height: 10, width: 10, x: 20, y: 20 } }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + ], + }), + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + objectStorage: adapter.objectStorage, + rasterizer: { + render: async () => { + renderCount += 1; + return renderCount === 1 + ? { body: new Uint8Array([1, 2, 3]), contentType: "image/png" } + : null; + }, + }, + tenantId: "tenant-1", + }); + + expect(result).toMatchObject({ + candidateCount: 2, + rasterizedCount: 0, + unresolvedCount: 2, + }); + expect(result.artifact.elements[0]?.metadata).not.toHaveProperty("assetRef"); + expect(result.artifact.elements[1]?.metadata).not.toHaveProperty("assetRef"); + }); + + it("compensates prior pages when a later image is unresolved", async () => { + const adapter = createNodePlatformAdapter({ env: {} }); + const storedKeys: string[] = []; + const deletedKeys: string[] = []; + + const result = await rasterizeDocumentPdfMultimodalAssets({ + artifact: artifact({ + elements: [ + { + id: "figure-1", + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + { + id: "figure-2", + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber: 2, + sectionPath: [], + type: "image", + }, + ], + }), + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + objectStorage: { + ...adapter.objectStorage, + deleteObject: async (key) => { + deletedKeys.push(key); + await adapter.objectStorage.deleteObject(key); + }, + putObject: async (input) => { + storedKeys.push(input.key); + return adapter.objectStorage.putObject(input); + }, + }, + rasterizer: { + render: async ({ pageNumber }) => + pageNumber === 1 ? { body: new Uint8Array([1, 2, 3]), contentType: "image/png" } : null, + }, + tenantId: "tenant-1", + }); + + expect(result).toMatchObject({ + candidateCount: 2, + rasterizedCount: 0, + unresolvedCount: 2, + }); + expect(result.artifact.elements.every((element) => !element.metadata.assetRef)).toBe(true); + expect(storedKeys).toHaveLength(1); + expect(deletedKeys).toEqual(storedKeys); + }); + + it("finishes rendering before object storage and preserves the render error cause", async () => { + const adapter = createNodePlatformAdapter({ env: {} }); + const renderFailure = new Error("second crop failed"); + let renderCount = 0; + let putCount = 0; + + const promise = rasterizeDocumentPdfMultimodalAssets({ + artifact: artifact({ + elements: [ + { + id: "figure-1", + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + { + id: "figure-2", + metadata: { boundingBox: { height: 10, width: 10, x: 20, y: 20 } }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + ], + }), + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + objectStorage: { + ...adapter.objectStorage, + putObject: async (input) => { + putCount += 1; + return adapter.objectStorage.putObject(input); + }, + }, + rasterizer: { + render: async () => { + renderCount += 1; + + if (renderCount === 2) { + throw renderFailure; + } + + return { body: new Uint8Array([1, 2, 3]), contentType: "image/png" }; + }, + }, + tenantId: "tenant-1", + }); + + const error = await promise.catch((caught: unknown) => caught); + expect(error).toBeInstanceOf(DocumentPdfRenderError); + expect((error as DocumentPdfRenderError).cause).toBe(renderFailure); + expect(renderCount).toBe(2); + expect(putCount).toBe(0); + }); + + it("renders one page at a time and compensates earlier pages after a later render failure", async () => { + const adapter = createNodePlatformAdapter({ env: {} }); + const renderFailure = new Error("page 2 failed"); + const batches: number[][] = []; + const storedKeys: string[] = []; + const deletedKeys: string[] = []; + const promise = rasterizeDocumentPdfMultimodalAssets({ + artifact: artifact({ + elements: [ + { + id: "figure-1a", + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + { + id: "figure-1b", + metadata: { boundingBox: { height: 10, width: 10, x: 20, y: 20 } }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + { + id: "figure-2", + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber: 2, + sectionPath: [], + type: "image", + }, + ], + }), + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + objectStorage: { + ...adapter.objectStorage, + deleteObject: async (key) => { + deletedKeys.push(key); + await adapter.objectStorage.deleteObject(key); + }, + putObject: async (input) => { + storedKeys.push(input.key); + return adapter.objectStorage.putObject(input); + }, + }, + rasterizer: { + render: async () => { + throw new Error("single render should not be used"); + }, + renderBatch: async ({ requests }) => { + batches.push(requests.map(({ pageNumber }) => pageNumber)); + + if (requests[0]?.pageNumber === 2) { + throw renderFailure; + } + + return requests.map((_, index) => ({ + body: new Uint8Array([index + 1]), + contentType: "image/png" as const, + })); + }, + }, + tenantId: "tenant-1", + }); + + const error = await promise.catch((caught: unknown) => caught); + expect(error).toBeInstanceOf(DocumentPdfRenderError); + expect((error as DocumentPdfRenderError).cause).toBe(renderFailure); + expect(batches).toEqual([[1, 1], [2]]); + expect(storedKeys).toHaveLength(2); + expect(deletedKeys).toEqual([...storedKeys].reverse()); + + for (const key of storedKeys) { + await expect(adapter.objectStorage.headObject(key)).resolves.toBeNull(); + } + }); + + it("does not expose a fallback-safe render error when compensation fails", async () => { + const adapter = createNodePlatformAdapter({ env: {} }); + const renderFailure = new Error("page 2 failed"); + let deleteCount = 0; + const promise = rasterizeDocumentPdfMultimodalAssets({ + artifact: artifact({ + elements: [ + { + id: "figure-1", + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + { + id: "figure-2", + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber: 2, + sectionPath: [], + type: "image", + }, + ], + }), + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + objectStorage: { + ...adapter.objectStorage, + deleteObject: async () => { + deleteCount += 1; + throw Object.assign(new Error("cleanup failed"), { retryable: true }); + }, + }, + rasterizer: { + render: async ({ pageNumber }) => { + if (pageNumber === 2) { + throw renderFailure; + } + + return { body: new Uint8Array([1]), contentType: "image/png" }; + }, + }, + tenantId: "tenant-1", + }); + + const error = await promise.catch((caught: unknown) => caught); + expect(error).toBeInstanceOf(AggregateError); + expect(error).not.toBeInstanceOf(DocumentPdfRenderError); + expect((error as AggregateError).cause).toBeInstanceOf(DocumentPdfRenderError); + expect(error).toMatchObject({ retryable: true }); + expect(deleteCount).toBe(1); + }); + + it("preserves storage errors and compensates objects written before a variant failure", async () => { + const adapter = createNodePlatformAdapter({ env: {} }); + const storageFailure = new Error("object storage unavailable"); + const storedKeys: string[] = []; + const deletedKeys: string[] = []; + let putCount = 0; + const promise = rasterizeDocumentPdfMultimodalAssets({ + artifact: artifact({ + elements: [ + { + id: "figure-1", + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + ], + }), + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + objectStorage: { + ...adapter.objectStorage, + deleteObject: async (key) => { + deletedKeys.push(key); + await adapter.objectStorage.deleteObject(key); + }, + putObject: async (input) => { + putCount += 1; + + if (putCount === 2) { + throw storageFailure; + } + + storedKeys.push(input.key); + return adapter.objectStorage.putObject(input); + }, + }, + rasterizer: { + render: async () => ({ + body: new Uint8Array([1, 2, 3]), + contentType: "image/png", + variants: { + thumbnail: { + body: new Uint8Array([4, 5, 6]), + contentType: "image/png", + }, + }, + }), + }, + tenantId: "tenant-1", + }); + + await expect(promise).rejects.toBe(storageFailure); + expect(storedKeys).toHaveLength(1); + expect(deletedKeys).toEqual(storedKeys); + await expect(adapter.objectStorage.headObject(storedKeys[0] ?? "")).resolves.toBeNull(); + }); + it("validates Poppler thumbnail rasterizer options", () => { expect(() => createPopplerPdfRasterizer({ thumbnailDpi: 0 })).toThrow( "Poppler PDF rasterizer thumbnailDpi must be at least 1", @@ -414,9 +933,568 @@ describe("rasterizeDocumentPdfMultimodalAssets", () => { geometry: { coordinateSystem: "pixel", sourceDpi: 144 }, }), ).toEqual({ height: 10, width: 20, x: 5, y: 2.5 }); + expect( + normalizePdfRasterBoundingBoxForDpi({ + boundingBox: { height: 72, width: 144, x: 36, y: 18 }, + dpi: 144, + geometry: { coordinateSystem: "pdf-point", pageHeight: 792, pageWidth: 612 }, + renderedPage: { height: 1_056, width: 816 }, + }), + ).toEqual({ height: 96, width: 192, x: 48, y: 24 }); + }); + + it("caps aggregate raster output across page-sized render batches", async () => { + const adapter = createNodePlatformAdapter({ env: {} }); + const storedKeys: string[] = []; + const deletedKeys: string[] = []; + const promise = rasterizeDocumentPdfMultimodalAssets({ + artifact: artifact({ + elements: [1, 2].map((pageNumber) => ({ + id: `figure-${pageNumber}`, + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber, + sectionPath: [], + type: "image" as const, + })), + }), + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + maxRasterizedBytes: 3, + objectStorage: { + ...adapter.objectStorage, + deleteObject: async (key) => { + deletedKeys.push(key); + }, + headObject: async () => null, + putObject: async (input) => { + storedKeys.push(input.key); + return { + ...(input.contentType ? { contentType: input.contentType } : {}), + key: input.key, + metadata: input.metadata ?? {}, + sizeBytes: input.body.byteLength, + }; + }, + }, + rasterizer: { + render: async () => ({ body: new Uint8Array([1, 2]), contentType: "image/png" }), + }, + tenantId: "tenant-1", + writeOwnerId: "compilation-1", + }); + + await expect(promise).rejects.toThrow("maxRasterizedBytes=3"); + expect(storedKeys).toHaveLength(1); + expect(storedKeys[0]).toContain("/assets/compilation-1/"); + expect(deletedKeys).toEqual(storedKeys); }); }); +describe("createPopplerPdfRasterizer batch rendering", () => { + it("holds a bounded slot across the full caller-owned materialization phase", async () => { + const rasterizer = createPopplerPdfRasterizer({ maxConcurrency: 2 }); + const releases: Array<() => void> = []; + let active = 0; + let maxActive = 0; + const operations = Array.from({ length: 4 }, () => + rasterizer.withMaterializationSlot?.(async () => { + active += 1; + maxActive = Math.max(maxActive, active); + await new Promise((resolve) => releases.push(resolve)); + active -= 1; + }), + ); + + await waitForCondition(() => releases.length === 2); + for (const release of releases.splice(0, 2)) release(); + await waitForCondition(() => releases.length === 2); + for (const release of releases.splice(0, 2)) release(); + await Promise.all(operations); + + expect(maxActive).toBe(2); + }); + + it("waits for uncancellable materialization work to settle before surfacing abort", async () => { + const rasterizer = createPopplerPdfRasterizer({ maxConcurrency: 1 }); + const controller = new AbortController(); + let release: (() => void) | undefined; + let settled = false; + const pending = rasterizer.withMaterializationSlot?.(async () => { + await new Promise((resolve) => { + release = resolve; + }); + }, controller.signal); + void pending?.then( + () => { + settled = true; + }, + () => { + settled = true; + }, + ); + await waitForCondition(() => release !== undefined); + controller.abort(); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(settled).toBe(false); + + release?.(); + await expect(pending).rejects.toMatchObject({ name: "AbortError" }); + expect(settled).toBe(true); + }); + + it("renders each distinct DPI once for multiple crops on the same page", async () => { + const fakePoppler = await createFakePopplerCommand(); + + try { + const rasterizer = createPopplerPdfRasterizer({ + command: fakePoppler.command, + dpi: 144, + pdfInfoCommand: fakePoppler.pdfInfoCommand, + thumbnailDpi: 48, + }); + const requests: readonly RenderDocumentPdfPageRequest[] = [ + { + boundingBox: { height: 10.1, width: 20.4, x: 10.8, y: 5.7 }, + boundingBoxGeometry: { coordinateSystem: "pixel" }, + elementId: "figure-1", + pageNumber: 2, + }, + { + boundingBox: { height: 50, width: 50, x: 90, y: 70 }, + boundingBoxGeometry: { coordinateSystem: "pixel" }, + elementId: "figure-2", + pageNumber: 2, + }, + { + boundingBox: { height: 400, width: 300, x: 100, y: 200 }, + boundingBoxGeometry: { + coordinateSystem: "pixel", + pageHeight: 800, + pageWidth: 1_000, + sourceDpi: 300, + }, + elementId: "pixel-space", + pageNumber: 2, + }, + { + boundingBox: { height: 0.5, width: 0.3, x: 0.1, y: 0.25 }, + boundingBoxGeometry: { + coordinateSystem: "relative", + pageHeight: 8_000, + pageWidth: 10_000, + }, + elementId: "relative-space", + pageNumber: 2, + }, + { + boundingBox: { height: 10, width: 10, x: 200, y: 200 }, + boundingBoxGeometry: { coordinateSystem: "pixel" }, + elementId: "outside-page", + pageNumber: 2, + }, + ]; + + const rendered = await rasterizer.renderBatch?.({ documentBody, requests }); + expect(rendered).toHaveLength(5); + + const invocations = await fakePoppler.readInvocations(); + expect( + invocations.map((args) => ({ + dpi: args[args.indexOf("-r") + 1], + pageNumber: args[args.indexOf("-f") + 1], + scaleTo: args.includes("-scale-to") ? args[args.indexOf("-scale-to") + 1] : undefined, + })), + ).toEqual([ + { dpi: "144", pageNumber: "2", scaleTo: undefined }, + { dpi: "48", pageNumber: "2", scaleTo: undefined }, + ]); + expect(invocations.every((args) => !args.includes("-x") && !args.includes("-W"))).toBe(true); + + const sharp = (await import("sharp")).default; + await expect(sharp(rendered?.[0]?.body).metadata()).resolves.toMatchObject({ + height: 11, + width: 22, + }); + await expect(sharp(rendered?.[1]?.body).metadata()).resolves.toMatchObject({ + height: 10, + width: 10, + }); + await expect(sharp(rendered?.[2]?.body).metadata()).resolves.toMatchObject({ + height: 40, + width: 30, + }); + expect(rendered?.[2]?.metadata).toMatchObject({ + crop: { + normalizedBoundingBox: { height: 40, width: 30, x: 10, y: 20 }, + }, + }); + await expect(sharp(rendered?.[3]?.body).metadata()).resolves.toMatchObject({ + height: 40, + width: 30, + }); + expect(rendered?.[3]?.metadata).toMatchObject({ + crop: { + normalizedBoundingBox: { height: 40, width: 30, x: 10, y: 20 }, + }, + }); + expect(rendered?.[4]).toBeNull(); + expect(rendered?.[0]?.metadata).toMatchObject({ + command: fakePoppler.command, + crop: { + boundingBox: { height: 10.1, width: 20.4, x: 10.8, y: 5.7 }, + geometry: { coordinateSystem: "pixel" }, + normalizedBoundingBox: { height: 10.1, width: 20.4, x: 10.8, y: 5.7 }, + }, + dpi: 144, + elementId: "figure-1", + pageNumber: 2, + thumbnailDpi: 48, + }); + expect(rendered?.[0]?.variants?.thumbnail?.metadata).toMatchObject({ + dpi: 48, + elementId: "figure-1", + pageNumber: 2, + variant: "thumbnail", + }); + } finally { + await fakePoppler.cleanup(); + } + }); + + it("renders separate pages independently while preserving batch result order", async () => { + const fakePoppler = await createFakePopplerCommand(); + + try { + const rasterizer = createPopplerPdfRasterizer({ + command: fakePoppler.command, + dpi: 144, + pdfInfoCommand: fakePoppler.pdfInfoCommand, + thumbnailDpi: 48, + }); + const requests: readonly RenderDocumentPdfPageRequest[] = [ + { elementId: "page-4", pageNumber: 4 }, + { + boundingBox: { height: 5, width: 6, x: 1, y: 2 }, + elementId: "figure-2", + pageNumber: 2, + }, + { + boundingBox: { height: 7, width: 8, x: 3, y: 4 }, + elementId: "figure-4", + pageNumber: 4, + }, + ]; + + const rendered = await rasterizer.renderBatch?.({ documentBody, requests }); + expect(rendered?.map((image) => image?.metadata?.elementId)).toEqual([ + "page-4", + "figure-2", + "figure-4", + ]); + expect( + (await fakePoppler.readInvocations()).map((args) => ({ + dpi: args[args.indexOf("-r") + 1], + pageNumber: args[args.indexOf("-f") + 1], + })), + ).toEqual([ + { dpi: "144", pageNumber: "4" }, + { dpi: "48", pageNumber: "4" }, + { dpi: "144", pageNumber: "2" }, + { dpi: "48", pageNumber: "2" }, + ]); + } finally { + await fakePoppler.cleanup(); + } + }); + + it("reuses the page render when the main and thumbnail DPI match", async () => { + const fakePoppler = await createFakePopplerCommand(); + + try { + const rasterizer = createPopplerPdfRasterizer({ + command: fakePoppler.command, + dpi: 96, + pdfInfoCommand: fakePoppler.pdfInfoCommand, + thumbnailDpi: 96, + }); + + const rendered = await rasterizer.renderBatch?.({ + documentBody, + requests: [ + { + boundingBox: { height: 10, width: 10, x: 0, y: 0 }, + elementId: "figure-1", + pageNumber: 1, + }, + ], + }); + + expect(rendered?.[0]?.variants?.thumbnail).toBeDefined(); + expect(await fakePoppler.readInvocations()).toHaveLength(1); + } finally { + await fakePoppler.cleanup(); + } + }); + + it("only applies proportional page caps when the natural render would exceed them", async () => { + const fakePoppler = await createFakePopplerCommand(); + + try { + const rasterizer = createPopplerPdfRasterizer({ + command: fakePoppler.command, + maxPageDimension: 1_000, + pdfInfoCommand: fakePoppler.pdfInfoCommand, + }); + + await rasterizer.render({ documentBody, elementId: "page-1", pageNumber: 1 }); + expect( + (await fakePoppler.readInvocations()).map((args) => + args.includes("-scale-to") ? args[args.indexOf("-scale-to") + 1] : undefined, + ), + ).toEqual(["1000", "333"]); + } finally { + await fakePoppler.cleanup(); + } + }); + + it("rejects encoded Poppler pages before reading oversized output files", async () => { + const fakePoppler = await createFakePopplerCommand(); + + try { + const rasterizer = createPopplerPdfRasterizer({ + command: fakePoppler.command, + maxEncodedPageBytes: 1, + pdfInfoCommand: fakePoppler.pdfInfoCommand, + }); + + await expect( + rasterizer.render({ documentBody, elementId: "page-1", pageNumber: 1 }), + ).rejects.toThrow("maxEncodedPageBytes=1"); + await expectTemporaryDirectoriesRemoved(await fakePoppler.readWorkDirs()); + } finally { + await fakePoppler.cleanup(); + } + }); + + it("caps aggregate encoded crop and variant bytes for one batch", async () => { + const fakePoppler = await createFakePopplerCommand(); + + try { + const rasterizer = createPopplerPdfRasterizer({ + command: fakePoppler.command, + maxEncodedCropBytes: 1, + pdfInfoCommand: fakePoppler.pdfInfoCommand, + thumbnailDpi: 144, + }); + + await expect( + rasterizer.renderBatch?.({ + documentBody, + requests: [ + { + boundingBox: { height: 10, width: 10, x: 0, y: 0 }, + elementId: "figure-1", + pageNumber: 1, + }, + ], + }), + ).rejects.toThrow("maxEncodedCropBytes=1"); + await expectTemporaryDirectoriesRemoved(await fakePoppler.readWorkDirs()); + } finally { + await fakePoppler.cleanup(); + } + }); + + it("refuses to guess absolute crop coordinates when page scaling lacks source dimensions", async () => { + const fakePoppler = await createFakePopplerCommand(); + + try { + const rasterizer = createPopplerPdfRasterizer({ + command: fakePoppler.command, + maxPageDimension: 100, + pdfInfoCommand: fakePoppler.pdfInfoCommand, + thumbnailDpi: 144, + }); + + await expect( + rasterizer.render({ + boundingBox: { height: 10, width: 10, x: 0, y: 0 }, + boundingBoxGeometry: { coordinateSystem: "pdf-point" }, + documentBody, + elementId: "figure-1", + pageNumber: 1, + }), + ).rejects.toThrow("cannot normalize coordinates after page dimension capping"); + await expectTemporaryDirectoriesRemoved(await fakePoppler.readWorkDirs()); + } finally { + await fakePoppler.cleanup(); + } + }); + + it("waits for an uncancellable image operation to settle before reporting abort", async () => { + const controller = new AbortController(); + let settleOperation: ((value: string) => void) | undefined; + const operation = new Promise((resolve) => { + settleOperation = resolve; + }); + let helperSettled = false; + const result = awaitUncancellablePdfRasterOperation(operation, controller.signal); + void result.then( + () => { + helperSettled = true; + }, + () => { + helperSettled = true; + }, + ); + + controller.abort(new DOMException("cancelled", "AbortError")); + await Promise.resolve(); + expect(helperSettled).toBe(false); + + settleOperation?.("done"); + await expect(result).rejects.toMatchObject({ name: "AbortError" }); + expect(helperSettled).toBe(true); + }); + + it("wraps Poppler failures and removes the shared temporary directory", async () => { + const fakePoppler = await createFakePopplerCommand("failure"); + + try { + const rasterizer = createPopplerPdfRasterizer({ + command: fakePoppler.command, + pdfInfoCommand: fakePoppler.pdfInfoCommand, + }); + const promise = rasterizer.render({ + documentBody, + elementId: "figure-1", + pageNumber: 1, + }); + const error = await promise.catch((caught: unknown) => caught); + + expect(error).toBeInstanceOf(DocumentPdfRenderError); + expect((error as DocumentPdfRenderError).cause).toBeInstanceOf(Error); + await expectTemporaryDirectoriesRemoved(await fakePoppler.readWorkDirs()); + } finally { + await fakePoppler.cleanup(); + } + }); + + it("honors Poppler timeouts and removes the shared temporary directory", async () => { + const fakePoppler = await createFakePopplerCommand("timeout"); + + try { + const rasterizer = createPopplerPdfRasterizer({ + command: fakePoppler.command, + pdfInfoCommand: fakePoppler.pdfInfoCommand, + timeoutMs: 1_000, + }); + const promise = rasterizer.render({ + documentBody, + elementId: "figure-1", + pageNumber: 1, + }); + const error = await promise.catch((caught: unknown) => caught); + + expect(error).toBeInstanceOf(DocumentPdfRenderError); + expect((error as DocumentPdfRenderError).cause).toBeInstanceOf(Error); + await expectTemporaryDirectoriesRemoved(await fakePoppler.readWorkDirs()); + } finally { + await fakePoppler.cleanup(); + } + }); + + it("classifies caller cancellation as an abort instead of a render fallback", async () => { + const fakePoppler = await createFakePopplerCommand("timeout"); + + try { + const controller = new AbortController(); + const rasterizer = createPopplerPdfRasterizer({ + command: fakePoppler.command, + pdfInfoCommand: fakePoppler.pdfInfoCommand, + timeoutMs: 5_000, + }); + const adapter = createNodePlatformAdapter({ env: {} }); + const promise = rasterizeDocumentPdfMultimodalAssets({ + artifact: artifact({ + elements: [ + { + id: "figure-1", + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + ], + }), + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + objectStorage: adapter.objectStorage, + rasterizer, + signal: controller.signal, + tenantId: "tenant-1", + }); + setTimeout(() => controller.abort(new DOMException("cancelled", "AbortError")), 100); + const error = await promise.catch((caught: unknown) => caught); + + expect(error).toBeInstanceOf(DOMException); + expect((error as Error).name).toBe("AbortError"); + expect(error).not.toBeInstanceOf(DocumentPdfRenderError); + } finally { + await fakePoppler.cleanup(); + } + }); + + it("enforces the whole-document rasterization duration cap as a render fallback", async () => { + const fakePoppler = await createFakePopplerCommand("timeout"); + + try { + const rasterizer = createPopplerPdfRasterizer({ + command: fakePoppler.command, + pdfInfoCommand: fakePoppler.pdfInfoCommand, + timeoutMs: 5_000, + }); + const adapter = createNodePlatformAdapter({ env: {} }); + const promise = rasterizeDocumentPdfMultimodalAssets({ + artifact: artifact({ + elements: [ + { + id: "figure-1", + metadata: { boundingBox: { height: 10, width: 10, x: 0, y: 0 } }, + pageNumber: 1, + sectionPath: [], + type: "image", + }, + ], + }), + documentBody, + documentMimeType: "application/pdf", + knowledgeSpaceId, + maxDurationMs: 150, + objectStorage: adapter.objectStorage, + rasterizer, + tenantId: "tenant-1", + }); + const error = await promise.catch((caught: unknown) => caught); + + expect(error).toBeInstanceOf(DocumentPdfRenderError); + expect((error as Error).message).toContain("maxDurationMs=150"); + } finally { + await fakePoppler.cleanup(); + } + }); +}); + +async function waitForCondition(condition: () => boolean): Promise { + for (let attempt = 0; attempt < 100; attempt += 1) { + if (condition()) return; + await new Promise((resolve) => setTimeout(resolve, 0)); + } + throw new Error("Timed out waiting for rasterizer test condition"); +} + function artifact(input: Pick): ParseArtifact { return { artifactHash: "a".repeat(64), diff --git a/knowledge-fs/packages/api/src/document-pdf-rasterizer.ts b/knowledge-fs/packages/api/src/document-pdf-rasterizer.ts index abba4d9931e..1afae65a19e 100644 --- a/knowledge-fs/packages/api/src/document-pdf-rasterizer.ts +++ b/knowledge-fs/packages/api/src/document-pdf-rasterizer.ts @@ -1,6 +1,6 @@ import { execFile } from "node:child_process"; import { createHash } from "node:crypto"; -import { mkdtemp, readFile, readdir, rm, writeFile } from "node:fs/promises"; +import { mkdtemp, readFile, readdir, rm, stat, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { promisify } from "node:util"; @@ -13,6 +13,7 @@ import { type PlatformAdapter, } from "@knowledge/core"; +import { createConcurrencyGate } from "./bounded-concurrency"; import { cloneJsonObject, isPlainObject } from "./json-utils"; import { createDocumentMultimodalAssetObjectKey, @@ -23,16 +24,31 @@ const execFileAsync = promisify(execFile); export interface DocumentPdfRasterizer { render(input: RenderDocumentPdfPageInput): Promise; + renderBatch?( + input: RenderDocumentPdfBatchInput, + ): Promise; + /** Holds a bounded slot across source loading, rendering, and object persistence. */ + withMaterializationSlot?(operation: () => Promise, signal?: AbortSignal): Promise; } -export interface RenderDocumentPdfPageInput { +export interface RenderDocumentPdfPageRequest { readonly boundingBox?: DocumentMultimodalBoundingBox | undefined; readonly boundingBoxGeometry?: DocumentPdfBoundingBoxGeometry | undefined; - readonly documentBody: Uint8Array; readonly elementId: string; readonly pageNumber: number; } +export interface RenderDocumentPdfPageInput extends RenderDocumentPdfPageRequest { + readonly documentBody: Uint8Array; + readonly signal?: AbortSignal | undefined; +} + +export interface RenderDocumentPdfBatchInput { + readonly documentBody: Uint8Array; + readonly requests: readonly RenderDocumentPdfPageRequest[]; + readonly signal?: AbortSignal | undefined; +} + export type DocumentPdfBoundingBoxCoordinateSystem = "pdf-point" | "pixel" | "relative"; export interface DocumentPdfBoundingBoxGeometry { @@ -64,26 +80,183 @@ export interface RasterizeDocumentPdfMultimodalAssetsInput { readonly documentBody: Uint8Array; readonly documentMimeType: string; readonly knowledgeSpaceId: string; + readonly maxDurationMs?: number | undefined; readonly maxRasterizedAssets?: number | undefined; + readonly maxRasterizedBytes?: number | undefined; readonly objectStorage: PlatformAdapter["objectStorage"]; readonly rasterizer?: DocumentPdfRasterizer | undefined; + readonly signal?: AbortSignal | undefined; readonly tenantId: string; + readonly writeOwnerId?: string | undefined; } export interface RasterizeDocumentPdfMultimodalAssetsResult { readonly artifact: ParseArtifact; + readonly candidateCount: number; readonly rasterizedCount: number; + readonly unresolvedCount: number; +} + +export class DocumentPdfRenderError extends Error { + constructor(message: string, options?: { readonly cause?: unknown }) { + super(message, { cause: options?.cause }); + this.name = "DocumentPdfRenderError"; + } +} + +class DocumentPdfObjectCleanupError extends AggregateError { + readonly retryable: boolean; + + constructor(errors: readonly unknown[], message: string, options?: { readonly cause?: unknown }) { + super(errors, message, { cause: options?.cause }); + this.name = "DocumentPdfObjectCleanupError"; + this.retryable = errors.some(isRetryableError); + } +} + +function isRetryableError(error: unknown): boolean { + if ( + error instanceof Error && + "retryable" in error && + (error as { readonly retryable?: unknown }).retryable === true + ) { + return true; + } + if (error instanceof AggregateError) { + return error.errors.some(isRetryableError); + } + return false; +} + +function createPdfRasterizationAbortScope({ + maxDurationMs, + parentSignal, +}: { + readonly maxDurationMs: number; + readonly parentSignal?: AbortSignal | undefined; +}): { readonly dispose: () => void; readonly signal: AbortSignal } { + const controller = new AbortController(); + const abortFromParent = () => { + if (parentSignal && !controller.signal.aborted) { + controller.abort(callerPdfRasterAbortReason(parentSignal)); + } + }; + + if (parentSignal?.aborted) { + abortFromParent(); + } else { + parentSignal?.addEventListener("abort", abortFromParent, { once: true }); + } + + const timeout = setTimeout(() => { + if (!controller.signal.aborted) { + controller.abort( + new DocumentPdfRenderError( + `Document PDF rasterization exceeded maxDurationMs=${maxDurationMs}`, + ), + ); + } + }, maxDurationMs); + timeout.unref(); + + return { + dispose: () => { + clearTimeout(timeout); + parentSignal?.removeEventListener("abort", abortFromParent); + }, + signal: controller.signal, + }; +} + +async function raceWithPdfRasterAbort( + operation: Promise, + signal: AbortSignal | undefined, +): Promise { + if (!signal) { + return operation; + } + + if (signal.aborted) { + void operation.catch(() => undefined); + throw pdfRasterAbortReason(signal); + } + + return new Promise((resolve, reject) => { + const onAbort = () => reject(pdfRasterAbortReason(signal)); + const cleanup = () => signal.removeEventListener("abort", onAbort); + signal.addEventListener("abort", onAbort, { once: true }); + + operation.then( + (value) => { + cleanup(); + resolve(value); + }, + (error: unknown) => { + cleanup(); + reject(error); + }, + ); + + if (signal.aborted) { + cleanup(); + reject(pdfRasterAbortReason(signal)); + } + }); +} + +function throwIfCallerAborted(signal: AbortSignal | undefined): void { + if (signal?.aborted) { + throw callerPdfRasterAbortReason(signal); + } +} + +function throwIfPdfRasterAborted(signal: AbortSignal | undefined): void { + if (signal?.aborted) { + throw pdfRasterAbortReason(signal); + } +} + +function callerPdfRasterAbortReason(signal: AbortSignal): Error { + if (signal.reason instanceof Error && !(signal.reason instanceof DocumentPdfRenderError)) { + return signal.reason; + } + + const error = new Error("Document PDF rasterization was aborted by the caller", { + cause: signal.reason, + }); + error.name = "AbortError"; + return error; +} + +function pdfRasterAbortReason(signal: AbortSignal): Error { + return signal.reason instanceof Error + ? signal.reason + : new Error("Document PDF rasterization was aborted"); } export interface PopplerPdfRasterizerOptions { readonly command?: string | undefined; readonly dpi?: number | undefined; + readonly pdfInfoCommand?: string | undefined; + readonly maxEncodedCropBytes?: number | undefined; + readonly maxEncodedImageBytes?: number | undefined; + readonly maxEncodedPageBytes?: number | undefined; + readonly maxConcurrency?: number | undefined; + readonly maxPageDimension?: number | undefined; + readonly maxPagePixels?: number | undefined; readonly thumbnailDpi?: number | undefined; readonly thumbnailVariantName?: string | undefined; readonly timeoutMs?: number | undefined; } const defaultMaxRasterizedAssets = 500; +const defaultMaxRasterizedBytes = 128 * 1024 * 1024; +const defaultMaxRasterizationDurationMs = 10 * 60 * 1_000; +const defaultMaxEncodedCropBytes = 64 * 1024 * 1024; +const defaultMaxEncodedImageBytes = 32 * 1024 * 1024; +const defaultMaxEncodedPageBytes = 32 * 1024 * 1024; +const defaultMaxPageDimension = 4_096; +const defaultMaxPagePixels = 20_000_000; const defaultThumbnailDpi = 48; const defaultThumbnailVariantName = "thumbnail"; @@ -92,115 +265,261 @@ export async function rasterizeDocumentPdfMultimodalAssets({ documentBody, documentMimeType, knowledgeSpaceId, + maxDurationMs = defaultMaxRasterizationDurationMs, maxRasterizedAssets = defaultMaxRasterizedAssets, + maxRasterizedBytes = defaultMaxRasterizedBytes, objectStorage, rasterizer, + signal, tenantId, + writeOwnerId, }: RasterizeDocumentPdfMultimodalAssetsInput): Promise { - if (!rasterizer || !isPdfMimeType(documentMimeType)) { - return { artifact, rasterizedCount: 0 }; + throwIfCallerAborted(signal); + + if (!isPdfMimeType(documentMimeType)) { + return { artifact, candidateCount: 0, rasterizedCount: 0, unresolvedCount: 0 }; } - if (!Number.isSafeInteger(maxRasterizedAssets) || maxRasterizedAssets < 1) { + if (rasterizer && (!Number.isSafeInteger(maxRasterizedAssets) || maxRasterizedAssets < 1)) { throw new Error("Document PDF rasterized asset max count must be at least 1"); } - let rasterizedCount = 0; - const elements = []; + if (rasterizer && (!Number.isSafeInteger(maxDurationMs) || maxDurationMs < 1)) { + throw new Error("Document PDF rasterization maxDurationMs must be at least 1"); + } - for (const element of artifact.elements) { + if (rasterizer && (!Number.isSafeInteger(maxRasterizedBytes) || maxRasterizedBytes < 1)) { + throw new Error("Document PDF rasterized byte max must be at least 1"); + } + + const candidates = artifact.elements.flatMap((element, elementIndex) => { const candidate = pdfRasterizationCandidate(element); if (!candidate) { - elements.push(element); - continue; + return []; } - if (rasterizedCount >= maxRasterizedAssets) { - throw new Error( - `Document PDF rasterized asset count exceeds maxRasterizedAssets=${maxRasterizedAssets}`, + return [ + { + candidate, + element, + elementIndex, + request: { + ...(candidate.boundingBox ? { boundingBox: candidate.boundingBox } : {}), + ...(candidate.boundingBoxGeometry + ? { boundingBoxGeometry: candidate.boundingBoxGeometry } + : {}), + elementId: element.id, + pageNumber: candidate.pageNumber, + } satisfies RenderDocumentPdfPageRequest, + }, + ]; + }); + const candidateCount = candidates.length; + const candidateElementIndexes = new Set(candidates.map(({ elementIndex }) => elementIndex)); + const pendingVisualElementIndexes = new Set( + artifact.elements.flatMap((element, elementIndex) => + isPendingPdfVisual(element) ? [elementIndex] : [], + ), + ); + const fallbackUnresolvedCount = new Set([ + ...candidateElementIndexes, + ...pendingVisualElementIndexes, + ]).size; + const hasUnrenderablePendingVisual = [...pendingVisualElementIndexes].some( + (elementIndex) => !candidateElementIndexes.has(elementIndex), + ); + + if (!rasterizer || candidateCount === 0 || hasUnrenderablePendingVisual) { + return { + artifact, + candidateCount, + rasterizedCount: 0, + unresolvedCount: fallbackUnresolvedCount, + }; + } + + if (candidateCount > maxRasterizedAssets) { + throw new DocumentPdfRenderError( + `Document PDF rasterized asset count exceeds maxRasterizedAssets=${maxRasterizedAssets}`, + ); + } + + const abortScope = createPdfRasterizationAbortScope({ maxDurationMs, parentSignal: signal }); + + const elements = [...artifact.elements]; + const candidatesByPage = new Map(); + + for (const entry of candidates) { + const pageCandidates = candidatesByPage.get(entry.candidate.pageNumber) ?? []; + pageCandidates.push(entry); + candidatesByPage.set(entry.candidate.pageNumber, pageCandidates); + } + + const createdObjectKeys = new Set(); + let rasterizedBytes = 0; + let rasterizedCount = 0; + + try { + // A page is the bounded unit of work: render all of its candidates while the shared page bitmap + // is resident, then persist those results before moving on. + for (const pageCandidates of candidatesByPage.values()) { + const renderedImages = await renderDocumentPdfRequests({ + documentBody, + rasterizer, + requests: pageCandidates.map(({ request }) => request), + signal: abortScope.signal, + }); + rasterizedBytes += renderedImages.reduce( + (total, image) => + total + + (image?.body.byteLength ?? 0) + + Object.values(image?.variants ?? {}).reduce( + (variantTotal, variant) => variantTotal + variant.body.byteLength, + 0, + ), + 0, + ); + + if (rasterizedBytes > maxRasterizedBytes) { + throw new DocumentPdfRenderError( + `Document PDF rasterizer output exceeds maxRasterizedBytes=${maxRasterizedBytes}`, + ); + } + const hasUnresolvedVisual = pageCandidates.some( + ({ element }, index) => + (element.type === "image" || element.type === "table") && renderedImages[index] === null, + ); + + if (hasUnresolvedVisual) { + const cleanupFailures = await deleteCreatedPdfRasterObjects({ + createdObjectKeys, + objectStorage, + }); + + if (cleanupFailures.length > 0) { + throw new DocumentPdfObjectCleanupError( + cleanupFailures, + `Document PDF rasterizer could not compensate ${cleanupFailures.length} object(s) before provider fallback`, + ); + } + + return { + artifact, + candidateCount, + rasterizedCount: 0, + unresolvedCount: fallbackUnresolvedCount, + }; + } + + for (const [ + candidateIndex, + { candidate, element, elementIndex }, + ] of pageCandidates.entries()) { + const rendered = renderedImages[candidateIndex]; + + if (!rendered) { + continue; + } + + const sha256 = sha256Hex(rendered.body); + const objectKey = createDocumentMultimodalAssetObjectKey({ + assetId: artifact.documentAssetId, + contentType: rendered.contentType, + elementId: element.id, + knowledgeSpaceId, + sha256, + tenantId, + ...(writeOwnerId ? { writeOwnerId } : {}), + }); + + await putTrackedPdfRasterObject({ + createdObjectKeys, + input: { + body: rendered.body, + contentType: rendered.contentType, + key: objectKey, + metadata: { + cropKind: candidate.cropKind, + documentAssetId: artifact.documentAssetId, + pageNumber: String(candidate.pageNumber), + parseArtifactId: artifact.id, + parseElementId: element.id, + sha256, + source: "pdf-raster", + tenantId, + }, + }, + objectStorage, + signal: abortScope.signal, + }); + const variants = await storeRenderedImageVariants({ + artifact, + createdObjectKeys, + element, + knowledgeSpaceId, + objectStorage, + pageNumber: candidate.pageNumber, + rendered, + signal: abortScope.signal, + tenantId, + writeOwnerId, + }); + + rasterizedCount += 1; + elements[elementIndex] = { + ...element, + metadata: { + ...cloneJsonObject(element.metadata), + assetRef: { + contentType: rendered.contentType, + cropKind: candidate.cropKind, + objectKey, + sha256, + source: "pdf-raster", + ...(Object.keys(variants).length > 0 ? { variants } : {}), + }, + pdfRaster: { + ...(candidate.boundingBox ? { boundingBox: candidate.boundingBox } : {}), + ...(candidate.boundingBoxGeometry ? { geometry: candidate.boundingBoxGeometry } : {}), + contentType: rendered.contentType, + cropKind: candidate.cropKind, + pageNumber: candidate.pageNumber, + ...(rendered.metadata ? { renderer: cloneJsonObject(rendered.metadata) } : {}), + sha256, + ...(Object.keys(variants).length > 0 ? { variants } : {}), + }, + }, + }; + } + } + } catch (error) { + if (error instanceof DocumentPdfObjectCleanupError) { + throw error; + } + + const cleanupFailures = await deleteCreatedPdfRasterObjects({ + createdObjectKeys, + objectStorage, + }); + + if (cleanupFailures.length > 0) { + throw new DocumentPdfObjectCleanupError( + [error, ...cleanupFailures], + `Document PDF raster asset processing failed and could not compensate ${cleanupFailures.length} object(s)`, + { cause: error }, ); } - const rendered = await rasterizer.render({ - ...(candidate.boundingBox ? { boundingBox: candidate.boundingBox } : {}), - ...(candidate.boundingBoxGeometry - ? { boundingBoxGeometry: candidate.boundingBoxGeometry } - : {}), - documentBody, - elementId: element.id, - pageNumber: candidate.pageNumber, - }); - - if (!rendered) { - elements.push(element); - continue; - } - - const sha256 = sha256Hex(rendered.body); - const objectKey = createDocumentMultimodalAssetObjectKey({ - assetId: artifact.documentAssetId, - contentType: rendered.contentType, - elementId: element.id, - knowledgeSpaceId, - sha256, - tenantId, - }); - - await objectStorage.putObject({ - body: rendered.body, - contentType: rendered.contentType, - key: objectKey, - metadata: { - cropKind: candidate.cropKind, - documentAssetId: artifact.documentAssetId, - pageNumber: String(candidate.pageNumber), - parseArtifactId: artifact.id, - parseElementId: element.id, - sha256, - source: "pdf-raster", - tenantId, - }, - }); - const variants = await storeRenderedImageVariants({ - artifact, - element, - knowledgeSpaceId, - objectStorage, - pageNumber: candidate.pageNumber, - rendered, - tenantId, - }); - - rasterizedCount += 1; - elements.push({ - ...element, - metadata: { - ...cloneJsonObject(element.metadata), - assetRef: { - contentType: rendered.contentType, - cropKind: candidate.cropKind, - objectKey, - sha256, - source: "pdf-raster", - ...(Object.keys(variants).length > 0 ? { variants } : {}), - }, - pdfRaster: { - ...(candidate.boundingBox ? { boundingBox: candidate.boundingBox } : {}), - ...(candidate.boundingBoxGeometry ? { geometry: candidate.boundingBoxGeometry } : {}), - contentType: rendered.contentType, - cropKind: candidate.cropKind, - pageNumber: candidate.pageNumber, - ...(rendered.metadata ? { renderer: cloneJsonObject(rendered.metadata) } : {}), - sha256, - ...(Object.keys(variants).length > 0 ? { variants } : {}), - }, - }, - }); + throw error; + } finally { + abortScope.dispose(); } + const unresolvedCount = candidateCount - rasterizedCount; + if (rasterizedCount === 0) { - return { artifact, rasterizedCount }; + return { artifact, candidateCount, rasterizedCount, unresolvedCount }; } return { @@ -215,26 +534,138 @@ export async function rasterizeDocumentPdfMultimodalAssets({ }, }, }), + candidateCount, rasterizedCount, + unresolvedCount, }; } +async function renderDocumentPdfRequests({ + documentBody, + rasterizer, + requests, + signal, +}: { + readonly documentBody: Uint8Array; + readonly rasterizer: DocumentPdfRasterizer; + readonly requests: readonly RenderDocumentPdfPageRequest[]; + readonly signal?: AbortSignal | undefined; +}): Promise { + try { + throwIfPdfRasterAborted(signal); + + if (rasterizer.renderBatch) { + const rendered = await raceWithPdfRasterAbort( + rasterizer.renderBatch({ documentBody, requests, signal }), + signal, + ); + + if (!Array.isArray(rendered) || rendered.length !== requests.length) { + throw new DocumentPdfRenderError( + `Document PDF rasterizer returned ${Array.isArray(rendered) ? rendered.length : "an invalid"} batch results for ${requests.length} requests`, + ); + } + + return rendered; + } + + const rendered: (RenderedDocumentPdfImage | null)[] = []; + + for (const request of requests) { + rendered.push( + await raceWithPdfRasterAbort( + rasterizer.render({ ...request, documentBody, signal }), + signal, + ), + ); + } + + return rendered; + } catch (error) { + if (signal?.aborted) { + throw pdfRasterAbortReason(signal); + } + + if (error instanceof DocumentPdfRenderError) { + throw error; + } + + throw new DocumentPdfRenderError("Document PDF rasterizer failed to render candidates", { + cause: error, + }); + } +} + +async function putTrackedPdfRasterObject({ + createdObjectKeys, + input, + objectStorage, + signal, +}: { + readonly createdObjectKeys: Set; + readonly input: Parameters[0]; + readonly objectStorage: PlatformAdapter["objectStorage"]; + readonly signal?: AbortSignal | undefined; +}): Promise { + throwIfPdfRasterAborted(signal); + const existedBefore = (await objectStorage.headObject(input.key)) !== null; + throwIfPdfRasterAborted(signal); + await objectStorage.putObject(input); + + if (!existedBefore) { + createdObjectKeys.add(input.key); + } + + throwIfPdfRasterAborted(signal); +} + +async function deleteCreatedPdfRasterObjects({ + createdObjectKeys, + objectStorage, +}: { + readonly createdObjectKeys: Set; + readonly objectStorage: PlatformAdapter["objectStorage"]; +}): Promise { + const failures: unknown[] = []; + + for (const key of [...createdObjectKeys].reverse()) { + try { + await objectStorage.deleteObject(key); + createdObjectKeys.delete(key); + } catch (error) { + failures.push( + new AggregateError([error], `Failed to compensate PDF raster object key=${key}`), + ); + // The worker-level execution owner keeps the complete receipt and performs the bounded retry. + break; + } + } + + return failures; +} + async function storeRenderedImageVariants({ artifact, + createdObjectKeys, element, knowledgeSpaceId, objectStorage, pageNumber, rendered, + signal, tenantId, + writeOwnerId, }: { readonly artifact: ParseArtifact; + readonly createdObjectKeys: Set; readonly element: ParseElement; readonly knowledgeSpaceId: string; readonly objectStorage: PlatformAdapter["objectStorage"]; readonly pageNumber: number; readonly rendered: RenderedDocumentPdfImage; + readonly signal?: AbortSignal | undefined; readonly tenantId: string; + readonly writeOwnerId?: string | undefined; }): Promise>> { const variants: Record> = {}; @@ -248,22 +679,28 @@ async function storeRenderedImageVariants({ sha256, tenantId, variant, + ...(writeOwnerId ? { writeOwnerId } : {}), }); - await objectStorage.putObject({ - body: image.body, - contentType: image.contentType, - key: objectKey, - metadata: { - documentAssetId: artifact.documentAssetId, - pageNumber: String(pageNumber), - parseArtifactId: artifact.id, - parseElementId: element.id, - sha256, - source: "pdf-raster", - tenantId, - variant, + await putTrackedPdfRasterObject({ + createdObjectKeys, + input: { + body: image.body, + contentType: image.contentType, + key: objectKey, + metadata: { + documentAssetId: artifact.documentAssetId, + pageNumber: String(pageNumber), + parseArtifactId: artifact.id, + parseElementId: element.id, + sha256, + source: "pdf-raster", + tenantId, + variant, + }, }, + objectStorage, + signal, }); variants[variant] = { @@ -281,6 +718,13 @@ async function storeRenderedImageVariants({ export function createPopplerPdfRasterizer({ command = "pdftoppm", dpi = 144, + maxEncodedCropBytes = defaultMaxEncodedCropBytes, + maxEncodedImageBytes = defaultMaxEncodedImageBytes, + maxEncodedPageBytes = defaultMaxEncodedPageBytes, + maxConcurrency = 2, + maxPageDimension = defaultMaxPageDimension, + maxPagePixels = defaultMaxPagePixels, + pdfInfoCommand = "pdfinfo", thumbnailDpi = defaultThumbnailDpi, thumbnailVariantName = defaultThumbnailVariantName, timeoutMs = 30_000, @@ -289,6 +733,34 @@ export function createPopplerPdfRasterizer({ throw new Error("Poppler PDF rasterizer dpi must be at least 1"); } + if (!Number.isSafeInteger(maxConcurrency) || maxConcurrency < 1) { + throw new Error("Poppler PDF rasterizer maxConcurrency must be at least 1"); + } + + if (!Number.isSafeInteger(maxPageDimension) || maxPageDimension < 1) { + throw new Error("Poppler PDF rasterizer maxPageDimension must be at least 1"); + } + + if (!Number.isSafeInteger(maxPagePixels) || maxPagePixels < 1) { + throw new Error("Poppler PDF rasterizer maxPagePixels must be at least 1"); + } + + if (!Number.isSafeInteger(maxEncodedPageBytes) || maxEncodedPageBytes < 1) { + throw new Error("Poppler PDF rasterizer maxEncodedPageBytes must be at least 1"); + } + + if (!Number.isSafeInteger(maxEncodedImageBytes) || maxEncodedImageBytes < 1) { + throw new Error("Poppler PDF rasterizer maxEncodedImageBytes must be at least 1"); + } + + if (!Number.isSafeInteger(maxEncodedCropBytes) || maxEncodedCropBytes < 1) { + throw new Error("Poppler PDF rasterizer maxEncodedCropBytes must be at least 1"); + } + + if (!pdfInfoCommand.trim()) { + throw new Error("Poppler PDF rasterizer pdfInfoCommand must be non-empty"); + } + if (!Number.isSafeInteger(thumbnailDpi) || thumbnailDpi < 1) { throw new Error("Poppler PDF rasterizer thumbnailDpi must be at least 1"); } @@ -301,121 +773,411 @@ export function createPopplerPdfRasterizer({ throw new Error("Poppler PDF rasterizer timeoutMs must be at least 1"); } - return { - render: async ({ boundingBox, boundingBoxGeometry, documentBody, elementId, pageNumber }) => { - const workDir = await mkdtemp(join(tmpdir(), "knowledge-fs-pdf-raster-")); - const inputPath = join(workDir, "input.pdf"); - const outputPrefix = join(workDir, "page"); - const thumbnailOutputPrefix = join(workDir, "thumbnail"); + const materializationGate = createConcurrencyGate(maxConcurrency); + const renderGate = createConcurrencyGate(maxConcurrency); + const renderBatch = async ({ + documentBody, + requests, + signal, + }: RenderDocumentPdfBatchInput): Promise => { + const operation = renderGate.run(async () => { + throwIfPdfRasterAborted(signal); try { - await writeFile(inputPath, documentBody); - await renderPopplerPng({ - boundingBox, - boundingBoxGeometry, + return await renderPopplerPdfBatch({ command, + documentBody, dpi, - inputPath, - outputPrefix, - pageNumber, + maxEncodedCropBytes, + maxEncodedImageBytes, + maxEncodedPageBytes, + maxPageDimension, + maxPagePixels, + pdfInfoCommand, + requests, + signal, + thumbnailDpi, + thumbnailVariantName, timeoutMs, }); - const outputPath = await findPopplerOutputPath(workDir, "page-"); - - if (!outputPath) { - return null; + } catch (error) { + if (signal?.aborted) { + throw pdfRasterAbortReason(signal); } - await renderPopplerPng({ - boundingBox, - boundingBoxGeometry, - command, - dpi: thumbnailDpi, - inputPath, - outputPrefix: thumbnailOutputPrefix, - pageNumber, - timeoutMs, - }); - const thumbnailOutputPath = await findPopplerOutputPath(workDir, "thumbnail-"); + if (error instanceof DocumentPdfRenderError) { + throw error; + } + throw new DocumentPdfRenderError("Poppler PDF rasterizer failed", { cause: error }); + } + }); - return { - body: new Uint8Array(await readFile(outputPath)), + return raceWithPdfRasterAbort(operation, signal); + }; + + return { + render: async ({ documentBody, signal, ...request }) => { + const rendered = await renderBatch({ documentBody, requests: [request], signal }); + + return rendered[0] ?? null; + }, + renderBatch, + withMaterializationSlot: async (operation: () => Promise, signal?: AbortSignal) => { + const pending = materializationGate.run(async () => { + throwIfPdfRasterAborted(signal); + const result = await operation(); + throwIfPdfRasterAborted(signal); + return result; + }); + + // The caller owns compensation. Wait for any uncancellable storage/database work to settle + // before surfacing cancellation so cleanup cannot race a late PUT or commit. + return pending; + }, + }; +} + +interface PopplerRenderedPage { + readonly body: Uint8Array; + readonly channels: 1 | 2 | 3 | 4; + readonly height: number; + readonly pixels: Uint8Array; + readonly width: number; + readonly wasDownscaled: boolean; +} + +interface PopplerPdfPageSize { + readonly heightPoints: number; + readonly widthPoints: number; +} + +async function renderPopplerPdfBatch({ + command, + documentBody, + dpi, + maxEncodedCropBytes, + maxEncodedImageBytes, + maxEncodedPageBytes, + maxPageDimension, + maxPagePixels, + pdfInfoCommand, + requests, + signal, + thumbnailDpi, + thumbnailVariantName, + timeoutMs, +}: { + readonly command: string; + readonly documentBody: Uint8Array; + readonly dpi: number; + readonly maxEncodedCropBytes: number; + readonly maxEncodedImageBytes: number; + readonly maxEncodedPageBytes: number; + readonly maxPageDimension: number; + readonly maxPagePixels: number; + readonly pdfInfoCommand: string; + readonly requests: readonly RenderDocumentPdfPageRequest[]; + readonly signal?: AbortSignal | undefined; + readonly thumbnailDpi: number; + readonly thumbnailVariantName: string; + readonly timeoutMs: number; +}): Promise { + if (requests.length === 0) { + return []; + } + + const workDir = await mkdtemp(join(tmpdir(), "knowledge-fs-pdf-raster-")); + const inputPath = join(workDir, "input.pdf"); + + try { + await writeFile(inputPath, documentBody, { signal }); + throwIfPdfRasterAborted(signal); + const rendered: (RenderedDocumentPdfImage | null)[] = Array.from( + { length: requests.length }, + () => null, + ); + const requestsByPage = new Map< + number, + { readonly index: number; readonly request: RenderDocumentPdfPageRequest }[] + >(); + let aggregateEncodedCropBytes = 0; + + for (const [index, request] of requests.entries()) { + const pageRequests = requestsByPage.get(request.pageNumber) ?? []; + pageRequests.push({ index, request }); + requestsByPage.set(request.pageNumber, pageRequests); + } + + for (const [pageNumber, pageRequests] of requestsByPage) { + throwIfPdfRasterAborted(signal); + const pageSize = await readPopplerPdfPageSize({ + command: pdfInfoCommand, + inputPath, + pageNumber, + signal, + timeoutMs, + }); + const page = await renderPopplerPage({ + command, + dpi, + inputPath, + maxEncodedPageBytes, + maxPageDimension, + maxPagePixels, + pageNumber, + pageSize, + signal, + timeoutMs, + workDir, + }); + + if (!page) { + continue; + } + + const thumbnailPage = + thumbnailDpi === dpi + ? page + : await renderPopplerPage({ + command, + dpi: thumbnailDpi, + inputPath, + maxEncodedPageBytes, + maxPageDimension: proportionalThumbnailMaxDimension({ + dpi, + maxPageDimension, + thumbnailDpi, + }), + maxPagePixels, + pageNumber, + pageSize, + signal, + timeoutMs, + workDir, + }); + + for (const { index, request } of pageRequests) { + const body = await cropPopplerPage({ + boundingBox: request.boundingBox, + boundingBoxGeometry: request.boundingBoxGeometry, + dpi, + page, + maxPagePixels, + signal, + }); + + if (!body) { + continue; + } + + const thumbnailBody = thumbnailPage + ? await cropPopplerPage({ + boundingBox: request.boundingBox, + boundingBoxGeometry: request.boundingBoxGeometry, + dpi: thumbnailDpi, + page: thumbnailPage, + maxPagePixels, + signal, + }) + : null; + + const encodedBytes = body.byteLength + (thumbnailBody?.byteLength ?? 0); + + if ( + body.byteLength > maxEncodedImageBytes || + (thumbnailBody?.byteLength ?? 0) > maxEncodedImageBytes + ) { + throw new DocumentPdfRenderError( + `Document PDF rasterizer encoded image exceeds maxEncodedImageBytes=${maxEncodedImageBytes}`, + ); + } + + aggregateEncodedCropBytes += encodedBytes; + + if (aggregateEncodedCropBytes > maxEncodedCropBytes) { + throw new DocumentPdfRenderError( + `Document PDF rasterizer encoded crop output exceeds maxEncodedCropBytes=${maxEncodedCropBytes}`, + ); + } + + rendered[index] = { + body, contentType: "image/png", - metadata: { + metadata: createPopplerImageMetadata({ command, - ...(boundingBox - ? { - crop: { - boundingBox, - normalizedBoundingBox: normalizePdfRasterBoundingBoxForDpi({ - boundingBox, - dpi, - geometry: boundingBoxGeometry, - }), - ...(boundingBoxGeometry ? { geometry: boundingBoxGeometry } : {}), - }, - } - : {}), dpi, - elementId, - pageNumber, + renderedPage: page, + request, thumbnailDpi, - }, - ...(thumbnailOutputPath + }), + ...(thumbnailBody ? { variants: { [thumbnailVariantName]: { - body: new Uint8Array(await readFile(thumbnailOutputPath)), + body: thumbnailBody, contentType: "image/png", - metadata: { + metadata: createPopplerImageMetadata({ command, - ...(boundingBox - ? { - crop: { - boundingBox, - normalizedBoundingBox: normalizePdfRasterBoundingBoxForDpi({ - boundingBox, - dpi: thumbnailDpi, - geometry: boundingBoxGeometry, - }), - ...(boundingBoxGeometry ? { geometry: boundingBoxGeometry } : {}), - }, - } - : {}), dpi: thumbnailDpi, - elementId, - pageNumber, + renderedPage: thumbnailPage ?? page, + request, variant: thumbnailVariantName, - }, + }), }, }, } : {}), }; - } finally { - await rm(workDir, { force: true, recursive: true }); } - }, + } + + return rendered; + } finally { + await rm(workDir, { force: true, recursive: true }); + } +} + +async function renderPopplerPage({ + command, + dpi, + inputPath, + maxEncodedPageBytes, + maxPageDimension, + maxPagePixels, + pageNumber, + pageSize, + signal, + timeoutMs, + workDir, +}: { + readonly command: string; + readonly dpi: number; + readonly inputPath: string; + readonly maxEncodedPageBytes: number; + readonly maxPageDimension: number; + readonly maxPagePixels: number; + readonly pageNumber: number; + readonly pageSize: PopplerPdfPageSize; + readonly signal?: AbortSignal | undefined; + readonly timeoutMs: number; + readonly workDir: string; +}): Promise { + const outputName = `page-${pageNumber}-dpi-${dpi}`; + const outputPrefix = join(workDir, outputName); + const scaleTo = popplerScaleToForPage({ dpi, maxPageDimension, pageSize }); + await renderPopplerPng({ + command, + dpi, + inputPath, + outputPrefix, + pageNumber, + scaleTo, + signal, + timeoutMs, + }); + throwIfPdfRasterAborted(signal); + const outputPath = await findPopplerOutputPath(workDir, `${outputName}-`); + + if (!outputPath) { + return null; + } + + const outputStat = await stat(outputPath); + + if (outputStat.size > maxEncodedPageBytes) { + throw new DocumentPdfRenderError( + `Poppler PDF rasterizer encoded page exceeds maxEncodedPageBytes=${maxEncodedPageBytes}`, + ); + } + + const sharp = (await import("sharp")).default; + const body = await readFile(outputPath, { signal }); + throwIfPdfRasterAborted(signal); + const { data: pixels, info } = await awaitUncancellablePdfRasterOperation( + sharp(outputPath, { limitInputPixels: maxPagePixels }) + .raw() + .toBuffer({ resolveWithObject: true }), + signal, + ); + + if (info.width * info.height > maxPagePixels) { + throw new DocumentPdfRenderError( + `Poppler PDF rasterizer decoded page exceeds maxPagePixels=${maxPagePixels}`, + ); + } + + return { + body: new Uint8Array(body), + channels: info.channels, + height: info.height, + pixels: new Uint8Array(pixels), + width: info.width, + wasDownscaled: scaleTo !== undefined, }; } +async function readPopplerPdfPageSize({ + command, + inputPath, + pageNumber, + signal, + timeoutMs, +}: { + readonly command: string; + readonly inputPath: string; + readonly pageNumber: number; + readonly signal?: AbortSignal | undefined; + readonly timeoutMs: number; +}): Promise { + const { stdout } = await execFileAsync( + command, + ["-f", String(pageNumber), "-l", String(pageNumber), "-box", inputPath], + { signal, timeout: timeoutMs, windowsHide: true }, + ); + const match = /Page(?:\s+\d+)?\s+size:\s*([\d.]+)\s+x\s+([\d.]+)\s+pts/iu.exec(String(stdout)); + const widthPoints = Number(match?.[1]); + const heightPoints = Number(match?.[2]); + + if (!(widthPoints > 0) || !(heightPoints > 0)) { + throw new DocumentPdfRenderError( + `Poppler PDF rasterizer could not determine page size for pageNumber=${pageNumber}`, + ); + } + + return { heightPoints, widthPoints }; +} + +function popplerScaleToForPage({ + dpi, + maxPageDimension, + pageSize, +}: { + readonly dpi: number; + readonly maxPageDimension: number; + readonly pageSize: PopplerPdfPageSize; +}): number | undefined { + const naturalMaxDimension = (Math.max(pageSize.widthPoints, pageSize.heightPoints) * dpi) / 72; + + return naturalMaxDimension > maxPageDimension ? maxPageDimension : undefined; +} + async function renderPopplerPng({ - boundingBox, - boundingBoxGeometry, command, dpi, inputPath, outputPrefix, pageNumber, + scaleTo, + signal, timeoutMs, }: { - readonly boundingBox: DocumentMultimodalBoundingBox | undefined; - readonly boundingBoxGeometry: DocumentPdfBoundingBoxGeometry | undefined; readonly command: string; readonly dpi: number; readonly inputPath: string; readonly outputPrefix: string; readonly pageNumber: number; + readonly scaleTo: number | undefined; + readonly signal?: AbortSignal | undefined; readonly timeoutMs: number; }): Promise { const args = [ @@ -426,24 +1188,196 @@ async function renderPopplerPng({ "-png", "-r", String(dpi), - ...popplerCropArgs( - boundingBox - ? normalizePdfRasterBoundingBoxForDpi({ - boundingBox, - dpi, - geometry: boundingBoxGeometry, - }) - : undefined, - ), + ...(scaleTo === undefined ? [] : ["-scale-to", String(scaleTo)]), inputPath, outputPrefix, ]; await execFileAsync(command, args, { + signal, timeout: timeoutMs, windowsHide: true, }); } +async function cropPopplerPage({ + boundingBox, + boundingBoxGeometry, + dpi, + maxPagePixels, + page, + signal, +}: { + readonly boundingBox: DocumentMultimodalBoundingBox | undefined; + readonly boundingBoxGeometry: DocumentPdfBoundingBoxGeometry | undefined; + readonly dpi: number; + readonly maxPagePixels: number; + readonly page: PopplerRenderedPage; + readonly signal?: AbortSignal | undefined; +}): Promise { + throwIfPdfRasterAborted(signal); + + if (!boundingBox) { + return page.body; + } + + assertScaledPdfRasterCoordinatesCanBeNormalized({ boundingBoxGeometry, page }); + + const normalizedBoundingBox = normalizePdfRasterBoundingBoxForDpi({ + boundingBox, + dpi, + geometry: boundingBoxGeometry, + renderedPage: page, + }); + const crop = clampPdfRasterCropToPage(normalizedBoundingBox, page); + + if (!crop) { + return null; + } + + const sharp = (await import("sharp")).default; + const body = await awaitUncancellablePdfRasterOperation( + sharp(page.pixels, { + limitInputPixels: maxPagePixels, + raw: { + channels: page.channels, + height: page.height, + width: page.width, + }, + }) + .extract(crop) + .png() + .toBuffer(), + signal, + ); + + return new Uint8Array(body); +} + +function assertScaledPdfRasterCoordinatesCanBeNormalized({ + boundingBoxGeometry, + page, +}: { + readonly boundingBoxGeometry: DocumentPdfBoundingBoxGeometry | undefined; + readonly page: PopplerRenderedPage; +}): void { + if (!page.wasDownscaled || boundingBoxGeometry?.coordinateSystem === "relative") { + return; + } + + if (boundingBoxGeometry?.pageWidth && boundingBoxGeometry.pageHeight) { + return; + } + + throw new DocumentPdfRenderError( + "Poppler PDF rasterizer cannot normalize coordinates after page dimension capping without source page dimensions", + ); +} + +export async function awaitUncancellablePdfRasterOperation( + operation: Promise, + signal: AbortSignal | undefined, +): Promise { + let aborted = signal?.aborted ?? false; + const onAbort = () => { + aborted = true; + }; + signal?.addEventListener("abort", onAbort, { once: true }); + + try { + const result = await operation; + + if ((aborted || signal?.aborted) && signal) { + throw pdfRasterAbortReason(signal); + } + + return result; + } catch (error) { + if ((aborted || signal?.aborted) && signal) { + throw pdfRasterAbortReason(signal); + } + + throw error; + } finally { + signal?.removeEventListener("abort", onAbort); + } +} + +function proportionalThumbnailMaxDimension({ + dpi, + maxPageDimension, + thumbnailDpi, +}: { + readonly dpi: number; + readonly maxPageDimension: number; + readonly thumbnailDpi: number; +}): number { + return Math.max(1, Math.floor(maxPageDimension * Math.min(1, thumbnailDpi / dpi))); +} + +function clampPdfRasterCropToPage( + boundingBox: DocumentMultimodalBoundingBox, + page: Pick, +): { + readonly height: number; + readonly left: number; + readonly top: number; + readonly width: number; +} | null { + const requestedLeft = Math.floor(boundingBox.x); + const requestedTop = Math.floor(boundingBox.y); + const requestedRight = Math.ceil(boundingBox.x + boundingBox.width); + const requestedBottom = Math.ceil(boundingBox.y + boundingBox.height); + const left = Math.max(0, requestedLeft); + const top = Math.max(0, requestedTop); + const right = Math.min(page.width, requestedRight); + const bottom = Math.min(page.height, requestedBottom); + + if (left >= right || top >= bottom) { + return null; + } + + return { height: bottom - top, left, top, width: right - left }; +} + +function createPopplerImageMetadata({ + command, + dpi, + renderedPage, + request, + thumbnailDpi, + variant, +}: { + readonly command: string; + readonly dpi: number; + readonly renderedPage: Pick; + readonly request: RenderDocumentPdfPageRequest; + readonly thumbnailDpi?: number | undefined; + readonly variant?: string | undefined; +}): Readonly> { + return { + command, + ...(request.boundingBox + ? { + crop: { + boundingBox: request.boundingBox, + normalizedBoundingBox: normalizePdfRasterBoundingBoxForDpi({ + boundingBox: request.boundingBox, + dpi, + geometry: request.boundingBoxGeometry, + renderedPage, + }), + ...(request.boundingBoxGeometry ? { geometry: request.boundingBoxGeometry } : {}), + }, + } + : {}), + dpi, + elementId: request.elementId, + pageNumber: request.pageNumber, + ...(thumbnailDpi !== undefined ? { thumbnailDpi } : {}), + ...(variant !== undefined ? { variant } : {}), + }; +} + function pdfRasterizationCandidate(element: ParseElement): { readonly boundingBox?: DocumentMultimodalBoundingBox; readonly boundingBoxGeometry?: DocumentPdfBoundingBoxGeometry; @@ -460,11 +1394,7 @@ function pdfRasterizationCandidate(element: ParseElement): { return null; } - const existingAssetRef = isPlainObject(element.metadata.assetRef) - ? element.metadata.assetRef - : {}; - - if (typeof existingAssetRef.objectKey === "string" || typeof existingAssetRef.uri === "string") { + if (hasPdfAssetReference(element)) { return null; } @@ -482,6 +1412,16 @@ function pdfRasterizationCandidate(element: ParseElement): { }; } +function isPendingPdfVisual(element: ParseElement): boolean { + return (element.type === "image" || element.type === "table") && !hasPdfAssetReference(element); +} + +function hasPdfAssetReference(element: ParseElement): boolean { + const assetRef = isPlainObject(element.metadata.assetRef) ? element.metadata.assetRef : {}; + + return typeof assetRef.objectKey === "string" || typeof assetRef.uri === "string"; +} + function inferPdfRasterCropKind(element: ParseElement): DocumentPdfRasterCropKind { if (element.type === "page-break") { return "page"; @@ -530,19 +1470,38 @@ export function normalizePdfRasterBoundingBoxForDpi({ boundingBox, dpi, geometry, + renderedPage, }: { readonly boundingBox: DocumentMultimodalBoundingBox; readonly dpi: number; readonly geometry?: DocumentPdfBoundingBoxGeometry | undefined; + readonly renderedPage?: { readonly height: number; readonly width: number } | undefined; }): DocumentMultimodalBoundingBox { const coordinateSystem = geometry?.coordinateSystem ?? "pixel"; if (coordinateSystem === "pdf-point") { + if (geometry?.pageWidth && geometry.pageHeight && renderedPage) { + return scaleBoundingBox( + boundingBox, + renderedPage.width / geometry.pageWidth, + renderedPage.height / geometry.pageHeight, + ); + } + const scale = dpi / 72; return scaleBoundingBox(boundingBox, scale, scale); } + if (coordinateSystem === "relative" && renderedPage) { + return { + height: boundingBox.height * renderedPage.height, + width: boundingBox.width * renderedPage.width, + x: boundingBox.x * renderedPage.width, + y: boundingBox.y * renderedPage.height, + }; + } + if (coordinateSystem === "relative" && geometry?.pageWidth && geometry.pageHeight) { const scale = dpi / 72; @@ -554,6 +1513,14 @@ export function normalizePdfRasterBoundingBoxForDpi({ }; } + if (coordinateSystem === "pixel" && geometry?.pageWidth && geometry.pageHeight && renderedPage) { + return scaleBoundingBox( + boundingBox, + renderedPage.width / geometry.pageWidth, + renderedPage.height / geometry.pageHeight, + ); + } + if (coordinateSystem === "pixel" && geometry?.sourceDpi && geometry.sourceDpi !== dpi) { const scale = dpi / geometry.sourceDpi; @@ -679,7 +1646,13 @@ function parseCoordinateSystem( return "relative"; } - if (normalized === "pixel" || normalized === "pixels" || normalized === "px") { + if ( + normalized === "pixel" || + normalized === "pixels" || + normalized === "pixelspace" || + normalized === "pixel-space" || + normalized === "px" + ) { return "pixel"; } @@ -825,23 +1798,6 @@ function metadataStringFromKeys( return undefined; } -function popplerCropArgs(boundingBox: DocumentMultimodalBoundingBox | undefined): string[] { - if (!boundingBox) { - return []; - } - - return [ - "-x", - String(Math.floor(boundingBox.x)), - "-y", - String(Math.floor(boundingBox.y)), - "-W", - String(Math.ceil(boundingBox.width)), - "-H", - String(Math.ceil(boundingBox.height)), - ]; -} - async function findPopplerOutputPath(workDir: string, prefix: string): Promise { const files = await readdir(workDir); const image = files.find((file) => file.startsWith(prefix) && file.endsWith(".png")); diff --git a/knowledge-fs/packages/api/src/gateway.test.ts b/knowledge-fs/packages/api/src/gateway.test.ts index 444e7a3be2e..96f93dfa8a2 100644 --- a/knowledge-fs/packages/api/src/gateway.test.ts +++ b/knowledge-fs/packages/api/src/gateway.test.ts @@ -6,6 +6,7 @@ import { type CacheAdapter, type DatabaseExecuteInput, type DatabaseExecuteResult, + type DatabaseTransactionCallback, type EmbeddingModel, EmbeddingModelSchema, EvidenceBundleSchema, @@ -682,7 +683,13 @@ function createFakeParseArtifactExecutor() { return { rows: [], rowsAffected: 0 }; }; - return { calls, executor, rows }; + return { + calls, + executor, + rows, + transaction: async (callback: DatabaseTransactionCallback) => + callback({ execute: executor }), + }; } function createFakeKnowledgeNodeExecutor() { @@ -6350,6 +6357,7 @@ describe("createKnowledgeGateway", () => { database: createSchemaDatabaseAdapter({ executor: fake.executor, kind: "postgres", + transaction: fake.transaction, }), }); @@ -6367,9 +6375,10 @@ describe("createKnowledgeGateway", () => { tableName: "parse_artifacts", }), ); - expect(fake.calls[0]?.sql).not.toContain("hello.md"); - expect(fake.calls[0]?.params).toContain(JSON.stringify(artifact.elements)); - expect(fake.calls[0]?.params).toContain(JSON.stringify(artifact.metadata)); + const parseArtifactInsert = fake.calls.find((call) => call.operation === "insert"); + expect(parseArtifactInsert?.sql).not.toContain("hello.md"); + expect(parseArtifactInsert?.params).toContain(JSON.stringify(artifact.elements)); + expect(parseArtifactInsert?.params).toContain(JSON.stringify(artifact.metadata)); expect(fake.calls).toContainEqual( expect.objectContaining({ maxRows: 1, @@ -6384,12 +6393,14 @@ describe("createKnowledgeGateway", () => { database: createSchemaDatabaseAdapter({ executor: tidbFake.executor, kind: "tidb", + transaction: tidbFake.transaction, }), }); await expect(tidbRepository.create(artifact)).resolves.toEqual(artifact); - expect(tidbFake.calls[0]?.sql).toContain("INSERT INTO `parse_artifacts`"); - expect(tidbFake.calls[0]?.sql).toContain("CAST(? AS JSON)"); - expect(tidbFake.calls[0]?.sql).not.toContain("RETURNING"); + const tidbParseArtifactInsert = tidbFake.calls.find((call) => call.operation === "insert"); + expect(tidbParseArtifactInsert?.sql).toContain("INSERT INTO `parse_artifacts`"); + expect(tidbParseArtifactInsert?.sql).toContain("CAST(? AS JSON)"); + expect(tidbParseArtifactInsert?.sql).not.toContain("RETURNING"); const cleanupRepository = createInMemoryParseArtifactRepository({ maxArtifacts: 4, diff --git a/knowledge-fs/packages/api/src/index-reindexer.test.ts b/knowledge-fs/packages/api/src/index-reindexer.test.ts index 9fa7ddcc1a8..1b3d8a056d2 100644 --- a/knowledge-fs/packages/api/src/index-reindexer.test.ts +++ b/knowledge-fs/packages/api/src/index-reindexer.test.ts @@ -205,7 +205,7 @@ describe("incremental reindexer", () => { }); }); - it("idempotently rebuilds unchanged artifacts so partial indexes can be repaired", async () => { + it("idempotently rebuilds changed artifacts so partial indexes can be repaired", async () => { const artifacts = createInMemoryParseArtifactRepository({ maxArtifacts: 4 }); const nodes = createInMemoryKnowledgeNodeRepository({ maxBatchSize: 4, diff --git a/knowledge-fs/packages/api/src/index-reindexer.ts b/knowledge-fs/packages/api/src/index-reindexer.ts index 11da6ca6d6c..dbd545c0b19 100644 --- a/knowledge-fs/packages/api/src/index-reindexer.ts +++ b/knowledge-fs/packages/api/src/index-reindexer.ts @@ -41,6 +41,7 @@ import { preflightLlmSemanticWindows, } from "./llm-semantic-chunker"; import { + type MaterializeParseArtifactResult, type ParseArtifactLookupInput, type ParseArtifactRepository, cloneParseArtifact, @@ -111,7 +112,7 @@ export interface UpdateIncrementalReindexProjectionStatusInput { } export interface IncrementalReindexer { - canonicalizeArtifact?(artifact: ParseArtifact): Promise; + canonicalizeArtifact?(artifact: ParseArtifact): Promise; failProjections?(input: UpdateIncrementalReindexProjectionStatusInput): Promise; getCanonicalArtifact?(input: ParseArtifactLookupInput): Promise; publishProjections?(input: UpdateIncrementalReindexProjectionStatusInput): Promise; @@ -193,16 +194,18 @@ export function createIncrementalReindexer({ }; return { - canonicalizeArtifact: async (artifact: ParseArtifact) => - cloneParseArtifact( - await artifacts.create(cloneParseArtifact(ParseArtifactSchema.parse(artifact))), - ), + canonicalizeArtifact: async (artifact: ParseArtifact) => { + const materialized = await artifacts.materialize( + cloneParseArtifact(ParseArtifactSchema.parse(artifact)), + ); + return { ...materialized, artifact: cloneParseArtifact(materialized.artifact) }; + }, getCanonicalArtifact: async (input: ParseArtifactLookupInput) => { const persisted = await artifacts.getByDocumentVersion(input); - return persisted - ? cloneParseArtifact(await artifacts.create(cloneParseArtifact(persisted))) - : null; + if (!persisted) return null; + const materialized = await artifacts.materialize(cloneParseArtifact(persisted)); + return cloneParseArtifact(materialized.artifact); }, ...(canUpdateProjectionStatuses ? { diff --git a/knowledge-fs/packages/api/src/knowledge-fs-errors.ts b/knowledge-fs/packages/api/src/knowledge-fs-errors.ts index f64de2a23e6..de482828d76 100644 --- a/knowledge-fs/packages/api/src/knowledge-fs-errors.ts +++ b/knowledge-fs/packages/api/src/knowledge-fs-errors.ts @@ -786,6 +786,7 @@ function normalizeKnowledgeFsErrorCode(code: string): string { DIFY_MODEL_RUNTIME_RESPONSE_INVALID: "MODEL_RUNTIME_FAILED", DIFY_MODEL_RUNTIME_RESPONSE_TOO_LARGE: "MODEL_RUNTIME_FAILED", DIFY_MODEL_RUNTIME_TIMEOUT: "MODEL_RUNTIME_TIMEOUT", + DIFY_OBJECT_STORAGE_REQUEST_FAILED: "DOCUMENT_COMPILATION_RETRYABLE", METADATA_PREPARE_FAILED: "DOCUMENT_COMPILATION_FAILED", OBJECT_VERIFICATION_FAILED: "UPLOAD_INTEGRITY_MISMATCH", PROVIDER_INPUT: "DOCUMENT_PARSER_INPUT_INVALID", diff --git a/knowledge-fs/packages/api/src/parse-artifact-repository.test.ts b/knowledge-fs/packages/api/src/parse-artifact-repository.test.ts index e0d1b93f940..9cf38551da5 100644 --- a/knowledge-fs/packages/api/src/parse-artifact-repository.test.ts +++ b/knowledge-fs/packages/api/src/parse-artifact-repository.test.ts @@ -3,6 +3,7 @@ import { type DatabaseExecuteInput, type DatabaseExecuteResult, type DatabaseRow, + type DatabaseTransactionCallback, type ParseArtifact, ParseArtifactSchema, } from "@knowledge/core"; @@ -34,6 +35,84 @@ const artifact = ParseArtifactSchema.parse({ }) satisfies ParseArtifact; describe("parse artifact repositories", () => { + it("atomically reports created, unchanged, and replaced materialization", async () => { + const memory = createInMemoryParseArtifactRepository({ maxArtifacts: 2 }); + const fake = createFakeParseArtifactExecutor(); + const database = createDatabaseParseArtifactRepository({ + database: createSchemaDatabaseAdapter({ + executor: fake.executor, + kind: "postgres", + transaction: fake.transaction, + }), + }); + + for (const repository of [memory, database]) { + const materialized = ParseArtifactSchema.parse({ + ...artifact, + elements: [ + { + ...artifact.elements[0], + metadata: { assetRef: { objectKey: "objects/first.png" } }, + }, + ], + }); + await expect(repository.materialize(materialized)).resolves.toEqual({ + artifact: materialized, + disposition: "created", + }); + + const sameHashRetry = ParseArtifactSchema.parse({ + ...materialized, + createdAt: "2026-05-09T11:05:01.000Z", + elements: [ + { + ...materialized.elements[0], + metadata: { assetRef: { objectKey: "objects/retry-must-not-win.png" } }, + }, + ], + id: "018f0d60-7a49-7cc2-9c1b-5b36f18f2c99", + }); + await expect(repository.materialize(sameHashRetry)).resolves.toEqual({ + artifact: materialized, + disposition: "unchanged", + }); + + const changedHash = ParseArtifactSchema.parse({ + ...sameHashRetry, + artifactHash: "e".repeat(64), + elements: [ + { + ...sameHashRetry.elements[0], + metadata: { assetRef: { objectKey: "objects/replacement.png" } }, + }, + ], + }); + await expect(repository.materialize(changedHash)).resolves.toEqual({ + artifact: expect.objectContaining({ + artifactHash: changedHash.artifactHash, + createdAt: materialized.createdAt, + elements: [ + expect.objectContaining({ + metadata: { assetRef: { objectKey: "objects/replacement.png" } }, + }), + ], + id: materialized.id, + }), + disposition: "replaced", + }); + await expect( + repository.getByDocumentVersion({ + documentAssetId: materialized.documentAssetId, + version: materialized.version, + }), + ).resolves.toMatchObject({ + artifactHash: changedHash.artifactHash, + elements: [{ metadata: { assetRef: { objectKey: "objects/replacement.png" } } }], + id: materialized.id, + }); + } + }); + it("keeps generated element ids bound to the first persisted artifact on retry", async () => { const retryArtifactId = "018f0d60-7a49-7cc2-9c1b-5b36f18f2c46"; const first = ParseArtifactSchema.parse({ @@ -59,7 +138,11 @@ describe("parse artifact repositories", () => { const memory = createInMemoryParseArtifactRepository({ maxArtifacts: 2 }); const fake = createFakeParseArtifactExecutor(); const database = createDatabaseParseArtifactRepository({ - database: createSchemaDatabaseAdapter({ executor: fake.executor, kind: "postgres" }), + database: createSchemaDatabaseAdapter({ + executor: fake.executor, + kind: "postgres", + transaction: fake.transaction, + }), }); for (const repository of [memory, database]) { @@ -83,7 +166,7 @@ describe("parse artifact repositories", () => { } }); - it("fails closed when a concurrent row change prevents generated id repair", async () => { + it("does not rewrite a canonical same-hash row when the retry uses new generated ids", async () => { const retryArtifactId = "018f0d60-7a49-7cc2-9c1b-5b36f18f2c46"; const first = ParseArtifactSchema.parse({ ...artifact, @@ -91,7 +174,11 @@ describe("parse artifact repositories", () => { }); const fake = createFakeParseArtifactExecutor({ failUpdates: true }); const repository = createDatabaseParseArtifactRepository({ - database: createSchemaDatabaseAdapter({ executor: fake.executor, kind: "postgres" }), + database: createSchemaDatabaseAdapter({ + executor: fake.executor, + kind: "postgres", + transaction: fake.transaction, + }), }); await repository.create(first); @@ -101,9 +188,14 @@ describe("parse artifact repositories", () => { id: retryArtifactId, }); - await expect(repository.create(retry)).rejects.toThrow( - "generated element ids could not be canonicalized", - ); + await expect(repository.materialize(retry)).resolves.toMatchObject({ + artifact: { + elements: [{ id: `${first.id}:element-1` }], + id: first.id, + }, + disposition: "unchanged", + }); + expect(fake.calls.filter((call) => call.operation === "update")).toEqual([]); }); it("stores clone-isolated artifacts and bounds in-memory capacity", async () => { @@ -150,7 +242,11 @@ describe("parse artifact repositories", () => { it("uses parameterized database writes and bounded deletes", async () => { const fake = createFakeParseArtifactExecutor(); const repository = createDatabaseParseArtifactRepository({ - database: createSchemaDatabaseAdapter({ executor: fake.executor, kind: "postgres" }), + database: createSchemaDatabaseAdapter({ + executor: fake.executor, + kind: "postgres", + transaction: fake.transaction, + }), }); await expect(repository.create(artifact)).resolves.toEqual(artifact); @@ -169,16 +265,26 @@ describe("parse artifact repositories", () => { }), ).resolves.toBe(0); - expect(fake.calls[0]).toEqual( + const insert = fake.calls.find((call) => call.operation === "insert"); + expect(insert).toEqual( expect.objectContaining({ maxRows: 1, operation: "insert", tableName: "parse_artifacts", }), ); - expect(fake.calls[0]?.sql).not.toContain("hello.md"); - expect(fake.calls[0]?.params).toContain(JSON.stringify(artifact.elements)); - expect(fake.calls[0]?.params).toContain(JSON.stringify(artifact.metadata)); + expect(insert?.sql).not.toContain("hello.md"); + expect(insert?.params).toContain(JSON.stringify(artifact.elements)); + expect(insert?.params).toContain(JSON.stringify(artifact.metadata)); + expect(fake.calls[0]).toEqual( + expect.objectContaining({ + operation: "select", + params: [artifact.documentAssetId], + tableName: "document_assets", + }), + ); + expect(fake.calls[0]?.sql).toContain("FOR UPDATE"); + expect(fake.calls[1]?.sql).toContain("FOR UPDATE"); expect(fake.calls).toContainEqual( expect.objectContaining({ maxRows: 1, @@ -198,24 +304,26 @@ describe("parse artifact repositories", () => { expect(fake.calls.at(-1)?.sql).not.toContain(artifact.documentAssetId); }); - it("fails closed when an upsert cannot resolve exactly one canonical artifact", async () => { + it("fails closed when a locked lookup resolves duplicate or mismatched artifacts", async () => { const row = parseArtifactRow(artifact); - const repositoryForRows = (rows: readonly DatabaseRow[]) => - createDatabaseParseArtifactRepository({ - database: createSchemaDatabaseAdapter({ - executor: async (input) => ({ - rows: input.operation === "select" ? [...rows] : [], - rowsAffected: input.operation === "insert" ? 1 : 0, - }), - kind: "tidb", - }), + const repositoryForRows = (rows: readonly DatabaseRow[]) => { + const executor = async (input: DatabaseExecuteInput): Promise => ({ + rows: + input.operation === "select" && input.tableName === "parse_artifacts" ? [...rows] : [], + rowsAffected: input.operation === "insert" ? 1 : 0, }); - await expect(repositoryForRows([]).create(artifact)).rejects.toThrow( - "Parse artifact upsert did not persist its logical row", - ); - await expect(repositoryForRows([row, { ...row }]).create(artifact)).rejects.toThrow( - "Parse artifact upsert resolved multiple persisted logical rows", + return createDatabaseParseArtifactRepository({ + database: createSchemaDatabaseAdapter({ + executor, + kind: "tidb", + transaction: async (callback) => callback({ execute: executor }), + }), + }); + }; + + await expect(repositoryForRows([row, { ...row }]).materialize(artifact)).rejects.toThrow( + "materialization resolved multiple persisted logical rows", ); await expect( repositoryForRows([ @@ -223,8 +331,8 @@ describe("parse artifact repositories", () => { ...row, document_asset_id: "018f0d60-7a49-7cc2-9c1b-5b36f18f2cff", }, - ]).create(artifact), - ).rejects.toThrow("Parse artifact upsert resolved a mismatched persisted logical row"); + ]).materialize(artifact), + ).rejects.toThrow("materialization resolved a mismatched persisted row"); }); it("guards memory prune overflow and database document deletes", async () => { @@ -250,7 +358,11 @@ describe("parse artifact repositories", () => { const fake = createFakeParseArtifactExecutor(); const database = createDatabaseParseArtifactRepository({ - database: createSchemaDatabaseAdapter({ executor: fake.executor, kind: "postgres" }), + database: createSchemaDatabaseAdapter({ + executor: fake.executor, + kind: "postgres", + transaction: fake.transaction, + }), }); await expect( @@ -342,6 +454,35 @@ function createFakeParseArtifactExecutor({ if (failUpdates) { return { rows: [], rowsAffected: 0 }; } + + if (input.params.length === 8) { + const [ + parser, + contentType, + artifactHash, + elements, + metadata, + id, + documentAssetId, + version, + ] = input.params; + const key = `${String(documentAssetId)}:${Number(version)}`; + const row = rows.get(key); + if (!row || row.id !== String(id)) { + return { rows: [], rowsAffected: 0 }; + } + rows.set(key, { + ...row, + artifact_hash: String(artifactHash), + content_type: String(contentType), + elements: typeof elements === "string" ? JSON.parse(elements) : elements, + metadata: typeof metadata === "string" ? JSON.parse(metadata) : metadata, + parser: String(parser), + }); + + return { rows: [], rowsAffected: 1 }; + } + const [elements, id, documentAssetId, version] = input.params; const key = `${String(documentAssetId)}:${Number(version)}`; const row = rows.get(key); @@ -359,7 +500,12 @@ function createFakeParseArtifactExecutor({ return { rows: [], rowsAffected: 0 }; }; - return { calls, executor }; + return { + calls, + executor, + transaction: async (callback: DatabaseTransactionCallback) => + callback({ execute: executor }), + }; } function parseArtifactRow(input: ParseArtifact): DatabaseRow { diff --git a/knowledge-fs/packages/api/src/parse-artifact-repository.ts b/knowledge-fs/packages/api/src/parse-artifact-repository.ts index 84df2229b74..d5bd0103837 100644 --- a/knowledge-fs/packages/api/src/parse-artifact-repository.ts +++ b/knowledge-fs/packages/api/src/parse-artifact-repository.ts @@ -1,5 +1,6 @@ import { type DatabaseAdapter, + type DatabaseExecutor, type DatabaseQueryValue, type DatabaseRow, type ParseArtifact, @@ -28,9 +29,17 @@ export interface ParseArtifactRepository { deleteByDocumentAsset(input: DeleteParseArtifactsByDocumentAssetInput): Promise; getById(input: ParseArtifactIdLookupInput): Promise; getByDocumentVersion(input: ParseArtifactLookupInput): Promise; + materialize(input: ParseArtifact): Promise; pruneDocumentVersions(input: PruneParseArtifactVersionsInput): Promise; } +export type ParseArtifactMaterializationDisposition = "created" | "replaced" | "unchanged"; + +export interface MaterializeParseArtifactResult { + readonly artifact: ParseArtifact; + readonly disposition: ParseArtifactMaterializationDisposition; +} + export interface DeleteParseArtifactsByDocumentAssetInput { readonly documentAssetId: string; readonly maxArtifacts: number; @@ -129,24 +138,40 @@ export function createInMemoryParseArtifactRepository({ const artifacts = new Map(); - return { - create: async (input) => { - const artifact = cloneParseArtifact(ParseArtifactSchema.parse(input)); - const key = parseArtifactKey(artifact.documentAssetId, artifact.version); - const existing = artifacts.get(key); + const materialize = async (input: ParseArtifact): Promise => { + const artifact = cloneParseArtifact(ParseArtifactSchema.parse(input)); + const key = parseArtifactKey(artifact.documentAssetId, artifact.version); + const existing = artifacts.get(key); - if (!existing && artifacts.size >= maxArtifacts) { - throw new ParseArtifactCapacityExceededError(maxArtifacts); - } + if (existing?.artifactHash === artifact.artifactHash) { + return { artifact: cloneParseArtifact(existing), disposition: "unchanged" }; + } + if (existing) { const stored = bindGeneratedElementIdsToArtifact( - existing ? { ...artifact, createdAt: existing.createdAt } : artifact, - existing?.id ?? artifact.id, + { ...artifact, createdAt: existing.createdAt }, + existing.id, ); artifacts.set(key, stored); - return cloneParseArtifact(stored); - }, + return { artifact: cloneParseArtifact(stored), disposition: "replaced" }; + } + + if (artifacts.size >= maxArtifacts) { + throw new ParseArtifactCapacityExceededError(maxArtifacts); + } + + const stored = bindGeneratedElementIdsToArtifact(artifact, artifact.id); + artifacts.set(key, stored); + + return { + artifact: cloneParseArtifact(stored), + disposition: "created", + }; + }; + + return { + create: async (input) => (await materialize(input)).artifact, getByDocumentVersion: async ({ documentAssetId, version }) => { const artifact = artifacts.get(parseArtifactKey(documentAssetId, version)); @@ -157,6 +182,7 @@ export function createInMemoryParseArtifactRepository({ return artifact ? cloneParseArtifact(artifact) : null; }, + materialize, deleteByDocumentAsset: async ({ documentAssetId, maxArtifacts }) => { if (!Number.isInteger(maxArtifacts) || maxArtifacts < 1) { throw new Error("Parse artifact delete maxArtifacts must be at least 1"); @@ -201,13 +227,16 @@ export function createDatabaseParseArtifactRepository({ database, }: DatabaseParseArtifactRepositoryOptions): ParseArtifactRepository { const tableName = "parse_artifacts"; - const repairGeneratedElementIds = async (persisted: ParseArtifact): Promise => { + const repairGeneratedElementIds = async ( + executor: DatabaseExecutor, + persisted: ParseArtifact, + ): Promise => { const canonical = bindGeneratedElementIdsToArtifact(persisted, persisted.id); if (JSON.stringify(canonical.elements) === JSON.stringify(persisted.elements)) { return canonical; } - const repaired = await database.execute({ + const repaired = await executor.execute({ maxRows: 1, operation: "update", params: [ @@ -242,80 +271,24 @@ export function createDatabaseParseArtifactRepository({ return canonical; }; - return { - create: async (input) => { - const artifact = ParseArtifactSchema.parse(input); - const elements = JSON.stringify(artifact.elements); - const metadata = JSON.stringify(artifact.metadata); - const params = [ - artifact.id, - artifact.documentAssetId, - artifact.version, - artifact.parser, - artifact.contentType, - artifact.artifactHash, - elements, - metadata, - artifact.createdAt, - ] satisfies readonly DatabaseQueryValue[]; - const columns = [ - "id", - "document_asset_id", - "version", - "parser", - "content_type", - "artifact_hash", - "elements", - "metadata", - "created_at", - ]; - const mutableColumns = columns.filter( - (column) => - column !== "id" && - column !== "document_asset_id" && - column !== "version" && - column !== "created_at", - ); - const upsertClause = - database.dialect === "postgres" - ? ` ON CONFLICT (${quoteDatabaseIdentifier( - database, - "document_asset_id", - )}, ${quoteDatabaseIdentifier(database, "version")}) DO UPDATE SET ${mutableColumns - .map( - (column) => - `${quoteDatabaseIdentifier(database, column)} = EXCLUDED.${quoteDatabaseIdentifier( - database, - column, - )}`, - ) - .join(", ")} RETURNING *` - : ` ON DUPLICATE KEY UPDATE ${mutableColumns - .map( - (column) => - `${quoteDatabaseIdentifier(database, column)} = VALUES(${quoteDatabaseIdentifier( - database, - column, - )})`, - ) - .join(", ")}`; - const result = await database.execute({ + const materialize = async (input: ParseArtifact): Promise => { + const artifact = ParseArtifactSchema.parse(input); + + return database.transaction(async (executor) => { + await executor.execute({ maxRows: 1, - operation: "insert", - params, - sql: `INSERT INTO ${quoteDatabaseIdentifier(database, tableName)} (${columns - .map((column) => quoteDatabaseIdentifier(database, column)) - .join(", ")}) VALUES (${params - .map((_, index) => jsonInsertPlaceholder(database, index + 1, columns[index])) - .join(", ")})${upsertClause};`, - tableName, + operation: "select", + params: [artifact.documentAssetId], + sql: `SELECT ${quoteDatabaseIdentifier(database, "id")} FROM ${quoteDatabaseIdentifier( + database, + "document_assets", + )} WHERE ${quoteDatabaseIdentifier(database, "id")} = ${databasePlaceholder( + database, + 1, + )} LIMIT 1 FOR UPDATE;`, + tableName: "document_assets", }); - - if (result.rows[0]) { - return repairGeneratedElementIds(mapParseArtifactRow(result.rows[0])); - } - - const stored = await database.execute({ + const locked = await executor.execute({ maxRows: 2, operation: "select", params: [artifact.documentAssetId, artifact.version], @@ -328,28 +301,133 @@ export function createDatabaseParseArtifactRepository({ )} = ${databasePlaceholder(database, 1)} AND ${quoteDatabaseIdentifier( database, "version", - )} = ${databasePlaceholder(database, 2)} LIMIT 2;`, + )} = ${databasePlaceholder(database, 2)} LIMIT 2 FOR UPDATE;`, + tableName, + }); + const [row, duplicate] = locked.rows; + + if (duplicate) { + throw new Error("Parse artifact materialization resolved multiple persisted logical rows"); + } + + if (row) { + const persisted = mapParseArtifactRow(row); + + if ( + persisted.documentAssetId !== artifact.documentAssetId || + persisted.version !== artifact.version + ) { + throw new Error("Parse artifact materialization resolved a mismatched persisted row"); + } + + if (persisted.artifactHash === artifact.artifactHash) { + return { + artifact: await repairGeneratedElementIds(executor, persisted), + disposition: "unchanged" as const, + }; + } + + const replaced = bindGeneratedElementIdsToArtifact( + { ...artifact, createdAt: persisted.createdAt }, + persisted.id, + ); + const updated = await executor.execute({ + maxRows: 1, + operation: "update", + params: [ + replaced.parser, + replaced.contentType, + replaced.artifactHash, + JSON.stringify(replaced.elements), + JSON.stringify(replaced.metadata), + replaced.id, + replaced.documentAssetId, + replaced.version, + ], + sql: `UPDATE ${quoteDatabaseIdentifier(database, tableName)} SET ${quoteDatabaseIdentifier( + database, + "parser", + )} = ${databasePlaceholder(database, 1)}, ${quoteDatabaseIdentifier( + database, + "content_type", + )} = ${databasePlaceholder(database, 2)}, ${quoteDatabaseIdentifier( + database, + "artifact_hash", + )} = ${databasePlaceholder(database, 3)}, ${quoteDatabaseIdentifier( + database, + "elements", + )} = ${jsonInsertPlaceholder( + database, + 4, + "elements", + )}, ${quoteDatabaseIdentifier(database, "metadata")} = ${jsonInsertPlaceholder( + database, + 5, + "metadata", + )} WHERE ${quoteDatabaseIdentifier(database, "id")} = ${databasePlaceholder( + database, + 6, + )} AND ${quoteDatabaseIdentifier( + database, + "document_asset_id", + )} = ${databasePlaceholder(database, 7)} AND ${quoteDatabaseIdentifier( + database, + "version", + )} = ${databasePlaceholder(database, 8)};`, + tableName, + }); + if (updated.rowsAffected !== 1) { + throw new Error("Parse artifact materialization could not replace its logical row"); + } + + return { artifact: replaced, disposition: "replaced" as const }; + } + + const created = bindGeneratedElementIdsToArtifact(artifact, artifact.id); + const columns = [ + "id", + "document_asset_id", + "version", + "parser", + "content_type", + "artifact_hash", + "elements", + "metadata", + "created_at", + ]; + const params = [ + created.id, + created.documentAssetId, + created.version, + created.parser, + created.contentType, + created.artifactHash, + JSON.stringify(created.elements), + JSON.stringify(created.metadata), + created.createdAt, + ] satisfies readonly DatabaseQueryValue[]; + const inserted = await executor.execute({ + maxRows: 1, + operation: "insert", + params, + sql: `INSERT INTO ${quoteDatabaseIdentifier(database, tableName)} (${columns + .map((column) => quoteDatabaseIdentifier(database, column)) + .join(", ")}) VALUES (${params + .map((_, index) => jsonInsertPlaceholder(database, index + 1, columns[index])) + .join(", ")});`, tableName, }); - const [row, duplicate] = stored.rows; - if (!row) { - throw new Error("Parse artifact upsert did not persist its logical row"); - } - if (duplicate) { - throw new Error("Parse artifact upsert resolved multiple persisted logical rows"); + if (inserted.rowsAffected !== 1) { + throw new Error("Parse artifact materialization did not create its logical row"); } - const persisted = mapParseArtifactRow(row); - if ( - persisted.documentAssetId !== artifact.documentAssetId || - persisted.version !== artifact.version - ) { - throw new Error("Parse artifact upsert resolved a mismatched persisted logical row"); - } + return { artifact: created, disposition: "created" as const }; + }); + }; - return repairGeneratedElementIds(persisted); - }, + return { + create: async (input) => (await materialize(input)).artifact, getByDocumentVersion: async ({ documentAssetId, version }) => { const result = await database.execute({ maxRows: 1, @@ -381,6 +459,7 @@ export function createDatabaseParseArtifactRepository({ return result.rows[0] ? mapParseArtifactRow(result.rows[0]) : null; }, + materialize, deleteByDocumentAsset: async ({ documentAssetId, maxArtifacts }) => { if (!Number.isInteger(maxArtifacts) || maxArtifacts < 1) { throw new Error("Parse artifact delete maxArtifacts must be at least 1"); diff --git a/knowledge-fs/packages/api/src/storage-path-utils.test.ts b/knowledge-fs/packages/api/src/storage-path-utils.test.ts index 117eefffc25..538cb2a5b78 100644 --- a/knowledge-fs/packages/api/src/storage-path-utils.test.ts +++ b/knowledge-fs/packages/api/src/storage-path-utils.test.ts @@ -1,6 +1,8 @@ import { describe, expect, it } from "vitest"; import { + createDocumentMultimodalAssetObjectKey, + createDocumentMultimodalAssetVariantObjectKey, createDocumentObjectKey, normalizeSourceFsPath, sanitizeFilename, @@ -65,4 +67,34 @@ describe("storage path utilities", () => { }), ).toBe("tenant-1/spaces/space-1/documents/asset-1/q2-report-final-.pdf"); }); + + it("isolates multimodal writes under an execution-owned namespace", () => { + expect( + createDocumentMultimodalAssetObjectKey({ + assetId: "asset-1", + contentType: "image/png", + elementId: "artifact:element-1", + knowledgeSpaceId: "space-1", + sha256: "a".repeat(64), + tenantId: "tenant-1", + writeOwnerId: "Attempt/One", + }), + ).toBe( + "tenant-1/spaces/space-1/documents/asset-1/assets/one/artifact-element-1-aaaaaaaaaaaa.png", + ); + expect( + createDocumentMultimodalAssetVariantObjectKey({ + assetId: "asset-1", + contentType: "image/png", + elementId: "artifact:element-1", + knowledgeSpaceId: "space-1", + sha256: "b".repeat(64), + tenantId: "tenant-1", + variant: "Thumb Nail", + writeOwnerId: "Attempt/One", + }), + ).toBe( + "tenant-1/spaces/space-1/documents/asset-1/assets/one/artifact-element-1-thumb-nail-bbbbbbbbbbbb.png", + ); + }); }); diff --git a/knowledge-fs/packages/api/src/storage-path-utils.ts b/knowledge-fs/packages/api/src/storage-path-utils.ts index e39d32468f5..3d4faf5f668 100644 --- a/knowledge-fs/packages/api/src/storage-path-utils.ts +++ b/knowledge-fs/packages/api/src/storage-path-utils.ts @@ -104,6 +104,7 @@ export function createDocumentMultimodalAssetObjectKey({ knowledgeSpaceId, sha256, tenantId, + writeOwnerId, }: { readonly assetId: string; readonly contentType: string; @@ -111,11 +112,13 @@ export function createDocumentMultimodalAssetObjectKey({ readonly knowledgeSpaceId: string; readonly sha256: string; readonly tenantId: string; + readonly writeOwnerId?: string | undefined; }): string { const extension = imageExtension(contentType); const safeElementId = sanitizeFilename(elementId).replace(/\.[a-z0-9]+$/u, "") || "asset"; + const ownerPrefix = multimodalWriteOwnerPrefix(writeOwnerId); - return `${tenantId}/spaces/${knowledgeSpaceId}/documents/${assetId}/assets/${safeElementId}-${sha256.slice(0, 12)}.${extension}`; + return `${tenantId}/spaces/${knowledgeSpaceId}/documents/${assetId}/assets/${ownerPrefix}${safeElementId}-${sha256.slice(0, 12)}.${extension}`; } export function createDocumentMultimodalAssetVariantObjectKey({ @@ -126,6 +129,7 @@ export function createDocumentMultimodalAssetVariantObjectKey({ sha256, tenantId, variant, + writeOwnerId, }: { readonly assetId: string; readonly contentType: string; @@ -134,12 +138,24 @@ export function createDocumentMultimodalAssetVariantObjectKey({ readonly sha256: string; readonly tenantId: string; readonly variant: string; + readonly writeOwnerId?: string | undefined; }): string { const extension = imageExtension(contentType); const safeElementId = sanitizeFilename(elementId).replace(/\.[a-z0-9]+$/u, "") || "asset"; const safeVariant = sanitizeFilename(variant).replace(/\.[a-z0-9]+$/u, "") || "variant"; + const ownerPrefix = multimodalWriteOwnerPrefix(writeOwnerId); - return `${tenantId}/spaces/${knowledgeSpaceId}/documents/${assetId}/assets/${safeElementId}-${safeVariant}-${sha256.slice(0, 12)}.${extension}`; + return `${tenantId}/spaces/${knowledgeSpaceId}/documents/${assetId}/assets/${ownerPrefix}${safeElementId}-${safeVariant}-${sha256.slice(0, 12)}.${extension}`; +} + +function multimodalWriteOwnerPrefix(writeOwnerId: string | undefined): string { + if (!writeOwnerId) { + return ""; + } + + const safeOwnerId = sanitizeFilename(writeOwnerId).replace(/\.[a-z0-9]+$/u, ""); + + return `${safeOwnerId || "write"}/`; } export function sanitizeFilename(filename: string): string { diff --git a/knowledge-fs/packages/parsers/src/index.ts b/knowledge-fs/packages/parsers/src/index.ts index c5d453d6bbd..c95e498228c 100644 --- a/knowledge-fs/packages/parsers/src/index.ts +++ b/knowledge-fs/packages/parsers/src/index.ts @@ -28,6 +28,8 @@ export interface ParseDocumentInput { } export interface ParserRouteHints { + /** A downstream handler can materialize PDF images from provider-supplied coordinates. */ + readonly imagesHandledExternally?: boolean; readonly language?: string; readonly layoutComplexity?: "complex" | "simple"; /** Request provider-side image extraction when no cheaper local extractor exists. */ @@ -49,6 +51,7 @@ export type ProviderErrorCode = export class ProviderError extends Error { readonly code: ProviderErrorCode; + readonly retryable: boolean; readonly status?: number; constructor( @@ -56,16 +59,19 @@ export class ProviderError extends Error { { cause, code, + retryable = false, status, }: { readonly cause?: unknown; readonly code: ProviderErrorCode; + readonly retryable?: boolean; readonly status?: number; }, ) { super(message, cause === undefined ? undefined : { cause }); this.name = "ProviderError"; this.code = code; + this.retryable = retryable; if (status !== undefined) { this.status = status; } @@ -84,7 +90,7 @@ export class ProviderRateLimitError extends ProviderError { message: string, options: { readonly cause?: unknown; readonly status?: number } = {}, ) { - super(message, { ...options, code: "provider_rate_limited" }); + super(message, { ...options, code: "provider_rate_limited", retryable: true }); this.name = "ProviderRateLimitError"; } } @@ -92,9 +98,19 @@ export class ProviderRateLimitError extends ProviderError { export class ProviderRequestError extends ProviderError { constructor( message: string, - options: { readonly cause?: unknown; readonly status?: number } = {}, + options: { + readonly cause?: unknown; + readonly retryable?: boolean; + readonly status?: number; + } = {}, ) { - super(message, { ...options, code: "provider_request_failed" }); + super(message, { + ...options, + code: "provider_request_failed", + retryable: + options.retryable ?? + (options.status !== undefined && isRetryableProviderStatus(options.status)), + }); this.name = "ProviderRequestError"; } } @@ -122,7 +138,9 @@ export interface UnstructuredParserClientOptions extends NativeParserOptions { readonly endpoint: string; readonly fetch?: typeof fetch; readonly maxResponseBytes?: number; + readonly maxConcurrency?: number; readonly maxRetries?: number; + readonly requestTimeoutMs?: number; readonly retryDelayMs?: number; readonly sleep?: (ms: number) => Promise; } @@ -166,7 +184,9 @@ const archivePathCollator = new Intl.Collator("en", { numeric: true, sensitivity // Image blocks are returned as base64 in the partition JSON. Keep the response bounded while // leaving enough headroom for the encoded images of ordinary PDF, Office, and presentation files. const defaultMaxResponseBytes = 32 * 1024 * 1024; +const defaultMaxConcurrency = 2; const defaultMaxRetries = 0; +const defaultRequestTimeoutMs = 120_000; const defaultMaxRows = 20_000; const defaultRetryDelayMs = 100; const defaultNow = () => new Date().toISOString(); @@ -265,104 +285,195 @@ export function createUnstructuredParserClient({ apiKey, endpoint, fetch: fetchImpl = fetch, + maxConcurrency = defaultMaxConcurrency, maxResponseBytes = defaultMaxResponseBytes, maxRetries = defaultMaxRetries, + requestTimeoutMs = defaultRequestTimeoutMs, retryDelayMs = defaultRetryDelayMs, sleep = sleepMs, ...options }: UnstructuredParserClientOptions): ParserAdapter { validateRetryOptions({ maxRetries, retryDelayMs }); + validateUnstructuredResourceOptions({ maxConcurrency, requestTimeoutMs }); + const requestGate = createAbortAwareConcurrencyGate(maxConcurrency); return { kind: "unstructured", - parse: async (input) => { - const parserVersion = options.parserVersion ?? "unstructured@4"; - assertInputBounds(input.body, options.maxInputBytes ?? defaultMaxInputBytes); - const response = await fetchWithRetries({ - buildRequest: () => { - const form = new FormData(); - const fileBody = input.body.buffer.slice( - input.body.byteOffset, - input.body.byteOffset + input.body.byteLength, - ) as ArrayBuffer; - form.set("files", new File([fileBody], input.filename, { type: input.mimeType })); - form.set("coordinates", "true"); - form.set("strategy", unstructuredPartitionStrategy(input)); - if (shouldRequestProviderImages(input)) { - form.append("extract_image_block_types", "Image"); - form.set("extract_image_block_to_payload", "true"); + parse: async (input) => + requestGate.run(async () => { + const deadline = createUnstructuredRequestDeadline(input.signal, requestTimeoutMs); + try { + const parserVersion = options.parserVersion ?? "unstructured@5"; + const partitionStrategy = unstructuredPartitionStrategy(input); + const providerImageBlockTypes = unstructuredProviderImageBlockTypes(input); + assertInputBounds(input.body, options.maxInputBytes ?? defaultMaxInputBytes); + const response = await fetchWithRetries({ + buildRequest: () => { + const form = new FormData(); + const fileBody = input.body.buffer.slice( + input.body.byteOffset, + input.body.byteOffset + input.body.byteLength, + ) as ArrayBuffer; + form.set("files", new File([fileBody], input.filename, { type: input.mimeType })); + form.set("coordinates", "true"); + form.set("strategy", partitionStrategy); + if (providerImageBlockTypes.length > 0) { + for (const blockType of providerImageBlockTypes) { + form.append("extract_image_block_types", blockType); + } + form.set("extract_image_block_to_payload", "true"); + } + + return new Request(unstructuredPartitionEndpoint(endpoint), { + body: form, + method: "POST", + ...(apiKey ? { headers: { authorization: `Bearer ${apiKey}` } } : {}), + signal: deadline.signal, + }); + }, + fetchImpl, + maxRetries, + retryDelayMs, + sleep, + signal: deadline.signal, + }); + + if (!response.ok) { + throw providerRequestError("Unstructured parser", response.status); } - return new Request(unstructuredPartitionEndpoint(endpoint), { - body: form, - method: "POST", - ...(apiKey ? { headers: { authorization: `Bearer ${apiKey}` } } : {}), - ...(input.signal ? { signal: input.signal } : {}), + const responseText = await boundedResponseText(response, maxResponseBytes); + let payload: unknown; + + try { + payload = JSON.parse(responseText); + } catch (error) { + throw new ProviderResponseError("Unstructured parser returned an invalid response", { + cause: error, + }); + } + + const parsed = UnstructuredResponseSchema.safeParse(payload); + + if (!parsed.success) { + throw new ProviderResponseError("Unstructured parser returned an invalid response"); + } + + const providerElements = unstructuredElementsToElements( + normalizeUnstructuredLayout(parsed.data), + ); + const elements = appendArchiveMediaFallbackElements(input, providerElements); + + const artifact = await createParseArtifact({ + artifactHashContext: unstructuredArtifactHashContext(input, { + partitionStrategy, + providerImageBlockTypes, + }), + elements, + input, + kind: "unstructured", + options, + parserVersion, }); - }, - fetchImpl, - maxRetries, - retryDelayMs, - sleep, - }); - - if (!response.ok) { - throw providerRequestError("Unstructured parser", response.status); - } - - const responseText = await boundedResponseText(response, maxResponseBytes); - let payload: unknown; - - try { - payload = JSON.parse(responseText); - } catch (error) { - throw new ProviderResponseError("Unstructured parser returned an invalid response", { - cause: error, - }); - } - - const parsed = UnstructuredResponseSchema.safeParse(payload); - - if (!parsed.success) { - throw new ProviderResponseError("Unstructured parser returned an invalid response"); - } - - const providerElements = unstructuredElementsToElements( - normalizeUnstructuredLayout(parsed.data), - ); - const elements = appendArchiveMediaFallbackElements(input, providerElements); - - return createParseArtifact({ - elements, - input, - kind: "unstructured", - options, - parserVersion, - }); - }, + deadline.throwIfExpired(); + return artifact; + } catch (error) { + if (deadline.expired()) { + throw new ProviderRequestError( + `Unstructured parser request timed out after requestTimeoutMs=${requestTimeoutMs}`, + { cause: error, retryable: true }, + ); + } + if (input.signal?.aborted) { + throw abortSignalReason(input.signal); + } + throw error; + } finally { + deadline.dispose(); + } + }, input.signal), }; } function unstructuredPartitionStrategy(input: ParseDocumentInput): "auto" | "fast" | "hi_res" { const hints = input.parserHints; - if (hints?.requiresOcr || hints?.layoutComplexity === "complex" || hints?.requiresTables) { + if ( + hints?.requiresOcr || + hints?.layoutComplexity === "complex" || + hints?.requiresTables || + shouldRequestProviderImages(input) || + (hints?.requiresImages === true && isPdf(input)) + ) { return "hi_res"; } if (hints?.layoutComplexity === "simple") { return "fast"; } - if (hints?.requiresImages && !imagesHandledOutsideUnstructured(input)) { - return "hi_res"; - } return "auto"; } function shouldRequestProviderImages(input: ParseDocumentInput): boolean { - return input.parserHints?.requiresImages === true && !imagesHandledOutsideUnstructured(input); + return unstructuredProviderImageBlockTypes(input).length > 0; } -function imagesHandledOutsideUnstructured(input: ParseDocumentInput): boolean { - return input.mimeType.toLowerCase() === "application/pdf" || archiveMediaRoots(input) !== null; +function unstructuredProviderImageBlockTypes( + input: ParseDocumentInput, +): readonly ("Image" | "Table")[] { + if ( + input.parserHints?.requiresImages !== true || + providerImagesHandledOutsideUnstructured(input) + ) { + return []; + } + + // PDF fallback mirrors the local rasterizer, which materializes both figures and tables. Other + // formats retain the narrower historical Image-only request to avoid increasing payload sizes. + return isPdf(input) ? ["Image", "Table"] : ["Image"]; +} + +function providerImagesHandledOutsideUnstructured(input: ParseDocumentInput): boolean { + if (archiveMediaRoots(input) !== null) { + return true; + } + + return isPdf(input) && input.parserHints?.imagesHandledExternally === true; +} + +function isPdf(input: ParseDocumentInput): boolean { + return normalizedMimeType(input.mimeType) === "application/pdf"; +} + +function normalizedMimeType(value: string): string { + return value.split(";", 1)[0]?.trim().toLowerCase() ?? ""; +} + +function unstructuredArtifactHashContext( + input: ParseDocumentInput, + request: { + readonly partitionStrategy: "auto" | "fast" | "hi_res"; + readonly providerImageBlockTypes: readonly ("Image" | "Table")[]; + }, +): string { + const hints = input.parserHints; + + return JSON.stringify({ + filename: input.filename, + mimeType: input.mimeType.trim().toLowerCase(), + parserHints: { + imagesHandledExternally: hints?.imagesHandledExternally === true, + language: hints?.language?.trim().toLowerCase() || null, + layoutComplexity: hints?.layoutComplexity ?? null, + requiresImages: hints?.requiresImages === true, + requiresOcr: hints?.requiresOcr === true, + requiresTables: hints?.requiresTables === true, + }, + request: { + coordinates: true, + imageBlockTypes: request.providerImageBlockTypes, + imagePayload: request.providerImageBlockTypes.length > 0, + strategy: request.partitionStrategy, + }, + }); } function unstructuredPartitionEndpoint(endpoint: string): string { @@ -667,6 +778,7 @@ function selectParser( } async function createParseArtifact({ + artifactHashContext, artifactMetadata, elements, input, @@ -674,6 +786,7 @@ async function createParseArtifact({ options, parserVersion, }: { + readonly artifactHashContext?: string | undefined; readonly artifactMetadata?: Readonly> | undefined; readonly elements: readonly ParseElementInput[]; readonly input: ParseDocumentInput; @@ -698,7 +811,7 @@ async function createParseArtifact({ ); return ParseArtifactSchema.parse({ - artifactHash: await artifactHash(parserVersion, input.body), + artifactHash: await artifactHash(parserVersion, input.body, artifactHashContext), contentType: inferContentType(materializedElements), createdAt: (options.now ?? defaultNow)(), documentAssetId: input.documentAssetId, @@ -1509,8 +1622,9 @@ function unstructuredParseElementMetadata({ const { image_base64: _imageBase64, page_number: _pageNumber, - ...parsed - } = cloneMetadata(metadata); + ...metadataWithoutInlineImage + } = metadata; + const parsed = cloneMetadata(metadataWithoutInlineImage); const assetRef = unstructuredAssetRef(metadata); const boundingBox = unstructuredBoundingBox(metadata.coordinates); const textAsHtml = metadataString(metadata, "text_as_html"); @@ -1930,8 +2044,14 @@ function assertInputBounds(body: Uint8Array, maxInputBytes: number): void { } } -async function artifactHash(parserVersion: string, body: Uint8Array): Promise { - const prefix = new TextEncoder().encode(`${parserVersion}\n`); +async function artifactHash( + parserVersion: string, + body: Uint8Array, + context?: string, +): Promise { + const prefix = new TextEncoder().encode( + context === undefined ? `${parserVersion}\n` : `${parserVersion}\n${context}\n`, + ); const bytes = new Uint8Array(prefix.byteLength + body.byteLength); bytes.set(prefix, 0); bytes.set(body, prefix.byteLength); @@ -1952,44 +2072,140 @@ async function boundedResponseText(response: Response, maxResponseBytes: number) const contentLength = response.headers.get("content-length"); if (contentLength && Number(contentLength) > maxResponseBytes) { + await cancelResponseBody(response.body); throw new ProviderResponseError( `Unstructured parser response exceeds maxResponseBytes=${maxResponseBytes}`, ); } - const body = new Uint8Array(await response.arrayBuffer()); + if (!response.body) { + return ""; + } - if (body.byteLength > maxResponseBytes) { - throw new ProviderResponseError( - `Unstructured parser response exceeds maxResponseBytes=${maxResponseBytes}`, - ); + const reader = response.body.getReader(); + const chunks: Uint8Array[] = []; + let totalBytes = 0; + + try { + while (true) { + const { done, value } = await reader.read(); + + if (done) { + break; + } + + totalBytes += value.byteLength; + if (totalBytes > maxResponseBytes) { + await cancelResponseReader(reader); + throw new ProviderResponseError( + `Unstructured parser response exceeds maxResponseBytes=${maxResponseBytes}`, + ); + } + chunks.push(value); + } + } finally { + reader.releaseLock(); + } + + const body = new Uint8Array(totalBytes); + let offset = 0; + for (const chunk of chunks) { + body.set(chunk, offset); + offset += chunk.byteLength; } return decodeUtf8(body); } +async function cancelResponseBody(body: ReadableStream | null): Promise { + if (!body) { + return; + } + + try { + await body.cancel(); + } catch { + // Preserve the bounded-response error even when the transport rejects cancellation. + } +} + +async function cancelResponseReader( + reader: ReadableStreamDefaultReader, +): Promise { + try { + await reader.cancel(); + } catch { + // Preserve the bounded-response error even when the transport rejects cancellation. + } +} + async function fetchWithRetries({ buildRequest, fetchImpl, maxRetries, retryDelayMs, + signal, sleep, }: { readonly buildRequest: () => Request; readonly fetchImpl: typeof fetch; readonly maxRetries: number; readonly retryDelayMs: number; + readonly signal: AbortSignal; readonly sleep: (ms: number) => Promise; }): Promise { for (let attempt = 0; ; attempt += 1) { - const response = await fetchImpl(buildRequest()); + signal.throwIfAborted(); + let response: Response; + try { + response = await fetchImpl(buildRequest()); + } catch (error) { + if (signal.aborted) { + throw abortSignalReason(signal); + } + if (attempt >= maxRetries) { + throw new ProviderRequestError("Unstructured parser request failed", { + cause: error, + retryable: true, + }); + } + await sleepWithAbort(sleep, retryDelayMs, signal); + continue; + } if (!isRetryableProviderStatus(response.status) || attempt >= maxRetries) { return response; } await response.body?.cancel().catch(() => undefined); - await sleep(retryDelayMs); + await sleepWithAbort(sleep, retryDelayMs, signal); + } +} + +async function sleepWithAbort( + sleep: (ms: number) => Promise, + delayMs: number, + signal: AbortSignal, +): Promise { + signal.throwIfAborted(); + if (delayMs === 0) { + return; + } + + let onAbort: (() => void) | undefined; + try { + await Promise.race([ + sleep(delayMs), + new Promise((_resolve, reject) => { + onAbort = () => reject(abortSignalReason(signal)); + signal.addEventListener("abort", onAbort, { once: true }); + if (signal.aborted) onAbort(); + }), + ]); + } finally { + if (onAbort) { + signal.removeEventListener("abort", onAbort); + } } } @@ -2009,6 +2225,158 @@ function validateRetryOptions({ } } +function validateUnstructuredResourceOptions({ + maxConcurrency, + requestTimeoutMs, +}: { + readonly maxConcurrency: number; + readonly requestTimeoutMs: number; +}): void { + if (!Number.isSafeInteger(maxConcurrency) || maxConcurrency < 1 || maxConcurrency > 32) { + throw new ProviderInputError( + "Unstructured parser maxConcurrency must be an integer between 1 and 32", + ); + } + if ( + !Number.isSafeInteger(requestTimeoutMs) || + requestTimeoutMs < 1 || + requestTimeoutMs > 600_000 + ) { + throw new ProviderInputError( + "Unstructured parser requestTimeoutMs must be an integer between 1 and 600000", + ); + } +} + +interface AbortAwareConcurrencyGate { + run(task: () => Promise, signal?: AbortSignal): Promise; +} + +function createAbortAwareConcurrencyGate(limit: number): AbortAwareConcurrencyGate { + let active = 0; + const waiters: Array<{ + readonly cleanup: () => void; + readonly reject: (error: unknown) => void; + readonly resolve: () => void; + readonly signal?: AbortSignal; + }> = []; + + const acquire = async (signal?: AbortSignal): Promise => { + if (signal?.aborted) { + throw abortSignalReason(signal); + } + if (active < limit) { + active += 1; + return; + } + + await new Promise((resolve, reject) => { + let onAbort: (() => void) | undefined; + const waiter = { + cleanup: () => { + if (onAbort) signal?.removeEventListener("abort", onAbort); + }, + reject, + resolve, + ...(signal ? { signal } : {}), + }; + onAbort = () => { + const index = waiters.indexOf(waiter); + if (index >= 0) { + waiters.splice(index, 1); + waiter.cleanup(); + reject(abortSignalReason(signal as AbortSignal)); + } + }; + signal?.addEventListener("abort", onAbort, { once: true }); + waiters.push(waiter); + if (signal?.aborted) { + onAbort(); + } + }); + }; + + const release = (): void => { + while (true) { + const next = waiters.shift(); + if (!next) { + active -= 1; + return; + } + if (next.signal?.aborted) { + next.cleanup(); + next.reject(abortSignalReason(next.signal)); + continue; + } + next.cleanup(); + next.resolve(); + return; + } + }; + + return { + run: async (task: () => Promise, signal?: AbortSignal): Promise => { + await acquire(signal); + try { + return await task(); + } finally { + release(); + } + }, + }; +} + +interface UnstructuredRequestDeadline { + readonly signal: AbortSignal; + dispose(): void; + expired(): boolean; + throwIfExpired(): void; +} + +function createUnstructuredRequestDeadline( + externalSignal: AbortSignal | undefined, + requestTimeoutMs: number, +): UnstructuredRequestDeadline { + const controller = new AbortController(); + const timeoutReason = new Error("Unstructured parser request deadline exceeded"); + let expired = false; + const onExternalAbort = () => controller.abort(abortSignalReason(externalSignal as AbortSignal)); + externalSignal?.addEventListener("abort", onExternalAbort, { once: true }); + if (externalSignal?.aborted) { + onExternalAbort(); + } + const timer = setTimeout(() => { + if (!controller.signal.aborted) { + expired = true; + controller.abort(timeoutReason); + } + }, requestTimeoutMs); + (timer as { unref?: () => void }).unref?.(); + + return { + signal: controller.signal, + dispose: () => { + clearTimeout(timer); + externalSignal?.removeEventListener("abort", onExternalAbort); + }, + expired: () => expired, + throwIfExpired: () => { + if (expired) { + throw timeoutReason; + } + }, + }; +} + +function abortSignalReason(signal: AbortSignal): unknown { + try { + signal.throwIfAborted(); + } catch (error) { + return error; + } + return new DOMException("The operation was aborted", "AbortError"); +} + function isRetryableProviderStatus(status: number): boolean { return status === 408 || status === 409 || status === 425 || status === 429 || status >= 500; } diff --git a/knowledge-fs/packages/parsers/src/parser.test.ts b/knowledge-fs/packages/parsers/src/parser.test.ts index 64fd7ea5d9e..a8fd425766b 100644 --- a/knowledge-fs/packages/parsers/src/parser.test.ts +++ b/knowledge-fs/packages/parsers/src/parser.test.ts @@ -4,6 +4,7 @@ import { describe, expect, it } from "vitest"; import { ProviderInputError, ProviderRateLimitError, + ProviderRequestError, ProviderResponseError, createNativeHtmlParser, createNativeMarkdownParser, @@ -925,7 +926,7 @@ describe("parser adapters", () => { metadata: { filename: "report.pdf", mimeType: "application/pdf", - parserVersion: "unstructured@4", + parserVersion: "unstructured@5", }, parser: "unstructured", version: 1, @@ -986,54 +987,94 @@ describe("parser adapters", () => { ]); }); - it.each([ - ["report.pdf", "application/pdf"], - ["handbook.docx", "application/vnd.openxmlformats-officedocument.wordprocessingml.document"], - ["briefing.pptx", "application/vnd.openxmlformats-officedocument.presentationml.presentation"], - ["forecast.xlsx", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"], - ])( - "uses local image extraction and adaptive strategy when parsing %s", - async (filename, mimeType) => { + it.each(["application/pdf", " Application/PDF; charset=binary "])( + "requests PDF image and table payloads for MIME %s without an external handler", + async (mimeType) => { const parser = createUnstructuredParserClient({ endpoint: "https://unstructured.example.test", fetch: async (request) => { - const parsedRequest = request instanceof Request ? request : new Request(request); - const form = await parsedRequest.formData(); + const form = await (request instanceof Request + ? request + : new Request(request) + ).formData(); - expect(form.get("strategy")).toBe("auto"); - expect(form.getAll("extract_image_block_types")).toEqual([]); - expect(form.get("extract_image_block_to_payload")).toBeNull(); - expect((form.get("files") as File).name).toBe(filename); + expect(form.get("strategy")).toBe("hi_res"); + expect(form.getAll("extract_image_block_types")).toEqual(["Image", "Table"]); + expect(form.get("extract_image_block_to_payload")).toBe("true"); - return new Response("[]", { - headers: { "content-type": "application/json" }, - status: 200, - }); + return new Response("[]", { status: 200 }); }, - generateId: () => "018f0d60-7a49-7cc2-9c1b-5b36f18f2c51", - now: () => createdAt, }); await expect( parser.parse({ body: new Uint8Array([1, 2, 3]), documentAssetId, - filename, + filename: "report.pdf", mimeType, + parserHints: { requiresImages: true }, version: 1, }), ).resolves.toMatchObject({ - metadata: { filename, mimeType, parserVersion: "unstructured@4" }, + metadata: { parserVersion: "unstructured@5" }, parser: "unstructured", }); }, ); + it("keeps hi_res PDF coordinates but suppresses payloads for an explicit external handler", async () => { + const parser = createUnstructuredParserClient({ + endpoint: "https://unstructured.example.test", + fetch: async (request) => { + const form = await (request instanceof Request ? request : new Request(request)).formData(); + + expect(form.get("strategy")).toBe("hi_res"); + expect(form.getAll("extract_image_block_types")).toEqual([]); + expect(form.get("extract_image_block_to_payload")).toBeNull(); + + return new Response("[]", { status: 200 }); + }, + }); + + await parser.parse({ + body: new Uint8Array([1, 2, 3]), + documentAssetId, + filename: "report.pdf", + mimeType: "application/pdf", + parserHints: { imagesHandledExternally: true, requiresImages: true }, + version: 1, + }); + }); + + it("does not request image payloads when the caller does not require images", async () => { + const parser = createUnstructuredParserClient({ + endpoint: "https://unstructured.example.test", + fetch: async (request) => { + const form = await (request instanceof Request ? request : new Request(request)).formData(); + + expect(form.get("strategy")).toBe("auto"); + expect(form.getAll("extract_image_block_types")).toEqual([]); + expect(form.get("extract_image_block_to_payload")).toBeNull(); + + return new Response("[]", { status: 200 }); + }, + }); + + await parser.parse({ + body: new Uint8Array([1, 2, 3]), + documentAssetId, + filename: "report.pdf", + mimeType: "application/pdf", + version: 1, + }); + }); + it.each([ [{ layoutComplexity: "simple" as const }, "fast", false], [{ requiresOcr: true }, "hi_res", false], [{ requiresTables: true }, "hi_res", false], [{ requiresImages: true }, "hi_res", true], + [{ imagesHandledExternally: true, requiresImages: true }, "hi_res", true], ])( "selects %s parsing hints without forcing hi_res for every document", async (parserHints, expectedStrategy, expectedImages) => { @@ -1061,6 +1102,46 @@ describe("parser adapters", () => { }, ); + it("includes Unstructured request strategy and parser hints in artifact hashes", async () => { + const parser = createUnstructuredParserClient({ + endpoint: "https://unstructured.example.test", + fetch: async () => new Response("[]", { status: 200 }), + }); + const parseWithHints = (parserHints: { + readonly imagesHandledExternally?: boolean; + readonly layoutComplexity?: "complex" | "simple"; + readonly requiresImages?: boolean; + readonly requiresOcr?: boolean; + readonly requiresTables?: boolean; + }) => + parser.parse({ + body: new Uint8Array([1, 2, 3]), + documentAssetId, + filename: "report.pdf", + mimeType: "application/pdf", + parserHints, + version: 1, + }); + + const [fast, ocr, tables, providerImages, externalImages] = await Promise.all([ + parseWithHints({ layoutComplexity: "simple" }), + parseWithHints({ requiresOcr: true }), + parseWithHints({ requiresTables: true }), + parseWithHints({ requiresImages: true }), + parseWithHints({ imagesHandledExternally: true, requiresImages: true }), + ]); + + expect( + new Set([ + fast.artifactHash, + ocr.artifactHash, + tables.artifactHash, + providerImages.artifactHash, + externalImages.artifactHash, + ]).size, + ).toBe(5); + }); + it.each([ [ "handbook.docx", @@ -1093,8 +1174,17 @@ describe("parser adapters", () => { ); const parser = createUnstructuredParserClient({ endpoint: "https://unstructured.example.test", - fetch: async () => - new Response( + fetch: async (request) => { + const form = await (request instanceof Request + ? request + : new Request(request) + ).formData(); + + expect(form.get("strategy")).toBe("auto"); + expect(form.getAll("extract_image_block_types")).toEqual([]); + expect(form.get("extract_image_block_to_payload")).toBeNull(); + + return new Response( JSON.stringify([ { metadata: { page_number: 1 }, @@ -1103,7 +1193,8 @@ describe("parser adapters", () => { }, ]), { headers: { "content-type": "application/json" }, status: 200 }, - ), + ); + }, generateId: () => "018f0d60-7a49-7cc2-9c1b-5b36f18f2c52", now: () => createdAt, }); @@ -1113,6 +1204,7 @@ describe("parser adapters", () => { documentAssetId, filename, mimeType, + parserHints: { requiresImages: true }, version: 1, }); @@ -1526,12 +1618,10 @@ describe("parser adapters", () => { expect(requestedUrl).toBe("https://unstructured.example.test/general/v0/general"); }); - it("retries retryable Unstructured failures and propagates AbortSignal", async () => { + it("retries retryable Unstructured failures", async () => { const statuses = [429, 200]; const delays: number[] = []; const seenAbortedSignals: boolean[] = []; - const controller = new AbortController(); - controller.abort(); const parser = createUnstructuredParserClient({ endpoint: "https://unstructured.example.test", fetch: async (request) => { @@ -1559,7 +1649,6 @@ describe("parser adapters", () => { documentAssetId, filename: "retry.pdf", mimeType: "application/pdf", - signal: controller.signal, version: 1, }), ).resolves.toMatchObject({ @@ -1572,7 +1661,107 @@ describe("parser adapters", () => { parser: "unstructured", }); expect(delays).toEqual([10]); - expect(seenAbortedSignals).toEqual([true, true]); + expect(seenAbortedSignals).toEqual([false, false]); + }); + + it("retries transient network failures without sleeping when retryDelayMs is zero", async () => { + let fetchCalls = 0; + let sleepCalls = 0; + const parser = createUnstructuredParserClient({ + endpoint: "https://unstructured.example.test", + fetch: async () => { + fetchCalls += 1; + if (fetchCalls === 1) throw new TypeError("connection reset"); + return new Response(JSON.stringify([{ text: "Recovered", type: "NarrativeText" }]), { + status: 200, + }); + }, + maxRetries: 1, + retryDelayMs: 0, + sleep: async () => { + sleepCalls += 1; + }, + }); + + await expect( + parser.parse({ + body: new Uint8Array([1]), + documentAssetId, + filename: "network-retry.pdf", + mimeType: "application/pdf", + version: 1, + }), + ).resolves.toMatchObject({ parser: "unstructured" }); + expect(fetchCalls).toBe(2); + expect(sleepCalls).toBe(0); + }); + + it("preserves caller cancellation while an Unstructured request is active", async () => { + const controller = new AbortController(); + let fetchStarted = false; + const parser = createUnstructuredParserClient({ + endpoint: "https://unstructured.example.test", + fetch: async (input) => { + const request = input instanceof Request ? input : new Request(input); + fetchStarted = true; + return await new Promise((_resolve, reject) => { + request.signal.addEventListener("abort", () => reject(request.signal.reason), { + once: true, + }); + }); + }, + }); + const pending = parser.parse({ + body: new Uint8Array([1]), + documentAssetId, + filename: "cancel-active.pdf", + mimeType: "application/pdf", + signal: controller.signal, + version: 1, + }); + + await waitForCondition(() => fetchStarted); + controller.abort(); + + await expect(pending).rejects.toMatchObject({ name: "AbortError" }); + }); + + it("validates Unstructured retry and resource bounds", () => { + const base = { + endpoint: "https://unstructured.example.test", + fetch: async () => new Response("[]", { status: 200 }), + }; + + expect(() => createUnstructuredParserClient({ ...base, maxRetries: -1 })).toThrow( + "maxRetries must be a non-negative integer", + ); + expect(() => createUnstructuredParserClient({ ...base, maxRetries: 0.5 })).toThrow( + "maxRetries must be a non-negative integer", + ); + expect(() => createUnstructuredParserClient({ ...base, retryDelayMs: -1 })).toThrow( + "retryDelayMs must be a non-negative integer", + ); + expect(() => createUnstructuredParserClient({ ...base, retryDelayMs: 0.5 })).toThrow( + "retryDelayMs must be a non-negative integer", + ); + expect(() => createUnstructuredParserClient({ ...base, maxConcurrency: 0 })).toThrow( + "maxConcurrency must be an integer between 1 and 32", + ); + expect(() => createUnstructuredParserClient({ ...base, maxConcurrency: 33 })).toThrow( + "maxConcurrency must be an integer between 1 and 32", + ); + expect(() => createUnstructuredParserClient({ ...base, maxConcurrency: 1.5 })).toThrow( + "maxConcurrency must be an integer between 1 and 32", + ); + expect(() => createUnstructuredParserClient({ ...base, requestTimeoutMs: 0 })).toThrow( + "requestTimeoutMs must be an integer between 1 and 600000", + ); + expect(() => createUnstructuredParserClient({ ...base, requestTimeoutMs: 600_001 })).toThrow( + "requestTimeoutMs must be an integer between 1 and 600000", + ); + expect(() => createUnstructuredParserClient({ ...base, requestTimeoutMs: 1.5 })).toThrow( + "requestTimeoutMs must be an integer between 1 and 600000", + ); }); it("rejects failed, invalid, and oversized Unstructured responses", async () => { @@ -1706,8 +1895,218 @@ describe("parser adapters", () => { }), ).rejects.toThrow("Unstructured parser response exceeds maxResponseBytes=3"); }); + + it("cancels an oversized streaming Unstructured response before reading later chunks", async () => { + let canceled = false; + let pullCount = 0; + const responseBody = new ReadableStream( + { + cancel: () => { + canceled = true; + }, + pull: (controller) => { + pullCount += 1; + if (pullCount === 1) { + controller.enqueue(new Uint8Array([91, 123])); + return; + } + if (pullCount === 2) { + controller.enqueue(new Uint8Array([125, 93])); + return; + } + controller.enqueue(new Uint8Array([32, 32])); + }, + }, + { highWaterMark: 0 }, + ); + const parser = createUnstructuredParserClient({ + endpoint: "https://unstructured.example.test", + fetch: async () => new Response(responseBody, { status: 200 }), + maxResponseBytes: 3, + }); + + await expect( + parser.parse({ + body: new Uint8Array([1]), + documentAssetId, + filename: "streamed.pdf", + mimeType: "application/pdf", + version: 1, + }), + ).rejects.toThrow("Unstructured parser response exceeds maxResponseBytes=3"); + expect(canceled).toBe(true); + expect(pullCount).toBe(2); + }); + + it("limits concurrent Unstructured requests and releases queued calls in FIFO order", async () => { + const releases: Array<() => void> = []; + let active = 0; + let maxActive = 0; + const parser = createUnstructuredParserClient({ + endpoint: "https://unstructured.example.test", + fetch: async () => { + active += 1; + maxActive = Math.max(maxActive, active); + await new Promise((resolve) => releases.push(resolve)); + active -= 1; + return new Response("[]", { status: 200 }); + }, + maxConcurrency: 2, + }); + const parses = Array.from({ length: 4 }, (_, index) => + parser.parse({ + body: new Uint8Array([index]), + documentAssetId, + filename: `concurrent-${index}.pdf`, + mimeType: "application/pdf", + version: 1, + }), + ); + + await waitForCondition(() => releases.length === 2); + expect(active).toBe(2); + for (const release of releases.splice(0, 2)) release(); + await waitForCondition(() => releases.length === 2); + for (const release of releases.splice(0, 2)) release(); + + await expect(Promise.all(parses)).resolves.toHaveLength(4); + expect(maxActive).toBe(2); + }); + + it("removes an aborted Unstructured request while it waits for the concurrency gate", async () => { + let fetchCalls = 0; + let releaseFirst: (() => void) | undefined; + const parser = createUnstructuredParserClient({ + endpoint: "https://unstructured.example.test", + fetch: async () => { + fetchCalls += 1; + await new Promise((resolve) => { + releaseFirst = resolve; + }); + return new Response("[]", { status: 200 }); + }, + maxConcurrency: 1, + }); + const first = parser.parse({ + body: new Uint8Array([1]), + documentAssetId, + filename: "first.pdf", + mimeType: "application/pdf", + version: 1, + }); + await waitForCondition(() => fetchCalls === 1); + const controller = new AbortController(); + const queued = parser.parse({ + body: new Uint8Array([2]), + documentAssetId, + filename: "queued.pdf", + mimeType: "application/pdf", + signal: controller.signal, + version: 1, + }); + controller.abort(); + + await expect(queued).rejects.toMatchObject({ name: "AbortError" }); + expect(fetchCalls).toBe(1); + releaseFirst?.(); + await expect(first).resolves.toMatchObject({ parser: "unstructured" }); + }); + + it("bounds stalled Unstructured response headers and response bodies", async () => { + const stalledHeaders = createUnstructuredParserClient({ + endpoint: "https://unstructured.example.test", + fetch: async (input) => { + const request = input instanceof Request ? input : new Request(input); + return await new Promise((_resolve, reject) => { + request.signal.addEventListener("abort", () => reject(request.signal.reason), { + once: true, + }); + }); + }, + requestTimeoutMs: 10, + }); + const stalledHeadersResult = stalledHeaders.parse({ + body: new Uint8Array([1]), + documentAssetId, + filename: "headers.pdf", + mimeType: "application/pdf", + version: 1, + }); + await expect(stalledHeadersResult).rejects.toMatchObject({ + code: "provider_request_failed", + retryable: true, + }); + + const stalledBody = createUnstructuredParserClient({ + endpoint: "https://unstructured.example.test", + fetch: async (input) => { + const request = input instanceof Request ? input : new Request(input); + return new Response( + new ReadableStream({ + start: (controller) => { + request.signal.addEventListener( + "abort", + () => controller.error(request.signal.reason), + { once: true }, + ); + }, + }), + { status: 200 }, + ); + }, + requestTimeoutMs: 10, + }); + const stalledBodyResult = stalledBody.parse({ + body: new Uint8Array([1]), + documentAssetId, + filename: "body.pdf", + mimeType: "application/pdf", + version: 1, + }); + await expect(stalledBodyResult).rejects.toMatchObject({ + code: "provider_request_failed", + retryable: true, + }); + }); + + it("classifies transient provider failures as retryable without retrying invalid input", async () => { + const networkFailure = createUnstructuredParserClient({ + endpoint: "https://unstructured.example.test", + fetch: async () => { + throw new TypeError("connection reset"); + }, + }).parse({ + body: new Uint8Array([1]), + documentAssetId, + filename: "network.pdf", + mimeType: "application/pdf", + version: 1, + }); + await expect(networkFailure).rejects.toBeInstanceOf(ProviderRequestError); + await expect(networkFailure).rejects.toMatchObject({ retryable: true }); + + const inputFailure = createUnstructuredParserClient({ + endpoint: "https://unstructured.example.test", + fetch: async () => new Response("bad request", { status: 400 }), + }).parse({ + body: new Uint8Array([1]), + documentAssetId, + filename: "input.pdf", + mimeType: "application/pdf", + version: 1, + }); + await expect(inputFailure).rejects.toMatchObject({ retryable: false, status: 400 }); + }); }); +async function waitForCondition(condition: () => boolean): Promise { + for (let attempt = 0; attempt < 100; attempt += 1) { + if (condition()) return; + await new Promise((resolve) => setTimeout(resolve, 0)); + } + throw new Error("Timed out waiting for parser test condition"); +} + describe("structured data parser coverage", () => { const structured = () => createNativeStructuredDataParser({ diff --git a/knowledge-fs/scripts/api-image-bundle-smoke.mjs b/knowledge-fs/scripts/api-image-bundle-smoke.mjs index d8da1480532..78658739c5b 100644 --- a/knowledge-fs/scripts/api-image-bundle-smoke.mjs +++ b/knowledge-fs/scripts/api-image-bundle-smoke.mjs @@ -51,6 +51,7 @@ try { const port = await dockerPort(containerId); const imageProcessing = await verifySharpRuntime(containerId); + const pdfRasterizer = await verifyPdfRasterizerRuntime(containerId); const health = await waitForHealth(`http://127.0.0.1:${port}/health`); console.log( @@ -61,6 +62,7 @@ try { imageTag, imageProcessing, ok: true, + pdfRasterizer, port, productionConfigValidated: false, runtime: health.runtime, @@ -73,6 +75,29 @@ try { } } +async function verifyPdfRasterizerRuntime(containerId) { + const [{ stderr, stdout }, concurrencyResult] = await Promise.all([ + execFileAsync(docker, ["exec", containerId, "pdftoppm", "-v"]), + execFileAsync(docker, [ + "exec", + containerId, + "printenv", + "KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY", + ]), + ]); + const version = `${stdout}${stderr}`.trim(); + const maxConcurrency = Number(concurrencyResult.stdout.trim()); + + if (!/^pdftoppm version\b/m.test(version)) { + throw new Error(`Unexpected Poppler PDF rasterizer version output: ${version}`); + } + if (maxConcurrency !== 2) { + throw new Error(`Unexpected Poppler PDF rasterizer max concurrency: ${maxConcurrency}`); + } + + return { command: "pdftoppm", maxConcurrency, version: version.split("\n")[0] }; +} + async function verifySharpRuntime(containerId) { const program = ` const sharp = (await import("sharp")).default; diff --git a/knowledge-fs/scripts/api-image-bundle-smoke.test.mjs b/knowledge-fs/scripts/api-image-bundle-smoke.test.mjs index c62f578ef96..166a78b352c 100644 --- a/knowledge-fs/scripts/api-image-bundle-smoke.test.mjs +++ b/knowledge-fs/scripts/api-image-bundle-smoke.test.mjs @@ -21,6 +21,10 @@ test("isolated API bundle smoke starts the container and checks compute health", assert.match(smokeScript, /components\?\.compute === true/); assert.match(smokeScript, /components\?\.objectStorage === false/); assert.match(smokeScript, /difyDependencyConnected/); + assert.match(smokeScript, /verifyPdfRasterizerRuntime/); + assert.match(smokeScript, /pdftoppm/); + assert.match(smokeScript, /KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY/); + assert.match(smokeScript, /maxConcurrency !== 2/); assert.match(smokeScript, /verifySharpRuntime/); assert.match(smokeScript, /await import\("sharp"\)/); assert.match(smokeScript, /sharp\.versions\.vips/); @@ -39,3 +43,15 @@ test("production API image carries and executes the target platform sharp runtim assert.match(apiDockerfile, /await import\('sharp'\)/); assert.match(apiDockerfile, /sharp native runtime smoke failed/); }); + +test("production API image carries and executes the Poppler PDF rasterizer", () => { + assert.match(apiDockerfile, /apt-get install --yes --no-install-recommends poppler-utils/); + assert.match(apiDockerfile, /KNOWLEDGE_PDF_RASTERIZER=poppler/); + assert.match(apiDockerfile, /KNOWLEDGE_PDF_RASTERIZER_DPI=144/); + assert.match(apiDockerfile, /KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI=48/); + assert.match(apiDockerfile, /KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS=30000/); + assert.match(apiDockerfile, /KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS=500/); + assert.match(apiDockerfile, /KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY=2/); + assert.match(apiDockerfile, /command -v pdftoppm/); + assert.match(apiDockerfile, /pdftoppm -v/); +}); diff --git a/knowledge-fs/scripts/compose-apps.test.mjs b/knowledge-fs/scripts/compose-apps.test.mjs index 0ca4640105d..2ca59290b9c 100644 --- a/knowledge-fs/scripts/compose-apps.test.mjs +++ b/knowledge-fs/scripts/compose-apps.test.mjs @@ -1,4 +1,5 @@ import assert from "node:assert/strict"; +import { spawnSync } from "node:child_process"; import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; @@ -43,6 +44,43 @@ function envVariableNames(source) { .map((line) => line.slice(0, line.indexOf("="))); } +function materializedDifyKnowledgeFsEnvironment(rootOverrides = {}) { + const dockerRoot = new URL("../../docker/", import.meta.url); + const env = { ...process.env }; + for (const name of [ + "KNOWLEDGE_PDF_RASTERIZER", + "KNOWLEDGE_PDF_RASTERIZER_DPI", + "KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS", + "KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY", + "KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI", + "KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS", + ]) { + delete env[name]; + } + Object.assign(env, rootOverrides); + + const result = spawnSync( + "docker", + [ + "compose", + "--project-directory", + dockerRoot.pathname, + "--env-file", + new URL(".env.example", dockerRoot).pathname, + "-f", + new URL("docker-compose.yaml", dockerRoot).pathname, + "config", + "--format", + "json", + ], + { encoding: "utf8", env }, + ); + + assert.ifError(result.error); + assert.equal(result.status, 0, result.stderr); + return JSON.parse(result.stdout).services.knowledge_fs.environment; +} + test("deployment Compose and Kubernetes artifacts are valid YAML", () => { for (const source of [compose, ...difyComposeFiles]) { const document = parse(source); @@ -111,7 +149,18 @@ test("app compose profile uses local middleware and the required Dify dependency compose, /^ {6}DATABASE_URL: postgresql:\/\/\$\{POSTGRES_USER:-knowledge_fs\}:\$\{POSTGRES_PASSWORD:-knowledge_fs\}@postgres:5432\/\$\{POSTGRES_DB:-knowledge_fs\}$/m, ); - assert.match(compose, /^ {6}UNSTRUCTURED_API_URL: http:\/\/unstructured:8000$/m); + assert.match( + compose, + /^ {6}UNSTRUCTURED_API_URL: \$\{UNSTRUCTURED_API_URL:-http:\/\/unstructured:8000\}$/m, + ); + assert.match( + compose, + /^ {6}UNSTRUCTURED_MAX_CONCURRENCY: \$\{UNSTRUCTURED_MAX_CONCURRENCY:-2\}$/m, + ); + assert.match( + compose, + /^ {6}UNSTRUCTURED_REQUEST_TIMEOUT_MS: \$\{UNSTRUCTURED_REQUEST_TIMEOUT_MS:-120000\}$/m, + ); assert.match( compose, /^ {6}DIFY_INNER_API_URL: \$\{DIFY_INNER_API_URL:-http:\/\/host\.docker\.internal:5001\}$/m, @@ -120,6 +169,30 @@ test("app compose profile uses local middleware and the required Dify dependency assert.doesNotMatch(compose, /^ {6}(?:MINIO|R2|OPENAI|ANTHROPIC|COHERE|GEMINI|VOYAGE)_/m); }); +test("app compose profile enables the bundled PDF rasterizer with bounded defaults", () => { + assert.match(compose, /^ {6}KNOWLEDGE_PDF_RASTERIZER: \$\{KNOWLEDGE_PDF_RASTERIZER:-poppler\}$/m); + assert.match( + compose, + /^ {6}KNOWLEDGE_PDF_RASTERIZER_DPI: \$\{KNOWLEDGE_PDF_RASTERIZER_DPI:-144\}$/m, + ); + assert.match( + compose, + /^ {6}KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI: \$\{KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI:-48\}$/m, + ); + assert.match( + compose, + /^ {6}KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS: \$\{KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS:-30000\}$/m, + ); + assert.match( + compose, + /^ {6}KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS: \$\{KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS:-500\}$/m, + ); + assert.match( + compose, + /^ {6}KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY: \$\{KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY:-2\}$/m, + ); +}); + test("app compose profile builds Admin as a production image after API readiness", () => { assert.match(compose, /^ {2}admin:$/m); assert.match(compose, /^ {6}dockerfile: apps\/admin\/Dockerfile$/m); @@ -150,6 +223,8 @@ test("Dify compose starts the integrated KnowledgeFS API by default and keeps it assert.match(knowledgeFs, /^ {4}expose:$/m); assert.match(knowledgeFs, /^ {6}- "8787"$/m); assert.doesNotMatch(knowledgeFs, /^ {4}ports:$/m); + assert.match(knowledgeFs, /^ {6}- path: \.\/envs\/core-services\/knowledge-fs\.env$/m); + assert.match(knowledgeFs, /whitelisted proxies/); assert.match( knowledgeFs, /^ {6}KNOWLEDGE_INTEGRATED_MODE_ENABLED: \$\{KNOWLEDGE_INTEGRATED_MODE_ENABLED:-true\}$/m, @@ -163,14 +238,59 @@ test("Dify compose starts the integrated KnowledgeFS API by default and keeps it assert.doesNotMatch(knowledgeFs, /^ {6}plugin_daemon:$/m); assert.match(knowledgeFs, /http:\/\/127\.0\.0\.1:8787\/ready/); assert.match(knowledgeFs, /^ {6}- default$/m); + assert.doesNotMatch(knowledgeFs, /^ {6}KNOWLEDGE_PDF_RASTERIZER:/m); + assert.doesNotMatch( + knowledgeFs, + /^ {6}KNOWLEDGE_PDF_RASTERIZER_(?:DPI|MAX_ASSETS|MAX_CONCURRENCY|TIMEOUT_MS|THUMBNAIL_DPI):/m, + ); + for (const suffix of [ + "", + "_DPI", + "_MAX_ASSETS", + "_MAX_CONCURRENCY", + "_THUMBNAIL_DPI", + "_TIMEOUT_MS", + ]) { + const canonicalName = `KNOWLEDGE_PDF_RASTERIZER${suffix}`; + assert.ok( + knowledgeFs.includes( + ` DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER${suffix}_OVERRIDE: \${${canonicalName}-}`, + ), + ); + } } assert.match(difyApiEnv, /^KNOWLEDGE_FS_ENABLED=\$\{KNOWLEDGE_FS_ENABLED:-false\}$/m); }); +test("Dify Compose whitelists root PDF overrides without shadowing service values when unset", () => { + const withoutRootOverrides = materializedDifyKnowledgeFsEnvironment({ + ROOT_ONLY_TEST_SECRET: "must-not-enter-container", + }); + assert.equal(withoutRootOverrides.DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_OVERRIDE, ""); + assert.equal( + withoutRootOverrides.DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY_OVERRIDE, + "", + ); + assert.equal(withoutRootOverrides.ROOT_ONLY_TEST_SECRET, undefined); + + const withRootOverrides = materializedDifyKnowledgeFsEnvironment({ + KNOWLEDGE_PDF_RASTERIZER: "off", + KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY: "6", + }); + assert.equal(withRootOverrides.DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_OVERRIDE, "off"); + assert.equal(withRootOverrides.DIFY_ROOT_KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY_OVERRIDE, "6"); +}); + test("KnowledgeFS deployment env contains only operator-owned runtime inputs", () => { assert.deepEqual(envVariableNames(difyKnowledgeFsEnv), [ "DATABASE_URL", "KNOWLEDGE_DOCUMENT_COMPILATION_RUNTIME", + "KNOWLEDGE_PDF_RASTERIZER", + "KNOWLEDGE_PDF_RASTERIZER_DPI", + "KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI", + "KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS", + "KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS", + "KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY", "KNOWLEDGE_OUTLINE_SUMMARY_MAX_CONCURRENCY", "KNOWLEDGE_OUTLINE_SUMMARY_BATCH_SIZE", "KNOWLEDGE_OUTLINE_SUMMARY_BATCH_MAX_INPUT_CHARS", @@ -187,8 +307,18 @@ test("KnowledgeFS deployment env contains only operator-owned runtime inputs", ( "KNOWLEDGE_QUERY_IMAGE_EXPANSION_TIMEOUT_MS", "UNSTRUCTURED_API_URL", "UNSTRUCTURED_API_KEY", + "UNSTRUCTURED_MAX_CONCURRENCY", + "UNSTRUCTURED_REQUEST_TIMEOUT_MS", + "UNSTRUCTURED_MAX_RESPONSE_BYTES", + "DIFY_OBJECT_STORAGE_REQUEST_TIMEOUT_MS", ]); assert.match(difyKnowledgeFsEnv, /^KNOWLEDGE_DOCUMENT_COMPILATION_RUNTIME=on$/m); + assert.match(difyKnowledgeFsEnv, /^KNOWLEDGE_PDF_RASTERIZER=poppler$/m); + assert.match(difyKnowledgeFsEnv, /^KNOWLEDGE_PDF_RASTERIZER_DPI=144$/m); + assert.match(difyKnowledgeFsEnv, /^KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI=48$/m); + assert.match(difyKnowledgeFsEnv, /^KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS=30000$/m); + assert.match(difyKnowledgeFsEnv, /^KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS=500$/m); + assert.match(difyKnowledgeFsEnv, /^KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY=2$/m); assert.match(difyKnowledgeFsEnv, /^KNOWLEDGE_FS_CAPABILITY_V2_ENABLED=false$/m); assert.doesNotMatch(difyKnowledgeFsEnv, /^MINIO_/m); }); @@ -205,6 +335,12 @@ test("deployment examples keep Dify KnowledgeFS rollout capabilities disabled", assert.match(kubernetesBaseline, /^ {2}KNOWLEDGE_LEGACY_AUTHORIZATION_REMOVED: "false"$/m); assert.match(kubernetesBaseline, /^ {2}KNOWLEDGE_DIRECT_UPLOAD_ENABLED: "off"$/m); assert.match(kubernetesBaseline, /^ {2}KNOWLEDGE_DIRECT_STREAM_ENABLED: "off"$/m); + assert.match(kubernetesBaseline, /^ {2}KNOWLEDGE_PDF_RASTERIZER: poppler$/m); + assert.match(kubernetesBaseline, /^ {2}KNOWLEDGE_PDF_RASTERIZER_DPI: "144"$/m); + assert.match(kubernetesBaseline, /^ {2}KNOWLEDGE_PDF_RASTERIZER_THUMBNAIL_DPI: "48"$/m); + assert.match(kubernetesBaseline, /^ {2}KNOWLEDGE_PDF_RASTERIZER_TIMEOUT_MS: "30000"$/m); + assert.match(kubernetesBaseline, /^ {2}KNOWLEDGE_PDF_RASTERIZER_MAX_ASSETS: "500"$/m); + assert.match(kubernetesBaseline, /^ {2}KNOWLEDGE_PDF_RASTERIZER_MAX_CONCURRENCY: "2"$/m); }); test("Kubernetes baseline starts at zero replicas with internal-only service and fail-closed probes", () => {