address comments

Author: icecrasher321

apps/sim/background/cleanup-soft-deletes.test.ts

@@ -0,0 +1,164 @@
+/**
+ * @vitest-environment node
+ */
+
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const {
+  mockBatchDeleteByWorkspaceAndTimestamp,
+  mockDeleteFileMetadata,
+  mockDeleteFiles,
+  mockDeleteRowsById,
+  mockIsUsingCloudStorage,
+  mockLimit,
+  mockOrderBy,
+  mockPrepareChatCleanup,
+  mockSelect,
+  mockSelectRowsByIdChunks,
+  mockTask,
+  mockWhere,
+} = vi.hoisted(() => {
+  const mockLimit = vi.fn(async () => [] as Array<{ key: string }>)
+  const mockOrderBy = vi.fn(() => ({ limit: mockLimit }))
+  const mockWhere = vi.fn(() => ({ orderBy: mockOrderBy, limit: mockLimit }))
+  const mockFrom = vi.fn(() => ({
+    where: mockWhere,
+    leftJoin: vi.fn(() => ({ where: mockWhere })),
+  }))
+  const mockSelect = vi.fn(() => ({ from: mockFrom }))
+
+  return {
+    mockBatchDeleteByWorkspaceAndTimestamp: vi.fn(async () => ({ deleted: 0, failed: 0 })),
+    mockDeleteFileMetadata: vi.fn(async () => true),
+    mockDeleteFiles: vi.fn(async () => ({ deleted: 0, failed: [] as Array<{ key: string }> })),
+    mockDeleteRowsById: vi.fn(async () => ({ deleted: 0, failed: 0 })),
+    mockIsUsingCloudStorage: vi.fn(() => true),
+    mockLimit,
+    mockOrderBy,
+    mockPrepareChatCleanup: vi.fn(async () => ({ execute: vi.fn(async () => undefined) })),
+    mockSelect,
+    mockSelectRowsByIdChunks: vi.fn(async () => [] as unknown[]),
+    mockTask: vi.fn((config: unknown) => config),
+    mockWhere,
+  }
+})
+
+vi.mock('@sim/db', () => ({ db: { select: mockSelect } }))
+
+vi.mock('@sim/db/schema', () => {
+  const table = (cols: string[]) =>
+    Object.fromEntries(cols.map((c) => [c, `col.${c}`])) as Record<string, string>
+  const wsFileCols = ['id', 'key', 'context', 'workspaceId', 'deletedAt', 'uploadedAt']
+  const softCols = ['id', 'archivedAt', 'deletedAt', 'workspaceId']
+  return {
+    a2aAgent: table(softCols),
+    copilotChats: table(['id', 'workflowId']),
+    document: table(['id', 'storageKey', 'knowledgeBaseId']),
+    knowledgeBase: table(softCols),
+    mcpServers: table(softCols),
+    memory: table(softCols),
+    userTableDefinitions: table(softCols),
+    workflow: table(softCols),
+    workflowFolder: table(softCols),
+    workflowMcpServer: table(softCols),
+    workspaceFile: table(wsFileCols),
+    workspaceFiles: table(wsFileCols),
+  }
+})
+
+vi.mock('@sim/logger', () => ({
+  createLogger: vi.fn(() => ({ error: vi.fn(), info: vi.fn(), warn: vi.fn() })),
+}))
+
+vi.mock('@trigger.dev/sdk', () => ({ task: mockTask }))
+
+vi.mock('drizzle-orm', () => ({
+  and: vi.fn((...args: unknown[]) => ({ op: 'and', args })),
+  asc: vi.fn((column: unknown) => ({ op: 'asc', column })),
+  eq: vi.fn((...args: unknown[]) => ({ op: 'eq', args })),
+  inArray: vi.fn((...args: unknown[]) => ({ op: 'inArray', args })),
+  isNotNull: vi.fn((...args: unknown[]) => ({ op: 'isNotNull', args })),
+  isNull: vi.fn((...args: unknown[]) => ({ op: 'isNull', args })),
+  lt: vi.fn((...args: unknown[]) => ({ op: 'lt', args })),
+  sql: vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })),
+}))
+
+vi.mock('@/lib/cleanup/batch-delete', () => ({
+  batchDeleteByWorkspaceAndTimestamp: mockBatchDeleteByWorkspaceAndTimestamp,
+  chunkArray: (items: string[], size: number) => {
+    const chunks: string[][] = []
+    for (let index = 0; index < items.length; index += size) {
+      chunks.push(items.slice(index, index + size))
+    }
+    return chunks
+  },
+  deleteRowsById: mockDeleteRowsById,
+  selectRowsByIdChunks: mockSelectRowsByIdChunks,
+}))
+
+vi.mock('@/lib/cleanup/chat-cleanup', () => ({ prepareChatCleanup: mockPrepareChatCleanup }))
+
+vi.mock('@/lib/uploads', () => ({
+  isUsingCloudStorage: mockIsUsingCloudStorage,
+  StorageService: { deleteFiles: mockDeleteFiles },
+}))
+
+vi.mock('@/lib/uploads/server/metadata', () => ({ deleteFileMetadata: mockDeleteFileMetadata }))
+
+import { runCleanupSoftDeletes } from '@/background/cleanup-soft-deletes'
+
+const basePayload = {
+  label: 'free/1',
+  plan: 'free' as const,
+  retentionHours: 720,
+  workspaceIds: ['ws-1'],
+}
+
+describe('cleanup soft deletes — orphan KB binding sweep', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockIsUsingCloudStorage.mockReturnValue(true)
+    mockLimit.mockResolvedValue([])
+  })
+
+  it('soft-deletes abandoned KB bindings and removes their storage objects', async () => {
+    mockLimit
+      .mockResolvedValueOnce([{ key: 'kb/orphan-1' }, { key: 'kb/orphan-2' }])
+      .mockResolvedValueOnce([])
+
+    await runCleanupSoftDeletes(basePayload)
+
+    expect(mockDeleteFiles).toHaveBeenCalledWith(['kb/orphan-1', 'kb/orphan-2'], 'knowledge-base')
+    expect(mockDeleteFileMetadata).toHaveBeenCalledWith('kb/orphan-1')
+    expect(mockDeleteFileMetadata).toHaveBeenCalledWith('kb/orphan-2')
+    expect(mockDeleteFileMetadata).toHaveBeenCalledTimes(2)
+  })
+
+  it('still removes bindings but skips object deletion without cloud storage', async () => {
+    mockIsUsingCloudStorage.mockReturnValue(false)
+    mockLimit.mockResolvedValueOnce([{ key: 'kb/orphan-1' }]).mockResolvedValueOnce([])
+
+    await runCleanupSoftDeletes(basePayload)
+
+    expect(mockDeleteFiles).not.toHaveBeenCalled()
+    expect(mockDeleteFileMetadata).toHaveBeenCalledWith('kb/orphan-1')
+  })
+
+  it('stops the batch loop when binding deletion makes no progress', async () => {
+    mockLimit.mockResolvedValue([{ key: 'kb/stuck' }])
+    mockDeleteFileMetadata.mockRejectedValue(new Error('db down'))
+
+    await runCleanupSoftDeletes(basePayload)
+
+    // One batch attempted, then the no-progress guard breaks the loop.
+    expect(mockDeleteFileMetadata).toHaveBeenCalledTimes(1)
+  })
+
+  it('does not run the sweep when there are no workspaces', async () => {
+    await runCleanupSoftDeletes({ ...basePayload, workspaceIds: [] })
+
+    expect(mockSelect).not.toHaveBeenCalled()
+    expect(mockDeleteFiles).not.toHaveBeenCalled()
+    expect(mockDeleteFileMetadata).not.toHaveBeenCalled()
+  })
+})

