improvement(db): name chat uploads in one indexed read and index workspace payer lookups

apps/sim/lib/uploads/contexts/workspace/__integration__/file-names.integration.ts

@@ -671,4 +671,137 @@ describe('workspace file names in PostgreSQL', () => {
       )?.name
     ).toBe('back\\slash.png')
   })
+
+  /**
+   * Counts every display-name write into the chat, including ones a unique violation rolls
+   * back: `nextval` is not transactional.
+   */
+  async function countChatNameWrites<T>(chatId: string, run: () => Promise<T>) {
+    const suffix = generateId().replaceAll('-', '')
+    const counterName = `chat_name_writes_${suffix}`
+    const counter = sql.identifier(counterName)
+    const fn = sql.identifier(`count_chat_name_write_${suffix}`)
+    await db.execute(sql`CREATE SEQUENCE ${counter}`)
+    await db.execute(sql`
+      CREATE FUNCTION ${fn}() RETURNS trigger LANGUAGE plpgsql AS $$
+      BEGIN
+        IF NEW.chat_id::text = TG_ARGV[0] THEN
+          PERFORM nextval(TG_ARGV[1]);
+        END IF;
+        RETURN NEW;
+      END $$
+    `)
+    await db.execute(sql`
+      CREATE TRIGGER ${fn} BEFORE INSERT OR UPDATE OF display_name ON ${workspaceFiles}
+      FOR EACH ROW EXECUTE FUNCTION ${fn}(${sql.raw(`'${chatId}'`)}, ${sql.raw(`'${counterName}'`)})
+    `)
+    try {
+      const result = await run()
+      const [{ writes }] = await db.execute<{ writes: number }>(
+        sql`SELECT CASE WHEN is_called THEN last_value ELSE 0 END::int AS writes FROM ${counter}`
+      )
+      return { result, writes }
+    } finally {
+      await db.execute(sql`DROP TRIGGER ${fn} ON ${workspaceFiles}`)
+      await db.execute(sql`DROP FUNCTION ${fn}()`)
+      await db.execute(sql`DROP SEQUENCE ${counter}`)
+    }
+  }
+
+  /** A chat upload the storage service registered before the chat message linked it. */
+  async function seedUnlinkedChatUpload(workspaceId: string, userId: string, name: string) {
+    const key = generateWorkspaceFileKey(workspaceId, name)
+    await db.insert(workspaceFiles).values({
+      id: `wf_${generateId()}`,
+      key,
+      userId,
+      workspaceId,
+      context: 'mothership',
+      originalName: name,
+      contentType: 'image/png',
+      sizeBytes: 10,
+    })
+    return key
+  }
+
+  function linkUpload(workspaceId: string, userId: string, chatId: string, key: string) {
+    return trackChatUpload(workspaceId, userId, chatId, key, 'image.png', 'image/png', 10)
+  }
+
+  it('names many same-named attachments with one write each', async () => {
+    const fixture = await seedWorkspace()
+    const chatId = await seedChat(fixture.workspaceId, fixture.aliceId)
+    const unlinked: string[] = []
+    for (let n = 0; n < 10; n++) {
+      unlinked.push(await seedUnlinkedChatUpload(fixture.workspaceId, fixture.aliceId, 'image.png'))
+    }
+
+    const { result, writes } = await countChatNameWrites(chatId, async () => {
+      const names = []
+      for (const key of unlinked) {
+        names.push(
+          (await linkUpload(fixture.workspaceId, fixture.aliceId, chatId, key)).displayName
+        )
+      }
+      for (let n = 0; n < 10; n++) {
+        names.push(
+          (await trackUpload(fixture.workspaceId, fixture.aliceId, chatId, 'image.png')).displayName
+        )
+      }
+      return names
+    })
+
+    expect(result).toEqual([
+      'image.png',
+      ...Array.from({ length: 19 }, (_, n) => `image (${n + 2}).png`),
+    ])
+    expect(writes).toBe(20)
+  })
+
+  it('keeps a re-linked upload on its name and treats soft-deleted siblings as taken', async () => {
+    const fixture = await seedWorkspace()
+    const chatId = await seedChat(fixture.workspaceId, fixture.aliceId)
+    const first = await trackUpload(fixture.workspaceId, fixture.aliceId, chatId, 'image.png')
+    const key = await seedUnlinkedChatUpload(fixture.workspaceId, fixture.aliceId, 'image.png')
+
+    expect((await linkUpload(fixture.workspaceId, fixture.aliceId, chatId, key)).displayName).toBe(
+      'image (2).png'
+    )
+    await db
+      .update(workspaceFiles)
+      .set({ deletedAt: new Date() })
+      .where(eq(workspaceFiles.id, first.id))
+    const { result, writes } = await countChatNameWrites(chatId, () =>
+      linkUpload(fixture.workspaceId, fixture.aliceId, chatId, key)
+    )
+    expect(result.displayName).toBe('image (2).png')
+    expect(writes).toBe(1)
+    expect(
+      (await trackUpload(fixture.workspaceId, fixture.aliceId, chatId, 'image.png')).displayName
+    ).toBe('image (3).png')
+  })
+
+  it.each([
+    ['notes 100%.png', 'notes 100% (2).png', 'notes 100% (3).png'],
+    ['a_b.png', 'a_b (2).png', 'a_b (3).png'],
+    ['report (2).pdf', 'report (2) (2).pdf', 'report (2) (3).pdf'],
+    ['README', 'README (2)', 'README (3)'],
+    ['archive.tar.gz', 'archive.tar (2).gz', 'archive.tar (3).gz'],
+    ['back\\slash.png', 'back\\slash (2).png', 'back\\slash (3).png'],
+    ['.env', '.env (2)', '.env (3)'],
+  ])('suffixes %s literally with one write per upload', async (...expected) => {
+    const fixture = await seedWorkspace()
+    const chatId = await seedChat(fixture.workspaceId, fixture.aliceId)
+    const { result, writes } = await countChatNameWrites(chatId, async () => {
+      const names = []
+      for (let n = 0; n < 3; n++) {
+        names.push(
+          (await trackUpload(fixture.workspaceId, fixture.aliceId, chatId, expected[0])).displayName
+        )
+      }
+      return names
+    })
+    expect(result).toEqual(expected)
+    expect(writes).toBe(3)
+  })
 })

