chore(cleanup): tighten the cleanup tests and comments

Drops tests the integration suites already cover, moves the non-string
attachment key case into the PostgreSQL collect test, and makes the image
checkpoint test finish chats after the slow one first so it proves the
lowest-unfinished checkpoint. The tables room test reads its key from
tableKeys.list. inlineChatImagePrefix returns null for an id outside the
key grammar, which the event schema admits, so the image sweep skips and
logs such an id instead of failing the event on every retry.

Author: waleedlatif1

apps/sim/app/workspace/[workspaceId]/tables/hooks/use-workspace-tables-room.test.tsx

@@ -43,12 +43,6 @@ describe('useWorkspaceTablesRoom', () => {
     vi.useRealTimers()
   })
 
-  /*
-   * The callback used to fire both invalidations and return nothing, so the room's delivery never
-   * waited: a broadcast arriving while the list refetch was still loading invalidated it again,
-   * which cancels and restarts the in-flight fetch. Under a burst of table writes a slow list never
-   * finished loading.
-   */
   it('does not restart a list refetch still in flight when the next broadcast arrives', async () => {
     const pending: Array<() => void> = []
     const fetchList = vi.fn(
@@ -58,11 +52,11 @@ describe('useWorkspaceTablesRoom', () => {
         })
     )
     const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } })
