fix(knowledge): finalize partial batch effects

Author: TheodoreSpeaks

apps/sim/lib/knowledge/application/add-workspace-files.test.ts

@@ -290,4 +290,53 @@ describe('add workspace files to knowledge base application command', () => {
     expect(mocks.createDocument).toHaveBeenCalledOnce()
     expect(mocks.recordAudit).toHaveBeenCalledOnce()
   })
+
+  it('audits completed documents before propagating a later infrastructure failure', async () => {
+    const failure = new Error('document store unavailable')
+    mocks.resolveFile
+      .mockResolvedValueOnce(workspaceFile)
+      .mockResolvedValueOnce({ ...workspaceFile, id: 'file-2', name: 'second.pdf' })
+    mocks.loadFileContext
+      .mockResolvedValueOnce({
+        fileId: 'file-1',
+        workspaceId: 'workspace-1',
+        workspaceOrganizationId: null,
+        allowPersonalApiKeys: true,
+        billedAccountUserId: 'billing-owner-1',
+      })
+      .mockResolvedValueOnce({
+        fileId: 'file-2',
+        workspaceId: 'workspace-1',
+        workspaceOrganizationId: null,
+        allowPersonalApiKeys: true,
+        billedAccountUserId: 'billing-owner-1',
+      })
+    mocks.createDocument
+      .mockResolvedValueOnce({
+        id: 'document-1',
+        filename: 'report.pdf',
+        fileUrl: 'https://storage.test/report.pdf',
+        fileSize: 100,
+        mimeType: 'application/pdf',
+      })
+      .mockRejectedValueOnce(failure)
+
+    await expect(
+      addWorkspaceFilesToKnowledgeBase.execute({
+        principal: delegatedPrincipal,
+        input: {
+          knowledgeBaseId: 'knowledge-1',
+          assertedWorkspaceId: 'workspace-1',
+          fileReferences: ['files/report.pdf', 'files/second.pdf'],
+        },
+      })
+    ).rejects.toBe(failure)
+
+    expect(mocks.recordAudit).toHaveBeenCalledOnce()
+    expect(mocks.recordAudit).toHaveBeenCalledWith(
+      expect.objectContaining({ resourceId: 'document-1' })
+    )
+    expect(mocks.platformUploaded).toHaveBeenCalledOnce()
+    expect(mocks.captureServerEvent).toHaveBeenCalledOnce()
+  })
 })

apps/sim/lib/knowledge/application/add-workspace-files.ts

@@ -10,7 +10,9 @@ import { knowledgeDelegationPolicy } from '@/lib/knowledge/application/authoriza
 import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case'
 import {
   ADD_WORKSPACE_FILES_COST_POLICY,
+  type KnowledgeBatchExecutionResult,
   requireBoundedKnowledgeBatch,
+  rethrowKnowledgeBatchTerminalFailure,
 } from '@/lib/knowledge/application/batch-policy'
 import {
   KnowledgeUsageLimitExceededError,
@@ -63,6 +65,10 @@ export interface AddWorkspaceFilesToKnowledgeBaseResult {
   cancelled: boolean
 }
 
+interface AddWorkspaceFilesExecutionResult
+  extends AddWorkspaceFilesToKnowledgeBaseResult,
+    KnowledgeBatchExecutionResult {}
+
 interface AddWorkspaceFilesContext extends ActiveKnowledgeBaseContext {
   fileReferences: string[]
 }
@@ -129,7 +135,7 @@ export const addWorkspaceFilesToKnowledgeBase = defineAuthorizedKnowledgeUseCase
       fileReferences,
     }
   },
-  async execute({ principal, input, context }): Promise<AddWorkspaceFilesToKnowledgeBaseResult> {
+  async execute({ principal, input, context }): Promise<AddWorkspaceFilesExecutionResult> {
     const prepared: PreparedWorkspaceFile[] = []
     const failed: string[] = []
     const canonicalFileIds = new Set<string>()
@@ -180,6 +186,7 @@ export const addWorkspaceFilesToKnowledgeBase = defineAuthorizedKnowledgeUseCase
     }
     const uploadedBy = resolveKnowledgeAttributedUserId(principal, context)
     const added: AddedWorkspaceFileDocument[] = []
+    let terminalFailure: KnowledgeBatchExecutionResult['terminalFailure']
 
     for (const candidate of prepared) {
       if (input.cancellationSignal?.aborted) break
@@ -231,7 +238,8 @@ export const addWorkspaceFilesToKnowledgeBase = defineAuthorizedKnowledgeUseCase
           failed.push(candidate.reference)
           continue
         }
-        throw error
+        terminalFailure = { error }
+        break
       }
     }
 
