fix(tables): snapshot every sandbox mount

Author: TheodoreSpeaks

apps/sim/lib/copilot/tools/handlers/function-execute.test.ts

@@ -13,7 +13,6 @@ import {
 const {
   mockGetTableById,
   mockListTables,
-  mockQueryRows,
   mockGetOrCreateTableSnapshot,
   mockDownloadFile,
   mockGeneratePresignedDownloadUrl,
@@ -32,12 +31,10 @@ const {
   mockMaterializeCopilotCodeSecrets,
   mockHasWorkspaceSandboxAccess,
   mockImportWorkspaceFileSecretProvenanceForRuntime,
-  mockLoadTableRowSecretProvenance,
   mockIsTableSnapshotSafeForModelMount,
 } = vi.hoisted(() => ({
   mockGetTableById: vi.fn(),
   mockListTables: vi.fn(),
-  mockQueryRows: vi.fn(),
   mockGetOrCreateTableSnapshot: vi.fn(),
   mockDownloadFile: vi.fn(),
   mockGeneratePresignedDownloadUrl: vi.fn(),
@@ -56,7 +53,6 @@ const {
   mockMaterializeCopilotCodeSecrets: vi.fn(),
   mockHasWorkspaceSandboxAccess: vi.fn(),
   mockImportWorkspaceFileSecretProvenanceForRuntime: vi.fn(),
-  mockLoadTableRowSecretProvenance: vi.fn(),
   mockIsTableSnapshotSafeForModelMount: vi.fn(),
 }))
 
@@ -65,10 +61,8 @@ vi.mock('@/lib/table/service', () => ({
   getTableById: mockGetTableById,
   listTables: mockListTables,
 }))
-vi.mock('@/lib/table/rows/service', () => ({ queryRows: mockQueryRows }))
 vi.mock('@/lib/table/rows/secret-provenance', () => ({
   isTableSnapshotSafeForModelMount: mockIsTableSnapshotSafeForModelMount,
-  loadTableRowSecretProvenance: mockLoadTableRowSecretProvenance,
 }))
 vi.mock('@/lib/table/snapshot-cache', () => ({
   getOrCreateTableSnapshot: mockGetOrCreateTableSnapshot,
@@ -125,13 +119,12 @@ vi.mock('@/lib/execution/remote-sandbox/workspace-sandboxes', () => ({
 import { projectToolResultForCopilot } from '@/lib/copilot/request/tools/resolved-secret-result'
 import { executeFunctionExecute } from '@/lib/copilot/tools/handlers/function-execute'
 import { executeRunCode } from '@/lib/copilot/tools/handlers/run-code'
-import { TABLE_LIMITS } from '@/lib/table/constants'
 import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry'
 
 const table = {
   id: 'tbl_1',
   workspaceId: 'ws_1',
-  rowCount: TABLE_LIMITS.DEFAULT_QUERY_LIMIT + 1,
+  rowCount: 1,
   schema: { columns: [{ id: 'col_name', name: 'name', type: 'string' }] },
 }
 
@@ -153,13 +146,7 @@ function resetExecutionMocks(): void {
   vi.clearAllMocks()
   mockExecuteTool.mockReset()
   mockMaterializeCopilotCodeSecrets.mockReset()
-  mockLoadTableRowSecretProvenance.mockReset()
   mockIsTableSnapshotSafeForModelMount.mockReset()
-  mockLoadTableRowSecretProvenance.mockResolvedValue({
-    version: 1,
-    complete: true,
-    entries: [],
-  })
   mockIsTableSnapshotSafeForModelMount.mockResolvedValue(true)
   mockListWorkspaceFiles.mockResolvedValue([])
   mockListWorkspaceFileFolders.mockResolvedValue([])
@@ -572,76 +559,12 @@ describe('executeFunctionExecute table mounts', () => {
     resetExecutionMocks()
     mockExecuteTool.mockResolvedValue({ success: true })
     mockGetTableById.mockResolvedValue(table)
-    // Row data is keyed by stable column id at rest, not display name.
-    mockQueryRows.mockResolvedValue({ rows: [{ data: { col_name: 'Ada' } }] })
     mockHasCloudStorage.mockReturnValue(true)
     mockGeneratePresignedDownloadUrl.mockResolvedValue('https://s3.example/presigned?sig=abc')
   })
 
-  it('mounts a table at the inline row limit without truncation', async () => {
-    mockGetTableById.mockResolvedValue({ ...table, rowCount: TABLE_LIMITS.DEFAULT_QUERY_LIMIT })
-
-    await executeFunctionExecute({ inputTables: ['tbl_1'] }, context as never)
-
-    expect(mockQueryRows).toHaveBeenCalledWith(
-      expect.objectContaining({ id: 'tbl_1' }),
-      { limit: TABLE_LIMITS.DEFAULT_QUERY_LIMIT },
-      'copilot-fn-exec'
-    )
-    expect(mockGetOrCreateTableSnapshot).not.toHaveBeenCalled()
-    const files = mountedFiles()
-    expect(files[0].path).toBe('/home/user/tables/tbl_1.csv')
-    expect(files[0].content).toBe('name\nAda')
-  })
-
-  it('mounts CSV with display-name headers and id-keyed values, never column ids', async () => {
-    mockGetTableById.mockResolvedValue({
-      id: 'tbl_2',
-      workspaceId: 'ws_1',
-      rowCount: 2,
-      schema: {
-        columns: [
-          { id: 'col_name', name: 'name', type: 'string' },
-          { id: 'col_company', name: 'company', type: 'string' },
-        ],
-      },
-    })
-    mockQueryRows.mockResolvedValue({
-      rows: [
-        { data: { col_name: 'Ada', col_company: 'Analytical Engine' } },
-        { data: { col_name: 'Grace', col_company: 'Navy, Inc' } },
-      ],
-    })
-
-    await executeFunctionExecute({ inputTables: ['tbl_2'] }, context as never)
-
-    const csv = mountedFiles()[0].content as string
-    const lines = csv.split('\n')
-    expect(lines[0]).toBe('name,company')
-    expect(lines[1]).toBe('Ada,Analytical Engine')
-    // Value containing a comma is quoted.
-    expect(lines[2]).toBe('Grace,"Navy, Inc"')
-    // No stable column id leaks into the mounted file.
-    expect(csv).not.toContain('col_name')
-    expect(csv).not.toContain('col_company')
-  })
-
-  it('reads values by column id for legacy name-keyed rows too', async () => {
-    // Legacy column with no id: getColumnId falls back to name, so name-keyed data is correct.
-    mockGetTableById.mockResolvedValue({
-      id: 'tbl_legacy',
-      workspaceId: 'ws_1',
-      rowCount: 1,
-      schema: { columns: [{ name: 'email', type: 'string' }] },
-    })
-    mockQueryRows.mockResolvedValue({ rows: [{ data: { email: 'a@b.com' } }] })
-
-    await executeFunctionExecute({ inputTables: ['tbl_legacy'] }, context as never)
-
-    expect(mountedFiles()[0].content).toBe('email\na@b.com')
-  })
-
-  it('mounts a table above the inline row limit by presigned snapshot URL', async () => {
+  it('mounts every table by presigned snapshot URL', async () => {
+    mockGetTableById.mockResolvedValue({ ...table, rowCount: 0 })
     mockGetOrCreateTableSnapshot.mockResolvedValue({
       key: 'table-snapshots/ws_1/tbl_1/v5.csv',
       size: 9,
@@ -651,7 +574,6 @@ describe('executeFunctionExecute table mounts', () => {
     await executeFunctionExecute({ inputTables: ['tbl_1'] }, context as never)
 
     expect(mockGetOrCreateTableSnapshot).toHaveBeenCalledTimes(1)
-    expect(mockQueryRows).not.toHaveBeenCalled()
     expect(mockDownloadFile).not.toHaveBeenCalled()
     expect(mockGeneratePresignedDownloadUrl).toHaveBeenCalledWith(
       'table-snapshots/ws_1/tbl_1/v5.csv',
@@ -715,35 +637,6 @@ describe('executeFunctionExecute table mounts', () => {
     expect(projectToolResultForCopilot(result, parentRegistry)).toEqual({ success: true })
   })
 
-  it('unknown inline row provenance still mounts and taints model egress', async () => {
-    mockGetTableById.mockResolvedValue({ ...table, rowCount: TABLE_LIMITS.DEFAULT_QUERY_LIMIT })
-    mockLoadTableRowSecretProvenance.mockResolvedValue({
-      version: 1,
-      complete: false,
-      entries: [],
-    })
-    mockExecuteTool.mockResolvedValue({ success: true, output: { result: 'raw output' } })
-    const parentRegistry = new ResolvedSecretTraceRegistry([], {
-      userId: 'u1',
-      workspaceId: 'ws_1',
-    })
-
-    const result = await executeFunctionExecute(
-      { inputTables: ['tbl_1'] },
-      { ...context, resolvedSecretTraceRegistry: parentRegistry }
-    )
-
-    expect(mountedFiles()[0].content).toBe('name\nAda')
-    expect(mockExecuteTool.mock.calls[0]?.[1]?.[PRIVATE_SECRET_PROVENANCE_FIELD]).toEqual({
-      version: 1,
-      complete: false,
-      selections: [],
-    })
-    expect(result).toEqual({ success: true, output: { result: 'raw output' } })
-    expect(parentRegistry.isComplete()).toBe(false)
-    expect(projectToolResultForCopilot(result, parentRegistry)).toEqual({ success: true })
-  })
-
   it('throws when a cloud snapshot exceeds the table mount limit', async () => {
     mockGetOrCreateTableSnapshot.mockResolvedValue({
       key: 'table-snapshots/ws_1/tbl_1/v5.csv',

apps/sim/lib/copilot/tools/handlers/function-execute.ts

@@ -11,7 +11,6 @@ import {
   materializeCopilotCodeSecrets,
 } from '@/lib/copilot/tools/secret-mount-materializer.server'
 import { decodeVfsPathSegments, encodeVfsPathSegments } from '@/lib/copilot/vfs/path-utils'
-import { neutralizeCsvFormula, toCsvRow } from '@/lib/core/utils/csv'
 import { isPayloadSizeLimitError } from '@/lib/core/utils/stream-limits'
 import type { PrivateSecretProvenanceBundleV1 } from '@/lib/execution/model-input-provenance'
 import {
@@ -20,14 +19,7 @@ import {
 } from '@/lib/execution/private-tool-metadata'
 import { MAX_PLAN_REQUIRED } from '@/lib/execution/remote-sandbox/workspace-sandboxes'
 import { recordSecretUsage } from '@/lib/secrets/usage/record'
-import { getColumnId } from '@/lib/table/column-keys'
-import { TABLE_LIMITS } from '@/lib/table/constants'
-import { formatCsvCell } from '@/lib/table/export-format'
-import {
-  isTableSnapshotSafeForModelMount,
-  loadTableRowSecretProvenance,
-} from '@/lib/table/rows/secret-provenance'
-import { queryRows } from '@/lib/table/rows/service'
+import { isTableSnapshotSafeForModelMount } from '@/lib/table/rows/secret-provenance'
 import { getTableById, listTables } from '@/lib/table/service'
 import { getOrCreateTableSnapshot, SNAPSHOT_MAX_BYTES } from '@/lib/table/snapshot-cache'
 import {
@@ -310,7 +302,6 @@ export async function resolveInputFiles(
   inputFiles?: unknown[],
   inputTables?: unknown[],
   inputDirectories?: unknown[],
-  provenanceUserId?: string,
   resolvedSecretTraceRegistry?: ResolvedSecretTraceRegistry,
   filePrincipal?: Principal
 ): Promise<SandboxFile[]> {
@@ -487,108 +478,61 @@ export async function resolveInputFiles(
           : undefined
       const mountPath = sandboxPath || `/home/user/tables/${table.id}.csv`
 
-      /** Tables that exceed the bounded inline query mount from a complete versioned snapshot. */
-      if (table.rowCount > TABLE_LIMITS.DEFAULT_QUERY_LIMIT) {
-        const snapshot = await getOrCreateTableSnapshot(table, 'copilot-fn-exec')
-        if (!resolvedSecretTraceRegistry) {
-          throw new Error(
-            `Input table "${tableId}" cannot be mounted because its secret provenance is unavailable.`
-          )
-        }
-        try {
-          const safeForModelMount = await isTableSnapshotSafeForModelMount({
-            tableId: table.id,
-            workspaceId,
-            rowsVersion: snapshot.version,
-          })
-          if (!safeForModelMount)
-            resolvedSecretTraceRegistry.markIncomplete('table-snapshot-unsafe-for-mount')
-        } catch {
-          resolvedSecretTraceRegistry.markIncomplete('table-snapshot-unsafe-for-mount')
-        }
-
-        if (hasCloudStorage()) {
-          // Mount by reference: the sandbox fetches the snapshot straight from storage via a
-          // presigned URL, so the bytes never pass through the web process — the only ceiling is
-          // sandbox disk (enforced at materialization by SNAPSHOT_MAX_BYTES).
-          if (snapshot.size > SNAPSHOT_MAX_BYTES) {
-            throw new Error(
-              `Input table "${tableId}" is ${Math.round(snapshot.size / 1024 / 1024)}MB, over the ${SNAPSHOT_MAX_BYTES / 1024 / 1024}MB table mount limit.`
-            )
-          }
-          const url = await generatePresignedDownloadUrl(
-            snapshot.key,
-            'execution',
-            MOUNT_URL_TTL_SECONDS
-          )
-          sandboxFiles.push({ type: 'url', path: mountPath, url })
-          continue
-        }
-
-        // Local storage: a presigned URL is an app-internal serve path a remote sandbox can't
-        // reach, so fall back to buffering the bytes through the web process (file-mount guards).
-        if (snapshot.size > MAX_FILE_SIZE) {
-          throw new Error(
-            `Input table "${tableId}" is ${Math.round(snapshot.size / 1024 / 1024)}MB, over the ${MAX_FILE_SIZE / 1024 / 1024}MB per-file mount limit.`
-          )
-        }
-        if (mounted.buffered + snapshot.size > MAX_TOTAL_SIZE) {
-          throw new Error(
-            `Mounting "${tableId}" would exceed the ${MAX_TOTAL_SIZE / 1024 / 1024}MB total mount limit. Mount fewer or smaller tables.`
-          )
-        }
-        const buffer = await downloadFile({
-          key: snapshot.key,
-          context: 'execution',
-          maxBytes: MAX_FILE_SIZE,
-        })
-        mounted.buffered += buffer.length
-        sandboxFiles.push({ path: mountPath, content: buffer.toString('utf-8') })
-        continue
-      }
-
-      /** The table fits completely within the bounded inline query. */
-      const rows = await queryRows(
-        table,
-        { limit: TABLE_LIMITS.DEFAULT_QUERY_LIMIT },
-        'copilot-fn-exec'
-      )
+      const snapshot = await getOrCreateTableSnapshot(table, 'copilot-fn-exec')
       if (!resolvedSecretTraceRegistry) {
         throw new Error(
           `Input table "${tableId}" cannot be mounted because its secret provenance is unavailable.`
         )
       }
       try {
-        const provenance = await loadTableRowSecretProvenance(rows.rows, {
-          userId: provenanceUserId ?? 'opaque-model-mount',
+        const safeForModelMount = await isTableSnapshotSafeForModelMount({
+          tableId: table.id,
           workspaceId,
+          rowsVersion: snapshot.version,
         })
-        if (
-          !provenance.complete ||
-          !(await resolvedSecretTraceRegistry.importProvenance(provenance, {
-            trusted: true,
-            origin: 'copilotFunctionExecute.result',
-          }))
-        ) {
-          resolvedSecretTraceRegistry.markIncomplete('source-provenance-incomplete', {
-            origin: 'copilotFunctionExecute.result',
-          })
-        }
+        if (!safeForModelMount)
+          resolvedSecretTraceRegistry.markIncomplete('table-snapshot-unsafe-for-mount')
       } catch {
-        resolvedSecretTraceRegistry.markIncomplete('source-provenance-incomplete', {
-          origin: 'copilotFunctionExecute.result',
-        })
+        resolvedSecretTraceRegistry.markIncomplete('table-snapshot-unsafe-for-mount')
       }
 
-      const columns = table.schema.columns
-      const csvLines = [toCsvRow(columns.map((column) => neutralizeCsvFormula(column.name)))]
-      for (const row of rows.rows) {
-        csvLines.push(
-          toCsvRow(columns.map((column) => formatCsvCell(column, row.data[getColumnId(column)])))
+      if (hasCloudStorage()) {
+        // Mount by reference: the sandbox fetches the snapshot straight from storage via a
+        // presigned URL, so the bytes never pass through the web process — the only ceiling is
+        // sandbox disk (enforced at materialization by SNAPSHOT_MAX_BYTES).
+        if (snapshot.size > SNAPSHOT_MAX_BYTES) {
+          throw new Error(
+            `Input table "${tableId}" is ${Math.round(snapshot.size / 1024 / 1024)}MB, over the ${SNAPSHOT_MAX_BYTES / 1024 / 1024}MB table mount limit.`
+          )
+        }
+        const url = await generatePresignedDownloadUrl(
+          snapshot.key,
+          'execution',
+          MOUNT_URL_TTL_SECONDS
+        )
+        sandboxFiles.push({ type: 'url', path: mountPath, url })
+        continue
+      }
+
+      // Local storage: a presigned URL is an app-internal serve path a remote sandbox can't
+      // reach, so fall back to buffering the bytes through the web process (file-mount guards).
+      if (snapshot.size > MAX_FILE_SIZE) {
+        throw new Error(
+          `Input table "${tableId}" is ${Math.round(snapshot.size / 1024 / 1024)}MB, over the ${MAX_FILE_SIZE / 1024 / 1024}MB per-file mount limit.`
         )
       }
-      const csvContent = csvLines.join('\n')
-      sandboxFiles.push({ path: mountPath, content: csvContent })
+      if (mounted.buffered + snapshot.size > MAX_TOTAL_SIZE) {
+        throw new Error(
+          `Mounting "${tableId}" would exceed the ${MAX_TOTAL_SIZE / 1024 / 1024}MB total mount limit. Mount fewer or smaller tables.`
+        )
+      }
+      const buffer = await downloadFile({
+        key: snapshot.key,
+        context: 'execution',
+        maxBytes: MAX_FILE_SIZE,
+      })
+      mounted.buffered += buffer.length
+      sandboxFiles.push({ path: mountPath, content: buffer.toString('utf-8') })
     }
   }
 
@@ -727,7 +671,6 @@ export async function executeFunctionExecute(
           inputFiles,
           inputTables,
           inputDirectories,
-          secretActorUserId ?? context.userId,
           mountedRegistry,
           inputFiles.length > 0 || inputDirectories.length > 0
             ? resolveCopilotFilePrincipal(context)