fix(mothership): tenant-check outputTable writes and route them through replaceTableRows (#5011)

* v0.6.29: login improvements, posthog telemetry (#4026) * feat(posthog): Add tracking on mothership abort (#4023) Co-authored-by: Theodore Li <theo@sim.ai> * fix(login): fix captcha headers for manual login (#4025) * fix(signup): fix turnstile key loading * fix(login): fix captcha header passing * Catch user already exists, remove login form captcha * fix(mothership): tenant-check outputTable writes and route them through replaceTableRows maybeWriteOutputToTable / maybeWriteReadCsvToTable accepted any table id without verifying it belongs to the caller's workspace, so a foreign table's rows could be wiped and replaced cross-tenant. They also wrote rows with raw drizzle keyed by column *name*, bypassing the service layer's job-slot lock, validation, plan row limits, rowCount maintenance, and the stable column-id keying every other writer uses. Both handlers now reject tables outside the caller's workspace and delegate to replaceTableRows with name→id remapped rows. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(mothership): fail outputTable writes per-row when any row matches no columns Review feedback: the all-rows-empty guard let mixed batches slip unmatched rows through as empty objects. Reject on the first row that maps to zero columns, naming the row. Adds the missing CSV-suite parity tests (no-matching-headers, service-error surfacing). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Waleed <walif6@gmail.com> Co-authored-by: Siddharth Ganesan <33737564+Sg312@users.noreply.github.com> Co-authored-by: Vikhyath Mondreti <vikhyathvikku@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-09-24 15:45:35 +08:00 · 2026-06-12 21:49:42 -04:00
parent fb9e481f32
commit c3c341632f
2 changed files with 268 additions and 63 deletions
@@ -0,0 +1,221 @@
+/**
+ * @vitest-environment node
+ */
+
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import type { TableDefinition } from '@/lib/table'
+
+const { mockGetTableById, mockReplaceTableRows } = vi.hoisted(() => ({
+  mockGetTableById: vi.fn(),
+  mockReplaceTableRows: vi.fn(),
+}))
+
+vi.mock('@/lib/table/service', () => ({
+  getTableById: mockGetTableById,
+  replaceTableRows: mockReplaceTableRows,
+}))
+
+vi.mock('@/lib/copilot/request/otel', () => ({
+  withCopilotSpan: (
+    _name: string,
+    _attrs: Record<string, unknown> | undefined,
+    fn: (span: unknown) => Promise<unknown>
+  ) => fn({ setAttribute: vi.fn(), setAttributes: vi.fn(), addEvent: vi.fn() }),
+}))
+
+import { FunctionExecute, Read as ReadTool } from '@/lib/copilot/generated/tool-catalog-v1'
+import {
+  maybeWriteOutputToTable,
+  maybeWriteReadCsvToTable,
+} from '@/lib/copilot/request/tools/tables'
+import type { ExecutionContext } from '@/lib/copilot/request/types'
+
+function buildTable(overrides: Partial<TableDefinition> = {}): TableDefinition {
+  return {
+    id: 'tbl_1',
+    name: 'People',
+    description: null,
+    schema: {
+      columns: [
+        { id: 'col_name', name: 'name', type: 'string' },
+        { id: 'col_age', name: 'age', type: 'number' },
+      ],
+    },
+    metadata: null,
+    rowCount: 0,
+    maxRows: 100,
+    workspaceId: 'workspace-1',
+    createdBy: 'user-1',
+    archivedAt: null,
+    createdAt: new Date('2024-01-01'),
+    updatedAt: new Date('2024-01-01'),
+    ...overrides,
+  } as TableDefinition
+}
+
+function buildContext(overrides: Partial<ExecutionContext> = {}): ExecutionContext {
+  return {
+    userId: 'user-1',
+    workflowId: 'wf-1',
+    workspaceId: 'workspace-1',
+    ...overrides,
+  }
+}
+
+describe('maybeWriteOutputToTable', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockGetTableById.mockResolvedValue(buildTable())
+    mockReplaceTableRows.mockResolvedValue({ deletedCount: 0, insertedCount: 2 })
+  })
+
+  it('rejects a table from another workspace without touching it', async () => {
+    mockGetTableById.mockResolvedValue(buildTable({ workspaceId: 'other-workspace' }))
+
+    const result = await maybeWriteOutputToTable(
+      FunctionExecute.id,
+      { outputTable: 'tbl_1' },
+      { success: true, output: { result: [{ name: 'Alice' }] } },
+      buildContext()
+    )
+
+    expect(result).toEqual({ success: false, error: 'Table "tbl_1" not found' })
+    expect(mockReplaceTableRows).not.toHaveBeenCalled()
+  })
+
+  it('replaces rows through the service with name keys remapped to column ids', async () => {
+    const result = await maybeWriteOutputToTable(
+      FunctionExecute.id,
+      { outputTable: 'tbl_1' },
+      {
+        success: true,
+        output: {
+          result: [
+            { name: 'Alice', age: 30 },
+            { name: 'Bob', age: 40 },
+          ],
+        },
+      },
+      buildContext()
+    )
+
+    expect(result.success).toBe(true)
+    expect(mockReplaceTableRows).toHaveBeenCalledTimes(1)
+    const [data, table] = mockReplaceTableRows.mock.calls[0]
+    expect(data).toMatchObject({
+      tableId: 'tbl_1',
+      workspaceId: 'workspace-1',
+      userId: 'user-1',
+      rows: [
+        { col_name: 'Alice', col_age: 30 },
+        { col_name: 'Bob', col_age: 40 },
+      ],
+    })
+    expect(table.id).toBe('tbl_1')
+  })
+
+  it('fails fast when no row keys match the table columns', async () => {
+    const result = await maybeWriteOutputToTable(
+      FunctionExecute.id,
+      { outputTable: 'tbl_1' },
+      { success: true, output: { result: [{ wrong: 1 }, { keys: 2 }] } },
+      buildContext()
+    )
+
+    expect(result.success).toBe(false)
+    expect(result.error).toContain('Row 1 has no keys matching columns')
+    expect(mockReplaceTableRows).not.toHaveBeenCalled()
+  })
+
+  it('fails fast when only some rows match instead of writing empty rows', async () => {
+    const result = await maybeWriteOutputToTable(
+      FunctionExecute.id,
+      { outputTable: 'tbl_1' },
+      { success: true, output: { result: [{ name: 'Alice' }, { wrong: 'x' }] } },
+      buildContext()
+    )
+
+    expect(result.success).toBe(false)
+    expect(result.error).toContain('Row 2 has no keys matching columns')
+    expect(mockReplaceTableRows).not.toHaveBeenCalled()
+  })
+
+  it('surfaces service validation failures as tool errors', async () => {
+    mockReplaceTableRows.mockRejectedValue(new Error('Row 1: name is required'))
+
+    const result = await maybeWriteOutputToTable(
+      FunctionExecute.id,
+      { outputTable: 'tbl_1' },
+      { success: true, output: { result: [{ age: 30 }] } },
+      buildContext()
+    )
+
+    expect(result.success).toBe(false)
+    expect(result.error).toContain('Row 1: name is required')
+  })
+})
+
+describe('maybeWriteReadCsvToTable', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockGetTableById.mockResolvedValue(buildTable())
+    mockReplaceTableRows.mockResolvedValue({ deletedCount: 0, insertedCount: 2 })
+  })
+
+  it('rejects a table from another workspace without touching it', async () => {
+    mockGetTableById.mockResolvedValue(buildTable({ workspaceId: 'other-workspace' }))
+
+    const result = await maybeWriteReadCsvToTable(
+      ReadTool.id,
+      { outputTable: 'tbl_1', path: 'files/people.csv' },
+      { success: true, output: { content: 'name,age\nAlice,30' } },
+      buildContext()
+    )
+
+    expect(result).toEqual({ success: false, error: 'Table "tbl_1" not found' })
+    expect(mockReplaceTableRows).not.toHaveBeenCalled()
+  })
+
+  it('imports CSV content through the service with id-keyed rows', async () => {
+    const result = await maybeWriteReadCsvToTable(
+      ReadTool.id,
+      { outputTable: 'tbl_1', path: 'files/people.csv' },
+      { success: true, output: { content: 'name,age\nAlice,30\nBob,40' } },
+      buildContext()
+    )
+
+    expect(result.success).toBe(true)
+    const [data] = mockReplaceTableRows.mock.calls[0]
+    expect(data.rows).toEqual([
+      { col_name: 'Alice', col_age: '30' },
+      { col_name: 'Bob', col_age: '40' },
+    ])
+  })
+
+  it('fails fast when the file headers match no table columns', async () => {
+    const result = await maybeWriteReadCsvToTable(
+      ReadTool.id,
+      { outputTable: 'tbl_1', path: 'files/people.csv' },
+      { success: true, output: { content: 'wrong,headers\n1,2' } },
+      buildContext()
+    )
+
+    expect(result.success).toBe(false)
+    expect(result.error).toContain('Row 1 has no keys matching columns')
+    expect(mockReplaceTableRows).not.toHaveBeenCalled()
+  })
+
+  it('surfaces service validation failures as tool errors', async () => {
+    mockReplaceTableRows.mockRejectedValue(new Error('Row 1: name is required'))
+
+    const result = await maybeWriteReadCsvToTable(
+      ReadTool.id,
+      { outputTable: 'tbl_1', path: 'files/people.csv' },
+      { success: true, output: { content: 'age\n30' } },
+      buildContext()
+    )
+
+    expect(result.success).toBe(false)
+    expect(result.error).toContain('Row 1: name is required')
+  })
+})
@@ -1,10 +1,7 @@
-import { db } from '@sim/db'
-import { userTableRows } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { toError } from '@sim/utils/errors'
 import { generateId } from '@sim/utils/id'
 import { parse as csvParse } from 'csv-parse/sync'
