improvement(knowledge): read inline document payloads lazily and scope chunk reads to the base

Follow-ups from an independent review pass over the export path.

- the document listing no longer selects fileUrl: a data: document can hold
  megabytes in that column, so the listing carries a flag and the archive
  reads one payload at a time when it reaches that entry
- iterateDocumentChunks filters on knowledgeBaseId as well as documentId, so
  a future caller cannot reach another base's chunks through the bundle
- tag definitions are validated by the single bundle gate, which answers 409
  like every other undescribable value rather than throwing a bare ZodError
- a consumer that abandons the download aborts the append loop and destroys
  the in-flight source instead of leaving it pending
- direct tests for the knowledge use-case builder's authorize() on the
  workspace, organization, and legacy personal branches

Author: waleedlatif1

apps/sim/lib/knowledge/application/authorized-knowledge-use-case.test.ts

@@ -0,0 +1,144 @@
+/**
+ * @vitest-environment node
+ */
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const mocks = vi.hoisted(() => ({
+  resolvePermission: vi.fn(),
+  authorizeOrganization: vi.fn(),
+  recordAudit: vi.fn(),
+}))
+
+vi.mock('@sim/audit', () => ({
+  AuditAction: { KNOWLEDGE_BASE_UPDATED: 'knowledge_base.updated' },
+  AuditResourceType: { KNOWLEDGE_BASE: 'knowledge_base' },
+  recordAudit: mocks.recordAudit,
+}))
+
+vi.mock('@sim/platform-authz/workspace', () => ({
+  permissionSatisfies: (actual: string | null, required: string) => {
+    const rank = { read: 1, write: 2, admin: 3 } as const
+    return (
+      actual !== null && rank[actual as keyof typeof rank] >= rank[required as keyof typeof rank]
+    )
+  },
+  resolveEffectiveWorkspacePermission: mocks.resolvePermission,
+}))
+
+vi.mock('@/lib/core/application/organization-authorization', () => ({
+  authorizeOrganizationOperation: mocks.authorizeOrganization,
+}))
+
+import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case'
+import { knowledgeOperations } from '@/lib/knowledge/application/operations'
+
+const workspaceContext = {
+  workspaceId: 'workspace-1',
+  workspaceOrganizationId: 'organization-1',
+  allowPersonalApiKeys: true,
+  billedAccountUserId: 'owner-1',
+  knowledgeBaseId: 'knowledge-1',
+}
+const organizationContext = {
+  workspaceId: undefined,
+  organizationId: 'organization-1',
+  knowledgeBaseId: 'knowledge-1',
+}
+const legacyContext = {
+  workspaceId: undefined,
+  legacyPersonalOwnerUserId: 'user-1',
+  knowledgeBaseId: 'knowledge-1',
+}
+
+const session = { kind: 'session', userId: 'user-1', sessionId: 'session-1' } as const
+
+function useCaseFor(context: object) {
+  const execute = vi.fn(async () => 'done')
+  const useCase = defineAuthorizedKnowledgeUseCase({
+    operation: knowledgeOperations.read,
+    resolveContext: () => context as never,
+    execute,
+    projectAudit: () => ({
+      action: 'knowledge_base.updated',
+      resourceType: 'knowledge_base',
+      resourceId: 'knowledge-1',
+      resourceName: 'Docs',
+      description: 'audited',
+      metadata: {},
+    }),
+  })
+  return { useCase, execute }
+}
+
+describe('defineAuthorizedKnowledgeUseCase', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mocks.resolvePermission.mockResolvedValue('read')
+    mocks.authorizeOrganization.mockResolvedValue(undefined)
+  })
+
+  /** `authorize` must run the same funnel `execute` does, and nothing else. */
+  it('authorizes a workspace base through the workspace funnel without executing', async () => {
+    const { useCase, execute } = useCaseFor(workspaceContext)
+
+    await useCase.authorize({ principal: session, input: {} })
+    expect(mocks.resolvePermission).toHaveBeenCalledTimes(1)
+    expect(execute).not.toHaveBeenCalled()
+    expect(mocks.recordAudit).not.toHaveBeenCalled()
+
+    mocks.resolvePermission.mockResolvedValue(null)
+    await expect(useCase.authorize({ principal: session, input: {} })).rejects.toMatchObject({
+      name: 'NoWorkspaceAccessError',
+    })
+  })
+
+  it('executes a workspace base and records its audit under the workspace', async () => {
+    const { useCase } = useCaseFor(workspaceContext)
+
+    await expect(useCase.execute({ principal: session, input: {} })).resolves.toBe('done')
+    expect(mocks.recordAudit).toHaveBeenCalledWith(
+      expect.objectContaining({ workspaceId: 'workspace-1', resourceId: 'knowledge-1' })
+    )
+  })
+
+  it('authorizes and executes an organization base through the organization operation', async () => {
+    const { useCase, execute } = useCaseFor(organizationContext)
+
+    await useCase.authorize({ principal: session, input: {} })
+    expect(mocks.authorizeOrganization).toHaveBeenCalledWith(
+      session,
+      knowledgeOperations.read.organizationOperation,
+      organizationContext
+    )
+    expect(execute).not.toHaveBeenCalled()
+    expect(mocks.resolvePermission).not.toHaveBeenCalled()
+
+    await expect(useCase.execute({ principal: session, input: {} })).resolves.toBe('done')
+    expect(mocks.recordAudit).toHaveBeenCalledWith(
+      expect.objectContaining({
+        workspaceId: undefined,
+        resourceId: 'knowledge-1',
+        metadata: expect.objectContaining({ organizationId: 'organization-1' }),
+      })
+    )
+  })
+
+  it('admits only the owner of a legacy personal base', async () => {
+    const { useCase, execute } = useCaseFor(legacyContext)
+
+    await useCase.authorize({ principal: session, input: {} })
+    expect(execute).not.toHaveBeenCalled()
+    await expect(useCase.execute({ principal: session, input: {} })).resolves.toBe('done')
+
+    const stranger = { kind: 'session', userId: 'user-2', sessionId: 'session-2' } as const
+    await expect(useCase.authorize({ principal: stranger, input: {} })).rejects.toMatchObject({
+      code: 'not_found',
+    })
+    await expect(
+      useCase.authorize({
+        principal: { kind: 'workspace_api_key', workspaceId: 'workspace-1', keyId: 'key-1' },
+        input: {},
+      })
+    ).rejects.toMatchObject({ code: 'not_found' })
+  })
+})

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