@@ -241,6 +249,7 @@ export const addWorkspaceFilesToKnowledgeBase = defineAuthorizedKnowledgeUseCase
       added,
       failed,
       cancelled: input.cancellationSignal?.aborted ?? false,
+      ...(terminalFailure && { terminalFailure }),
     }
   },
   projectAudit: ({ input, context, result }) =>
@@ -260,29 +269,33 @@ export const addWorkspaceFilesToKnowledgeBase = defineAuthorizedKnowledgeUseCase
       },
     })),
   afterSuccess: ({ principal, context, result }) => {
-    const userId = resolveKnowledgeAttributedUserId(principal, context)
-    for (const document of result.added) {
-      PlatformEvents.knowledgeBaseDocumentsUploaded({
-        knowledgeBaseId: context.knowledgeBaseId,
-        documentsCount: 1,
-        uploadType: 'single',
-        mimeType: document.mimeType,
-        fileSize: document.fileSize,
-      })
-      captureServerEvent(
-        userId,
-        'knowledge_base_document_uploaded',
-        {
-          knowledge_base_id: context.knowledgeBaseId,
-          workspace_id: context.workspaceId,
-          document_count: 1,
-          upload_type: 'single',
-        },
-        {
-          groups: { workspace: context.workspaceId },
-          setOnce: { first_document_uploaded_at: new Date().toISOString() },
-        }
-      )
+    try {
+      const userId = resolveKnowledgeAttributedUserId(principal, context)
+      for (const document of result.added) {
+        PlatformEvents.knowledgeBaseDocumentsUploaded({
+          knowledgeBaseId: context.knowledgeBaseId,
+          documentsCount: 1,
+          uploadType: 'single',
+          mimeType: document.mimeType,
+          fileSize: document.fileSize,
+        })
+        captureServerEvent(
+          userId,
+          'knowledge_base_document_uploaded',
+          {
+            knowledge_base_id: context.knowledgeBaseId,
+            workspace_id: context.workspaceId,
+            document_count: 1,
+            upload_type: 'single',
+          },
+          {
+            groups: { workspace: context.workspaceId },
+            setOnce: { first_document_uploaded_at: new Date().toISOString() },
+          }
+        )
+      }
+    } finally {
+      rethrowKnowledgeBatchTerminalFailure(result)
     }
   },
 })

apps/sim/lib/knowledge/application/batch-policy.ts

@@ -17,6 +17,18 @@ export const BULK_DELETE_KNOWLEDGE_DOCUMENTS_COST_POLICY = {
   execution: 'sequential_best_effort',
 } as const
 
+export interface KnowledgeBatchTerminalFailure {
+  error: unknown
+}
+
+export interface KnowledgeBatchExecutionResult {
+  terminalFailure?: KnowledgeBatchTerminalFailure
+}
+
+export function rethrowKnowledgeBatchTerminalFailure(result: KnowledgeBatchExecutionResult): void {
+  if (result.terminalFailure) throw result.terminalFailure.error
+}
+
 export function requireBoundedKnowledgeBatch(
   items: readonly string[],
   resource: string,

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

@@ -460,4 +460,31 @@ describe('knowledge document application use cases', () => {
     expect(mocks.deleteDocument).toHaveBeenCalledOnce()
     expect(mocks.recordAudit).toHaveBeenCalledOnce()
   })
+
+  it('audits completed document deletions before propagating infrastructure failure', async () => {
+    const failure = new Error('document store unavailable')
+    mocks.resolveCanonicalDocument.mockImplementation(async ({ documentId }) => ({
+      ...context,
+      documentId,
+      document: { ...document, id: documentId },
+    }))
+    mocks.deleteDocument.mockResolvedValueOnce(undefined).mockRejectedValueOnce(failure)
+
+    await expect(
+      bulkDeleteKnowledgeDocuments.execute({
+        principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' },
+        input: {
+          knowledgeBaseId: 'knowledge-1',
+          assertedWorkspaceId: 'workspace-1',
+          documentIds: ['document-1', 'document-2'],
+        },
+      })
+    ).rejects.toBe(failure)
+
+    expect(mocks.recordAudit).toHaveBeenCalledOnce()
+    expect(mocks.recordAudit).toHaveBeenCalledWith(
+      expect.objectContaining({ resourceId: 'document-1' })
+    )
+    expect(mocks.captureServerEvent).toHaveBeenCalledOnce()
+  })
 })

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

