improvement(cleanup): sweep chat inline images through their own outbox event

Inline image cleanup gets CHAT_INLINE_IMAGE_CLEANUP_EVENT ({ chatIds }) and
its own handler: it sweeps the chat-images/<chatId>/ prefix of each chat
whose row is gone, ten chats at a time, checkpoints how many leading chats
are finished (the lowest unfinished one), and continues or defers past its
deadline. enqueueChatBackendCleanup and prepareChatCleanup each queue it for
the chats they cover; the backend event drops its copilotBackend flag and
imagesSwept checkpoint and keeps accepting the shipped { chatIds, files }
payload. Completed image events are pruned with the backend ones.

The reference check behind deletableChatFiles now runs one query per scope
that returns only the matched keys (passed as one JSON parameter), after an
EXISTS that skips a scope with no chat left. ChatFileRef carries one
optional scope (absent on legacy events, null when no other chat can
reference the key); the event schema still maps a legacy organizationId to
it. Legacy key-scope recovery reuses extractWorkspaceIdFromStorageKey, and
inlineChatImagePrefix validates the chat id for both the key and the sweep.

Author: waleedlatif1

apps/sim/lib/auth/admin-removal-chat-cleanup.ts

@@ -14,17 +14,9 @@ import {
 const ADMIN_REMOVAL_CLEANUP_DELAY_MS = 60 * 60_000
 
 /**
- * The user delete's `before` hook: `/admin/remove-user` cascades the account's chats, and those in
- * the workspaces it owns, away with no transaction to queue their cleanup in, so the cleanup is
- * queued now, before the delete, and delayed past it. The handler purges only chats whose rows
- * are gone, so a delete that fails or is refused purges nothing. A failed read or enqueue throws,
- * which aborts the removal rather than erasing chats whose cleanup nothing recorded. Better Auth's
- * delete runs outside any transaction this hook can join, so what is written between this read
- * and that delete is covered only in part: a chat created in that window is not purged at all, and
- * an attachment (a `workspace_files` row or message key) added to a chat already read stays in
- * storage, while that chat's inline images still go, because the handler sweeps each read chat's
- * whole image prefix. Account deletion, which owns its transaction, re-reads under the account
- * lock instead.
+ * The user delete's `before` hook: queues, delayed past the delete, the cleanup of the chats
+ * `/admin/remove-user` cascades away; a failed read or enqueue throws and aborts the removal. A
+ * chat or attachment written between this read and Better Auth's delete is not covered.
  */
 export async function enqueueAdminRemovedChatCleanup(user: { id: string }): Promise<void> {
   const targets = await readAccountRemovalChats({

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

@@ -1 +1,2 @@
 export const CHAT_BACKEND_CLEANUP_EVENT = 'copilot.chat.backend-cleanup'
+export const CHAT_INLINE_IMAGE_CLEANUP_EVENT = 'copilot.chat.inline-image-cleanup'

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

@@ -0,0 +1,106 @@
+import { copilotChats } from '@sim/db/schema'
+import { queueTableRows, resetDbChainMock } from '@sim/testing/mocks/database.mock'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { mockSweepStoragePrefix, deadline } = vi.hoisted(() => ({
+  mockSweepStoragePrefix: vi.fn(),
+  deadline: { near: false },
+}))
+
+vi.mock('@/lib/uploads/core/prefix-sweep', () => ({
+  sweepStoragePrefix: mockSweepStoragePrefix,
+  outboxDeadlineNear: () => deadline.near,
+}))
+
+import { CHAT_INLINE_IMAGE_CLEANUP_EVENT } from '@/lib/cleanup/chat-backend-cleanup-event'
+import { chatBackendCleanupOutboxHandlers } from '@/lib/cleanup/chat-backend-cleanup-outbox'
+import type { OutboxEventContext } from '@/lib/core/outbox/service'
+
+const sweepImages = chatBackendCleanupOutboxHandlers[CHAT_INLINE_IMAGE_CLEANUP_EVENT]
+const chatIds = Array.from({ length: 25 }, (_, index) => `chat-${String(index).padStart(2, '0')}`)
+const prefixOf = (chatId: string) => `chat-images/${chatId}/`
+
+function outboxContext() {
+  return {
+    eventId: 'event-1',
+    eventType: CHAT_INLINE_IMAGE_CLEANUP_EVENT,
+    attempts: 0,
+    maxAttempts: 10,
+    signal: new AbortController().signal,
+    checkpointPayload: vi.fn(async (_patch: Record<string, unknown>) => {}),
+  } satisfies OutboxEventContext
+}
+
+const sweptPrefixes = () => mockSweepStoragePrefix.mock.calls.map(([prefix]) => prefix)
+
+describe('chat inline image cleanup', () => {
+  beforeEach(() => {
+    resetDbChainMock()
+    deadline.near = false
+    mockSweepStoragePrefix.mockReset().mockResolvedValue('done')
+  })
+
+  it('sweeps only the chats whose rows are gone', async () => {
+    queueTableRows(copilotChats, [{ id: chatIds[1] }])
+
+    await expect(sweepImages({ chatIds: chatIds.slice(0, 3) }, outboxContext())).resolves.toBe(
+      undefined
+    )
+    expect(sweptPrefixes().sort()).toEqual([prefixOf(chatIds[0]), prefixOf(chatIds[2])])
+  })
+
+  it('sweeps several chats at once, never more than ten', async () => {
+    queueTableRows(copilotChats, [])
+    let inFlight = 0
+    let peak = 0
+    mockSweepStoragePrefix.mockImplementation(async () => {
+      peak = Math.max(peak, ++inFlight)
+      await new Promise((resolve) => setImmediate(resolve))
+      inFlight--
+      return 'done'
+    })
+
+    await sweepImages({ chatIds }, outboxContext())
+
+    expect(peak).toBe(10)
+    expect(new Set(sweptPrefixes())).toEqual(new Set(chatIds.map(prefixOf)))
+  })
+
+  it('checkpoints the chats finished before the lowest unfinished one and resumes after them', async () => {
+    queueTableRows(copilotChats, [])
+    const slow = chatIds[3]
+    mockSweepStoragePrefix.mockImplementation(async (prefix: string) => {
+      if (prefix !== prefixOf(slow)) return 'done'
+      deadline.near = true
+      return 'continue'
+    })
+    const context = outboxContext()
+
+    const result = await sweepImages({ chatIds: chatIds.slice(0, 6) }, context)
+
+    expect(result).toMatchObject({ outcome: 'deferred', consumeAttempt: false })
+    expect(context.checkpointPayload).toHaveBeenCalledWith({ swept: 3 })
+
+    queueTableRows(copilotChats, [])
+    deadline.near = false
+    mockSweepStoragePrefix.mockReset().mockResolvedValue('done')
+    await expect(
+      sweepImages({ chatIds: chatIds.slice(0, 6), swept: 3 }, outboxContext())
+    ).resolves.toBe(undefined)
+    expect(sweptPrefixes().sort()).toEqual(chatIds.slice(3, 6).map(prefixOf))
+  })
+
+  it('spends an attempt when a chat’s listing stops advancing', async () => {
+    queueTableRows(copilotChats, [])
+    mockSweepStoragePrefix.mockImplementation(async (prefix: string) =>
+      prefix === prefixOf(chatIds[0]) ? 'stalled' : 'done'
+    )
+    const context = outboxContext()
+
+    const result = await sweepImages({ chatIds: chatIds.slice(0, 2) }, context)
+
+    expect(result).toMatchObject({ outcome: 'deferred' })
+    expect(result).not.toMatchObject({ consumeAttempt: false })
+    expect(context.checkpointPayload).not.toHaveBeenCalled()
+  })
+})

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

@@ -2,8 +2,14 @@ import { db } from '@sim/db'
 import { copilotChats } from '@sim/db/schema'
 import { describeError, toError } from '@sim/utils/errors'
 import { inArray } from 'drizzle-orm'
-import { chatBackendCleanupPayloadSchema } from '@/lib/cleanup/chat-backend-cleanup'
-import { CHAT_BACKEND_CLEANUP_EVENT } from '@/lib/cleanup/chat-backend-cleanup-event'
+import {
+  chatBackendCleanupPayloadSchema,
+  chatInlineImageCleanupPayloadSchema,
+} from '@/lib/cleanup/chat-backend-cleanup'
+import {
+  CHAT_BACKEND_CLEANUP_EVENT,
+  CHAT_INLINE_IMAGE_CLEANUP_EVENT,
+} from '@/lib/cleanup/chat-backend-cleanup-event'
 import { cleanupCopilotBackend, deletableChatFiles } from '@/lib/cleanup/chat-cleanup'
 import { env } from '@/lib/core/config/env'
 import {
@@ -13,80 +19,57 @@ import {
   type OutboxHandlerRegistry,
 } from '@/lib/core/outbox/service'
 import { mapWithConcurrency } from '@/lib/core/utils/concurrency'
-import {
-  INLINE_CHAT_IMAGE_PREFIX,
-  inlineImageRequestIdSchema,
-} from '@/lib/mothership/chat/inline-image-reference'
+import { inlineChatImagePrefix } from '@/lib/mothership/chat/inline-image-key'
 import { outboxDeadlineNear, sweepStoragePrefix } from '@/lib/uploads/core/prefix-sweep'
 import { deleteFile } from '@/lib/uploads/core/storage-service'
 
-const FILE_DELETE_CONCURRENCY = 10
+/** Most object deletes, or chat prefix sweeps, one handler run has in flight at once. */
+const STORAGE_CONCURRENCY = 10
+
+/** The ids among `chatIds` that still have a `copilot_chats` row. */
+async function survivingChatIds(chatIds: readonly string[]): Promise<Set<string>> {
+  if (chatIds.length === 0) return new Set()
+  const rows = await db
+    .select({ id: copilotChats.id })
+    .from(copilotChats)
+    .where(inArray(copilotChats.id, [...chatIds]))
+  return new Set(rows.map(({ id }) => id))
+}
 
 /**
- * Deletes every inline image under each deleted chat's `chat-images/<chatId>/` prefix, from the
- * `from`th chat on. `swept` is how many leading chats are done: all of them unless the deadline
- * neared (`continue`) or a listing stopped advancing (`stalled`).
+ * Sweeps `chatId`'s inline image prefix until it is empty, the deadline nears (`continue`), or a
+ * listing stops advancing (`stalled`).
  */
-async function sweepInlineImages(
-  chatIds: readonly string[],
-  deleted: ReadonlySet<string>,
-  from: number,
+async function sweepChatImages(
+  chatId: string,
   context: OutboxEventContext
-): Promise<{ swept: number; outcome: 'done' | 'continue' | 'stalled' }> {
-  for (let index = from; index < chatIds.length; index++) {
-    const chatId = chatIds[index]
-    // A chat id outside the image key grammar never had an image stored under it.
-    if (!deleted.has(chatId) || !inlineImageRequestIdSchema.safeParse(chatId).success) continue
-    for (;;) {
-      context.signal.throwIfAborted()
-      const outcome = await sweepStoragePrefix(
-        `${INLINE_CHAT_IMAGE_PREFIX}${chatId}/`,
-        'mothership',
-        context
-      )
-      if (outcome === 'done') break
-      if (outcome === 'stalled' || outboxDeadlineNear(context)) return { swept: index, outcome }
-    }
-    if (outboxDeadlineNear(context)) return { swept: index + 1, outcome: 'continue' }
+): Promise<'done' | 'continue' | 'stalled'> {
+  const prefix = inlineChatImagePrefix(chatId)
+  for (;;) {
+    context.signal.throwIfAborted()
+    if (outboxDeadlineNear(context)) return 'continue'
+    const outcome = await sweepStoragePrefix(prefix, 'mothership', context)
+    if (outcome !== 'continue') return outcome
   }
-  return { swept: chatIds.length, outcome: 'done' }
 }
 
 export const chatBackendCleanupOutboxHandlers = {
   /**
-   * Purges the inline images, copilot backend data and stored files of chats whose rows are gone.
-   * A chat whose row survived (its delete was spared, rolled back, or never happened) keeps all
-   * three. Inline images go by the chat's prefix rather than by key, so one a turn published after
-   * the delete read the transcript goes too; the sweep checkpoints how many chats it finished and
-   * continues in a later run once its deadline nears. Any failure throws so the outbox retries it:
-   * a failed backend call, a failed object delete, or chat ids queued for the backend reaching a
-   * processor without `COPILOT_API_KEY`. An object already gone counts as deleted, so a retry
-   * repeats only what is left. A stored file another chat of its scope still references, as a
-   * fork's uncopied attachment does, is kept ({@link deletableChatFiles}).
+   * Purges the copilot backend data and stored files of chats whose rows are gone; a surviving
+   * chat keeps both, and a file another chat of its scope still references is kept
+   * ({@link deletableChatFiles}). Any failure throws so the outbox retries; an object already gone
+   * counts as deleted.
    */
   [CHAT_BACKEND_CLEANUP_EVENT]: async (rawPayload, context) => {
-    const { chatIds, files, copilotBackend, imagesSwept } =
-      chatBackendCleanupPayloadSchema.parse(rawPayload)
+    const { chatIds, files } = chatBackendCleanupPayloadSchema.parse(rawPayload)
     context.signal.throwIfAborted()
-    const named = [...new Set([...chatIds, ...files.map((file) => file.chatId)])]
-    const surviving = await db
-      .select({ id: copilotChats.id })
-      .from(copilotChats)
-      .where(inArray(copilotChats.id, named))
-    const survivingIds = new Set(surviving.map(({ id }) => id))
+    const survivingIds = await survivingChatIds([
+      ...new Set([...chatIds, ...files.map((file) => file.chatId)]),
+    ])
     const deletedChatIds = chatIds.filter((id) => !survivingIds.has(id))
     const deletedFiles = files.filter((file) => !survivingIds.has(file.chatId))
 
-    if (imagesSwept < chatIds.length) {
-      const sweep = await sweepInlineImages(chatIds, new Set(deletedChatIds), imagesSwept, context)
-      if (sweep.swept > imagesSwept) await context.checkpointPayload({ imagesSwept: sweep.swept })
-      if (sweep.outcome === 'continue')
-        return continueOutboxHandler('Continuing chat inline image cleanup')
-      if (sweep.outcome === 'stalled')
-        return deferOutboxHandler('Chat inline image cleanup removed nothing from its listing')
-    }
-
-    if (copilotBackend && deletedChatIds.length > 0) {
+    if (deletedChatIds.length > 0) {
       context.signal.throwIfAborted()
       if (!env.COPILOT_API_KEY) throw new Error('Copilot cleanup is not configured')
       const result = await cleanupCopilotBackend(
@@ -99,7 +82,7 @@ export const chatBackendCleanupOutboxHandlers = {
     if (deletedFiles.length === 0) return
     const failures = await mapWithConcurrency(
       await deletableChatFiles(deletedFiles),
-      FILE_DELETE_CONCURRENCY,
+      STORAGE_CONCURRENCY,
       async ({ key, context: storageContext }) => {
         context.signal.throwIfAborted()
         try {
@@ -112,4 +95,26 @@ export const chatBackendCleanupOutboxHandlers = {
     const failure = failures.find((error) => error !== undefined)
     if (failure) throw failure
   },
+
+  /**
+   * Deletes everything under the `chat-images/<chatId>/` prefix of each chat whose row is gone, so
+   * an image a turn published after the delete read the transcript goes too. Sweeps several chats
+   * at once, checkpoints how many leading chats are finished, and continues in a later run once
+   * the deadline nears.
+   */
+  [CHAT_INLINE_IMAGE_CLEANUP_EVENT]: async (rawPayload, context) => {
+    const { chatIds, swept } = chatInlineImageCleanupPayloadSchema.parse(rawPayload)
+    context.signal.throwIfAborted()
+    const pending = chatIds.slice(swept)
+    const survivingIds = await survivingChatIds(pending)
+    const outcomes = await mapWithConcurrency(pending, STORAGE_CONCURRENCY, (chatId) =>
+      survivingIds.has(chatId) ? Promise.resolve('done' as const) : sweepChatImages(chatId, context)
+    )
+    const unfinished = outcomes.findIndex((outcome) => outcome !== 'done')
+    if (unfinished === -1) return
+    if (unfinished > 0) await context.checkpointPayload({ swept: swept + unfinished })
+    if (outcomes.includes('continue'))
+      return continueOutboxHandler('Continuing chat inline image cleanup')
+    return deferOutboxHandler('Chat inline image cleanup removed nothing from its listing')
+  },
 } satisfies OutboxHandlerRegistry