@@ -135,14 +135,14 @@ describe('exportKnowledgeBase', () => {
       input: { knowledgeBaseId: 'knowledge-1', vectors: true },
     })
     withVectors.chunks('doc-1')
-    expect(mocks.iterateChunks).toHaveBeenLastCalledWith('doc-1', 1536)
+    expect(mocks.iterateChunks).toHaveBeenLastCalledWith('knowledge-1', 'doc-1', 1536)
 
     const textOnly = await exportKnowledgeBase.execute({
       principal,
       input: { knowledgeBaseId: 'knowledge-1', vectors: false },
     })
     textOnly.chunks('doc-1')
-    expect(mocks.iterateChunks).toHaveBeenLastCalledWith('doc-1', null)
+    expect(mocks.iterateChunks).toHaveBeenLastCalledWith('knowledge-1', 'doc-1', null)
     expect(textOnly.embedding.vectorsIncluded).toBe(false)
   })
 
@@ -196,6 +196,19 @@ describe('exportKnowledgeBase', () => {
     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' },
+    ])
+
+    await expect(
+      exportKnowledgeBase.execute({
+        principal,
+        input: { knowledgeBaseId: 'knowledge-1', vectors: true },
+      })
+    ).rejects.toMatchObject({ code: 'conflict' })
+  })
+
   it('propagates an oversized base without recording audit', async () => {
     mocks.listDocuments.mockRejectedValueOnce(
       new OrchestrationError('payload_too_large', 'Knowledge base has 2001 documents')

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

@@ -6,10 +6,10 @@ import { knowledgeOperations } from '@/lib/knowledge/application/operations'
 import { KNOWLEDGE_BUNDLE_VERSION } from '@/lib/knowledge/constants'
 import { toKbEmbeddingDimensions } from '@/lib/knowledge/embedding-models'
 import {
-  assertDescribableByBundle,
   bundleEntryPaths,
   type KnowledgeBundleManifest,
   type KnowledgeBundleTag,
+  parseDescribableBundle,
   toManifestDocument,
 } from '@/lib/knowledge/transfer/bundle'
 import {
@@ -57,32 +57,32 @@ export const exportKnowledgeBase = defineAuthorizedKnowledgeUseCase({
       listExportableTags(knowledgeBase.id),
       listExportableDocuments(knowledgeBase.id),
     ])
-    const bundle: KnowledgeBaseExportBundle = {
-      knowledgeBase: {
-        name: knowledgeBase.name,
-        description: knowledgeBase.description,
-        chunkingConfig: knowledgeBase.chunkingConfig,
-      },
+    const manifest = parseDescribableBundle({
+      version: KNOWLEDGE_BUNDLE_VERSION,
+      exportedAt: new Date().toISOString(),
       embedding: {
         model: knowledgeBase.embeddingModel,
         dimension,
         vectorsIncluded: input.vectors,
       },
-      tags,
-      documents,
-      chunks: (documentId) => iterateDocumentChunks(documentId, input.vectors ? dimension : null),
-    }
-    assertDescribableByBundle({
-      version: KNOWLEDGE_BUNDLE_VERSION,
-      exportedAt: new Date().toISOString(),
-      embedding: bundle.embedding,
-      knowledgeBase: bundle.knowledgeBase,
+      knowledgeBase: {
+        name: knowledgeBase.name,
+        description: knowledgeBase.description,
+        chunkingConfig: knowledgeBase.chunkingConfig,
+      },
       tags,
       documents: documents.map((document) =>
         toManifestDocument(document, bundleEntryPaths(document), 0)
       ),
     })
-    return bundle
+    return {
+      knowledgeBase: manifest.knowledgeBase,
+      embedding: manifest.embedding,
+      tags: manifest.tags,
+      documents,
+      chunks: (documentId) =>
+        iterateDocumentChunks(knowledgeBase.id, documentId, input.vectors ? dimension : null),
+    }
   },
   projectAudit: ({ context, input, result }) => ({
     action: AuditAction.KNOWLEDGE_BASE_EXPORTED,

apps/sim/lib/knowledge/transfer/bundle.ts

@@ -220,13 +220,14 @@ export function bundleEntryPaths(
 }
 
 /**
- * Refuses an export whose stored values the bundle format cannot describe,
- * before any byte streams: the manifest is written last, so a value the import
- * side would reject must surface as a clear error rather than a truncated archive.
+ * Validates an export's stored values against the bundle format before any byte
+ * streams, and returns them in their wire shape. The manifest is written last,
+ * so a value the import side would reject must surface as a clear error rather
+ * than a truncated archive.
  */
-export function assertDescribableByBundle(manifest: unknown): void {
+export function parseDescribableBundle(manifest: unknown): KnowledgeBundleManifest {
   const result = knowledgeBundleManifestSchema.safeParse(manifest)
-  if (result.success) return
+  if (result.success) return result.data
   const [issue] = result.error.issues
   throw new OrchestrationError(
     'conflict',

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

@@ -2,17 +2,23 @@
  * @vitest-environment node
  */
 import { Readable } from 'node:stream'
+import { sleep } from '@sim/utils/helpers'
 import JSZip from 'jszip'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 const mocks = vi.hoisted(() => ({
   downloadFileStream: vi.fn(),
+  readInlineFileUrl: vi.fn(),
 }))
 
 vi.mock('@/lib/uploads/core/storage-service', () => ({
   downloadFileStream: mocks.downloadFileStream,
 }))
 
+vi.mock('@/lib/knowledge/transfer/export-source', () => ({
+  readInlineFileUrl: mocks.readInlineFileUrl,
+}))
+
 import type { KnowledgeBaseExportBundle } from '@/lib/knowledge/application/exports'
 import { decodeVectorBase64, knowledgeBundleManifestSchema } from '@/lib/knowledge/transfer/bundle'
 import {
@@ -72,10 +78,7 @@ function bundle(overrides: Partial<KnowledgeBaseExportBundle> = {}): KnowledgeBa
         id: INLINE_ID,
         filename: 'note.txt',
         mimeType: 'text/plain',
-        file: {
-          kind: 'data-uri',
-          fileUrl: `data:text/plain;base64,${Buffer.from('hi').toString('base64')}`,
-        },
+        file: { kind: 'data-uri', documentId: INLINE_ID },
         hasChunks: false,
       }),
       exportableDocument({ id: TEXT_ONLY_ID, filename: 'wiki page', file: null, hasChunks: true }),
@@ -98,6 +101,9 @@ describe('buildKnowledgeBundleArchive', () => {
   beforeEach(() => {
     vi.clearAllMocks()
     mocks.downloadFileStream.mockImplementation(async () => Readable.from([Buffer.from('pdf')]))
+    mocks.readInlineFileUrl.mockResolvedValue(
+      `data:text/plain;base64,${Buffer.from('hi').toString('base64')}`
+    )
   })
 
   it('writes files, chunk lines, and a manifest that validates against the bundle schema', async () => {
@@ -190,6 +196,20 @@ describe('buildKnowledgeBundleArchive', () => {
       'entry:manifest.json',
     ])
   })
+
+  /** A browser that abandons the download must not leave the append loop or its blob stream hanging. */
+  it('releases the in-flight source and stops appending when the consumer goes away', async () => {
+    const blob = new Readable({ read() {} })
+    mocks.downloadFileStream.mockResolvedValue(blob)
+    const archive = buildKnowledgeBundleArchive(bundle())
+    await sleep(1)
+    expect(mocks.downloadFileStream).toHaveBeenCalledTimes(1)
+
+    archive.destroy()
+    await sleep(1)
+    expect(blob.destroyed).toBe(true)
+    expect(mocks.readInlineFileUrl).not.toHaveBeenCalled()
+  })
 })
 
 describe('knowledgeBundleFileName', () => {