-import { eq } from 'drizzle-orm'
 import { FunctionExecute, Read as ReadTool } from '@/lib/copilot/generated/tool-catalog-v1'
 import { CopilotTableOutcome } from '@/lib/copilot/generated/trace-attribute-values-v1'
 import { TraceAttr } from '@/lib/copilot/generated/trace-attributes-v1'
@@ -12,14 +9,45 @@ import { TraceEvent } from '@/lib/copilot/generated/trace-events-v1'
 import { TraceSpan } from '@/lib/copilot/generated/trace-spans-v1'
 import { withCopilotSpan } from '@/lib/copilot/request/otel'
 import type { ExecutionContext, ToolCallResult } from '@/lib/copilot/request/types'
-import type { RowData } from '@/lib/table'
-import { nKeysBetween } from '@/lib/table/order-key'
-import { buildOrderedRowValues, getTableById } from '@/lib/table/service'
+import type { RowData, TableDefinition } from '@/lib/table'
+import { buildIdByName, rowDataNameToId } from '@/lib/table/column-keys'
+import { getTableById, replaceTableRows } from '@/lib/table/service'

 const logger = createLogger('CopilotToolResultTables')

 const MAX_OUTPUT_TABLE_ROWS = 10_000
-const BATCH_CHUNK_SIZE = 500
+
+/**
+ * Replaces a table's rows with wire rows keyed by column name. Translates the
+ * keys to stable column ids (unknown keys are dropped, matching every other
+ * name-translating boundary) and delegates to `replaceTableRows`, which owns
+ * locking, validation, plan row limits, batching, and rowCount maintenance.
+ */
+async function replaceTableRowsFromWire(
+  table: TableDefinition,
+  rows: Array<Record<string, unknown>>,
+  context: ExecutionContext
+): Promise<{ error?: string }> {
+  const idByName = buildIdByName(table.schema)
+  const idKeyedRows = rows.map((row) => rowDataNameToId(row as RowData, idByName))
+  const emptyIndex = idKeyedRows.findIndex((row) => Object.keys(row).length === 0)
+  if (emptyIndex !== -1) {
+    return {
+      error: `Row ${emptyIndex + 1} has no keys matching columns on table "${table.name}" (columns: ${table.schema.columns.map((c) => c.name).join(', ')})`,
+    }
+  }
+  await replaceTableRows(
+    {
+      tableId: table.id,
+      rows: idKeyedRows,
+      workspaceId: table.workspaceId,
+      userId: context.userId,
+    },
+    table,
+    generateId().slice(0, 8)
+  )
+  return {}
+}

 export async function maybeWriteOutputToTable(
  toolName: string,
@@ -44,7 +72,7 @@ export async function maybeWriteOutputToTable(
    async (span) => {
      try {
        const table = await getTableById(outputTable)
-        if (!table) {
+        if (!table || table.workspaceId !== context.workspaceId) {
          span.setAttribute(TraceAttr.CopilotTableOutcome, CopilotTableOutcome.TableNotFound)
          return {
            success: false,
@@ -97,33 +125,11 @@ export async function maybeWriteOutputToTable(
        if (context.abortSignal?.aborted) {
          throw new Error('Request aborted before tool mutation could be applied')
        }
-        await db.transaction(async (tx) => {
-          if (context.abortSignal?.aborted) {
-            throw new Error('Request aborted before tool mutation could be applied')
-          }
-          await tx.delete(userTableRows).where(eq(userTableRows.tableId, outputTable))
-
-          const now = new Date()
-          // Replace-all: table was just cleared — mint a fresh contiguous key run.
-          const orderKeys = nKeysBetween(null, null, rows.length)
-          for (let i = 0; i < rows.length; i += BATCH_CHUNK_SIZE) {
-            if (context.abortSignal?.aborted) {
-              throw new Error('Request aborted before tool mutation could be applied')
-            }
-            const chunk = rows.slice(i, i + BATCH_CHUNK_SIZE)
-            const values = buildOrderedRowValues({
-              tableId: outputTable,
-              workspaceId: context.workspaceId!,
-              rows: chunk as RowData[],
-              startPosition: i,
-              orderKeys: orderKeys.slice(i, i + BATCH_CHUNK_SIZE),
-              now,
-              createdBy: context.userId,
-              makeId: () => `row_${generateId().replace(/-/g, '')}`,
-            })
-            await tx.insert(userTableRows).values(values)
-          }
-        })
+        const replaceResult = await replaceTableRowsFromWire(table, rows, context)
+        if (replaceResult.error) {
+          span.setAttribute(TraceAttr.CopilotTableOutcome, CopilotTableOutcome.InvalidShape)
+          return { success: false, error: replaceResult.error }
+        }

        logger.info('Tool output written to table', {
          toolName,
@@ -181,7 +187,7 @@ export async function maybeWriteReadCsvToTable(
    async (span) => {
      try {
        const table = await getTableById(outputTable)
-        if (!table) {
+        if (!table || table.workspaceId !== context.workspaceId) {
          span.setAttribute(TraceAttr.CopilotTableOutcome, CopilotTableOutcome.TableNotFound)
          return { success: false, error: `Table "${outputTable}" not found` }
        }
@@ -243,33 +249,11 @@ export async function maybeWriteReadCsvToTable(
        if (context.abortSignal?.aborted) {
          throw new Error('Request aborted before tool mutation could be applied')
        }
-        await db.transaction(async (tx) => {
-          if (context.abortSignal?.aborted) {
-            throw new Error('Request aborted before tool mutation could be applied')
-          }
-          await tx.delete(userTableRows).where(eq(userTableRows.tableId, outputTable))
-
-          const now = new Date()
-          // Replace-all: table was just cleared — mint a fresh contiguous key run.
-          const orderKeys = nKeysBetween(null, null, rows.length)
-          for (let i = 0; i < rows.length; i += BATCH_CHUNK_SIZE) {
-            if (context.abortSignal?.aborted) {
-              throw new Error('Request aborted before tool mutation could be applied')
-            }
-            const chunk = rows.slice(i, i + BATCH_CHUNK_SIZE)
-            const values = buildOrderedRowValues({
-              tableId: outputTable,
-              workspaceId: context.workspaceId!,
-              rows: chunk as RowData[],
-              startPosition: i,
-              orderKeys: orderKeys.slice(i, i + BATCH_CHUNK_SIZE),
-              now,
-              createdBy: context.userId,
-              makeId: () => `row_${generateId().replace(/-/g, '')}`,
-            })
-            await tx.insert(userTableRows).values(values)
-          }
-        })
+        const replaceResult = await replaceTableRowsFromWire(table, rows, context)
+        if (replaceResult.error) {
+          span.setAttribute(TraceAttr.CopilotTableOutcome, CopilotTableOutcome.InvalidShape)
+          return { success: false, error: replaceResult.error }
+        }

        logger.info('Read output written to table', {
          toolName,