apps/sim/background/cleanup-soft-deletes.ts

@@ -2,6 +2,7 @@ import { db } from '@sim/db'
 import {
   a2aAgent,
   copilotChats,
+  document,
   knowledgeBase,
   mcpServers,
   memory,
@@ -14,19 +15,31 @@ import {
 } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { task } from '@trigger.dev/sdk'
-import { and, inArray, isNotNull, lt } from 'drizzle-orm'
+import { and, asc, eq, inArray, isNotNull, isNull, lt, sql } from 'drizzle-orm'
 import type { CleanupJobPayload } from '@/lib/billing/cleanup-dispatcher'
 import {
   batchDeleteByWorkspaceAndTimestamp,
+  chunkArray,
   deleteRowsById,
   selectRowsByIdChunks,
 } from '@/lib/cleanup/batch-delete'
 import { prepareChatCleanup } from '@/lib/cleanup/chat-cleanup'
 import type { StorageContext } from '@/lib/uploads'
 import { isUsingCloudStorage, StorageService } from '@/lib/uploads'
+import { deleteFileMetadata } from '@/lib/uploads/server/metadata'
 
 const logger = createLogger('CleanupSoftDeletes')
 
+const KB_ORPHAN_BINDING_BATCH_SIZE = 500
+const KB_ORPHAN_BINDING_TOTAL_LIMIT = 5_000
+/**
+ * Grace window before an unreferenced KB binding is swept. Comfortably longer
+ * than any presign → upload → document-insert flow, so an in-flight upload is
+ * never mistaken for an abandoned one.
+ */
+const KB_ORPHAN_BINDING_GRACE_HOURS = 7 * 24
+const KB_ORPHAN_BINDING_WORKSPACE_CHUNK = 50
+
 interface WorkspaceFileScope {
   /** Rows from `workspace_file` (singular, legacy workspace-context only). */
   legacyRows: Array<{ id: string; key: string }>
@@ -162,6 +175,87 @@ const CLEANUP_TARGETS = [
   },
 ] as const
 
+/**
+ * Sweep abandoned knowledge-base ownership bindings. The presigned upload flow
+ * writes a `workspace_files` binding when it hands out an upload URL, before the
+ * object is stored and before any document is created. If the upload is never
+ * completed, that binding is orphaned — no `document.storageKey` ever references
+ * its key. Such bindings are inert (read access requires a live document, and
+ * the move re-point only follows referenced keys), but they accumulate, so we
+ * drop the best-effort object and soft-delete the binding once they are older
+ * than the grace window.
+ */
+async function cleanupOrphanedKnowledgeBaseBindings(
+  workspaceIds: string[],
+  label: string
+): Promise<{ total: number; deleted: number; failed: number }> {
+  const stats = { total: 0, deleted: 0, failed: 0 }
+  if (workspaceIds.length === 0) return stats
+
+  const orphanCutoff = new Date(Date.now() - KB_ORPHAN_BINDING_GRACE_HOURS * 60 * 60 * 1000)
+
+  for (const chunkIds of chunkArray(workspaceIds, KB_ORPHAN_BINDING_WORKSPACE_CHUNK)) {
+    let attempted = 0
+    while (attempted < KB_ORPHAN_BINDING_TOTAL_LIMIT) {
+      const limit = Math.min(
+        KB_ORPHAN_BINDING_BATCH_SIZE,
+        KB_ORPHAN_BINDING_TOTAL_LIMIT - attempted
+      )
+      const rows = await db
+        .select({ key: workspaceFiles.key })
+        .from(workspaceFiles)
+        .where(
+          and(
+            inArray(workspaceFiles.workspaceId, chunkIds),
+            eq(workspaceFiles.context, 'knowledge-base'),
+            isNull(workspaceFiles.deletedAt),
+            lt(workspaceFiles.uploadedAt, orphanCutoff),
+            sql`NOT EXISTS (
+              SELECT 1 FROM ${document} AS doc
+              WHERE doc.storage_key = ${workspaceFiles.key}
+            )`
+          )
+        )
+        .orderBy(asc(workspaceFiles.uploadedAt), asc(workspaceFiles.key))
+        .limit(limit)
+
+      if (rows.length === 0) break
+
+      const keys = rows.map((row) => row.key)
+      stats.total += keys.length
+      attempted += keys.length
+
+      if (isUsingCloudStorage()) {
+        const result = await StorageService.deleteFiles(keys, 'knowledge-base')
+        stats.failed += result.failed.length
+        for (const { key, error } of result.failed) {
+          logger.error(`[${label}] Failed to delete orphan KB object ${key}:`, { error })
+        }
+      }
+
+      let deletedThisBatch = 0
+      for (const key of keys) {
+        try {
+          await deleteFileMetadata(key)
+          deletedThisBatch++
+        } catch (error) {
+          stats.failed++
+          logger.error(`[${label}] Failed to delete orphan KB binding ${key}:`, { error })
+        }
+      }
+      stats.deleted += deletedThisBatch
+
+      // No progress (every delete failed) — stop rather than reselect the same rows.
+      if (deletedThisBatch === 0) break
+    }
+  }
+
+  logger.info(
+    `[${label}/kb_orphan_bindings] Complete: ${stats.deleted}/${stats.total} bindings cleaned, ${stats.failed} failed`
+  )
+  return stats
+}
+
 export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise<void> {
   const startTime = Date.now()
   const { workspaceIds, retentionHours, label } = payload
@@ -256,8 +350,10 @@ export async function runCleanupSoftDeletes(payload: CleanupJobPayload): Promise
     totalDeleted += result.deleted
   }
 
+  const orphanBindingStats = await cleanupOrphanedKnowledgeBaseBindings(workspaceIds, label)
+
   logger.info(
-    `[${label}] Complete: ${totalDeleted} rows deleted, ${fileStats.filesDeleted} files cleaned`
+    `[${label}] Complete: ${totalDeleted} rows deleted, ${fileStats.filesDeleted} files cleaned, ${orphanBindingStats.deleted} orphan KB bindings cleaned`
   )
 
   // Clean up copilot backend + chat storage files after DB rows are gone

apps/sim/lib/knowledge/connectors/sync-engine.ts

@@ -974,6 +974,7 @@ async function addDocument(
     const storageKey = extractStorageKey(urlPath)
     if (storageKey && storageKey !== urlPath) {
       await deleteFile({ key: storageKey, context: 'knowledge-base' }).catch(() => undefined)
+      await deleteFileMetadata(storageKey).catch(() => undefined)
     }
     throw error
   }
@@ -1069,6 +1070,7 @@ async function updateDocument(
     const storageKey = extractStorageKey(urlPath)
     if (storageKey && storageKey !== urlPath) {
       await deleteFile({ key: storageKey, context: 'knowledge-base' }).catch(() => undefined)
+      await deleteFileMetadata(storageKey).catch(() => undefined)
     }
     throw error
   }