diff --git a/apps/sim/lib/copilot/tools/server/table/user-table.test.ts b/apps/sim/lib/copilot/tools/server/table/user-table.test.ts index 4d223e842df..8ad2e0d6b2b 100644 --- a/apps/sim/lib/copilot/tools/server/table/user-table.test.ts +++ b/apps/sim/lib/copilot/tools/server/table/user-table.test.ts @@ -12,6 +12,9 @@ const { mockBatchInsertRows, mockReplaceTableRows, mockAddWorkflowGroup, + mockCreateTable, + mockDeleteTable, + mockGetWorkspaceTableLimits, fakeEnrichment, } = vi.hoisted(() => ({ mockResolveWorkspaceFileReference: vi.fn(), @@ -20,6 +23,9 @@ const { mockBatchInsertRows: vi.fn(), mockReplaceTableRows: vi.fn(), mockAddWorkflowGroup: vi.fn(), + mockCreateTable: vi.fn(), + mockDeleteTable: vi.fn(), + mockGetWorkspaceTableLimits: vi.fn(), fakeEnrichment: { id: 'work-email', name: 'Work Email', @@ -54,13 +60,13 @@ vi.mock('@/lib/table/service', () => ({ addWorkflowGroup: mockAddWorkflowGroup, batchInsertRows: mockBatchInsertRows, batchUpdateRows: vi.fn(), - createTable: vi.fn(), + createTable: mockCreateTable, deleteColumn: vi.fn(), deleteColumns: vi.fn(), deleteRow: vi.fn(), deleteRowsByFilter: vi.fn(), deleteRowsByIds: vi.fn(), - deleteTable: vi.fn(), + deleteTable: mockDeleteTable, getRowById: vi.fn(), getTableById: mockGetTableById, insertRow: vi.fn(), @@ -74,6 +80,10 @@ vi.mock('@/lib/table/service', () => ({ updateRowsByFilter: vi.fn(), })) +vi.mock('@/lib/table/billing', () => ({ + getWorkspaceTableLimits: mockGetWorkspaceTableLimits, +})) + import { userTableServerTool } from '@/lib/copilot/tools/server/table/user-table' function buildTable(overrides: Partial = {}): TableDefinition { @@ -232,6 +242,93 @@ describe('userTableServerTool.import_file', () => { }) }) +describe('userTableServerTool.create_from_file', () => { + beforeEach(() => { + vi.clearAllMocks() + mockResolveWorkspaceFileReference.mockResolvedValue({ name: 'people.csv', type: 'text/csv' }) + mockDownloadWorkspaceFile.mockResolvedValue(Buffer.from('name,age\nAlice,30\nBob,40')) + mockGetWorkspaceTableLimits.mockResolvedValue({ maxRowsPerTable: 1000, maxTables: 3 }) + mockCreateTable.mockResolvedValue(buildTable({ id: 'tbl_new', name: 'people' })) + mockDeleteTable.mockResolvedValue(undefined) + mockBatchInsertRows.mockImplementation(async (data: { rows: unknown[] }) => + data.rows.map((_, i) => ({ id: `row_${i}` })) + ) + }) + + it('stamps the workspace plan limits on the created table', async () => { + const result = await userTableServerTool.execute( + { operation: 'create_from_file', args: { fileId: 'file-1' } }, + { userId: 'user-1', workspaceId: 'workspace-1' } + ) + + expect(result.success).toBe(true) + expect(mockGetWorkspaceTableLimits).toHaveBeenCalledWith('workspace-1') + expect(mockCreateTable).toHaveBeenCalledTimes(1) + const createArgs = mockCreateTable.mock.calls[0][0] as { maxRows: number; maxTables: number } + expect(createArgs.maxRows).toBe(1000) + expect(createArgs.maxTables).toBe(3) + }) + + it('truncates to the plan row limit and reports dropped rows', async () => { + // File has 2 data rows (Alice, Bob); plan cap is 1. + mockGetWorkspaceTableLimits.mockResolvedValueOnce({ maxRowsPerTable: 1, maxTables: 3 }) + + const result = await userTableServerTool.execute( + { operation: 'create_from_file', args: { fileId: 'file-1' } }, + { userId: 'user-1', workspaceId: 'workspace-1' } + ) + + expect(result.success).toBe(true) + expect(mockCreateTable).toHaveBeenCalledTimes(1) + const insertCall = mockBatchInsertRows.mock.calls[0][0] as { rows: unknown[] } + expect(insertCall.rows).toHaveLength(1) + expect(result.data?.rowCount).toBe(1) + expect(result.message).toMatch(/dropped 1 row/i) + expect(mockDeleteTable).not.toHaveBeenCalled() + }) + + it('rolls back the created table and reports the reason when row insertion fails', async () => { + mockBatchInsertRows.mockRejectedValueOnce(new Error('Row 2: Column "email" must be unique')) + + const result = await userTableServerTool.execute( + { operation: 'create_from_file', args: { fileId: 'file-1' } }, + { userId: 'user-1', workspaceId: 'workspace-1' } + ) + + expect(result.success).toBe(false) + expect(mockDeleteTable).toHaveBeenCalledWith('tbl_new', expect.any(String)) + expect(result.message).toMatch(/rolled back/i) + expect(result.message).toMatch(/must be unique/i) + }) +}) + +describe('userTableServerTool.create', () => { + beforeEach(() => { + vi.clearAllMocks() + mockGetWorkspaceTableLimits.mockResolvedValue({ maxRowsPerTable: 1000, maxTables: 3 }) + mockCreateTable.mockResolvedValue(buildTable({ id: 'tbl_new', name: 'People' })) + }) + + it('stamps the workspace plan limits on the created table', async () => { + const result = await userTableServerTool.execute( + { + operation: 'create', + args: { + name: 'People', + schema: { columns: [{ name: 'name', type: 'string', required: true }] }, + }, + }, + { userId: 'user-1', workspaceId: 'workspace-1' } + ) + + expect(result.success).toBe(true) + expect(mockGetWorkspaceTableLimits).toHaveBeenCalledWith('workspace-1') + const createArgs = mockCreateTable.mock.calls[0][0] as { maxRows: number; maxTables: number } + expect(createArgs.maxRows).toBe(1000) + expect(createArgs.maxTables).toBe(3) + }) +}) + describe('userTableServerTool.list_enrichments', () => { beforeEach(() => { vi.clearAllMocks() diff --git a/apps/sim/lib/copilot/tools/server/table/user-table.ts b/apps/sim/lib/copilot/tools/server/table/user-table.ts index 578c4e909bc..d20d172711a 100644 --- a/apps/sim/lib/copilot/tools/server/table/user-table.ts +++ b/apps/sim/lib/copilot/tools/server/table/user-table.ts @@ -14,6 +14,7 @@ import { type CsvHeaderMapping, CsvImportValidationError, coerceRowsForTable, + getWorkspaceTableLimits, inferSchemaFromCsv, parseCsvBuffer, sanitizeName, @@ -263,6 +264,7 @@ export const userTableServerTool: BaseServerTool const requestId = generateId().slice(0, 8) assertNotAborted() + const planLimits = await getWorkspaceTableLimits(workspaceId) const table = await createTable( { name: args.name, @@ -270,6 +272,8 @@ export const userTableServerTool: BaseServerTool schema: args.schema, workspaceId, userId: context.userId, + maxRows: planLimits.maxRowsPerTable, + maxTables: planLimits.maxTables, }, requestId ) @@ -761,6 +765,11 @@ export const userTableServerTool: BaseServerTool const tableName = args.name || file.name.replace(/\.[^.]+$/, '') const requestId = generateId().slice(0, 8) assertNotAborted() + const planLimits = await getWorkspaceTableLimits(workspaceId) + + const droppedRows = Math.max(0, rows.length - planLimits.maxRowsPerTable) + const rowsToImport = droppedRows > 0 ? rows.slice(0, planLimits.maxRowsPerTable) : rows + const table = await createTable( { name: tableName, @@ -768,24 +777,59 @@ export const userTableServerTool: BaseServerTool schema: { columns }, workspaceId, userId: context.userId, + maxRows: planLimits.maxRowsPerTable, + maxTables: planLimits.maxTables, }, requestId ) - const coerced = coerceRowsForTable(rows, { columns }, headerToColumn) - const inserted = await batchInsertAll(table.id, coerced, table, workspaceId, context) + const coerced = coerceRowsForTable(rowsToImport, { columns }, headerToColumn) + let inserted: number + try { + inserted = await batchInsertAll(table.id, coerced, table, workspaceId, context) + } catch (insertError) { + const cleanupRequestId = generateId().slice(0, 8) + await deleteTable(table.id, cleanupRequestId).catch((cleanupError) => { + logger.error('Failed to roll back table after import failure', { + tableId: table.id, + error: toError(cleanupError).message, + }) + }) + const reason = toError(insertError).message + const cause = + insertError instanceof Error && insertError.cause + ? toError(insertError.cause).message + : undefined + logger.error('Failed to import rows into new table', { + tableId: table.id, + fileName: file.name, + error: reason, + cause, + }) + return { + success: false, + message: `Failed to import rows from "${file.name}" — the table was rolled back. ${cause ? `${reason} (${cause})` : reason}`, + } + } logger.info('Table created from file', { tableId: table.id, fileName: file.name, columns: columns.length, rows: inserted, + droppedRows, userId: context.userId, }) + const createdMessage = `Created table "${table.name}" with ${columns.length} columns and ${inserted.toLocaleString()} rows from "${file.name}"` + const message = + droppedRows > 0 + ? `${createdMessage}. Dropped ${droppedRows.toLocaleString()} row(s) that exceed this plan's limit of ${planLimits.maxRowsPerTable.toLocaleString()} rows per table.` + : createdMessage + return { success: true, - message: `Created table "${table.name}" with ${columns.length} columns and ${inserted} rows from "${file.name}"`, + message, data: { tableId: table.id, tableName: table.name,