fix(tables): reject stale sandbox snapshots

Author: TheodoreSpeaks

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

@@ -31,7 +31,7 @@ const {
   mockMaterializeCopilotCodeSecrets,
   mockHasWorkspaceSandboxAccess,
   mockImportWorkspaceFileSecretProvenanceForRuntime,
-  mockIsTableSnapshotSafeForModelMount,
+  mockGetTableSnapshotModelMountSafety,
 } = vi.hoisted(() => ({
   mockGetTableById: vi.fn(),
   mockListTables: vi.fn(),
@@ -53,7 +53,7 @@ const {
   mockMaterializeCopilotCodeSecrets: vi.fn(),
   mockHasWorkspaceSandboxAccess: vi.fn(),
   mockImportWorkspaceFileSecretProvenanceForRuntime: vi.fn(),
-  mockIsTableSnapshotSafeForModelMount: vi.fn(),
+  mockGetTableSnapshotModelMountSafety: vi.fn(),
 }))
 
 vi.mock('@/lib/core/security/encryption', () => encryptionMock)
@@ -62,7 +62,7 @@ vi.mock('@/lib/table/service', () => ({
   listTables: mockListTables,
 }))
 vi.mock('@/lib/table/rows/secret-provenance', () => ({
-  isTableSnapshotSafeForModelMount: mockIsTableSnapshotSafeForModelMount,
+  getTableSnapshotModelMountSafety: mockGetTableSnapshotModelMountSafety,
 }))
 vi.mock('@/lib/table/snapshot-cache', () => ({
   getOrCreateTableSnapshot: mockGetOrCreateTableSnapshot,
@@ -146,8 +146,8 @@ function resetExecutionMocks(): void {
   vi.clearAllMocks()
   mockExecuteTool.mockReset()
   mockMaterializeCopilotCodeSecrets.mockReset()
-  mockIsTableSnapshotSafeForModelMount.mockReset()
-  mockIsTableSnapshotSafeForModelMount.mockResolvedValue(true)
+  mockGetTableSnapshotModelMountSafety.mockReset()
+  mockGetTableSnapshotModelMountSafety.mockResolvedValue('safe')
   mockListWorkspaceFiles.mockResolvedValue([])
   mockListWorkspaceFileFolders.mockResolvedValue([])
   mockListAllWorkspaceFiles.mockImplementation(async () => {
@@ -609,7 +609,7 @@ describe('executeFunctionExecute table mounts', () => {
   })
 
   it('unknown snapshot provenance still mounts and taints model egress', async () => {
-    mockIsTableSnapshotSafeForModelMount.mockResolvedValue(false)
+    mockGetTableSnapshotModelMountSafety.mockResolvedValue('unsafe-provenance')
     mockGetOrCreateTableSnapshot.mockResolvedValue({
       key: 'table-snapshots/ws_1/tbl_1/v5.csv',
       size: 9,
@@ -637,6 +637,21 @@ describe('executeFunctionExecute table mounts', () => {
     expect(projectToolResultForCopilot(result, parentRegistry)).toEqual({ success: true })
   })
 
+  it('rejects a snapshot that becomes stale before mounting', async () => {
+    mockGetTableSnapshotModelMountSafety.mockResolvedValue('stale')
+    mockGetOrCreateTableSnapshot.mockResolvedValue({
+      key: 'table-snapshots/ws_1/tbl_1/v5.csv',
+      size: 9,
+      version: 5,
+    })
+
+    await expect(
+      executeFunctionExecute({ inputTables: ['tbl_1'] }, context as never)
+    ).rejects.toThrow(/changed while preparing its snapshot/)
+    expect(mockGeneratePresignedDownloadUrl).not.toHaveBeenCalled()
+    expect(mockExecuteTool).not.toHaveBeenCalled()
+  })
+
   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

@@ -19,7 +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 { isTableSnapshotSafeForModelMount } from '@/lib/table/rows/secret-provenance'
+import { getTableSnapshotModelMountSafety } from '@/lib/table/rows/secret-provenance'
 import { getTableById, listTables } from '@/lib/table/service'
 import { getOrCreateTableSnapshot, SNAPSHOT_MAX_BYTES } from '@/lib/table/snapshot-cache'
 import {
@@ -484,15 +484,15 @@ export async function resolveInputFiles(
           `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 {
+      const mountSafety = await getTableSnapshotModelMountSafety({
+        tableId: table.id,
+        workspaceId,
+        rowsVersion: snapshot.version,
+      })
+      if (mountSafety === 'stale') {
+        throw new Error(`Input table "${tableId}" changed while preparing its snapshot. Retry.`)
+      }
+      if (mountSafety === 'unsafe-provenance') {
         resolvedSecretTraceRegistry.markIncomplete('table-snapshot-unsafe-for-mount')
       }
 

apps/sim/lib/table/rows/secret-provenance.test.ts

@@ -20,7 +20,7 @@ vi.mock('@/lib/execution/durable-secret-provenance-enforcement', () => ({
 import type { DbTransaction } from '@/lib/table/planner'
 import {
   classifyTableRowSecretProvenanceForCopy,
-  isTableSnapshotSafeForModelMount,
+  getTableSnapshotModelMountSafety,
   loadTableRowSecretProvenance,
   mutateTableRowsWithSecretProvenance,
   updateTableRowsWithDerivedSecretProvenance,
@@ -67,30 +67,31 @@ describe('table row secret provenance', () => {
     queueTableRows(userTableRows, [])
 
     await expect(
-      isTableSnapshotSafeForModelMount({
+      getTableSnapshotModelMountSafety({
         tableId: 'table-1',
         workspaceId: 'workspace-1',
         rowsVersion: 7,
       })
-    ).resolves.toBe(true)
+    ).resolves.toBe('safe')
 
     expect(dbChainMockFns.limit).toHaveBeenCalledTimes(3)
     expect(dbChainMockFns.orderBy).not.toHaveBeenCalled()
   })
 
-  it('rejects the snapshot as soon as an unsafe row exists', async () => {
+  it('classifies unsafe provenance after confirming the snapshot remains current', async () => {
     queueTableRows(userTableDefinitions, [{ rowsVersion: 7 }])
     queueTableRows(userTableRows, [{ id: 'unsafe-row' }])
+    queueTableRows(userTableDefinitions, [{ rowsVersion: 7 }])
 
     await expect(
-      isTableSnapshotSafeForModelMount({
+      getTableSnapshotModelMountSafety({
         tableId: 'table-1',
         workspaceId: 'workspace-1',
         rowsVersion: 7,
       })
-    ).resolves.toBe(false)
+    ).resolves.toBe('unsafe-provenance')
 
-    expect(dbChainMockFns.limit).toHaveBeenCalledTimes(2)
+    expect(dbChainMockFns.limit).toHaveBeenCalledTimes(3)
   })
 
   it('rejects a snapshot when the table changes during the safety check', async () => {
@@ -99,12 +100,12 @@ describe('table row secret provenance', () => {
     queueTableRows(userTableRows, [])
 
     await expect(
-      isTableSnapshotSafeForModelMount({
+      getTableSnapshotModelMountSafety({
         tableId: 'table-1',
         workspaceId: 'workspace-1',
         rowsVersion: 7,
       })
-    ).resolves.toBe(false)
+    ).resolves.toBe('stale')
   })
 
   it('keeps untouched legacy rows readable with exact-empty provenance', async () => {

apps/sim/lib/table/rows/secret-provenance.ts

@@ -705,17 +705,19 @@ async function readTableRowsVersion(tableId: string, workspaceId: string): Promi
   return table?.rowsVersion ?? null
 }
 
+export type TableSnapshotModelMountSafety = 'safe' | 'unsafe-provenance' | 'stale'
+
 /**
- * Verifies that a version-pinned table contains no secret-bearing or unknown
- * cells before its opaque CSV bytes cross into a model-controlled sandbox.
+ * Classifies whether a version-pinned table snapshot can cross into a model-controlled sandbox.
+ * A version change takes precedence over provenance because stale bytes must never be mounted.
  */
-export async function isTableSnapshotSafeForModelMount(options: {
+export async function getTableSnapshotModelMountSafety(options: {
   tableId: string
   workspaceId: string
   rowsVersion: number
-}): Promise<boolean> {
+}): Promise<TableSnapshotModelMountSafety> {
   if ((await readTableRowsVersion(options.tableId, options.workspaceId)) !== options.rowsVersion) {
-    return false
+    return 'stale'
   }
 
   const [unsafeRow] = await db
@@ -748,9 +750,11 @@ export async function isTableSnapshotSafeForModelMount(options: {
     )
     .limit(1)
 
-  if (unsafeRow) return false
+  if ((await readTableRowsVersion(options.tableId, options.workspaceId)) !== options.rowsVersion) {
+    return 'stale'
+  }
 
-  return (await readTableRowsVersion(options.tableId, options.workspaceId)) === options.rowsVersion
+  return unsafeRow ? 'unsafe-provenance' : 'safe'
 }
 
 /**