Skip to content

Commit c3c3416

Browse files
TheodoreSpeakswaleedlatif1Sg312icecrasher321claude
authored
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>
1 parent fb9e481 commit c3c3416

2 files changed

Lines changed: 268 additions & 63 deletions

File tree

‎apps/sim/lib/copilot/request/tools/tables.test.ts‎

Lines changed: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,221 @@
1+
/**
2+
* @vitest-environment node
3+
*/
4+
5+
import { beforeEach, describe, expect, it, vi } from 'vitest'
6+
import type { TableDefinition } from '@/lib/table'
7+
8+
const { mockGetTableById, mockReplaceTableRows } = vi.hoisted(() => ({
9+
mockGetTableById: vi.fn(),
10+
mockReplaceTableRows: vi.fn(),
11+
}))
12+
13+
vi.mock('@/lib/table/service', () => ({
14+
getTableById: mockGetTableById,
15+
replaceTableRows: mockReplaceTableRows,
16+
}))
17+
18+
vi.mock('@/lib/copilot/request/otel', () => ({
19+
withCopilotSpan: (
20+
_name: string,
21+
_attrs: Record<string, unknown> | undefined,
22+
fn: (span: unknown) => Promise<unknown>
23+
) => fn({ setAttribute: vi.fn(), setAttributes: vi.fn(), addEvent: vi.fn() }),
24+
}))
25+
26+
import { FunctionExecute, Read as ReadTool } from '@/lib/copilot/generated/tool-catalog-v1'
27+
import {
28+
maybeWriteOutputToTable,
29+
maybeWriteReadCsvToTable,
30+
} from '@/lib/copilot/request/tools/tables'
31+
import type { ExecutionContext } from '@/lib/copilot/request/types'
32+
33+
function buildTable(overrides: Partial<TableDefinition> = {}): TableDefinition {
34+
return {
35+
id: 'tbl_1',
36+
name: 'People',
37+
description: null,
38+
schema: {
39+
columns: [
40+
{ id: 'col_name', name: 'name', type: 'string' },
41+
{ id: 'col_age', name: 'age', type: 'number' },
42+
],
43+
},
44+
metadata: null,
45+
rowCount: 0,
46+
maxRows: 100,
47+
workspaceId: 'workspace-1',
48+
createdBy: 'user-1',
49+
archivedAt: null,
50+
createdAt: new Date('2024-01-01'),
51+
updatedAt: new Date('2024-01-01'),
52+
...overrides,
53+
} as TableDefinition
54+
}
55+
56+
function buildContext(overrides: Partial<ExecutionContext> = {}): ExecutionContext {
57+
return {
58+
userId: 'user-1',
59+
workflowId: 'wf-1',
60+
workspaceId: 'workspace-1',
61+
...overrides,
62+
}
63+
}
64+
65+
describe('maybeWriteOutputToTable', () => {
66+
beforeEach(() => {
67+
vi.clearAllMocks()
68+
mockGetTableById.mockResolvedValue(buildTable())
69+
mockReplaceTableRows.mockResolvedValue({ deletedCount: 0, insertedCount: 2 })
70+
})
71+
72+
it('rejects a table from another workspace without touching it', async () => {
73+
mockGetTableById.mockResolvedValue(buildTable({ workspaceId: 'other-workspace' }))
74+
75+
const result = await maybeWriteOutputToTable(
76+
FunctionExecute.id,
77+
{ outputTable: 'tbl_1' },
78+
{ success: true, output: { result: [{ name: 'Alice' }] } },
79+
buildContext()
80+
)
81+
82+
expect(result).toEqual({ success: false, error: 'Table "tbl_1" not found' })
83+
expect(mockReplaceTableRows).not.toHaveBeenCalled()
84+
})
85+
86+
it('replaces rows through the service with name keys remapped to column ids', async () => {
87+
const result = await maybeWriteOutputToTable(
88+
FunctionExecute.id,
89+
{ outputTable: 'tbl_1' },
90+
{
91+
success: true,
92+
output: {
93+
result: [
94+
{ name: 'Alice', age: 30 },
95+
{ name: 'Bob', age: 40 },
96+
],
97+
},
98+
},
99+
buildContext()
100+
)
101+
102+
expect(result.success).toBe(true)
103+
expect(mockReplaceTableRows).toHaveBeenCalledTimes(1)
104+
const [data, table] = mockReplaceTableRows.mock.calls[0]
105+
expect(data).toMatchObject({
106+
tableId: 'tbl_1',
107+
workspaceId: 'workspace-1',
108+
userId: 'user-1',
109+
rows: [
110+
{ col_name: 'Alice', col_age: 30 },
111+
{ col_name: 'Bob', col_age: 40 },
112+
],
113+
})
114+
expect(table.id).toBe('tbl_1')
115+
})
116+
117+
it('fails fast when no row keys match the table columns', async () => {
118+
const result = await maybeWriteOutputToTable(
119+
FunctionExecute.id,
120+
{ outputTable: 'tbl_1' },
121+
{ success: true, output: { result: [{ wrong: 1 }, { keys: 2 }] } },
122+
buildContext()
123+
)
124+
125+
expect(result.success).toBe(false)
126+
expect(result.error).toContain('Row 1 has no keys matching columns')
127+
expect(mockReplaceTableRows).not.toHaveBeenCalled()
128+
})
129+
130+
it('fails fast when only some rows match instead of writing empty rows', async () => {
131+
const result = await maybeWriteOutputToTable(
132+
FunctionExecute.id,
133+
{ outputTable: 'tbl_1' },
134+
{ success: true, output: { result: [{ name: 'Alice' }, { wrong: 'x' }] } },
135+
buildContext()
136+
)
137+
138+
expect(result.success).toBe(false)
139+
expect(result.error).toContain('Row 2 has no keys matching columns')
140+
expect(mockReplaceTableRows).not.toHaveBeenCalled()
141+
})
142+
143+
it('surfaces service validation failures as tool errors', async () => {
144+
mockReplaceTableRows.mockRejectedValue(new Error('Row 1: name is required'))
145+
146+
const result = await maybeWriteOutputToTable(
147+
FunctionExecute.id,
148+
{ outputTable: 'tbl_1' },
149+
{ success: true, output: { result: [{ age: 30 }] } },
150+
buildContext()
151+
)
152+
153+
expect(result.success).toBe(false)
154+
expect(result.error).toContain('Row 1: name is required')
155+
})
156+
})
157+
158+
describe('maybeWriteReadCsvToTable', () => {
159+
beforeEach(() => {
160+
vi.clearAllMocks()
161+
mockGetTableById.mockResolvedValue(buildTable())
162+
mockReplaceTableRows.mockResolvedValue({ deletedCount: 0, insertedCount: 2 })
163+
})
164+
165+
it('rejects a table from another workspace without touching it', async () => {
166+
mockGetTableById.mockResolvedValue(buildTable({ workspaceId: 'other-workspace' }))
167+
168+
const result = await maybeWriteReadCsvToTable(
169+
ReadTool.id,
170+
{ outputTable: 'tbl_1', path: 'files/people.csv' },
171+
{ success: true, output: { content: 'name,age\nAlice,30' } },
172+
buildContext()
173+
)
174+
175+
expect(result).toEqual({ success: false, error: 'Table "tbl_1" not found' })
176+
expect(mockReplaceTableRows).not.toHaveBeenCalled()
177+
})
178+
179+
it('imports CSV content through the service with id-keyed rows', async () => {
180+
const result = await maybeWriteReadCsvToTable(
181+
ReadTool.id,
182+
{ outputTable: 'tbl_1', path: 'files/people.csv' },
183+
{ success: true, output: { content: 'name,age\nAlice,30\nBob,40' } },
184+
buildContext()
185+
)
186+
187+
expect(result.success).toBe(true)
188+
const [data] = mockReplaceTableRows.mock.calls[0]
189+
expect(data.rows).toEqual([
190+
{ col_name: 'Alice', col_age: '30' },
191+
{ col_name: 'Bob', col_age: '40' },
192+
])
193+
})
194+
195+
it('fails fast when the file headers match no table columns', async () => {
196+
const result = await maybeWriteReadCsvToTable(
197+
ReadTool.id,
198+
{ outputTable: 'tbl_1', path: 'files/people.csv' },
199+
{ success: true, output: { content: 'wrong,headers\n1,2' } },
200+
buildContext()
201+
)
202+
203+
expect(result.success).toBe(false)
204+
expect(result.error).toContain('Row 1 has no keys matching columns')
205+
expect(mockReplaceTableRows).not.toHaveBeenCalled()
206+
})
207+
208+
it('surfaces service validation failures as tool errors', async () => {
209+
mockReplaceTableRows.mockRejectedValue(new Error('Row 1: name is required'))
210+
211+
const result = await maybeWriteReadCsvToTable(
212+
ReadTool.id,
213+
{ outputTable: 'tbl_1', path: 'files/people.csv' },
214+
{ success: true, output: { content: 'age\n30' } },
215+
buildContext()
216+
)
217+
218+
expect(result.success).toBe(false)
219+
expect(result.error).toContain('Row 1: name is required')
220+
})
221+
})

