fix(tables): enforce plan limits in mothership user_table tool (#4832)

* fix(tables): enforce plan limits in mothership user_table tool

* improvement(tables): truncate over-limit CSV imports to the plan cap instead of rejecting

* fix(tables): log rollback failures and surface a clear reason on failed CSV import

Author: TheodoreSpeaks

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> = {}): 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()

apps/sim/lib/copilot/tools/server/table/user-table.ts

@@ -14,6 +14,7 @@ import {
   type CsvHeaderMapping,
   CsvImportValidationError,
   coerceRowsForTable,
+  getWorkspaceTableLimits,
   inferSchemaFromCsv,
   parseCsvBuffer,
   sanitizeName,
@@ -263,13 +264,16 @@ export const userTableServerTool: BaseServerTool<UserTableArgs, UserTableResult>
 
           const requestId = generateId().slice(0, 8)
           assertNotAborted()
+          const planLimits = await getWorkspaceTableLimits(workspaceId)
           const table = await createTable(
             {
               name: args.name,
               description: args.description,
               schema: args.schema,
               workspaceId,
               userId: context.userId,
+              maxRows: planLimits.maxRowsPerTable,
+              maxTables: planLimits.maxTables,
             },
             requestId
           )
@@ -761,31 +765,71 @@ export const userTableServerTool: BaseServerTool<UserTableArgs, UserTableResult>
           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,
               description: args.description || `Imported from ${file.name}`,
               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,