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 c28bfa4..1f51bac 100644 --- a/src/app/api/sources/reconcile/route.ts +++ b/src/app/api/sources/reconcile/route.ts @@ -1,50 +1,27 @@ import { serve } from "@upstash/workflow/nextjs" -import { reconcileSourcesForWorkspace } from "@/domains/sources/reconcile" -import { makeKnowhereClient } from "@/integrations/knowhere" -import { logger } from "@/lib/logger" - -type ReconcilePayload = { - readonly workspaceId: string - readonly sourceId: string - readonly apiKey: string -} - -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 workspace = { id: workspaceId } - const client = makeKnowhereClient(apiKey) - let delay = INITIAL_DELAY_S - - 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 +import { sourceReconcileRouteWorkflow } from "@/domains/sources/source-reconcile-route-workflow" + +type ReconcilePayload = Parameters< + typeof sourceReconcileRouteWorkflow.normalizeReconcilePayload +>[0] + +export const { POST } = serve( + async (context) => { + const payload = sourceReconcileRouteWorkflow.normalizeReconcilePayload( + context.requestPayload, + ) + await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context, + payload, }) - - if (resolved.done) { - logger.info("workflow: source resolved", { - sourceId, - status: resolved.status, - attempts: attempt + 1, - }) - return - } - - await context.sleep(`wait-${attempt}`, delay) - delay = Math.min(Math.round(delay * 1.5), MAX_DELAY_S) - } - - logger.error("workflow: exhausted poll attempts", { - sourceId, - maxAttempts: MAX_POLL_ATTEMPTS, - }) -}) + }, + { + failureFunction: async ({ context, failResponse }) => { + await sourceReconcileRouteWorkflow.markSourceFailedAfterWorkflowFailure( + context.requestPayload, + failResponse, + ) + }, + }, +) 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 2006be6..ea75943 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -11,10 +11,11 @@ import { type ChatTurnValue, } from "@/domains/chat/service" import { chatTurnPersistence } from "@/domains/chat/chat-turn-persistence" -import { reconcileSourcesForWorkspace } from "@/domains/sources/reconcile" -import { localizeRemoteLibrarySources } from "@/domains/sources/remote-library" +import { startBackgroundReconciliation } from "@/domains/sources/background-reconcile" 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,19 +55,19 @@ 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, - client, - localSources: sources, - localizeDocument: (document) => - sourceService.localizeRemoteDocument(workspace.id, document), - }) const loadSourceAssetUrls = (source: (typeof sources)[number]) => sourceService.getParseAssetUrls(workspace.id, source.id) @@ -74,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, @@ -84,7 +85,7 @@ const answerChatEffect = (input: AnswerChatInput) => hardenMediaAssetUrls: ({ results, artifacts }) => hardenChatMediaAssetUrls({ workspaceId: workspace.id, - sources: compatibleSources, + sources, results, artifacts, loadSourceAssetUrls, @@ -143,6 +144,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/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/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/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/parsed-result-assets.test.ts b/src/domains/sources/parsed-result-assets.test.ts index 8e466ac..bbada92 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,521 @@ 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", + expect.objectContaining({ + method: "HEAD", + signal: expect.any(AbortSignal), + }), + ) + expect(fetchResult).toHaveBeenCalledWith( + "https://knowhere.example/result.zip", + 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( + 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") + }) + + 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", () => { + 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, + }) + }) + + 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 { return { jobId: "job_1", diff --git a/src/domains/sources/parsed-result-assets.ts b/src/domains/sources/parsed-result-assets.ts index 6a36732..09a4179 100644 --- a/src/domains/sources/parsed-result-assets.ts +++ b/src/domains/sources/parsed-result-assets.ts @@ -1,15 +1,27 @@ 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" +import { logger } from "@/lib/logger" + export type StoredParsedResultAssets = { resultBlobUrl: string assetUrlsByFilePath: Readonly> } +export type MultipartUploadPart = { + readonly etag: string + readonly partNumber: number +} + export type ParsedResultBlobStore = { put( pathname: string, @@ -23,6 +35,139 @@ 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 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 + 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 +184,44 @@ type ParsedResultWithAssets = { readonly tableChunks?: readonly ParsedTableChunk[] } +type ParsedAssetUpload = { + readonly filePath: string + readonly body: Buffer | string + 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 @@ -52,6 +235,14 @@ export type StoreParsedResultAssetsInput = { } 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 +const parsedAssetBlobOperationTimeoutMs = 30_000 // --------------------------------------------------------------------------- // Effect core @@ -116,6 +307,118 @@ export const storeParsedResultAssetsEffect = Effect.fn( }, ) +export const prepareParsedResultAssetBatchEffect = Effect.fn( + "prepareParsedResultAssetBatch", +)( + function* ({ + workspaceId, + sourceId, + jobId, + documentId, + 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 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 + 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(() => + withTimeout( + `Parsed asset upload for ${asset.filePath}`, + parsedAssetBlobOperationTimeoutMs, + 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 +441,150 @@ export async function storeParsedResultAssets({ ) } +export async function createResultZipMultipartUploadPlan({ + workspaceId, + sourceId, + jobId, + client, + blobStore = vercelMultipartBlobStore, + 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 withTimeout( + "Create result ZIP multipart upload", + resultZipBlobOperationTimeoutMs, + blobStore.createMultipartUpload(pathname, getResultZipBlobOptions()), + ) + + const uploadPlan: ResultZipMultipartUploadPlan = { + pathname, + key: upload.key, + uploadId: upload.uploadId, + sizeBytes, + 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( + 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 startedAt = Date.now() + const job = await client.jobs.get(jobId) + const resultUrl = requireJobResultUrl(job, jobId) + const body = await fetchResultZipRange({ + resultUrl, + startByte, + endByte, + fetchResult, + }) + + 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({ + pathname, + key, + uploadId, + parts, + blobStore = vercelMultipartBlobStore, +}: CompleteResultZipMultipartUploadInput): Promise<{ readonly url: string }> { + 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( + input: PrepareParsedResultAssetBatchInput, +): Promise { + return Effect.runPromise(prepareParsedResultAssetBatchEffect(input)) +} + // --------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------- @@ -158,6 +605,115 @@ 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 +} + +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(/^\.\/+/, "") @@ -179,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" @@ -189,6 +755,173 @@ function getContentTypeForPath(filePath: string): string { return "application/octet-stream" } +async function getResultZipSize( + resultUrl: string, + fetchResult: ResultZipFetch, +): Promise { + 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, + "content-length", + ) + if (contentLength !== null) return contentLength + } + + 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() + + 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 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) { + 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 +} + +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, +): 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 +931,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/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 5c867bf..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,38 +106,103 @@ 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 { 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( @@ -140,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, ) @@ -156,6 +262,7 @@ function normalizeRemoteDocument( documentId, namespace, status: getRemoteSourceStatus(status), + ...(revisionKey ? { revisionKey } : {}), ...(title ? { title } : {}), ...(mimeType ? { mimeType } : {}), ...(sizeBytes !== undefined ? { sizeBytes } : {}), @@ -164,12 +271,68 @@ 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" 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" } @@ -249,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 62accd6..c7c4b6f 100644 --- a/src/domains/sources/repository.ts +++ b/src/domains/sources/repository.ts @@ -14,10 +14,13 @@ 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 readonly saveParseResultEffect: typeof sourceParseResultRepository.saveParseResultEffect + readonly mergeParseAssetUrlsEffect: typeof sourceParseResultRepository.mergeParseAssetUrlsEffect + readonly getParseResultProgressEffect: typeof sourceParseResultRepository.getParseResultProgressEffect readonly getParseAssetUrlsEffect: typeof sourceParseResultRepository.getParseAssetUrlsEffect } @@ -33,9 +36,14 @@ export const sourceRepository: SourceRepository = { demoSourceRepository.upsertMaterializedDemoSourceEffect, markParsingEffect: sourceRowRepository.markParsingEffect, markReadyEffect: sourceRowRepository.markReadyEffect, + updateRevisionKeyEffect: sourceRowRepository.updateRevisionKeyEffect, markFailedEffect: sourceRowRepository.markFailedEffect, 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-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 d21ac90..9476285 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 { localizeRemoteLibrarySources } from "./remote-library" +import { + startBackgroundReconciliation as defaultStartBackgroundReconciliation, +} from "./background-reconcile" +import { listRemoteLibrarySourceViews } 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,41 +88,28 @@ 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({ + 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) 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, @@ -148,19 +139,39 @@ const listSourcesEffect = ( sourceOptions.get(source.id), ), ), + ...remoteSourceViews, ], }) }) 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..34b8855 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,14 +110,16 @@ 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); }); - 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", @@ -136,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", @@ -143,6 +157,12 @@ describe("source route service", () => { sourceFileName: "local-duplicate.pdf", }, ], + pagination: { + page: 2, + pageSize: 200, + total: 2, + totalPages: 2, + }, }) .mockResolvedValueOnce({ documents: [ @@ -153,6 +173,12 @@ describe("source route service", () => { sourceFileName: "legacy.pdf", }, ], + pagination: { + page: 1, + pageSize: 200, + total: 1, + totalPages: 1, + }, }); const knowhereClient = { documents: { @@ -174,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), @@ -230,83 +226,60 @@ 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, + }, ]); }); - 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 +290,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 +320,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 +336,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 +347,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/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-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 c3912db..8a5711e 100644 --- a/src/domains/sources/source-parse-result-repository.ts +++ b/src/domains/sources/source-parse-result-repository.ts @@ -15,18 +15,38 @@ 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, ) => 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* () { @@ -58,6 +78,67 @@ 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 [result] = yield* Effect.promise(() => + db + .insert(sourceParseResults) + .values({ + sourceId, + resultBlobUrl: input.resultBlobUrl, + assetUrls: input.assetUrlsByFilePath, + }) + .onConflictDoUpdate({ + target: sourceParseResults.sourceId, + set: { + resultBlobUrl: input.resultBlobUrl, + assetUrls: buildAtomicAssetUrlsMergeSql( + input.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 +161,7 @@ const getParseAssetUrlsEffect: SourceParseResultRepository["getParseAssetUrlsEff export const sourceParseResultRepository: SourceParseResultRepository = { saveParseResultEffect, + mergeParseAssetUrlsEffect, + getParseResultProgressEffect, getParseAssetUrlsEffect, } 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..0841c54 --- /dev/null +++ b/src/domains/sources/source-reconcile-route-workflow.test.ts @@ -0,0 +1,194 @@ +import { afterEach, describe, expect, it, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + loggerError: vi.fn(), + loggerInfo: vi.fn(), + loggerWarn: vi.fn(), + makeKnowhereClient: vi.fn(), + markFailed: vi.fn(), + markSourceReadyAfterReconciliation: vi.fn(), + pollSourceReconciliation: vi.fn(), +})) + +vi.mock("@/domains/sources/source-reconcile-workflow", () => ({ + markSourceReadyAfterReconciliation: mocks.markSourceReadyAfterReconciliation, + pollSourceReconciliation: mocks.pollSourceReconciliation, +})) + +vi.mock("@/domains/sources/workflow-runtime", () => ({ + sourceWorkflowRuntime: { + markFailed: mocks.markFailed, + }, +})) + +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("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-ready", + segmentIndex: 0, + }) + }) + + it("marks the source ready when Knowhere publishes a document id", 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.markSourceReadyAfterReconciliation.mockResolvedValue({ + status: "ready", + }) + + try { + await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context, + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_1", + }), + }) + } finally { + restore() + } + + 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 () => { + 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-ready", + segmentIndex: 5, + }, + workflowRunId: sourceReconcileRouteWorkflow.getPollWorkflowRunId( + "source_1", + 5, + ), + }, + ]) + }) + + 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< + typeof sourceReconcileRouteWorkflow.setContinuationTriggerForTesting +>[0] extends (input: infer Input) => Promise + ? Input + : never + +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: 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 new file mode 100644 index 0000000..a1e732a --- /dev/null +++ b/src/domains/sources/source-reconcile-route-workflow.ts @@ -0,0 +1,217 @@ +import "server-only" + +import { Client, type WorkflowContext } from "@upstash/workflow" + +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-ready" | "poll-and-mirror" | "asset-batches" + +type NormalizedReconcilePayload = { + readonly workspaceId: string + readonly sourceId: string + readonly apiKey: string + readonly phase: "poll-and-ready" + 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 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; readying source", { + 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-ready", + segmentIndex: nextSegmentIndex, + }, + workflowRunId: getPollWorkflowRunId(sourceId, nextSegmentIndex), + }), + ) + logger.info("workflow: poll continuation triggered", { + sourceId, + maxAttempts: maxPollAttempts, + segmentIndex: nextSegmentIndex, + }) + return + } + + const ready = await context.run("source-ready", async () => + markSourceReadyAfterReconciliation({ + workspaceId, + sourceId, + documentId: jobToPrepare.documentId, + }), + ) + logger.info("workflow: source parse reconciliation finished", { + sourceId, + jobId: jobToPrepare.jobId, + status: ready.status, + }) +} + +function normalizeReconcilePayload( + payload: ReconcilePayload, +): NormalizedReconcilePayload { + return { + workspaceId: payload.workspaceId, + sourceId: payload.sourceId, + apiKey: payload.apiKey, + phase: "poll-and-ready", + 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 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 = { + getPollWorkflowRunId, + markSourceFailedAfterWorkflowFailure, + normalizeReconcilePayload, + runPollAndMirrorWorkflow, + setContinuationTriggerForTesting, +} 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..3f0632e --- /dev/null +++ b/src/domains/sources/source-reconcile-workflow.test.ts @@ -0,0 +1,199 @@ +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("accepts done jobs without a result URL when a document id is published", 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: "ready-to-prepare", + jobId: "job_1", + documentId: "doc_1", + }) + expect(repository.markFailed).not.toHaveBeenCalled() + }) +}) + +describe("markSourceReadyAfterReconciliation", () => { + it("marks ready after Knowhere publishes a document id 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..0355350 --- /dev/null +++ b/src/domains/sources/source-reconcile-workflow.ts @@ -0,0 +1,217 @@ +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" } + } + + 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/source-row-repository.ts b/src/domains/sources/source-row-repository.ts index bfb613f..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 = { @@ -64,12 +66,18 @@ type SourceRowRepository = { sourceId: string, jobId: string, documentId?: string, + requiredStatus?: string, ) => Effect.Effect readonly markReadyEffect: ( workspaceId: string, sourceId: string, documentId: string, ) => Effect.Effect + readonly updateRevisionKeyEffect: ( + workspaceId: string, + sourceId: string, + revisionKey: string, + ) => Effect.Effect readonly markFailedEffect: ( workspaceId: string, sourceId: string, @@ -178,13 +186,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, @@ -197,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, @@ -333,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, @@ -395,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 e2c6b5f..bbec158 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, @@ -56,6 +57,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, @@ -71,11 +79,21 @@ type SourceWorkflowRuntime = UploadRepositoryRuntime & { sourceId: string, documentId: string, ) => Promise + readonly updateRevisionKey: ( + workspaceId: string, + sourceId: string, + revisionKey: string, + ) => Promise readonly saveParseResult: ( workspaceId: string, sourceId: string, input: SaveSourceParseResultInput, ) => Promise + readonly mergeParseAssetUrls: ( + workspaceId: string, + sourceId: string, + input: SaveSourceParseResultInput, + ) => Promise readonly softDelete: ( workspaceId: string, sourceId: string, @@ -132,6 +150,7 @@ const markParsing: SourceWorkflowRuntime["markParsing"] = ( sourceId: string, jobId: string, documentId?: string, + requiredStatus?: string, ) => databaseRuntime.runPromise( sourceRepository.markParsingEffect( @@ -139,6 +158,7 @@ const markParsing: SourceWorkflowRuntime["markParsing"] = ( sourceId, jobId, documentId, + requiredStatus, ), ) @@ -151,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, @@ -192,6 +225,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 +286,7 @@ export const sourceWorkflowRuntime: SourceWorkflowRuntime = { createUploading, findInWorkspace, getParseAssetUrls, + getParseResultProgress, hideDemoSource, listForWorkspace, listHiddenDemoSourceIds, @@ -245,6 +294,8 @@ export const sourceWorkflowRuntime: SourceWorkflowRuntime = { markFailed, markParsing, markReady, + updateRevisionKey, + mergeParseAssetUrls, saveParseResult, softDelete, upsertMaterializedDemoSource, diff --git a/src/domains/workspace/initial-state.test.ts b/src/domains/workspace/initial-state.test.ts index 6361411..9ea24e3 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,8 @@ 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 startBackgroundReconciliation = vi.fn(async () => undefined) const deps = createDependencies({ getClientForWorkspace: vi.fn(async () => ({ client, apiKey: "sk_test" })), getOptionalAuthenticated: vi.fn(async () => ({ @@ -441,20 +443,16 @@ 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(state.sources).toEqual([ expect.objectContaining({ @@ -466,8 +464,8 @@ describe("loadWorkspaceShellInitialState", () => { kind: "workspace", title: "uploaded.pdf", mimeType: "application/pdf", - status: "ready", - documentId: "doc_uploaded", + status: "parsing", + documentId: undefined, }, ]) }) @@ -534,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 1a73f83..368d43b 100644 --- a/src/domains/workspace/initial-state.ts +++ b/src/domains/workspace/initial-state.ts @@ -13,10 +13,12 @@ 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 { 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,10 +145,7 @@ type WorkspaceShellInitialStateDependencies = { workspace: Workspace, client: WorkspaceShellInitialStateClient, ) => Promise - readonly localizeRemoteDocument: ( - workspaceId: string, - input: Parameters[1], - ) => Promise + readonly startBackgroundReconciliation?: typeof defaultStartBackgroundReconciliation readonly sourceViewOptionsBySourceId: ( sources: readonly Source[], client: WorkspaceShellInitialStateClient, @@ -163,7 +163,7 @@ const defaultDependencies: WorkspaceShellInitialStateDependencies = { listMessages: chatThreadService.listMessages, listSourcesForWorkspace: sourceWorkflowRuntime.listForWorkspace, reconcileSourcesForWorkspace: reconcileDefaultSourcesForWorkspace, - localizeRemoteDocument: sourceService.localizeRemoteDocument, + startBackgroundReconciliation: defaultStartBackgroundReconciliation, sourceViewOptionsBySourceId: getSourceViewOptionsBySourceId, } @@ -302,12 +302,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, @@ -322,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 = @@ -342,19 +338,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, @@ -386,6 +379,7 @@ export const loadWorkspaceShellInitialStateEffect = ( sourceOptions.get(source.id), ), ), + ...remoteSourceViews, ], officialLibrarySources: toOfficialLibrarySourceViews(demoCatalog), chatThreads: chatThreads.map(toChatThreadView), @@ -412,15 +406,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(