From 3e95efdde3088bca2b942cd438a0afcd3b697413 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 30 Jun 2026 02:26:44 +0800 Subject: [PATCH 1/6] Make source reconciliation timeout-safe --- src/app/api/sources/reconcile/route.ts | 140 +++++- .../sources/parsed-result-assets.test.ts | 264 +++++++++- src/domains/sources/parsed-result-assets.ts | 470 +++++++++++++++++- src/domains/sources/remote-library.ts | 25 +- src/domains/sources/repository.ts | 6 + src/domains/sources/route-listing.ts | 71 +-- src/domains/sources/route-service.test.ts | 55 +- .../sources/source-parse-result-repository.ts | 90 ++++ .../sources/source-reconcile-workflow.test.ts | 200 ++++++++ .../sources/source-reconcile-workflow.ts | 228 +++++++++ src/domains/sources/workflow-runtime.ts | 29 ++ src/domains/workspace/initial-state.test.ts | 41 +- src/domains/workspace/initial-state.ts | 70 ++- 13 files changed, 1575 insertions(+), 114 deletions(-) create mode 100644 src/domains/sources/source-reconcile-workflow.test.ts create mode 100644 src/domains/sources/source-reconcile-workflow.ts diff --git a/src/app/api/sources/reconcile/route.ts b/src/app/api/sources/reconcile/route.ts index c28bfa4..a9822b0 100644 --- a/src/app/api/sources/reconcile/route.ts +++ b/src/app/api/sources/reconcile/route.ts @@ -1,6 +1,18 @@ import { serve } from "@upstash/workflow/nextjs" -import { reconcileSourcesForWorkspace } from "@/domains/sources/reconcile" +import { + completeResultZipMultipartUpload, + createResultZipMultipartUploadPlan, + getResultZipPartRange, + prepareParsedResultAssetBatch, + uploadResultZipPart, + type MultipartUploadPart, +} from "@/domains/sources/parsed-result-assets" +import { + markSourceReadyAfterReconciliation, + pollSourceReconciliation, +} from "@/domains/sources/source-reconcile-workflow" +import { sourceWorkflowRuntime } from "@/domains/sources/workflow-runtime" import { makeKnowhereClient } from "@/integrations/knowhere" import { logger } from "@/lib/logger" @@ -16,24 +28,39 @@ const MAX_DELAY_S = 30 export const { POST } = serve(async (context) => { const { workspaceId, sourceId, apiKey } = context.requestPayload - const workspace = { id: workspaceId } const client = makeKnowhereClient(apiKey) let delay = INITIAL_DELAY_S + let completedJob: { + readonly jobId: string + readonly documentId: string + } | null = null for (let attempt = 0; attempt < MAX_POLL_ATTEMPTS; attempt++) { - const resolved = await context.run(`poll-${attempt}`, async () => { - const sources = await reconcileSourcesForWorkspace(workspace, client) - const source = sources.find((s) => s.id === sourceId) - if (!source || source.status !== "parsing") { - return { done: true, status: source?.status ?? "gone" } as const - } - return { done: false } as const + const poll = await context.run(`poll-${attempt}`, async () => { + return pollSourceReconciliation({ + workspaceId, + sourceId, + client, + }) }) - if (resolved.done) { + if (poll.kind === "ready-to-prepare") { + completedJob = { + jobId: poll.jobId, + documentId: poll.documentId, + } + logger.info("workflow: source parse completed; preparing artifacts", { + sourceId, + jobId: poll.jobId, + attempts: attempt + 1, + }) + break + } + + if (poll.kind === "resolved") { logger.info("workflow: source resolved", { sourceId, - status: resolved.status, + status: poll.status, attempts: attempt + 1, }) return @@ -43,8 +70,95 @@ export const { POST } = serve(async (context) => { delay = Math.min(Math.round(delay * 1.5), MAX_DELAY_S) } - logger.error("workflow: exhausted poll attempts", { + const jobToPrepare = completedJob + if (!jobToPrepare) { + logger.error("workflow: exhausted poll attempts", { + sourceId, + maxAttempts: MAX_POLL_ATTEMPTS, + }) + return + } + + const uploadPlan = await context.run("mirror-zip-start", async () => + createResultZipMultipartUploadPlan({ + workspaceId, + sourceId, + jobId: jobToPrepare.jobId, + client, + }), + ) + const parts: MultipartUploadPart[] = [] + + for (let partNumber = 1; partNumber <= uploadPlan.partCount; partNumber++) { + const range = getResultZipPartRange( + partNumber, + uploadPlan.partSizeBytes, + uploadPlan.sizeBytes, + ) + const part = await context.run(`mirror-zip-part-${partNumber}`, async () => + uploadResultZipPart({ + pathname: uploadPlan.pathname, + key: uploadPlan.key, + uploadId: uploadPlan.uploadId, + partNumber, + jobId: jobToPrepare.jobId, + client, + ...range, + }), + ) + parts.push(part) + } + + const resultBlobUrl = await context.run("mirror-zip-complete", async () => { + const result = await completeResultZipMultipartUpload({ + pathname: uploadPlan.pathname, + key: uploadPlan.key, + uploadId: uploadPlan.uploadId, + parts, + }) + return result.url + }) + + await context.run("parse-result-save-zip", async () => { + const saved = await sourceWorkflowRuntime.mergeParseAssetUrls( + workspaceId, + sourceId, + { + resultBlobUrl, + assetUrlsByFilePath: {}, + }, + ) + if (!saved) throw new Error("Source disappeared before saving result ZIP.") + }) + + let batchIndex = 0 + for (;;) { + const batch = await context.run( + `parse-assets-batch-${batchIndex}`, + async () => + prepareParsedResultAssetBatch({ + workspaceId, + sourceId, + jobId: jobToPrepare.jobId, + resultBlobUrl, + client, + repository: sourceWorkflowRuntime, + }), + ) + if (!batch.hasMore) break + batchIndex += 1 + } + + const ready = await context.run("source-ready", async () => + markSourceReadyAfterReconciliation({ + workspaceId, + sourceId, + documentId: jobToPrepare.documentId, + }), + ) + logger.info("workflow: source artifact preparation finished", { sourceId, - maxAttempts: MAX_POLL_ATTEMPTS, + status: ready.status, + assetBatches: batchIndex + 1, }) }) diff --git a/src/domains/sources/parsed-result-assets.test.ts b/src/domains/sources/parsed-result-assets.test.ts index 8e466ac..6db63c7 100644 --- a/src/domains/sources/parsed-result-assets.test.ts +++ b/src/domains/sources/parsed-result-assets.test.ts @@ -1,7 +1,14 @@ import { describe, expect, it, vi } from "vitest" import type { JobResult } from "@ontos-ai/knowhere-sdk" -import { storeParsedResultAssets } from "./parsed-result-assets" +import { + completeResultZipMultipartUpload, + createResultZipMultipartUploadPlan, + getResultZipPartRange, + prepareParsedResultAssetBatch, + storeParsedResultAssets, + uploadResultZipPart, +} from "./parsed-result-assets" describe("storeParsedResultAssets", () => { it("stores the result zip and extracted media assets in blob storage", async () => { @@ -91,6 +98,261 @@ describe("storeParsedResultAssets", () => { }) }) +describe("result ZIP multipart mirroring", () => { + it("uploads expected 8 MB byte ranges and completes with the final Blob URL", async () => { + const partSizeBytes = 8 * 1024 * 1024 + const sizeBytes = partSizeBytes * 2 + 3 + const fetchResult = vi.fn(async (_url: string | URL, init?: RequestInit) => { + if (init?.method === "HEAD") { + return new Response(null, { + status: 200, + headers: { + "content-length": String(sizeBytes), + }, + }) + } + + const range = init?.headers + ? new Headers(init.headers).get("range") + : null + const match = /^bytes=(\d+)-(\d+)$/u.exec(range ?? "") + if (!match) { + return new Response("missing range", { status: 400 }) + } + + const startByte = Number.parseInt(match[1] ?? "", 10) + const endByte = Number.parseInt(match[2] ?? "", 10) + return new Response(Buffer.alloc(endByte - startByte + 1), { + status: 206, + }) + }) + const client = { + jobs: { + get: vi.fn(async () => ({ + resultUrl: "https://knowhere.example/result.zip", + })), + }, + } + const blobStore = { + createMultipartUpload: vi.fn(async () => ({ + key: "blob-key", + uploadId: "upload-1", + })), + uploadPart: vi.fn( + async ( + _pathname: string, + _body: Buffer, + options: { readonly partNumber: number }, + ) => ({ + etag: `etag-${options.partNumber}`, + partNumber: options.partNumber, + }), + ), + completeMultipartUpload: vi.fn(async () => ({ + url: "https://blob.example/result.zip", + })), + } + + const plan = await createResultZipMultipartUploadPlan({ + workspaceId: "workspace_1", + sourceId: "source_1", + jobId: "job_1", + client, + blobStore, + fetchResult, + partSizeBytes, + }) + const parts: Array<{ readonly etag: string; readonly partNumber: number }> = + [] + for (let partNumber = 1; partNumber <= plan.partCount; partNumber++) { + parts.push( + await uploadResultZipPart({ + pathname: plan.pathname, + key: plan.key, + uploadId: plan.uploadId, + partNumber, + jobId: "job_1", + client, + blobStore, + fetchResult, + ...getResultZipPartRange(partNumber, partSizeBytes, sizeBytes), + }), + ) + } + const completed = await completeResultZipMultipartUpload({ + pathname: plan.pathname, + key: plan.key, + uploadId: plan.uploadId, + parts, + blobStore, + }) + + expect(plan).toEqual({ + pathname: + "workspaces/workspace_1/sources/source_1/parsed-result/result.zip", + key: "blob-key", + uploadId: "upload-1", + sizeBytes, + partSizeBytes, + partCount: 3, + }) + expect(fetchResult).toHaveBeenCalledWith( + "https://knowhere.example/result.zip", + { method: "HEAD" }, + ) + expect(fetchResult).toHaveBeenCalledWith( + "https://knowhere.example/result.zip", + { headers: { Range: `bytes=0-${partSizeBytes - 1}` } }, + ) + expect(fetchResult).toHaveBeenCalledWith( + "https://knowhere.example/result.zip", + { + headers: { + Range: `bytes=${partSizeBytes}-${partSizeBytes * 2 - 1}`, + }, + }, + ) + expect(fetchResult).toHaveBeenCalledWith( + "https://knowhere.example/result.zip", + { + headers: { + Range: `bytes=${partSizeBytes * 2}-${partSizeBytes * 2 + 2}`, + }, + }, + ) + expect(blobStore.uploadPart).toHaveBeenCalledTimes(3) + expect(blobStore.completeMultipartUpload).toHaveBeenCalledWith( + plan.pathname, + [ + { etag: "etag-1", partNumber: 1 }, + { etag: "etag-2", partNumber: 2 }, + { etag: "etag-3", partNumber: 3 }, + ], + expect.objectContaining({ + key: "blob-key", + uploadId: "upload-1", + }), + ) + expect(completed.url).toBe("https://blob.example/result.zip") + }) +}) + +describe("prepareParsedResultAssetBatch", () => { + it("skips saved assets, persists one bounded batch, and reports remaining work", async () => { + const client = { + jobs: { + load: vi.fn(async () => ({ + imageChunks: [ + { + filePath: "images/image-1.jpg", + data: Buffer.from("already uploaded"), + }, + { + filePath: "images/image-2.png", + data: Buffer.from("new image"), + }, + ], + tableChunks: [ + { + filePath: "tables/table-1.html", + html: "
", + }, + ], + })), + }, + } + const repository = { + getParseResultProgress: vi.fn(async () => ({ + resultBlobUrl: "https://blob.example/result.zip", + assetUrlsByFilePath: { + "images/image-1.jpg": "https://blob.example/images/image-1.jpg", + }, + })), + mergeParseAssetUrls: vi.fn(async () => undefined), + } + const blobStore = { + put: vi.fn(async (pathname: string) => ({ + url: `https://blob.example/${pathname}`, + })), + } + + const result = await prepareParsedResultAssetBatch({ + workspaceId: "workspace_1", + sourceId: "source_1", + jobId: "job_1", + resultBlobUrl: "https://blob.example/result.zip", + client, + repository, + blobStore, + batchLimit: 1, + }) + + expect(client.jobs.load).toHaveBeenCalledWith("job_1") + expect(blobStore.put).toHaveBeenCalledWith( + "workspaces/workspace_1/sources/source_1/parsed-result/images/image-2.png", + Buffer.from("new image"), + expect.objectContaining({ contentType: "image/png" }), + ) + expect(repository.mergeParseAssetUrls).toHaveBeenCalledWith( + "workspace_1", + "source_1", + { + resultBlobUrl: "https://blob.example/result.zip", + assetUrlsByFilePath: { + "images/image-2.png": + "https://blob.example/workspaces/workspace_1/sources/source_1/parsed-result/images/image-2.png", + }, + }, + ) + expect(result).toEqual({ + uploadedCount: 1, + remainingCount: 1, + hasMore: true, + }) + }) + + it("marks the asset phase complete when every image and table is saved", async () => { + const client = { + jobs: { + load: vi.fn(async () => ({ + imageChunks: [ + { + filePath: "images/image-1.jpg", + data: Buffer.from("already uploaded"), + }, + ], + tableChunks: [], + })), + }, + } + const repository = { + getParseResultProgress: vi.fn(async () => ({ + resultBlobUrl: "https://blob.example/result.zip", + assetUrlsByFilePath: { + "images/image-1.jpg": "https://blob.example/images/image-1.jpg", + }, + })), + mergeParseAssetUrls: vi.fn(async () => undefined), + } + + const result = await prepareParsedResultAssetBatch({ + workspaceId: "workspace_1", + sourceId: "source_1", + jobId: "job_1", + resultBlobUrl: "https://blob.example/result.zip", + client, + repository, + }) + + expect(repository.mergeParseAssetUrls).not.toHaveBeenCalled() + expect(result).toEqual({ + uploadedCount: 0, + remainingCount: 0, + hasMore: false, + }) + }) +}) + function makeJobResult(): JobResult { return { jobId: "job_1", diff --git a/src/domains/sources/parsed-result-assets.ts b/src/domains/sources/parsed-result-assets.ts index 6a36732..97ee62a 100644 --- a/src/domains/sources/parsed-result-assets.ts +++ b/src/domains/sources/parsed-result-assets.ts @@ -1,7 +1,12 @@ import "server-only" import path from "node:path" -import { put } from "@vercel/blob" +import { + completeMultipartUpload as completeVercelMultipartUpload, + createMultipartUpload as createVercelMultipartUpload, + put, + uploadPart as uploadVercelPart, +} from "@vercel/blob" import { Effect } from "effect" import type { JobResult } from "@ontos-ai/knowhere-sdk" @@ -10,6 +15,11 @@ export type StoredParsedResultAssets = { assetUrlsByFilePath: Readonly> } +export type MultipartUploadPart = { + readonly etag: string + readonly partNumber: number +} + export type ParsedResultBlobStore = { put( pathname: string, @@ -23,6 +33,132 @@ export type ParsedResultBlobStore = { ): Promise<{ url: string }> } +export type ResultZipMultipartBlobStore = { + createMultipartUpload( + pathname: string, + options: ResultZipBlobOptions, + ): Promise<{ + readonly key: string + readonly uploadId: string + }> + uploadPart( + pathname: string, + body: Buffer, + options: ResultZipUploadPartOptions, + ): Promise + completeMultipartUpload( + pathname: string, + parts: readonly MultipartUploadPart[], + options: ResultZipCompleteOptions, + ): Promise<{ readonly url: string }> +} + +export type ParsedResultAssetProgressRepository = { + getParseResultProgress( + workspaceId: string, + sourceId: string, + ): Promise + mergeParseAssetUrls( + workspaceId: string, + sourceId: string, + input: StoredParsedResultAssets, + ): Promise +} + +export type CreateResultZipMultipartUploadPlanInput = { + readonly workspaceId: string + readonly sourceId: string + readonly jobId: string + readonly client: ResultZipJobClient + readonly blobStore?: ResultZipMultipartBlobStore + readonly fetchResult?: ResultZipFetch + readonly partSizeBytes?: number +} + +export type ResultZipMultipartUploadPlan = { + readonly pathname: string + readonly key: string + readonly uploadId: string + readonly sizeBytes: number + readonly partSizeBytes: number + readonly partCount: number +} + +export type ResultZipPartRange = { + readonly startByte: number + readonly endByte: number +} + +export type UploadResultZipPartInput = ResultZipPartRange & { + readonly pathname: string + readonly key: string + readonly uploadId: string + readonly partNumber: number + readonly jobId: string + readonly client: ResultZipJobClient + readonly blobStore?: ResultZipMultipartBlobStore + readonly fetchResult?: ResultZipFetch +} + +export type CompleteResultZipMultipartUploadInput = { + readonly pathname: string + readonly key: string + readonly uploadId: string + readonly parts: readonly MultipartUploadPart[] + readonly blobStore?: ResultZipMultipartBlobStore +} + +export type PrepareParsedResultAssetBatchInput = { + readonly workspaceId: string + readonly sourceId: string + readonly jobId: string + readonly resultBlobUrl: string + readonly client: { + readonly jobs: { + load(jobId: string): Promise + } + } + readonly repository: ParsedResultAssetProgressRepository + readonly blobStore?: ParsedResultBlobStore + readonly batchLimit?: number + readonly maxDurationMs?: number +} + +export type PrepareParsedResultAssetBatchResult = { + readonly uploadedCount: number + readonly remainingCount: number + readonly hasMore: boolean +} + +type ResultZipJobClient = { + readonly jobs: { + get(jobId: string): Promise> + } +} + +type ResultZipFetch = ( + input: string | URL, + init?: RequestInit, +) => Promise + +type ResultZipBlobOptions = { + readonly access: "public" + readonly addRandomSuffix: false + readonly allowOverwrite: true + readonly contentType: "application/zip" +} + +type ResultZipUploadPartOptions = ResultZipBlobOptions & { + readonly key: string + readonly uploadId: string + readonly partNumber: number +} + +type ResultZipCompleteOptions = ResultZipBlobOptions & { + readonly key: string + readonly uploadId: string +} + type ParsedImageChunk = { readonly filePath?: string readonly data?: Buffer @@ -39,6 +175,12 @@ type ParsedResultWithAssets = { readonly tableChunks?: readonly ParsedTableChunk[] } +type ParsedAssetUpload = { + readonly filePath: string + readonly body: Buffer | string + readonly contentType: string +} + export type StoreParsedResultAssetsInput = { workspaceId: string sourceId: string @@ -52,6 +194,9 @@ export type StoreParsedResultAssetsInput = { } const parsedResultDirectoryName = "parsed-result" +const resultZipPartSizeBytes = 8 * 1024 * 1024 +const parsedAssetBatchLimit = 10 +const parsedAssetBatchMaxDurationMs = 45_000 // --------------------------------------------------------------------------- // Effect core @@ -116,6 +261,75 @@ export const storeParsedResultAssetsEffect = Effect.fn( }, ) +export const prepareParsedResultAssetBatchEffect = Effect.fn( + "prepareParsedResultAssetBatch", +)( + function* ({ + workspaceId, + sourceId, + jobId, + resultBlobUrl, + client, + repository, + blobStore = vercelBlobStore, + batchLimit = parsedAssetBatchLimit, + maxDurationMs = parsedAssetBatchMaxDurationMs, + }: PrepareParsedResultAssetBatchInput) { + const startedAt = Date.now() + const progress = yield* Effect.tryPromise(() => + repository.getParseResultProgress(workspaceId, sourceId), + ) + const existingAssetUrls = progress?.assetUrlsByFilePath ?? {} + const parseResult = (yield* Effect.tryPromise(() => + client.jobs.load(jobId), + )) as ParsedResultWithAssets + const assets = getUploadableParsedAssets(parseResult) + const missingAssets = assets.filter( + (asset) => existingAssetUrls[asset.filePath] === undefined, + ) + const uploadedAssetUrls: Record = {} + const blobPrefix = getParsedResultBlobPrefix(workspaceId, sourceId) + + for (const asset of missingAssets) { + const uploadedCount = Object.keys(uploadedAssetUrls).length + if (uploadedCount >= batchLimit) break + if (uploadedCount > 0 && Date.now() - startedAt >= maxDurationMs) break + + const blob = yield* Effect.tryPromise(() => + blobStore.put( + `${blobPrefix}/${asset.filePath}`, + asset.body, + getBlobPutOptions(asset.contentType), + ), + ) + uploadedAssetUrls[asset.filePath] = blob.url + } + + if (Object.keys(uploadedAssetUrls).length > 0) { + yield* Effect.tryPromise(() => + repository.mergeParseAssetUrls(workspaceId, sourceId, { + resultBlobUrl: progress?.resultBlobUrl ?? resultBlobUrl, + assetUrlsByFilePath: uploadedAssetUrls, + }), + ) + } + + const assetUrlsAfterBatch = { + ...existingAssetUrls, + ...uploadedAssetUrls, + } + const remainingCount = assets.filter( + (asset) => assetUrlsAfterBatch[asset.filePath] === undefined, + ).length + + return { + uploadedCount: Object.keys(uploadedAssetUrls).length, + remainingCount, + hasMore: remainingCount > 0, + } + }, +) + // --------------------------------------------------------------------------- // Async wrapper (backward-compatible) // --------------------------------------------------------------------------- @@ -138,6 +352,115 @@ export async function storeParsedResultAssets({ ) } +export async function createResultZipMultipartUploadPlan({ + workspaceId, + sourceId, + jobId, + client, + blobStore = vercelMultipartBlobStore, + fetchResult = fetch, + partSizeBytes = resultZipPartSizeBytes, +}: CreateResultZipMultipartUploadPlanInput): Promise { + const job = await client.jobs.get(jobId) + const resultUrl = requireJobResultUrl(job, jobId) + const sizeBytes = await getResultZipSize(resultUrl, fetchResult) + const pathname = `${getParsedResultBlobPrefix(workspaceId, sourceId)}/result.zip` + const upload = await blobStore.createMultipartUpload( + pathname, + getResultZipBlobOptions(), + ) + + return { + pathname, + key: upload.key, + uploadId: upload.uploadId, + sizeBytes, + partSizeBytes, + partCount: Math.max(1, Math.ceil(sizeBytes / partSizeBytes)), + } +} + +export function getResultZipPartRange( + partNumber: number, + partSizeBytes: number, + sizeBytes: number, +): ResultZipPartRange { + if (!Number.isInteger(partNumber) || partNumber < 1) { + throw new Error(`Invalid result ZIP part number: ${partNumber}.`) + } + if (!Number.isInteger(partSizeBytes) || partSizeBytes <= 0) { + throw new Error(`Invalid result ZIP part size: ${partSizeBytes}.`) + } + if (!Number.isInteger(sizeBytes) || sizeBytes <= 0) { + throw new Error(`Invalid result ZIP size: ${sizeBytes}.`) + } + + const startByte = (partNumber - 1) * partSizeBytes + if (startByte >= sizeBytes) { + throw new Error( + `Result ZIP part ${partNumber} starts beyond ${sizeBytes} bytes.`, + ) + } + + return { + startByte, + endByte: Math.min(sizeBytes - 1, startByte + partSizeBytes - 1), + } +} + +export async function uploadResultZipPart({ + pathname, + key, + uploadId, + partNumber, + jobId, + startByte, + endByte, + client, + blobStore = vercelMultipartBlobStore, + fetchResult = fetch, +}: UploadResultZipPartInput): Promise { + const job = await client.jobs.get(jobId) + const resultUrl = requireJobResultUrl(job, jobId) + const body = await fetchResultZipRange({ + resultUrl, + startByte, + endByte, + fetchResult, + }) + + return blobStore.uploadPart(pathname, body, { + ...getResultZipBlobOptions(), + key, + uploadId, + partNumber, + }) +} + +export async function completeResultZipMultipartUpload({ + pathname, + key, + uploadId, + parts, + blobStore = vercelMultipartBlobStore, +}: CompleteResultZipMultipartUploadInput): Promise<{ readonly url: string }> { + return blobStore.completeMultipartUpload( + pathname, + [...parts].sort((a, b) => a.partNumber - b.partNumber), + { + ...getResultZipBlobOptions(), + key, + uploadId, + }, + ) +} + +export async function prepareParsedResultAssetBatch( + input: PrepareParsedResultAssetBatchInput, +): Promise { + return Effect.runPromise(prepareParsedResultAssetBatchEffect(input)) +} + // --------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------- @@ -158,6 +481,45 @@ function getBlobPutOptions(contentType: string) { } } +function getResultZipBlobOptions(): ResultZipBlobOptions { + return { + access: "public", + addRandomSuffix: false, + allowOverwrite: true, + contentType: "application/zip", + } +} + +function getUploadableParsedAssets( + parseResult: ParsedResultWithAssets, +): readonly ParsedAssetUpload[] { + const assets: ParsedAssetUpload[] = [] + + for (const image of parseResult.imageChunks ?? []) { + const filePath = normalizeParsedAssetPath(image.filePath) + if (!filePath || !image.data) continue + + assets.push({ + filePath, + body: image.data, + contentType: getContentTypeForPath(filePath), + }) + } + + for (const table of parseResult.tableChunks ?? []) { + const filePath = normalizeParsedAssetPath(table.filePath) + if (!filePath || typeof table.html !== "string") continue + + assets.push({ + filePath, + body: table.html, + contentType: "text/html; charset=utf-8", + }) + } + + return assets +} + function normalizeParsedAssetPath(value: string | undefined): string | null { if (!value) return null const normalized = value.replaceAll("\\", "/").replace(/^\.\/+/, "") @@ -189,6 +551,103 @@ function getContentTypeForPath(filePath: string): string { return "application/octet-stream" } +async function getResultZipSize( + resultUrl: string, + fetchResult: ResultZipFetch, +): Promise { + const headResponse = await fetchResult(resultUrl, { method: "HEAD" }) + if (headResponse.ok) { + const contentLength = getPositiveIntegerHeader( + headResponse.headers, + "content-length", + ) + if (contentLength !== null) return contentLength + } + + const rangeResponse = await fetchResult(resultUrl, { + headers: { + Range: "bytes=0-0", + }, + }) + await rangeResponse.body?.cancel() + + const contentRange = rangeResponse.headers.get("content-range") + const rangeSize = parseContentRangeSize(contentRange) + if (rangeSize !== null) return rangeSize + + if (rangeResponse.ok) { + const contentLength = getPositiveIntegerHeader( + rangeResponse.headers, + "content-length", + ) + if (contentLength !== null) return contentLength + } + + throw new Error("Unable to determine Knowhere result ZIP size.") +} + +async function fetchResultZipRange(input: { + readonly resultUrl: string + readonly startByte: number + readonly endByte: number + readonly fetchResult: ResultZipFetch +}): Promise { + const expectedByteLength = input.endByte - input.startByte + 1 + const response = await input.fetchResult(input.resultUrl, { + headers: { + Range: `bytes=${input.startByte}-${input.endByte}`, + }, + }) + + if (!response.ok) { + throw new Error( + `Knowhere result ZIP range fetch failed with ${response.status}.`, + ) + } + if (response.status !== 206 && input.startByte !== 0) { + throw new Error("Knowhere result ZIP server ignored a non-initial range.") + } + + const body = Buffer.from(await response.arrayBuffer()) + if (body.byteLength !== expectedByteLength) { + throw new Error( + `Knowhere result ZIP range returned ${body.byteLength} bytes, expected ${expectedByteLength}.`, + ) + } + return body +} + +function getPositiveIntegerHeader( + headers: Headers, + name: string, +): number | null { + const value = headers.get(name) + if (!value) return null + + const parsed = Number.parseInt(value, 10) + return Number.isSafeInteger(parsed) && parsed > 0 ? parsed : null +} + +function parseContentRangeSize(value: string | null): number | null { + if (!value) return null + + const match = /^bytes\s+\d+-\d+\/(\d+)$/iu.exec(value.trim()) + if (!match) return null + + const parsed = Number.parseInt(match[1] ?? "", 10) + return Number.isSafeInteger(parsed) && parsed > 0 ? parsed : null +} + +function requireJobResultUrl( + job: Pick, + jobId: string, +): string { + if (typeof job.resultUrl === "string" && job.resultUrl.length > 0) { + return job.resultUrl + } + throw new Error(`Knowhere job ${jobId} does not expose a result ZIP URL.`) +} + const vercelBlobStore: ParsedResultBlobStore = { put: (pathname, body, options) => put(pathname, body, { @@ -198,3 +657,12 @@ const vercelBlobStore: ParsedResultBlobStore = { multipart: options.multipart, }), } + +const vercelMultipartBlobStore: ResultZipMultipartBlobStore = { + createMultipartUpload: (pathname, options) => + createVercelMultipartUpload(pathname, options), + uploadPart: (pathname, body, options) => + uploadVercelPart(pathname, body, options), + completeMultipartUpload: (pathname, parts, options) => + completeVercelMultipartUpload(pathname, [...parts], options), +} diff --git a/src/domains/sources/remote-library.ts b/src/domains/sources/remote-library.ts index 5c867bf..4cc3934 100644 --- a/src/domains/sources/remote-library.ts +++ b/src/domains/sources/remote-library.ts @@ -104,7 +104,10 @@ export function localizeRemoteLibrarySources( input: RemoteLibraryLocalizationInput, ): Effect.Effect { return Effect.gen(function* () { - const remoteDocuments = yield* listRemoteLibraryDocuments(input) + const remoteDocuments = (yield* listRemoteLibraryDocuments(input)).filter( + (document) => + !matchesActiveNotebookParsingSource(document, input.localSources), + ) if (remoteDocuments.length === 0) return input.localSources const localizedSources = yield* Effect.all( @@ -170,6 +173,26 @@ function getRemoteSourceStatus(status: string): SourceStatus { return "parsing" } +function matchesActiveNotebookParsingSource( + document: RemoteDocument, + localSources: readonly Source[], +): boolean { + if (document.documentMetadata?.createdByClient !== "notebook") return false + if (!document.title || !document.mimeType || document.sizeBytes === undefined) { + return false + } + + return localSources.some( + (source) => + source.status === "parsing" && + Boolean(source.knowhereJobId) && + !source.knowhereDocumentId && + source.title === document.title && + source.mimeType === document.mimeType && + source.sizeBytes === document.sizeBytes, + ) +} + function isActiveDocumentStatus(status: string): boolean { return status === "active" || status === "ready" || status === "done" } diff --git a/src/domains/sources/repository.ts b/src/domains/sources/repository.ts index 62accd6..52f34cb 100644 --- a/src/domains/sources/repository.ts +++ b/src/domains/sources/repository.ts @@ -18,6 +18,8 @@ type SourceRepository = { readonly clearStagedBlobEffect: typeof sourceRowRepository.clearStagedBlobEffect readonly softDeleteEffect: typeof sourceRowRepository.softDeleteEffect readonly saveParseResultEffect: typeof sourceParseResultRepository.saveParseResultEffect + readonly mergeParseAssetUrlsEffect: typeof sourceParseResultRepository.mergeParseAssetUrlsEffect + readonly getParseResultProgressEffect: typeof sourceParseResultRepository.getParseResultProgressEffect readonly getParseAssetUrlsEffect: typeof sourceParseResultRepository.getParseAssetUrlsEffect } @@ -37,5 +39,9 @@ export const sourceRepository: SourceRepository = { clearStagedBlobEffect: sourceRowRepository.clearStagedBlobEffect, softDeleteEffect: sourceRowRepository.softDeleteEffect, saveParseResultEffect: sourceParseResultRepository.saveParseResultEffect, + mergeParseAssetUrlsEffect: + sourceParseResultRepository.mergeParseAssetUrlsEffect, + getParseResultProgressEffect: + sourceParseResultRepository.getParseResultProgressEffect, getParseAssetUrlsEffect: sourceParseResultRepository.getParseAssetUrlsEffect, } diff --git a/src/domains/sources/route-listing.ts b/src/domains/sources/route-listing.ts index d21ac90..0052500 100644 --- a/src/domains/sources/route-listing.ts +++ b/src/domains/sources/route-listing.ts @@ -10,8 +10,11 @@ import { routeResult } from "@/lib/route-result" import { logger } from "@/lib/logger" import { knowhereDemoApi } from "@/integrations/knowhere-demo" import { toSourceView } from "./view" -import { startBackgroundReconciliation } from "./background-reconcile" +import { + startBackgroundReconciliation as defaultStartBackgroundReconciliation, +} from "./background-reconcile" import { localizeRemoteLibrarySources } from "./remote-library" +import type { Source } from "@/infrastructure/db/schema" import type { JsonRouteResult, ListSourcesBody, @@ -39,6 +42,7 @@ type RouteListingDependencies = Pick< readonly reconcileSourcesForWorkspace: SourceRouteServiceDependencies[ "reconcileSourcesForWorkspace" ] + readonly startBackgroundReconciliation?: typeof defaultStartBackgroundReconciliation } type RouteListing = { @@ -84,12 +88,7 @@ const listSourcesEffect = ( deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), ) const client = deps.makeKnowhereClient(apiKey) - const sources = yield* Effect.tryPromise(() => { - if (!hasParsingSources(listedSources)) { - return Promise.resolve(listedSources) - } - return deps.reconcileSourcesForWorkspace(workspace, client) - }) + const sources = listedSources const demoSourceResolution = resolveWorkspaceDemoSources(sources, catalog) const workspaceSources = yield* localizeRemoteLibrarySources({ workspace, @@ -102,23 +101,16 @@ const listSourcesEffect = ( getWorkspaceSourcesNeedingKnowhereChunkCount(workspaceSources) const materializedDemoSourceOptions = getMaterializedDemoSourceViewOptionsBySourceId(workspaceSources, catalog) - const parsingSources = sources.filter( - (source) => source.status === "parsing" && source.knowhereJobId, - ) - if (parsingSources.length > 0) { - logger.info("route-listing: re-triggering reconciliation for parsing sources", { + yield* Effect.sync(() => + triggerBackgroundReconciliationForParsingSources({ workspaceId: workspace.id, - count: parsingSources.length, - sourceIds: parsingSources.map((s) => s.id), - }) - } - for (const source of parsingSources) { - yield* Effect.fork( - Effect.tryPromise(() => - startBackgroundReconciliation(workspace.id, source.id, apiKey), - ), - ) - } + sources, + apiKey, + startBackgroundReconciliation: + deps.startBackgroundReconciliation ?? + defaultStartBackgroundReconciliation, + }), + ) const sourceOptions = yield* deps.getSourceViewOptionsBySourceId( sourcesNeedingKnowhereChunkCount, client, @@ -154,13 +146,32 @@ const listSourcesEffect = ( export { createRouteListing } -function hasParsingSources( - sources: readonly { - readonly status: string - readonly knowhereJobId: string | null - }[], -): boolean { - return sources.some( +function triggerBackgroundReconciliationForParsingSources(input: { + readonly workspaceId: string + readonly sources: readonly Source[] + readonly apiKey: string + readonly startBackgroundReconciliation: typeof defaultStartBackgroundReconciliation +}): void { + const parsingSources = input.sources.filter( (source) => source.status === "parsing" && source.knowhereJobId, ) + if (parsingSources.length === 0) return + + logger.info("route-listing: re-triggering reconciliation for parsing sources", { + workspaceId: input.workspaceId, + count: parsingSources.length, + sourceIds: parsingSources.map((source) => source.id), + }) + + for (const source of parsingSources) { + void input + .startBackgroundReconciliation(input.workspaceId, source.id, input.apiKey) + .catch((error: unknown) => { + logger.warn("route-listing: background reconciliation trigger failed", { + workspaceId: input.workspaceId, + sourceId: source.id, + error: error instanceof Error ? error.message : String(error), + }) + }) + } } diff --git a/src/domains/sources/route-service.test.ts b/src/domains/sources/route-service.test.ts index 4ac60cc..9901d8f 100644 --- a/src/domains/sources/route-service.test.ts +++ b/src/domains/sources/route-service.test.ts @@ -37,7 +37,7 @@ const source: Source = { const localizeNoRemoteDocuments = vi.fn(async () => source); describe("source route service", () => { - it("reconciles authenticated sources through the listing route workflow", async () => { + it("lists authenticated sources and triggers background reconciliation", async () => { const knowhereClient = { documents: { archive: vi.fn(async () => undefined), @@ -63,6 +63,7 @@ describe("source route service", () => { ); const listSourcesForWorkspace = vi.fn(async () => [source]); const reconcileSourcesForWorkspace = vi.fn(async () => [source]); + const startBackgroundReconciliation = vi.fn(async () => undefined); const listHiddenDemoSourceIds = vi.fn(async () => []); const listing = createRouteListing({ demoApi: { @@ -79,6 +80,7 @@ describe("source route service", () => { makeKnowhereClient: vi.fn(() => knowhereClient), listSourcesForWorkspace, reconcileSourcesForWorkspace, + startBackgroundReconciliation, sourceService: { listHiddenDemoSourceIds, localizeRemoteDocument: localizeNoRemoteDocuments, @@ -108,9 +110,11 @@ describe("source route service", () => { "session=abc", ); expect(listSourcesForWorkspace).toHaveBeenCalledWith(workspace.id); - expect(reconcileSourcesForWorkspace).toHaveBeenCalledWith( - workspace, - knowhereClient, + expect(reconcileSourcesForWorkspace).not.toHaveBeenCalled(); + expect(startBackgroundReconciliation).toHaveBeenCalledWith( + workspace.id, + source.id, + "jwt_123", ); expect(listHiddenDemoSourceIds).toHaveBeenCalledWith(workspace.id); }); @@ -299,14 +303,16 @@ describe("source route service", () => { ]); }); - it("reconciles parsing sources before localizing matching Knowhere documents", async () => { - const reconciledSource: Source = { + it("keeps matching Notebook uploads parsing until artifacts are ready", async () => { + const parsingSource: Source = { ...source, id: "source_1", title: "uploaded.pdf", - status: "ready", - knowhereJobId: null, - knowhereDocumentId: "doc_uploaded", + mimeType: "application/pdf", + sizeBytes: 100, + status: "parsing", + knowhereJobId: "job_1", + knowhereDocumentId: null, }; const listDocuments = vi .fn() @@ -317,6 +323,12 @@ describe("source route service", () => { namespace: "default", status: "active", sourceFileName: "uploaded.pdf", + documentMetadata: { + createdByClient: "notebook", + title: "uploaded.pdf", + mimeType: "application/pdf", + sizeBytes: 100, + }, }, ], }) @@ -341,8 +353,9 @@ describe("source route service", () => { upload: vi.fn(), }, }; - const reconcileSourcesForWorkspace = vi.fn(async () => [reconciledSource]); - const localizeRemoteDocument = vi.fn(async () => reconciledSource); + const reconcileSourcesForWorkspace = vi.fn(async () => [parsingSource]); + const localizeRemoteDocument = vi.fn(async () => parsingSource); + const startBackgroundReconciliation = vi.fn(async () => undefined); const listing = createRouteListing({ demoApi: { fetchCatalog: vi.fn(async () => emptyDemoCatalog), @@ -356,8 +369,9 @@ describe("source route service", () => { })), getSourceViewOptionsBySourceId: vi.fn(() => Effect.succeed(new Map())), makeKnowhereClient: vi.fn(() => knowhereClient), - listSourcesForWorkspace: vi.fn(async () => [source]), + listSourcesForWorkspace: vi.fn(async () => [parsingSource]), reconcileSourcesForWorkspace, + startBackgroundReconciliation, sourceService: { listHiddenDemoSourceIds: vi.fn(async () => []), localizeRemoteDocument, @@ -366,22 +380,19 @@ describe("source route service", () => { const result = await listing.listSources({ cookieHeader: "session=abc" }); - expect(reconcileSourcesForWorkspace).toHaveBeenCalledWith( - workspace, - knowhereClient, - ); - expect(localizeRemoteDocument).toHaveBeenCalledWith( + expect(reconcileSourcesForWorkspace).not.toHaveBeenCalled(); + expect(startBackgroundReconciliation).toHaveBeenCalledWith( workspace.id, - expect.objectContaining({ - documentId: "doc_uploaded", - }), + "source_1", + "jwt_123", ); + expect(localizeRemoteDocument).not.toHaveBeenCalled(); expect(result.body.sources).toEqual([ expect.objectContaining({ id: "source_1", - documentId: "doc_uploaded", title: "uploaded.pdf", - status: "ready", + status: "parsing", + documentId: undefined, }), ]); }); diff --git a/src/domains/sources/source-parse-result-repository.ts b/src/domains/sources/source-parse-result-repository.ts index c3912db..ee899c7 100644 --- a/src/domains/sources/source-parse-result-repository.ts +++ b/src/domains/sources/source-parse-result-repository.ts @@ -15,12 +15,26 @@ type SaveSourceParseResultInput = { readonly assetUrlsByFilePath: Readonly> } +type SourceParseResultProgress = { + readonly resultBlobUrl: string + readonly assetUrlsByFilePath: Readonly> +} + type SourceParseResultRepository = { readonly saveParseResultEffect: ( workspaceId: string, sourceId: string, input: SaveSourceParseResultInput, ) => Effect.Effect + readonly mergeParseAssetUrlsEffect: ( + workspaceId: string, + sourceId: string, + input: SaveSourceParseResultInput, + ) => Effect.Effect + readonly getParseResultProgressEffect: ( + workspaceId: string, + sourceId: string, + ) => Effect.Effect readonly getParseAssetUrlsEffect: ( workspaceId: string, sourceId: string, @@ -58,6 +72,80 @@ const saveParseResultEffect: SourceParseResultRepository["saveParseResultEffect" return result ?? null }) +const mergeParseAssetUrlsEffect: SourceParseResultRepository["mergeParseAssetUrlsEffect"] = + (workspaceId: string, sourceId: string, input: SaveSourceParseResultInput) => + Effect.gen(function* () { + const db = yield* DbClient + const source = yield* Effect.promise(() => + sourceRowRepository.findInWorkspaceWithDb(db, workspaceId, sourceId), + ) + if (!source) return null + + const [current] = yield* Effect.promise(() => + db + .select({ + resultBlobUrl: sourceParseResults.resultBlobUrl, + assetUrls: sourceParseResults.assetUrls, + }) + .from(sourceParseResults) + .where(eq(sourceParseResults.sourceId, sourceId)) + .limit(1), + ) + const assetUrlsByFilePath = { + ...(current?.assetUrls ?? {}), + ...input.assetUrlsByFilePath, + } + + const [result] = yield* Effect.promise(() => + db + .insert(sourceParseResults) + .values({ + sourceId, + resultBlobUrl: input.resultBlobUrl, + assetUrls: assetUrlsByFilePath, + }) + .onConflictDoUpdate({ + target: sourceParseResults.sourceId, + set: { + resultBlobUrl: input.resultBlobUrl, + assetUrls: assetUrlsByFilePath, + updatedAt: sql`now()`, + }, + }) + .returning(), + ) + + return result ?? null + }) + +const getParseResultProgressEffect: SourceParseResultRepository["getParseResultProgressEffect"] = + (workspaceId: string, sourceId: string) => + Effect.gen(function* () { + const db = yield* DbClient + const source = yield* Effect.promise(() => + sourceRowRepository.findInWorkspaceWithDb(db, workspaceId, sourceId), + ) + if (!source) return null + + const row = yield* Effect.promise(() => + db + .select({ + resultBlobUrl: sourceParseResults.resultBlobUrl, + assetUrls: sourceParseResults.assetUrls, + }) + .from(sourceParseResults) + .where(eq(sourceParseResults.sourceId, sourceId)) + .limit(1), + ) + const progress = row[0] + if (!progress) return null + + return { + resultBlobUrl: progress.resultBlobUrl, + assetUrlsByFilePath: progress.assetUrls, + } + }) + const getParseAssetUrlsEffect: SourceParseResultRepository["getParseAssetUrlsEffect"] = (workspaceId: string, sourceId: string) => Effect.gen(function* () { @@ -80,5 +168,7 @@ const getParseAssetUrlsEffect: SourceParseResultRepository["getParseAssetUrlsEff export const sourceParseResultRepository: SourceParseResultRepository = { saveParseResultEffect, + mergeParseAssetUrlsEffect, + getParseResultProgressEffect, getParseAssetUrlsEffect, } diff --git a/src/domains/sources/source-reconcile-workflow.test.ts b/src/domains/sources/source-reconcile-workflow.test.ts new file mode 100644 index 0000000..9548ea0 --- /dev/null +++ b/src/domains/sources/source-reconcile-workflow.test.ts @@ -0,0 +1,200 @@ +import { describe, expect, it, vi } from "vitest" +import type { JobResult } from "@ontos-ai/knowhere-sdk" + +import type { Source, Workspace } from "@/infrastructure/db/schema" +import { + markSourceReadyAfterReconciliation, + pollSourceReconciliation, +} from "./source-reconcile-workflow" + +const workspace: Workspace = { + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-06T00:00:00Z"), +} + +function makeSource(overrides: Partial = {}): Source { + return { + id: "source_1", + workspaceId: workspace.id, + title: "notes.pdf", + mimeType: "application/pdf", + sizeBytes: 1, + status: "parsing", + failureReason: null, + knowhereJobId: "job_1", + knowhereDocumentId: null, + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-05-06T00:00:00Z"), + updatedAt: new Date("2026-05-06T00:00:00Z"), + deletedAt: null, + ...overrides, + } +} + +function makeDoneJob(overrides: Partial = {}): JobResult { + return { + jobId: "job_1", + status: "done", + sourceType: "file", + namespace: workspace.namespace, + documentId: "doc_1", + resultUrl: "https://knowhere.example/result.zip", + createdAt: new Date("2026-05-06T00:00:00Z"), + isDone: true, + isFailed: false, + isTerminal: true, + ...overrides, + } +} + +function createRepository(source: Source | null = makeSource()) { + return { + findInWorkspace: vi.fn(async () => source), + markFailed: vi.fn(async () => null), + markReady: vi.fn(async () => + source ? { ...source, status: "ready" as const } : null, + ), + clearStagedBlob: vi.fn(async () => null), + } +} + +describe("pollSourceReconciliation", () => { + it("leaves a completed Knowhere job parsing until artifacts can be prepared", async () => { + const repository = createRepository() + const client = { + jobs: { + get: vi.fn(async () => makeDoneJob()), + }, + } + + const result = await pollSourceReconciliation({ + workspaceId: workspace.id, + sourceId: "source_1", + client, + repository, + }) + + expect(result).toEqual({ + kind: "ready-to-prepare", + jobId: "job_1", + documentId: "doc_1", + }) + expect(repository.markReady).not.toHaveBeenCalled() + expect(repository.markFailed).not.toHaveBeenCalled() + }) + + it("marks failed Knowhere jobs failed and cleans staged uploads", async () => { + const source = makeSource({ + stagedBlobPathname: "source-uploads/upload_1/document.pdf", + }) + const repository = createRepository(source) + const blobStore = { + deleteStagedSourceBlob: vi.fn(async () => undefined), + } + const client = { + jobs: { + get: vi.fn(async () => + makeDoneJob({ + status: "failed", + isDone: false, + isFailed: true, + error: { + code: "parser_failed", + message: "Parser rejected this document.", + requestId: "request_1", + }, + }), + ), + }, + } + + const result = await pollSourceReconciliation({ + workspaceId: workspace.id, + sourceId: "source_1", + client, + repository, + blobStore, + }) + + expect(result).toEqual({ kind: "resolved", status: "failed" }) + expect(repository.markFailed).toHaveBeenCalledWith( + workspace.id, + "source_1", + "Parser rejected this document.", + "parsing", + ) + expect(blobStore.deleteStagedSourceBlob).toHaveBeenCalledWith( + "source-uploads/upload_1/document.pdf", + ) + expect(repository.clearStagedBlob).toHaveBeenCalledWith( + workspace.id, + "source_1", + ) + }) + + it("marks done jobs without a result URL failed", async () => { + const repository = createRepository() + const jobWithoutResultUrl = makeDoneJob() + delete jobWithoutResultUrl.resultUrl + const client = { + jobs: { + get: vi.fn(async () => jobWithoutResultUrl), + }, + } + + const result = await pollSourceReconciliation({ + workspaceId: workspace.id, + sourceId: "source_1", + client, + repository, + }) + + expect(result).toEqual({ kind: "resolved", status: "failed" }) + expect(repository.markFailed).toHaveBeenCalledWith( + workspace.id, + "source_1", + "Parsing finished but Knowhere did not publish a result ZIP.", + "parsing", + ) + }) +}) + +describe("markSourceReadyAfterReconciliation", () => { + it("marks ready after artifact preparation and cleans staged uploads", async () => { + const source = makeSource({ + stagedBlobPathname: "source-uploads/upload_1/document.pdf", + }) + const repository = createRepository(source) + const blobStore = { + deleteStagedSourceBlob: vi.fn(async () => undefined), + } + + const result = await markSourceReadyAfterReconciliation({ + workspaceId: workspace.id, + sourceId: "source_1", + documentId: "doc_1", + repository, + blobStore, + }) + + expect(result).toEqual({ status: "ready" }) + expect(repository.markReady).toHaveBeenCalledWith( + workspace.id, + "source_1", + "doc_1", + ) + expect(blobStore.deleteStagedSourceBlob).toHaveBeenCalledWith( + "source-uploads/upload_1/document.pdf", + ) + expect(repository.clearStagedBlob).toHaveBeenCalledWith( + workspace.id, + "source_1", + ) + }) +}) diff --git a/src/domains/sources/source-reconcile-workflow.ts b/src/domains/sources/source-reconcile-workflow.ts new file mode 100644 index 0000000..97e5ebe --- /dev/null +++ b/src/domains/sources/source-reconcile-workflow.ts @@ -0,0 +1,228 @@ +import "server-only" + +import { del } from "@vercel/blob" +import type { JobResult } from "@ontos-ai/knowhere-sdk" + +import type { Source } from "@/infrastructure/db/schema" +import { logger } from "@/lib/logger" +import { sourceWorkflowRuntime } from "./workflow-runtime" + +type SourceReconcileWorkflowClient = { + readonly jobs: { + get(jobId: string): Promise + } +} + +type SourceReconcileWorkflowRepository = { + readonly findInWorkspace: ( + workspaceId: string, + sourceId: string, + ) => Promise + readonly markFailed: ( + workspaceId: string, + sourceId: string, + reason: string, + requiredStatus?: string, + ) => Promise + readonly markReady: ( + workspaceId: string, + sourceId: string, + documentId: string, + ) => Promise + readonly clearStagedBlob: ( + workspaceId: string, + sourceId: string, + ) => Promise +} + +type SourceReconcileWorkflowBlobStore = { + readonly deleteStagedSourceBlob: (pathname: string) => Promise +} + +export type PollSourceReconciliationInput = { + readonly workspaceId: string + readonly sourceId: string + readonly client: SourceReconcileWorkflowClient + readonly repository?: SourceReconcileWorkflowRepository + readonly blobStore?: SourceReconcileWorkflowBlobStore +} + +export type MarkSourceReadyAfterReconciliationInput = { + readonly workspaceId: string + readonly sourceId: string + readonly documentId: string + readonly repository?: SourceReconcileWorkflowRepository + readonly blobStore?: SourceReconcileWorkflowBlobStore +} + +export type PollSourceReconciliationResult = + | { + readonly kind: "waiting" + readonly jobId: string + readonly jobStatus: string + } + | { + readonly kind: "ready-to-prepare" + readonly jobId: string + readonly documentId: string + } + | { + readonly kind: "resolved" + readonly status: string + } + +export type MarkSourceReadyAfterReconciliationResult = { + readonly status: string +} + +export async function pollSourceReconciliation({ + workspaceId, + sourceId, + client, + repository = sourceWorkflowRuntime, + blobStore = vercelBlobStore, +}: PollSourceReconciliationInput): Promise { + const source = await repository.findInWorkspace(workspaceId, sourceId) + if (!source) return { kind: "resolved", status: "gone" } + if (source.status !== "parsing") { + return { kind: "resolved", status: source.status } + } + + const jobId = source.knowhereJobId + if (!jobId) { + await failSourceAndCleanup({ + workspaceId, + source, + reason: "Parsing source is missing a Knowhere job id.", + repository, + blobStore, + }) + return { kind: "resolved", status: "failed" } + } + + const job = await client.jobs.get(jobId) + if (isFailedJob(job)) { + await failSourceAndCleanup({ + workspaceId, + source, + reason: job.error?.message ?? "Parsing failed.", + repository, + blobStore, + }) + return { kind: "resolved", status: "failed" } + } + + if (!isDoneJob(job)) { + return { + kind: "waiting", + jobId, + jobStatus: job.status, + } + } + + if (!job.documentId) { + await failSourceAndCleanup({ + workspaceId, + source, + reason: "Parsing finished but Knowhere did not publish a document id.", + repository, + blobStore, + }) + return { kind: "resolved", status: "failed" } + } + + if (!job.resultUrl) { + await failSourceAndCleanup({ + workspaceId, + source, + reason: "Parsing finished but Knowhere did not publish a result ZIP.", + repository, + blobStore, + }) + return { kind: "resolved", status: "failed" } + } + + return { + kind: "ready-to-prepare", + jobId, + documentId: job.documentId, + } +} + +export async function markSourceReadyAfterReconciliation({ + workspaceId, + sourceId, + documentId, + repository = sourceWorkflowRuntime, + blobStore = vercelBlobStore, +}: MarkSourceReadyAfterReconciliationInput): Promise { + const source = await repository.findInWorkspace(workspaceId, sourceId) + if (!source) return { status: "gone" } + if (source.status !== "parsing") return { status: source.status } + + const readySource = await repository.markReady( + workspaceId, + sourceId, + documentId, + ) + if (!readySource) return { status: "unchanged" } + + await cleanupStagedBlob({ + workspaceId, + source, + repository, + blobStore, + }) + return { status: readySource.status } +} + +async function failSourceAndCleanup(input: { + readonly workspaceId: string + readonly source: Source + readonly reason: string + readonly repository: SourceReconcileWorkflowRepository + readonly blobStore: SourceReconcileWorkflowBlobStore +}): Promise { + await input.repository.markFailed( + input.workspaceId, + input.source.id, + input.reason, + "parsing", + ) + await cleanupStagedBlob(input) +} + +async function cleanupStagedBlob(input: { + readonly workspaceId: string + readonly source: Source + readonly repository: SourceReconcileWorkflowRepository + readonly blobStore: SourceReconcileWorkflowBlobStore +}): Promise { + if (!input.source.stagedBlobPathname) return + + try { + await input.blobStore.deleteStagedSourceBlob(input.source.stagedBlobPathname) + await input.repository.clearStagedBlob( + input.workspaceId, + input.source.id, + ) + } catch (error) { + logger.warn("source-reconcile-workflow: staged blob cleanup failed", { + workspaceId: input.workspaceId, + sourceId: input.source.id, + error: error instanceof Error ? error.message : String(error), + }) + } +} + +function isDoneJob(job: JobResult): boolean { + return job.isDone || job.status === "done" +} + +function isFailedJob(job: JobResult): boolean { + return job.isFailed || job.status === "failed" +} + +const vercelBlobStore: SourceReconcileWorkflowBlobStore = { + deleteStagedSourceBlob: del, +} diff --git a/src/domains/sources/workflow-runtime.ts b/src/domains/sources/workflow-runtime.ts index e2c6b5f..534c276 100644 --- a/src/domains/sources/workflow-runtime.ts +++ b/src/domains/sources/workflow-runtime.ts @@ -56,6 +56,13 @@ type SourceWorkflowRuntime = UploadRepositoryRuntime & { workspaceId: string, sourceId: string, ) => Promise>> + readonly getParseResultProgress: ( + workspaceId: string, + sourceId: string, + ) => Promise<{ + readonly resultBlobUrl: string + readonly assetUrlsByFilePath: Readonly> + } | null> readonly listForWorkspace: (workspaceId: string) => Promise readonly localizeRemoteDocument: ( workspaceId: string, @@ -76,6 +83,11 @@ type SourceWorkflowRuntime = UploadRepositoryRuntime & { sourceId: string, input: SaveSourceParseResultInput, ) => Promise + readonly mergeParseAssetUrls: ( + workspaceId: string, + sourceId: string, + input: SaveSourceParseResultInput, + ) => Promise readonly softDelete: ( workspaceId: string, sourceId: string, @@ -192,6 +204,21 @@ const saveParseResult: SourceWorkflowRuntime["saveParseResult"] = ( sourceRepository.saveParseResultEffect(workspaceId, sourceId, input), ) +const mergeParseAssetUrls: SourceWorkflowRuntime["mergeParseAssetUrls"] = ( + workspaceId: string, + sourceId: string, + input: SaveSourceParseResultInput, +) => + databaseRuntime.runPromise( + sourceRepository.mergeParseAssetUrlsEffect(workspaceId, sourceId, input), + ) + +const getParseResultProgress: SourceWorkflowRuntime["getParseResultProgress"] = + (workspaceId: string, sourceId: string) => + databaseRuntime.runPromise( + sourceRepository.getParseResultProgressEffect(workspaceId, sourceId), + ) + const getParseAssetUrls: SourceWorkflowRuntime["getParseAssetUrls"] = ( workspaceId: string, sourceId: string, @@ -238,6 +265,7 @@ export const sourceWorkflowRuntime: SourceWorkflowRuntime = { createUploading, findInWorkspace, getParseAssetUrls, + getParseResultProgress, hideDemoSource, listForWorkspace, listHiddenDemoSourceIds, @@ -245,6 +273,7 @@ export const sourceWorkflowRuntime: SourceWorkflowRuntime = { markFailed, markParsing, markReady, + mergeParseAssetUrls, saveParseResult, softDelete, upsertMaterializedDemoSource, diff --git a/src/domains/workspace/initial-state.test.ts b/src/domains/workspace/initial-state.test.ts index 6361411..544b74e 100644 --- a/src/domains/workspace/initial-state.test.ts +++ b/src/domains/workspace/initial-state.test.ts @@ -383,20 +383,16 @@ describe("loadWorkspaceShellInitialState", () => { ]) }) - it("reconciles parsing sources before localizing matching Knowhere documents", async () => { + it("keeps matching Notebook uploads parsing while background reconciliation prepares artifacts", async () => { const workspace = makeWorkspace() const parsingSource = makeSource(workspace.id, { + title: "uploaded.pdf", + mimeType: "application/pdf", + sizeBytes: 100, status: "parsing", knowhereJobId: "job_123", knowhereDocumentId: null, }) - const reconciledSource = makeSource(workspace.id, { - id: parsingSource.id, - title: "uploaded.pdf", - status: "ready", - knowhereJobId: null, - knowhereDocumentId: "doc_uploaded", - }) const client = { documents: { list: vi @@ -408,6 +404,12 @@ describe("loadWorkspaceShellInitialState", () => { namespace: "default", status: "active", sourceFileName: "uploaded.pdf", + documentMetadata: { + createdByClient: "notebook", + title: "uploaded.pdf", + mimeType: "application/pdf", + sizeBytes: 100, + }, }, ], }) @@ -427,8 +429,9 @@ describe("loadWorkspaceShellInitialState", () => { load: vi.fn(), }, } as unknown as InitialStateClient - const reconcileSourcesForWorkspace = vi.fn(async () => [reconciledSource]) - const localizeRemoteDocument = vi.fn(async () => reconciledSource) + const reconcileSourcesForWorkspace = vi.fn(async () => [parsingSource]) + const localizeRemoteDocument = vi.fn(async () => parsingSource) + const startBackgroundReconciliation = vi.fn(async () => undefined) const deps = createDependencies({ getClientForWorkspace: vi.fn(async () => ({ client, apiKey: "sk_test" })), getOptionalAuthenticated: vi.fn(async () => ({ @@ -442,20 +445,18 @@ describe("loadWorkspaceShellInitialState", () => { listSourcesForWorkspace: vi.fn(async () => [parsingSource]), reconcileSourcesForWorkspace, localizeRemoteDocument, + startBackgroundReconciliation, }) const state = await loadWorkspaceShellInitialState(deps) - expect(reconcileSourcesForWorkspace).toHaveBeenCalledWith( - workspace, - client, - ) - expect(localizeRemoteDocument).toHaveBeenCalledWith( + expect(reconcileSourcesForWorkspace).not.toHaveBeenCalled() + expect(startBackgroundReconciliation).toHaveBeenCalledWith( workspace.id, - expect.objectContaining({ - documentId: "doc_uploaded", - }), + parsingSource.id, + "sk_test", ) + expect(localizeRemoteDocument).not.toHaveBeenCalled() expect(state.sources).toEqual([ expect.objectContaining({ id: "demo-tsla-q4-2025", @@ -466,8 +467,8 @@ describe("loadWorkspaceShellInitialState", () => { kind: "workspace", title: "uploaded.pdf", mimeType: "application/pdf", - status: "ready", - documentId: "doc_uploaded", + status: "parsing", + documentId: undefined, }, ]) }) diff --git a/src/domains/workspace/initial-state.ts b/src/domains/workspace/initial-state.ts index 1a73f83..abd7278 100644 --- a/src/domains/workspace/initial-state.ts +++ b/src/domains/workspace/initial-state.ts @@ -16,7 +16,9 @@ import { sourceViewOptionsBySourceId as getSourceViewOptionsBySourceId } from "@ import { localizeRemoteLibrarySources } from "@/domains/sources/remote-library" import { reconcileSourcesForWorkspace as reconcileDefaultSourcesForWorkspace } from "@/domains/sources/reconcile" import { sourceService } from "@/domains/sources/service" -import { startBackgroundReconciliation } from "@/domains/sources/background-reconcile" +import { + startBackgroundReconciliation as defaultStartBackgroundReconciliation, +} from "@/domains/sources/background-reconcile" import { sourceWorkflowRuntime } from "@/domains/sources/workflow-runtime" import type { @@ -37,6 +39,7 @@ import { type OfficialLibrarySource, } from "@/integrations/knowhere-demo" import { effectOperation } from "@/lib/effect-operation" +import { logger } from "@/lib/logger" import { notebookRequestContext } from "./request-context" type WorkspaceShellInitialState = { @@ -142,6 +145,7 @@ type WorkspaceShellInitialStateDependencies = { workspace: Workspace, client: WorkspaceShellInitialStateClient, ) => Promise + readonly startBackgroundReconciliation?: typeof defaultStartBackgroundReconciliation readonly localizeRemoteDocument: ( workspaceId: string, input: Parameters[1], @@ -163,6 +167,7 @@ const defaultDependencies: WorkspaceShellInitialStateDependencies = { listMessages: chatThreadService.listMessages, listSourcesForWorkspace: sourceWorkflowRuntime.listForWorkspace, reconcileSourcesForWorkspace: reconcileDefaultSourcesForWorkspace, + startBackgroundReconciliation: defaultStartBackgroundReconciliation, localizeRemoteDocument: sourceService.localizeRemoteDocument, sourceViewOptionsBySourceId: getSourceViewOptionsBySourceId, } @@ -302,12 +307,9 @@ export const loadWorkspaceShellInitialStateEffect = ( const sources = yield* effectOperation.tryPromise( { context: workspaceInitialStateContext, - operation: "reconcileSourcesForWorkspace", + operation: "useListedSourcesForWorkspace", }, - () => - hasParsingSources(listedSources) - ? deps.reconcileSourcesForWorkspace(workspace, client) - : Promise.resolve(listedSources), + () => Promise.resolve(listedSources), ) const demoSourceResolution = resolveWorkspaceDemoSources( sources, @@ -342,19 +344,16 @@ export const loadWorkspaceShellInitialStateEffect = ( workspaceSources, demoCatalog, ) - for (const source of sources) { - if (source.status === "parsing" && source.knowhereJobId) { - yield* Effect.fork( - effectOperation.tryPromise( - { - context: workspaceInitialStateContext, - operation: "startBackgroundReconciliation", - }, - () => startBackgroundReconciliation(workspace.id, source.id, apiKey), - ), - ) - } - } + yield* Effect.sync(() => + triggerBackgroundReconciliationForParsingSources({ + workspaceId: workspace.id, + sources, + apiKey, + startBackgroundReconciliation: + deps.startBackgroundReconciliation ?? + defaultStartBackgroundReconciliation, + }), + ) const sourceOptions = yield* effectOperation.addContext( { context: workspaceInitialStateContext, @@ -412,15 +411,34 @@ function resolveDashboardUrl(): string | undefined { return process.env.DASHBOARD_ORIGIN } -function hasParsingSources( - sources: readonly { - readonly status: string - readonly knowhereJobId: string | null - }[], -): boolean { - return sources.some( +function triggerBackgroundReconciliationForParsingSources(input: { + readonly workspaceId: string + readonly sources: readonly Source[] + readonly apiKey: string + readonly startBackgroundReconciliation: typeof defaultStartBackgroundReconciliation +}): void { + const parsingSources = input.sources.filter( (source) => source.status === "parsing" && source.knowhereJobId, ) + if (parsingSources.length === 0) return + + logger.info("initial-state: re-triggering reconciliation for parsing sources", { + workspaceId: input.workspaceId, + count: parsingSources.length, + sourceIds: parsingSources.map((source) => source.id), + }) + + for (const source of parsingSources) { + void input + .startBackgroundReconciliation(input.workspaceId, source.id, input.apiKey) + .catch((error: unknown) => { + logger.warn("initial-state: background reconciliation trigger failed", { + workspaceId: input.workspaceId, + sourceId: source.id, + error: error instanceof Error ? error.message : String(error), + }) + }) + } } function toOfficialLibrarySourceViews( From 50990f6fe01ff35658fdf472bb0a4954ac7e508a Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 30 Jun 2026 10:57:02 +0800 Subject: [PATCH 2/6] Fix source reconciliation workflow step limits --- src/app/api/sources/reconcile/route.ts | 181 +------ src/domains/chat/route-answer.ts | 46 +- src/domains/chat/route-service.test.ts | 64 ++- .../sources/parsed-result-assets.test.ts | 76 ++- src/domains/sources/parsed-result-assets.ts | 169 ++++++- .../source-parse-result-repository.test.ts | 20 + .../sources/source-parse-result-repository.ts | 27 +- .../source-reconcile-route-workflow.test.ts | 357 +++++++++++++ .../source-reconcile-route-workflow.ts | 467 ++++++++++++++++++ src/domains/sources/source-row-repository.ts | 4 +- src/domains/sources/workflow-runtime.ts | 3 + 11 files changed, 1195 insertions(+), 219 deletions(-) create mode 100644 src/domains/sources/source-parse-result-repository.test.ts create mode 100644 src/domains/sources/source-reconcile-route-workflow.test.ts create mode 100644 src/domains/sources/source-reconcile-route-workflow.ts diff --git a/src/app/api/sources/reconcile/route.ts b/src/app/api/sources/reconcile/route.ts index a9822b0..223d380 100644 --- a/src/app/api/sources/reconcile/route.ts +++ b/src/app/api/sources/reconcile/route.ts @@ -1,164 +1,35 @@ import { serve } from "@upstash/workflow/nextjs" -import { - completeResultZipMultipartUpload, - createResultZipMultipartUploadPlan, - getResultZipPartRange, - prepareParsedResultAssetBatch, - uploadResultZipPart, - type MultipartUploadPart, -} from "@/domains/sources/parsed-result-assets" -import { - markSourceReadyAfterReconciliation, - pollSourceReconciliation, -} from "@/domains/sources/source-reconcile-workflow" -import { sourceWorkflowRuntime } from "@/domains/sources/workflow-runtime" -import { makeKnowhereClient } from "@/integrations/knowhere" -import { logger } from "@/lib/logger" +import { sourceReconcileRouteWorkflow } from "@/domains/sources/source-reconcile-route-workflow" -type ReconcilePayload = { - readonly workspaceId: string - readonly sourceId: string - readonly apiKey: string -} +type ReconcilePayload = Parameters< + typeof sourceReconcileRouteWorkflow.normalizeReconcilePayload +>[0] -const MAX_POLL_ATTEMPTS = 60 -const INITIAL_DELAY_S = 3 -const MAX_DELAY_S = 30 - -export const { POST } = serve(async (context) => { - const { workspaceId, sourceId, apiKey } = context.requestPayload - const client = makeKnowhereClient(apiKey) - let delay = INITIAL_DELAY_S - let completedJob: { - readonly jobId: string - readonly documentId: string - } | null = null - - for (let attempt = 0; attempt < MAX_POLL_ATTEMPTS; attempt++) { - const poll = await context.run(`poll-${attempt}`, async () => { - return pollSourceReconciliation({ - workspaceId, - sourceId, - client, - }) - }) - - if (poll.kind === "ready-to-prepare") { - completedJob = { - jobId: poll.jobId, - documentId: poll.documentId, - } - logger.info("workflow: source parse completed; preparing artifacts", { - sourceId, - jobId: poll.jobId, - attempts: attempt + 1, - }) - break - } - - if (poll.kind === "resolved") { - logger.info("workflow: source resolved", { - sourceId, - status: poll.status, - attempts: attempt + 1, +export const { POST } = serve( + async (context) => { + const payload = sourceReconcileRouteWorkflow.normalizeReconcilePayload( + context.requestPayload, + ) + if (payload.phase === "asset-batches") { + await sourceReconcileRouteWorkflow.runAssetBatchWorkflow({ + context, + payload, }) return } - await context.sleep(`wait-${attempt}`, delay) - delay = Math.min(Math.round(delay * 1.5), MAX_DELAY_S) - } - - const jobToPrepare = completedJob - if (!jobToPrepare) { - logger.error("workflow: exhausted poll attempts", { - sourceId, - maxAttempts: MAX_POLL_ATTEMPTS, + await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context, + payload, }) - return - } - - const uploadPlan = await context.run("mirror-zip-start", async () => - createResultZipMultipartUploadPlan({ - workspaceId, - sourceId, - jobId: jobToPrepare.jobId, - client, - }), - ) - const parts: MultipartUploadPart[] = [] - - for (let partNumber = 1; partNumber <= uploadPlan.partCount; partNumber++) { - const range = getResultZipPartRange( - partNumber, - uploadPlan.partSizeBytes, - uploadPlan.sizeBytes, - ) - const part = await context.run(`mirror-zip-part-${partNumber}`, async () => - uploadResultZipPart({ - pathname: uploadPlan.pathname, - key: uploadPlan.key, - uploadId: uploadPlan.uploadId, - partNumber, - jobId: jobToPrepare.jobId, - client, - ...range, - }), - ) - parts.push(part) - } - - const resultBlobUrl = await context.run("mirror-zip-complete", async () => { - const result = await completeResultZipMultipartUpload({ - pathname: uploadPlan.pathname, - key: uploadPlan.key, - uploadId: uploadPlan.uploadId, - parts, - }) - return result.url - }) - - await context.run("parse-result-save-zip", async () => { - const saved = await sourceWorkflowRuntime.mergeParseAssetUrls( - workspaceId, - sourceId, - { - resultBlobUrl, - assetUrlsByFilePath: {}, - }, - ) - if (!saved) throw new Error("Source disappeared before saving result ZIP.") - }) - - let batchIndex = 0 - for (;;) { - const batch = await context.run( - `parse-assets-batch-${batchIndex}`, - async () => - prepareParsedResultAssetBatch({ - workspaceId, - sourceId, - jobId: jobToPrepare.jobId, - resultBlobUrl, - client, - repository: sourceWorkflowRuntime, - }), - ) - if (!batch.hasMore) break - batchIndex += 1 - } - - const ready = await context.run("source-ready", async () => - markSourceReadyAfterReconciliation({ - workspaceId, - sourceId, - documentId: jobToPrepare.documentId, - }), - ) - logger.info("workflow: source artifact preparation finished", { - sourceId, - status: ready.status, - assetBatches: batchIndex + 1, - }) -}) + }, + { + failureFunction: async ({ context, failResponse }) => { + await sourceReconcileRouteWorkflow.markSourceFailedAfterWorkflowFailure( + context.requestPayload, + failResponse, + ) + }, + }, +) diff --git a/src/domains/chat/route-answer.ts b/src/domains/chat/route-answer.ts index 2006be6..06c643a 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -11,10 +11,12 @@ import { type ChatTurnValue, } from "@/domains/chat/service" import { chatTurnPersistence } from "@/domains/chat/chat-turn-persistence" -import { reconcileSourcesForWorkspace } from "@/domains/sources/reconcile" +import { startBackgroundReconciliation } from "@/domains/sources/background-reconcile" import { localizeRemoteLibrarySources } from "@/domains/sources/remote-library" import { sourceService } from "@/domains/sources/service" +import { sourceWorkflowRuntime } from "@/domains/sources/workflow-runtime" import { notebookRequestContext } from "@/domains/workspace/request-context" +import type { Source } from "@/infrastructure/db/schema" import { isAuthError } from "@/integrations/dashboard/api-key-service" import { summarizeUnknownError } from "@/lib/format-log-value" import { logger } from "@/lib/logger" @@ -54,11 +56,18 @@ const answerChatEffect = (input: AnswerChatInput) => return routeResult.error(body.status, body.message) } - const { workspace, client } = yield* Effect.tryPromise(() => + const { workspace, client, apiKey } = yield* Effect.tryPromise(() => notebookRequestContext.getAuthenticatedWithClient(), ) const sources = yield* Effect.tryPromise(() => - reconcileSourcesForWorkspace(workspace, client), + sourceWorkflowRuntime.listForWorkspace(workspace.id), + ) + yield* Effect.sync(() => + triggerBackgroundReconciliationForParsingSources({ + workspaceId: workspace.id, + sources, + apiKey, + }), ) const compatibleSources = yield* localizeRemoteLibrarySources({ workspace, @@ -143,6 +152,37 @@ export const chatAnswerRouteService: ChatAnswerRouteService = { answerChat, } +function triggerBackgroundReconciliationForParsingSources(input: { + readonly workspaceId: string + readonly sources: readonly Source[] + readonly apiKey: string +}): void { + const parsingSources = input.sources.filter( + (source) => source.status === "parsing" && source.knowhereJobId, + ) + if (parsingSources.length === 0) return + + logger.info("chat: re-triggering reconciliation for parsing sources", { + workspaceId: input.workspaceId, + count: parsingSources.length, + sourceIds: parsingSources.map((source) => source.id), + }) + + for (const source of parsingSources) { + void startBackgroundReconciliation( + input.workspaceId, + source.id, + input.apiKey, + ).catch((error: unknown) => { + logger.warn("chat: background reconciliation trigger failed", { + workspaceId: input.workspaceId, + sourceId: source.id, + error: error instanceof Error ? error.message : String(error), + }) + }) + } +} + function toChatRouteFailure(detail: string): ChatRouteFailure { const routeDetail = getSafeRouteDetail(detail) if (isAuthError({ message: routeDetail })) { diff --git a/src/domains/chat/route-service.test.ts b/src/domains/chat/route-service.test.ts index f24e5cd..4358e8a 100644 --- a/src/domains/chat/route-service.test.ts +++ b/src/domains/chat/route-service.test.ts @@ -17,8 +17,9 @@ const mocks = vi.hoisted(() => ({ loggerError: vi.fn(), loggerInfo: vi.fn(), loggerWarn: vi.fn(), - reconcileSourcesForWorkspace: vi.fn(), + listSourcesForWorkspace: vi.fn(), softDeleteChatThread: vi.fn(), + startBackgroundReconciliation: vi.fn(), })) vi.mock("@/domains/chat", async (importOriginal) => { @@ -33,8 +34,14 @@ vi.mock("@/domains/chat/service", () => ({ handleChatTurn: mocks.handleChatTurn, })) -vi.mock("@/domains/sources/reconcile", () => ({ - reconcileSourcesForWorkspace: mocks.reconcileSourcesForWorkspace, +vi.mock("@/domains/sources/background-reconcile", () => ({ + startBackgroundReconciliation: mocks.startBackgroundReconciliation, +})) + +vi.mock("@/domains/sources/workflow-runtime", () => ({ + sourceWorkflowRuntime: { + listForWorkspace: mocks.listSourcesForWorkspace, + }, })) vi.mock("@/domains/workspace/request-context", () => ({ @@ -82,7 +89,7 @@ describe("chat route services", () => { apiKey: "jwt_123", client, }) - mocks.reconcileSourcesForWorkspace.mockResolvedValue([readySource]) + mocks.listSourcesForWorkspace.mockResolvedValue([readySource]) mocks.handleChatTurn.mockResolvedValue( Either.right({ threadId: "thread_1", @@ -111,10 +118,8 @@ describe("chat route services", () => { ], }, }) - expect(mocks.reconcileSourcesForWorkspace).toHaveBeenCalledWith( - workspace, - client, - ) + expect(mocks.listSourcesForWorkspace).toHaveBeenCalledWith(workspace.id) + expect(mocks.startBackgroundReconciliation).not.toHaveBeenCalled() expect(mocks.handleChatTurn).toHaveBeenCalledWith( expect.objectContaining({ workspace, @@ -135,6 +140,45 @@ describe("chat route services", () => { ) }) + it("triggers background reconciliation for parsing sources without blocking chat", async () => { + const workspace = makeWorkspace() + const client = { retrieval: { query: vi.fn() } } + const parsingSource = makeSource({ + status: "parsing", + knowhereDocumentId: null, + }) + mocks.getAuthenticatedWithClient.mockResolvedValue({ + user: { id: "user_1" }, + workspace, + apiKey: "jwt_123", + client, + }) + mocks.listSourcesForWorkspace.mockResolvedValue([parsingSource]) + mocks.startBackgroundReconciliation.mockResolvedValue(undefined) + mocks.handleChatTurn.mockResolvedValue( + Either.right({ + threadId: "thread_1", + messages: [], + }), + ) + + const result = await chatAnswerRouteService.answerChat({ + body: { message: "Summarize it" }, + }) + + expect(result.status).toBe(200) + expect(mocks.startBackgroundReconciliation).toHaveBeenCalledWith( + workspace.id, + parsingSource.id, + "jwt_123", + ) + expect(mocks.handleChatTurn).toHaveBeenCalledWith( + expect.objectContaining({ + sources: [parsingSource], + }), + ) + }) + it("returns an explicit generation failure instead of a fake session error", async () => { const workspace = makeWorkspace() const client = { retrieval: { query: vi.fn() } } @@ -144,7 +188,7 @@ describe("chat route services", () => { apiKey: "jwt_123", client, }) - mocks.reconcileSourcesForWorkspace.mockResolvedValue([makeSource()]) + mocks.listSourcesForWorkspace.mockResolvedValue([makeSource()]) mocks.handleChatTurn.mockRejectedValue( new Error("Gateway rejected tool schema: dataType enum invalid"), ) @@ -178,7 +222,7 @@ describe("chat route services", () => { apiKey: "jwt_123", client, }) - mocks.reconcileSourcesForWorkspace.mockResolvedValue([makeSource()]) + mocks.listSourcesForWorkspace.mockResolvedValue([makeSource()]) mocks.handleChatTurn.mockRejectedValue( new Error("HTTP 401: invalid API key"), ) diff --git a/src/domains/sources/parsed-result-assets.test.ts b/src/domains/sources/parsed-result-assets.test.ts index 6db63c7..db7795e 100644 --- a/src/domains/sources/parsed-result-assets.test.ts +++ b/src/domains/sources/parsed-result-assets.test.ts @@ -198,27 +198,35 @@ describe("result ZIP multipart mirroring", () => { }) expect(fetchResult).toHaveBeenCalledWith( "https://knowhere.example/result.zip", - { method: "HEAD" }, + expect.objectContaining({ + method: "HEAD", + signal: expect.any(AbortSignal), + }), ) expect(fetchResult).toHaveBeenCalledWith( "https://knowhere.example/result.zip", - { headers: { Range: `bytes=0-${partSizeBytes - 1}` } }, + expect.objectContaining({ + headers: { Range: `bytes=0-${partSizeBytes - 1}` }, + signal: expect.any(AbortSignal), + }), ) expect(fetchResult).toHaveBeenCalledWith( "https://knowhere.example/result.zip", - { + expect.objectContaining({ headers: { Range: `bytes=${partSizeBytes}-${partSizeBytes * 2 - 1}`, }, - }, + signal: expect.any(AbortSignal), + }), ) expect(fetchResult).toHaveBeenCalledWith( "https://knowhere.example/result.zip", - { + expect.objectContaining({ headers: { Range: `bytes=${partSizeBytes * 2}-${partSizeBytes * 2 + 2}`, }, - }, + signal: expect.any(AbortSignal), + }), ) expect(blobStore.uploadPart).toHaveBeenCalledTimes(3) expect(blobStore.completeMultipartUpload).toHaveBeenCalledWith( @@ -235,6 +243,62 @@ describe("result ZIP multipart mirroring", () => { ) expect(completed.url).toBe("https://blob.example/result.zip") }) + + it("falls back to a one-byte range request when HEAD does not expose ZIP size", async () => { + const sizeBytes = 12_345 + const fetchResult = vi.fn(async (_url: string | URL, init?: RequestInit) => { + if (init?.method === "HEAD") { + return new Response(null, { status: 405 }) + } + + return new Response(Buffer.from("x"), { + status: 206, + headers: { + "content-range": `bytes 0-0/${sizeBytes}`, + }, + }) + }) + const client = { + jobs: { + get: vi.fn(async () => ({ + resultUrl: "https://knowhere.example/result.zip", + })), + }, + } + const blobStore = { + createMultipartUpload: vi.fn(async () => ({ + key: "blob-key", + uploadId: "upload-1", + })), + uploadPart: vi.fn(), + completeMultipartUpload: vi.fn(), + } + + const plan = await createResultZipMultipartUploadPlan({ + workspaceId: "workspace_1", + sourceId: "source_1", + jobId: "job_1", + client, + blobStore, + fetchResult, + }) + + expect(fetchResult).toHaveBeenCalledWith( + "https://knowhere.example/result.zip", + expect.objectContaining({ + method: "HEAD", + signal: expect.any(AbortSignal), + }), + ) + expect(fetchResult).toHaveBeenCalledWith( + "https://knowhere.example/result.zip", + expect.objectContaining({ + headers: { Range: "bytes=0-0" }, + signal: expect.any(AbortSignal), + }), + ) + expect(plan.sizeBytes).toBe(sizeBytes) + }) }) describe("prepareParsedResultAssetBatch", () => { diff --git a/src/domains/sources/parsed-result-assets.ts b/src/domains/sources/parsed-result-assets.ts index 97ee62a..6a6a727 100644 --- a/src/domains/sources/parsed-result-assets.ts +++ b/src/domains/sources/parsed-result-assets.ts @@ -10,6 +10,8 @@ import { import { Effect } from "effect" import type { JobResult } from "@ontos-ai/knowhere-sdk" +import { logger } from "@/lib/logger" + export type StoredParsedResultAssets = { resultBlobUrl: string assetUrlsByFilePath: Readonly> @@ -197,6 +199,10 @@ const parsedResultDirectoryName = "parsed-result" const resultZipPartSizeBytes = 8 * 1024 * 1024 const parsedAssetBatchLimit = 10 const parsedAssetBatchMaxDurationMs = 45_000 +const resultZipMetadataTimeoutMs = 15_000 +const resultZipRangeFetchTimeoutMs = 30_000 +const resultZipBlobOperationTimeoutMs = 30_000 +const parsedAssetBlobOperationTimeoutMs = 30_000 // --------------------------------------------------------------------------- // Effect core @@ -296,10 +302,14 @@ export const prepareParsedResultAssetBatchEffect = Effect.fn( if (uploadedCount > 0 && Date.now() - startedAt >= maxDurationMs) break const blob = yield* Effect.tryPromise(() => - blobStore.put( - `${blobPrefix}/${asset.filePath}`, - asset.body, - getBlobPutOptions(asset.contentType), + withTimeout( + `Parsed asset upload for ${asset.filePath}`, + parsedAssetBlobOperationTimeoutMs, + blobStore.put( + `${blobPrefix}/${asset.filePath}`, + asset.body, + getBlobPutOptions(asset.contentType), + ), ), ) uploadedAssetUrls[asset.filePath] = blob.url @@ -361,16 +371,18 @@ export async function createResultZipMultipartUploadPlan({ fetchResult = fetch, partSizeBytes = resultZipPartSizeBytes, }: CreateResultZipMultipartUploadPlanInput): Promise { + const startedAt = Date.now() const job = await client.jobs.get(jobId) const resultUrl = requireJobResultUrl(job, jobId) const sizeBytes = await getResultZipSize(resultUrl, fetchResult) const pathname = `${getParsedResultBlobPrefix(workspaceId, sourceId)}/result.zip` - const upload = await blobStore.createMultipartUpload( - pathname, - getResultZipBlobOptions(), + const upload = await withTimeout( + "Create result ZIP multipart upload", + resultZipBlobOperationTimeoutMs, + blobStore.createMultipartUpload(pathname, getResultZipBlobOptions()), ) - return { + const uploadPlan: ResultZipMultipartUploadPlan = { pathname, key: upload.key, uploadId: upload.uploadId, @@ -378,6 +390,15 @@ export async function createResultZipMultipartUploadPlan({ partSizeBytes, partCount: Math.max(1, Math.ceil(sizeBytes / partSizeBytes)), } + logger.info("parsed-result-assets: result ZIP upload plan created", { + workspaceId, + sourceId, + jobId, + sizeBytes, + partCount: uploadPlan.partCount, + durationMs: Date.now() - startedAt, + }) + return uploadPlan } export function getResultZipPartRange( @@ -420,6 +441,7 @@ export async function uploadResultZipPart({ blobStore = vercelMultipartBlobStore, fetchResult = fetch, }: UploadResultZipPartInput): Promise { + const startedAt = Date.now() const job = await client.jobs.get(jobId) const resultUrl = requireJobResultUrl(job, jobId) const body = await fetchResultZipRange({ @@ -429,12 +451,25 @@ export async function uploadResultZipPart({ fetchResult, }) - return blobStore.uploadPart(pathname, body, { - ...getResultZipBlobOptions(), - key, - uploadId, + const part = await withTimeout( + `Upload result ZIP part ${partNumber}`, + resultZipBlobOperationTimeoutMs, + blobStore.uploadPart(pathname, body, { + ...getResultZipBlobOptions(), + key, + uploadId, + partNumber, + }), + ) + logger.info("parsed-result-assets: result ZIP part uploaded", { + jobId, partNumber, + startByte, + endByte, + byteLength: body.byteLength, + durationMs: Date.now() - startedAt, }) + return part } export async function completeResultZipMultipartUpload({ @@ -444,15 +479,25 @@ export async function completeResultZipMultipartUpload({ parts, blobStore = vercelMultipartBlobStore, }: CompleteResultZipMultipartUploadInput): Promise<{ readonly url: string }> { - return blobStore.completeMultipartUpload( - pathname, - [...parts].sort((a, b) => a.partNumber - b.partNumber), - { - ...getResultZipBlobOptions(), - key, - uploadId, - }, + const startedAt = Date.now() + const result = await withTimeout( + "Complete result ZIP multipart upload", + resultZipBlobOperationTimeoutMs, + blobStore.completeMultipartUpload( + pathname, + [...parts].sort((a, b) => a.partNumber - b.partNumber), + { + ...getResultZipBlobOptions(), + key, + uploadId, + }, + ), ) + logger.info("parsed-result-assets: result ZIP multipart upload completed", { + partCount: parts.length, + durationMs: Date.now() - startedAt, + }) + return result } export async function prepareParsedResultAssetBatch( @@ -555,7 +600,13 @@ async function getResultZipSize( resultUrl: string, fetchResult: ResultZipFetch, ): Promise { - const headResponse = await fetchResult(resultUrl, { method: "HEAD" }) + const headResponse = await fetchWithTimeout({ + description: "Knowhere result ZIP HEAD request", + input: resultUrl, + init: { method: "HEAD" }, + timeoutMs: resultZipMetadataTimeoutMs, + fetchResult, + }) if (headResponse.ok) { const contentLength = getPositiveIntegerHeader( headResponse.headers, @@ -564,10 +615,16 @@ async function getResultZipSize( if (contentLength !== null) return contentLength } - const rangeResponse = await fetchResult(resultUrl, { - headers: { - Range: "bytes=0-0", + const rangeResponse = await fetchWithTimeout({ + description: "Knowhere result ZIP size range request", + input: resultUrl, + init: { + headers: { + Range: "bytes=0-0", + }, }, + timeoutMs: resultZipMetadataTimeoutMs, + fetchResult, }) await rangeResponse.body?.cancel() @@ -593,10 +650,16 @@ async function fetchResultZipRange(input: { readonly fetchResult: ResultZipFetch }): Promise { const expectedByteLength = input.endByte - input.startByte + 1 - const response = await input.fetchResult(input.resultUrl, { - headers: { - Range: `bytes=${input.startByte}-${input.endByte}`, + const response = await fetchWithTimeout({ + description: "Knowhere result ZIP range fetch", + input: input.resultUrl, + init: { + headers: { + Range: `bytes=${input.startByte}-${input.endByte}`, + }, }, + timeoutMs: resultZipRangeFetchTimeoutMs, + fetchResult: input.fetchResult, }) if (!response.ok) { @@ -617,6 +680,58 @@ async function fetchResultZipRange(input: { return body } +async function fetchWithTimeout(input: { + readonly description: string + readonly input: string | URL + readonly init?: RequestInit + readonly timeoutMs: number + readonly fetchResult: ResultZipFetch +}): Promise { + const controller = new AbortController() + const timeoutId = setTimeout(() => { + controller.abort(new Error(`${input.description} timed out.`)) + }, input.timeoutMs) + + try { + return await input.fetchResult(input.input, { + ...input.init, + signal: controller.signal, + }) + } catch (error) { + if (controller.signal.aborted) { + throw new Error( + `${input.description} timed out after ${input.timeoutMs}ms.`, + ) + } + throw error + } finally { + clearTimeout(timeoutId) + } +} + +function withTimeout( + description: string, + timeoutMs: number, + promise: Promise, +): Promise { + return new Promise((resolve, reject) => { + const timeoutId = setTimeout(() => { + reject(new Error(`${description} timed out after ${timeoutMs}ms.`)) + }, timeoutMs) + + promise.then( + (value) => { + clearTimeout(timeoutId) + resolve(value) + }, + (error: unknown) => { + clearTimeout(timeoutId) + reject(error) + }, + ) + }) +} + function getPositiveIntegerHeader( headers: Headers, name: string, diff --git a/src/domains/sources/source-parse-result-repository.test.ts b/src/domains/sources/source-parse-result-repository.test.ts new file mode 100644 index 0000000..299ceea --- /dev/null +++ b/src/domains/sources/source-parse-result-repository.test.ts @@ -0,0 +1,20 @@ +import { describe, expect, it } from "vitest" +import { PgDialect } from "drizzle-orm/pg-core" + +import { buildAtomicAssetUrlsMergeSql } from "./source-parse-result-repository" + +describe("sourceParseResultRepository", () => { + it("builds an atomic JSONB merge expression for asset progress", () => { + const mergeSql = buildAtomicAssetUrlsMergeSql({ + "images/image-2.png": "https://blob.example/images/image-2.png", + }) + const query = new PgDialect().sqlToQuery(mergeSql) + + expect(query.sql).toContain( + '"source_parse_results"."asset_urls" || $1::jsonb', + ) + expect(query.params).toContain( + '{"images/image-2.png":"https://blob.example/images/image-2.png"}', + ) + }) +}) diff --git a/src/domains/sources/source-parse-result-repository.ts b/src/domains/sources/source-parse-result-repository.ts index ee899c7..8a5711e 100644 --- a/src/domains/sources/source-parse-result-repository.ts +++ b/src/domains/sources/source-parse-result-repository.ts @@ -41,6 +41,12 @@ type SourceParseResultRepository = { ) => Effect.Effect>, never, DbClient> } +export function buildAtomicAssetUrlsMergeSql( + assetUrlsByFilePath: Readonly>, +) { + return sql`${sourceParseResults.assetUrls} || ${JSON.stringify(assetUrlsByFilePath)}::jsonb` +} + const saveParseResultEffect: SourceParseResultRepository["saveParseResultEffect"] = (workspaceId: string, sourceId: string, input: SaveSourceParseResultInput) => Effect.gen(function* () { @@ -81,34 +87,21 @@ const mergeParseAssetUrlsEffect: SourceParseResultRepository["mergeParseAssetUrl ) if (!source) return null - const [current] = yield* Effect.promise(() => - db - .select({ - resultBlobUrl: sourceParseResults.resultBlobUrl, - assetUrls: sourceParseResults.assetUrls, - }) - .from(sourceParseResults) - .where(eq(sourceParseResults.sourceId, sourceId)) - .limit(1), - ) - const assetUrlsByFilePath = { - ...(current?.assetUrls ?? {}), - ...input.assetUrlsByFilePath, - } - const [result] = yield* Effect.promise(() => db .insert(sourceParseResults) .values({ sourceId, resultBlobUrl: input.resultBlobUrl, - assetUrls: assetUrlsByFilePath, + assetUrls: input.assetUrlsByFilePath, }) .onConflictDoUpdate({ target: sourceParseResults.sourceId, set: { resultBlobUrl: input.resultBlobUrl, - assetUrls: assetUrlsByFilePath, + assetUrls: buildAtomicAssetUrlsMergeSql( + input.assetUrlsByFilePath, + ), updatedAt: sql`now()`, }, }) diff --git a/src/domains/sources/source-reconcile-route-workflow.test.ts b/src/domains/sources/source-reconcile-route-workflow.test.ts new file mode 100644 index 0000000..8cb69f7 --- /dev/null +++ b/src/domains/sources/source-reconcile-route-workflow.test.ts @@ -0,0 +1,357 @@ +import { afterEach, describe, expect, it, vi } from "vitest" + +import type { Source } from "@/infrastructure/db/schema" + +const mocks = vi.hoisted(() => ({ + completeResultZipMultipartUpload: vi.fn(), + createResultZipMultipartUploadPlan: vi.fn(), + getResultZipPartRange: vi.fn(), + loggerError: vi.fn(), + loggerInfo: vi.fn(), + loggerWarn: vi.fn(), + makeKnowhereClient: vi.fn(), + markSourceReadyAfterReconciliation: vi.fn(), + pollSourceReconciliation: vi.fn(), + prepareParsedResultAssetBatch: vi.fn(), + uploadResultZipPart: vi.fn(), + findInWorkspace: vi.fn(), + getParseResultProgress: vi.fn(), + markFailed: vi.fn(), + markParsing: vi.fn(), + mergeParseAssetUrls: vi.fn(), +})) + +vi.mock("@/domains/sources/parsed-result-assets", () => ({ + completeResultZipMultipartUpload: mocks.completeResultZipMultipartUpload, + createResultZipMultipartUploadPlan: mocks.createResultZipMultipartUploadPlan, + getResultZipPartRange: mocks.getResultZipPartRange, + prepareParsedResultAssetBatch: mocks.prepareParsedResultAssetBatch, + uploadResultZipPart: mocks.uploadResultZipPart, +})) + +vi.mock("@/domains/sources/source-reconcile-workflow", () => ({ + markSourceReadyAfterReconciliation: mocks.markSourceReadyAfterReconciliation, + pollSourceReconciliation: mocks.pollSourceReconciliation, +})) + +vi.mock("@/domains/sources/workflow-runtime", () => ({ + sourceWorkflowRuntime: { + findInWorkspace: mocks.findInWorkspace, + getParseResultProgress: mocks.getParseResultProgress, + markFailed: mocks.markFailed, + markParsing: mocks.markParsing, + mergeParseAssetUrls: mocks.mergeParseAssetUrls, + }, +})) + +vi.mock("@/integrations/knowhere", () => ({ + makeKnowhereClient: mocks.makeKnowhereClient, +})) + +vi.mock("@/lib/logger", () => ({ + logger: { + error: mocks.loggerError, + info: mocks.loggerInfo, + warn: mocks.loggerWarn, + }, +})) + +import { sourceReconcileRouteWorkflow } from "./source-reconcile-route-workflow" + +describe("sourceReconcileRouteWorkflow", () => { + afterEach(() => { + vi.clearAllMocks() + }) + + it("defaults legacy payloads to the poll-and-mirror phase", () => { + expect( + sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + }), + ).toEqual({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + phase: "poll-and-mirror", + segmentIndex: 0, + }) + }) + + it("mirrors the ZIP and triggers asset continuation without running asset batches inline", async () => { + const context = createWorkflowContext() + const continuations: ContinuationTriggerInput[] = [] + const restore = + sourceReconcileRouteWorkflow.setContinuationTriggerForTesting( + async (input) => { + continuations.push(input) + }, + ) + const client = { jobs: {} } + mocks.makeKnowhereClient.mockReturnValue(client) + mocks.pollSourceReconciliation.mockResolvedValue({ + kind: "ready-to-prepare", + jobId: "job_1", + documentId: "doc_1", + }) + mocks.createResultZipMultipartUploadPlan.mockResolvedValue({ + pathname: "result.zip", + key: "blob-key", + uploadId: "upload-1", + sizeBytes: 10, + partSizeBytes: 10, + partCount: 1, + }) + mocks.getResultZipPartRange.mockReturnValue({ + startByte: 0, + endByte: 9, + }) + mocks.uploadResultZipPart.mockResolvedValue({ + etag: "etag-1", + partNumber: 1, + }) + mocks.completeResultZipMultipartUpload.mockResolvedValue({ + url: "https://blob.example/result.zip", + }) + mocks.mergeParseAssetUrls.mockResolvedValue({ id: "parse_result_1" }) + mocks.markParsing.mockResolvedValue(makeSource()) + + try { + await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context, + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + }), + }) + } finally { + restore() + } + + expect(mocks.mergeParseAssetUrls).toHaveBeenCalledWith( + "workspace_1", + "source_1", + { + resultBlobUrl: "https://blob.example/result.zip", + assetUrlsByFilePath: {}, + }, + ) + expect(mocks.markParsing).toHaveBeenCalledWith( + "workspace_1", + "source_1", + "job_1", + "doc_1", + "parsing", + ) + expect(mocks.prepareParsedResultAssetBatch).not.toHaveBeenCalled() + expect(continuations).toEqual([ + { + url: "https://notebook.example/api/sources/reconcile", + payload: { + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + phase: "asset-batches", + segmentIndex: 0, + }, + workflowRunId: sourceReconcileRouteWorkflow.getAssetWorkflowRunId( + "source_1", + 0, + ), + }, + ]) + }) + + it("limits asset continuation work and triggers the next segment when assets remain", async () => { + const context = createWorkflowContext() + const continuations: ContinuationTriggerInput[] = [] + const restore = + sourceReconcileRouteWorkflow.setContinuationTriggerForTesting( + async (input) => { + continuations.push(input) + }, + ) + const client = { jobs: {} } + mocks.makeKnowhereClient.mockReturnValue(client) + mocks.findInWorkspace.mockResolvedValue( + makeSource({ + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + }), + ) + mocks.getParseResultProgress.mockResolvedValue({ + resultBlobUrl: "https://blob.example/result.zip", + assetUrlsByFilePath: {}, + }) + mocks.prepareParsedResultAssetBatch.mockResolvedValue({ + uploadedCount: 10, + remainingCount: 737, + hasMore: true, + }) + + try { + await sourceReconcileRouteWorkflow.runAssetBatchWorkflow({ + context, + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + phase: "asset-batches", + segmentIndex: 2, + }), + }) + } finally { + restore() + } + + expect(mocks.prepareParsedResultAssetBatch).toHaveBeenCalledTimes(20) + expect(mocks.markSourceReadyAfterReconciliation).not.toHaveBeenCalled() + expect(continuations).toEqual([ + { + url: "https://notebook.example/api/sources/reconcile", + payload: { + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + phase: "asset-batches", + segmentIndex: 3, + }, + workflowRunId: sourceReconcileRouteWorkflow.getAssetWorkflowRunId( + "source_1", + 3, + ), + }, + ]) + }) + + it("marks the source ready when the asset continuation has no remaining work", async () => { + const context = createWorkflowContext() + mocks.makeKnowhereClient.mockReturnValue({ jobs: {} }) + mocks.findInWorkspace.mockResolvedValue( + makeSource({ + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + }), + ) + mocks.getParseResultProgress.mockResolvedValue({ + resultBlobUrl: "https://blob.example/result.zip", + assetUrlsByFilePath: {}, + }) + mocks.prepareParsedResultAssetBatch.mockResolvedValue({ + uploadedCount: 0, + remainingCount: 0, + hasMore: false, + }) + mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ + status: "ready", + }) + + await sourceReconcileRouteWorkflow.runAssetBatchWorkflow({ + context, + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + phase: "asset-batches", + }), + }) + + expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", + documentId: "doc_1", + }) + }) + + it("triggers a fresh poll run when Knowhere is still running after the segment budget", async () => { + const context = createWorkflowContext() + const continuations: ContinuationTriggerInput[] = [] + const restore = + sourceReconcileRouteWorkflow.setContinuationTriggerForTesting( + async (input) => { + continuations.push(input) + }, + ) + mocks.makeKnowhereClient.mockReturnValue({ jobs: {} }) + mocks.pollSourceReconciliation.mockResolvedValue({ + kind: "waiting", + jobId: "job_1", + jobStatus: "running", + }) + + try { + await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context, + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + segmentIndex: 4, + }), + }) + } finally { + restore() + } + + expect(mocks.pollSourceReconciliation).toHaveBeenCalledTimes(25) + expect(continuations).toEqual([ + { + url: "https://notebook.example/api/sources/reconcile", + payload: { + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + phase: "poll-and-mirror", + segmentIndex: 5, + }, + workflowRunId: sourceReconcileRouteWorkflow.getPollWorkflowRunId( + "source_1", + 5, + ), + }, + ]) + }) +}) + +type ContinuationTriggerInput = Parameters< + typeof sourceReconcileRouteWorkflow.setContinuationTriggerForTesting +>[0] extends (input: infer Input) => Promise + ? Input + : never + +function createWorkflowContext(): { + readonly run: (stepName: string, step: () => Promise | T) => Promise + readonly sleep: (stepName: string, duration: number | string) => Promise + readonly url: string +} { + return { + run: async (_stepName, step) => step(), + sleep: vi.fn(async () => undefined), + url: "https://notebook.example/api/sources/reconcile", + } +} + +function makeSource(overrides: Partial = {}): Source { + return { + id: "source_1", + workspaceId: "workspace_1", + title: "source.pdf", + mimeType: "application/pdf", + sizeBytes: 1024, + status: "parsing", + failureReason: null, + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-06-30T00:00:00.000Z"), + updatedAt: new Date("2026-06-30T00:00:00.000Z"), + deletedAt: null, + ...overrides, + } +} diff --git a/src/domains/sources/source-reconcile-route-workflow.ts b/src/domains/sources/source-reconcile-route-workflow.ts new file mode 100644 index 0000000..c4763df --- /dev/null +++ b/src/domains/sources/source-reconcile-route-workflow.ts @@ -0,0 +1,467 @@ +import "server-only" + +import { Client, type WorkflowContext } from "@upstash/workflow" + +import { + completeResultZipMultipartUpload, + createResultZipMultipartUploadPlan, + getResultZipPartRange, + prepareParsedResultAssetBatch, + uploadResultZipPart, + type MultipartUploadPart, +} from "@/domains/sources/parsed-result-assets" +import { + markSourceReadyAfterReconciliation, + pollSourceReconciliation, +} from "@/domains/sources/source-reconcile-workflow" +import { makeKnowhereClient } from "@/integrations/knowhere" +import { logger } from "@/lib/logger" +import { sourceWorkflowRuntime } from "./workflow-runtime" + +type ReconcilePayload = { + readonly workspaceId: string + readonly sourceId: string + readonly apiKey: string + readonly phase?: ReconcilePhase + readonly segmentIndex?: number +} + +type ReconcilePhase = "poll-and-mirror" | "asset-batches" + +type NormalizedReconcilePayload = { + readonly workspaceId: string + readonly sourceId: string + readonly apiKey: string + readonly phase: ReconcilePhase + readonly segmentIndex: number +} + +type ReconcileWorkflowContext = Pick< + WorkflowContext, + "run" | "sleep" | "url" +> + +type ContinuationTriggerInput = { + readonly url: string + readonly payload: ReconcilePayload + readonly workflowRunId: string +} + +const maxPollAttempts = 25 +const maxAssetBatchesPerWorkflowRun = 20 +const initialDelaySeconds = 3 +const maxDelaySeconds = 30 + +let triggerContinuation: typeof triggerReconcileContinuation = + triggerReconcileContinuation + +async function runPollAndMirrorWorkflow(input: { + readonly context: ReconcileWorkflowContext + readonly payload: NormalizedReconcilePayload +}): Promise { + const { context, payload } = input + const { workspaceId, sourceId, apiKey } = payload + const client = makeKnowhereClient(apiKey) + let delay = initialDelaySeconds + let completedJob: { + readonly jobId: string + readonly documentId: string + } | null = null + + for (let attempt = 0; attempt < maxPollAttempts; attempt++) { + const poll = await context.run(`poll-${attempt}`, async () => { + return pollSourceReconciliation({ + workspaceId, + sourceId, + client, + }) + }) + + if (poll.kind === "ready-to-prepare") { + completedJob = { + jobId: poll.jobId, + documentId: poll.documentId, + } + logger.info("workflow: source parse completed; preparing artifacts", { + sourceId, + jobId: poll.jobId, + attempts: attempt + 1, + }) + break + } + + if (poll.kind === "resolved") { + logger.info("workflow: source resolved", { + sourceId, + status: poll.status, + attempts: attempt + 1, + }) + return + } + + await context.sleep(`wait-${attempt}`, delay) + delay = Math.min(Math.round(delay * 1.5), maxDelaySeconds) + } + + const jobToPrepare = completedJob + if (!jobToPrepare) { + const nextSegmentIndex = payload.segmentIndex + 1 + await context.run(`trigger-poll-continuation-${nextSegmentIndex}`, async () => + triggerContinuation({ + url: context.url, + payload: { + workspaceId, + sourceId, + apiKey, + phase: "poll-and-mirror", + segmentIndex: nextSegmentIndex, + }, + workflowRunId: getPollWorkflowRunId(sourceId, nextSegmentIndex), + }), + ) + logger.info("workflow: poll continuation triggered", { + sourceId, + maxAttempts: maxPollAttempts, + segmentIndex: nextSegmentIndex, + }) + return + } + + const uploadPlan = await context.run("mirror-zip-start", async () => + createResultZipMultipartUploadPlan({ + workspaceId, + sourceId, + jobId: jobToPrepare.jobId, + client, + }), + ) + logger.info("workflow: result ZIP multipart upload planned", { + sourceId, + jobId: jobToPrepare.jobId, + sizeBytes: uploadPlan.sizeBytes, + partCount: uploadPlan.partCount, + }) + const parts: MultipartUploadPart[] = [] + + for (let partNumber = 1; partNumber <= uploadPlan.partCount; partNumber++) { + const range = getResultZipPartRange( + partNumber, + uploadPlan.partSizeBytes, + uploadPlan.sizeBytes, + ) + const part = await context.run(`mirror-zip-part-${partNumber}`, async () => + uploadResultZipPart({ + pathname: uploadPlan.pathname, + key: uploadPlan.key, + uploadId: uploadPlan.uploadId, + partNumber, + jobId: jobToPrepare.jobId, + client, + ...range, + }), + ) + parts.push(part) + logger.info("workflow: result ZIP part mirrored", { + sourceId, + jobId: jobToPrepare.jobId, + partNumber, + partCount: uploadPlan.partCount, + startByte: range.startByte, + endByte: range.endByte, + }) + } + + const resultBlobUrl = await context.run("mirror-zip-complete", async () => { + const result = await completeResultZipMultipartUpload({ + pathname: uploadPlan.pathname, + key: uploadPlan.key, + uploadId: uploadPlan.uploadId, + parts, + }) + return result.url + }) + logger.info("workflow: result ZIP multipart upload completed", { + sourceId, + jobId: jobToPrepare.jobId, + partCount: parts.length, + }) + + await context.run("parse-result-save-zip", async () => { + const saved = await sourceWorkflowRuntime.mergeParseAssetUrls( + workspaceId, + sourceId, + { + resultBlobUrl, + assetUrlsByFilePath: {}, + }, + ) + if (!saved) throw new Error("Source disappeared before saving result ZIP.") + }) + logger.info("workflow: result ZIP parse-result row saved", { + sourceId, + jobId: jobToPrepare.jobId, + }) + + await context.run("source-save-document-id", async () => { + const saved = await sourceWorkflowRuntime.markParsing( + workspaceId, + sourceId, + jobToPrepare.jobId, + jobToPrepare.documentId, + "parsing", + ) + if (!saved) throw new Error("Source disappeared before saving document id.") + }) + logger.info("workflow: source document id saved for asset continuation", { + sourceId, + jobId: jobToPrepare.jobId, + }) + + await context.run("trigger-asset-continuation-0", async () => + triggerContinuation({ + url: context.url, + payload: { + workspaceId, + sourceId, + apiKey, + phase: "asset-batches", + segmentIndex: 0, + }, + workflowRunId: getAssetWorkflowRunId(sourceId, 0), + }), + ) + logger.info("workflow: parsed asset continuation triggered", { + sourceId, + jobId: jobToPrepare.jobId, + segmentIndex: 0, + }) +} + +async function runAssetBatchWorkflow(input: { + readonly context: ReconcileWorkflowContext + readonly payload: NormalizedReconcilePayload +}): Promise { + const { context, payload } = input + const { workspaceId, sourceId, apiKey, segmentIndex } = payload + const client = makeKnowhereClient(apiKey) + const source = await context.run("load-source", async () => + sourceWorkflowRuntime.findInWorkspace(workspaceId, sourceId), + ) + if (!source) { + logger.info("workflow: asset continuation resolved missing source", { + sourceId, + segmentIndex, + }) + return + } + if (source.status !== "parsing") { + logger.info("workflow: asset continuation resolved non-parsing source", { + sourceId, + segmentIndex, + status: source.status, + }) + return + } + + const jobId = source.knowhereJobId + const documentId = source.knowhereDocumentId + if (!jobId || !documentId) { + await context.run("asset-continuation-failed-missing-job", async () => + sourceWorkflowRuntime.markFailed( + workspaceId, + sourceId, + "Source artifact preparation is missing Knowhere job metadata.", + "parsing", + ), + ) + logger.error("workflow: asset continuation missing Knowhere metadata", { + sourceId, + segmentIndex, + hasJobId: Boolean(jobId), + hasDocumentId: Boolean(documentId), + }) + return + } + + const progress = await context.run("load-parse-progress", async () => + sourceWorkflowRuntime.getParseResultProgress(workspaceId, sourceId), + ) + if (!progress) { + await context.run("asset-continuation-failed-missing-progress", async () => + sourceWorkflowRuntime.markFailed( + workspaceId, + sourceId, + "Source artifact preparation is missing the mirrored parse result ZIP.", + "parsing", + ), + ) + logger.error("workflow: asset continuation missing parse-result progress", { + sourceId, + segmentIndex, + jobId, + }) + return + } + + let batchIndex = 0 + let lastBatch: { + readonly uploadedCount: number + readonly remainingCount: number + readonly hasMore: boolean + } | null = null + while (batchIndex < maxAssetBatchesPerWorkflowRun) { + const batch = await context.run( + `parse-assets-batch-${segmentIndex}-${batchIndex}`, + async () => + prepareParsedResultAssetBatch({ + workspaceId, + sourceId, + jobId, + resultBlobUrl: progress.resultBlobUrl, + client, + repository: sourceWorkflowRuntime, + }), + ) + lastBatch = batch + logger.info("workflow: parsed asset batch prepared", { + sourceId, + jobId, + segmentIndex, + batchIndex, + uploadedCount: batch.uploadedCount, + remainingCount: batch.remainingCount, + hasMore: batch.hasMore, + }) + if (!batch.hasMore) break + batchIndex += 1 + } + + if (lastBatch?.hasMore) { + const nextSegmentIndex = segmentIndex + 1 + await context.run(`trigger-asset-continuation-${nextSegmentIndex}`, async () => + triggerContinuation({ + url: context.url, + payload: { + workspaceId, + sourceId, + apiKey, + phase: "asset-batches", + segmentIndex: nextSegmentIndex, + }, + workflowRunId: getAssetWorkflowRunId(sourceId, nextSegmentIndex), + }), + ) + logger.info("workflow: parsed asset continuation triggered", { + sourceId, + jobId, + segmentIndex: nextSegmentIndex, + remainingCount: lastBatch.remainingCount, + }) + return + } + + const ready = await context.run("source-ready", async () => + markSourceReadyAfterReconciliation({ + workspaceId, + sourceId, + documentId, + }), + ) + logger.info("workflow: source artifact preparation finished", { + sourceId, + jobId, + status: ready.status, + segmentIndex, + assetBatches: batchIndex + 1, + }) +} + +function normalizeReconcilePayload( + payload: ReconcilePayload, +): NormalizedReconcilePayload { + return { + workspaceId: payload.workspaceId, + sourceId: payload.sourceId, + apiKey: payload.apiKey, + phase: payload.phase ?? "poll-and-mirror", + segmentIndex: + typeof payload.segmentIndex === "number" && + Number.isInteger(payload.segmentIndex) && + payload.segmentIndex >= 0 + ? payload.segmentIndex + : 0, + } +} + +async function triggerReconcileContinuation( + input: ContinuationTriggerInput, +): Promise { + const token = process.env.QSTASH_TOKEN + if (!token) { + throw new Error("QSTASH_TOKEN is required to continue source reconciliation.") + } + + await new Client({ token }).trigger({ + url: input.url, + body: input.payload, + workflowRunId: input.workflowRunId, + retries: 3, + }) +} + +function getAssetWorkflowRunId(sourceId: string, segmentIndex: number): string { + return `${sourceId}-assets-${segmentIndex}` +} + +function getPollWorkflowRunId(sourceId: string, segmentIndex: number): string { + return `${sourceId}-poll-${segmentIndex}` +} + +function setContinuationTriggerForTesting( + trigger: typeof triggerReconcileContinuation, +): () => void { + const previous = triggerContinuation + triggerContinuation = trigger + return () => { + triggerContinuation = previous + } +} + +async function markSourceFailedAfterWorkflowFailure( + payload: ReconcilePayload, + failResponse: string, +): Promise { + const normalized = normalizeReconcilePayload(payload) + const reason = `Source reconciliation workflow failed: ${getSafeFailureReason( + failResponse, + )}` + const source = await sourceWorkflowRuntime.markFailed( + normalized.workspaceId, + normalized.sourceId, + reason, + "parsing", + ) + logger.error("workflow: marked source failed after workflow failure", { + sourceId: normalized.sourceId, + workspaceId: normalized.workspaceId, + phase: normalized.phase, + segmentIndex: normalized.segmentIndex, + markedFailed: Boolean(source), + }) +} + +function getSafeFailureReason(value: string): string { + const normalized = value.replace(/\s+/g, " ").trim() + if (normalized.length === 0) return "retry attempts were exhausted." + return normalized.slice(0, 500) +} + +export const sourceReconcileRouteWorkflow = { + getAssetWorkflowRunId, + getPollWorkflowRunId, + markSourceFailedAfterWorkflowFailure, + normalizeReconcilePayload, + runAssetBatchWorkflow, + runPollAndMirrorWorkflow, + setContinuationTriggerForTesting, +} diff --git a/src/domains/sources/source-row-repository.ts b/src/domains/sources/source-row-repository.ts index bfb613f..f501db0 100644 --- a/src/domains/sources/source-row-repository.ts +++ b/src/domains/sources/source-row-repository.ts @@ -64,6 +64,7 @@ type SourceRowRepository = { sourceId: string, jobId: string, documentId?: string, + requiredStatus?: string, ) => Effect.Effect readonly markReadyEffect: ( workspaceId: string, @@ -178,13 +179,14 @@ const markParsingEffect: SourceRowRepository["markParsingEffect"] = ( sourceId: string, jobId: string, documentId?: string, + requiredStatus?: string, ) => updateInWorkspaceEffect(workspaceId, sourceId, { status: "parsing", knowhereJobId: jobId, knowhereDocumentId: documentId, failureReason: null, - }) + }, requiredStatus) const markReadyEffect: SourceRowRepository["markReadyEffect"] = ( workspaceId: string, diff --git a/src/domains/sources/workflow-runtime.ts b/src/domains/sources/workflow-runtime.ts index 534c276..a9bf32b 100644 --- a/src/domains/sources/workflow-runtime.ts +++ b/src/domains/sources/workflow-runtime.ts @@ -31,6 +31,7 @@ type UploadRepositoryRuntime = { sourceId: string, jobId: string, documentId?: string, + requiredStatus?: string, ) => Promise readonly markFailed: ( workspaceId: string, @@ -144,6 +145,7 @@ const markParsing: SourceWorkflowRuntime["markParsing"] = ( sourceId: string, jobId: string, documentId?: string, + requiredStatus?: string, ) => databaseRuntime.runPromise( sourceRepository.markParsingEffect( @@ -151,6 +153,7 @@ const markParsing: SourceWorkflowRuntime["markParsing"] = ( sourceId, jobId, documentId, + requiredStatus, ), ) From f2c8a50c6dd97eaf51160e48c5ecfb0973582d97 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 30 Jun 2026 11:45:00 +0800 Subject: [PATCH 3/6] Fix reconciliation resume trigger behavior --- .../sources/background-reconcile.test.ts | 91 ++++++++ src/domains/sources/background-reconcile.ts | 30 +-- .../sources/parsed-result-assets.test.ts | 196 ++++++++++++++++++ src/domains/sources/parsed-result-assets.ts | 159 ++++++++++++++ .../source-reconcile-route-workflow.test.ts | 143 ++++++++++++- .../source-reconcile-route-workflow.ts | 58 ++++++ 6 files changed, 663 insertions(+), 14 deletions(-) create mode 100644 src/domains/sources/background-reconcile.test.ts diff --git a/src/domains/sources/background-reconcile.test.ts b/src/domains/sources/background-reconcile.test.ts new file mode 100644 index 0000000..f6ad6f3 --- /dev/null +++ b/src/domains/sources/background-reconcile.test.ts @@ -0,0 +1,91 @@ +import { afterEach, describe, expect, it, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + loggerError: vi.fn(), + loggerInfo: vi.fn(), + loggerWarn: vi.fn(), + trigger: vi.fn(), +})) + +vi.mock("@upstash/workflow", () => ({ + Client: class { + trigger = mocks.trigger + }, +})) + +vi.mock("@/lib/logger", () => ({ + logger: { + error: mocks.loggerError, + info: mocks.loggerInfo, + warn: mocks.loggerWarn, + }, +})) + +describe("startBackgroundReconciliation", () => { + afterEach(() => { + vi.clearAllMocks() + vi.useRealTimers() + delete process.env.QSTASH_TOKEN + delete process.env.NOTEBOOK_PUBLIC_URL + vi.resetModules() + }) + + it("deduplicates workflow triggers only within a bounded cooldown", async () => { + const triggerCooldownMs: number = 5 * 60_000 + vi.useFakeTimers() + vi.setSystemTime(new Date("2026-06-30T00:00:00.000Z")) + process.env.QSTASH_TOKEN = "qstash_token" + process.env.NOTEBOOK_PUBLIC_URL = "https://notebook.example" + mocks.trigger.mockResolvedValue({}) + + const { startBackgroundReconciliation } = await import( + "./background-reconcile" + ) + + await startBackgroundReconciliation( + "workspace_1", + "source_1", + "knowhere_key", + ) + await startBackgroundReconciliation( + "workspace_1", + "source_1", + "knowhere_key", + ) + + expect(mocks.trigger).toHaveBeenCalledTimes(1) + expect(mocks.trigger).toHaveBeenLastCalledWith({ + url: "https://notebook.example/api/sources/reconcile", + body: { + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "knowhere_key", + }, + workflowRunId: `source_1-${Math.floor( + new Date("2026-06-30T00:00:00.000Z").getTime() / triggerCooldownMs, + )}`, + retries: 3, + }) + + vi.setSystemTime(new Date("2026-06-30T00:05:01.000Z")) + await startBackgroundReconciliation( + "workspace_1", + "source_1", + "knowhere_key", + ) + + expect(mocks.trigger).toHaveBeenCalledTimes(2) + expect(mocks.trigger).toHaveBeenLastCalledWith({ + url: "https://notebook.example/api/sources/reconcile", + body: { + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "knowhere_key", + }, + workflowRunId: `source_1-${Math.floor( + new Date("2026-06-30T00:05:01.000Z").getTime() / triggerCooldownMs, + )}`, + retries: 3, + }) + }) +}) diff --git a/src/domains/sources/background-reconcile.ts b/src/domains/sources/background-reconcile.ts index e5c509b..604ad64 100644 --- a/src/domains/sources/background-reconcile.ts +++ b/src/domains/sources/background-reconcile.ts @@ -5,15 +5,14 @@ import { Client } from "@upstash/workflow" import { logger } from "@/lib/logger" -// Re-trigger protection: Layers 1 & 2. +// Re-trigger protection: bounded process guard plus bucketed workflow IDs. // -// Layer 1 — Upstash idempotency via workflowRunId=sourceId ensures at most one -// running workflow per source, even across process restarts or multiple instances. -// -// Layer 2 — In-memory Set avoids the network call entirely when the same process -// already triggered a workflow for this source. +// Continuation workflows use deterministic run IDs, but the initial trigger +// must stay retryable after an old completed workflow run. The cooldown avoids +// duplicate same-process trigger calls without permanently blocking a source. -const triggeredSourceIds = new Set() +const triggerCooldownMs: number = 5 * 60_000 +const lastTriggeredAtBySourceId: Map = new Map() function resolveBaseURL(): string { return process.env.NOTEBOOK_PUBLIC_URL ?? "http://localhost:3000" @@ -29,8 +28,15 @@ const startBackgroundReconciliationEffect = ( apiKey: string, ): Effect.Effect => Effect.gen(function* () { - if (triggeredSourceIds.has(sourceId)) return - triggeredSourceIds.add(sourceId) + const now = Date.now() + const lastTriggeredAt = lastTriggeredAtBySourceId.get(sourceId) + if ( + lastTriggeredAt !== undefined && + now - lastTriggeredAt < triggerCooldownMs + ) { + return + } + lastTriggeredAtBySourceId.set(sourceId, now) const token = process.env.QSTASH_TOKEN if (!token) { @@ -38,7 +44,7 @@ const startBackgroundReconciliationEffect = ( sourceId, workspaceId, }) - triggeredSourceIds.delete(sourceId) + lastTriggeredAtBySourceId.delete(sourceId) return } @@ -53,7 +59,7 @@ const startBackgroundReconciliationEffect = ( return await new Client({ token }).trigger({ url, body: { workspaceId, sourceId, apiKey }, - workflowRunId: sourceId, + workflowRunId: `${sourceId}-${Math.floor(now / triggerCooldownMs)}`, retries: 3, }) } catch (err) { @@ -72,7 +78,7 @@ const startBackgroundReconciliationEffect = ( }).pipe( Effect.catchAllCause((cause) => Effect.sync(() => { - triggeredSourceIds.delete(sourceId) + lastTriggeredAtBySourceId.delete(sourceId) logger.error("background-reconcile: failed to trigger workflow", { sourceId, workspaceId, diff --git a/src/domains/sources/parsed-result-assets.test.ts b/src/domains/sources/parsed-result-assets.test.ts index db7795e..bbada92 100644 --- a/src/domains/sources/parsed-result-assets.test.ts +++ b/src/domains/sources/parsed-result-assets.test.ts @@ -415,6 +415,202 @@ describe("prepareParsedResultAssetBatch", () => { hasMore: false, }) }) + + it("uses indexed document assets to complete without loading the result ZIP", async () => { + const client = { + jobs: { + load: vi.fn(), + }, + documents: { + listChunks: vi.fn(async (_documentId: string, params: { + readonly chunkType: "image" | "table" + }) => ({ + chunks: + params.chunkType === "image" + ? [ + { + filePath: "images/image-1.jpg", + sourceChunkPath: null, + metadata: {}, + }, + ] + : [ + { + filePath: "tables/table-1.html", + sourceChunkPath: null, + metadata: {}, + }, + ], + pagination: { + totalPages: 1, + }, + })), + }, + } + const repository = { + getParseResultProgress: vi.fn(async () => ({ + resultBlobUrl: "https://blob.example/result.zip", + assetUrlsByFilePath: { + "images/image-1.jpg": "https://blob.example/images/image-1.jpg", + "tables/table-1.html": "https://blob.example/tables/table-1.html", + }, + })), + mergeParseAssetUrls: vi.fn(async () => undefined), + } + + const result = await prepareParsedResultAssetBatch({ + workspaceId: "workspace_1", + sourceId: "source_1", + jobId: "job_1", + documentId: "doc_1", + resultBlobUrl: "https://blob.example/result.zip", + client, + repository, + }) + + expect(client.documents.listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 200, + chunkType: "image", + includeAssetUrls: false, + }) + expect(client.documents.listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 200, + chunkType: "table", + includeAssetUrls: false, + }) + expect(client.jobs.load).not.toHaveBeenCalled() + expect(repository.mergeParseAssetUrls).not.toHaveBeenCalled() + expect(result).toEqual({ + uploadedCount: 0, + remainingCount: 0, + hasMore: false, + }) + }) + + it("falls back to loading the result ZIP when indexed assets are missing", async () => { + const client = { + jobs: { + load: vi.fn(async () => ({ + imageChunks: [ + { + filePath: "images/image-2.png", + data: Buffer.from("new image"), + }, + ], + tableChunks: [], + })), + }, + documents: { + listChunks: vi.fn(async () => ({ + chunks: [ + { + filePath: "images/image-2.png", + sourceChunkPath: null, + metadata: {}, + }, + ], + pagination: { + totalPages: 1, + }, + })), + }, + } + const repository = { + getParseResultProgress: vi.fn(async () => ({ + resultBlobUrl: "https://blob.example/result.zip", + assetUrlsByFilePath: { + "images/image-1.jpg": "https://blob.example/images/image-1.jpg", + }, + })), + mergeParseAssetUrls: vi.fn(async () => undefined), + } + const blobStore = { + put: vi.fn(async (pathname: string) => ({ + url: `https://blob.example/${pathname}`, + })), + } + + const result = await prepareParsedResultAssetBatch({ + workspaceId: "workspace_1", + sourceId: "source_1", + jobId: "job_1", + documentId: "doc_1", + resultBlobUrl: "https://blob.example/result.zip", + client, + repository, + blobStore, + }) + + expect(client.jobs.load).toHaveBeenCalledWith("job_1") + expect(blobStore.put).toHaveBeenCalledWith( + "workspaces/workspace_1/sources/source_1/parsed-result/images/image-2.png", + Buffer.from("new image"), + expect.objectContaining({ contentType: "image/png" }), + ) + expect(result).toEqual({ + uploadedCount: 1, + remainingCount: 0, + hasMore: false, + }) + }) + + it("falls back to loading the result ZIP when the asset index check fails", async () => { + const client = { + jobs: { + load: vi.fn(async () => ({ + imageChunks: [ + { + filePath: "images/image-3.png", + data: Buffer.from("new image"), + }, + ], + tableChunks: [], + })), + }, + documents: { + listChunks: vi.fn(async () => { + throw new Error("index unavailable") + }), + }, + } + const repository = { + getParseResultProgress: vi.fn(async () => ({ + resultBlobUrl: "https://blob.example/result.zip", + assetUrlsByFilePath: {}, + })), + mergeParseAssetUrls: vi.fn(async () => undefined), + } + const blobStore = { + put: vi.fn(async (pathname: string) => ({ + url: `https://blob.example/${pathname}`, + })), + } + + const result = await prepareParsedResultAssetBatch({ + workspaceId: "workspace_1", + sourceId: "source_1", + jobId: "job_1", + documentId: "doc_1", + resultBlobUrl: "https://blob.example/result.zip", + client, + repository, + blobStore, + }) + + expect(client.jobs.load).toHaveBeenCalledWith("job_1") + expect(blobStore.put).toHaveBeenCalledWith( + "workspaces/workspace_1/sources/source_1/parsed-result/images/image-3.png", + Buffer.from("new image"), + expect.objectContaining({ contentType: "image/png" }), + ) + expect(result).toEqual({ + uploadedCount: 1, + remainingCount: 0, + hasMore: false, + }) + }) }) function makeJobResult(): JobResult { diff --git a/src/domains/sources/parsed-result-assets.ts b/src/domains/sources/parsed-result-assets.ts index 6a6a727..09a4179 100644 --- a/src/domains/sources/parsed-result-assets.ts +++ b/src/domains/sources/parsed-result-assets.ts @@ -114,11 +114,18 @@ export type PrepareParsedResultAssetBatchInput = { readonly workspaceId: string readonly sourceId: string readonly jobId: string + readonly documentId?: string readonly resultBlobUrl: string readonly client: { readonly jobs: { load(jobId: string): Promise } + readonly documents?: { + listChunks( + documentId: string, + params: ParsedAssetChunkListParams, + ): Promise + } } readonly repository: ParsedResultAssetProgressRepository readonly blobStore?: ParsedResultBlobStore @@ -183,6 +190,38 @@ type ParsedAssetUpload = { readonly contentType: string } +type ParsedAssetChunkType = "image" | "table" + +type ParsedAssetChunkListParams = { + readonly page: number + readonly pageSize: number + readonly chunkType: ParsedAssetChunkType + readonly includeAssetUrls?: boolean +} + +type ParsedAssetChunkListResponse = { + readonly chunks: readonly ParsedAssetChunkIndexItem[] + readonly pagination?: { + readonly totalPages?: number + } +} + +type ParsedAssetChunkIndexItem = { + readonly filePath?: string | null + readonly sourceChunkPath?: string | null + readonly metadata: Readonly> +} + +type ParsedAssetIndexCompletion = + | { + readonly kind: "complete" + readonly assetCount: number + } + | { + readonly kind: "missing" + readonly missingCount: number + } + export type StoreParsedResultAssetsInput = { workspaceId: string sourceId: string @@ -199,6 +238,7 @@ const parsedResultDirectoryName = "parsed-result" const resultZipPartSizeBytes = 8 * 1024 * 1024 const parsedAssetBatchLimit = 10 const parsedAssetBatchMaxDurationMs = 45_000 +const parsedAssetChunkPageSize = 200 const resultZipMetadataTimeoutMs = 15_000 const resultZipRangeFetchTimeoutMs = 30_000 const resultZipBlobOperationTimeoutMs = 30_000 @@ -274,6 +314,7 @@ export const prepareParsedResultAssetBatchEffect = Effect.fn( workspaceId, sourceId, jobId, + documentId, resultBlobUrl, client, repository, @@ -286,6 +327,44 @@ export const prepareParsedResultAssetBatchEffect = Effect.fn( repository.getParseResultProgress(workspaceId, sourceId), ) const existingAssetUrls = progress?.assetUrlsByFilePath ?? {} + const indexedCompletion = yield* Effect.tryPromise(() => + getIndexedAssetCompletion({ + documentId, + existingAssetUrls, + client, + }), + ).pipe( + Effect.catchAll((error) => + Effect.sync(() => { + logger.warn( + "parsed-result-assets: asset index check failed; falling back to result ZIP load", + { + workspaceId, + sourceId, + jobId, + documentId, + error: error instanceof Error ? error.message : String(error), + }, + ) + return null + }), + ), + ) + if (indexedCompletion?.kind === "complete") { + logger.info("parsed-result-assets: asset index already complete", { + workspaceId, + sourceId, + jobId, + documentId, + assetCount: indexedCompletion.assetCount, + }) + return { + uploadedCount: 0, + remainingCount: 0, + hasMore: false, + } + } + const parseResult = (yield* Effect.tryPromise(() => client.jobs.load(jobId), )) as ParsedResultWithAssets @@ -565,6 +644,76 @@ function getUploadableParsedAssets( return assets } +async function getIndexedAssetCompletion(input: { + readonly documentId?: string + readonly existingAssetUrls: Readonly> + readonly client: PrepareParsedResultAssetBatchInput["client"] +}): Promise { + if (!input.documentId || !input.client.documents) return null + + const indexedAssetPaths = await listIndexedAssetPaths( + input.client.documents, + input.documentId, + ) + const missingCount = indexedAssetPaths.filter( + (filePath) => input.existingAssetUrls[filePath] === undefined, + ).length + if (missingCount > 0) return { kind: "missing", missingCount } + + return { + kind: "complete", + assetCount: indexedAssetPaths.length, + } +} + +async function listIndexedAssetPaths( + documents: NonNullable, + documentId: string, +): Promise { + const [imagePaths, tablePaths] = await Promise.all([ + listIndexedAssetPathsByType(documents, documentId, "image"), + listIndexedAssetPathsByType(documents, documentId, "table"), + ]) + + return [...new Set([...imagePaths, ...tablePaths])] +} + +async function listIndexedAssetPathsByType( + documents: NonNullable, + documentId: string, + chunkType: ParsedAssetChunkType, +): Promise { + const paths: string[] = [] + let page = 1 + let totalPages = 1 + + do { + const response = await documents.listChunks(documentId, { + page, + pageSize: parsedAssetChunkPageSize, + chunkType, + includeAssetUrls: false, + }) + + for (const chunk of response.chunks) { + const filePath = normalizeParsedAssetPath( + getFirstString([ + chunk.filePath, + chunk.metadata["filePath"], + chunk.metadata["file_path"], + chunk.sourceChunkPath, + ]), + ) + if (filePath) paths.push(filePath) + } + + totalPages = getSafeTotalPages(response.pagination?.totalPages) + page += 1 + } while (page <= totalPages) + + return paths +} + function normalizeParsedAssetPath(value: string | undefined): string | null { if (!value) return null const normalized = value.replaceAll("\\", "/").replace(/^\.\/+/, "") @@ -586,6 +735,16 @@ function normalizeParsedAssetPath(value: string | undefined): string | null { return parts.join("/") } +function getFirstString(values: readonly unknown[]): string | undefined { + return values.find((value): value is string => typeof value === "string") +} + +function getSafeTotalPages(value: number | undefined): number { + return typeof value === "number" && Number.isFinite(value) && value > 0 + ? Math.floor(value) + : 1 +} + function getContentTypeForPath(filePath: string): string { const extension = path.extname(filePath).toLowerCase() if (extension === ".jpg" || extension === ".jpeg") return "image/jpeg" diff --git a/src/domains/sources/source-reconcile-route-workflow.test.ts b/src/domains/sources/source-reconcile-route-workflow.test.ts index 8cb69f7..2307260 100644 --- a/src/domains/sources/source-reconcile-route-workflow.test.ts +++ b/src/domains/sources/source-reconcile-route-workflow.test.ts @@ -164,6 +164,140 @@ describe("sourceReconcileRouteWorkflow", () => { ]) }) + it("resumes asset continuation without remirroring an already saved ZIP", async () => { + vi.setSystemTime(new Date("2026-06-30T03:20:00.000Z")) + const context = createWorkflowContext() + const continuations: ContinuationTriggerInput[] = [] + const restore = + sourceReconcileRouteWorkflow.setContinuationTriggerForTesting( + async (input) => { + continuations.push(input) + }, + ) + mocks.makeKnowhereClient.mockReturnValue({ jobs: {} }) + mocks.pollSourceReconciliation.mockResolvedValue({ + kind: "ready-to-prepare", + jobId: "job_1", + documentId: "doc_1", + }) + mocks.getParseResultProgress.mockResolvedValue({ + resultBlobUrl: "https://blob.example/result.zip", + assetUrlsByFilePath: { + "images/image-1.png": "https://blob.example/images/image-1.png", + }, + }) + mocks.markParsing.mockResolvedValue(makeSource()) + + try { + await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context, + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + }), + }) + } finally { + restore() + vi.useRealTimers() + } + + const resumeSegmentIndex = new Date("2026-06-30T03:20:00.000Z").getTime() + expect(mocks.createResultZipMultipartUploadPlan).not.toHaveBeenCalled() + expect(mocks.uploadResultZipPart).not.toHaveBeenCalled() + expect(mocks.completeResultZipMultipartUpload).not.toHaveBeenCalled() + expect(mocks.mergeParseAssetUrls).not.toHaveBeenCalled() + expect(mocks.markParsing).toHaveBeenCalledWith( + "workspace_1", + "source_1", + "job_1", + "doc_1", + "parsing", + ) + expect(continuations).toEqual([ + { + url: "https://notebook.example/api/sources/reconcile", + payload: { + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + phase: "asset-batches", + segmentIndex: resumeSegmentIndex, + }, + workflowRunId: sourceReconcileRouteWorkflow.getAssetWorkflowRunId( + "source_1", + resumeSegmentIndex, + ), + }, + ]) + }) + + it("reuses the asset resume segment during workflow replay", async () => { + vi.setSystemTime(new Date("2026-06-30T03:20:00.000Z")) + const stepResults = new Map() + const context = createWorkflowContext(stepResults) + const continuations: ContinuationTriggerInput[] = [] + const restore = + sourceReconcileRouteWorkflow.setContinuationTriggerForTesting( + async (input) => { + continuations.push(input) + }, + ) + mocks.makeKnowhereClient.mockReturnValue({ jobs: {} }) + mocks.pollSourceReconciliation.mockResolvedValue({ + kind: "ready-to-prepare", + jobId: "job_1", + documentId: "doc_1", + }) + mocks.getParseResultProgress.mockResolvedValue({ + resultBlobUrl: "https://blob.example/result.zip", + assetUrlsByFilePath: { + "images/image-1.png": "https://blob.example/images/image-1.png", + }, + }) + mocks.markParsing.mockResolvedValue(makeSource()) + + try { + await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context, + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + }), + }) + vi.setSystemTime(new Date("2026-06-30T03:20:14.000Z")) + await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context: createWorkflowContext(stepResults), + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + }), + }) + } finally { + restore() + vi.useRealTimers() + } + + const resumeSegmentIndex = new Date("2026-06-30T03:20:00.000Z").getTime() + expect(continuations).toHaveLength(1) + expect(continuations[0]?.workflowRunId).toBe( + sourceReconcileRouteWorkflow.getAssetWorkflowRunId( + "source_1", + resumeSegmentIndex, + ), + ) + expect([...stepResults.keys()]).toContain( + `trigger-asset-continuation-${resumeSegmentIndex}`, + ) + expect([...stepResults.keys()]).not.toContain( + `trigger-asset-continuation-${new Date( + "2026-06-30T03:20:14.000Z", + ).getTime()}`, + ) + }) + it("limits asset continuation work and triggers the next segment when assets remain", async () => { const context = createWorkflowContext() const continuations: ContinuationTriggerInput[] = [] @@ -321,13 +455,18 @@ type ContinuationTriggerInput = Parameters< ? Input : never -function createWorkflowContext(): { +function createWorkflowContext(stepResults = new Map()): { readonly run: (stepName: string, step: () => Promise | T) => Promise readonly sleep: (stepName: string, duration: number | string) => Promise readonly url: string } { return { - run: async (_stepName, step) => step(), + run: async (stepName: string, step: () => Promise | T) => { + if (stepResults.has(stepName)) return stepResults.get(stepName) as T + const result = await step() + stepResults.set(stepName, result) + return result + }, sleep: vi.fn(async () => undefined), url: "https://notebook.example/api/sources/reconcile", } diff --git a/src/domains/sources/source-reconcile-route-workflow.ts b/src/domains/sources/source-reconcile-route-workflow.ts index c4763df..2d94375 100644 --- a/src/domains/sources/source-reconcile-route-workflow.ts +++ b/src/domains/sources/source-reconcile-route-workflow.ts @@ -127,6 +127,63 @@ async function runPollAndMirrorWorkflow(input: { return } + const existingProgress = await context.run( + "parse-result-check-existing", + async () => + sourceWorkflowRuntime.getParseResultProgress(workspaceId, sourceId), + ) + if (existingProgress) { + const resumeAssetSegmentIndex = await context.run( + "create-asset-resume-segment", + async () => Date.now(), + ) + const savedAssetCount: number = Object.keys( + existingProgress.assetUrlsByFilePath, + ).length + + await context.run("source-save-document-id-for-existing-zip", async () => { + const saved = await sourceWorkflowRuntime.markParsing( + workspaceId, + sourceId, + jobToPrepare.jobId, + jobToPrepare.documentId, + "parsing", + ) + if (!saved) throw new Error("Source disappeared before saving document id.") + }) + logger.info("workflow: result ZIP already mirrored; resuming asset batches", { + sourceId, + jobId: jobToPrepare.jobId, + savedAssetCount, + segmentIndex: resumeAssetSegmentIndex, + }) + + await context.run( + `trigger-asset-continuation-${resumeAssetSegmentIndex}`, + async () => + triggerContinuation({ + url: context.url, + payload: { + workspaceId, + sourceId, + apiKey, + phase: "asset-batches", + segmentIndex: resumeAssetSegmentIndex, + }, + workflowRunId: getAssetWorkflowRunId( + sourceId, + resumeAssetSegmentIndex, + ), + }), + ) + logger.info("workflow: parsed asset continuation triggered", { + sourceId, + jobId: jobToPrepare.jobId, + segmentIndex: resumeAssetSegmentIndex, + }) + return + } + const uploadPlan = await context.run("mirror-zip-start", async () => createResultZipMultipartUploadPlan({ workspaceId, @@ -317,6 +374,7 @@ async function runAssetBatchWorkflow(input: { workspaceId, sourceId, jobId, + documentId, resultBlobUrl: progress.resultBlobUrl, client, repository: sourceWorkflowRuntime, From 3ad074ba6b5174b1a6d7bb23d1943ff104efe8a0 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 30 Jun 2026 15:08:29 +0800 Subject: [PATCH 4/6] Implement lazy Knowhere chunk cache --- next.config.ts | 3 + package.json | 2 +- pnpm-lock.yaml | 12 +- .../sources/[sourceId]/chunks/route.test.ts | 126 ++- src/app/api/sources/reconcile/route.ts | 8 - src/components/source-row.tsx | 13 +- src/components/workspace-chat-workflow.ts | 15 +- src/components/workspace-source-workflow.ts | 6 +- src/domains/chat/route-answer.ts | 12 +- src/domains/chunks/index.ts | 4 + src/domains/chunks/server.test.ts | 440 +++++++++++ src/domains/chunks/server.ts | 739 ++++++++++++++++++ src/domains/sources/lifecycle.ts | 38 +- src/domains/sources/reconcile.test.ts | 87 +-- src/domains/sources/reconcile.ts | 20 - src/domains/sources/remote-library.ts | 199 ++++- src/domains/sources/repository.ts | 2 + src/domains/sources/route-chunks.ts | 97 ++- src/domains/sources/route-dependencies.ts | 3 +- src/domains/sources/route-listing.ts | 8 +- src/domains/sources/route-service.test.ts | 127 ++- src/domains/sources/route-types.ts | 29 +- src/domains/sources/service.ts | 6 + .../source-reconcile-route-workflow.test.ts | 364 +-------- .../source-reconcile-route-workflow.ts | 324 +------- .../sources/source-reconcile-workflow.test.ts | 17 +- .../sources/source-reconcile-workflow.ts | 11 - src/domains/sources/source-row-repository.ts | 19 +- src/domains/sources/types.ts | 2 +- src/domains/sources/workflow-runtime.ts | 19 + src/domains/workspace/initial-state.test.ts | 14 - src/domains/workspace/initial-state.ts | 17 +- 32 files changed, 1797 insertions(+), 986 deletions(-) create mode 100644 src/domains/chunks/server.test.ts create mode 100644 src/domains/chunks/server.ts diff --git a/next.config.ts b/next.config.ts index 490c3cc..e294c70 100644 --- a/next.config.ts +++ b/next.config.ts @@ -14,6 +14,9 @@ const nextConfig: NextConfig = { "notebook.127.0.0.1.nip.io", "dashboard.127.0.0.1.nip.io", ], + turbopack: { + root: process.cwd(), + }, }; export default nextConfig; diff --git a/package.json b/package.json index 8449b3b..ee4221c 100644 --- a/package.json +++ b/package.json @@ -24,7 +24,7 @@ "@antv/chart-visualization-skills": "0.1.3", "@effect/platform": "^0.96.1", "@neondatabase/serverless": "^1.1.0", - "@ontos-ai/knowhere-sdk": "^0.6.0", + "@ontos-ai/knowhere-sdk": "^0.10.0", "@radix-ui/react-alert-dialog": "^1.1.15", "@radix-ui/react-checkbox": "^1.3.3", "@radix-ui/react-dialog": "^1.1.15", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 484f2fd..66cd259 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -21,8 +21,8 @@ importers: specifier: ^1.1.0 version: 1.1.0 '@ontos-ai/knowhere-sdk': - specifier: ^0.6.0 - version: 0.6.0 + specifier: ^0.10.0 + version: 0.10.0 '@radix-ui/react-alert-dialog': specifier: ^1.1.15 version: 1.1.15(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4) @@ -1409,9 +1409,9 @@ packages: resolution: {integrity: sha512-nn5ozdjYQpUCZlWGuxcJY/KpxkWQs4DcbMCmKojjyrYDEAGy4Ce19NN4v5MduafTwJlbKc99UA8YhSVqq9yPZA==} engines: {node: '>=12.4.0'} - '@ontos-ai/knowhere-sdk@0.6.0': - resolution: {integrity: sha512-YqDgg+BKwGp7NepD50aquhEOmUgdDAQxpoyttqh0d+4tWjOgkSo74vEGiJVxt8v+LAcA/5CmxPeghX1Hx0v3wQ==} - engines: {node: '>=20.19.0', npm: '>=10.0.0'} + '@ontos-ai/knowhere-sdk@0.10.0': + resolution: {integrity: sha512-wb5wNnnp4YvdlZ6w2UOQ2gw/9anGsmn2hw2JD+uCya4VUyXK78HKMi1LCI6zNnLeZHyGAm0q6+d2wlqDXZ6haQ==} + engines: {node: '>=20.19.0', npm: '>=10.0.0', pnpm: '>=9.0.0'} '@open-draft/deferred-promise@2.2.0': resolution: {integrity: sha512-CecwLWx3rhxVQF6V4bAgPS5t+So2sTbPgAzafKkVizyi7tlwpcFpdFqq+wqF2OwNBmqFuu6tOyouTuxgpMfzmA==} @@ -6397,7 +6397,7 @@ snapshots: '@nolyfill/is-core-module@1.0.39': {} - '@ontos-ai/knowhere-sdk@0.6.0': + '@ontos-ai/knowhere-sdk@0.10.0': dependencies: axios: 1.16.0 jszip: 3.10.1 diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index 6074abf..524a657 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -2,14 +2,19 @@ import { NextRequest } from "next/server" import { beforeEach, describe, expect, it, vi } from "vitest" const mocks = vi.hoisted(() => ({ + blobGet: vi.fn(), + blobPut: vi.fn(), + deleteBlob: vi.fn(), ensureApiKeyForWorkspace: vi.fn(), ensureWorkspace: vi.fn(), fetchDemoChunkPage: vi.fn(), findSourceInWorkspace: vi.fn(), getCurrentUser: vi.fn(), getSourceParseAssetUrls: vi.fn(), + localizeRemoteDocument: vi.fn(), makeKnowhereClient: vi.fn(), requireUser: vi.fn(), + updateSourceRevisionKey: vi.fn(), })) vi.mock("next/headers", () => ({ @@ -36,10 +41,18 @@ vi.mock("@/integrations/knowhere", () => ({ makeKnowhereClient: mocks.makeKnowhereClient, })) +vi.mock("@vercel/blob", () => ({ + del: mocks.deleteBlob, + get: mocks.blobGet, + put: mocks.blobPut, +})) + vi.mock("@/domains/sources/service", () => ({ sourceService: { findInWorkspace: mocks.findSourceInWorkspace, getParseAssetUrls: mocks.getSourceParseAssetUrls, + localizeRemoteDocument: mocks.localizeRemoteDocument, + updateSourceRevisionKey: mocks.updateSourceRevisionKey, }, })) @@ -54,6 +67,11 @@ import { GET } from "./route" describe("GET /api/sources/[sourceId]/chunks", () => { beforeEach(() => { vi.clearAllMocks() + mocks.blobGet.mockResolvedValue(null) + mocks.blobPut.mockImplementation(async (pathname: string) => ({ + url: `https://blob.example/${pathname}`, + })) + mocks.updateSourceRevisionKey.mockResolvedValue(null) }) it("serves API-owned demo chunks for anonymous canonical demo sources", async () => { @@ -502,7 +520,48 @@ describe("GET /api/sources/[sourceId]/chunks", () => { }) }) - it("does not load chunks from unlocalized remote source ids", async () => { + it("materializes a remote source id on open before loading chunks", async () => { + const knowhereClient = { + documents: { + list: vi.fn(async () => ({ + documents: [ + { + documentId: "doc_remote", + namespace: "default", + status: "active", + currentJobResultId: "job_result_1", + sourceFileName: "remote.pdf", + documentMetadata: { + mimeType: "application/pdf", + }, + }, + ], + })), + listChunks: vi.fn(async () => ({ + documentId: "doc_remote", + jobResultId: "job_result_1", + chunks: [ + { + id: "dchk_remote", + chunkId: "parser_remote", + chunkType: "text", + content: "Remote chunk", + sectionPath: "Summary", + sourceChunkPath: "Default_Root/remote.pdf/Summary", + filePath: null, + metadata: {}, + sortOrder: 0, + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + })), + }, + } mocks.getCurrentUser.mockResolvedValue({ id: "user_1", email: null, @@ -515,6 +574,27 @@ describe("GET /api/sources/[sourceId]/chunks", () => { createdAt: new Date("2026-05-10T00:00:00.000Z"), }) mocks.fetchDemoChunkPage.mockRejectedValue(new Error("not a demo")) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.makeKnowhereClient.mockReturnValue(knowhereClient) + mocks.localizeRemoteDocument.mockResolvedValue({ + id: "00000000-0000-0000-0000-000000000009", + workspaceId: "workspace_1", + title: "remote.pdf", + mimeType: "application/pdf", + sizeBytes: 0, + status: "ready", + failureReason: null, + knowhereJobId: "job_result_1", + knowhereDocumentId: "doc_remote", + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-05-10T00:00:00.000Z"), + updatedAt: new Date("2026-05-10T00:00:00.000Z"), + deletedAt: null, + }) const response = await GET( new NextRequest( @@ -527,12 +607,46 @@ describe("GET /api/sources/[sourceId]/chunks", () => { }, ) - await expect(response.json()).resolves.toEqual({ - message: "Source not found.", + await expect(response.json()).resolves.toMatchObject({ + chunks: [ + { + chunkId: "dchk_remote", + parserChunkId: "parser_remote", + documentId: "doc_remote", + sourceTitle: "remote.pdf", + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + }, }) - expect(response.status).toBe(404) + expect(response.status).toBe(200) expect(mocks.findSourceInWorkspace).not.toHaveBeenCalled() - expect(mocks.ensureApiKeyForWorkspace).not.toHaveBeenCalled() - expect(mocks.makeKnowhereClient).not.toHaveBeenCalled() + expect(mocks.ensureApiKeyForWorkspace).toHaveBeenCalledWith( + "workspace_1", + "session=abc", + ) + expect(mocks.localizeRemoteDocument).toHaveBeenCalledWith( + "workspace_1", + { + documentId: "doc_remote", + namespace: "default", + status: "ready", + title: "remote.pdf", + mimeType: "application/pdf", + sizeBytes: undefined, + revisionKey: "job_result_1", + }, + ) + expect(knowhereClient.documents.listChunks).toHaveBeenCalledWith( + "doc_remote", + { + page: 1, + pageSize: 1, + includeAssetUrls: true, + }, + ) }) }) diff --git a/src/app/api/sources/reconcile/route.ts b/src/app/api/sources/reconcile/route.ts index 223d380..1f51bac 100644 --- a/src/app/api/sources/reconcile/route.ts +++ b/src/app/api/sources/reconcile/route.ts @@ -11,14 +11,6 @@ export const { POST } = serve( const payload = sourceReconcileRouteWorkflow.normalizeReconcilePayload( context.requestPayload, ) - if (payload.phase === "asset-batches") { - await sourceReconcileRouteWorkflow.runAssetBatchWorkflow({ - context, - payload, - }) - return - } - await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ context, payload, diff --git a/src/components/source-row.tsx b/src/components/source-row.tsx index 48be677..77cf9d7 100644 --- a/src/components/source-row.tsx +++ b/src/components/source-row.tsx @@ -37,6 +37,7 @@ export function SourceRow({ const isBusy = source.status === "uploading" || source.status === "parsing"; const isFailed = source.status === "failed"; const isLibrarySource = source.officialLibrary !== undefined; + const isRemoteSource = source.kind === "remote"; const iconBg = fileIconTint(source.title); @@ -57,7 +58,7 @@ export function SourceRow({ > onToggleIncluded?.(source.id, checked === true) } @@ -96,9 +97,7 @@ export function SourceRow({ }`} > {isReady - ? `${isLibrarySource ? "Official Library" : "Processed"} · ${ - source.chunkCount ?? 0 - } chunks` + ? `${getReadySourceLabel(source)} · ${source.chunkCount ?? 0} chunks` : source.status === "parsing" ? "Preparing" : source.status === "uploading" @@ -162,6 +161,12 @@ export function SourceRow({ ); } +function getReadySourceLabel(source: SourceView): string { + if (source.officialLibrary !== undefined) return "Official Library"; + if (source.kind === "remote") return "Remote"; + return "Processed"; +} + function fileIconTint(title: string): { bg: string; fg: string } { const ext = title.split(".").pop()?.toLowerCase(); switch (ext) { diff --git a/src/components/workspace-chat-workflow.ts b/src/components/workspace-chat-workflow.ts index c108134..2943a23 100644 --- a/src/components/workspace-chat-workflow.ts +++ b/src/components/workspace-chat-workflow.ts @@ -250,7 +250,8 @@ export function useWorkspaceChatWorkflow({ async function handleChatSend(text: string): Promise { const sendStart = Date.now() const selectedSourcesCount = sources.filter( - (source) => source.status === "ready" && !source.excludedFromQuery, + (source) => + isQueryableReadySource(source) && !source.excludedFromQuery, ).length const demoSourceIds = getMaterializableDemoSourceIds(sources) if (demoSourceIds.length > 0) { @@ -396,10 +397,14 @@ function getMaterializableDemoSourceIds( } function hasQueryableReadySource(sources: readonly SourceView[]): boolean { - return sources.some( - (source) => - source.status === "ready" && - !isUnmaterializedOfficialLibrarySource(source), + return sources.some(isQueryableReadySource) +} + +function isQueryableReadySource(source: SourceView): boolean { + return ( + source.status === "ready" && + !isUnmaterializedOfficialLibrarySource(source) && + source.kind !== "remote" ) } diff --git a/src/components/workspace-source-workflow.ts b/src/components/workspace-source-workflow.ts index 341eda5..a646323 100644 --- a/src/components/workspace-source-workflow.ts +++ b/src/components/workspace-source-workflow.ts @@ -221,13 +221,17 @@ export function useWorkspaceSourceWorkflow({ function isQueryableReadySource(source: SourceView): boolean { if (source.status !== "ready") return false - return !isUnmaterializedOfficialLibrarySource(source) + return !isUnmaterializedOfficialLibrarySource(source) && !isRemoteSource(source) } function isUnmaterializedOfficialLibrarySource(source: SourceView): boolean { return source.kind === "demo" && source.officialLibrary !== undefined } +function isRemoteSource(source: SourceView): boolean { + return source.kind === "remote" +} + function archiveSourceMutation( _key: string, { arg: sourceId }: { readonly arg: string }, diff --git a/src/domains/chat/route-answer.ts b/src/domains/chat/route-answer.ts index 06c643a..ea75943 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -12,7 +12,6 @@ import { } from "@/domains/chat/service" import { chatTurnPersistence } from "@/domains/chat/chat-turn-persistence" import { startBackgroundReconciliation } from "@/domains/sources/background-reconcile" -import { localizeRemoteLibrarySources } from "@/domains/sources/remote-library" import { sourceService } from "@/domains/sources/service" import { sourceWorkflowRuntime } from "@/domains/sources/workflow-runtime" import { notebookRequestContext } from "@/domains/workspace/request-context" @@ -69,13 +68,6 @@ const answerChatEffect = (input: AnswerChatInput) => apiKey, }), ) - const compatibleSources = yield* localizeRemoteLibrarySources({ - workspace, - client, - localSources: sources, - localizeDocument: (document) => - sourceService.localizeRemoteDocument(workspace.id, document), - }) const loadSourceAssetUrls = (source: (typeof sources)[number]) => sourceService.getParseAssetUrls(workspace.id, source.id) @@ -83,7 +75,7 @@ const answerChatEffect = (input: AnswerChatInput) => yield* Effect.tryPromise(() => handleChatTurn({ workspace, - sources: compatibleSources, + sources, question: body.value.question, threadId: body.value.threadId, excludedSourceIds: body.value.excludedSourceIds, @@ -93,7 +85,7 @@ const answerChatEffect = (input: AnswerChatInput) => hardenMediaAssetUrls: ({ results, artifacts }) => hardenChatMediaAssetUrls({ workspaceId: workspace.id, - sources: compatibleSources, + sources, results, artifacts, loadSourceAssetUrls, diff --git a/src/domains/chunks/index.ts b/src/domains/chunks/index.ts index d774708..035d3ef 100644 --- a/src/domains/chunks/index.ts +++ b/src/domains/chunks/index.ts @@ -20,6 +20,10 @@ export type ChunkKnowhereClient = { includeAssetUrls: boolean }, ): Promise<{ + documentId?: string + namespace?: string + jobId?: string | null + jobResultId?: string | null chunks: DocumentChunk[] pagination?: { page?: number diff --git a/src/domains/chunks/server.test.ts b/src/domains/chunks/server.test.ts new file mode 100644 index 0000000..f58f5ea --- /dev/null +++ b/src/domains/chunks/server.test.ts @@ -0,0 +1,440 @@ +import { describe, expect, it, vi, type Mock } from "vitest" +import { Effect } from "effect" +import type { DocumentChunk } from "@ontos-ai/knowhere-sdk" + +import type { Source } from "@/infrastructure/db/schema" +import { + loadChunkPageForSource, + loadChunksForSource, +} from "./server" + +describe("server chunk cache", () => { + it("returns upstream chunks on a visible cache miss and warms mirrored assets in the background", async () => { + const warmTasks: Array<() => Promise> = [] + const cacheStore = createCacheStore() + const listChunks = vi.fn(async () => ({ + documentId: "doc_1", + jobResultId: "revision_1", + chunks: [ + makeDocumentChunk({ + id: "image_1", + chunkId: "parser_image_1", + chunkType: "image", + filePath: "images/image-1.png", + assetUrl: "https://knowhere.example/assets/image-1.png", + }), + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + })) + const fetchAsset = vi.fn(async () => + new Response("image-body", { + headers: { "content-type": "image/png" }, + }), + ) + + const page = await Effect.runPromise( + loadChunkPageForSource( + makeSource({ knowhereJobId: "revision_1" }), + { documents: { listChunks } }, + { page: 1, pageSize: 1 }, + { + cacheStore, + fetchAsset, + scheduleWarm: (task) => warmTasks.push(task), + workspaceId: "workspace_1", + }, + ), + ) + + expect(page.chunks[0]?.assetUrl).toBe( + "https://knowhere.example/assets/image-1.png", + ) + expect(cacheStore.putMock).not.toHaveBeenCalled() + expect(warmTasks).toHaveLength(1) + + await warmTasks[0]?.() + + expect(fetchAsset).toHaveBeenCalledWith( + "https://knowhere.example/assets/image-1.png", + ) + expect(cacheStore.putMock).toHaveBeenCalledWith( + expect.stringContaining("/chunk-assets/revision_1/"), + Buffer.from("image-body"), + expect.objectContaining({ contentType: "image/png" }), + ) + const cachedPagePut = cacheStore.putMock.mock.calls.find( + ([pathname]) => + typeof pathname === "string" && pathname.endsWith(".json"), + ) + expect(cachedPagePut).toBeDefined() + const cachedPage = JSON.parse(String(cachedPagePut?.[1])) as { + readonly chunks: readonly { readonly assetUrl?: string }[] + } + expect(cachedPage.chunks[0]?.assetUrl).toContain( + "https://blob.example/workspaces/workspace_1/sources/source_1/chunk-assets/revision_1/", + ) + }) + + it("returns a cached visible page after verifying the current Knowhere revision", async () => { + const cachedPage = { + chunks: [ + { + chunkId: "image_1", + documentId: "doc_1", + type: "image", + content: "", + assetUrl: "https://blob.example/image-1.png", + sourceTitle: "notes.pdf", + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + } + const listChunks = vi.fn(async () => ({ + documentId: "doc_1", + jobResultId: "revision_1", + chunks: [], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + })) + const cacheStore = createCacheStore({ + get: vi.fn(async () => ({ + statusCode: 200, + stream: createTextStream(JSON.stringify(cachedPage)), + })), + }) + + const page = await Effect.runPromise( + loadChunkPageForSource( + makeSource({ knowhereJobId: "revision_1" }), + { documents: { listChunks } }, + { page: 1, pageSize: 1 }, + { + cacheStore, + workspaceId: "workspace_1", + }, + ), + ) + + expect(page).toEqual(cachedPage) + expect(listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 1, + includeAssetUrls: false, + }) + expect(cacheStore.getMock).toHaveBeenCalledWith( + expect.stringContaining("/revision_1/visible/page-1-size-1.json"), + { access: "public" }, + ) + }) + + it("ignores old cached pages when Knowhere reports a new job id", async () => { + const warmTasks: Array<() => Promise> = [] + const staleCachedPage = { + chunks: [ + { + chunkId: "stale_1", + documentId: "doc_1", + type: "text", + content: "Old chunk", + sourceTitle: "notes.pdf", + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + } + const cacheStore = createCacheStore({ + get: vi.fn(async (pathname: string) => + pathname.includes("/job_old/") + ? { + statusCode: 200, + stream: createTextStream(JSON.stringify(staleCachedPage)), + } + : null, + ), + }) + const listChunks = vi.fn(async ( + _documentId: string, + params: { readonly includeAssetUrls: boolean }, + ) => ({ + documentId: "doc_1", + jobId: "job_new", + chunks: params.includeAssetUrls + ? [ + makeDocumentChunk({ + id: "text_new", + chunkId: "parser_text_new", + content: "New chunk", + }), + ] + : [], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + })) + const onRevisionKey = vi.fn(async () => undefined) + + const page = await Effect.runPromise( + loadChunkPageForSource( + makeSource({ knowhereJobId: "job_old" }), + { documents: { listChunks } }, + { page: 1, pageSize: 1 }, + { + cacheStore, + onRevisionKey, + scheduleWarm: (task) => warmTasks.push(task), + workspaceId: "workspace_1", + }, + ), + ) + + expect(page.chunks[0]?.content).toBe("New chunk") + expect(cacheStore.getMock).toHaveBeenCalledWith( + expect.stringContaining("/job_new/visible/page-1-size-1.json"), + { access: "public" }, + ) + expect(cacheStore.getMock).not.toHaveBeenCalledWith( + expect.stringContaining("/job_old/visible/page-1-size-1.json"), + expect.anything(), + ) + expect(onRevisionKey).toHaveBeenCalledWith("job_new") + expect(warmTasks).toHaveLength(1) + }) + + it("uses structure-only chunk loading for full-tree requests", async () => { + const warmTasks: Array<() => Promise> = [] + const cacheStore = createCacheStore() + const listChunks = vi.fn(async () => ({ + documentId: "doc_1", + jobResultId: "revision_1", + chunks: [ + makeDocumentChunk({ + id: "text_1", + chunkId: "parser_text_1", + content: "A text chunk", + }), + ], + pagination: { + page: 1, + pageSize: 200, + total: 1, + totalPages: 1, + }, + })) + const fetchAsset = vi.fn() + + const chunks = await Effect.runPromise( + loadChunksForSource( + makeSource({ knowhereJobId: "revision_1" }), + { documents: { listChunks } }, + { + cacheStore, + fetchAsset, + scheduleWarm: (task) => warmTasks.push(task), + workspaceId: "workspace_1", + }, + ), + ) + + expect(chunks).toMatchObject([{ chunkId: "text_1" }]) + expect(listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 200, + includeAssetUrls: false, + }) + expect(fetchAsset).not.toHaveBeenCalled() + expect(warmTasks).toHaveLength(1) + await warmTasks[0]?.() + expect(cacheStore.putMock).toHaveBeenCalledWith( + expect.stringContaining("/structure/page-1-size-200.json"), + expect.any(String), + expect.objectContaining({ contentType: "application/json; charset=utf-8" }), + ) + }) + + it("caches media chunks without previews when Knowhere has no upstream asset URL", async () => { + const warmTasks: Array<() => Promise> = [] + const cacheStore = createCacheStore() + const listChunks = vi.fn(async () => ({ + documentId: "doc_1", + jobResultId: "revision_1", + chunks: [ + makeDocumentChunk({ + id: "image_1", + chunkId: "parser_image_1", + chunkType: "image", + filePath: "images/image-1.png", + assetUrl: null, + }), + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + })) + const fetchAsset = vi.fn() + + const page = await Effect.runPromise( + loadChunkPageForSource( + makeSource({ knowhereJobId: "revision_1" }), + { documents: { listChunks } }, + { page: 1, pageSize: 1 }, + { + cacheStore, + fetchAsset, + scheduleWarm: (task) => warmTasks.push(task), + workspaceId: "workspace_1", + }, + ), + ) + + expect(page.chunks[0]?.assetUrl).toBeUndefined() + await warmTasks[0]?.() + expect(fetchAsset).not.toHaveBeenCalled() + const cachedPagePut = cacheStore.putMock.mock.calls.find( + ([pathname]) => + typeof pathname === "string" && pathname.endsWith(".json"), + ) + const cachedPage = JSON.parse(String(cachedPagePut?.[1])) as { + readonly chunks: readonly { readonly assetUrl?: string }[] + } + expect(cachedPage.chunks[0]?.assetUrl).toBeUndefined() + }) +}) + +type TestCacheGetResult = + | { + readonly statusCode: 200 + readonly stream: ReadableStream + } + | { + readonly statusCode: 304 + readonly stream: null + } + +type TestCachePutOptions = { + readonly access: "public" + readonly allowOverwrite: boolean + readonly contentType: string + readonly multipart?: boolean +} + +type TestCacheStore = { + readonly get: ( + pathname: string, + options: { readonly access: "public" }, + ) => Promise + readonly put: ( + pathname: string, + body: string | Buffer, + options: TestCachePutOptions, + ) => Promise<{ readonly url: string }> + readonly getMock: TestCacheGetMock + readonly putMock: TestCachePutMock +} + +type TestCacheGetMock = Mock< + ( + pathname: string, + options: { readonly access: "public" }, + ) => Promise +> + +type TestCachePutMock = Mock< + ( + pathname: string, + body: string | Buffer, + options: TestCachePutOptions, + ) => Promise<{ readonly url: string }> +> + +function createCacheStore(overrides: Partial<{ + readonly get: TestCacheGetMock + readonly put: TestCachePutMock +}> = {}): TestCacheStore { + const getMock = + overrides.get ?? + vi.fn(async () => null) + const putMock = + overrides.put ?? + vi.fn(async (pathname: string) => ({ + url: `https://blob.example/${pathname}`, + })) + + return { + get: (pathname, options) => getMock(pathname, options), + put: (pathname, body, options) => putMock(pathname, body, options), + getMock, + putMock, + } +} + +function makeDocumentChunk( + overrides: Partial = {}, +): DocumentChunk { + return { + id: "document_chunk_1", + chunkId: "parser_chunk_1", + chunkType: "text", + content: "Chunk content", + sectionId: null, + sectionPath: null, + sourceChunkPath: null, + filePath: null, + sortOrder: 1, + metadata: {}, + assetUrl: null, + ...overrides, + } +} + +function makeSource(overrides: Partial = {}): Source { + return { + id: "source_1", + workspaceId: "workspace_1", + title: "notes.pdf", + mimeType: "application/pdf", + sizeBytes: 100, + status: "ready", + failureReason: null, + knowhereJobId: "revision_1", + knowhereDocumentId: "doc_1", + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-05-06T00:00:00Z"), + updatedAt: new Date("2026-05-06T00:00:00Z"), + deletedAt: null, + ...overrides, + } +} + +function createTextStream(text: string): ReadableStream { + const stream = new Response(text).body + if (!stream) throw new Error("Response body stream was not created.") + return stream +} diff --git a/src/domains/chunks/server.ts b/src/domains/chunks/server.ts new file mode 100644 index 0000000..a28ee3b --- /dev/null +++ b/src/domains/chunks/server.ts @@ -0,0 +1,739 @@ +import "server-only" + +import path from "node:path" +import { createHash } from "node:crypto" +import { get as getBlob, put } from "@vercel/blob" +import { after } from "next/server" +import { Effect } from "effect" +import type { + DocumentChunk, + DocumentChunkListResponse, +} from "@ontos-ai/knowhere-sdk" + +import { + resolveChunkConnectionTargets, + toParsedChunkView, + type ChunkKnowhereClient, + type ChunkPage, + type ChunkPageParams, + type LoadChunksOptions, +} from "@/domains/chunks" +import type { ParsedChunkView } from "@/domains/chunks/types" +import type { Source } from "@/infrastructure/db/schema" +import { logger } from "@/lib/logger" + +type ChunkPageMode = "visible" | "structure" + +type ChunkPageBlobGetResult = + | { + readonly statusCode: 200 + readonly stream: ReadableStream + } + | { + readonly statusCode: 304 + readonly stream: null + } + +type ChunkPageBlobPutOptions = { + readonly access: "public" + readonly allowOverwrite: boolean + readonly contentType: string + readonly multipart?: boolean +} + +type ChunkPageBlobStore = { + readonly get: ( + pathname: string, + options: { readonly access: "public" }, + ) => Promise + readonly put: ( + pathname: string, + body: string | Buffer, + options: ChunkPageBlobPutOptions, + ) => Promise<{ readonly url: string }> +} + +type FetchChunkAsset = (assetUrl: string) => Promise + +type ChunkPageWarmScheduler = (task: () => Promise) => void + +type ServerLoadChunksOptions = LoadChunksOptions & { + readonly workspaceId?: string + readonly cacheStore?: ChunkPageBlobStore + readonly fetchAsset?: FetchChunkAsset + readonly mode?: ChunkPageMode + readonly onRevisionKey?: (revisionKey: string) => Promise + readonly scheduleWarm?: ChunkPageWarmScheduler +} + +type WarmChunkPageCacheInput = { + readonly source: Source + readonly client: ChunkKnowhereClient + readonly params: ChunkPageParams + readonly revisionKey: string + readonly workspaceId: string + readonly cacheStore: ChunkPageBlobStore + readonly fetchAsset: FetchChunkAsset + readonly scheduleWarm: ChunkPageWarmScheduler + readonly startAssetIndex?: number + readonly mirroredAssetUrlsByOriginalUrl?: Readonly> +} + +type MirrorableChunkAsset = { + readonly chunkId: string + readonly chunkType: "image" | "table" + readonly assetUrl: string + readonly filePath?: string +} + +const documentChunkPageSize = 200 +const visibleChunkPageMode: ChunkPageMode = "visible" +const structureChunkPageMode: ChunkPageMode = "structure" +const maximumMirroredAssetsPerWarmStep = 50 +const maximumWarmStepDurationMs = 45_000 +const assetMirrorConcurrency = 10 + +const defaultBlobStore: ChunkPageBlobStore = { + get: (pathname, options) => getBlob(pathname, options), + put: (pathname, body, options) => + put(pathname, body, { + access: options.access, + allowOverwrite: options.allowOverwrite, + contentType: options.contentType, + multipart: options.multipart, + }), +} + +const defaultFetchAsset: FetchChunkAsset = (assetUrl: string) => fetch(assetUrl) + +const defaultScheduleWarm: ChunkPageWarmScheduler = ( + task: () => Promise, +) => { + try { + after(task) + } catch { + void task() + } +} + +export const loadChunksForSource = ( + source: Source, + client: ChunkKnowhereClient, + options: ServerLoadChunksOptions = {}, +) => + Effect.gen(function* () { + if (source.status !== "ready" || !source.knowhereDocumentId) return [] + + const chunks: ParsedChunkView[] = [] + let page = 1 + let totalPages = 1 + + do { + const chunkPage = yield* loadChunkPageForSource(source, client, { + page, + pageSize: documentChunkPageSize, + }, { + ...options, + mode: structureChunkPageMode, + }) + chunks.push(...chunkPage.chunks) + totalPages = chunkPage.pagination.totalPages + page += 1 + } while (page <= totalPages) + + return resolveChunkConnectionTargets(chunks) + }) + +export const loadChunkPageForSource = ( + source: Source, + client: ChunkKnowhereClient, + params: ChunkPageParams, + options: ServerLoadChunksOptions = {}, +) => + Effect.gen(function* () { + const emptyPage = createEmptyChunkPage(params) + if (source.status !== "ready" || !source.knowhereDocumentId) { + return emptyPage + } + + const mode = options.mode ?? visibleChunkPageMode + const workspaceId = options.workspaceId ?? source.workspaceId + const cacheStore = options.cacheStore ?? defaultBlobStore + const includeAssetUrls = mode === visibleChunkPageMode + const revisionProbeResponse = yield* Effect.promise(() => + client.documents.listChunks(source.knowhereDocumentId!, { + page: params.page, + pageSize: params.pageSize, + includeAssetUrls: false, + }), + ) + const probeRevisionKey = getRevisionKey(revisionProbeResponse, source) + if (probeRevisionKey) { + scheduleRevisionKeyUpdate(source, probeRevisionKey, options.onRevisionKey) + const cachedPage = yield* Effect.promise(() => + readCachedChunkPage({ + cacheStore, + documentId: source.knowhereDocumentId!, + mode, + params, + revisionKey: probeRevisionKey, + workspaceId, + }), + ) + if (cachedPage) return cachedPage + } + + const response = includeAssetUrls + ? yield* Effect.promise(() => + client.documents.listChunks(source.knowhereDocumentId!, { + page: params.page, + pageSize: params.pageSize, + includeAssetUrls, + }), + ) + : revisionProbeResponse + const revisionKey = getRevisionKey(response, source) ?? probeRevisionKey + if (revisionKey && revisionKey !== probeRevisionKey) { + scheduleRevisionKeyUpdate(source, revisionKey, options.onRevisionKey) + } + + const chunkPage = createChunkPageFromResponse({ + response, + source, + params, + options: + mode === visibleChunkPageMode + ? { assetUrlsByFilePath: options.assetUrlsByFilePath } + : {}, + }) + + if (revisionKey) { + if (mode === visibleChunkPageMode) { + scheduleChunkPageWarm({ + source, + client, + params, + revisionKey, + workspaceId, + cacheStore, + fetchAsset: options.fetchAsset ?? defaultFetchAsset, + scheduleWarm: options.scheduleWarm ?? defaultScheduleWarm, + }) + } else { + scheduleStructurePageCacheWrite({ + cacheStore, + chunkPage, + documentId: source.knowhereDocumentId, + mode, + params, + revisionKey, + scheduleWarm: options.scheduleWarm ?? defaultScheduleWarm, + workspaceId, + }) + } + } + + return chunkPage + }) + +export async function warmChunkPageCache( + input: WarmChunkPageCacheInput, +): Promise { + const response = await input.client.documents.listChunks( + input.source.knowhereDocumentId!, + { + page: input.params.page, + pageSize: input.params.pageSize, + includeAssetUrls: true, + }, + ) + const assets = collectMirrorableChunkAssets(response.chunks) + const startAssetIndex = input.startAssetIndex ?? 0 + const mirroredAssetUrls = new Map( + Object.entries(input.mirroredAssetUrlsByOriginalUrl ?? {}), + ) + + const stepStartedAt = Date.now() + const assetBatch = assets.slice( + startAssetIndex, + startAssetIndex + maximumMirroredAssetsPerWarmStep, + ) + const mirroredBatch = await mirrorChunkAssets({ + assets: assetBatch, + cacheStore: input.cacheStore, + fetchAsset: input.fetchAsset, + revisionKey: input.revisionKey, + source: input.source, + }) + for (const mirroredAsset of mirroredBatch) { + mirroredAssetUrls.set(mirroredAsset.assetUrl, mirroredAsset.blobUrl) + } + + const nextAssetIndex = startAssetIndex + assetBatch.length + if ( + nextAssetIndex < assets.length && + Date.now() - stepStartedAt < maximumWarmStepDurationMs + ) { + input.scheduleWarm(() => + warmChunkPageCache({ + ...input, + startAssetIndex: nextAssetIndex, + mirroredAssetUrlsByOriginalUrl: Object.fromEntries(mirroredAssetUrls), + }), + ) + return + } + + if (nextAssetIndex < assets.length) { + input.scheduleWarm(() => + warmChunkPageCache({ + ...input, + startAssetIndex: nextAssetIndex, + mirroredAssetUrlsByOriginalUrl: Object.fromEntries(mirroredAssetUrls), + }), + ) + return + } + + const rewrittenChunks = response.chunks.map((chunk) => + rewriteChunkAssetUrl(chunk, mirroredAssetUrls), + ) + const chunkPage = createChunkPageFromResponse({ + response: { + ...response, + chunks: rewrittenChunks, + }, + source: input.source, + params: input.params, + options: {}, + }) + + await writeCachedChunkPage({ + cacheStore: input.cacheStore, + chunkPage, + documentId: input.source.knowhereDocumentId!, + mode: visibleChunkPageMode, + params: input.params, + revisionKey: input.revisionKey, + workspaceId: input.workspaceId, + }) +} + +function scheduleChunkPageWarm(input: WarmChunkPageCacheInput): void { + input.scheduleWarm(async () => { + try { + await warmChunkPageCache(input) + } catch (error) { + logger.warn("chunks: chunk page cache warm failed", { + sourceId: input.source.id, + documentId: input.source.knowhereDocumentId, + page: input.params.page, + pageSize: input.params.pageSize, + revisionKey: input.revisionKey, + error: getErrorMessage(error), + }) + } + }) +} + +function scheduleStructurePageCacheWrite(input: { + readonly cacheStore: ChunkPageBlobStore + readonly chunkPage: ChunkPage + readonly documentId: string + readonly mode: ChunkPageMode + readonly params: ChunkPageParams + readonly revisionKey: string + readonly scheduleWarm: ChunkPageWarmScheduler + readonly workspaceId: string +}): void { + input.scheduleWarm(async () => { + try { + await writeCachedChunkPage(input) + } catch (error) { + logger.warn("chunks: structure chunk page cache write failed", { + documentId: input.documentId, + page: input.params.page, + pageSize: input.params.pageSize, + revisionKey: input.revisionKey, + error: getErrorMessage(error), + }) + } + }) +} + +async function readCachedChunkPage(input: { + readonly cacheStore: ChunkPageBlobStore + readonly documentId: string + readonly mode: ChunkPageMode + readonly params: ChunkPageParams + readonly revisionKey: string + readonly workspaceId: string +}): Promise { + const pathname = getChunkPageCachePathname(input) + const result = await input.cacheStore.get(pathname, { access: "public" }) + if (!result || result.statusCode !== 200) return null + + const text = await new Response(result.stream).text() + return parseCachedChunkPage(text) +} + +async function writeCachedChunkPage(input: { + readonly cacheStore: ChunkPageBlobStore + readonly chunkPage: ChunkPage + readonly documentId: string + readonly mode: ChunkPageMode + readonly params: ChunkPageParams + readonly revisionKey: string + readonly workspaceId: string +}): Promise { + await input.cacheStore.put( + getChunkPageCachePathname(input), + JSON.stringify(input.chunkPage), + { + access: "public", + allowOverwrite: true, + contentType: "application/json; charset=utf-8", + }, + ) +} + +function createChunkPageFromResponse(input: { + readonly response: { + readonly chunks: readonly DocumentChunk[] + readonly pagination?: { + readonly page?: number + readonly pageSize?: number + readonly total?: number + readonly totalPages?: number + } + } + readonly source: Source + readonly params: ChunkPageParams + readonly options: LoadChunksOptions +}): ChunkPage { + const chunks = input.response.chunks.map((chunk) => + toParsedChunkView( + chunk, + input.source.title, + input.source.knowhereDocumentId ?? undefined, + input.options, + ), + ) + + return { + chunks, + pagination: { + page: getFiniteNonNegativeNumber( + input.response.pagination?.page, + input.params.page, + ), + pageSize: getFiniteNonNegativeNumber( + input.response.pagination?.pageSize, + input.params.pageSize, + ), + total: getFiniteNonNegativeNumber( + input.response.pagination?.total, + chunks.length, + ), + totalPages: getFiniteNonNegativeNumber( + input.response.pagination?.totalPages, + Math.ceil(chunks.length / input.params.pageSize), + ), + }, + } +} + +function createEmptyChunkPage(params: ChunkPageParams): ChunkPage { + return { + chunks: [], + pagination: { + page: params.page, + pageSize: params.pageSize, + total: 0, + totalPages: 0, + }, + } +} + +function parseCachedChunkPage(text: string): ChunkPage | null { + try { + const value: unknown = JSON.parse(text) + if (!isRecord(value)) return null + const chunks = value["chunks"] + const pagination = value["pagination"] + if (!Array.isArray(chunks) || !isRecord(pagination)) return null + const page = getNumber(pagination["page"]) + const pageSize = getNumber(pagination["pageSize"]) + const total = getNumber(pagination["total"]) + const totalPages = getNumber(pagination["totalPages"]) + if ( + page === undefined || + pageSize === undefined || + total === undefined || + totalPages === undefined + ) { + return null + } + + return { + chunks: chunks.filter(isParsedChunkView), + pagination: { + page, + pageSize, + total, + totalPages, + }, + } + } catch { + return null + } +} + +async function mirrorChunkAssets(input: { + readonly assets: readonly MirrorableChunkAsset[] + readonly cacheStore: ChunkPageBlobStore + readonly fetchAsset: FetchChunkAsset + readonly revisionKey: string + readonly source: Source +}): Promise< + readonly { + readonly assetUrl: string + readonly blobUrl: string + }[] +> { + return mapWithConcurrency( + input.assets, + assetMirrorConcurrency, + async (asset) => { + const response = await input.fetchAsset(asset.assetUrl) + if (!response.ok) { + throw new Error( + `Chunk asset fetch failed with status ${response.status}.`, + ) + } + + const body = Buffer.from(await response.arrayBuffer()) + const blob = await input.cacheStore.put( + getMirroredAssetPathname({ + asset, + revisionKey: input.revisionKey, + source: input.source, + }), + body, + { + access: "public", + allowOverwrite: true, + contentType: getMirroredAssetContentType(asset, response), + multipart: true, + }, + ) + return { + assetUrl: asset.assetUrl, + blobUrl: blob.url, + } + }, + ) +} + +async function mapWithConcurrency( + inputs: readonly Input[], + concurrency: number, + mapInput: (input: Input) => Promise, +): Promise { + const results: Array = [] + let nextIndex = 0 + const workerCount = Math.min(concurrency, inputs.length) + + async function runWorker(): Promise { + while (nextIndex < inputs.length) { + const index = nextIndex + nextIndex += 1 + const input = inputs[index] + if (input === undefined) return + results[index] = await mapInput(input) + } + } + + await Promise.all( + Array.from({ length: workerCount }, () => runWorker()), + ) + + return results.filter( + (result): result is Output => result !== undefined, + ) +} + +function collectMirrorableChunkAssets( + chunks: readonly DocumentChunk[], +): readonly MirrorableChunkAsset[] { + return chunks.flatMap((chunk): readonly MirrorableChunkAsset[] => { + if (chunk.chunkType !== "image" && chunk.chunkType !== "table") return [] + if (typeof chunk.assetUrl !== "string" || chunk.assetUrl.length === 0) { + return [] + } + + return [ + { + chunkId: chunk.id, + chunkType: chunk.chunkType, + assetUrl: chunk.assetUrl, + ...(chunk.filePath ? { filePath: chunk.filePath } : {}), + }, + ] + }) +} + +function rewriteChunkAssetUrl( + chunk: DocumentChunk, + mirroredAssetUrls: ReadonlyMap, +): DocumentChunk { + if (typeof chunk.assetUrl !== "string") return chunk + const mirroredAssetUrl = mirroredAssetUrls.get(chunk.assetUrl) + return mirroredAssetUrl ? { ...chunk, assetUrl: mirroredAssetUrl } : chunk +} + +function getChunkPageCachePathname(input: { + readonly documentId: string + readonly mode: ChunkPageMode + readonly params: ChunkPageParams + readonly revisionKey: string + readonly workspaceId: string +}): string { + return [ + "workspaces", + encodePathSegment(input.workspaceId), + "chunk-pages", + encodePathSegment(input.documentId), + encodePathSegment(input.revisionKey), + input.mode, + `page-${input.params.page}-size-${input.params.pageSize}.json`, + ].join("/") +} + +function getMirroredAssetPathname(input: { + readonly asset: MirrorableChunkAsset + readonly revisionKey: string + readonly source: Source +}): string { + const basename = getMirroredAssetBasename(input.asset) + return [ + "workspaces", + encodePathSegment(input.source.workspaceId), + "sources", + encodePathSegment(input.source.id), + "chunk-assets", + encodePathSegment(input.revisionKey), + `${hashValue(input.asset.assetUrl)}-${basename}`, + ].join("/") +} + +function getMirroredAssetBasename(asset: MirrorableChunkAsset): string { + const candidate = + asset.filePath ?? + getUrlPathname(asset.assetUrl).split("/").filter(Boolean).at(-1) ?? + `${asset.chunkId}.bin` + return candidate.replace(/[^a-zA-Z0-9._-]+/g, "-") +} + +function getMirroredAssetContentType( + asset: MirrorableChunkAsset, + response: Response, +): string { + const headerContentType = response.headers.get("content-type") + if (headerContentType) return headerContentType + if (asset.chunkType === "table") return "text/html; charset=utf-8" + + const extension = path.extname(asset.filePath ?? getUrlPathname(asset.assetUrl)) + .toLowerCase() + if (extension === ".jpg" || extension === ".jpeg") return "image/jpeg" + if (extension === ".png") return "image/png" + if (extension === ".webp") return "image/webp" + if (extension === ".gif") return "image/gif" + return "application/octet-stream" +} + +function getRevisionKey( + response: Pick, + source: Source, +): string | null { + return ( + getNonEmptyString(response.jobId) ?? + getNonEmptyString(response.jobResultId) ?? + getFallbackRevisionKey(source) + ) +} + +function getFallbackRevisionKey(source: Source): string | null { + return ( + getNonEmptyString(source.knowhereJobId) ?? + getNonEmptyString(source.knowhereDocumentId) + ) +} + +function scheduleRevisionKeyUpdate( + source: Source, + revisionKey: string, + onRevisionKey: ((revisionKey: string) => Promise) | undefined, +): void { + if (!onRevisionKey || source.knowhereJobId === revisionKey) return + + void onRevisionKey(revisionKey).catch((error: unknown) => { + logger.warn("chunks: source revision key update failed", { + sourceId: source.id, + documentId: source.knowhereDocumentId, + revisionKey, + error: getErrorMessage(error), + }) + }) +} + +function getFiniteNonNegativeNumber( + value: number | undefined, + fallback: number, +): number { + return typeof value === "number" && Number.isFinite(value) && value >= 0 + ? value + : fallback +} + +function isParsedChunkView(value: unknown): value is ParsedChunkView { + if (!isRecord(value)) return false + return ( + typeof value["chunkId"] === "string" && + typeof value["type"] === "string" && + typeof value["content"] === "string" && + typeof value["sourceTitle"] === "string" + ) +} + +function getNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) ? value : undefined +} + +function getNonEmptyString(value: unknown): string | null { + return typeof value === "string" && value.length > 0 ? value : null +} + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null && !Array.isArray(value) +} + +function getUrlPathname(value: string): string { + try { + return new URL(value).pathname + } catch { + return value.split("?")[0] ?? value + } +} + +function encodePathSegment(value: string): string { + return encodeURIComponent(value) +} + +function hashValue(value: string): string { + return createHash("sha256").update(value).digest("hex").slice(0, 16) +} + +function getErrorMessage(error: unknown): string { + return error instanceof Error ? error.message : String(error) +} diff --git a/src/domains/sources/lifecycle.ts b/src/domains/sources/lifecycle.ts index 98c11fb..dbf4c9b 100644 --- a/src/domains/sources/lifecycle.ts +++ b/src/domains/sources/lifecycle.ts @@ -5,17 +5,8 @@ import type { JobResult } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" import { logger } from "@/lib/logger" -import type { - StoreParsedResultAssetsInput, - StoredParsedResultAssets, -} from "./parsed-result-assets" type SourceLifecycleRepository = { - saveSourceParseResult( - workspaceId: string, - sourceId: string, - input: StoredParsedResultAssets, - ): Promise markSourceReady( workspaceId: string, sourceId: string, @@ -30,12 +21,6 @@ type SourceLifecycleRepository = { clearSourceStagedBlob(workspaceId: string, sourceId: string): Promise } -type SourceLifecycleParsedResultStore = { - storeParsedResultAssets( - input: Omit, - ): Promise -} - type SourceLifecycleBlobStore = { deleteStagedSourceBlob(pathname: string): Promise } @@ -44,9 +29,7 @@ type ApplyKnowhereJobToSourceInput = { workspaceId: string source: Source job: JobResult - client: StoreParsedResultAssetsInput["client"] repository: SourceLifecycleRepository - parsedResultStore: SourceLifecycleParsedResultStore blobStore: SourceLifecycleBlobStore } @@ -61,13 +44,11 @@ export const applyKnowhereJobToSourceEffect = Effect.fn( workspaceId, source, job, - client, repository, - parsedResultStore, blobStore, }: ApplyKnowhereJobToSourceInput) { - // Best-effort early exit: skip expensive asset uploads when the source has - // already been resolved. The atomic guard (Layer 3) is in the DB UPDATE below. + // Best-effort early exit: skip duplicate status transitions when the source + // has already been resolved. The atomic guard is in the DB UPDATE below. if (source.status !== "parsing") return if (job.isDone || job.status === "done") { @@ -77,17 +58,6 @@ export const applyKnowhereJobToSourceEffect = Effect.fn( jobId: source.knowhereJobId ?? "", documentId: job.documentId, }) - const stored = yield* Effect.tryPromise(() => - parsedResultStore.storeParsedResultAssets({ - workspaceId, - sourceId: source.id, - job, - client, - }), - ) - yield* Effect.tryPromise(() => - repository.saveSourceParseResult(workspaceId, source.id, stored), - ) yield* Effect.tryPromise(() => repository.markSourceReady(workspaceId, source.id, job.documentId!), ) @@ -154,9 +124,7 @@ export async function applyKnowhereJobToSource({ workspaceId, source, job, - client, repository, - parsedResultStore, blobStore, }: ApplyKnowhereJobToSourceInput): Promise { return Effect.runPromise( @@ -164,9 +132,7 @@ export async function applyKnowhereJobToSource({ workspaceId, source, job, - client, repository, - parsedResultStore, blobStore, }), ) diff --git a/src/domains/sources/reconcile.test.ts b/src/domains/sources/reconcile.test.ts index 808fd82..4c90172 100644 --- a/src/domains/sources/reconcile.test.ts +++ b/src/domains/sources/reconcile.test.ts @@ -38,17 +38,10 @@ async function loadReconcile({ listSourcesForWorkspace, markSourceFailed = vi.fn().mockResolvedValue(undefined), markSourceReady = vi.fn().mockResolvedValue(undefined), - saveSourceParseResult = vi.fn().mockResolvedValue(undefined), - storeParsedResultAssets = vi.fn().mockResolvedValue({ - resultBlobUrl: "https://blob.example/result.zip", - assetUrlsByFilePath: {}, - }), }: { listSourcesForWorkspace: ReturnType markSourceFailed?: ReturnType markSourceReady?: ReturnType - saveSourceParseResult?: ReturnType - storeParsedResultAssets?: ReturnType }): Promise { vi.resetModules() vi.doMock("./workflow-runtime", () => ({ @@ -56,13 +49,9 @@ async function loadReconcile({ listForWorkspace: listSourcesForWorkspace, markFailed: markSourceFailed, markReady: markSourceReady, - saveParseResult: saveSourceParseResult, clearStagedBlob: vi.fn(), }, })) - vi.doMock("./parsed-result-assets", () => ({ - storeParsedResultAssets, - })) return await import("./reconcile") } @@ -103,7 +92,7 @@ describe("reconcileSourcesForWorkspace", () => { expect(result).toEqual([ready]) }) - it("stores parsed result assets before marking a completed source ready", async () => { + it("does not store parsed result assets before marking a completed source ready", async () => { const parsing = makeSource({ id: "source_1", knowhereJobId: "job_1" }) const ready = makeSource({ id: "source_1", @@ -125,12 +114,6 @@ describe("reconcileSourcesForWorkspace", () => { }) const storeParsedResultAssets = vi.fn(async () => { calls.push("store") - return { - resultBlobUrl: "https://blob.example/result.zip", - assetUrlsByFilePath: { - "images/image-1.jpg": "https://blob.example/image-1.jpg", - }, - } }) const saveSourceParseResult = vi.fn(async () => { calls.push("save") @@ -146,28 +129,11 @@ describe("reconcileSourcesForWorkspace", () => { }, } as unknown as import("@ontos-ai/knowhere-sdk").default - await reconcileSourcesForWorkspace(workspace, mockClient, { - storeParsedResultAssets, - saveSourceParseResult, - }) + await reconcileSourcesForWorkspace(workspace, mockClient) - expect(storeParsedResultAssets).toHaveBeenCalledWith({ - workspaceId: workspace.id, - sourceId: "source_1", - job, - client: mockClient, - }) - expect(saveSourceParseResult).toHaveBeenCalledWith( - workspace.id, - "source_1", - { - resultBlobUrl: "https://blob.example/result.zip", - assetUrlsByFilePath: { - "images/image-1.jpg": "https://blob.example/image-1.jpg", - }, - }, - ) - expect(calls).toEqual(["store", "save", "ready"]) + expect(storeParsedResultAssets).not.toHaveBeenCalled() + expect(saveSourceParseResult).not.toHaveBeenCalled() + expect(calls).toEqual(["ready"]) }) it("keeps original public Blob uploads after completed URL parsing jobs", async () => { @@ -280,7 +246,7 @@ describe("reconcileSourcesForWorkspace", () => { }) describe("applyKnowhereJobToSource", () => { - it("stores completed parsed results before readying the source and then cleans staged blobs", async () => { + it("readies completed sources by document id and then cleans staged blobs", async () => { const parsing = makeSource({ id: "source_1", knowhereJobId: "job_1", @@ -298,21 +264,7 @@ describe("applyKnowhereJobToSource", () => { isTerminal: true, } const calls: string[] = [] - const parsedAssets = { - resultBlobUrl: "https://blob.example/result.zip", - assetUrlsByFilePath: { - "images/image-1.jpg": "https://blob.example/image-1.jpg", - }, - } - const client = { - jobs: { - load: vi.fn(), - }, - } const repository = { - saveSourceParseResult: vi.fn(async () => { - calls.push("save") - }), markSourceReady: vi.fn(async () => { calls.push("ready") }), @@ -323,12 +275,6 @@ describe("applyKnowhereJobToSource", () => { calls.push("clear-staged") }), } - const parsedResultStore = { - storeParsedResultAssets: vi.fn(async () => { - calls.push("store-assets") - return parsedAssets - }), - } const blobStore = { deleteStagedSourceBlob: vi.fn(async () => { calls.push("delete-staged") @@ -339,23 +285,10 @@ describe("applyKnowhereJobToSource", () => { workspaceId: workspace.id, source: parsing, job, - client, repository, - parsedResultStore, blobStore, }) - expect(parsedResultStore.storeParsedResultAssets).toHaveBeenCalledWith({ - workspaceId: workspace.id, - sourceId: "source_1", - job, - client, - }) - expect(repository.saveSourceParseResult).toHaveBeenCalledWith( - workspace.id, - "source_1", - parsedAssets, - ) expect(repository.markSourceReady).toHaveBeenCalledWith( workspace.id, "source_1", @@ -368,12 +301,6 @@ describe("applyKnowhereJobToSource", () => { workspace.id, "source_1", ) - expect(calls).toEqual([ - "store-assets", - "save", - "ready", - "delete-staged", - "clear-staged", - ]) + expect(calls).toEqual(["ready", "delete-staged", "clear-staged"]) }) }) diff --git a/src/domains/sources/reconcile.ts b/src/domains/sources/reconcile.ts index 023eacd..ebcd1a2 100644 --- a/src/domains/sources/reconcile.ts +++ b/src/domains/sources/reconcile.ts @@ -7,23 +7,10 @@ import type { JobResult } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" import { logger } from "@/lib/logger" -import { - storeParsedResultAssets, - type StoreParsedResultAssetsInput, - type StoredParsedResultAssets, -} from "./parsed-result-assets" import { applyKnowhereJobToSource } from "./lifecycle" import { sourceWorkflowRuntime } from "./workflow-runtime" type SourceReconcileDependencies = { - storeParsedResultAssets?: ( - input: Omit, - ) => Promise - saveSourceParseResult?: ( - workspaceId: string, - sourceId: string, - input: StoredParsedResultAssets, - ) => Promise deleteStagedSourceBlob?: (pathname: string) => Promise clearSourceStagedBlob?: ( workspaceId: string, @@ -128,19 +115,12 @@ async function updateSourceFromJob( workspaceId, source, job, - client, repository: { - saveSourceParseResult: - deps.saveSourceParseResult ?? sourceWorkflowRuntime.saveParseResult, markSourceReady: sourceWorkflowRuntime.markReady, markSourceFailed: sourceWorkflowRuntime.markFailed, clearSourceStagedBlob: deps.clearSourceStagedBlob ?? sourceWorkflowRuntime.clearStagedBlob, }, - parsedResultStore: { - storeParsedResultAssets: - deps.storeParsedResultAssets ?? storeParsedResultAssets, - }, blobStore: { deleteStagedSourceBlob: deps.deleteStagedSourceBlob ?? del, }, diff --git a/src/domains/sources/remote-library.ts b/src/domains/sources/remote-library.ts index 4cc3934..87d3fd8 100644 --- a/src/domains/sources/remote-library.ts +++ b/src/domains/sources/remote-library.ts @@ -1,6 +1,7 @@ import { Effect } from "effect" import type { Source } from "@/infrastructure/db/schema" +import type { SourceView } from "./types" import type { SourceStatus } from "./types" import { getCompatibleNamespaces, sharedLibraryNamespace } from "./namespace" @@ -8,6 +9,7 @@ type RemoteDocument = { readonly documentId: string readonly namespace: string readonly status: SourceStatus + readonly revisionKey?: string readonly title?: string readonly mimeType?: string readonly sizeBytes?: number @@ -19,19 +21,38 @@ type RemoteDocumentCandidate = RemoteDocument | { readonly documentId?: string | null readonly namespace?: string | null readonly status?: string | null + readonly currentJobResultId?: string | null + readonly jobResultId?: string | null + readonly jobId?: string | null readonly sourceFileName?: string | null readonly documentMetadata?: Record } type RemoteDocumentListResponse = { readonly documents: readonly RemoteDocumentCandidate[] + readonly pagination?: RemoteDocumentListPagination +} + +type RemoteDocumentListPagination = { + readonly page?: number + readonly pageSize?: number + readonly total?: number + readonly totalPages?: number + readonly page_size?: number + readonly total_pages?: number +} + +type RemoteDocumentListParams = { + readonly namespace?: string + readonly page?: number + readonly pageSize?: number } type RemoteDocumentClient = { readonly documents?: { - readonly list?: (params?: { - readonly namespace?: string - }) => Promise + readonly list?: ( + params?: RemoteDocumentListParams, + ) => Promise } } @@ -56,6 +77,12 @@ type RemoteDocumentRaw = { readonly documentId?: unknown readonly namespace?: unknown readonly status?: unknown + readonly current_job_result_id?: unknown + readonly currentJobResultId?: unknown + readonly job_result_id?: unknown + readonly jobResultId?: unknown + readonly job_id?: unknown + readonly jobId?: unknown readonly source_file_name?: unknown readonly sourceFileName?: unknown readonly document_metadata?: unknown @@ -64,6 +91,12 @@ type RemoteDocumentRaw = { export type RemoteLibraryDocument = RemoteDocument +const remoteSourceIdPrefix = "knowhere-doc" +const remoteDocumentPageSize = 200 +const emptyRemoteDocumentListResponse: RemoteDocumentListResponse = { + documents: [], +} + export function listRemoteLibraryDocuments( input: RemoteLibraryProjectionInput, ): Effect.Effect { @@ -73,33 +106,95 @@ export function listRemoteLibraryDocuments( for (const namespace of getCompatibleNamespaces(input.workspace)) { if (!input.client.documents?.list) continue - const response = yield* Effect.tryPromise(() => - input.client.documents?.list?.({ - namespace, - }) ?? - Promise.resolve({ documents: [] }), - ).pipe( - Effect.catchAll(() => - Effect.succeed({ - documents: [], - } satisfies RemoteDocumentListResponse), - ), - ) - - for (const rawDocument of response.documents ?? []) { - const document = normalizeRemoteDocument(rawDocument) - if (!document) continue - if (seenDocumentIds.has(document.documentId)) continue - - seenDocumentIds.add(document.documentId) - documents.push(document) - } + let page = 1 + let totalPages = 1 + + do { + const response = yield* Effect.tryPromise(() => + input.client.documents?.list?.({ + namespace, + page, + pageSize: remoteDocumentPageSize, + }) ?? + Promise.resolve(emptyRemoteDocumentListResponse), + ).pipe( + Effect.catchAll(() => + Effect.succeed(emptyRemoteDocumentListResponse), + ), + ) + + for (const rawDocument of response.documents ?? []) { + const document = normalizeRemoteDocument(rawDocument) + if (!document) continue + if (seenDocumentIds.has(document.documentId)) continue + + seenDocumentIds.add(document.documentId) + documents.push(document) + } + + totalPages = getRemoteDocumentTotalPages(response.pagination) + page += 1 + } while (page <= totalPages) } return documents }) } +export function listRemoteLibrarySourceViews( + input: RemoteLibraryProjectionInput, +): Effect.Effect { + return Effect.gen(function* () { + const localDocumentIds = new Set( + input.localSources.flatMap((source): string[] => + source.knowhereDocumentId ? [source.knowhereDocumentId] : [], + ), + ) + const remoteDocuments = (yield* listRemoteLibraryDocuments(input)).filter( + (document) => + !localDocumentIds.has(document.documentId) && + !matchesActiveNotebookParsingSource(document, input.localSources), + ) + + return remoteDocuments.map(toRemoteSourceView) + }) +} + +export function decodeRemoteSourceId(sourceId: string): { + readonly namespace: string + readonly documentId: string +} | null { + const parts = sourceId.split(":") + if (parts.length !== 3 || parts[0] !== remoteSourceIdPrefix) return null + + const namespace = decodeRemoteSourceIdSegment(parts[1] ?? "") + const documentId = decodeRemoteSourceIdSegment(parts[2] ?? "") + if (!namespace || !documentId) return null + + return { namespace, documentId } +} + +export function findRemoteLibraryDocumentBySourceId(input: { + readonly sourceId: string + readonly workspace: RemoteLibraryWorkspace + readonly client: RemoteDocumentClient + readonly localSources: readonly Source[] +}): Effect.Effect { + return Effect.gen(function* () { + const decoded = decodeRemoteSourceId(input.sourceId) + if (!decoded) return null + + const documents = yield* listRemoteLibraryDocuments(input) + return ( + documents.find( + (document) => + document.documentId === decoded.documentId && + document.namespace === decoded.namespace, + ) ?? null + ) + }) +} + export function localizeRemoteLibrarySources( input: RemoteLibraryLocalizationInput, ): Effect.Effect { @@ -143,6 +238,14 @@ function normalizeRemoteDocument( const status = getString(raw.status) ?? "ready" if (status === "archived") return null + const revisionKey = getString( + raw.jobId ?? + raw.job_id ?? + raw.currentJobResultId ?? + raw.current_job_result_id ?? + raw.jobResultId ?? + raw.job_result_id, + ) const documentMetadata = getRecord( raw.documentMetadata ?? raw.document_metadata, ) @@ -159,6 +262,7 @@ function normalizeRemoteDocument( documentId, namespace, status: getRemoteSourceStatus(status), + ...(revisionKey ? { revisionKey } : {}), ...(title ? { title } : {}), ...(mimeType ? { mimeType } : {}), ...(sizeBytes !== undefined ? { sizeBytes } : {}), @@ -167,6 +271,42 @@ function normalizeRemoteDocument( } } +function toRemoteSourceView(document: RemoteDocument): SourceView { + return { + id: encodeRemoteSourceId(document), + kind: "remote", + namespace: document.namespace, + title: document.title ?? document.documentId, + mimeType: document.mimeType ?? "application/octet-stream", + status: document.status, + documentId: document.documentId, + excludedFromQuery: true, + } +} + +function encodeRemoteSourceId( + document: Pick, +): string { + return [ + remoteSourceIdPrefix, + encodeRemoteSourceIdSegment(document.namespace), + encodeRemoteSourceIdSegment(document.documentId), + ].join(":") +} + +function encodeRemoteSourceIdSegment(value: string): string { + return encodeURIComponent(value) +} + +function decodeRemoteSourceIdSegment(value: string): string | null { + try { + const decoded = decodeURIComponent(value) + return decoded.length > 0 ? decoded : null + } catch { + return null + } +} + function getRemoteSourceStatus(status: string): SourceStatus { if (isActiveDocumentStatus(status)) return "ready" if (status === "failed") return "failed" @@ -272,6 +412,17 @@ function getNumber(value: unknown): number | undefined { return typeof value === "number" && Number.isFinite(value) ? value : undefined } +function getRemoteDocumentTotalPages( + pagination: RemoteDocumentListPagination | undefined, +): number { + const totalPages = + getNumber(pagination?.totalPages) ?? + getNumber(pagination?.total_pages) + return totalPages !== undefined && totalPages > 0 + ? Math.floor(totalPages) + : 1 +} + function getRecord(value: unknown): Record { if (!value || typeof value !== "object" || Array.isArray(value)) return {} return value as Record diff --git a/src/domains/sources/repository.ts b/src/domains/sources/repository.ts index 52f34cb..c7c4b6f 100644 --- a/src/domains/sources/repository.ts +++ b/src/domains/sources/repository.ts @@ -14,6 +14,7 @@ type SourceRepository = { readonly upsertMaterializedDemoSourceEffect: typeof demoSourceRepository.upsertMaterializedDemoSourceEffect readonly markParsingEffect: typeof sourceRowRepository.markParsingEffect readonly markReadyEffect: typeof sourceRowRepository.markReadyEffect + readonly updateRevisionKeyEffect: typeof sourceRowRepository.updateRevisionKeyEffect readonly markFailedEffect: typeof sourceRowRepository.markFailedEffect readonly clearStagedBlobEffect: typeof sourceRowRepository.clearStagedBlobEffect readonly softDeleteEffect: typeof sourceRowRepository.softDeleteEffect @@ -35,6 +36,7 @@ export const sourceRepository: SourceRepository = { demoSourceRepository.upsertMaterializedDemoSourceEffect, markParsingEffect: sourceRowRepository.markParsingEffect, markReadyEffect: sourceRowRepository.markReadyEffect, + updateRevisionKeyEffect: sourceRowRepository.updateRevisionKeyEffect, markFailedEffect: sourceRowRepository.markFailedEffect, clearStagedBlobEffect: sourceRowRepository.clearStagedBlobEffect, softDeleteEffect: sourceRowRepository.softDeleteEffect, diff --git a/src/domains/sources/route-chunks.ts b/src/domains/sources/route-chunks.ts index 959bde4..733ccc1 100644 --- a/src/domains/sources/route-chunks.ts +++ b/src/domains/sources/route-chunks.ts @@ -4,6 +4,10 @@ import { demoView } from "@/domains/demo/view" import type { DemoChunkPage } from "@/integrations/knowhere-demo" import { logger } from "@/lib/logger" import { routeResult } from "@/lib/route-result" +import { + decodeRemoteSourceId, + findRemoteLibraryDocumentBySourceId, +} from "./remote-library" import { getClientForWorkspace } from "./route-dependencies" import { sourceRowRepository } from "./source-row-repository" import type { @@ -49,7 +53,10 @@ const loadSourceChunksEffect = ( Effect.gen(function* () { if (!sourceRowRepository.isWorkspaceSourceId(input.sourceId)) { const demoResult = yield* loadDemoChunkPageEffect(input, deps) - return demoResult ?? sourceNotFound() + if (demoResult) return demoResult + + const remoteResult = yield* loadRemoteChunkPageEffect(input, deps) + return remoteResult ?? sourceNotFound() } const user = yield* Effect.tryPromise(() => deps.getCurrentUser()) @@ -81,13 +88,88 @@ const loadSourceChunksEffect = ( const client = yield* Effect.tryPromise(() => getClientForWorkspace(workspace.id, input.cookieHeader, deps), ) + if (input.shouldLoadAll) { + const chunks = yield* deps.loadChunksForSource(source, client, { + workspaceId: workspace.id, + onRevisionKey: async (revisionKey) => { + await deps.sourceService.updateSourceRevisionKey( + workspace.id, + source.id, + revisionKey, + ) + }, + }) + return routeResult.ok({ chunks }) + } + const assetUrlsByFilePath = yield* Effect.tryPromise(() => deps.sourceService.getParseAssetUrls(workspace.id, source.id), ) + const chunkPage = yield* deps.loadChunkPageForSource( + source, + client, + input.pageParams, + { + assetUrlsByFilePath, + workspaceId: workspace.id, + onRevisionKey: async (revisionKey) => { + await deps.sourceService.updateSourceRevisionKey( + workspace.id, + source.id, + revisionKey, + ) + }, + }, + ) + return routeResult.ok(chunkPage) + }) + +const loadRemoteChunkPageEffect = ( + input: LoadSourceChunksInput, + deps: RouteChunksDependencies, +) => + Effect.gen(function* () { + if (!decodeRemoteSourceId(input.sourceId)) return null + + const user = yield* Effect.tryPromise(() => deps.getCurrentUser()) + if (!user) return null + + const workspace = yield* Effect.tryPromise(() => + deps.ensureWorkspace(user.id), + ) + const client = yield* Effect.tryPromise(() => + getClientForWorkspace(workspace.id, input.cookieHeader, deps), + ) + const remoteDocument = yield* findRemoteLibraryDocumentBySourceId({ + sourceId: input.sourceId, + workspace, + client, + localSources: [], + }) + if (!remoteDocument) return null + + const source = yield* Effect.tryPromise(() => + deps.sourceService.localizeRemoteDocument(workspace.id, { + documentId: remoteDocument.documentId, + namespace: remoteDocument.namespace, + status: remoteDocument.status, + title: remoteDocument.title, + mimeType: remoteDocument.mimeType, + sizeBytes: remoteDocument.sizeBytes, + revisionKey: remoteDocument.revisionKey ?? null, + }), + ) if (input.shouldLoadAll) { const chunks = yield* deps.loadChunksForSource(source, client, { - assetUrlsByFilePath, + workspaceId: workspace.id, + onRevisionKey: async (revisionKey) => { + await deps.sourceService.updateSourceRevisionKey( + workspace.id, + source.id, + revisionKey, + ) + }, }) return routeResult.ok({ chunks }) } @@ -96,7 +178,16 @@ const loadSourceChunksEffect = ( source, client, input.pageParams, - { assetUrlsByFilePath }, + { + workspaceId: workspace.id, + onRevisionKey: async (revisionKey) => { + await deps.sourceService.updateSourceRevisionKey( + workspace.id, + source.id, + revisionKey, + ) + }, + }, ) return routeResult.ok(chunkPage) }) diff --git a/src/domains/sources/route-dependencies.ts b/src/domains/sources/route-dependencies.ts index 724411e..d056e1c 100644 --- a/src/domains/sources/route-dependencies.ts +++ b/src/domains/sources/route-dependencies.ts @@ -5,7 +5,7 @@ import { del } from "@vercel/blob" import { loadChunkPageForSource, loadChunksForSource, -} from "@/domains/chunks" +} from "@/domains/chunks/server" import { ensureApiKeyForWorkspace } from "@/integrations/dashboard/api-key-service" import { makeKnowhereClient as makeDefaultKnowhereClient } from "@/integrations/knowhere" import { knowhereDemoApi } from "@/integrations/knowhere-demo" @@ -49,6 +49,7 @@ const defaultDependencies: SourceRouteServiceDependencies = { hideDemoSource: defaultSourceService.hideDemoSource, listHiddenDemoSourceIds: defaultSourceService.listHiddenDemoSourceIds, localizeRemoteDocument: defaultSourceService.localizeRemoteDocument, + updateSourceRevisionKey: defaultSourceService.updateSourceRevisionKey, softDelete: defaultSourceService.softDelete, upsertMaterializedDemoSource: defaultSourceService.upsertMaterializedDemoSource, diff --git a/src/domains/sources/route-listing.ts b/src/domains/sources/route-listing.ts index 0052500..9476285 100644 --- a/src/domains/sources/route-listing.ts +++ b/src/domains/sources/route-listing.ts @@ -13,7 +13,7 @@ import { toSourceView } from "./view" import { startBackgroundReconciliation as defaultStartBackgroundReconciliation, } from "./background-reconcile" -import { localizeRemoteLibrarySources } from "./remote-library" +import { listRemoteLibrarySourceViews } from "./remote-library" import type { Source } from "@/infrastructure/db/schema" import type { JsonRouteResult, @@ -90,12 +90,11 @@ const listSourcesEffect = ( const client = deps.makeKnowhereClient(apiKey) const sources = listedSources const demoSourceResolution = resolveWorkspaceDemoSources(sources, catalog) - const workspaceSources = yield* localizeRemoteLibrarySources({ + const workspaceSources = demoSourceResolution.workspaceSources + const remoteSourceViews = yield* listRemoteLibrarySourceViews({ workspace, client, localSources: demoSourceResolution.workspaceSources, - localizeDocument: (document) => - deps.sourceService.localizeRemoteDocument(workspace.id, document), }) const sourcesNeedingKnowhereChunkCount = getWorkspaceSourcesNeedingKnowhereChunkCount(workspaceSources) @@ -140,6 +139,7 @@ const listSourcesEffect = ( sourceOptions.get(source.id), ), ), + ...remoteSourceViews, ], }) }) diff --git a/src/domains/sources/route-service.test.ts b/src/domains/sources/route-service.test.ts index 9901d8f..34b8855 100644 --- a/src/domains/sources/route-service.test.ts +++ b/src/domains/sources/route-service.test.ts @@ -119,7 +119,7 @@ describe("source route service", () => { expect(listHiddenDemoSourceIds).toHaveBeenCalledWith(workspace.id); }); - it("localizes shared default and legacy namespace documents into workspace sources", async () => { + it("lists shared default and legacy namespace documents as lightweight remote sources", async () => { const localReadySource: Source = { ...source, id: "source_ready", @@ -140,6 +140,16 @@ describe("source route service", () => { mimeType: "application/pdf", }, }, + ], + pagination: { + page: 1, + page_size: 200, + total: 2, + total_pages: 2, + }, + }) + .mockResolvedValueOnce({ + documents: [ { documentId: "doc_local", namespace: "default", @@ -147,6 +157,12 @@ describe("source route service", () => { sourceFileName: "local-duplicate.pdf", }, ], + pagination: { + page: 2, + pageSize: 200, + total: 2, + totalPages: 2, + }, }) .mockResolvedValueOnce({ documents: [ @@ -157,6 +173,12 @@ describe("source route service", () => { sourceFileName: "legacy.pdf", }, ], + pagination: { + page: 1, + pageSize: 200, + total: 1, + totalPages: 1, + }, }); const knowhereClient = { documents: { @@ -178,37 +200,7 @@ describe("source route service", () => { upload: vi.fn(), }, }; - const localizedDefaultSource: Source = { - ...source, - id: "00000000-0000-0000-0000-000000000101", - title: "cli.pdf", - mimeType: "application/pdf", - status: "ready", - knowhereJobId: null, - knowhereDocumentId: "doc_default", - }; - const refreshedLocalSource: Source = { - ...localReadySource, - title: "local-duplicate.pdf", - mimeType: "application/octet-stream", - status: "ready", - knowhereJobId: null, - knowhereDocumentId: "doc_local", - }; - const localizedLegacySource: Source = { - ...source, - id: "00000000-0000-0000-0000-000000000102", - title: "legacy.pdf", - mimeType: "application/octet-stream", - status: "ready", - knowhereJobId: null, - knowhereDocumentId: "doc_legacy", - }; - const localizeRemoteDocument = vi - .fn() - .mockResolvedValueOnce(localizedDefaultSource) - .mockResolvedValueOnce(refreshedLocalSource) - .mockResolvedValueOnce(localizedLegacySource); + const localizeRemoteDocument = vi.fn(); const listing = createRouteListing({ demoApi: { fetchCatalog: vi.fn(async () => emptyDemoCatalog), @@ -234,72 +226,47 @@ describe("source route service", () => { expect(listDocuments).toHaveBeenNthCalledWith(1, { namespace: "default", + page: 1, + pageSize: 200, }); expect(listDocuments).toHaveBeenNthCalledWith(2, { + namespace: "default", + page: 2, + pageSize: 200, + }); + expect(listDocuments).toHaveBeenNthCalledWith(3, { namespace: workspace.namespace, + page: 1, + pageSize: 200, }); - expect(localizeRemoteDocument).toHaveBeenNthCalledWith( - 1, - workspace.id, - { - documentId: "doc_default", - namespace: "default", - status: "ready", - title: "cli.pdf", - mimeType: "application/pdf", - sourceFileName: "cli.pdf", - documentMetadata: { - mimeType: "application/pdf", - }, - }, - ); - expect(localizeRemoteDocument).toHaveBeenNthCalledWith( - 2, - workspace.id, - { - documentId: "doc_local", - namespace: "default", - status: "ready", - title: "local-duplicate.pdf", - sourceFileName: "local-duplicate.pdf", - documentMetadata: {}, - }, - ); - expect(localizeRemoteDocument).toHaveBeenNthCalledWith( - 3, - workspace.id, - { - documentId: "doc_legacy", - namespace: workspace.namespace, - status: "ready", - title: "legacy.pdf", - sourceFileName: "legacy.pdf", - documentMetadata: {}, - }, - ); + expect(localizeRemoteDocument).not.toHaveBeenCalled(); expect(result.body.sources).toEqual([ expect.objectContaining({ id: "source_ready", documentId: "doc_local", - title: "local-duplicate.pdf", + title: "notes.pdf", status: "ready", }), - expect.objectContaining({ - id: "00000000-0000-0000-0000-000000000101", - kind: "workspace", + { + id: "knowhere-doc:default:doc_default", + kind: "remote", + namespace: "default", title: "cli.pdf", mimeType: "application/pdf", status: "ready", documentId: "doc_default", - }), - expect.objectContaining({ - id: "00000000-0000-0000-0000-000000000102", - kind: "workspace", + excludedFromQuery: true, + }, + { + id: "knowhere-doc:notebook-workspace_1:doc_legacy", + kind: "remote", + namespace: workspace.namespace, title: "legacy.pdf", mimeType: "application/octet-stream", status: "ready", documentId: "doc_legacy", - }), + excludedFromQuery: true, + }, ]); }); diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index c56335b..998e792 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -2,9 +2,11 @@ import type { ChunkKnowhereClient, ChunkPage, ChunkPageParams, +} from "@/domains/chunks" +import type { loadChunkPageForSource, loadChunksForSource, -} from "@/domains/chunks" +} from "@/domains/chunks/server" import type { ParsedChunkView } from "@/domains/chunks/types" import type { SourceStatus, SourceView } from "@/domains/sources/types" import type { AuthUser } from "@/infrastructure/auth" @@ -23,14 +25,32 @@ type SourceRouteKnowhereClient = UploadKnowhereClient & readonly documents: ChunkKnowhereClient["documents"] & { list?(params?: { readonly namespace?: string + readonly page?: number + readonly pageSize?: number }): Promise<{ readonly documents: readonly { readonly documentId: string readonly namespace: string readonly status: string + readonly currentJobResultId?: string | null readonly sourceFileName?: string | null readonly documentMetadata?: Record }[] + readonly pagination?: { + readonly page?: number + readonly pageSize?: number + readonly total?: number + readonly totalPages?: number + readonly page_size?: number + readonly total_pages?: number + } + }> + get?(documentId: string): Promise<{ + readonly documentId: string + readonly namespace: string + readonly status: string + readonly currentJobResultId?: string | null + readonly sourceFileName?: string | null }> archive(documentId: string): Promise } @@ -151,12 +171,19 @@ type SourceWorkflowService = { workspaceId: string, input: { readonly documentId: string + readonly namespace?: string readonly title?: string readonly mimeType?: string readonly sizeBytes?: number readonly status: SourceStatus + readonly revisionKey?: string | null }, ) => Promise + readonly updateSourceRevisionKey: ( + workspaceId: string, + sourceId: string, + revisionKey: string, + ) => Promise readonly upsertMaterializedDemoSource: ( workspaceId: string, input: { diff --git a/src/domains/sources/service.ts b/src/domains/sources/service.ts index b515300..b5d39b3 100644 --- a/src/domains/sources/service.ts +++ b/src/domains/sources/service.ts @@ -25,6 +25,11 @@ type SourceService = { workspaceId: string, input: Parameters[1], ) => Promise + readonly updateSourceRevisionKey: ( + workspaceId: string, + sourceId: string, + revisionKey: string, + ) => Promise readonly listHiddenDemoSourceIds: (workspaceId: string) => Promise readonly hideDemoSource: ( workspaceId: string, @@ -84,6 +89,7 @@ export const sourceService: SourceService = { listHiddenDemoSourceIds: sourceWorkflowRuntime.listHiddenDemoSourceIds, listForWorkspace: sourceWorkflowRuntime.listForWorkspace, localizeRemoteDocument: sourceWorkflowRuntime.localizeRemoteDocument, + updateSourceRevisionKey: sourceWorkflowRuntime.updateRevisionKey, softDelete: sourceWorkflowRuntime.softDelete, upsertMaterializedDemoSource: sourceWorkflowRuntime.upsertMaterializedDemoSource, diff --git a/src/domains/sources/source-reconcile-route-workflow.test.ts b/src/domains/sources/source-reconcile-route-workflow.test.ts index 2307260..0841c54 100644 --- a/src/domains/sources/source-reconcile-route-workflow.test.ts +++ b/src/domains/sources/source-reconcile-route-workflow.test.ts @@ -1,32 +1,13 @@ import { afterEach, describe, expect, it, vi } from "vitest" -import type { Source } from "@/infrastructure/db/schema" - const mocks = vi.hoisted(() => ({ - completeResultZipMultipartUpload: vi.fn(), - createResultZipMultipartUploadPlan: vi.fn(), - getResultZipPartRange: vi.fn(), loggerError: vi.fn(), loggerInfo: vi.fn(), loggerWarn: vi.fn(), makeKnowhereClient: vi.fn(), + markFailed: vi.fn(), markSourceReadyAfterReconciliation: vi.fn(), pollSourceReconciliation: vi.fn(), - prepareParsedResultAssetBatch: vi.fn(), - uploadResultZipPart: vi.fn(), - findInWorkspace: vi.fn(), - getParseResultProgress: vi.fn(), - markFailed: vi.fn(), - markParsing: vi.fn(), - mergeParseAssetUrls: vi.fn(), -})) - -vi.mock("@/domains/sources/parsed-result-assets", () => ({ - completeResultZipMultipartUpload: mocks.completeResultZipMultipartUpload, - createResultZipMultipartUploadPlan: mocks.createResultZipMultipartUploadPlan, - getResultZipPartRange: mocks.getResultZipPartRange, - prepareParsedResultAssetBatch: mocks.prepareParsedResultAssetBatch, - uploadResultZipPart: mocks.uploadResultZipPart, })) vi.mock("@/domains/sources/source-reconcile-workflow", () => ({ @@ -36,11 +17,7 @@ vi.mock("@/domains/sources/source-reconcile-workflow", () => ({ vi.mock("@/domains/sources/workflow-runtime", () => ({ sourceWorkflowRuntime: { - findInWorkspace: mocks.findInWorkspace, - getParseResultProgress: mocks.getParseResultProgress, markFailed: mocks.markFailed, - markParsing: mocks.markParsing, - mergeParseAssetUrls: mocks.mergeParseAssetUrls, }, })) @@ -63,23 +40,24 @@ describe("sourceReconcileRouteWorkflow", () => { vi.clearAllMocks() }) - it("defaults legacy payloads to the poll-and-mirror phase", () => { + it("normalizes legacy mirror and asset payloads to poll-and-ready", () => { expect( sourceReconcileRouteWorkflow.normalizeReconcilePayload({ workspaceId: "workspace_1", sourceId: "source_1", apiKey: "jwt_1", + phase: "asset-batches", }), ).toEqual({ workspaceId: "workspace_1", sourceId: "source_1", apiKey: "jwt_1", - phase: "poll-and-mirror", + phase: "poll-and-ready", segmentIndex: 0, }) }) - it("mirrors the ZIP and triggers asset continuation without running asset batches inline", async () => { + it("marks the source ready when Knowhere publishes a document id", async () => { const context = createWorkflowContext() const continuations: ContinuationTriggerInput[] = [] const restore = @@ -95,167 +73,9 @@ describe("sourceReconcileRouteWorkflow", () => { jobId: "job_1", documentId: "doc_1", }) - mocks.createResultZipMultipartUploadPlan.mockResolvedValue({ - pathname: "result.zip", - key: "blob-key", - uploadId: "upload-1", - sizeBytes: 10, - partSizeBytes: 10, - partCount: 1, - }) - mocks.getResultZipPartRange.mockReturnValue({ - startByte: 0, - endByte: 9, - }) - mocks.uploadResultZipPart.mockResolvedValue({ - etag: "etag-1", - partNumber: 1, - }) - mocks.completeResultZipMultipartUpload.mockResolvedValue({ - url: "https://blob.example/result.zip", - }) - mocks.mergeParseAssetUrls.mockResolvedValue({ id: "parse_result_1" }) - mocks.markParsing.mockResolvedValue(makeSource()) - - try { - await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ - context, - payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ - workspaceId: "workspace_1", - sourceId: "source_1", - apiKey: "jwt_1", - }), - }) - } finally { - restore() - } - - expect(mocks.mergeParseAssetUrls).toHaveBeenCalledWith( - "workspace_1", - "source_1", - { - resultBlobUrl: "https://blob.example/result.zip", - assetUrlsByFilePath: {}, - }, - ) - expect(mocks.markParsing).toHaveBeenCalledWith( - "workspace_1", - "source_1", - "job_1", - "doc_1", - "parsing", - ) - expect(mocks.prepareParsedResultAssetBatch).not.toHaveBeenCalled() - expect(continuations).toEqual([ - { - url: "https://notebook.example/api/sources/reconcile", - payload: { - workspaceId: "workspace_1", - sourceId: "source_1", - apiKey: "jwt_1", - phase: "asset-batches", - segmentIndex: 0, - }, - workflowRunId: sourceReconcileRouteWorkflow.getAssetWorkflowRunId( - "source_1", - 0, - ), - }, - ]) - }) - - it("resumes asset continuation without remirroring an already saved ZIP", async () => { - vi.setSystemTime(new Date("2026-06-30T03:20:00.000Z")) - const context = createWorkflowContext() - const continuations: ContinuationTriggerInput[] = [] - const restore = - sourceReconcileRouteWorkflow.setContinuationTriggerForTesting( - async (input) => { - continuations.push(input) - }, - ) - mocks.makeKnowhereClient.mockReturnValue({ jobs: {} }) - mocks.pollSourceReconciliation.mockResolvedValue({ - kind: "ready-to-prepare", - jobId: "job_1", - documentId: "doc_1", - }) - mocks.getParseResultProgress.mockResolvedValue({ - resultBlobUrl: "https://blob.example/result.zip", - assetUrlsByFilePath: { - "images/image-1.png": "https://blob.example/images/image-1.png", - }, - }) - mocks.markParsing.mockResolvedValue(makeSource()) - - try { - await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ - context, - payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ - workspaceId: "workspace_1", - sourceId: "source_1", - apiKey: "jwt_1", - }), - }) - } finally { - restore() - vi.useRealTimers() - } - - const resumeSegmentIndex = new Date("2026-06-30T03:20:00.000Z").getTime() - expect(mocks.createResultZipMultipartUploadPlan).not.toHaveBeenCalled() - expect(mocks.uploadResultZipPart).not.toHaveBeenCalled() - expect(mocks.completeResultZipMultipartUpload).not.toHaveBeenCalled() - expect(mocks.mergeParseAssetUrls).not.toHaveBeenCalled() - expect(mocks.markParsing).toHaveBeenCalledWith( - "workspace_1", - "source_1", - "job_1", - "doc_1", - "parsing", - ) - expect(continuations).toEqual([ - { - url: "https://notebook.example/api/sources/reconcile", - payload: { - workspaceId: "workspace_1", - sourceId: "source_1", - apiKey: "jwt_1", - phase: "asset-batches", - segmentIndex: resumeSegmentIndex, - }, - workflowRunId: sourceReconcileRouteWorkflow.getAssetWorkflowRunId( - "source_1", - resumeSegmentIndex, - ), - }, - ]) - }) - - it("reuses the asset resume segment during workflow replay", async () => { - vi.setSystemTime(new Date("2026-06-30T03:20:00.000Z")) - const stepResults = new Map() - const context = createWorkflowContext(stepResults) - const continuations: ContinuationTriggerInput[] = [] - const restore = - sourceReconcileRouteWorkflow.setContinuationTriggerForTesting( - async (input) => { - continuations.push(input) - }, - ) - mocks.makeKnowhereClient.mockReturnValue({ jobs: {} }) - mocks.pollSourceReconciliation.mockResolvedValue({ - kind: "ready-to-prepare", - jobId: "job_1", - documentId: "doc_1", - }) - mocks.getParseResultProgress.mockResolvedValue({ - resultBlobUrl: "https://blob.example/result.zip", - assetUrlsByFilePath: { - "images/image-1.png": "https://blob.example/images/image-1.png", - }, + mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ + status: "ready", }) - mocks.markParsing.mockResolvedValue(makeSource()) try { await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ @@ -266,137 +86,16 @@ describe("sourceReconcileRouteWorkflow", () => { apiKey: "jwt_1", }), }) - vi.setSystemTime(new Date("2026-06-30T03:20:14.000Z")) - await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ - context: createWorkflowContext(stepResults), - payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ - workspaceId: "workspace_1", - sourceId: "source_1", - apiKey: "jwt_1", - }), - }) - } finally { - restore() - vi.useRealTimers() - } - - const resumeSegmentIndex = new Date("2026-06-30T03:20:00.000Z").getTime() - expect(continuations).toHaveLength(1) - expect(continuations[0]?.workflowRunId).toBe( - sourceReconcileRouteWorkflow.getAssetWorkflowRunId( - "source_1", - resumeSegmentIndex, - ), - ) - expect([...stepResults.keys()]).toContain( - `trigger-asset-continuation-${resumeSegmentIndex}`, - ) - expect([...stepResults.keys()]).not.toContain( - `trigger-asset-continuation-${new Date( - "2026-06-30T03:20:14.000Z", - ).getTime()}`, - ) - }) - - it("limits asset continuation work and triggers the next segment when assets remain", async () => { - const context = createWorkflowContext() - const continuations: ContinuationTriggerInput[] = [] - const restore = - sourceReconcileRouteWorkflow.setContinuationTriggerForTesting( - async (input) => { - continuations.push(input) - }, - ) - const client = { jobs: {} } - mocks.makeKnowhereClient.mockReturnValue(client) - mocks.findInWorkspace.mockResolvedValue( - makeSource({ - knowhereJobId: "job_1", - knowhereDocumentId: "doc_1", - }), - ) - mocks.getParseResultProgress.mockResolvedValue({ - resultBlobUrl: "https://blob.example/result.zip", - assetUrlsByFilePath: {}, - }) - mocks.prepareParsedResultAssetBatch.mockResolvedValue({ - uploadedCount: 10, - remainingCount: 737, - hasMore: true, - }) - - try { - await sourceReconcileRouteWorkflow.runAssetBatchWorkflow({ - context, - payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ - workspaceId: "workspace_1", - sourceId: "source_1", - apiKey: "jwt_1", - phase: "asset-batches", - segmentIndex: 2, - }), - }) } finally { restore() } - expect(mocks.prepareParsedResultAssetBatch).toHaveBeenCalledTimes(20) - expect(mocks.markSourceReadyAfterReconciliation).not.toHaveBeenCalled() - expect(continuations).toEqual([ - { - url: "https://notebook.example/api/sources/reconcile", - payload: { - workspaceId: "workspace_1", - sourceId: "source_1", - apiKey: "jwt_1", - phase: "asset-batches", - segmentIndex: 3, - }, - workflowRunId: sourceReconcileRouteWorkflow.getAssetWorkflowRunId( - "source_1", - 3, - ), - }, - ]) - }) - - it("marks the source ready when the asset continuation has no remaining work", async () => { - const context = createWorkflowContext() - mocks.makeKnowhereClient.mockReturnValue({ jobs: {} }) - mocks.findInWorkspace.mockResolvedValue( - makeSource({ - knowhereJobId: "job_1", - knowhereDocumentId: "doc_1", - }), - ) - mocks.getParseResultProgress.mockResolvedValue({ - resultBlobUrl: "https://blob.example/result.zip", - assetUrlsByFilePath: {}, - }) - mocks.prepareParsedResultAssetBatch.mockResolvedValue({ - uploadedCount: 0, - remainingCount: 0, - hasMore: false, - }) - mocks.markSourceReadyAfterReconciliation.mockResolvedValue({ - status: "ready", - }) - - await sourceReconcileRouteWorkflow.runAssetBatchWorkflow({ - context, - payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ - workspaceId: "workspace_1", - sourceId: "source_1", - apiKey: "jwt_1", - phase: "asset-batches", - }), - }) - expect(mocks.markSourceReadyAfterReconciliation).toHaveBeenCalledWith({ workspaceId: "workspace_1", sourceId: "source_1", documentId: "doc_1", }) + expect(continuations).toEqual([]) }) it("triggers a fresh poll run when Knowhere is still running after the segment budget", async () => { @@ -437,7 +136,7 @@ describe("sourceReconcileRouteWorkflow", () => { workspaceId: "workspace_1", sourceId: "source_1", apiKey: "jwt_1", - phase: "poll-and-mirror", + phase: "poll-and-ready", segmentIndex: 5, }, workflowRunId: sourceReconcileRouteWorkflow.getPollWorkflowRunId( @@ -447,6 +146,28 @@ describe("sourceReconcileRouteWorkflow", () => { }, ]) }) + + it("marks a parsing source failed after workflow retry exhaustion", async () => { + mocks.markFailed.mockResolvedValue({ id: "source_1" }) + + await sourceReconcileRouteWorkflow.markSourceFailedAfterWorkflowFailure( + { + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + phase: "asset-batches", + segmentIndex: 2, + }, + "Retry attempts exhausted.", + ) + + expect(mocks.markFailed).toHaveBeenCalledWith( + "workspace_1", + "source_1", + "Source reconciliation workflow failed: Retry attempts exhausted.", + "parsing", + ) + }) }) type ContinuationTriggerInput = Parameters< @@ -471,26 +192,3 @@ function createWorkflowContext(stepResults = new Map()): { url: "https://notebook.example/api/sources/reconcile", } } - -function makeSource(overrides: Partial = {}): Source { - return { - id: "source_1", - workspaceId: "workspace_1", - title: "source.pdf", - mimeType: "application/pdf", - sizeBytes: 1024, - status: "parsing", - failureReason: null, - knowhereJobId: "job_1", - knowhereDocumentId: "doc_1", - stagedBlobPathname: null, - stagedBlobUrl: null, - originalBlobPathname: null, - originalBlobUrl: null, - demoKey: null, - createdAt: new Date("2026-06-30T00:00:00.000Z"), - updatedAt: new Date("2026-06-30T00:00:00.000Z"), - deletedAt: null, - ...overrides, - } -} diff --git a/src/domains/sources/source-reconcile-route-workflow.ts b/src/domains/sources/source-reconcile-route-workflow.ts index 2d94375..a1e732a 100644 --- a/src/domains/sources/source-reconcile-route-workflow.ts +++ b/src/domains/sources/source-reconcile-route-workflow.ts @@ -2,14 +2,6 @@ import "server-only" import { Client, type WorkflowContext } from "@upstash/workflow" -import { - completeResultZipMultipartUpload, - createResultZipMultipartUploadPlan, - getResultZipPartRange, - prepareParsedResultAssetBatch, - uploadResultZipPart, - type MultipartUploadPart, -} from "@/domains/sources/parsed-result-assets" import { markSourceReadyAfterReconciliation, pollSourceReconciliation, @@ -26,13 +18,13 @@ type ReconcilePayload = { readonly segmentIndex?: number } -type ReconcilePhase = "poll-and-mirror" | "asset-batches" +type ReconcilePhase = "poll-and-ready" | "poll-and-mirror" | "asset-batches" type NormalizedReconcilePayload = { readonly workspaceId: string readonly sourceId: string readonly apiKey: string - readonly phase: ReconcilePhase + readonly phase: "poll-and-ready" readonly segmentIndex: number } @@ -48,7 +40,6 @@ type ContinuationTriggerInput = { } const maxPollAttempts = 25 -const maxAssetBatchesPerWorkflowRun = 20 const initialDelaySeconds = 3 const maxDelaySeconds = 30 @@ -82,7 +73,7 @@ async function runPollAndMirrorWorkflow(input: { jobId: poll.jobId, documentId: poll.documentId, } - logger.info("workflow: source parse completed; preparing artifacts", { + logger.info("workflow: source parse completed; readying source", { sourceId, jobId: poll.jobId, attempts: attempt + 1, @@ -113,7 +104,7 @@ async function runPollAndMirrorWorkflow(input: { workspaceId, sourceId, apiKey, - phase: "poll-and-mirror", + phase: "poll-and-ready", segmentIndex: nextSegmentIndex, }, workflowRunId: getPollWorkflowRunId(sourceId, nextSegmentIndex), @@ -127,310 +118,17 @@ async function runPollAndMirrorWorkflow(input: { return } - const existingProgress = await context.run( - "parse-result-check-existing", - async () => - sourceWorkflowRuntime.getParseResultProgress(workspaceId, sourceId), - ) - if (existingProgress) { - const resumeAssetSegmentIndex = await context.run( - "create-asset-resume-segment", - async () => Date.now(), - ) - const savedAssetCount: number = Object.keys( - existingProgress.assetUrlsByFilePath, - ).length - - await context.run("source-save-document-id-for-existing-zip", async () => { - const saved = await sourceWorkflowRuntime.markParsing( - workspaceId, - sourceId, - jobToPrepare.jobId, - jobToPrepare.documentId, - "parsing", - ) - if (!saved) throw new Error("Source disappeared before saving document id.") - }) - logger.info("workflow: result ZIP already mirrored; resuming asset batches", { - sourceId, - jobId: jobToPrepare.jobId, - savedAssetCount, - segmentIndex: resumeAssetSegmentIndex, - }) - - await context.run( - `trigger-asset-continuation-${resumeAssetSegmentIndex}`, - async () => - triggerContinuation({ - url: context.url, - payload: { - workspaceId, - sourceId, - apiKey, - phase: "asset-batches", - segmentIndex: resumeAssetSegmentIndex, - }, - workflowRunId: getAssetWorkflowRunId( - sourceId, - resumeAssetSegmentIndex, - ), - }), - ) - logger.info("workflow: parsed asset continuation triggered", { - sourceId, - jobId: jobToPrepare.jobId, - segmentIndex: resumeAssetSegmentIndex, - }) - return - } - - const uploadPlan = await context.run("mirror-zip-start", async () => - createResultZipMultipartUploadPlan({ - workspaceId, - sourceId, - jobId: jobToPrepare.jobId, - client, - }), - ) - logger.info("workflow: result ZIP multipart upload planned", { - sourceId, - jobId: jobToPrepare.jobId, - sizeBytes: uploadPlan.sizeBytes, - partCount: uploadPlan.partCount, - }) - const parts: MultipartUploadPart[] = [] - - for (let partNumber = 1; partNumber <= uploadPlan.partCount; partNumber++) { - const range = getResultZipPartRange( - partNumber, - uploadPlan.partSizeBytes, - uploadPlan.sizeBytes, - ) - const part = await context.run(`mirror-zip-part-${partNumber}`, async () => - uploadResultZipPart({ - pathname: uploadPlan.pathname, - key: uploadPlan.key, - uploadId: uploadPlan.uploadId, - partNumber, - jobId: jobToPrepare.jobId, - client, - ...range, - }), - ) - parts.push(part) - logger.info("workflow: result ZIP part mirrored", { - sourceId, - jobId: jobToPrepare.jobId, - partNumber, - partCount: uploadPlan.partCount, - startByte: range.startByte, - endByte: range.endByte, - }) - } - - const resultBlobUrl = await context.run("mirror-zip-complete", async () => { - const result = await completeResultZipMultipartUpload({ - pathname: uploadPlan.pathname, - key: uploadPlan.key, - uploadId: uploadPlan.uploadId, - parts, - }) - return result.url - }) - logger.info("workflow: result ZIP multipart upload completed", { - sourceId, - jobId: jobToPrepare.jobId, - partCount: parts.length, - }) - - await context.run("parse-result-save-zip", async () => { - const saved = await sourceWorkflowRuntime.mergeParseAssetUrls( - workspaceId, - sourceId, - { - resultBlobUrl, - assetUrlsByFilePath: {}, - }, - ) - if (!saved) throw new Error("Source disappeared before saving result ZIP.") - }) - logger.info("workflow: result ZIP parse-result row saved", { - sourceId, - jobId: jobToPrepare.jobId, - }) - - await context.run("source-save-document-id", async () => { - const saved = await sourceWorkflowRuntime.markParsing( - workspaceId, - sourceId, - jobToPrepare.jobId, - jobToPrepare.documentId, - "parsing", - ) - if (!saved) throw new Error("Source disappeared before saving document id.") - }) - logger.info("workflow: source document id saved for asset continuation", { - sourceId, - jobId: jobToPrepare.jobId, - }) - - await context.run("trigger-asset-continuation-0", async () => - triggerContinuation({ - url: context.url, - payload: { - workspaceId, - sourceId, - apiKey, - phase: "asset-batches", - segmentIndex: 0, - }, - workflowRunId: getAssetWorkflowRunId(sourceId, 0), - }), - ) - logger.info("workflow: parsed asset continuation triggered", { - sourceId, - jobId: jobToPrepare.jobId, - segmentIndex: 0, - }) -} - -async function runAssetBatchWorkflow(input: { - readonly context: ReconcileWorkflowContext - readonly payload: NormalizedReconcilePayload -}): Promise { - const { context, payload } = input - const { workspaceId, sourceId, apiKey, segmentIndex } = payload - const client = makeKnowhereClient(apiKey) - const source = await context.run("load-source", async () => - sourceWorkflowRuntime.findInWorkspace(workspaceId, sourceId), - ) - if (!source) { - logger.info("workflow: asset continuation resolved missing source", { - sourceId, - segmentIndex, - }) - return - } - if (source.status !== "parsing") { - logger.info("workflow: asset continuation resolved non-parsing source", { - sourceId, - segmentIndex, - status: source.status, - }) - return - } - - const jobId = source.knowhereJobId - const documentId = source.knowhereDocumentId - if (!jobId || !documentId) { - await context.run("asset-continuation-failed-missing-job", async () => - sourceWorkflowRuntime.markFailed( - workspaceId, - sourceId, - "Source artifact preparation is missing Knowhere job metadata.", - "parsing", - ), - ) - logger.error("workflow: asset continuation missing Knowhere metadata", { - sourceId, - segmentIndex, - hasJobId: Boolean(jobId), - hasDocumentId: Boolean(documentId), - }) - return - } - - const progress = await context.run("load-parse-progress", async () => - sourceWorkflowRuntime.getParseResultProgress(workspaceId, sourceId), - ) - if (!progress) { - await context.run("asset-continuation-failed-missing-progress", async () => - sourceWorkflowRuntime.markFailed( - workspaceId, - sourceId, - "Source artifact preparation is missing the mirrored parse result ZIP.", - "parsing", - ), - ) - logger.error("workflow: asset continuation missing parse-result progress", { - sourceId, - segmentIndex, - jobId, - }) - return - } - - let batchIndex = 0 - let lastBatch: { - readonly uploadedCount: number - readonly remainingCount: number - readonly hasMore: boolean - } | null = null - while (batchIndex < maxAssetBatchesPerWorkflowRun) { - const batch = await context.run( - `parse-assets-batch-${segmentIndex}-${batchIndex}`, - async () => - prepareParsedResultAssetBatch({ - workspaceId, - sourceId, - jobId, - documentId, - resultBlobUrl: progress.resultBlobUrl, - client, - repository: sourceWorkflowRuntime, - }), - ) - lastBatch = batch - logger.info("workflow: parsed asset batch prepared", { - sourceId, - jobId, - segmentIndex, - batchIndex, - uploadedCount: batch.uploadedCount, - remainingCount: batch.remainingCount, - hasMore: batch.hasMore, - }) - if (!batch.hasMore) break - batchIndex += 1 - } - - if (lastBatch?.hasMore) { - const nextSegmentIndex = segmentIndex + 1 - await context.run(`trigger-asset-continuation-${nextSegmentIndex}`, async () => - triggerContinuation({ - url: context.url, - payload: { - workspaceId, - sourceId, - apiKey, - phase: "asset-batches", - segmentIndex: nextSegmentIndex, - }, - workflowRunId: getAssetWorkflowRunId(sourceId, nextSegmentIndex), - }), - ) - logger.info("workflow: parsed asset continuation triggered", { - sourceId, - jobId, - segmentIndex: nextSegmentIndex, - remainingCount: lastBatch.remainingCount, - }) - return - } - const ready = await context.run("source-ready", async () => markSourceReadyAfterReconciliation({ workspaceId, sourceId, - documentId, + documentId: jobToPrepare.documentId, }), ) - logger.info("workflow: source artifact preparation finished", { + logger.info("workflow: source parse reconciliation finished", { sourceId, - jobId, + jobId: jobToPrepare.jobId, status: ready.status, - segmentIndex, - assetBatches: batchIndex + 1, }) } @@ -441,7 +139,7 @@ function normalizeReconcilePayload( workspaceId: payload.workspaceId, sourceId: payload.sourceId, apiKey: payload.apiKey, - phase: payload.phase ?? "poll-and-mirror", + phase: "poll-and-ready", segmentIndex: typeof payload.segmentIndex === "number" && Number.isInteger(payload.segmentIndex) && @@ -467,10 +165,6 @@ async function triggerReconcileContinuation( }) } -function getAssetWorkflowRunId(sourceId: string, segmentIndex: number): string { - return `${sourceId}-assets-${segmentIndex}` -} - function getPollWorkflowRunId(sourceId: string, segmentIndex: number): string { return `${sourceId}-poll-${segmentIndex}` } @@ -515,11 +209,9 @@ function getSafeFailureReason(value: string): string { } export const sourceReconcileRouteWorkflow = { - getAssetWorkflowRunId, getPollWorkflowRunId, markSourceFailedAfterWorkflowFailure, normalizeReconcilePayload, - runAssetBatchWorkflow, runPollAndMirrorWorkflow, setContinuationTriggerForTesting, } diff --git a/src/domains/sources/source-reconcile-workflow.test.ts b/src/domains/sources/source-reconcile-workflow.test.ts index 9548ea0..3f0632e 100644 --- a/src/domains/sources/source-reconcile-workflow.test.ts +++ b/src/domains/sources/source-reconcile-workflow.test.ts @@ -138,7 +138,7 @@ describe("pollSourceReconciliation", () => { ) }) - it("marks done jobs without a result URL failed", async () => { + it("accepts done jobs without a result URL when a document id is published", async () => { const repository = createRepository() const jobWithoutResultUrl = makeDoneJob() delete jobWithoutResultUrl.resultUrl @@ -155,18 +155,17 @@ describe("pollSourceReconciliation", () => { repository, }) - expect(result).toEqual({ kind: "resolved", status: "failed" }) - expect(repository.markFailed).toHaveBeenCalledWith( - workspace.id, - "source_1", - "Parsing finished but Knowhere did not publish a result ZIP.", - "parsing", - ) + expect(result).toEqual({ + kind: "ready-to-prepare", + jobId: "job_1", + documentId: "doc_1", + }) + expect(repository.markFailed).not.toHaveBeenCalled() }) }) describe("markSourceReadyAfterReconciliation", () => { - it("marks ready after artifact preparation and cleans staged uploads", async () => { + it("marks ready after Knowhere publishes a document id and cleans staged uploads", async () => { const source = makeSource({ stagedBlobPathname: "source-uploads/upload_1/document.pdf", }) diff --git a/src/domains/sources/source-reconcile-workflow.ts b/src/domains/sources/source-reconcile-workflow.ts index 97e5ebe..0355350 100644 --- a/src/domains/sources/source-reconcile-workflow.ts +++ b/src/domains/sources/source-reconcile-workflow.ts @@ -131,17 +131,6 @@ export async function pollSourceReconciliation({ return { kind: "resolved", status: "failed" } } - if (!job.resultUrl) { - await failSourceAndCleanup({ - workspaceId, - source, - reason: "Parsing finished but Knowhere did not publish a result ZIP.", - repository, - blobStore, - }) - return { kind: "resolved", status: "failed" } - } - return { kind: "ready-to-prepare", jobId, diff --git a/src/domains/sources/source-row-repository.ts b/src/domains/sources/source-row-repository.ts index f501db0..77c56dd 100644 --- a/src/domains/sources/source-row-repository.ts +++ b/src/domains/sources/source-row-repository.ts @@ -37,10 +37,12 @@ type SourceUpdate = Partial< type LocalizeRemoteDocumentInput = { readonly documentId: string + readonly namespace?: string readonly title?: string readonly mimeType?: string readonly sizeBytes?: number readonly status: SourceStatus + readonly revisionKey?: string | null } type SourceRowRepository = { @@ -71,6 +73,11 @@ type SourceRowRepository = { sourceId: string, documentId: string, ) => Effect.Effect + readonly updateRevisionKeyEffect: ( + workspaceId: string, + sourceId: string, + revisionKey: string, + ) => Effect.Effect readonly markFailedEffect: ( workspaceId: string, sourceId: string, @@ -199,6 +206,15 @@ const markReadyEffect: SourceRowRepository["markReadyEffect"] = ( failureReason: null, }, "parsing") +const updateRevisionKeyEffect: SourceRowRepository["updateRevisionKeyEffect"] = ( + workspaceId: string, + sourceId: string, + revisionKey: string, +) => + updateInWorkspaceEffect(workspaceId, sourceId, { + knowhereJobId: revisionKey, + }, "ready") + const markFailedEffect: SourceRowRepository["markFailedEffect"] = ( workspaceId: string, sourceId: string, @@ -335,7 +351,7 @@ async function localizeRemoteDocumentWithDb( status: input.status, failureReason: input.status === "failed" ? "Knowhere document failed." : null, - knowhereJobId: null, + knowhereJobId: input.revisionKey ?? null, knowhereDocumentId: input.documentId, stagedBlobPathname: null, stagedBlobUrl: null, @@ -397,6 +413,7 @@ export const sourceRowRepository: SourceRowRepository = { localizeRemoteDocumentEffect, markParsingEffect, markReadyEffect, + updateRevisionKeyEffect, markFailedEffect, clearStagedBlobEffect, softDeleteEffect, diff --git a/src/domains/sources/types.ts b/src/domains/sources/types.ts index 94f97b5..e6246be 100644 --- a/src/domains/sources/types.ts +++ b/src/domains/sources/types.ts @@ -8,7 +8,7 @@ export type SourceOriginalFileView = { readonly pdfPreviewMode?: "browser" } -export type SourceKind = "workspace" | "demo" +export type SourceKind = "workspace" | "demo" | "remote" export type SourceOfficialLibraryView = { readonly librarySourceId: string diff --git a/src/domains/sources/workflow-runtime.ts b/src/domains/sources/workflow-runtime.ts index a9bf32b..bbec158 100644 --- a/src/domains/sources/workflow-runtime.ts +++ b/src/domains/sources/workflow-runtime.ts @@ -79,6 +79,11 @@ type SourceWorkflowRuntime = UploadRepositoryRuntime & { sourceId: string, documentId: string, ) => Promise + readonly updateRevisionKey: ( + workspaceId: string, + sourceId: string, + revisionKey: string, + ) => Promise readonly saveParseResult: ( workspaceId: string, sourceId: string, @@ -166,6 +171,19 @@ const markReady: SourceWorkflowRuntime["markReady"] = ( sourceRepository.markReadyEffect(workspaceId, sourceId, documentId), ) +const updateRevisionKey: SourceWorkflowRuntime["updateRevisionKey"] = ( + workspaceId: string, + sourceId: string, + revisionKey: string, +) => + databaseRuntime.runPromise( + sourceRepository.updateRevisionKeyEffect( + workspaceId, + sourceId, + revisionKey, + ), + ) + const markFailed: SourceWorkflowRuntime["markFailed"] = ( workspaceId: string, sourceId: string, @@ -276,6 +294,7 @@ export const sourceWorkflowRuntime: SourceWorkflowRuntime = { markFailed, markParsing, markReady, + updateRevisionKey, mergeParseAssetUrls, saveParseResult, softDelete, diff --git a/src/domains/workspace/initial-state.test.ts b/src/domains/workspace/initial-state.test.ts index 544b74e..9ea24e3 100644 --- a/src/domains/workspace/initial-state.test.ts +++ b/src/domains/workspace/initial-state.test.ts @@ -430,7 +430,6 @@ describe("loadWorkspaceShellInitialState", () => { }, } as unknown as InitialStateClient const reconcileSourcesForWorkspace = vi.fn(async () => [parsingSource]) - const localizeRemoteDocument = vi.fn(async () => parsingSource) const startBackgroundReconciliation = vi.fn(async () => undefined) const deps = createDependencies({ getClientForWorkspace: vi.fn(async () => ({ client, apiKey: "sk_test" })), @@ -444,7 +443,6 @@ describe("loadWorkspaceShellInitialState", () => { })), listSourcesForWorkspace: vi.fn(async () => [parsingSource]), reconcileSourcesForWorkspace, - localizeRemoteDocument, startBackgroundReconciliation, }) @@ -456,7 +454,6 @@ describe("loadWorkspaceShellInitialState", () => { parsingSource.id, "sk_test", ) - expect(localizeRemoteDocument).not.toHaveBeenCalled() expect(state.sources).toEqual([ expect.objectContaining({ id: "demo-tsla-q4-2025", @@ -535,17 +532,6 @@ function createDependencies( listMessages: vi.fn(async () => []), listSourcesForWorkspace: vi.fn(async () => []), reconcileSourcesForWorkspace: vi.fn(async () => []), - localizeRemoteDocument: vi.fn(async (workspaceId, input) => - makeSource(workspaceId, { - id: `source_${input.documentId}`, - title: input.title, - mimeType: input.mimeType, - sizeBytes: input.sizeBytes, - status: input.status, - knowhereJobId: null, - knowhereDocumentId: input.documentId, - }), - ), sourceViewOptionsBySourceId: vi.fn(() => Effect.succeed(new Map())), ...overrides, } diff --git a/src/domains/workspace/initial-state.ts b/src/domains/workspace/initial-state.ts index abd7278..368d43b 100644 --- a/src/domains/workspace/initial-state.ts +++ b/src/domains/workspace/initial-state.ts @@ -13,7 +13,7 @@ import { import { chatThreadService } from "@/domains/chat/thread-service" import { toChatMessageView, toChatThreadView } from "@/domains/chat/view" import { sourceViewOptionsBySourceId as getSourceViewOptionsBySourceId } from "@/domains/sources/counts" -import { localizeRemoteLibrarySources } from "@/domains/sources/remote-library" +import { listRemoteLibrarySourceViews } from "@/domains/sources/remote-library" import { reconcileSourcesForWorkspace as reconcileDefaultSourcesForWorkspace } from "@/domains/sources/reconcile" import { sourceService } from "@/domains/sources/service" import { @@ -146,10 +146,6 @@ type WorkspaceShellInitialStateDependencies = { client: WorkspaceShellInitialStateClient, ) => Promise readonly startBackgroundReconciliation?: typeof defaultStartBackgroundReconciliation - readonly localizeRemoteDocument: ( - workspaceId: string, - input: Parameters[1], - ) => Promise readonly sourceViewOptionsBySourceId: ( sources: readonly Source[], client: WorkspaceShellInitialStateClient, @@ -168,7 +164,6 @@ const defaultDependencies: WorkspaceShellInitialStateDependencies = { listSourcesForWorkspace: sourceWorkflowRuntime.listForWorkspace, reconcileSourcesForWorkspace: reconcileDefaultSourcesForWorkspace, startBackgroundReconciliation: defaultStartBackgroundReconciliation, - localizeRemoteDocument: sourceService.localizeRemoteDocument, sourceViewOptionsBySourceId: getSourceViewOptionsBySourceId, } @@ -324,17 +319,16 @@ export const loadWorkspaceShellInitialStateEffect = ( ) .filter((source) => !hiddenDemoSourceIds.has(source.demoSourceId)) const demoSources = visibleDemoCatalogSources.map(demoView.toSourceView) - const workspaceSources = yield* effectOperation.addContext( + const workspaceSources = demoSourceResolution.workspaceSources + const remoteSourceViews = yield* effectOperation.addContext( { context: workspaceInitialStateContext, - operation: "localizeRemoteLibrarySources", + operation: "listRemoteLibrarySourceViews", }, - localizeRemoteLibrarySources({ + listRemoteLibrarySourceViews({ workspace, client, localSources: demoSourceResolution.workspaceSources, - localizeDocument: (document) => - deps.localizeRemoteDocument(workspace.id, document), }), ) const sourcesNeedingKnowhereChunkCount = @@ -385,6 +379,7 @@ export const loadWorkspaceShellInitialStateEffect = ( sourceOptions.get(source.id), ), ), + ...remoteSourceViews, ], officialLibrarySources: toOfficialLibrarySourceViews(demoCatalog), chatThreads: chatThreads.map(toChatThreadView), From b7f5a50167acd8bae98225e16c4f09cc4f8fb352 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 30 Jun 2026 22:56:02 +0800 Subject: [PATCH 5/6] Fix chunk image rendering after full-tree loads --- src/components/chunks-panel.tsx | 4 + src/components/parsed-chunk-card.test.ts | 29 ++++++ src/components/parsed-chunk-card.tsx | 22 ++++- .../workspace-selected-chunks.test.ts | 90 ++++++++++++++++++- src/components/workspace-selected-chunks.ts | 42 +++++++-- 5 files changed, 174 insertions(+), 13 deletions(-) diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index 41aaf32..4a686a7 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -401,6 +401,7 @@ export function ChunksPanel({ hasOriginalFile ? handleChunkSelected : undefined } onReferenceClick={requestChunkFocus} + selectedSourceFile={selectedSourceFile} /> ))} @@ -1129,6 +1130,7 @@ function VirtualChunkRow({ measureElement, onChunkClick, onReferenceClick, + selectedSourceFile, }: { virtualItem: VirtualItem; chunk: ParsedChunkView | undefined; @@ -1137,6 +1139,7 @@ function VirtualChunkRow({ measureElement: (node: HTMLDivElement | null) => void; onChunkClick?: (chunk: ParsedChunkView) => void; onReferenceClick: (chunkId: string) => void; + selectedSourceFile: SourceOriginalFileView | null; }): ReactNode { if (!chunk) { return null; @@ -1163,6 +1166,7 @@ function VirtualChunkRow({ isOriginalPreviewAvailable={isOriginalPreviewAvailable} onChunkClick={onChunkClick} onReferenceClick={onReferenceClick} + sourceOriginalFile={selectedSourceFile} /> ); diff --git a/src/components/parsed-chunk-card.test.ts b/src/components/parsed-chunk-card.test.ts index b269c08..2e34ab7 100644 --- a/src/components/parsed-chunk-card.test.ts +++ b/src/components/parsed-chunk-card.test.ts @@ -182,4 +182,33 @@ describe("ParsedChunkCard", () => { expect(table.innerHTML).not.toContain("script"); expect(table.innerHTML).not.toContain("onclick"); }); + + it("renders the original image file when an image chunk has no asset URL", () => { + render( + React.createElement(ParsedChunkCard, { + chunk: { + chunkId: "image_1", + type: "image", + content: "Logo summary", + summary: "A black and white logo.", + sourceTitle: "logo.png", + }, + isFocused: false, + onReferenceClick: vi.fn(), + sourceOriginalFile: { + url: "https://blob.example/source-uploads/logo.png", + mimeType: "image/png", + }, + }), + ); + + const image = screen.getByRole("img", { + name: "A black and white logo.", + }); + + expect(image.getAttribute("src")).toBe( + "https://blob.example/source-uploads/logo.png", + ); + expect(screen.queryByText("Image content is not available in this view.")).toBeNull(); + }); }); diff --git a/src/components/parsed-chunk-card.tsx b/src/components/parsed-chunk-card.tsx index 386d596..53178f5 100644 --- a/src/components/parsed-chunk-card.tsx +++ b/src/components/parsed-chunk-card.tsx @@ -8,6 +8,7 @@ import { Button } from "@/components/ui/button"; import { Card, CardContent } from "@/components/ui/card"; import { parsedChunkCardModel } from "@/components/parsed-chunk-card-model"; import type { ParsedChunkView } from "@/domains/chunks/types"; +import type { SourceOriginalFileView } from "@/domains/sources/types"; import { cn } from "@/lib/utils"; const keywordPanelClassName = @@ -25,12 +26,14 @@ export function ParsedChunkCard({ isOriginalPreviewAvailable = false, onChunkClick, onReferenceClick, + sourceOriginalFile = null, }: { readonly chunk: ParsedChunkView; readonly isFocused: boolean; readonly isOriginalPreviewAvailable?: boolean; readonly onChunkClick?: (chunk: ParsedChunkView) => void; readonly onReferenceClick: (chunkId: string) => void; + readonly sourceOriginalFile?: SourceOriginalFileView | null; }): ReactNode { if (chunk.type === "image") { return ( @@ -40,6 +43,7 @@ export function ParsedChunkCard({ isFocused={isFocused} isOriginalPreviewAvailable={isOriginalPreviewAvailable} onChunkClick={onChunkClick} + sourceOriginalFile={sourceOriginalFile} /> ); @@ -362,12 +366,16 @@ function ImageChunkCard({ isFocused, isOriginalPreviewAvailable, onChunkClick, + sourceOriginalFile, }: { readonly chunk: ParsedChunkView; readonly isFocused: boolean; readonly isOriginalPreviewAvailable: boolean; readonly onChunkClick?: (chunk: ParsedChunkView) => void; + readonly sourceOriginalFile: SourceOriginalFileView | null; }): ReactNode { + const imageAssetUrl = getImageChunkAssetUrl(chunk, sourceOriginalFile); + return ( - {chunk.assetUrl ? ( + {imageAssetUrl ? (
{/* eslint-disable-next-line @next/next/no-img-element -- Parsed artifact dimensions are not known before render. */} {chunk.summary @@ -407,6 +415,16 @@ function ImageChunkCard({ ); } +function getImageChunkAssetUrl( + chunk: ParsedChunkView, + sourceOriginalFile: SourceOriginalFileView | null, +): string | null { + if (chunk.assetUrl) return chunk.assetUrl; + if (!sourceOriginalFile?.mimeType.startsWith("image/")) return null; + + return sourceOriginalFile.url; +} + function renderTextChunkContent( chunk: ParsedChunkView, onReferenceClick: (chunkId: string) => void, diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index 9a20c73..2772c85 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -1,13 +1,22 @@ // @vitest-environment jsdom -import { renderHook } from "@testing-library/react"; +import { renderHook, waitFor } from "@testing-library/react"; import React from "react"; import { SWRConfig } from "swr"; -import { describe, expect, it } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import type { ReactNode } from "react"; import { useWorkspaceSelectedChunks } from "./workspace-selected-chunks"; +import type { ParsedChunkView } from "@/domains/chunks/types"; import type { SourceView } from "@/domains/sources/types"; +const fetchChunkPageMock = vi.hoisted(() => vi.fn()); + +vi.mock("@/domains/workspace/client", () => ({ + workspaceClient: { + fetchChunkPage: fetchChunkPageMock, + }, +})); + const readySource: SourceView = { id: "source_1", title: "lecture.pdf", @@ -17,7 +26,20 @@ const readySource: SourceView = { }; describe("useWorkspaceSelectedChunks", () => { - it("returns prefetched chunks without asking SWR to page the source", () => { + beforeEach(() => { + fetchChunkPageMock.mockReset(); + fetchChunkPageMock.mockResolvedValue({ + chunks: [], + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 1, + }, + }); + }); + + it("returns prefetched chunks while checking the visible page for media", async () => { const { result } = renderHook( () => useWorkspaceSelectedChunks({ @@ -37,6 +59,9 @@ describe("useWorkspaceSelectedChunks", () => { { wrapper: createSWRWrapper }, ); + await waitFor(() => + expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1), + ); expect(result.current.selectedSource?.id).toBe("source_1"); expect(result.current.selectedChunks.map((chunk) => chunk.chunkId)).toEqual([ "chunk_1", @@ -45,6 +70,65 @@ describe("useWorkspaceSelectedChunks", () => { expect(result.current.isSelectedChunksLoading).toBe(false); }); + it("keeps visible page asset URLs when full-tree chunks arrive without media", async () => { + const pagedImageChunk: ParsedChunkView = { + chunkId: "image_1", + type: "image", + content: "Image summary", + sourceTitle: "logo.png", + assetUrl: "https://blob.example/chunk-assets/image-1.png", + }; + const structureOnlyImageChunk: ParsedChunkView = { + chunkId: "image_1", + type: "image", + content: "Image summary", + sourceTitle: "logo.png", + }; + fetchChunkPageMock.mockResolvedValue({ + chunks: [pagedImageChunk], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }); + + const { result, rerender } = renderHook( + (input: { + readonly prefetchedChunksBySourceId: Readonly< + Record + >; + }) => + useWorkspaceSelectedChunks({ + selectedSourceId: "source_1", + sources: [readySource], + prefetchedChunksBySourceId: input.prefetchedChunksBySourceId, + }), + { + initialProps: { prefetchedChunksBySourceId: {} }, + wrapper: createSWRWrapper, + }, + ); + + await waitFor(() => + expect(result.current.selectedChunks[0]?.assetUrl).toBe( + "https://blob.example/chunk-assets/image-1.png", + ), + ); + + rerender({ + prefetchedChunksBySourceId: { + source_1: [structureOnlyImageChunk], + }, + }); + + expect(result.current.selectedChunks[0]).toMatchObject({ + chunkId: "image_1", + assetUrl: "https://blob.example/chunk-assets/image-1.png", + }); + }); + it("returns an empty chunk list when no source is selected", () => { const { result } = renderHook( () => diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index e1f9efb..4ec9f97 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -38,7 +38,7 @@ export function useWorkspaceSelectedChunks({ ? prefetchedChunksBySourceId[selectedSourceId] : undefined const selectedChunkSourceId = - selectedSource && selectedSource.status === "ready" && !prefetchedSelectedChunks + selectedSource && selectedSource.status === "ready" ? selectedSource.id : null const { @@ -59,13 +59,6 @@ export function useWorkspaceSelectedChunks({ keepPreviousData: false, }, ) - const resolvedPrefetchedChunks = useMemo( - () => - prefetchedSelectedChunks - ? resolveChunkConnectionTargets(prefetchedSelectedChunks) - : undefined, - [prefetchedSelectedChunks], - ) const pagedSelectedChunks = useMemo( () => resolveChunkConnectionTargets( @@ -73,6 +66,16 @@ export function useWorkspaceSelectedChunks({ ), [selectedChunkPages], ) + const resolvedPrefetchedChunks = useMemo( + () => + prefetchedSelectedChunks + ? mergeVisibleChunkAssetUrls( + resolveChunkConnectionTargets(prefetchedSelectedChunks), + pagedSelectedChunks, + ) + : undefined, + [pagedSelectedChunks, prefetchedSelectedChunks], + ) const selectedChunks = selectedSourceId ? (resolvedPrefetchedChunks ?? pagedSelectedChunks) : [] @@ -114,3 +117,26 @@ function fetchChunksByKey([ ]: SourceChunksKey): Promise { return workspaceClient.fetchChunkPage(sourceId, page) } + +function mergeVisibleChunkAssetUrls( + chunks: readonly ParsedChunkView[], + visibleChunks: readonly ParsedChunkView[], +): ParsedChunkView[] { + if (visibleChunks.length === 0) return [...chunks] + + const visibleChunksById = new Map( + visibleChunks.map((chunk) => [chunk.chunkId, chunk]), + ) + + return chunks.map((chunk) => { + if (chunk.assetUrl) return chunk + + const visibleChunk = visibleChunksById.get(chunk.chunkId) + if (!visibleChunk?.assetUrl) return chunk + + return { + ...chunk, + assetUrl: visibleChunk.assetUrl, + } + }) +} From 430548b5e9402e20cac851b1abec9059773eb4af Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 30 Jun 2026 23:47:52 +0800 Subject: [PATCH 6/6] feat: retry failed source parsing --- src/app/api/sources/[sourceId]/route.test.ts | 127 +++++++++++++++ src/app/api/sources/[sourceId]/route.ts | 13 +- src/components/source-row.test.ts | 93 +++++++++++ src/components/source-row.tsx | 32 +++- src/components/sources-panel.test.ts | 34 +++++ src/components/sources-panel.tsx | 7 + src/components/workspace-shell-layout.tsx | 9 ++ src/components/workspace-shell.tsx | 2 + .../workspace-source-workflow.test.ts | 51 +++++++ src/components/workspace-source-workflow.ts | 38 +++++ src/domains/sources/failure-message.test.ts | 40 +++++ src/domains/sources/failure-message.ts | 125 +++++++++++++++ src/domains/sources/knowhere-upload.ts | 13 +- src/domains/sources/retry.test.ts | 139 +++++++++++++++++ src/domains/sources/retry.ts | 144 ++++++++++++++++++ src/domains/sources/route-dependencies.ts | 1 + src/domains/sources/route-request.test.ts | 28 +++- src/domains/sources/route-request.ts | 96 ++++++++++-- src/domains/sources/route-retry.ts | 78 ++++++++++ src/domains/sources/route-service.test.ts | 125 +++++++++++++++ src/domains/sources/route-service.ts | 4 + src/domains/sources/route-types.ts | 23 +++ src/domains/sources/route-upload.ts | 8 +- src/domains/sources/service.ts | 22 +++ src/domains/sources/types.ts | 2 + src/domains/sources/upload.test.ts | 10 +- src/domains/sources/view.test.ts | 24 ++- src/domains/sources/view.ts | 9 +- src/domains/workspace/client.test.ts | 40 +++++ src/domains/workspace/client.ts | 23 +++ src/domains/workspace/route-client.ts | 14 +- 31 files changed, 1337 insertions(+), 37 deletions(-) create mode 100644 src/domains/sources/failure-message.test.ts create mode 100644 src/domains/sources/failure-message.ts create mode 100644 src/domains/sources/retry.test.ts create mode 100644 src/domains/sources/retry.ts create mode 100644 src/domains/sources/route-retry.ts diff --git a/src/app/api/sources/[sourceId]/route.test.ts b/src/app/api/sources/[sourceId]/route.test.ts index b82d620..6f82574 100644 --- a/src/app/api/sources/[sourceId]/route.test.ts +++ b/src/app/api/sources/[sourceId]/route.test.ts @@ -14,7 +14,9 @@ const mocks = vi.hoisted(() => { hideDemoSource: vi.fn(), makeKnowhereClient: vi.fn(), requireUser: vi.fn(), + retrySourceToKnowhere: vi.fn(), softDeleteSource: vi.fn(), + startBackgroundReconciliation: vi.fn(), }; }); @@ -46,10 +48,15 @@ vi.mock("@/integrations/knowhere", () => ({ makeKnowhereClient: mocks.makeKnowhereClient, })); +vi.mock("@/domains/sources/background-reconcile", () => ({ + startBackgroundReconciliation: mocks.startBackgroundReconciliation, +})); + vi.mock("@/domains/sources/service", () => ({ sourceService: { findInWorkspace: mocks.findSourceInWorkspace, hideDemoSource: mocks.hideDemoSource, + retrySourceToKnowhere: mocks.retrySourceToKnowhere, softDelete: mocks.softDeleteSource, }, })); @@ -255,4 +262,124 @@ describe("PATCH /api/sources/[sourceId]", () => { expect(mocks.ensureApiKeyForWorkspace).not.toHaveBeenCalled(); expect(mocks.archive).not.toHaveBeenCalled(); }); + + it("retries a failed source and starts background reconciliation", async () => { + mocks.requireUser.mockResolvedValue({ id: "user_1" }); + mocks.ensureWorkspace.mockResolvedValue({ id: "workspace_1" }); + mocks.findSourceInWorkspace.mockResolvedValue({ + id: "source_1", + workspaceId: "workspace_1", + title: "lecture.pdf", + mimeType: "application/pdf", + sizeBytes: 5, + status: "failed", + failureReason: "Knowhere upload failed.", + knowhereJobId: null, + knowhereDocumentId: null, + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: "source-uploads/upload_1/document.pdf", + originalBlobUrl: + "https://store.public.blob.vercel-storage.com/source-uploads/upload_1/document.pdf", + demoKey: null, + createdAt: new Date("2026-05-10T00:00:00Z"), + updatedAt: new Date("2026-05-10T00:00:00Z"), + deletedAt: null, + }); + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123"); + const knowhereClient = { + jobs: { + create: vi.fn(), + get: vi.fn(), + upload: vi.fn(), + }, + documents: { + archive: mocks.archive, + }, + }; + mocks.makeKnowhereClient.mockReturnValue(knowhereClient); + mocks.retrySourceToKnowhere.mockResolvedValue({ + id: "source_1", + workspaceId: "workspace_1", + title: "lecture.pdf", + mimeType: "application/pdf", + sizeBytes: 5, + status: "parsing", + failureReason: null, + knowhereJobId: "job_retry", + knowhereDocumentId: null, + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: "source-uploads/upload_1/document.pdf", + originalBlobUrl: + "https://store.public.blob.vercel-storage.com/source-uploads/upload_1/document.pdf", + demoKey: null, + createdAt: new Date("2026-05-10T00:00:00Z"), + updatedAt: new Date("2026-05-10T00:00:00Z"), + deletedAt: null, + }); + mocks.startBackgroundReconciliation.mockResolvedValue(undefined); + + const response = await PATCH( + new NextRequest("http://localhost:3001/api/sources/source_1", { + method: "PATCH", + body: JSON.stringify({ retry: true }), + }), + { params: Promise.resolve({ sourceId: "source_1" }) }, + ); + + await expect(response.json()).resolves.toEqual({ + source: { + id: "source_1", + kind: "workspace", + title: "lecture.pdf", + mimeType: "application/pdf", + status: "parsing", + documentId: undefined, + originalFile: { + url: "https://store.public.blob.vercel-storage.com/source-uploads/upload_1/document.pdf", + mimeType: "application/pdf", + sizeBytes: 5, + }, + }, + }); + expect(response.status).toBe(200); + expect(mocks.retrySourceToKnowhere).toHaveBeenCalledWith( + { id: "workspace_1" }, + expect.objectContaining({ id: "source_1", status: "failed" }), + knowhereClient, + ); + expect(mocks.startBackgroundReconciliation).toHaveBeenCalledWith( + "workspace_1", + "source_1", + "jwt_123", + ); + }); + + it("rejects retry requests for failed rows without a saved original Blob", async () => { + mocks.requireUser.mockResolvedValue({ id: "user_1" }); + mocks.ensureWorkspace.mockResolvedValue({ id: "workspace_1" }); + mocks.findSourceInWorkspace.mockResolvedValue({ + id: "source_1", + status: "failed", + originalBlobPathname: null, + originalBlobUrl: null, + }); + + const response = await PATCH( + new NextRequest("http://localhost:3001/api/sources/source_1", { + method: "PATCH", + body: JSON.stringify({ retry: true }), + }), + { params: Promise.resolve({ sourceId: "source_1" }) }, + ); + + await expect(response.json()).resolves.toEqual({ + message: + "This source cannot be retried because its original file is unavailable.", + }); + expect(response.status).toBe(409); + expect(mocks.ensureApiKeyForWorkspace).not.toHaveBeenCalled(); + expect(mocks.retrySourceToKnowhere).not.toHaveBeenCalled(); + }); }); diff --git a/src/app/api/sources/[sourceId]/route.ts b/src/app/api/sources/[sourceId]/route.ts index 9d30567..9d05d81 100644 --- a/src/app/api/sources/[sourceId]/route.ts +++ b/src/app/api/sources/[sourceId]/route.ts @@ -19,18 +19,19 @@ export async function PATCH( ): Promise { const { sourceId } = await context.params const routeContext = await nextRouteContext.read() - const archiveRequest = await sourceRouteRequest.readArchiveSource({ + const mutationRequest = await sourceRouteRequest.readSourceMutation({ cookieHeader: routeContext.cookieHeader, request, sourceId, }) - if (!archiveRequest.ok) { - return nextRouteResponse.toNextResponse(archiveRequest.result) + if (!mutationRequest.ok) { + return nextRouteResponse.toNextResponse(mutationRequest.result) } - const result = await sourceRouteService.archiveSource({ - ...archiveRequest.input, - }) + const result = + mutationRequest.mutation.kind === "archive" + ? await sourceRouteService.archiveSource(mutationRequest.mutation.input) + : await sourceRouteService.retrySource(mutationRequest.mutation.input) return nextRouteResponse.toNextResponse(result) } diff --git a/src/components/source-row.test.ts b/src/components/source-row.test.ts index 9faa567..26e413f 100644 --- a/src/components/source-row.test.ts +++ b/src/components/source-row.test.ts @@ -75,6 +75,99 @@ describe("SourceRow", () => { .toBeTruthy(); }); + it("shows failed source retry with a brief error message", () => { + const onRetryClick = vi.fn(); + const onSelect = vi.fn(); + + render( + React.createElement(SourceRow, { + isArchiving: false, + isSelected: false, + onRetryClick, + onSelect, + source: { + id: "source_1", + mimeType: "application/pdf", + title: "lecture.pdf", + status: "failed", + failureMessage: + "Too many concurrent requests (2/2 active). Please retry after 30 seconds.", + originalFile: { + url: "https://store.public.blob.vercel-storage.com/source-uploads/upload_1/document.pdf", + mimeType: "application/pdf", + }, + }, + }), + ); + + fireEvent.click( + screen.getByRole("button", { + name: "Retry lecture.pdf processing", + }), + ); + + expect( + screen.getByText( + "Too many concurrent requests (2/2 active). Please retry after 30 seconds.", + ), + ).toBeTruthy(); + expect(onRetryClick).toHaveBeenCalledWith("source_1"); + expect(onSelect).not.toHaveBeenCalled(); + }); + + it("shows failed source retry loading locally", () => { + render( + React.createElement(SourceRow, { + isArchiving: false, + isRetrying: true, + isSelected: false, + onRetryClick: vi.fn(), + onSelect: vi.fn(), + source: { + id: "source_1", + mimeType: "application/pdf", + title: "lecture.pdf", + status: "failed", + originalFile: { + url: "https://store.public.blob.vercel-storage.com/source-uploads/upload_1/document.pdf", + mimeType: "application/pdf", + }, + }, + }), + ); + + const retryButton = screen.getByRole("button", { + name: "Retry lecture.pdf processing", + }); + + expect((retryButton as HTMLButtonElement).disabled).toBe(true); + expect(within(retryButton).getByRole("status", { name: "Loading" })) + .toBeTruthy(); + }); + + it("hides retry when a failed source has no saved original file", () => { + render( + React.createElement(SourceRow, { + isArchiving: false, + isSelected: false, + onRetryClick: vi.fn(), + onSelect: vi.fn(), + source: { + id: "source_1", + mimeType: "application/pdf", + title: "legacy.pdf", + status: "failed", + }, + }), + ); + + expect( + screen.queryByRole("button", { + name: "Retry legacy.pdf processing", + }), + ).toBeNull(); + }); + it("links ready sources to the document chunk tree route", () => { const onSelect = vi.fn(); diff --git a/src/components/source-row.tsx b/src/components/source-row.tsx index 77cf9d7..abfe8cb 100644 --- a/src/components/source-row.tsx +++ b/src/components/source-row.tsx @@ -2,7 +2,7 @@ import type { ReactElement } from "react"; import Link from "next/link"; -import { FileText, ListTree, Plus, Trash2 } from "lucide-react"; +import { FileText, ListTree, Plus, RotateCcw, Trash2 } from "lucide-react"; import { Checkbox } from "@/components/ui/checkbox"; import { Spinner } from "@/components/ui/spinner"; @@ -13,9 +13,11 @@ export type SourceRowProps = { readonly isArchiving: boolean; readonly isAdding?: boolean; readonly isNarrow?: boolean; + readonly isRetrying?: boolean; readonly isSelected: boolean; readonly onAddClick?: (sourceId: string) => void; readonly onArchiveClick?: (sourceId: string) => void; + readonly onRetryClick?: (sourceId: string) => void; readonly onSelect: () => void; readonly onToggleIncluded?: (sourceId: string, included: boolean) => void; readonly source: SourceView; @@ -30,12 +32,15 @@ export function SourceRow({ onSelect, onToggleIncluded, onArchiveClick, + onRetryClick, chunkTreeHref, isArchiving, + isRetrying = false, }: SourceRowProps): ReactElement { const isReady = source.status === "ready"; const isBusy = source.status === "uploading" || source.status === "parsing"; const isFailed = source.status === "failed"; + const canRetry = isFailed && source.originalFile !== undefined; const isLibrarySource = source.officialLibrary !== undefined; const isRemoteSource = source.kind === "remote"; @@ -104,6 +109,11 @@ export function SourceRow({ ? "Uploading" : "Failed"}

+ {isFailed && source.failureMessage ? ( +

+ {source.failureMessage} +

+ ) : null}
@@ -137,6 +147,26 @@ export function SourceRow({ {isNarrow ? null : "Add"} )} + {canRetry && onRetryClick ? ( + + ) : null} {onArchiveClick && (