fix(cleanup): delete an organization attachment once no chat references it

Organization Chat attachments (assistant/ keys) have no workspace_files row and
a fork carries the same key, so the purge now deletes one under its own
storage context only after confirming no remaining chat of that organization
still references it.

Author: waleedlatif1

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

@@ -15,15 +15,25 @@ function attachment(key: string) {
   return { id: 'wf_x', key, filename: 'x.png', media_type: 'image/png', size: 1 }
 }
 
-/** Purges one deleted chat whose message rows are `messages`; returns the deleted keys by context. */
-async function purge(messages: Record<string, unknown>[]) {
+/**
+ * Purges one deleted chat whose message rows are `messages`, while `remaining` are the message
+ * rows other chats of the organization still hold; returns the deleted keys by context.
+ */
+async function purge(
+  messages: Record<string, unknown>[],
+  remaining: Record<string, unknown>[] = []
+) {
   queueTableRows(workspaceFiles, [])
   queueTableRows(
     copilotMessages,
     messages.map((content) => ({ chatId, content }))
   )
   const cleanup = await prepareChatCleanup([chatId], 'test')
   queueTableRows(copilotChats, [])
+  queueTableRows(
+    copilotMessages,
+    remaining.map((content) => ({ content }))
+  )
   await cleanup.execute()
   return Object.fromEntries(
     storageServiceMockFns.mockDeleteFiles.mock.calls.map(([keys, context]) => [context, keys])
@@ -38,7 +48,7 @@ describe('chat purge storage', () => {
     storageServiceMockFns.mockDeleteFiles.mockResolvedValue({ deleted: 1, failed: [] })
   })
 
-  it('deletes only copilot-storage attachment keys as copilot storage', async () => {
+  it('never deletes a workspace attachment as copilot storage', async () => {
     // A fork carries its source's key when the file's copy failed or the file was deleted, and
     // the copilot bucket can be the workspace bucket, so a workspace key here is another chat's file.
     const deleted = await purge([
@@ -47,14 +57,29 @@ describe('chat purge storage', () => {
         content: 'look',
         fileAttachments: [
           attachment('workspace/ws-1/1700-abc-shared.png'),
-          attachment('assistant/org-1/user-1/u1/shared.png'),
           attachment('copilot/1234/legacy.png'),
         ],
       },
     ])
     expect(deleted).toEqual({ copilot: ['copilot/1234/legacy.png'] })
   })
 
+  it('deletes an organization attachment only once no remaining chat references it', async () => {
+    const shared = 'assistant/org-1/user-1/u1/shared.png'
+    const own = 'assistant/org-1/user-1/u2/own.png'
+    const message = {
+      role: 'user',
+      content: 'look',
+      fileAttachments: [attachment(shared), attachment(own)],
+    }
+    // A fork of this chat still holds the shared upload.
+    const deleted = await purge(
+      [message],
+      [{ role: 'user', content: 'fork', fileAttachments: [attachment(shared)] }]
+    )
+    expect(deleted).toEqual({ mothership: [own] })
+  })
+
   it('deletes the inline images an assistant message published', async () => {
     const deleted = await purge([
       {

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

@@ -3,7 +3,7 @@ import { copilotChats, copilotMessages, workspaceFiles } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { chunkArray } from '@sim/utils/helpers'
 import { isRecordLike } from '@sim/utils/object'
-import { and, inArray, isNull } from 'drizzle-orm'
+import { and, eq, inArray, isNull, or, sql } from 'drizzle-orm'
 import { env } from '@/lib/core/config/env'
 import {
   inlineChatImageKey,
@@ -36,6 +36,19 @@ interface FileRef {
   key: string
   context: ChatScopedContext
   chatId: string
+  /**
+   * An organization Chat attachment (`assistant/<orgId>/…`): no `workspace_files` row owns it,
+   * and a fork carries the same key, so it is deleted only once no remaining chat of that
+   * organization references it.
+   */
+  organizationId?: string
+}
+
+/** The organization an `assistant/<orgId>/…` attachment key was uploaded under. */
+function organizationAttachmentOwner(key: string): string | undefined {
+  if (tryInferContextFromKey(key) !== 'mothership') return undefined
+  const [, organizationId] = key.split('/')
+  return organizationId || undefined
 }
 
 /**
@@ -71,15 +84,15 @@ function inlineChatImageKeys(chatId: string, content: Record<string, unknown>):
 /**
  * Collect all file storage keys for the given chat IDs from three sources:
  * 1. workspaceFiles rows with chatId FK (chat-scoped contexts only)
- * 2. fileAttachments[].key inside each copilot_messages.content, for copilot-storage keys only
+ * 2. fileAttachments[].key inside each copilot_messages.content: copilot keys, and
+ *    organization attachments under their own context (see {@link FileRef.organizationId})
  * 3. the chat-scoped inline images each assistant message published
  *
- * An attachment whose key belongs to another storage context is owned by its
- * `workspace_files` row (source 1, or the workspace file lifecycle), never by the message:
- * a fork carries the same attachment when its copy failed or its file was deleted, and the
- * copilot bucket falls back to the workspace bucket on GCS (and may be configured to it on
- * S3), so deleting such a key as copilot storage could delete a file another chat or the
- * workspace still uses.
+ * A workspace attachment is owned by its `workspace_files` row (source 1, or the workspace
+ * file lifecycle), never by the message: a fork carries the same attachment when its copy
+ * failed or its file was deleted, and the copilot bucket falls back to the workspace bucket
+ * on GCS (and may be configured to it on S3), so deleting such a key as copilot storage
+ * could delete a file another chat or the workspace still uses.
  */
 export async function collectChatFiles(chatIds: string[]): Promise<FileRef[]> {
   const files: FileRef[] = []
@@ -136,8 +149,12 @@ export async function collectChatFiles(chatIds: string[]): Promise<FileRef[]> {
           (attachment as Record<string, unknown>).key
         ) {
           const key = (attachment as Record<string, unknown>).key as string
-          if (tryInferContextFromKey(key) !== 'copilot') continue
-          if (!seen.has(key)) {
+          if (seen.has(key)) continue
+          const organizationId = organizationAttachmentOwner(key)
+          if (organizationId) {
+            seen.add(key)
+            files.push({ key, context: 'mothership', chatId: row.chatId, organizationId })
+          } else if (tryInferContextFromKey(key) === 'copilot') {
             seen.add(key)
             files.push({ key, context: 'copilot', chatId: row.chatId })
           }
@@ -149,6 +166,55 @@ export async function collectChatFiles(chatIds: string[]): Promise<FileRef[]> {
   return files
 }
 
+/**
+ * Organization attachment keys that a remaining chat still references, so they outlive the
+ * chats being purged. Runs after the caller deleted those chats' rows, so any match is
+ * another chat: a fork, or a chat that is only soft-deleted and may be restored.
+ */
+async function organizationAttachmentsStillReferenced(files: FileRef[]): Promise<Set<string>> {
+  const keysByOrganization = new Map<string, string[]>()
+  for (const file of files) {
+    if (!file.organizationId) continue
+    const keys = keysByOrganization.get(file.organizationId)
+    if (keys) keys.push(file.key)
+    else keysByOrganization.set(file.organizationId, [file.key])
+  }
+  const referenced = new Set<string>()
+  for (const [organizationId, keys] of keysByOrganization) {
+    for (const chunk of chunkArray(keys, CHAT_FILE_COLLECT_CHUNK_SIZE)) {
+      const wanted = new Set(chunk)
+      const rows = await cleanupDb
+        .select({ content: copilotMessages.content })
+        .from(copilotMessages)
+        .innerJoin(copilotChats, eq(copilotChats.id, copilotMessages.chatId))
+        .where(
+          and(
+            eq(copilotChats.organizationId, organizationId),
+            or(
+              ...chunk.map(
+                (key) =>
+                  sql`${copilotMessages.content} @> ${JSON.stringify({ fileAttachments: [{ key }] })}::jsonb`
+              )
+            )
+          )
+        )
+      for (const { content } of rows) {
+        const attachments = isRecordLike(content) ? content.fileAttachments : undefined
+        if (!Array.isArray(attachments)) continue
+        for (const attachment of attachments) {
+          if (
+            isRecordLike(attachment) &&
+            typeof attachment.key === 'string' &&
+            wanted.has(attachment.key)
+          )
+            referenced.add(attachment.key)
+        }
+      }
+    }
+  }
+  return referenced
+}
+
 /** Groups files by storage context so each context can use one batch DELETE call. */
 export async function deleteStorageFiles(
   files: FileRef[],
@@ -271,7 +337,12 @@ export async function prepareChatCleanup(
         )
       }
       const confirmedChatIds = chatIds.filter((id) => !survivors.has(id))
-      const confirmedFiles = files.filter((file) => !survivors.has(file.chatId))
+      const sharedKeys = await organizationAttachmentsStillReferenced(
+        files.filter((file) => !survivors.has(file.chatId))
+      )
+      const confirmedFiles = files.filter(
+        (file) => !survivors.has(file.chatId) && !sharedKeys.has(file.key)
+      )
 
       // Call copilot backend
       if (confirmedChatIds.length > 0) {