Skip to content

Commit 26e6206

Browse files
committed
fix(files): uniquify materialized upload names
1 parent fbd02bc commit 26e6206

2 files changed

Lines changed: 242 additions & 40 deletions

File tree

apps/sim/lib/copilot/tools/handlers/materialize-file.test.ts

Lines changed: 161 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,22 +5,26 @@ import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing'
55
import { beforeEach, describe, expect, it, vi } from 'vitest'
66

77
const {
8+
mockAllocateUniqueWorkspaceFileName,
89
mockCheckStorageQuotaForBillingContext,
910
mockDecompress,
1011
mockFetchBuffer,
1112
mockFindFolder,
1213
mockFindUpload,
14+
mockGetWorkspaceFile,
1315
mockHasCloudStorage,
1416
mockHeadObject,
1517
mockIncrementStorageUsageForBillingContextInTx,
1618
mockMaybeNotifyStorageLimitForBillingContext,
1719
mockResolveStorageBillingContext,
1820
} = vi.hoisted(() => ({
21+
mockAllocateUniqueWorkspaceFileName: vi.fn(),
1922
mockCheckStorageQuotaForBillingContext: vi.fn(),
2023
mockDecompress: vi.fn(),
2124
mockFetchBuffer: vi.fn(),
2225
mockFindFolder: vi.fn(),
2326
mockFindUpload: vi.fn(),
27+
mockGetWorkspaceFile: vi.fn(),
2428
mockHasCloudStorage: vi.fn(),
2529
mockHeadObject: vi.fn(),
2630
mockIncrementStorageUsageForBillingContextInTx: vi.fn(),
@@ -41,7 +45,9 @@ vi.mock('@/lib/uploads', () => ({
4145
}))
4246

4347
vi.mock('@/lib/uploads/contexts/workspace/workspace-file-manager', () => ({
48+
allocateUniqueWorkspaceFileName: mockAllocateUniqueWorkspaceFileName,
4449
fetchWorkspaceFileBuffer: mockFetchBuffer,
50+
getWorkspaceFile: mockGetWorkspaceFile,
4551
}))
4652

4753
vi.mock('@/lib/uploads/contexts/workspace/workspace-file-folder-manager', () => ({
@@ -76,7 +82,9 @@ vi.mock('@/lib/billing/storage', () => ({
7682
}))
7783

7884
vi.mock('@/lib/copilot/vfs/path-utils', () => ({
79-
canonicalWorkspaceFilePath: vi.fn(() => 'files/report.txt'),
85+
canonicalWorkspaceFilePath: vi.fn(
86+
({ name }: { name: string }) => `files/${encodeURIComponent(name)}`
87+
),
8088
encodeVfsPathSegments: (segments: string[]) =>
8189
segments.map((s) => encodeURIComponent(s)).join('/'),
8290
}))
@@ -237,6 +245,8 @@ describe('executeMaterializeFile - save storage transition', () => {
237245
vi.clearAllMocks()
238246
resetDbChainMock()
239247
mockFindUpload.mockResolvedValue(mothershipRow)
248+
mockAllocateUniqueWorkspaceFileName.mockResolvedValue('report.txt')
249+
mockGetWorkspaceFile.mockResolvedValue({ id: 'file-1', name: 'report.txt' })
240250
mockHeadObject.mockResolvedValue({ size: 250, contentType: 'text/plain' })
241251
mockHasCloudStorage.mockReturnValue(true)
242252
mockResolveStorageBillingContext.mockResolvedValue(STORAGE_CONTEXT)
@@ -278,6 +288,11 @@ describe('executeMaterializeFile - save storage transition', () => {
278288
expect(result.success).toBe(true)
279289
expect(mockHeadObject).toHaveBeenCalledWith('mothership/file-1', 'mothership')
280290
expect(mockCheckStorageQuotaForBillingContext).toHaveBeenCalledWith(STORAGE_CONTEXT, 250)
291+
expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenCalledWith(
292+
context.workspaceId,
293+
'report.txt',
294+
null
295+
)
281296
expect(dbChainMockFns.set).toHaveBeenCalledWith(
282297
expect.objectContaining({ context: 'workspace', chatId: null, size: 250 })
283298
)
@@ -287,15 +302,160 @@ describe('executeMaterializeFile - save storage transition', () => {
287302
)
288303
})
289304

305+
it('materializes with an available root-level copy name', async () => {
306+
mockFindUpload.mockResolvedValueOnce({
307+
...mothershipRow,
308+
originalName: 'image.png',
309+
displayName: 'image.png',
310+
})
311+
mockAllocateUniqueWorkspaceFileName.mockResolvedValueOnce('image (1).png')
312+
dbChainMockFns.returning.mockResolvedValueOnce([
313+
{ id: 'file-1', originalName: 'image (1).png' },
314+
])
315+
316+
const result = await executeMaterializeFile(
317+
{ fileNames: ['image.png'], operation: 'save' },
318+
context
319+
)
320+
321+
expect(result.success).toBe(true)
322+
expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenCalledWith(
323+
context.workspaceId,
324+
'image.png',
325+
null
326+
)
327+
expect(dbChainMockFns.set).toHaveBeenCalledWith(
328+
expect.objectContaining({ context: 'workspace', originalName: 'image (1).png' })
329+
)
330+
expect(result.resources).toEqual([{ type: 'file', id: 'file-1', title: 'image (1).png' }])
331+
})
332+
333+
it('reallocates and retries when a concurrent root-level write claims the name', async () => {
334+
const nameCollision = Object.assign(new Error('duplicate workspace file name'), {
335+
code: '23505',
336+
constraint_name: 'workspace_files_workspace_folder_name_active_unique',
337+
})
338+
mockFindUpload.mockResolvedValueOnce({
339+
...mothershipRow,
340+
originalName: 'image.png',
341+
displayName: 'image.png',
342+
})
343+
mockAllocateUniqueWorkspaceFileName
344+
.mockResolvedValueOnce('image (1).png')
345+
.mockResolvedValueOnce('image (2).png')
346+
dbChainMockFns.returning
347+
.mockRejectedValueOnce(nameCollision)
348+
.mockResolvedValueOnce([{ id: 'file-1', originalName: 'image (2).png' }])
349+
350+
const result = await executeMaterializeFile(
351+
{ fileNames: ['image.png'], operation: 'save' },
352+
context
353+
)
354+
355+
expect(result.success).toBe(true)
356+
expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenCalledTimes(2)
357+
expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenNthCalledWith(
358+
1,
359+
context.workspaceId,
360+
'image.png',
361+
null
362+
)
363+
expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenNthCalledWith(
364+
2,
365+
context.workspaceId,
366+
'image.png',
367+
null
368+
)
369+
expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(2)
370+
expect(dbChainMockFns.set).toHaveBeenNthCalledWith(
371+
1,
372+
expect.objectContaining({ originalName: 'image (1).png' })
373+
)
374+
expect(dbChainMockFns.set).toHaveBeenNthCalledWith(
375+
2,
376+
expect.objectContaining({ originalName: 'image (2).png' })
377+
)
378+
expect(mockIncrementStorageUsageForBillingContextInTx).toHaveBeenCalledTimes(1)
379+
expect(result.resources).toEqual([{ type: 'file', id: 'file-1', title: 'image (2).png' }])
380+
})
381+
382+
it('stops after the bounded number of root-level name collisions', async () => {
383+
const nameCollision = Object.assign(new Error('duplicate workspace file name'), {
384+
code: '23505',
385+
constraint_name: 'workspace_files_workspace_folder_name_active_unique',
386+
})
387+
dbChainMockFns.returning.mockRejectedValue(nameCollision)
388+
389+
const result = await executeMaterializeFile(
390+
{ fileNames: ['report.txt'], operation: 'save' },
391+
context
392+
)
393+
394+
expect(result.success).toBe(false)
395+
expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenCalledTimes(8)
396+
expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(8)
397+
expect(mockIncrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled()
398+
expect(mockMaybeNotifyStorageLimitForBillingContext).not.toHaveBeenCalled()
399+
})
400+
401+
it('does not retry unique violations from a different constraint', async () => {
402+
const keyCollision = Object.assign(new Error('duplicate workspace file key'), {
403+
code: '23505',
404+
constraint_name: 'workspace_files_key_active_unique',
405+
})
406+
dbChainMockFns.returning.mockRejectedValueOnce(keyCollision)
407+
408+
const result = await executeMaterializeFile(
409+
{ fileNames: ['report.txt'], operation: 'save' },
410+
context
411+
)
412+
413+
expect(result.success).toBe(false)
414+
expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenCalledTimes(1)
415+
expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(1)
416+
expect(mockIncrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled()
417+
})
418+
290419
it('treats a lost conditional transition as a replay no-op', async () => {
291420
dbChainMockFns.returning.mockResolvedValueOnce([])
421+
mockGetWorkspaceFile.mockResolvedValueOnce({ id: 'file-1', name: 'report (1).txt' })
292422

293423
const result = await executeMaterializeFile(
294424
{ fileNames: ['report.txt'], operation: 'save' },
295425
context
296426
)
297427

298428
expect(result.success).toBe(true)
429+
expect(mockGetWorkspaceFile).toHaveBeenCalledWith(context.workspaceId, 'file-1', {
430+
throwOnError: true,
431+
})
432+
expect(result.resources).toEqual([{ type: 'file', id: 'file-1', title: 'report (1).txt' }])
433+
expect(mockIncrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled()
434+
expect(mockMaybeNotifyStorageLimitForBillingContext).not.toHaveBeenCalled()
435+
})
436+
437+
it('fails a replay when the materialized workspace file no longer exists', async () => {
438+
dbChainMockFns.returning.mockResolvedValueOnce([])
439+
mockGetWorkspaceFile.mockResolvedValueOnce(null)
440+
441+
const result = await executeMaterializeFile(
442+
{ fileNames: ['report.txt'], operation: 'save' },
443+
context
444+
)
445+
446+
expect(result.success).toBe(false)
447+
expect(result.output).toEqual({
448+
succeeded: [],
449+
failed: [
450+
{
451+
fileName: 'report.txt',
452+
error: 'Upload no longer available: "report.txt".',
453+
},
454+
],
455+
})
456+
expect(mockGetWorkspaceFile).toHaveBeenCalledWith(context.workspaceId, 'file-1', {
457+
throwOnError: true,
458+
})
299459
expect(mockIncrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled()
300460
expect(mockMaybeNotifyStorageLimitForBillingContext).not.toHaveBeenCalled()
301461
})

apps/sim/lib/copilot/tools/handlers/materialize-file.ts

Lines changed: 81 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,12 @@ import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit'
22
import { db } from '@sim/db'
33
import { folder as folderTable, workflow, workspaceFiles } from '@sim/db/schema'
44
import { createLogger } from '@sim/logger'
5-
import { getErrorMessage, toError } from '@sim/utils/errors'
5+
import {
6+
getErrorMessage,
7+
getPostgresConstraintName,
8+
getPostgresErrorCode,
9+
toError,
10+
} from '@sim/utils/errors'
611
import { generateId } from '@sim/utils/id'
712
import { and, eq, isNull, sql } from 'drizzle-orm'
813
import {
@@ -23,7 +28,11 @@ import {
2328
MAX_ARCHIVE_BYTES,
2429
} from '@/lib/uploads/archive'
2530
import { findWorkspaceFileFolderIdByPath } from '@/lib/uploads/contexts/workspace/workspace-file-folder-manager'
26-
import { fetchWorkspaceFileBuffer } from '@/lib/uploads/contexts/workspace/workspace-file-manager'
31+
import {
32+
allocateUniqueWorkspaceFileName,
33+
fetchWorkspaceFileBuffer,
34+
getWorkspaceFile,
35+
} from '@/lib/uploads/contexts/workspace/workspace-file-manager'
2736
import { hasCloudStorage, headObject } from '@/lib/uploads/core/storage-service'
2837
import { isArchiveFileName } from '@/lib/uploads/utils/file-utils'
2938
import { parseWorkflowJson } from '@/lib/workflows/operations/import-export'
@@ -32,6 +41,8 @@ import { deduplicateWorkflowName } from '@/lib/workflows/utils'
3241
import { extractWorkflowMetadata } from '@/app/api/v1/admin/types'
3342

3443
const logger = createLogger('MaterializeFile')
44+
const MAX_MATERIALIZE_NAME_RETRIES = 8
45+
const WORKSPACE_FILE_NAME_UNIQUE_INDEX = 'workspace_files_workspace_folder_name_active_unique'
3546

3647
function toFileRecord(row: typeof workspaceFiles.$inferSelect) {
3748
const pathPrefix = getServePathPrefix()
@@ -107,46 +118,75 @@ async function executeSave(
107118
* workspace lock before locking its payer. Any quota/stale-payer failure
108119
* rolls back the row transition.
109120
*/
110-
const transition = await db.transaction(async (tx) => {
111-
await tx.execute(sql`SELECT 1 FROM workspace WHERE id = ${workspaceId} FOR UPDATE`)
112-
113-
const [updated] = await tx
114-
.update(workspaceFiles)
115-
.set({
116-
context: 'workspace',
117-
// A workspace file has no birth chat or message — clear both provenance
118-
// fields so the row reads as workspace-owned, not stale chat-owned.
119-
chatId: null,
120-
messageId: null,
121-
originalName: row.displayName ?? row.originalName,
122-
size: verifiedSize,
123-
})
124-
.where(
125-
and(
126-
eq(workspaceFiles.id, row.id),
127-
eq(workspaceFiles.workspaceId, workspaceId),
128-
eq(workspaceFiles.chatId, chatId),
129-
eq(workspaceFiles.context, 'mothership'),
130-
isNull(workspaceFiles.deletedAt)
131-
)
132-
)
133-
.returning({ id: workspaceFiles.id, originalName: workspaceFiles.originalName })
121+
let transition: {
122+
updated: { id: string; originalName: string }
123+
updatedUsage: number | undefined
124+
} | null = null
134125

135-
if (!updated) {
136-
return null
137-
}
126+
for (let attempt = 0; attempt < MAX_MATERIALIZE_NAME_RETRIES; attempt++) {
127+
const materializedName = await allocateUniqueWorkspaceFileName(workspaceId, displayName, null)
138128

139-
const updatedUsage = await incrementStorageUsageForBillingContextInTx(
140-
tx,
141-
billingContext,
142-
verifiedSize
143-
)
144-
return { updated, updatedUsage }
145-
})
129+
try {
130+
transition = await db.transaction(async (tx) => {
131+
await tx.execute(sql`SELECT 1 FROM workspace WHERE id = ${workspaceId} FOR UPDATE`)
132+
133+
const [updated] = await tx
134+
.update(workspaceFiles)
135+
.set({
136+
context: 'workspace',
137+
// A workspace file has no birth chat or message — clear both provenance
138+
// fields so the row reads as workspace-owned, not stale chat-owned.
139+
chatId: null,
140+
messageId: null,
141+
originalName: materializedName,
142+
size: verifiedSize,
143+
})
144+
.where(
145+
and(
146+
eq(workspaceFiles.id, row.id),
147+
eq(workspaceFiles.workspaceId, workspaceId),
148+
eq(workspaceFiles.chatId, chatId),
149+
eq(workspaceFiles.context, 'mothership'),
150+
isNull(workspaceFiles.deletedAt)
151+
)
152+
)
153+
.returning({ id: workspaceFiles.id, originalName: workspaceFiles.originalName })
146154

147-
const updated = transition?.updated ?? {
148-
id: row.id,
149-
originalName: row.displayName ?? row.originalName,
155+
if (!updated) {
156+
return null
157+
}
158+
159+
const updatedUsage = await incrementStorageUsageForBillingContextInTx(
160+
tx,
161+
billingContext,
162+
verifiedSize
163+
)
164+
return { updated, updatedUsage }
165+
})
166+
break
167+
} catch (error) {
168+
const isNameCollision =
169+
getPostgresErrorCode(error) === '23505' &&
170+
getPostgresConstraintName(error) === WORKSPACE_FILE_NAME_UNIQUE_INDEX
171+
if (!isNameCollision || attempt === MAX_MATERIALIZE_NAME_RETRIES - 1) {
172+
throw error
173+
}
174+
logger.warn('Workspace file name was claimed during materialization; retrying', {
175+
fileName,
176+
materializedName,
177+
attempt: attempt + 1,
178+
})
179+
}
180+
}
181+
182+
const replayedFile = transition
183+
? null
184+
: await getWorkspaceFile(workspaceId, row.id, { throwOnError: true })
185+
const updated =
186+
transition?.updated ??
187+
(replayedFile ? { id: replayedFile.id, originalName: replayedFile.name } : null)
188+
if (!updated) {
189+
return { success: false, error: `Upload no longer available: "${fileName}".` }
150190
}
151191
if (transition?.updatedUsage !== undefined) {
152192
void maybeNotifyStorageLimitForBillingContext(billingContext, transition.updatedUsage)
@@ -541,6 +581,8 @@ export async function executeMaterializeFile(
541581
operation,
542582
chatId: context.chatId,
543583
error: toError(err).message,
584+
postgresCode: getPostgresErrorCode(err),
585+
postgresConstraint: getPostgresConstraintName(err),
544586
})
545587
failed.push({
546588
fileName,

0 commit comments

Comments
 (0)