diff --git a/packages/components/nodes/documentloaders/AAIDomains/AAIDomains.ts b/packages/components/nodes/documentloaders/AAIDomains/AAIDomains.ts index 16feafd87c8..3df934eb6cc 100644 --- a/packages/components/nodes/documentloaders/AAIDomains/AAIDomains.ts +++ b/packages/components/nodes/documentloaders/AAIDomains/AAIDomains.ts @@ -273,36 +273,13 @@ class AAIDomains_DocumentLoaders implements INode { } // Apply metadata - if (metadata) { - const parsedMetadata = typeof metadata === 'object' ? metadata : JSON.parse(metadata) - docs = docs.map((doc) => ({ - ...doc, - metadata: - _omitMetadataKeys === '*' - ? { - ...parsedMetadata - } - : omit( - { - ...doc.metadata, - ...parsedMetadata - }, - omitMetadataKeys - ) - })) - } else { - docs = docs.map((doc) => ({ - ...doc, - metadata: - _omitMetadataKeys === '*' - ? {} - : omit( - { - ...doc.metadata - }, - omitMetadataKeys - ) - })) + const parsedMetadata = metadata ? (typeof metadata === 'object' ? metadata : JSON.parse(metadata)) : null + for (const doc of docs) { + if (_omitMetadataKeys === '*') { + doc.metadata = parsedMetadata ? { ...parsedMetadata } : {} + } else { + doc.metadata = omit({ ...doc.metadata, ...(parsedMetadata || {}) }, omitMetadataKeys) + } } if (output === 'text') { @@ -375,7 +352,7 @@ class AAIDomainsLoader extends BaseDocumentLoader { // Use larger page size since we're only selecting essential fields const pageSize = Math.min(this.limit, 100) - let allDomains: any[] = [] + let allDocs: IDocument[] = [] let currentPage = 0 console.info('[AAIDomains] Starting load with params:', { @@ -388,8 +365,8 @@ class AAIDomainsLoader extends BaseDocumentLoader { hasAnalysis: this.hasAnalysis }) - while (allDomains.length < this.limit) { - const remainingItems = this.limit - allDomains.length + while (allDocs.length < this.limit) { + const remainingItems = this.limit - allDocs.length const currentPageSize = Math.min(pageSize, remainingItems) console.info(`[AAIDomains] Fetching page ${currentPage}, size ${currentPageSize}`) @@ -513,23 +490,24 @@ class AAIDomainsLoader extends BaseDocumentLoader { break } - // Transform domain_tags array to flat tags array - const domainsWithTags = data.map((domain: any) => ({ + // Transform domain_tags array to flat tags array (non-mutating) + const domainsWithTags = (data as any[]).map((domain: any) => ({ ...domain, tags: domain.domain_tags?.map((dt: any) => dt.tags).filter(Boolean) || [] })) // Apply tag filtering if specified - let filteredDomains = domainsWithTags + let filteredDomains: any[] = domainsWithTags if (this.includeTags.length > 0 || this.excludeTags.length > 0) { filteredDomains = this.filterByTags(domainsWithTags) } - allDomains.push(...filteredDomains) + const pageDocs = filteredDomains.map((d: any) => this.createDocumentFromDomain(d)) + allDocs.push(...pageDocs) currentPage++ // Stop if we've fetched enough - if (allDomains.length >= this.limit || data.length < currentPageSize) { + if (allDocs.length >= this.limit || data.length < currentPageSize) { shouldStopPagination = true break } @@ -563,14 +541,14 @@ class AAIDomainsLoader extends BaseDocumentLoader { await new Promise((resolve) => setTimeout(resolve, 200)) } - console.info(`[AAIDomains] Load complete. Total domains: ${allDomains.length}`) + console.info(`[AAIDomains] Load complete. Total documents: ${allDocs.length}`) // Truncate to exact limit - if (allDomains.length > this.limit) { - allDomains = allDomains.slice(0, this.limit) + if (allDocs.length > this.limit) { + allDocs = allDocs.slice(0, this.limit) } - return allDomains.map((domain) => this.createDocumentFromDomain(domain)) + return allDocs } private filterByTags(domains: any[]): any[] { diff --git a/packages/server/src/Interface.DocumentStore.ts b/packages/server/src/Interface.DocumentStore.ts index b68e4b1efbc..113a3318f2e 100644 --- a/packages/server/src/Interface.DocumentStore.ts +++ b/packages/server/src/Interface.DocumentStore.ts @@ -79,6 +79,7 @@ export interface IDocumentStoreLoaderForPreview extends IDocumentStoreLoader { rehydrated?: boolean preview?: boolean previewChunkCount?: number + previewChunkOffset?: number user: IUser } diff --git a/packages/server/src/services/documentstore/index.ts b/packages/server/src/services/documentstore/index.ts index 02b451b9a80..4e09a3a6221 100644 --- a/packages/server/src/services/documentstore/index.ts +++ b/packages/server/src/services/documentstore/index.ts @@ -15,7 +15,7 @@ import { import { StatusCodes } from 'http-status-codes' import { cloneDeep, omit } from 'lodash' import * as path from 'path' -import { DataSource, In } from 'typeorm' +import { DataSource, In, MoreThan } from 'typeorm' import { v4 as uuidv4 } from 'uuid' import { addLoaderSource, @@ -57,6 +57,10 @@ import { checkStorage, updateStorageUsage } from '../../utils/quotaUsage' import { Telemetry } from '../../utils/telemetry' import nodesService from '../nodes' +// Batch sizes for chunk DB operations and vector store upsert +const SAVE_BATCH_SIZE = 500 +const UPSERT_BATCH_SIZE = 500 + const createDocumentStore = async (newDocumentStore: DocumentStore, orgId: string) => { try { const appServer = getRunningExpressApp() @@ -392,6 +396,12 @@ const getDocumentStoreFileChunks = async ( } // Sync and refresh chunks for a specific loader or store +// DURABILITY CONTRACT: syncAndRefreshChunks +// Success path: SYNCING -> SYNC, totalChunks/totalChars = actual persisted count +// Partial failure (batch N>1 fails): SYNCING -> STALE, counts = actual persisted rows +// First-batch failure: SYNCING -> STALE, old chunks untouched (delete deferred until full save) +// Idempotent rerun: safe — delete of old chunks happens only after new set is fully persisted +// Mechanism: safe delete timing — old chunks deleted AFTER all new chunks confirmed saved const syncAndRefreshChunks = async (storeId: string, fileId: string, userId: string, organizationId: string, workspaceId: string) => { try { const appServer = getRunningExpressApp() @@ -432,40 +442,67 @@ const syncAndRefreshChunks = async (storeId: string, fileId: string, userId: str entity.loaders = JSON.stringify(loaders) await appServer.AppDataSource.getRepository(DocumentStore).save(entity) - // Delete existing chunks - await appServer.AppDataSource.getRepository(DocumentStoreFileChunk).delete({ - docId: fileId, - userId, - organizationId - }) + const chunkRepository = appServer.AppDataSource.getRepository(DocumentStoreFileChunk) + const existingChunkIds = ( + await chunkRepository.find({ + where: { + docId: fileId, + userId, + organizationId + } + }) + ).map((chunk) => chunk.id) // Get fresh documents from Google Drive const docs = await _splitIntoChunks(appServer.AppDataSource, componentNodes, data, userId, organizationId) - // Save new chunks - let totalChars = 0 - for (let i = 0; i < docs.length; i++) { - const chunk = docs[i] - totalChars += chunk.pageContent.length + // Save new chunks first (safe delete timing — delete AFTER all saved) + let persistedChunks = 0 + let persistedChars = 0 + let saveFailed = false - const docChunk: DocumentStoreFileChunk = { - userId, - organizationId, - docId: fileId, - storeId: storeId, - id: uuidv4(), - chunkNo: i + 1, - pageContent: chunk.pageContent, - metadata: JSON.stringify(chunk.metadata) + for (let i = 0; i < docs.length; i += SAVE_BATCH_SIZE) { + const batch = docs.slice(i, i + SAVE_BATCH_SIZE) + try { + await Promise.all( + batch.map(async (chunk: IDocument, localIndex: number) => { + const globalIndex = i + localIndex + const docChunk: DocumentStoreFileChunk = { + userId, + organizationId, + docId: fileId, + storeId: storeId, + id: uuidv4(), + chunkNo: globalIndex + 1, + pageContent: chunk.pageContent, + metadata: JSON.stringify(chunk.metadata) + } + const dChunk = chunkRepository.create(docChunk) + await chunkRepository.save(dChunk) + }) + ) + persistedChunks += batch.length + persistedChars += batch.reduce((acc: number, chunk: IDocument) => acc + (chunk.pageContent?.length ?? 0), 0) + } catch (batchError) { + saveFailed = true + break } - - const dChunk = appServer.AppDataSource.getRepository(DocumentStoreFileChunk).create(docChunk) - await appServer.AppDataSource.getRepository(DocumentStoreFileChunk).save(dChunk) } - loader.totalChunks = docs.length - loader.totalChars = totalChars - loader.status = DocumentStoreStatus.SYNC + if (!saveFailed) { + // Delete old chunks only after all new chunks are confirmed saved + if (existingChunkIds.length > 0) { + await chunkRepository.delete({ id: In(existingChunkIds) }) + } + loader.totalChunks = persistedChunks + loader.totalChars = persistedChars + loader.status = DocumentStoreStatus.SYNC + } else { + // Partial failure: record actual persisted count, mark STALE + loader.totalChunks = persistedChunks + loader.totalChars = persistedChars + loader.status = DocumentStoreStatus.STALE + } // Check if all loaders are synced const allSynced = loaders.every((ldr: IDocumentStoreLoader) => ldr.status === DocumentStoreStatus.SYNC) @@ -918,14 +955,25 @@ export const previewChunks = async ({ appDataSource, componentNodes, data, orgId } let docs = await _splitIntoChunks(appDataSource, componentNodes, data, data.userId, data.organizationId) const totalChunks = docs.length + const previewChunkOffset = Math.max(0, data.previewChunkOffset || 0) + // if -1, return all chunks if (data.previewChunkCount === -1) data.previewChunkCount = totalChunks // return all docs if the user ask for more than we have if (totalChunks <= (data.previewChunkCount || 0)) data.previewChunkCount = totalChunks - // return only the first n chunks - if (totalChunks > (data.previewChunkCount || 0)) docs = docs.slice(0, data.previewChunkCount) + const previewChunkCount = Math.max(0, data.previewChunkCount || 0) + if (previewChunkCount > 0) { + docs = docs.slice(previewChunkOffset, previewChunkOffset + previewChunkCount) + } else { + docs = [] + } - return { chunks: docs, totalChunks: totalChunks, previewChunkCount: data.previewChunkCount } + return { + chunks: docs, + totalChunks: totalChunks, + previewChunkCount: previewChunkCount, + previewChunkOffset: previewChunkOffset + } } catch (error) { throw new InternalFlowiseError( StatusCodes.INTERNAL_SERVER_ERROR, @@ -1119,6 +1167,12 @@ const processLoaderMiddleware = async ( } } +// DURABILITY CONTRACT: _saveChunksToStorage +// Success path: SYNCING -> SYNC, totalChunks/totalChars = actual persisted count +// Partial failure (batch N>1 fails): SYNCING -> STALE, counts = actual persisted rows +// First-batch failure: SYNCING -> STALE, old chunks untouched (delete deferred until full save) +// Idempotent rerun: safe — delete of old chunks happens only after new set is fully persisted +// Mechanism: safe delete timing — old chunks deleted AFTER all new chunks confirmed saved const _saveChunksToStorage = async ( appDataSource: DataSource, componentNodes: IComponentNodes, @@ -1227,44 +1281,68 @@ const _saveChunksToStorage = async ( existingLoaders.push(loader) } - //step 7: remove all previous chunks - await appDataSource.getRepository(DocumentStoreFileChunk).delete({ docId: newLoaderId }) + const chunkRepository = appDataSource.getRepository(DocumentStoreFileChunk) + const existingChunkIds = ( + await chunkRepository.find({ + where: { + docId: newLoaderId, + userId: data.userId, + organizationId: data.organizationId + } + }) + ).map((chunk) => chunk.id) + if (response.chunks) { - //step 8: now save the new chunks - const totalChars = response.chunks.reduce((acc, chunk) => { - if (chunk.pageContent) { - return acc + chunk.pageContent.length + //step 7: save new chunks first (safe delete timing — delete AFTER all saved) + let persistedChunks = 0 + let persistedChars = 0 + let saveFailed = false + + for (let i = 0; i < response.chunks.length; i += SAVE_BATCH_SIZE) { + const batch = response.chunks.slice(i, i + SAVE_BATCH_SIZE) + try { + await Promise.all( + batch.map(async (chunk: IDocument, localIndex: number) => { + const globalIndex = i + localIndex + const docChunk: DocumentStoreFileChunk = { + docId: newLoaderId, + storeId: data.storeId || '', + id: uuidv4(), + chunkNo: globalIndex + 1, + pageContent: sanitizeChunkContent(chunk.pageContent), + metadata: JSON.stringify(chunk.metadata), + userId: data.userId, + organizationId: data.organizationId + } + const dChunk = chunkRepository.create(docChunk) + await chunkRepository.save(dChunk) + }) + ) + persistedChunks += batch.length + persistedChars += batch.reduce((acc: number, chunk: IDocument) => acc + (chunk.pageContent?.length ?? 0), 0) + } catch (batchError) { + saveFailed = true + break } - return acc - }, 0) - await Promise.all( - response.chunks.map(async (chunk: IDocument, index: number) => { - try { - const docChunk: DocumentStoreFileChunk = { - docId: newLoaderId, - storeId: data.storeId || '', - id: uuidv4(), - chunkNo: index + 1, - pageContent: sanitizeChunkContent(chunk.pageContent), - metadata: JSON.stringify(chunk.metadata), - userId: data.userId, - organizationId: data.organizationId - } - const dChunk = appDataSource.getRepository(DocumentStoreFileChunk).create(docChunk) - await appDataSource.getRepository(DocumentStoreFileChunk).save(dChunk) - } catch (chunkError) { - throw new InternalFlowiseError( - StatusCodes.INTERNAL_SERVER_ERROR, - `Error: documentStoreServices._saveChunksToStorage - ${getErrorMessage(chunkError)}` - ) - } - }) - ) - // update the loader with the new metrics - loader.totalChunks = response.totalChunks - loader.totalChars = totalChars + } + + if (!saveFailed) { + //step 8: delete old chunks only after all new chunks are confirmed saved + if (existingChunkIds.length > 0) { + await chunkRepository.delete({ id: In(existingChunkIds) }) + } + loader.totalChunks = persistedChunks + loader.totalChars = persistedChars + loader.status = 'SYNC' + } else { + // Partial failure: record actual persisted count, mark STALE + loader.totalChunks = persistedChunks + loader.totalChars = persistedChars + loader.status = DocumentStoreStatus.STALE + } + } else { + loader.status = 'SYNC' } - loader.status = 'SYNC' // have a flag and iterate over the loaders and update the entity status to SYNC const allSynced = existingLoaders.every((ldr: IDocumentStoreLoader) => ldr.status === 'SYNC') entity.status = allSynced ? DocumentStoreStatus.SYNC : DocumentStoreStatus.STALE @@ -1565,25 +1643,112 @@ const _insertIntoVectorStoreWorkerThread = async ( // Prepare docs for upserting const filterOptions: ICommonObject = { - storeId: data.storeId + storeId: data.storeId, + organizationId: entity.organizationId } if (data.docId) { filterOptions['docId'] = data.docId } - const chunks = await appDataSource.getRepository(DocumentStoreFileChunk).find({ - where: filterOptions - }) - const docs: Document[] = chunks.map((chunk: DocumentStoreFileChunk) => { - return new Document({ - pageContent: chunk.pageContent, - metadata: JSON.parse(chunk.metadata) - }) - }) - vStoreNodeData.inputs.document = docs + const recordManagerConfig = + typeof data.recordManagerConfig === 'string' ? JSON.parse(data.recordManagerConfig) : data.recordManagerConfig + const isFullCleanup = recordManagerObj && recordManagerConfig?.cleanup === 'full' + const parseChunkMetadata = (metadata?: string) => { + try { + return metadata ? JSON.parse(metadata) : {} + } catch { + return {} + } + } - // Get Vector Store Instance - const vectorStoreObj = await _createVectorStoreObject(componentNodes, data, vStoreNodeData, upsertHistory) - const indexResult = await vectorStoreObj.vectorStoreMethods.upsert(vStoreNodeData, options) + let indexResult: ICommonObject | undefined + // IMPORTANT: Full-cleanup mode must call upsert() exactly once with ALL docs accumulated. + // The record manager captures indexStartDt = getTime() at the start of each upsert() call + // and deletes all keys with timestamp < indexStartDt after the call completes. + // Splitting into multiple upsert() calls would cause each call to delete the previous + // call's keys, resulting in data loss. Paginate the DB reads, but keep a single upsert(). + if (isFullCleanup) { + const docs: Document[] = [] + // Paginate chunks using PK keyset cursor (id ASC). The chunk dataset is treated as + // stable during this worker iteration — no concurrent chunk writes are expected. + // If concurrent writes are introduced in future, migrate to a monotonic sequence cursor. + let lastId = '' + let hasMore = true + while (hasMore) { + const batchFilter: ICommonObject = lastId ? { ...filterOptions, id: MoreThan(lastId) } : { ...filterOptions } + const chunks = await appDataSource.getRepository(DocumentStoreFileChunk).find({ + where: batchFilter, + take: UPSERT_BATCH_SIZE, + order: { id: 'ASC' } + }) + if (chunks.length === 0) { + hasMore = false + continue + } + for (const chunk of chunks) { + docs.push( + new Document({ + pageContent: chunk.pageContent, + metadata: parseChunkMetadata(chunk.metadata) + }) + ) + } + if (chunks.length < UPSERT_BATCH_SIZE) { + hasMore = false + continue + } + lastId = chunks[chunks.length - 1].id + } + vStoreNodeData.inputs.document = docs + + const vectorStoreObj = await _createVectorStoreObject(componentNodes, data, vStoreNodeData, upsertHistory) + indexResult = await vectorStoreObj.vectorStoreMethods.upsert(vStoreNodeData, options) + } else { + const vectorStoreObj = await _createVectorStoreObject(componentNodes, data, vStoreNodeData, upsertHistory) + // Paginate chunks using PK keyset cursor (id ASC). The chunk dataset is treated as + // stable during this worker iteration — no concurrent chunk writes are expected. + // If concurrent writes are introduced in future, migrate to a monotonic sequence cursor. + let lastId = '' + let hasMore = true + while (hasMore) { + const batchFilter: ICommonObject = lastId ? { ...filterOptions, id: MoreThan(lastId) } : { ...filterOptions } + const chunks = await appDataSource.getRepository(DocumentStoreFileChunk).find({ + where: batchFilter, + take: UPSERT_BATCH_SIZE, + order: { id: 'ASC' } + }) + if (chunks.length === 0) { + hasMore = false + continue + } + const docs: Document[] = chunks.map((chunk: DocumentStoreFileChunk) => { + return new Document({ + pageContent: chunk.pageContent, + metadata: parseChunkMetadata(chunk.metadata) + }) + }) + vStoreNodeData.inputs.document = docs + + const batchResult = await vectorStoreObj.vectorStoreMethods.upsert(vStoreNodeData, options) + if (batchResult) { + if (!indexResult) { + indexResult = { ...batchResult, addedDocs: batchResult.addedDocs ?? [] } + } else { + indexResult.numAdded = (indexResult.numAdded ?? 0) + (batchResult.numAdded ?? 0) + indexResult.numDeleted = (indexResult.numDeleted ?? 0) + (batchResult.numDeleted ?? 0) + indexResult.numUpdated = (indexResult.numUpdated ?? 0) + (batchResult.numUpdated ?? 0) + indexResult.numSkipped = (indexResult.numSkipped ?? 0) + (batchResult.numSkipped ?? 0) + indexResult.totalKeys = batchResult.totalKeys ?? indexResult.totalKeys + indexResult.addedDocs = [...(indexResult.addedDocs ?? []), ...(batchResult.addedDocs ?? [])] + } + } + + if (chunks.length < UPSERT_BATCH_SIZE) { + hasMore = false + continue + } + lastId = chunks[chunks.length - 1].id + } + } // Save to DB if (indexResult) { @@ -1597,16 +1762,18 @@ const _insertIntoVectorStoreWorkerThread = async ( await appDataSource.getRepository(UpsertHistory).save(upsertHistoryItem) } - await telemetry.sendTelemetry( - 'vector_upserted', - { - version: await getAppVersion(), - chatlowId: chatflowid, - type: ChatType.INTERNAL, - flowGraph: omit(indexResult['result'], ['totalKeys', 'addedDocs']) - }, - orgId - ) + if (indexResult) { + await telemetry.sendTelemetry( + 'vector_upserted', + { + version: await getAppVersion(), + chatflowId: chatflowid, + type: ChatType.INTERNAL, + flowGraph: omit(indexResult['result'], ['totalKeys', 'addedDocs']) + }, + orgId + ) + } entity.status = DocumentStoreStatus.UPSERTED await appDataSource.getRepository(DocumentStore).save(entity) diff --git a/packages/server/test/services/documentstore/index.test.ts b/packages/server/test/services/documentstore/index.test.ts new file mode 100644 index 00000000000..b46dce7827e --- /dev/null +++ b/packages/server/test/services/documentstore/index.test.ts @@ -0,0 +1,452 @@ +// RED tests: _saveChunksToStorage durability contract (must FAIL before implementation) + +let mockReturnChunks: any[] = [] + +jest.mock( + '/virtual/test/mock-loader', + () => ({ + nodeClass: class MockLoader { + async init() { + return mockReturnChunks + } + } + }), + { virtual: true } +) + +const mockChunkSave = jest.fn() +const mockChunkCreate = jest.fn((data: any) => data) +const mockChunkDelete = jest.fn() +const mockChunkCount = jest.fn() +const mockChunkFind = jest.fn() +const mockDocStoreSave = jest.fn() +const mockDocStoreFindOneBy = jest.fn() + +// jest.mock calls MUST precede imports (jest hoisting requirement) +jest.mock('flowise-components', () => ({ + addArrayFilesToStorage: jest.fn(), + addSingleFileToStorage: jest.fn(), + getFileFromStorage: jest.fn(), + getFileFromUpload: jest.fn(), + ICommonObject: {}, + IDocument: {}, + mapExtToInputField: jest.fn(), + mapMimeTypeToInputField: jest.fn(), + removeFilesFromStorage: jest.fn(), + removeSpecificFileFromStorage: jest.fn().mockResolvedValue({ totalSize: 0 }), + removeSpecificFileFromUpload: jest.fn() +})) + +jest.mock('@langchain/core/documents', () => ({ + Document: class Document {} +})) + +jest.mock('../../../src/database/entities/DocumentStore', () => { + class DocumentStoreMock {} + return { DocumentStore: DocumentStoreMock } +}) + +jest.mock('../../../src/database/entities/DocumentStoreFileChunk', () => { + class DocumentStoreFileChunkMock {} + return { DocumentStoreFileChunk: DocumentStoreFileChunkMock } +}) + +jest.mock('../../../src/database/entities/ChatFlow', () => ({ + ChatFlow: class ChatFlow {} +})) + +jest.mock('../../../src/database/entities/UpsertHistory', () => ({ + UpsertHistory: class UpsertHistory {} +})) + +jest.mock('../../../src/enterprise/database/entities/workspace.entity', () => ({ + Workspace: class Workspace {} +})) + +jest.mock('../../../src/enterprise/utils/ControllerServiceUtils', () => ({ + getWorkspaceSearchOptions: jest.fn() +})) + +jest.mock('../../../src/errors/internalFlowiseError', () => ({ + InternalFlowiseError: class InternalFlowiseError extends Error { + statusCode: number + constructor(statusCode: number, message: string) { + super(message) + this.statusCode = statusCode + } + } +})) + +jest.mock('../../../src/errors/utils', () => ({ + getErrorMessage: (e: any) => e?.message || String(e) +})) + +jest.mock('../../../src/Interface', () => ({ + addLoaderSource: jest.fn().mockReturnValue('mock-source'), + ChatType: { CHATFLOW: 'CHATFLOW' }, + DocumentStoreDTO: jest.fn(), + DocumentStoreStatus: { + EMPTY_SYNC: 'EMPTY', + SYNC: 'SYNC', + SYNCING: 'SYNCING', + STALE: 'STALE', + NEW: 'NEW', + UPSERTING: 'UPSERTING', + UPSERTED: 'UPSERTED' + }, + MODE: { QUEUE: 'queue' } +})) + +jest.mock('../../../src/UsageCacheManager', () => ({ + UsageCacheManager: class UsageCacheManager {} +})) + +jest.mock('../../../src/utils/getRunningExpressApp', () => ({ + getRunningExpressApp: jest.fn() +})) + +jest.mock('../../../src/utils/logger', () => ({ + default: { + info: jest.fn(), + error: jest.fn(), + debug: jest.fn(), + warn: jest.fn() + } +})) + +jest.mock('../../../src/utils', () => ({ + databaseEntities: {}, + getAppVersion: jest.fn(), + saveUpsertFlowData: jest.fn(), + constructGraphs: jest.fn(), + getEndingNodes: jest.fn(), + getTelemetryFlowObj: jest.fn(), + isFlowValidForStream: jest.fn() +})) + +jest.mock('../../../src/utils/constants', () => ({ + DOCUMENT_STORE_BASE_FOLDER: 'document-store', + INPUT_PARAMS_TYPE: {}, + OMIT_QUEUE_JOB_DATA: [] +})) + +jest.mock('../../../src/utils/quotaUsage', () => ({ + checkStorage: jest.fn(), + updateStorageUsage: jest.fn() +})) + +jest.mock('../../../src/utils/prompt', () => ({ + DOCUMENTSTORE_TOOL_DESCRIPTION_PROMPT_GENERATOR: '' +})) + +jest.mock('../../../src/utils/telemetry', () => ({ + Telemetry: class Telemetry {} +})) + +jest.mock('../../../src/services/nodes', () => ({ + default: {} +})) + +let uuidCounter = 0 +jest.mock('uuid', () => ({ + v4: jest.fn(() => `mock-uuid-${uuidCounter++}`) +})) + +import { processLoader } from '../../../src/services/documentstore/index' +import { DocumentStore } from '../../../src/database/entities/DocumentStore' +import { DocumentStoreFileChunk } from '../../../src/database/entities/DocumentStoreFileChunk' + +const LOADER_ID = 'loader-1' +const STORE_ID = 'store-1' + +const makeChunks = (n: number) => + Array.from({ length: n }, (_, i) => ({ + pageContent: `chunk-content-${i}`, + metadata: { source: 'test', index: i } + })) + +const makeEntity = () => ({ + id: STORE_ID, + name: 'Test Store', + description: 'test', + loaders: JSON.stringify([ + { + id: LOADER_ID, + loaderId: 'mockLoader', + loaderName: 'Mock Loader', + loaderConfig: { text: 'hello' }, + splitterId: '', + splitterConfig: {}, + credential: '', + status: 'SYNCING', + totalChunks: 0, + totalChars: 0, + files: [], + source: '' + } + ]), + whereUsed: '[]', + status: 'SYNCING', + vectorStoreConfig: null, + embeddingConfig: null, + recordManagerConfig: null, + userId: 'user-1', + organizationId: 'org-1', + workspaceId: 'ws-1', + createdDate: new Date(), + updatedDate: new Date() +}) + +const makeData = () => ({ + id: LOADER_ID, + storeId: STORE_ID, + loaderId: 'mockLoader', + loaderName: 'Mock Loader', + loaderConfig: { text: 'hello' }, + splitterId: '', + splitterConfig: {}, + credential: '', + previewChunkCount: -1, + userId: 'user-1', + organizationId: 'org-1', + user: { id: 'user-1', organizationId: 'org-1' } +}) + +const componentNodes: any = { + mockLoader: { filePath: '/virtual/test/mock-loader' } +} + +const buildMockAppDataSource = () => { + const docStoreRepo = { + save: mockDocStoreSave, + findOneBy: mockDocStoreFindOneBy, + create: jest.fn((d: any) => d), + find: jest.fn().mockResolvedValue([]), + count: jest.fn().mockResolvedValue(0), + delete: jest.fn() + } + const chunkRepo = { + save: mockChunkSave, + create: mockChunkCreate, + delete: mockChunkDelete, + find: mockChunkFind, + count: mockChunkCount, + findOneBy: jest.fn() + } + + return { + getRepository: jest.fn((entity: any) => { + if (entity === DocumentStore) return docStoreRepo + if (entity === DocumentStoreFileChunk) return chunkRepo + return {} + }) + } as any +} + +describe('_saveChunksToStorage durability contract', () => { + let appDataSource: any + + beforeEach(() => { + jest.clearAllMocks() + uuidCounter = 0 + appDataSource = buildMockAppDataSource() + mockChunkSave.mockResolvedValue(undefined) + mockChunkDelete.mockResolvedValue(undefined) + mockChunkCount.mockResolvedValue(0) + mockChunkFind.mockResolvedValue([{ id: 'existing-chunk-1' }]) + }) + + describe('happy path', () => { + it('saves all chunks, sets loader status SYNC, counts match actual persisted', async () => { + const chunks = makeChunks(3) + mockReturnChunks = chunks + + const entity = makeEntity() + mockDocStoreFindOneBy.mockResolvedValue(entity) + mockDocStoreSave.mockResolvedValue(entity) + + await processLoader({ + appDataSource, + componentNodes, + data: makeData(), + docLoaderId: LOADER_ID, + orgId: 'org-1', + workspaceId: 'ws-1', + subscriptionId: 'sub-1', + usageCacheManager: {} as any + } as any) + + expect(mockChunkSave).toHaveBeenCalledTimes(3) + + expect(mockDocStoreSave).toHaveBeenCalled() + const savedEntity = mockDocStoreSave.mock.calls[0][0] + const loaders = JSON.parse(savedEntity.loaders) + expect(loaders[0].status).toBe('SYNC') + expect(loaders[0].totalChunks).toBe(3) + + const expectedChars = chunks.reduce((acc, c) => acc + c.pageContent.length, 0) + expect(loaders[0].totalChars).toBe(expectedChars) + expect(mockChunkDelete).toHaveBeenCalledTimes(1) + expect(mockChunkDelete).not.toHaveBeenCalledWith({ docId: LOADER_ID }) + }) + }) + + describe('first-batch failure', () => { + it('does NOT delete old chunks when first save fails, sets loader status STALE', async () => { + const chunks = makeChunks(3) + mockReturnChunks = chunks + + const entity = makeEntity() + mockDocStoreFindOneBy.mockResolvedValue(entity) + mockDocStoreSave.mockResolvedValue(entity) + mockChunkSave.mockRejectedValue(new Error('DB connection lost')) + + let thrownError: any = null + try { + await processLoader({ + appDataSource, + componentNodes, + data: makeData(), + docLoaderId: LOADER_ID, + orgId: 'org-1', + workspaceId: 'ws-1', + subscriptionId: 'sub-1', + usageCacheManager: {} as any + } as any) + } catch (e) { + thrownError = e + } + + // CONTRACT: handle error internally, don't throw — BUG: current code throws → RED + expect(thrownError).toBeNull() + + // CONTRACT: old chunks preserved (delete deferred) — BUG: L1247 deletes BEFORE save → RED + expect(mockChunkDelete).not.toHaveBeenCalled() + + // CONTRACT: entity saved with STALE status + const savedEntity = mockDocStoreSave.mock.calls[0]?.[0] + if (savedEntity) { + const loaders = JSON.parse(savedEntity.loaders) + expect(loaders[0].status).toBe('STALE') + } + }) + }) + + describe('mid-run failure (batch N>1)', () => { + it('preserves completed batches, sets status STALE, counts = actual persisted', async () => { + const chunks = makeChunks(501) // 2 batches (SAVE_BATCH_SIZE=500) + mockReturnChunks = chunks + + const entity = makeEntity() + mockDocStoreFindOneBy.mockResolvedValue(entity) + mockDocStoreSave.mockResolvedValue(entity) + + let saveCount = 0 + mockChunkSave.mockImplementation(() => { + saveCount++ + if (saveCount > 500) { + return Promise.reject(new Error('DB connection lost mid-run')) + } + return Promise.resolve(undefined) + }) + + let thrownError: any = null + try { + await processLoader({ + appDataSource, + componentNodes, + data: makeData(), + docLoaderId: LOADER_ID, + orgId: 'org-1', + workspaceId: 'ws-1', + subscriptionId: 'sub-1', + usageCacheManager: {} as any + } as any) + } catch (e) { + thrownError = e + } + + // CONTRACT: handle partial failure internally — BUG: current code throws → RED + expect(thrownError).toBeNull() + + expect(mockDocStoreSave).toHaveBeenCalled() + const savedEntity = mockDocStoreSave.mock.calls[0]?.[0] + if (savedEntity) { + const loaders = JSON.parse(savedEntity.loaders) + expect(loaders[0].status).toBe('STALE') + // BUG: L1284 sets totalChunks = response.totalChunks (501) not actual (500) → RED + expect(loaders[0].totalChunks).toBe(500) + } + }) + }) + + describe('idempotent rerun', () => { + it('after partial failure, rerun reaches SYNC with correct counts, no duplication', async () => { + const chunks = makeChunks(3) + mockReturnChunks = chunks + + const entity1 = makeEntity() + mockDocStoreFindOneBy.mockResolvedValue(entity1) + mockDocStoreSave.mockResolvedValue(entity1) + + let firstRunSaveCount = 0 + mockChunkSave.mockImplementation(() => { + firstRunSaveCount++ + if (firstRunSaveCount === 2) { + return Promise.reject(new Error('Transient DB error')) + } + return Promise.resolve(undefined) + }) + + try { + await processLoader({ + appDataSource, + componentNodes, + data: makeData(), + docLoaderId: LOADER_ID, + orgId: 'org-1', + workspaceId: 'ws-1', + subscriptionId: 'sub-1', + usageCacheManager: {} as any + } as any) + } catch (_) { + // current code throws on failure — expected for now + } + + jest.clearAllMocks() + uuidCounter = 0 + appDataSource = buildMockAppDataSource() + + const entity2 = makeEntity() + mockDocStoreFindOneBy.mockResolvedValue(entity2) + mockDocStoreSave.mockResolvedValue(entity2) + mockChunkSave.mockResolvedValue(undefined) + mockChunkDelete.mockResolvedValue(undefined) + mockChunkCount.mockResolvedValue(0) + mockChunkFind.mockResolvedValue([{ id: 'existing-chunk-2' }]) + + await processLoader({ + appDataSource, + componentNodes, + data: makeData(), + docLoaderId: LOADER_ID, + orgId: 'org-1', + workspaceId: 'ws-1', + subscriptionId: 'sub-1', + usageCacheManager: {} as any + } as any) + + expect(mockDocStoreSave).toHaveBeenCalled() + const savedEntity = mockDocStoreSave.mock.calls[0][0] + const loaders = JSON.parse(savedEntity.loaders) + expect(loaders[0].status).toBe('SYNC') + expect(loaders[0].totalChunks).toBe(3) + + // CONTRACT: delete happens AFTER save (safe delete timing) + // BUG: L1247 deletes BEFORE L1256 save → RED + const deleteCallOrder = mockChunkDelete.mock.invocationCallOrder[0] + const firstSaveCallOrder = mockChunkSave.mock.invocationCallOrder[0] + expect(deleteCallOrder).toBeGreaterThan(firstSaveCallOrder) + }) + }) +})