-    queryClient.setQueryData([...tableKeys.lists(), 'ws-1'], [])
+    queryClient.setQueryData(tableKeys.list('ws-1'), [])
 
     function Probe() {
       useQuery({
-        queryKey: [...tableKeys.lists(), 'ws-1'],
+        queryKey: tableKeys.list('ws-1'),
         queryFn: fetchList,
         staleTime: Number.POSITIVE_INFINITY,
       })

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

@@ -40,15 +40,6 @@ describe('chat inline image cleanup', () => {
     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
@@ -71,6 +62,8 @@ describe('chat inline image cleanup', () => {
     const slow = chatIds[3]
     mockSweepStoragePrefix.mockImplementation(async (prefix: string) => {
       if (prefix !== prefixOf(slow)) return 'done'
+      // Finishes after the chats behind it, so only "lowest unfinished" can yield 3 here.
+      await new Promise((resolve) => setImmediate(resolve))
       deadline.near = true
       return 'continue'
     })
@@ -90,6 +83,15 @@ describe('chat inline image cleanup', () => {
     expect(sweptPrefixes().sort()).toEqual(chatIds.slice(3, 6).map(prefixOf))
   })
 
+  it('skips a chat id no image key can hold instead of failing the event on every retry', async () => {
+    queueTableRows(copilotChats, [])
+
+    await expect(
+      sweepImages({ chatIds: ['../escape', chatIds[0]] }, outboxContext())
+    ).resolves.toBe(undefined)
+    expect(sweptPrefixes()).toEqual([prefixOf(chatIds[0])])
+  })
+
   it('spends an attempt when a chat’s listing stops advancing', async () => {
     queueTableRows(copilotChats, [])
     mockSweepStoragePrefix.mockImplementation(async (prefix: string) =>

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

@@ -1,5 +1,6 @@
 import { db } from '@sim/db'
 import { copilotChats } from '@sim/db/schema'
+import { createLogger } from '@sim/logger'
 import { describeError, toError } from '@sim/utils/errors'
 import { inArray } from 'drizzle-orm'
 import {
@@ -23,6 +24,8 @@ 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 logger = createLogger('ChatBackendCleanupOutbox')
+
 /** Most object deletes, or chat prefix sweeps, one handler run has in flight at once. */
 const STORAGE_CONCURRENCY = 10
 
@@ -45,6 +48,13 @@ async function sweepChatImages(
   context: OutboxEventContext
 ): Promise<'done' | 'continue' | 'stalled'> {
   const prefix = inlineChatImagePrefix(chatId)
+  if (!prefix) {
+    logger.warn('Skipping an inline image sweep for a chat id no image key can hold', {
+      eventId: context.eventId,
+      chatId,
+    })
+    return 'done'
+  }
   for (;;) {
     context.signal.throwIfAborted()
     if (outboxDeadlineNear(context)) return 'continue'

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

@@ -364,15 +364,14 @@ describe('cleanup of a permanently deleted chat', () => {
     expect(await objectExists(ownKey)).toBe(false)
   })
 
-  it('keeps what an event queued before refs carried their scope cannot place', async () => {
+  it('places a legacy upload ref by its workspace key, keeping one a remaining chat attaches', async () => {
     const ownerId = await seedUser()
     const workspaceId = await seedWorkspace(ownerId)
     const deletedChat = await seedChat(ownerId, workspaceId)
     const remainingChat = await seedChat(ownerId, workspaceId)
     const referenced = `workspace/${workspaceId}/1700000000000-abcd1234-kept.png`
     const unreferenced = `workspace/${workspaceId}/1700000000000-abcd1234-gone.png`
-    const unplaceable = `copilot/${generateId()}/legacy.png`
-    for (const key of [referenced, unreferenced, unplaceable]) await putObject(key)
+    for (const key of [referenced, unreferenced]) await putObject(key)
     await db.insert(copilotMessages).values({
       chatId: remainingChat,
       messageId: generateId(),
@@ -392,19 +391,10 @@ describe('cleanup of a permanently deleted chat', () => {
         })),
       },
     })
-    await db.insert(outboxEvent).values({
-      id: generateId(),
-      eventType: CHAT_BACKEND_CLEANUP_EVENT,
-      payload: {
-        chatIds: [],
-        files: [{ key: unplaceable, context: 'copilot', chatId: deletedChat }],
-      },
-    })
 
-    expect((await deliverCleanups()).outcomes).toEqual(['completed', 'completed'])
+    expect((await deliverCleanups()).outcomes).toEqual(['completed'])
     expect(await objectExists(referenced)).toBe(true)
     expect(await objectExists(unreferenced)).toBe(false)
-    expect(await objectExists(unplaceable)).toBe(true)
   })
 
   it('accepts the payloads queued before refs carried their scope', async () => {

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

@@ -36,12 +36,14 @@ const chatFileSchema = z
       ])
       .nullable()
       .optional(),
-    /** Legacy: an organization attachment's scope on events queued before refs carried `scope`. */
+    /**
+     * Legacy: an organization attachment's scope on events queued before refs carried `scope`.
+     * Remove once those events have drained.
+     */
     organizationId: chatIdSchema.optional(),
   })
   .strict()
   .transform(({ organizationId, ...file }): ChatFileRef => {
-    // Legacy: remove with `organizationId` once events queued before `scope` have drained.
     if (file.scope === undefined && organizationId)
       return { ...file, scope: { kind: 'organization', id: organizationId } }
     return file

apps/sim/lib/cleanup/chat-cleanup.integration.ts

@@ -3,8 +3,8 @@ import { db, dbFor } from '@sim/db'
 import { copilotChats, copilotMessages, organization, user } from '@sim/db/schema'
 import { generateId } from '@sim/utils/id'
 import { compareStrings } from '@sim/utils/string'
-import { eq, inArray } from 'drizzle-orm'
-import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'
+import { eq } from 'drizzle-orm'
+import { afterAll, beforeAll, describe, expect, it } from 'vitest'
 import { deletableChatFiles } from '@/lib/cleanup/chat-cleanup'
 import { type ChatFileRef, collectChatFiles } from '@/lib/cleanup/chat-files'
 
@@ -13,8 +13,7 @@ function attachment(key: string) {
 }
 
 const userId = generateId()
-/** One organization with chats, and one whose only chats were the deleted ones. */
-const organizationIds = { withChats: generateId(), empty: generateId() }
+const organizationId = generateId()
 
 beforeAll(async () => {
   const now = new Date()
@@ -26,20 +25,18 @@ beforeAll(async () => {
     createdAt: now,
     updatedAt: now,
   })
-  await db.insert(organization).values(
-    Object.values(organizationIds).map((id) => ({
-      id,
-      name: 'Chat cleanup fixture',
-      slug: id,
-      createdAt: now,
-      updatedAt: now,
-    }))
-  )
+  await db.insert(organization).values({
+    id: organizationId,
+    name: 'Chat cleanup fixture',
+    slug: organizationId,
+    createdAt: now,
+    updatedAt: now,
+  })
 })
 
 afterAll(async () => {
   await db.delete(copilotChats).where(eq(copilotChats.userId, userId))
-  await db.delete(organization).where(inArray(organization.id, Object.values(organizationIds)))
+  await db.delete(organization).where(eq(organization.id, organizationId))
   await db.delete(user).where(eq(user.id, userId))
   await Promise.all([db.$client.end(), dbFor('cleanup').$client.end()])
 })
@@ -80,6 +77,7 @@ describe('chat cleanup file collection in PostgreSQL', () => {
             'not an attachment',
             null,
             { key: '' },
+            { key: 42 },
           ],
         },
         {
@@ -149,11 +147,11 @@ describe('deletable chat files against PostgreSQL message rows', () => {
     await seedChat(
       [{ role: 'user', content: 'fork', fileAttachments: [attachment(shared), 'x', { key: 7 }] }],
       [],
-      organizationIds.withChats
+      organizationId
     )
     // Attached only by a chat outside the scope, which cannot be a fork of the deleted one.
     await seedChat([{ role: 'user', content: 'other', fileAttachments: [attachment(elsewhere)] }])
-    const scope = { kind: 'organization', id: organizationIds.withChats } as const
+    const scope = { kind: 'organization', id: organizationId } as const
 
     const deletable = await deletableChatFiles([
       ref(shared, scope),
@@ -166,27 +164,14 @@ describe('deletable chat files against PostgreSQL message rows', () => {
     )
   })
 
-  it('scans no messages for a scope with no chat left', async () => {
-    const scan = vi.spyOn(dbFor('cleanup'), 'execute')
-    const key = `copilot/${generateId()}/gone.png`
-
-    const deletable = await deletableChatFiles([
-      ref(key, { kind: 'organization', id: organizationIds.empty }),
-    ])
-
-    expect(deletable.map((file) => file.key)).toEqual([key])
-    expect(scan).not.toHaveBeenCalled()
-    scan.mockRestore()
-  })
-
   it('places a legacy ref by its key, keeping one it cannot place', async () => {
-    const referenced = `assistant/${organizationIds.withChats}/${userId}/u1/kept.png`
-    const unreferenced = `assistant/${organizationIds.withChats}/${userId}/u2/gone.png`
+    const referenced = `assistant/${organizationId}/${userId}/u1/kept.png`
+    const unreferenced = `assistant/${organizationId}/${userId}/u2/gone.png`
     const unplaceable = `copilot/${generateId()}/legacy.png`
     await seedChat(
       [{ role: 'user', content: 'fork', fileAttachments: [attachment(referenced)] }],
       [],
-      organizationIds.withChats
+      organizationId
     )
 
     const deletable = await deletableChatFiles(

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

@@ -52,15 +52,4 @@ describe('chat purge storage', () => {
     ])
     expect(deleted).toEqual({ copilot: ['copilot/1234/legacy.png'] })
   })
-
-  it('skips an attachment whose key is not a string instead of failing the purge', async () => {
-    const deleted = await purge([
-      {
-        role: 'user',
-        content: 'look',
-        fileAttachments: [{ ...attachment('x'), key: 42 }, attachment('copilot/1234/kept.png')],
-      },
-    ])
-    expect(deleted).toEqual({ copilot: ['copilot/1234/kept.png'] })
-  })
 })

apps/sim/lib/cleanup/chat-files.ts

@@ -71,7 +71,8 @@ export function chatFileReferenceScope(
   const workspaceId = extractWorkspaceIdFromStorageKey(file.key)
   if (workspaceId) return { kind: 'workspace', id: workspaceId }
   // A fork copies the source's inline images under its own id, so the chat's own are unshared.
-  if (file.key.startsWith(inlineChatImagePrefix(file.chatId))) return 'unshared'
+  const ownImages = inlineChatImagePrefix(file.chatId)
+  if (ownImages && file.key.startsWith(ownImages)) return 'unshared'
   return 'unknown'
 }
 

apps/sim/lib/mothership/chat/inline-image-key.ts

@@ -8,17 +8,24 @@ import {
 import { collectMarkdownImageSources } from '@/lib/mothership/chat/markdown-images'
 import type { PersistedMessage } from '@/lib/mothership/chat/persisted-message'
 
-/** The `chat-images/<chatId>/` prefix every inline image a chat published is stored under. */
-export function inlineChatImagePrefix(chatId: string): string {
-  inlineImageRequestIdSchema.parse(chatId)
-  return `${INLINE_CHAT_IMAGE_PREFIX}${chatId}/`
+function chatImagePrefix(validChatId: string): string {
+  return `${INLINE_CHAT_IMAGE_PREFIX}${validChatId}/`
+}
+
+/**
+ * The `chat-images/<chatId>/` prefix every inline image a chat published is stored under, or null
+ * for an id outside the key grammar, under which no image can have been stored.
+ */
+export function inlineChatImagePrefix(chatId: string): string | null {
+  return inlineImageRequestIdSchema.safeParse(chatId).success ? chatImagePrefix(chatId) : null
 }
 
 /** Stable per-message source identity makes replay independent of sandbox lifetime. */
 export function inlineChatImageKey(chatId: string, requestId: string, reference: string): string {
+  const prefix = chatImagePrefix(inlineImageRequestIdSchema.parse(chatId))
   inlineImageRequestIdSchema.parse(requestId)
   const digest = createHash('sha256').update(normalizeInlineFileReference(reference)).digest('hex')
-  return `${inlineChatImagePrefix(chatId)}${requestId}/${digest}.webp`
+  return `${prefix}${requestId}/${digest}.webp`
 }
 
 /**

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

@@ -17,7 +17,6 @@ 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' } })