diff --git a/apps/sim/lib/copilot/tools/handlers/materialize-file.test.ts b/apps/sim/lib/copilot/tools/handlers/materialize-file.test.ts index bdad3967c45..d3e2751c277 100644 --- a/apps/sim/lib/copilot/tools/handlers/materialize-file.test.ts +++ b/apps/sim/lib/copilot/tools/handlers/materialize-file.test.ts @@ -5,22 +5,26 @@ import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' import { beforeEach, describe, expect, it, vi } from 'vitest' const { + mockAllocateUniqueWorkspaceFileName, mockCheckStorageQuotaForBillingContext, mockDecompress, mockFetchBuffer, mockFindFolder, mockFindUpload, + mockGetWorkspaceFile, mockHasCloudStorage, mockHeadObject, mockIncrementStorageUsageForBillingContextInTx, mockMaybeNotifyStorageLimitForBillingContext, mockResolveStorageBillingContext, } = vi.hoisted(() => ({ + mockAllocateUniqueWorkspaceFileName: vi.fn(), mockCheckStorageQuotaForBillingContext: vi.fn(), mockDecompress: vi.fn(), mockFetchBuffer: vi.fn(), mockFindFolder: vi.fn(), mockFindUpload: vi.fn(), + mockGetWorkspaceFile: vi.fn(), mockHasCloudStorage: vi.fn(), mockHeadObject: vi.fn(), mockIncrementStorageUsageForBillingContextInTx: vi.fn(), @@ -41,7 +45,9 @@ vi.mock('@/lib/uploads', () => ({ })) vi.mock('@/lib/uploads/contexts/workspace/workspace-file-manager', () => ({ + allocateUniqueWorkspaceFileName: mockAllocateUniqueWorkspaceFileName, fetchWorkspaceFileBuffer: mockFetchBuffer, + getWorkspaceFile: mockGetWorkspaceFile, })) vi.mock('@/lib/uploads/contexts/workspace/workspace-file-folder-manager', () => ({ @@ -76,7 +82,9 @@ vi.mock('@/lib/billing/storage', () => ({ })) vi.mock('@/lib/copilot/vfs/path-utils', () => ({ - canonicalWorkspaceFilePath: vi.fn(() => 'files/report.txt'), + canonicalWorkspaceFilePath: vi.fn( + ({ name }: { name: string }) => `files/${encodeURIComponent(name)}` + ), encodeVfsPathSegments: (segments: string[]) => segments.map((s) => encodeURIComponent(s)).join('/'), })) @@ -237,6 +245,8 @@ describe('executeMaterializeFile - save storage transition', () => { vi.clearAllMocks() resetDbChainMock() mockFindUpload.mockResolvedValue(mothershipRow) + mockAllocateUniqueWorkspaceFileName.mockResolvedValue('report.txt') + mockGetWorkspaceFile.mockResolvedValue({ id: 'file-1', name: 'report.txt' }) mockHeadObject.mockResolvedValue({ size: 250, contentType: 'text/plain' }) mockHasCloudStorage.mockReturnValue(true) mockResolveStorageBillingContext.mockResolvedValue(STORAGE_CONTEXT) @@ -278,6 +288,11 @@ describe('executeMaterializeFile - save storage transition', () => { expect(result.success).toBe(true) expect(mockHeadObject).toHaveBeenCalledWith('mothership/file-1', 'mothership') expect(mockCheckStorageQuotaForBillingContext).toHaveBeenCalledWith(STORAGE_CONTEXT, 250) + expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenCalledWith( + context.workspaceId, + 'report.txt', + null + ) expect(dbChainMockFns.set).toHaveBeenCalledWith( expect.objectContaining({ context: 'workspace', chatId: null, size: 250 }) ) @@ -287,8 +302,129 @@ describe('executeMaterializeFile - save storage transition', () => { ) }) + it('materializes with an available root-level copy name', async () => { + mockFindUpload.mockResolvedValueOnce({ + ...mothershipRow, + originalName: 'image.png', + displayName: 'image.png', + }) + mockAllocateUniqueWorkspaceFileName.mockResolvedValueOnce('image (1).png') + dbChainMockFns.returning.mockResolvedValueOnce([ + { id: 'file-1', originalName: 'image (1).png' }, + ]) + + const result = await executeMaterializeFile( + { fileNames: ['image.png'], operation: 'save' }, + context + ) + + expect(result.success).toBe(true) + expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenCalledWith( + context.workspaceId, + 'image.png', + null + ) + expect(dbChainMockFns.set).toHaveBeenCalledWith( + expect.objectContaining({ + context: 'workspace', + originalName: 'image (1).png', + displayName: 'image (1).png', + }) + ) + expect(result.output).toEqual({ succeeded: ['image (1).png'], failed: [] }) + expect(result.resources).toEqual([{ type: 'file', id: 'file-1', title: 'image (1).png' }]) + }) + + it('reallocates and retries when a concurrent root-level write claims the name', async () => { + const nameCollision = Object.assign(new Error('duplicate workspace file name'), { + code: '23505', + constraint_name: 'workspace_files_workspace_folder_name_active_unique', + }) + mockFindUpload.mockResolvedValueOnce({ + ...mothershipRow, + originalName: 'image.png', + displayName: 'image.png', + }) + mockAllocateUniqueWorkspaceFileName + .mockResolvedValueOnce('image (1).png') + .mockResolvedValueOnce('image (2).png') + dbChainMockFns.returning + .mockRejectedValueOnce(nameCollision) + .mockResolvedValueOnce([{ id: 'file-1', originalName: 'image (2).png' }]) + + const result = await executeMaterializeFile( + { fileNames: ['image.png'], operation: 'save' }, + context + ) + + expect(result.success).toBe(true) + expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenCalledTimes(2) + expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenNthCalledWith( + 1, + context.workspaceId, + 'image.png', + null + ) + expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenNthCalledWith( + 2, + context.workspaceId, + 'image.png', + null + ) + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(2) + expect(dbChainMockFns.set).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ originalName: 'image (1).png' }) + ) + expect(dbChainMockFns.set).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ originalName: 'image (2).png' }) + ) + expect(mockIncrementStorageUsageForBillingContextInTx).toHaveBeenCalledTimes(1) + expect(result.output).toEqual({ succeeded: ['image (2).png'], failed: [] }) + expect(result.resources).toEqual([{ type: 'file', id: 'file-1', title: 'image (2).png' }]) + }) + + it('stops after the bounded number of root-level name collisions', async () => { + const nameCollision = Object.assign(new Error('duplicate workspace file name'), { + code: '23505', + constraint_name: 'workspace_files_workspace_folder_name_active_unique', + }) + dbChainMockFns.returning.mockRejectedValue(nameCollision) + + const result = await executeMaterializeFile( + { fileNames: ['report.txt'], operation: 'save' }, + context + ) + + expect(result.success).toBe(false) + expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenCalledTimes(8) + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(8) + expect(mockIncrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled() + expect(mockMaybeNotifyStorageLimitForBillingContext).not.toHaveBeenCalled() + }) + + it('does not retry unique violations from a different constraint', async () => { + const keyCollision = Object.assign(new Error('duplicate workspace file key'), { + code: '23505', + constraint_name: 'workspace_files_key_active_unique', + }) + dbChainMockFns.returning.mockRejectedValueOnce(keyCollision) + + const result = await executeMaterializeFile( + { fileNames: ['report.txt'], operation: 'save' }, + context + ) + + expect(result.success).toBe(false) + expect(mockAllocateUniqueWorkspaceFileName).toHaveBeenCalledTimes(1) + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(1) + expect(mockIncrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled() + }) + it('treats a lost conditional transition as a replay no-op', async () => { dbChainMockFns.returning.mockResolvedValueOnce([]) + mockGetWorkspaceFile.mockResolvedValueOnce({ id: 'file-1', name: 'report (1).txt' }) const result = await executeMaterializeFile( { fileNames: ['report.txt'], operation: 'save' }, @@ -296,6 +432,37 @@ describe('executeMaterializeFile - save storage transition', () => { ) expect(result.success).toBe(true) + expect(mockGetWorkspaceFile).toHaveBeenCalledWith(context.workspaceId, 'file-1', { + throwOnError: true, + }) + expect(result.output).toEqual({ succeeded: ['report (1).txt'], failed: [] }) + expect(result.resources).toEqual([{ type: 'file', id: 'file-1', title: 'report (1).txt' }]) + expect(mockIncrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled() + expect(mockMaybeNotifyStorageLimitForBillingContext).not.toHaveBeenCalled() + }) + + it('fails a replay when the materialized workspace file no longer exists', async () => { + dbChainMockFns.returning.mockResolvedValueOnce([]) + mockGetWorkspaceFile.mockResolvedValueOnce(null) + + const result = await executeMaterializeFile( + { fileNames: ['report.txt'], operation: 'save' }, + context + ) + + expect(result.success).toBe(false) + expect(result.output).toEqual({ + succeeded: [], + failed: [ + { + fileName: 'report.txt', + error: 'Upload no longer available: "report.txt".', + }, + ], + }) + expect(mockGetWorkspaceFile).toHaveBeenCalledWith(context.workspaceId, 'file-1', { + throwOnError: true, + }) expect(mockIncrementStorageUsageForBillingContextInTx).not.toHaveBeenCalled() expect(mockMaybeNotifyStorageLimitForBillingContext).not.toHaveBeenCalled() }) diff --git a/apps/sim/lib/copilot/tools/handlers/materialize-file.ts b/apps/sim/lib/copilot/tools/handlers/materialize-file.ts index 3e2abda31a2..69314264ec9 100644 --- a/apps/sim/lib/copilot/tools/handlers/materialize-file.ts +++ b/apps/sim/lib/copilot/tools/handlers/materialize-file.ts @@ -2,7 +2,12 @@ import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit' import { db } from '@sim/db' import { folder as folderTable, workflow, workspaceFiles } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { getErrorMessage, toError } from '@sim/utils/errors' +import { + getErrorMessage, + getPostgresConstraintName, + getPostgresErrorCode, + toError, +} from '@sim/utils/errors' import { generateId } from '@sim/utils/id' import { and, eq, isNull, sql } from 'drizzle-orm' import { @@ -23,7 +28,11 @@ import { MAX_ARCHIVE_BYTES, } from '@/lib/uploads/archive' import { findWorkspaceFileFolderIdByPath } from '@/lib/uploads/contexts/workspace/workspace-file-folder-manager' -import { fetchWorkspaceFileBuffer } from '@/lib/uploads/contexts/workspace/workspace-file-manager' +import { + allocateUniqueWorkspaceFileName, + fetchWorkspaceFileBuffer, + getWorkspaceFile, +} from '@/lib/uploads/contexts/workspace/workspace-file-manager' import { hasCloudStorage, headObject } from '@/lib/uploads/core/storage-service' import { isArchiveFileName } from '@/lib/uploads/utils/file-utils' import { parseWorkflowJson } from '@/lib/workflows/operations/import-export' @@ -32,6 +41,8 @@ import { deduplicateWorkflowName } from '@/lib/workflows/utils' import { extractWorkflowMetadata } from '@/app/api/v1/admin/types' const logger = createLogger('MaterializeFile') +const MAX_MATERIALIZE_NAME_RETRIES = 8 +const WORKSPACE_FILE_NAME_UNIQUE_INDEX = 'workspace_files_workspace_folder_name_active_unique' function toFileRecord(row: typeof workspaceFiles.$inferSelect) { const pathPrefix = getServePathPrefix() @@ -107,46 +118,76 @@ async function executeSave( * workspace lock before locking its payer. Any quota/stale-payer failure * rolls back the row transition. */ - const transition = await db.transaction(async (tx) => { - await tx.execute(sql`SELECT 1 FROM workspace WHERE id = ${workspaceId} FOR UPDATE`) - - const [updated] = await tx - .update(workspaceFiles) - .set({ - context: 'workspace', - // A workspace file has no birth chat or message — clear both provenance - // fields so the row reads as workspace-owned, not stale chat-owned. - chatId: null, - messageId: null, - originalName: row.displayName ?? row.originalName, - size: verifiedSize, - }) - .where( - and( - eq(workspaceFiles.id, row.id), - eq(workspaceFiles.workspaceId, workspaceId), - eq(workspaceFiles.chatId, chatId), - eq(workspaceFiles.context, 'mothership'), - isNull(workspaceFiles.deletedAt) - ) - ) - .returning({ id: workspaceFiles.id, originalName: workspaceFiles.originalName }) + let transition: { + updated: { id: string; originalName: string } + updatedUsage: number | undefined + } | null = null - if (!updated) { - return null - } + for (let attempt = 0; attempt < MAX_MATERIALIZE_NAME_RETRIES; attempt++) { + const materializedName = await allocateUniqueWorkspaceFileName(workspaceId, displayName, null) - const updatedUsage = await incrementStorageUsageForBillingContextInTx( - tx, - billingContext, - verifiedSize - ) - return { updated, updatedUsage } - }) + try { + transition = await db.transaction(async (tx) => { + await tx.execute(sql`SELECT 1 FROM workspace WHERE id = ${workspaceId} FOR UPDATE`) + + const [updated] = await tx + .update(workspaceFiles) + .set({ + context: 'workspace', + // A workspace file has no birth chat or message — clear both provenance + // fields so the row reads as workspace-owned, not stale chat-owned. + chatId: null, + messageId: null, + originalName: materializedName, + displayName: materializedName, + size: verifiedSize, + }) + .where( + and( + eq(workspaceFiles.id, row.id), + eq(workspaceFiles.workspaceId, workspaceId), + eq(workspaceFiles.chatId, chatId), + eq(workspaceFiles.context, 'mothership'), + isNull(workspaceFiles.deletedAt) + ) + ) + .returning({ id: workspaceFiles.id, originalName: workspaceFiles.originalName }) - const updated = transition?.updated ?? { - id: row.id, - originalName: row.displayName ?? row.originalName, + if (!updated) { + return null + } + + const updatedUsage = await incrementStorageUsageForBillingContextInTx( + tx, + billingContext, + verifiedSize + ) + return { updated, updatedUsage } + }) + break + } catch (error) { + const isNameCollision = + getPostgresErrorCode(error) === '23505' && + getPostgresConstraintName(error) === WORKSPACE_FILE_NAME_UNIQUE_INDEX + if (!isNameCollision || attempt === MAX_MATERIALIZE_NAME_RETRIES - 1) { + throw error + } + logger.warn('Workspace file name was claimed during materialization; retrying', { + fileName, + materializedName, + attempt: attempt + 1, + }) + } + } + + const replayedFile = transition + ? null + : await getWorkspaceFile(workspaceId, row.id, { throwOnError: true }) + const updated = + transition?.updated ?? + (replayedFile ? { id: replayedFile.id, originalName: replayedFile.name } : null) + if (!updated) { + return { success: false, error: `Upload no longer available: "${fileName}".` } } if (transition?.updatedUsage !== undefined) { void maybeNotifyStorageLimitForBillingContext(billingContext, transition.updatedUsage) @@ -530,7 +571,11 @@ export async function executeMaterializeFile( } if (result.success) { - succeeded.push(fileName) + const materializedName = + operation === 'save' + ? result.resources?.find((resource) => resource.type === 'file')?.title + : undefined + succeeded.push(materializedName ?? fileName) if (result.resources) resources.push(...result.resources) } else { failed.push({ fileName, error: result.error ?? 'Failed to materialize file' }) @@ -541,6 +586,8 @@ export async function executeMaterializeFile( operation, chatId: context.chatId, error: toError(err).message, + postgresCode: getPostgresErrorCode(err), + postgresConstraint: getPostgresConstraintName(err), }) failed.push({ fileName,