@@ -12,7 +12,9 @@ import { knowledgeDelegationPolicy } from '@/lib/knowledge/application/authoriza
 import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case'
 import {
   BULK_DELETE_KNOWLEDGE_DOCUMENTS_COST_POLICY,
+  type KnowledgeBatchExecutionResult,
   requireBoundedKnowledgeBatch,
+  rethrowKnowledgeBatchTerminalFailure,
 } from '@/lib/knowledge/application/batch-policy'
 import {
   KnowledgeUsageLimitExceededError,
@@ -142,6 +144,10 @@ export interface BulkDeleteKnowledgeDocumentsResult {
   cancelled: boolean
 }
 
+interface BulkDeleteKnowledgeDocumentsExecutionResult
+  extends BulkDeleteKnowledgeDocumentsResult,
+    KnowledgeBatchExecutionResult {}
+
 interface BulkDeleteKnowledgeDocumentsContext extends ActiveKnowledgeBaseContext {
   documentIds: string[]
 }
@@ -585,9 +591,14 @@ export const bulkDeleteKnowledgeDocuments = defineAuthorizedKnowledgeUseCase({
       documentIds,
     }
   },
-  async execute({ principal, input, context }): Promise<BulkDeleteKnowledgeDocumentsResult> {
+  async execute({
+    principal,
+    input,
+    context,
+  }): Promise<BulkDeleteKnowledgeDocumentsExecutionResult> {
     const deletedDocuments: DeletedKnowledgeDocument[] = []
     const failed: string[] = []
+    let terminalFailure: KnowledgeBatchExecutionResult['terminalFailure']
 
     for (const documentId of context.documentIds) {
       if (input.cancellationSignal?.aborted) break
@@ -621,7 +632,8 @@ export const bulkDeleteKnowledgeDocuments = defineAuthorizedKnowledgeUseCase({
           failed.push(documentId)
           continue
         }
-        throw error
+        terminalFailure = { error }
+        break
       }
     }
 
@@ -631,6 +643,7 @@ export const bulkDeleteKnowledgeDocuments = defineAuthorizedKnowledgeUseCase({
       failed,
       deletedDocuments,
       cancelled: input.cancellationSignal?.aborted ?? false,
+      ...(terminalFailure && { terminalFailure }),
     }
   },
   projectAudit: ({ input, context, result }) =>
@@ -650,17 +663,21 @@ export const bulkDeleteKnowledgeDocuments = defineAuthorizedKnowledgeUseCase({
       },
     })),
   afterSuccess: ({ principal, context, result }) => {
-    const userId = resolveKnowledgeAttributedUserId(principal, context)
-    for (const _document of result.deletedDocuments) {
-      captureServerEvent(
-        userId,
-        'knowledge_base_document_deleted',
-        {
-          knowledge_base_id: context.knowledgeBaseId,
-          workspace_id: context.workspaceId,
-        },
-        { groups: { workspace: context.workspaceId } }
-      )
+    try {
+      const userId = resolveKnowledgeAttributedUserId(principal, context)
+      for (const _document of result.deletedDocuments) {
+        captureServerEvent(
+          userId,
+          'knowledge_base_document_deleted',
+          {
+            knowledge_base_id: context.knowledgeBaseId,
+            workspace_id: context.workspaceId,
+          },
+          { groups: { workspace: context.workspaceId } }
+        )
+      }
+    } finally {
+      rethrowKnowledgeBatchTerminalFailure(result)
     }
   },
 })

apps/sim/lib/knowledge/application/knowledge-bases.test.ts

@@ -510,4 +510,30 @@ describe('knowledge base application use cases', () => {
     expect(mocks.deleteRecord).toHaveBeenCalledOnce()
     expect(mocks.recordAudit).toHaveBeenCalledOnce()
   })
+
+  it('audits completed knowledge base deletions before propagating infrastructure failure', async () => {
+    const failure = new Error('knowledge store unavailable')
+    mocks.resolveKnowledgeBase.mockImplementation(async ({ knowledgeBaseId }) => ({
+      ...context,
+      knowledgeBaseId,
+      knowledgeBase: { ...knowledgeBase, id: knowledgeBaseId, name: knowledgeBaseId },
+    }))
+    mocks.deleteRecord.mockResolvedValueOnce(undefined).mockRejectedValueOnce(failure)
+
+    await expect(
+      bulkDeleteKnowledgeBases.execute({
+        principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' },
+        input: {
+          assertedWorkspaceId: 'workspace-1',
+          knowledgeBaseIds: ['knowledge-1', 'knowledge-2'],
+        },
+      })
+    ).rejects.toBe(failure)
+
+    expect(mocks.recordAudit).toHaveBeenCalledOnce()
+    expect(mocks.recordAudit).toHaveBeenCalledWith(
+      expect.objectContaining({ resourceId: 'knowledge-1' })
+    )
+    expect(mocks.knowledgeBaseDeleted).toHaveBeenCalledOnce()
+  })
 })