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

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,86 @@ 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' }))
+    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('rejects a file exceeding the plan row limit without creating a table', async () => {
+    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(false)
+    expect(result.message).toMatch(/exceeds this plan's limit/i)
+    expect(mockCreateTable).not.toHaveBeenCalled()
+    expect(mockDeleteTable).not.toHaveBeenCalled()
+  })
+
+  it('deletes the created table when row insertion fails', async () => {
+    mockBatchInsertRows.mockRejectedValueOnce(new Error('Maximum row limit (1000) reached'))
+
+    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))
+  })
+})
+
+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,19 +765,36 @@ 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)
+
+          if (rows.length > planLimits.maxRowsPerTable) {
+            return {
+              success: false,
+              message: `"${file.name}" has ${rows.length.toLocaleString()} rows, which exceeds this plan's limit of ${planLimits.maxRowsPerTable.toLocaleString()} rows per table.`,
+            }
+          }
+
           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)
+          let inserted: number
+          try {
+            inserted = await batchInsertAll(table.id, coerced, table, workspaceId, context)
+          } catch (insertError) {
+            await deleteTable(table.id, requestId).catch(() => {})
+            throw insertError
+          }
 
           logger.info('Table created from file', {
             tableId: table.id,