11import { dbFor } from '@sim/db'
2- import { copilotChats , copilotMessages , workspaceFiles } from '@sim/db/schema'
2+ import {
3+ copilotChats ,
4+ copilotMessages ,
5+ workspaceFile ,
6+ workspaceFiles ,
7+ workspaceFileVersion ,
8+ } from '@sim/db/schema'
39import { chunkArray } from '@sim/utils/helpers'
410import { isRecordLike } from '@sim/utils/object'
511import { and , inArray , isNull , sql } from 'drizzle-orm'
@@ -86,11 +92,30 @@ function fileAttachmentKeys(value: unknown): string[] {
8692 )
8793}
8894
95+ /**
96+ * The keys among `keys` that no workspace file row, legacy file row or file version names, live
97+ * or soft-deleted, in one query.
98+ */
99+ async function keysWithoutFileRow ( executor : DbOrTx , keys : readonly string [ ] ) : Promise < string [ ] > {
100+ if ( keys . length === 0 ) return [ ]
101+ const rows = await executor . execute < { key : string } > ( sql `
102+ SELECT candidate AS key
103+ FROM jsonb_array_elements_text(${ JSON . stringify ( keys ) } ::jsonb) AS candidate
104+ WHERE NOT EXISTS (SELECT 1 FROM ${ workspaceFiles } WHERE ${ workspaceFiles . key } = candidate)
105+ AND NOT EXISTS (SELECT 1 FROM ${ workspaceFile } WHERE ${ workspaceFile . key } = candidate)
106+ AND NOT EXISTS (
107+ SELECT 1 FROM ${ workspaceFileVersion } WHERE ${ workspaceFileVersion . key } = candidate
108+ )
109+ ` )
110+ return rows . map ( ( row ) => row . key )
111+ }
112+
89113/**
90114 * Collect all file storage keys for the given chat IDs from two sources:
91115 * 1. workspaceFiles rows with chatId FK (chat-scoped contexts only)
92- * 2. fileAttachments[].key inside each copilot_messages.content: copilot keys, and
93- * organization attachments (`assistant/<orgId>/…`) under their own context
116+ * 2. fileAttachments[].key inside each copilot_messages.content: copilot keys, organization
117+ * attachments (`assistant/<orgId>/…`) under their own context, and `workspace/<ws>/…` uploads
118+ * no file row names at all
94119 *
95120 * Each ref carries the scope whose chats may still reference it (see {@link ChatFileRef}).
96121 *
@@ -102,7 +127,10 @@ function fileAttachmentKeys(value: unknown): string[] {
102127 * file lifecycle), never by the message: a fork carries the same attachment when its copy
103128 * failed or its file was deleted, and the copilot bucket falls back to the workspace bucket
104129 * on GCS (and may be configured to it on S3), so deleting such a key as copilot storage
105- * could delete a file another chat or the workspace still uses.
130+ * could delete a file another chat or the workspace still uses. Only once no row names the key
131+ * at all (a soft-deleted row can still be restored) does the message own it: a fork kept the
132+ * upload past its source chat, whose row cascaded away, so the last chat attaching it collects
133+ * it, and the reference check keeps it while another chat of that workspace still attaches it.
106134 *
107135 * Reads through the cleanup pool unless `executor` is given, which a deletion passes to read the
108136 * keys inside its transaction once it holds the chat rows, so no row naming a key can commit
@@ -177,9 +205,18 @@ export async function collectChatFiles(
177205 }
178206 }
179207
208+ const unownedCandidates = new Map < string , { chatId : string ; workspaceId : string } > ( )
180209 for ( const row of messageRows ) {
181210 for ( const key of fileAttachmentKeys ( row . fileAttachments ) ) {
182211 if ( seen . has ( key ) ) continue
212+ const uploadWorkspaceId = key . startsWith ( 'workspace/' )
213+ ? extractWorkspaceIdFromStorageKey ( key )
214+ : null
215+ if ( uploadWorkspaceId ) {
216+ if ( ! unownedCandidates . has ( key ) )
217+ unownedCandidates . set ( key , { chatId : row . chatId , workspaceId : uploadWorkspaceId } )
218+ continue
219+ }
183220 const organizationId = organizationAttachmentOwner ( key )
184221 if ( organizationId ) {
185222 seen . add ( key )
@@ -195,6 +232,18 @@ export async function collectChatFiles(
195232 }
196233 }
197234 }
235+
236+ for ( const key of await keysWithoutFileRow ( executor , [ ...unownedCandidates . keys ( ) ] ) ) {
237+ const candidate = unownedCandidates . get ( key )
238+ if ( ! candidate || seen . has ( key ) ) continue
239+ seen . add ( key )
240+ files . push ( {
241+ key,
242+ context : 'mothership' ,
243+ chatId : candidate . chatId ,
244+ scope : { kind : 'workspace' , id : candidate . workspaceId } ,
245+ } )
246+ }
198247 }
199248
200249 return files
0 commit comments