refactor(uploads): keep sweepStoragePrefix beside the listing it pages

Moves the prefix sweep step into lib/uploads/core/prefix-sweep.ts with the
listing and delete helpers it drives, and deletes sweep-prefix.ts. The
provider paging cases (a short or empty non-final page, a stalled listing)
now test the step itself against a mocked S3 listing, since a same-module
call to listStorageKeyPage cannot be intercepted.

Author: waleedlatif1

apps/sim/lib/cleanup/chat-backend-cleanup-outbox.ts

@@ -17,9 +17,8 @@ import {
   INLINE_CHAT_IMAGE_PREFIX,
   inlineImageRequestIdSchema,
 } from '@/lib/mothership/chat/inline-image-reference'
-import { outboxDeadlineNear } from '@/lib/uploads/core/prefix-sweep'
+import { outboxDeadlineNear, sweepStoragePrefix } from '@/lib/uploads/core/prefix-sweep'
 import { deleteFile } from '@/lib/uploads/core/storage-service'
-import { sweepStoragePrefix } from '@/lib/uploads/core/sweep-prefix'
 
 const FILE_DELETE_CONCURRENCY = 10
 

apps/sim/lib/projects/files/prefix-cleanup.test.ts

@@ -10,15 +10,10 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 vi.mock('@/lib/uploads/core/setup.server', () => uploadsSetupMock)
 vi.mock('@/lib/uploads/config', () => uploadsConfigMock)
 vi.mock('@/lib/uploads/server/metadata', () => uploadsMetadataMock)
-vi.mock('@/lib/uploads/core/prefix-sweep', async (importOriginal) => {
-  const actual = await importOriginal<typeof import('@/lib/uploads/core/prefix-sweep')>()
-  return { ...actual, listStorageKeyPage: vi.fn(actual.listStorageKeyPage) }
-})
 
 import type { OutboxEventContext } from '@/lib/core/outbox/service'
 import { PROJECT_FILE_PREFIX_CLEANUP_EVENT } from '@/lib/projects/files/outbox-events'
 import { projectFilePrefixCleanupOutboxHandlers } from '@/lib/projects/files/prefix-cleanup'
-import { listStorageKeyPage } from '@/lib/uploads/core/prefix-sweep'
 
 const cleanup = projectFilePrefixCleanupOutboxHandlers[PROJECT_FILE_PREFIX_CLEANUP_EVENT]
 
@@ -83,46 +78,4 @@ describe('Project prefix cleanup on local storage', () => {
 
     expect(await exists('project/p1/zz.txt')).toBe(false)
   })
-
-  it('keeps sweeping after a short page that is not the last', async () => {
-    await put('project/p1/a.txt')
-    await put('project/p1/b.txt')
-    vi.mocked(listStorageKeyPage).mockResolvedValueOnce({
-      keys: ['project/p1/a.txt'],
-      done: false,
-      nextPageToken: 'next',
-    })
-
-    const result = await cleanup({ projectId: 'p1' }, outboxContext())
-
-    expect(result).toMatchObject({ outcome: 'deferred', consumeAttempt: false })
-    expect(await exists('project/p1/a.txt')).toBe(false)
-    expect(await exists('project/p1/b.txt')).toBe(true)
-  })
-
-  it('follows an empty page that is not the last to the objects after it', async () => {
-    await put('project/p1/a.txt')
-    vi.mocked(listStorageKeyPage)
-      .mockResolvedValueOnce({ keys: [], done: false, nextPageToken: 'next' })
-      .mockResolvedValueOnce({ keys: ['project/p1/a.txt'], done: true })
-
-    await expect(cleanup({ projectId: 'p1' }, outboxContext())).resolves.toBeUndefined()
-
-    expect(vi.mocked(listStorageKeyPage).mock.calls[1]?.[3]).toBe('next')
-    expect(await exists('project/p1/a.txt')).toBe(false)
-  })
-
-  it('spends an attempt when its listing stops advancing without removing anything', async () => {
-    vi.mocked(listStorageKeyPage).mockResolvedValue({
-      keys: [],
-      done: false,
-      nextPageToken: 'stuck',
-    })
-
-    const result = await cleanup({ projectId: 'p1' }, outboxContext())
-
-    expect(result).toMatchObject({ outcome: 'deferred' })
-    expect(result).not.toMatchObject({ consumeAttempt: false })
-    expect(listStorageKeyPage).toHaveBeenCalledTimes(2)
-  })
 })

apps/sim/lib/projects/files/prefix-cleanup.ts

@@ -16,8 +16,11 @@ import {
   PROJECT_STORAGE_RECONCILE_EVENT,
 } from '@/lib/projects/files/outbox-events'
 import { isObjectNotFoundError } from '@/lib/uploads/core/errors'
-import { outboxDeadlineNear, storageCleanupAvailableAt } from '@/lib/uploads/core/prefix-sweep'
-import { sweepStoragePrefix } from '@/lib/uploads/core/sweep-prefix'
+import {
+  outboxDeadlineNear,
+  storageCleanupAvailableAt,
+  sweepStoragePrefix,
+} from '@/lib/uploads/core/prefix-sweep'
 import { abortProviderUpload } from '@/lib/uploads/upload-session/provider'
 import type { UploadSessionRecord } from '@/lib/uploads/upload-session/service'
 import { PROJECT_FILE_UPLOAD_BINDING_KEY } from '@/lib/uploads/upload-session/types'

