fix(cleanup): keep shared chat files, sweep inline images by prefix, and make room refetches wait

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

@@ -29,7 +29,7 @@ const roots: Root[] = []
 function mount(
   workspaceId: string,
   roomType: RoomType,
-  onChanged: () => unknown = () => {},
+  onChanged: () => Promise<unknown> = async () => {},
   dedupeKey?: string
 ): Root {
   function Probe() {
@@ -119,7 +119,7 @@ describe('useWorkspaceInvalidationRoom', () => {
     })
 
     it('delivers a burst of broadcasts once at once and once after it, not once per broadcast', async () => {
-      const onChanged = vi.fn()
+      const onChanged = vi.fn(async () => {})
       mount('ws-1', 'workspace-files', onChanged)
 
       for (let index = 0; index < 20; index++) broadcast('workspace-files-changed', 'ws-1')

apps/sim/app/workspace/[workspaceId]/hooks/use-workspace-invalidation-room.ts

@@ -36,10 +36,11 @@ interface JoinErrorPayload {
 }
 
 /**
- * Runs when the room's resources changed. A returned promise is awaited, up to
- * {@link DELIVERY_SETTLE_CAP_MS}, before the next delivery starts.
+ * Runs when the room's resources changed and returns the invalidations it started. The promise is
+ * awaited, up to {@link DELIVERY_SETTLE_CAP_MS}, before the next delivery starts, so a callback
+ * that drops it lets the next broadcast cancel and restart refetches still in flight.
  */
-type RoomChangedCallback = () => unknown
+type RoomChangedCallback = () => Promise<unknown>
 
 interface SharedRoomSubscription {
   callbacks: Map<string | symbol, Set<RoomChangedCallback>>

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

@@ -0,0 +1,89 @@
+/**
+ * @vitest-environment jsdom
+ */
+
+import { act } from 'react'
+import { QueryClient, QueryClientProvider, useQuery } from '@tanstack/react-query'
+import { createRoot, type Root } from 'react-dom/client'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { mockUseSocket } = vi.hoisted(() => ({ mockUseSocket: vi.fn() }))
+
+vi.mock('@/app/workspace/providers/socket-provider', () => ({
+  useSocket: () => mockUseSocket(),
+}))
+
+import { INVALIDATION_COALESCE_MS } from '@/app/workspace/[workspaceId]/hooks/use-workspace-invalidation-room'
+import { useWorkspaceTablesRoom } from '@/app/workspace/[workspaceId]/tables/hooks/use-workspace-tables-room'
+import { tableKeys } from '@/hooks/queries/utils/table-keys'
+
+function fakeSocket() {
+  return { connected: true, emit: vi.fn(), on: vi.fn(), off: vi.fn() }
+}
+
+let socket: ReturnType<typeof fakeSocket>
+let root: Root | null = null
+
+function broadcast(workspaceId: string) {
+  for (const [name, handler] of socket.on.mock.calls) {
+    if (name === 'workspace-tables-changed') handler({ workspaceId })
+  }
+}
+
+describe('useWorkspaceTablesRoom', () => {
+  beforeEach(() => {
+    vi.useFakeTimers()
+    socket = fakeSocket()
+    mockUseSocket.mockImplementation(() => ({ socket }))
+  })
+
+  afterEach(() => {
+    if (root) act(() => root?.unmount())
+    root = null
+    vi.useRealTimers()
+  })
+
+  it('does not restart a list refetch still in flight when the next broadcast arrives', async () => {
+    const pending: Array<() => void> = []
+    const fetchList = vi.fn(
+      () =>
+        new Promise<string[]>((resolve) => {
+          pending.push(() => resolve([]))
+        })
+    )
+    const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } })
+    queryClient.setQueryData(tableKeys.list('ws-1'), [])
+
+    function Probe() {
+      useQuery({
+        queryKey: tableKeys.list('ws-1'),
+        queryFn: fetchList,
+        staleTime: Number.POSITIVE_INFINITY,
+      })
+      useWorkspaceTablesRoom('ws-1')
+      return null
+    }
+    const container = document.createElement('div')
+    document.body.appendChild(container)
+    root = createRoot(container)
+    act(() =>
+      root?.render(
+        <QueryClientProvider client={queryClient}>
+          <Probe />
+        </QueryClientProvider>
+      )
+    )
+
+    broadcast('ws-1')
+    await vi.advanceTimersByTimeAsync(0)
+    expect(fetchList).toHaveBeenCalledTimes(1)
+
+    broadcast('ws-1')
+    await vi.advanceTimersByTimeAsync(5 * INVALIDATION_COALESCE_MS)
+    expect(fetchList).toHaveBeenCalledTimes(1)
+
+    for (const resolve of pending.splice(0)) resolve()
+    await vi.advanceTimersByTimeAsync(INVALIDATION_COALESCE_MS)
+    expect(fetchList).toHaveBeenCalledTimes(2)
+  })
+})

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

@@ -16,8 +16,10 @@ import { tableKeys } from '@/hooks/queries/utils/table-keys'
  */
 export function useWorkspaceTablesRoom(workspaceId: string): void {
   const queryClient = useQueryClient()
-  useWorkspaceInvalidationRoom(workspaceId, ROOM_TYPES.WORKSPACE_TABLES, () => {
-    queryClient.invalidateQueries({ queryKey: tableKeys.lists() })
-    queryClient.invalidateQueries({ queryKey: folderKeys.resource('table') })
-  })
+  useWorkspaceInvalidationRoom(workspaceId, ROOM_TYPES.WORKSPACE_TABLES, () =>
+    Promise.all([
+      queryClient.invalidateQueries({ queryKey: tableKeys.lists() }),
+      queryClient.invalidateQueries({ queryKey: folderKeys.resource('table') }),
+    ])
+  )
 }

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/hooks/use-resource-folders.ts

@@ -77,7 +77,8 @@ export function useResourceFolders(
   useWorkspaceInvalidationRoom(
     isFileResource ? (workspaceId ?? '') : '',
     ROOM_TYPES.WORKSPACE_FILES,
-    () => (workspaceId ? invalidateWorkspaceFileBrowsers(queryClient, workspaceId) : undefined),
+    () =>
+      workspaceId ? invalidateWorkspaceFileBrowsers(queryClient, workspaceId) : Promise.resolve(),
     WORKSPACE_FILE_BROWSER_INVALIDATION_KEY
   )
 

apps/sim/app/workspace/[workspaceId]/w/components/sidebar/hooks/use-workspace-workflows-room.ts

@@ -21,9 +21,11 @@ import { invalidateWorkflowLists } from '@/hooks/queries/utils/invalidate-workfl
  */
 export function useWorkspaceWorkflowsRoom(workspaceId: string): void {
   const queryClient = useQueryClient()
-  useWorkspaceInvalidationRoom(workspaceId, ROOM_TYPES.WORKSPACE_WORKFLOWS, () => {
-    invalidateWorkflowLists(queryClient, workspaceId, ['active', 'archived'])
-    queryClient.invalidateQueries({ queryKey: folderKeys.list(workspaceId, 'active') })
-    queryClient.invalidateQueries({ queryKey: folderKeys.list(workspaceId, 'archived') })
-  })
+  useWorkspaceInvalidationRoom(workspaceId, ROOM_TYPES.WORKSPACE_WORKFLOWS, () =>
+    Promise.all([
+      invalidateWorkflowLists(queryClient, workspaceId, ['active', 'archived']),
+      queryClient.invalidateQueries({ queryKey: folderKeys.list(workspaceId, 'active') }),
+      queryClient.invalidateQueries({ queryKey: folderKeys.list(workspaceId, 'archived') }),
+    ])
+  )
 }

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

@@ -14,14 +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 a chat or attachment written between
- * this read and that delete is not covered; 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,124 @@
+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 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'
+      // 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'
+    })
+    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('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) =>
+      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()
+  })
+
+  it('spends an attempt when the lowest unfinished chat stalls, though a later one continues', async () => {
+    queueTableRows(copilotChats, [])
+    mockSweepStoragePrefix.mockImplementation(async (prefix: string) => {
+      if (prefix === prefixOf(chatIds[0])) return 'stalled'
+      deadline.near = true
+      return 'continue'
+    })
+    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()
+  })
+})