apps/sim/lib/uploads/contexts/workspace/file-manager.ts

@@ -24,7 +24,8 @@ import {
 } from '@sim/utils/errors'
 import { generateShortId } from '@sim/utils/id'
 import { omit } from '@sim/utils/object'
-import { and, desc, eq, inArray, isNotNull, isNull, or, type SQL, sql } from 'drizzle-orm'
+import { escapeLikePattern } from '@sim/utils/string'
+import { and, desc, eq, inArray, isNotNull, isNull, ne, or, type SQL, sql } from 'drizzle-orm'
 import type { ShareRecord } from '@/lib/api/contracts/public-shares'
 import type { V2FileSortBy } from '@/lib/api/contracts/v2/files'
 import {
@@ -453,16 +454,21 @@ async function cleanupWorkspaceStorageObject(key: string, reason: string): Promi
   }
 }
 
+/** Splits a name where {@link withCopySuffix} inserts its ` (n)`: before the last extension. */
+function copySuffixParts(fileName: string): { stem: string; extension: string } {
+  const lastDot = fileName.lastIndexOf('.')
+  const hasExtension = lastDot > 0 && lastDot < fileName.length - 1
+  return hasExtension
+    ? { stem: fileName.slice(0, lastDot), extension: fileName.slice(lastDot) }
+    : { stem: fileName, extension: '' }
+}
+
 /**
  * Inserts ` (n)` before the last extension (e.g. `a.pdf` → `a (1).pdf`), or appends for names without.
  */
 function withCopySuffix(fileName: string, n: number | string): string {
-  const lastDot = fileName.lastIndexOf('.')
-  const hasExtension = lastDot > 0 && lastDot < fileName.length - 1
-  if (hasExtension) {
-    return `${fileName.slice(0, lastDot)} (${n})${fileName.slice(lastDot)}`
-  }
-  return `${fileName} (${n})`
+  const { stem, extension } = copySuffixParts(fileName)
+  return `${stem} (${n})${extension}`
 }
 
 /**
@@ -1239,6 +1245,45 @@ const MAX_CHAT_DISPLAY_NAME_RETRIES = 1000
 /** Postgres constraint name for the partial unique index on `(chat_id, display_name)`. */
 export const CHAT_DISPLAY_NAME_INDEX = 'workspace_files_chat_display_name_unique'
 
