preserve exact large-value access through workflow materialization

Author: icecrasher321

apps/sim/app/api/function/execute/route.ts

@@ -10,6 +10,7 @@ import {
 } from '@/lib/copilot/request/tools/files'
 import { isE2bEnabled } from '@/lib/core/config/feature-flags'
 import { generateRequestId } from '@/lib/core/utils/request'
+import { collectUserFileKeys } from '@/lib/core/utils/user-file'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { executeInE2B, executeShellInE2B } from '@/lib/execution/e2b'
 import { executeInIsolatedVM, type IsolatedVMBrokerHandler } from '@/lib/execution/isolated-vm'
@@ -18,6 +19,7 @@ import {
   isLargeArrayManifest,
   materializeLargeArrayManifest,
 } from '@/lib/execution/payloads/large-array-manifest'
+import { collectLargeValueKeys } from '@/lib/execution/payloads/large-execution-value'
 import { containsLargeValueRef, isLargeValueRef } from '@/lib/execution/payloads/large-value-ref'
 import {
   MAX_FUNCTION_INLINE_BYTES,
@@ -703,6 +705,8 @@ interface FunctionRouteExecutionContext {
   workspaceId?: string
   executionId?: string
   largeValueExecutionIds?: string[]
+  largeValueKeys?: string[]
+  fileKeys?: string[]
   allowLargeValueWorkflowScope?: boolean
   userId?: string
   requestId: string
@@ -745,17 +749,42 @@ function getBrokerFileArgs(args: unknown): {
 function createFunctionRuntimeBrokers(
   context: FunctionRouteExecutionContext
 ): Record<string, IsolatedVMBrokerHandler> {
+  context.largeValueKeys ??= []
+  context.fileKeys ??= []
+  const largeValueKeys = context.largeValueKeys
+  const fileKeys = context.fileKeys
   const base = {
     requestId: context.requestId,
     workflowId: context.workflowId,
     workspaceId: context.workspaceId,
     executionId: context.executionId,
     largeValueExecutionIds: context.largeValueExecutionIds,
+    largeValueKeys,
+    fileKeys,
     allowLargeValueWorkflowScope: context.allowLargeValueWorkflowScope,
     userId: context.userId,
     logger,
   }
 
+  const recordMaterializedKeys = (value: unknown) => {
+    const keys = collectLargeValueKeys(value)
+    const existingKeys = new Set(largeValueKeys)
+    for (const key of keys) {
+      if (!existingKeys.has(key)) {
+        existingKeys.add(key)
+        largeValueKeys.push(key)
+      }
+    }
+    const keysForFiles = collectUserFileKeys(value)
+    const existingFileKeys = new Set(fileKeys)
+    for (const key of keysForFiles) {
+      if (!existingFileKeys.has(key)) {
+        existingFileKeys.add(key)
+        fileKeys.push(key)
+      }
+    }
+  }
+
   const readFile = async (args: unknown, encoding: 'base64' | 'text', chunked = false) => {
     const fileArgs = getBrokerFileArgs(args)
     return readUserFileContent(fileArgs.file, {
@@ -790,6 +819,7 @@ function createFunctionRuntimeBrokers(
       if (value === undefined) {
         throw unavailableLargeValueError(ref)
       }
+      recordMaterializedKeys(value)
       return value
     },
     'sim.values.readArray': async (args) => {
@@ -802,10 +832,12 @@ function createFunctionRuntimeBrokers(
       if (!context.executionId) {
         throw new Error('Large array manifests require an execution context.')
       }
-      return materializeLargeArrayManifest(manifest, {
+      const value = await materializeLargeArrayManifest(manifest, {
         ...base,
         maxBytes: clampInlineBytes(options.maxBytes, MAX_INLINE_MATERIALIZATION_BYTES),
       })
+      recordMaterializedKeys(value)
+      return value
     },
   }
 }
@@ -829,7 +861,17 @@ async function functionJsonResponse<T>(
   context: FunctionRouteExecutionContext,
   init?: ResponseInit
 ) {
-  return NextResponse.json(await compactFunctionRouteBody(body, context), init)
+  return NextResponse.json(
+    await compactFunctionRouteBody(
+      {
+        ...body,
+        largeValueKeys: context.largeValueKeys,
+        fileKeys: context.fileKeys,
+      },
+      context
+    ),
+    init
+  )
 }
 
 async function maybeExportSandboxFileToWorkspace(args: {
@@ -974,6 +1016,8 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
       workflowId,
       executionId,
       largeValueExecutionIds,
+      largeValueKeys,
+      fileKeys,
       allowLargeValueWorkflowScope = false,
       workspaceId,
       isCustomTool = false,
@@ -998,6 +1042,8 @@ export const POST = withRouteHandler(async (req: NextRequest) => {
       workspaceId,
       executionId,
       largeValueExecutionIds,
+      largeValueKeys,
+      fileKeys,
       allowLargeValueWorkflowScope,
       userId: auth.userId,
       requestId,

apps/sim/app/api/workflows/[id]/execute/response-block.test.ts

@@ -8,13 +8,17 @@
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 import { AuthType } from '@/lib/auth/hybrid'
 import { clearLargeValueCacheForTests } from '@/lib/execution/payloads/cache'
+import { createLargeArrayManifest } from '@/lib/execution/payloads/large-array-manifest'
 import { compactExecutionPayload } from '@/lib/execution/payloads/serializer'
+import { storeLargeValue } from '@/lib/execution/payloads/store'
 import { EXECUTION_RESOURCE_LIMIT_CODE } from '@/lib/execution/resource-errors'
 import type { ExecutionResult } from '@/lib/workflows/types'
 import { createHttpResponseFromBlock, workflowHasResponseBlock } from '@/lib/workflows/utils'
 
-const { mockUploadFile } = vi.hoisted(() => ({
+const { mockDownloadFile, mockUploadFile, uploadedFiles } = vi.hoisted(() => ({
+  mockDownloadFile: vi.fn(),
   mockUploadFile: vi.fn(),
+  uploadedFiles: new Map<string, Buffer>(),
 }))
 
 const MATERIALIZATION_CONTEXT = {
@@ -26,6 +30,7 @@ const MATERIALIZATION_CONTEXT = {
 
 vi.mock('@/lib/uploads', () => ({
   StorageService: {
+    downloadFile: mockDownloadFile,
     uploadFile: mockUploadFile,
   },
 }))
@@ -60,7 +65,14 @@ describe('Response block gating by auth type', () => {
   beforeEach(() => {
     vi.clearAllMocks()
     clearLargeValueCacheForTests()
-    mockUploadFile.mockImplementation(async ({ customKey }) => ({ key: customKey }))
+    uploadedFiles.clear()
+    mockUploadFile.mockImplementation(async ({ customKey, file }) => {
+      uploadedFiles.set(customKey, file)
+      return { key: customKey }
+    })
+    mockDownloadFile.mockImplementation(
+      async ({ key }) => uploadedFiles.get(key) ?? Buffer.from('{}')
+    )
     resultWithResponseBlock = buildExecutionResult()
   })
 
@@ -165,6 +177,211 @@ describe('Response block gating by auth type', () => {
     expect(body.success).toBeUndefined()
   })
 
+  it('should materialize Response block manifests from an allowed source execution', async () => {
+    const rows = [{ key: 'SIM-1' }, { key: 'SIM-2' }]
+    const manifest = await createLargeArrayManifest(rows, {
+      ...MATERIALIZATION_CONTEXT,
+      executionId: 'source-execution-1',
+    })
+
+    const response = await createHttpResponseFromBlock(
+      buildExecutionResult({
+        output: {
+          data: { rows: manifest },
+          status: 200,
+          headers: {},
+        },
+      }),
+      {
+        ...MATERIALIZATION_CONTEXT,
+        largeValueExecutionIds: ['source-execution-1'],
+      }
+    )
+    const body = await response.json()
+
+    expect(body.rows).toEqual(rows)
+  })
+
+  it('should reject Response block manifests from non-source same-workflow executions', async () => {
+    const manifest = await createLargeArrayManifest([{ key: 'SIM-stale' }], {
+      ...MATERIALIZATION_CONTEXT,
+      executionId: 'stale-execution-1',
+    })
+
+    await expect(
+      createHttpResponseFromBlock(
+        buildExecutionResult({
+          output: {
+            data: { rows: manifest },
+            status: 200,
+            headers: {},
+          },
+        }),
+        {
+          ...MATERIALIZATION_CONTEXT,
+          largeValueExecutionIds: ['source-execution-1'],
+        }
+      )
+    ).rejects.toThrow('Large execution value is not available in this execution')
+  })
+
+  it('should materialize Response block manifests inherited by the source snapshot', async () => {
+    const rows = [{ key: 'SIM-inherited' }]
+    const manifest = await createLargeArrayManifest(rows, {
+      ...MATERIALIZATION_CONTEXT,
+      executionId: 'original-execution-1',
+    })
+
+    const response = await createHttpResponseFromBlock(
+      buildExecutionResult({
+        output: {
+          data: { rows: manifest },
+          status: 200,
+          headers: {},
+        },
+      }),
+      {
+        ...MATERIALIZATION_CONTEXT,
+        largeValueExecutionIds: ['source-execution-1', 'original-execution-1'],
+      }
+    )
+
+    const body = await response.json()
+
+    expect(body.rows).toEqual(rows)
+  })
+
+  it('should recursively materialize refs inside Response block manifest rows', async () => {
+    const text = 'nested'.repeat(2 * 1024 * 1024)
+    const nestedOutput = await compactExecutionPayload(
+      { text },
+      {
+        ...MATERIALIZATION_CONTEXT,
+        executionId: 'original-execution-1',
+        requireDurable: true,
+        preserveRoot: true,
+      }
+    )
+    const nestedRef = (nestedOutput as unknown as { text: unknown }).text
+    const manifest = await createLargeArrayManifest([{ nested: nestedRef }], {
+      ...MATERIALIZATION_CONTEXT,
+      executionId: 'source-execution-1',
+    })
+    const response = await createHttpResponseFromBlock(
+      buildExecutionResult({
+        output: {
+          data: { rows: manifest },
+          status: 200,
+          headers: {},
+        },
+      }),
+      {
+        ...MATERIALIZATION_CONTEXT,
+        largeValueExecutionIds: ['source-execution-1'],
+      }
+    )
+
+    const body = await response.json()
+
+    expect(body.rows).toEqual([{ nested: text }])
+  })
+
+  it('should recursively materialize refs inside stored Response block objects', async () => {
+    const text = 'nested'.repeat(2 * 1024 * 1024)
+    const nestedOutput = await compactExecutionPayload(
+      { text },
+      {
+        ...MATERIALIZATION_CONTEXT,
+        executionId: 'original-execution-1',
+        requireDurable: true,
+        preserveRoot: true,
+      }
+    )
+    const nestedRef = (nestedOutput as unknown as { text: unknown }).text
+    const storedValue = {
+      wrapper: {
+        nested: nestedRef,
+        padding: 'x'.repeat(2048),
+      },
+    }
+    const storedJson = JSON.stringify(storedValue)
+    const storedOutput = await storeLargeValue(
+      storedValue,
+      storedJson,
+      Buffer.byteLength(storedJson),
+      {
+        ...MATERIALIZATION_CONTEXT,
+        executionId: 'source-execution-1',
+        requireDurable: true,
+      }
+    )
+
+    const response = await createHttpResponseFromBlock(
+      buildExecutionResult({
+        output: {
+          data: storedOutput,
+          status: 200,
+          headers: {},
+        },
+      }),
+      {
+        ...MATERIALIZATION_CONTEXT,
+        largeValueExecutionIds: ['source-execution-1'],
+      }
+    )
+
+    const body = await response.json()
+
+    expect(body.wrapper.nested).toEqual(text)
+  })
+
+  it('should memoize repeated materialized objects while resolving nested refs', async () => {
+    const text = 'nested'.repeat(2 * 1024 * 1024)
+    const nestedOutput = await compactExecutionPayload(
+      { text },
+      {
+        ...MATERIALIZATION_CONTEXT,
+        executionId: 'original-execution-1',
+        requireDurable: true,
+        preserveRoot: true,
+      }
+    )
+    const nestedRef = (nestedOutput as unknown as { text: unknown }).text
+    const sourceValue = { nested: nestedRef }
+    const sourceJson = JSON.stringify(sourceValue)
+    const sourceRef = await storeLargeValue(
+      sourceValue,
+      sourceJson,
+      Buffer.byteLength(sourceJson),
+      {
+        ...MATERIALIZATION_CONTEXT,
+        executionId: 'source-execution-1',
+        requireDurable: true,
+      }
+    )
+
+    const response = await createHttpResponseFromBlock(
+      buildExecutionResult({
+        output: {
+          data: { first: sourceRef, second: sourceRef },
+          status: 200,
+          headers: {},
+        },
+      }),
+      {
+        ...MATERIALIZATION_CONTEXT,
+        largeValueKeys: sourceRef.key ? [sourceRef.key] : [],
+      }
+    )
+
+    const body = await response.json()
+
+    expect(body).toEqual({
+      first: { nested: text },
+      second: { nested: text },
+    })
+  })
+
   it('should materialize large string refs for Response block HTTP output', async () => {
     const text = 'x'.repeat(9 * 1024 * 1024)
     const output = await compactExecutionPayload(