‎apps/sim/lib/copilot/request/tools/tables.ts‎

Lines changed: 47 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,53 @@
1-
import { db } from '@sim/db'
2-
import { userTableRows } from '@sim/db/schema'
31
import { createLogger } from '@sim/logger'
42
import { toError } from '@sim/utils/errors'
53
import { generateId } from '@sim/utils/id'
64
import { parse as csvParse } from 'csv-parse/sync'
7-
import { eq } from 'drizzle-orm'
85
import { FunctionExecute, Read as ReadTool } from '@/lib/copilot/generated/tool-catalog-v1'
96
import { CopilotTableOutcome } from '@/lib/copilot/generated/trace-attribute-values-v1'
107
import { TraceAttr } from '@/lib/copilot/generated/trace-attributes-v1'
118
import { TraceEvent } from '@/lib/copilot/generated/trace-events-v1'
129
import { TraceSpan } from '@/lib/copilot/generated/trace-spans-v1'
1310
import { withCopilotSpan } from '@/lib/copilot/request/otel'
1411
import type { ExecutionContext, ToolCallResult } from '@/lib/copilot/request/types'
15-
import type { RowData } from '@/lib/table'
16-
import { nKeysBetween } from '@/lib/table/order-key'
17-
import { buildOrderedRowValues, getTableById } from '@/lib/table/service'
12+
import type { RowData, TableDefinition } from '@/lib/table'
13+
import { buildIdByName, rowDataNameToId } from '@/lib/table/column-keys'
14+
import { getTableById, replaceTableRows } from '@/lib/table/service'
1815

