improvement(knowledge): validate stored chunk counts and scope inline reads to the base

- the bundle gate checks each document's stored chunk count against the
  format's per-document ceiling, so a base holding more chunks than a bundle
  can describe is refused up front rather than written into an invalid manifest
- the inline payload read carries its knowledge base id, matching the chunk
  reads, so neither can reach a row outside the base being exported
- an append is refused once the consumer has aborted, since a destroyed
  archive has no listener left to receive the error it would emit

Author: waleedlatif1

apps/sim/lib/knowledge/application/exports.test.ts

@@ -85,6 +85,7 @@ const documents = [
     characterCount: 40,
     tags: {},
     file: { kind: 'storage', key: 'kb/handbook.pdf' },
+    storedChunkCount: 2,
     hasChunks: true,
   },
 ]
@@ -196,6 +197,19 @@ describe('exportKnowledgeBase', () => {
     expect(mocks.recordAudit).not.toHaveBeenCalled()
   })
 
+  /** The written manifest carries the streamed count, so the gate must check the stored one. */
+  it('refuses a document holding more chunks than the bundle format describes', async () => {
+    mocks.listDocuments.mockResolvedValueOnce([{ ...documents[0], storedChunkCount: 5_001 }])
+
+    await expect(
+      exportKnowledgeBase.execute({
+        principal,
+        input: { knowledgeBaseId: 'knowledge-1', vectors: true },
+      })
+    ).rejects.toMatchObject({ code: 'conflict' })
+    expect(mocks.recordAudit).not.toHaveBeenCalled()
+  })
+
   it('refuses a stored tag definition the bundle format cannot describe', async () => {
     mocks.listTags.mockResolvedValueOnce([
       { slot: 'tag1', displayName: 'Product', fieldType: 'mystery' },

apps/sim/lib/knowledge/application/exports.ts

@@ -72,7 +72,7 @@ export const exportKnowledgeBase = defineAuthorizedKnowledgeUseCase({
       },
       tags,
       documents: documents.map((document) =>
-        toManifestDocument(document, bundleEntryPaths(document), 0)
+        toManifestDocument(document, bundleEntryPaths(document), document.storedChunkCount)
       ),
     })
     return {

apps/sim/lib/knowledge/transfer/export-archive.test.ts

@@ -42,6 +42,7 @@ function exportableDocument(overrides: Partial<ExportableDocument>): ExportableD
     characterCount: 40,
     tags: { tag1: 'Billing' },
     file: { kind: 'storage', key: 'kb/handbook.pdf' },
+    storedChunkCount: 2,
     hasChunks: true,
     ...overrides,
   }
@@ -78,7 +79,7 @@ function bundle(overrides: Partial<KnowledgeBaseExportBundle> = {}): KnowledgeBa
         id: INLINE_ID,
         filename: 'note.txt',
         mimeType: 'text/plain',
-        file: { kind: 'data-uri', documentId: INLINE_ID },
+        file: { kind: 'data-uri', knowledgeBaseId: 'kb-1', documentId: INLINE_ID },
         hasChunks: false,
       }),
       exportableDocument({ id: TEXT_ONLY_ID, filename: 'wiki page', file: null, hasChunks: true }),

apps/sim/lib/knowledge/transfer/export-archive.ts

@@ -39,7 +39,7 @@ async function openFileSource(source: ExportableFileSource): Promise<Readable |
     return downloadFileStream({ key: source.key, context: 'knowledge-base' })
   }
   return decodeDataUriWithinLimit(
-    await readInlineFileUrl(source.documentId),
+    await readInlineFileUrl(source.knowledgeBaseId, source.documentId),
     MAX_KNOWLEDGE_DOCUMENT_FILE_SIZE
   ).buffer
 }
@@ -65,14 +65,20 @@ function toChunkLine(chunk: ExportableChunk, vectors: boolean): KnowledgeBundleC
  * entry at a time and emits `entry` exactly once per append, or `error` in its
  * place, which `once` turns into a rejection. A consumer that goes away
  * destroys the archive without either event, so `closed` aborts the wait and