apps/sim/lib/uploads/core/prefix-sweep.test.ts

@@ -0,0 +1,54 @@
+import { storageServiceMock, storageServiceMockFns } from '@sim/testing/mocks/storage-service.mock'
+import { setUploadsConfig, uploadsConfigMock } from '@sim/testing/mocks/uploads-config.mock'
+import { uploadsSetupMock } from '@sim/testing/mocks/uploads-setup.mock'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { mockListS3ObjectKeyPage } = vi.hoisted(() => ({ mockListS3ObjectKeyPage: vi.fn() }))
+
+vi.mock('@/lib/uploads/config', () => uploadsConfigMock)
+vi.mock('@/lib/uploads/core/setup.server', () => uploadsSetupMock)
+vi.mock('@/lib/uploads/core/storage-service', () => storageServiceMock)
+vi.mock('@/lib/uploads/providers/s3/client', () => ({
+  listS3ObjectKeyPage: mockListS3ObjectKeyPage,
+}))
+
+import { sweepStoragePrefix } from '@/lib/uploads/core/prefix-sweep'
+
+const outbox = { signal: new AbortController().signal }
+const deletedKeys = () => storageServiceMockFns.mockDeleteFile.mock.calls.map(([{ key }]) => key)
+
+/** A provider (Azure Blob, GCS, S3-compatibles) may return a short or empty page that is not the last. */
+describe('sweepStoragePrefix over a provider listing', () => {
+  beforeEach(() => {
+    setUploadsConfig({ USE_S3_STORAGE: true, S3_CONFIG: { bucket: 'bucket', region: 'r' } })
+    mockListS3ObjectKeyPage.mockReset()
+    storageServiceMockFns.mockDeleteFile.mockReset().mockResolvedValue(undefined)
+  })
+
+  it('continues after a short page that is not the last', async () => {
+    mockListS3ObjectKeyPage.mockResolvedValueOnce({
+      keys: ['project/p1/a.txt'],
+      nextContinuationToken: 'next',
+    })
+
+    expect(await sweepStoragePrefix('project/p1/', 'project', outbox)).toBe('continue')
+    expect(deletedKeys()).toEqual(['project/p1/a.txt'])
+  })
+
+  it('follows an empty page that is not the last to the objects after it', async () => {
+    mockListS3ObjectKeyPage
+      .mockResolvedValueOnce({ keys: [], nextContinuationToken: 'next' })
+      .mockResolvedValueOnce({ keys: ['project/p1/a.txt'] })
+
+    expect(await sweepStoragePrefix('project/p1/', 'project', outbox)).toBe('done')
+    expect(mockListS3ObjectKeyPage.mock.calls[1]?.[0]).toMatchObject({ continuationToken: 'next' })
+    expect(deletedKeys()).toEqual(['project/p1/a.txt'])
+  })
+
+  it('stalls when its listing stops advancing without removing anything', async () => {
+    mockListS3ObjectKeyPage.mockResolvedValue({ keys: [], nextContinuationToken: 'stuck' })
+
+    expect(await sweepStoragePrefix('project/p1/', 'project', outbox)).toBe('stalled')
+    expect(mockListS3ObjectKeyPage).toHaveBeenCalledTimes(2)
+  })
+})

apps/sim/lib/uploads/core/prefix-sweep.ts

@@ -215,3 +215,34 @@ export async function deleteStorageKeysUnderPrefix(
   }
   return { deleted, complete: true }
 }
+
+/**
+ * One outbox run's step of deleting every object under `prefix`: deletes one listing page's
+ * objects, re-listing from the prefix start on the next step. A provider (Azure Blob, GCS,
+ * S3-compatibles) may return a short or empty page that is not the last, so only `page.done` ends
+ * the sweep; a non-final page that removes nothing is followed to the next one within the step, so
+ * objects behind it are still reached. `continue` after a step that removed something with more
+ * left; `stalled` when no page removed anything before the listing stopped advancing or the
+ * deadline neared.
+ */
+export async function sweepStoragePrefix(
+  prefix: string,
+  context: StorageContext,
+  outbox: Pick<OutboxEventContext, 'signal' | 'deadlineAt'>
+): Promise<'done' | 'continue' | 'stalled'> {
+  let pageToken: string | undefined
+  for (;;) {
+    const page = await listStorageKeyPage(prefix, context, outbox.signal, pageToken)
+    const { deleted, complete } = await deleteStorageKeysUnderPrefix(
+      prefix,
+      context,
+      page.keys,
+      outbox
+    )
+    if (complete && page.done) return 'done'
+    if (deleted > 0) return 'continue'
+    const advanced = page.nextPageToken !== undefined && page.nextPageToken !== pageToken
+    if (!complete || !advanced || outboxDeadlineNear(outbox)) return 'stalled'
+    pageToken = page.nextPageToken
+  }
+}