1916
const logger = createLogger('CopilotToolResultTables')
2017

2118
const MAX_OUTPUT_TABLE_ROWS = 10_000
22-
const BATCH_CHUNK_SIZE = 500
19+
20+
/**
21+
* Replaces a table's rows with wire rows keyed by column name. Translates the
22+
* keys to stable column ids (unknown keys are dropped, matching every other
23+
* name-translating boundary) and delegates to `replaceTableRows`, which owns
24+
* locking, validation, plan row limits, batching, and rowCount maintenance.
25+
*/
26+
async function replaceTableRowsFromWire(
27+
table: TableDefinition,
28+
rows: Array<Record<string, unknown>>,
29+
context: ExecutionContext
30+
): Promise<{ error?: string }> {
31+
const idByName = buildIdByName(table.schema)
32+
const idKeyedRows = rows.map((row) => rowDataNameToId(row as RowData, idByName))
33+
const emptyIndex = idKeyedRows.findIndex((row) => Object.keys(row).length === 0)
34+
if (emptyIndex !== -1) {
35+
return {
36+
error: `Row ${emptyIndex + 1} has no keys matching columns on table "${table.name}" (columns: ${table.schema.columns.map((c) => c.name).join(', ')})`,
37+
}
38+
}
39+
await replaceTableRows(
40+
{
41+
tableId: table.id,
42+
rows: idKeyedRows,
43+
workspaceId: table.workspaceId,
44+
userId: context.userId,
45+
},
46+
table,
47+
generateId().slice(0, 8)
48+
)
49+
return {}
50+
}
2351