+/**
+ * The first `suffixedName(fileName, n)` free in the chat under {@link CHAT_DISPLAY_NAME_INDEX}.
+ *
+ * Reads with the index's own predicate, so soft-deleted rows keep their slot, and skips the row
+ * being re-linked so it keeps its current name. One indexed read replaces probing each
+ * candidate with a write; `knownTaken` carries names that lost a race since the read.
+ */
+async function nextFreeChatDisplayName(
+  tx: DbTransaction,
+  chatId: string,
+  fileName: string,
+  relinkedFileId: string | null,
+  knownTaken: ReadonlySet<string>
+): Promise<string> {
+  const { stem, extension } = copySuffixParts(fileName)
+  const suffixedPattern = `${escapeLikePattern(stem)} (%)${escapeLikePattern(extension)}`
+  const rows = await tx
+    .select({ displayName: workspaceFiles.displayName })
+    .from(workspaceFiles)
+    .where(
+      and(
+        eq(workspaceFiles.chatId, chatId),
+        eq(workspaceFiles.context, 'mothership'),
+        or(
+          eq(workspaceFiles.displayName, fileName),
+          sql`${workspaceFiles.displayName} LIKE ${suffixedPattern} ESCAPE '\\'`
+        ),
+        relinkedFileId ? ne(workspaceFiles.id, relinkedFileId) : undefined
+      )
+    )
+  const taken = new Set(knownTaken)
+  for (const row of rows) if (row.displayName) taken.add(row.displayName)
+  for (let n = 1; n <= MAX_CHAT_DISPLAY_NAME_RETRIES; n++) {
+    const candidate = suffixedName(fileName, n)
+    if (!taken.has(candidate)) return candidate
+  }
+  throw new FileConflictError(fileName)
+}
+
 /**
  * Raised when a caller-supplied storage key may not be bound to a chat upload —
  * it addresses another workspace, or it already has a `workspace_files` record
@@ -1368,12 +1413,14 @@ export async function trackChatUpload(
     }
   }
 
-  for (let n = 1; n <= MAX_CHAT_DISPLAY_NAME_RETRIES; n++) {
-    const candidate = suffixedName(fileName, n)
+  const lostRaces = new Set<string>()
+  for (let attempt = 1; attempt <= MAX_CHAT_DISPLAY_NAME_RETRIES; attempt++) {
+    let candidate = fileName
     try {
       if (claimable.kind === 'update') {
         await db.transaction(async (tx) => {
           await lockWorkspaceProject(tx, workspaceId)
+          candidate = await nextFreeChatDisplayName(tx, chatId, fileName, claimable.id, lostRaces)
           const updated = await tx
             .update(workspaceFiles)
             .set({
@@ -1420,6 +1467,7 @@ export async function trackChatUpload(
 
       await db.transaction(async (tx) => {
         await lockWorkspaceProject(tx, workspaceId)
+        candidate = await nextFreeChatDisplayName(tx, chatId, fileName, null, lostRaces)
         const [inserted] = await tx
           .insert(workspaceFiles)
           .values({
@@ -1452,8 +1500,9 @@ export async function trackChatUpload(
         getPostgresConstraintName(error) === CHAT_DISPLAY_NAME_INDEX
       ) {
         logger.warn(
-          `Chat upload displayName collision on attempt ${n} for "${candidate}" in chat ${chatId}, retrying with suffix`
+          `Chat upload displayName collision on attempt ${attempt} for "${candidate}" in chat ${chatId}, retrying with suffix`
         )
+        lostRaces.add(candidate)
         continue
       }
       throw error

packages/db/migrations/0435_workspace_billed_account_user_index.sql

@@ -0,0 +1,26 @@
+-- Payer lookups for a user's workspaces; the user FK's delete check scans through it too.
+DO $$
+BEGIN
+  IF EXISTS (
+    SELECT 1 FROM pg_index WHERE indexrelid = to_regclass('workspace_billed_account_user_id_idx_failed')
+      AND (indisvalid OR indrelid <> to_regclass('workspace'))
+  ) THEN
+    RAISE EXCEPTION 'Refusing to drop unexpected recovery index workspace_billed_account_user_id_idx_failed';
+  END IF;
+  IF EXISTS (
+    SELECT 1 FROM pg_index WHERE indexrelid = to_regclass('workspace_billed_account_user_id_idx')
+      AND indrelid = to_regclass('workspace') AND NOT indisvalid
+  ) THEN
+    ALTER INDEX "workspace_billed_account_user_id_idx" RENAME TO "workspace_billed_account_user_id_idx_failed";
+  END IF;
+END $$;
+--> statement-breakpoint
+COMMIT;
+--> statement-breakpoint
+SET lock_timeout = 0;
+--> statement-breakpoint
+DROP INDEX CONCURRENTLY IF EXISTS "workspace_billed_account_user_id_idx_failed";
+--> statement-breakpoint
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "workspace_billed_account_user_id_idx" ON "workspace" USING btree ("billed_account_user_id");
+--> statement-breakpoint
+SET lock_timeout = '5s';