- * the in-flight source is released instead of leaking.
+ * the in-flight source is released instead of leaking; an append is refused
+ * outright once it has fired, since a destroyed archive has no listener left
+ * to receive the error it would emit.
  */
 async function appendEntry(
   archive: ZipArchive,
   source: Readable | Buffer | string,
   name: string,
   closed: AbortSignal
 ): Promise<void> {
+  if (closed.aborted) {
+    if (source instanceof Readable) source.destroy()
+    closed.throwIfAborted()
+  }
   const consumed = once(archive, 'entry', { signal: closed })
   archive.append(source, { name })
   try {

apps/sim/lib/knowledge/transfer/export-source.ts

@@ -27,10 +27,16 @@ const CHUNK_PAGE_SIZE = { text: 500, vectors: 100 } as const
 /** Where a document's original bytes come from, when it has any. */
 export type ExportableFileSource =
   | { kind: 'storage'; key: string }
-  | { kind: 'data-uri'; documentId: string }
+  | { kind: 'data-uri'; knowledgeBaseId: string; documentId: string }
 
 export interface ExportableDocument extends ExportableDocumentRecord {
   file: ExportableFileSource | null
+  /**
+   * Chunks the document reports holding. The archive writes what its chunk
+   * stream actually produced, which can only be lower; this is what the bundle
+   * gate checks against the format's per-document ceiling before any byte streams.
+   */
+  storedChunkCount: number
   /** True when the document finished processing and holds chunks worth exporting. */
   hasChunks: boolean
 }
@@ -50,13 +56,12 @@ function exportableDocumentCondition(knowledgeBaseId: string) {
   )
 }
 
-function fileSourceFor(row: {
-  id: string
-  storageKey: string | null
-  hasInlineFile: boolean
-}): ExportableFileSource | null {
+function fileSourceFor(
+  knowledgeBaseId: string,
+  row: { id: string; storageKey: string | null; hasInlineFile: boolean }
+): ExportableFileSource | null {
   if (row.storageKey) return { kind: 'storage', key: row.storageKey }
-  if (row.hasInlineFile) return { kind: 'data-uri', documentId: row.id }
+  if (row.hasInlineFile) return { kind: 'data-uri', knowledgeBaseId, documentId: row.id }
   return null
 }
 
@@ -77,11 +82,14 @@ export async function listExportableTags(
  * reached: the column can hold megabytes per row, so the listing carries a flag
  * and the archive fetches one payload at a time.
  */
-export async function readInlineFileUrl(documentId: string): Promise<string> {
+export async function readInlineFileUrl(
+  knowledgeBaseId: string,
+  documentId: string
+): Promise<string> {
   const [row] = await db
     .select({ fileUrl: document.fileUrl })
     .from(document)
-    .where(eq(document.id, documentId))
+    .where(and(eq(document.knowledgeBaseId, knowledgeBaseId), eq(document.id, documentId)))
     .limit(1)
   if (!row) throw new OrchestrationError('not_found', 'Document not found')
   return row.fileUrl
@@ -140,7 +148,7 @@ export async function listExportableDocuments(
 
   const documents: ExportableDocument[] = []
   for (const row of rows) {
-    const file = fileSourceFor(row)
+    const file = fileSourceFor(knowledgeBaseId, row)
     const hasChunks = row.processingStatus === 'completed' && row.chunkCount > 0
     if (!file && !hasChunks) continue
     documents.push({
@@ -152,6 +160,7 @@ export async function listExportableDocuments(
       tokenCount: row.tokenCount,
       characterCount: row.characterCount,
       file,
+      storedChunkCount: row.chunkCount,
       hasChunks,
       tags: Object.fromEntries(ALL_TAG_SLOTS.map((slot) => [slot, row[slot]])),
     })