2452
export async function maybeWriteOutputToTable(
2553
toolName: string,
@@ -44,7 +72,7 @@ export async function maybeWriteOutputToTable(
4472
async (span) => {
4573
try {
4674
const table = await getTableById(outputTable)
47-
if (!table) {
75+
if (!table || table.workspaceId !== context.workspaceId) {
4876
span.setAttribute(TraceAttr.CopilotTableOutcome, CopilotTableOutcome.TableNotFound)
4977
return {
5078
success: false,
@@ -97,33 +125,11 @@ export async function maybeWriteOutputToTable(
97125
if (context.abortSignal?.aborted) {
98126
throw new Error('Request aborted before tool mutation could be applied')
99127
}
100-
await db.transaction(async (tx) => {
101-
if (context.abortSignal?.aborted) {
102-
throw new Error('Request aborted before tool mutation could be applied')
103-
}
104-
await tx.delete(userTableRows).where(eq(userTableRows.tableId, outputTable))
105-
106-
const now = new Date()
107-
// Replace-all: table was just cleared — mint a fresh contiguous key run.
108-
const orderKeys = nKeysBetween(null, null, rows.length)
109-
for (let i = 0; i < rows.length; i += BATCH_CHUNK_SIZE) {
110-
if (context.abortSignal?.aborted) {
111-
throw new Error('Request aborted before tool mutation could be applied')
112-
}
113-
const chunk = rows.slice(i, i + BATCH_CHUNK_SIZE)
114-
const values = buildOrderedRowValues({
115-
tableId: outputTable,
116-
workspaceId: context.workspaceId!,
117-
rows: chunk as RowData[],
118-
startPosition: i,
119-
orderKeys: orderKeys.slice(i, i + BATCH_CHUNK_SIZE),
120-
now,
121-
createdBy: context.userId,
122-
makeId: () => `row_${generateId().replace(/-/g, '')}`,
123-
})
124-
await tx.insert(userTableRows).values(values)
125-
}
126-
})
128+
const replaceResult = await replaceTableRowsFromWire(table, rows, context)
129+
if (replaceResult.error) {
130+
span.setAttribute(TraceAttr.CopilotTableOutcome, CopilotTableOutcome.InvalidShape)
131+
return { success: false, error: replaceResult.error }
132+
}
127133

128134
logger.info('Tool output written to table', {
129135
toolName,
@@ -181,7 +187,7 @@ export async function maybeWriteReadCsvToTable(
181187
async (span) => {
182188
try {
183189
const table = await getTableById(outputTable)
184-
if (!table) {
190+
if (!table || table.workspaceId !== context.workspaceId) {
185191
span.setAttribute(TraceAttr.CopilotTableOutcome, CopilotTableOutcome.TableNotFound)
186192
return { success: false, error: `Table "${outputTable}" not found` }
187193
}
@@ -243,33 +249,11 @@ export async function maybeWriteReadCsvToTable(
243249
if (context.abortSignal?.aborted) {
244250
throw new Error('Request aborted before tool mutation could be applied')
245251
}
246-
await db.transaction(async (tx) => {
247-
if (context.abortSignal?.aborted) {
248-
throw new Error('Request aborted before tool mutation could be applied')
249-
}
250-
await tx.delete(userTableRows).where(eq(userTableRows.tableId, outputTable))
251-
252-
const now = new Date()
253-
// Replace-all: table was just cleared — mint a fresh contiguous key run.
254-
const orderKeys = nKeysBetween(null, null, rows.length)
255-
for (let i = 0; i < rows.length; i += BATCH_CHUNK_SIZE) {
256-
if (context.abortSignal?.aborted) {
257-
throw new Error('Request aborted before tool mutation could be applied')
258-
}
259-
const chunk = rows.slice(i, i + BATCH_CHUNK_SIZE)
260-
const values = buildOrderedRowValues({
261-
tableId: outputTable,
262-
workspaceId: context.workspaceId!,
263-
rows: chunk as RowData[],
264-
startPosition: i,
265-
orderKeys: orderKeys.slice(i, i + BATCH_CHUNK_SIZE),
266-
now,
267-
createdBy: context.userId,
268-
makeId: () => `row_${generateId().replace(/-/g, '')}`,
269-
})
270-
await tx.insert(userTableRows).values(values)
271-
}
272-
})
252+
const replaceResult = await replaceTableRowsFromWire(table, rows, context)
253+
if (replaceResult.error) {
254+
span.setAttribute(TraceAttr.CopilotTableOutcome, CopilotTableOutcome.InvalidShape)
255+
return { success: false, error: replaceResult.error }
256+
}
273257

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

0 commit comments

Comments
 (0)