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
2026-09-24 15:45:35 +08:00 · 2026-06-01 20:11:59 -04:00
parent 8aa74a2d8d
commit fc7e35edc0
2 changed files with 146 additions and 5 deletions
@@ -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()
@@ -14,6 +14,7 @@ import {
  type CsvHeaderMapping,
  CsvImportValidationError,
  coerceRowsForTable,
+  getWorkspaceTableLimits,
  inferSchemaFromCsv,
  parseCsvBuffer,
  sanitizeName,
@@ -263,6 +264,7 @@ 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,
@@ -270,6 +272,8 @@ export const userTableServerTool: BaseServerTool<UserTableArgs, UserTableResult>
              schema: args.schema,
              workspaceId,
              userId: context.userId,
+              maxRows: planLimits.maxRowsPerTable,
+              maxTables: planLimits.maxTables,
            },
            requestId
          )
@@ -761,6 +765,11 @@ 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,
@@ -768,24 +777,59 @@ export const userTableServerTool: BaseServerTool<UserTableArgs, UserTableResult>
              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,