fix(mothership): export chat sandbox files through workspace storage

Author: Sg312

apps/sim/lib/function-execution/execute-request.test.ts

@@ -3,7 +3,6 @@
  */
 import { spawnSync } from 'node:child_process'
 import { readFileSync } from 'node:fs'
-import { resolve } from 'node:path'
 import {
   createMockRequest,
   envFlagsMock,
@@ -2085,6 +2084,81 @@ describe('Function execution request', () => {
       expect(runtimePayload).not.toContain('__simSandboxFileMount')
     })
 
+    it.each(['python', 'javascript', 'shell'])(
+      'exports a chat %s harvest through workspace policy without workflow context',
+      async (language) => {
+        envFlagsMock.isMothershipSandboxEnabled = true
+        hybridAuthMockFns.mockCheckInternalAuth.mockResolvedValue({
+          success: true,
+          userId: 'user-123',
+          authType: 'internal_jwt',
+          sandboxProfile: 'mothership',
+        })
+        const bytes = Buffer.from([0, 255, 13, 10, 42])
+        const runtimeResult = {
+          result: 'kept-result',
+          stdout: 'kept-stdout',
+          sandboxId: 'sbx',
+          collectedFiles: [
+            {
+              path: '/tmp/sim/outputs/call-test/report.txt',
+              relativePath: 'report.txt',
+              contentBase64: bytes.toString('base64'),
+              byteLength: bytes.length,
+            },
+          ],
+        }
+        const sandbox = language === 'shell' ? mockExecuteShellInSandbox : mockExecuteInSandbox
+        sandbox.mockResolvedValueOnce(runtimeResult)
+        const response = await POST(
+          createMockRequest('POST', {
+            code: 'x',
+            language,
+            workspaceId: 'workspace-1',
+            sandboxSessionKey: 'same-chat',
+          })
+        )
+        const data = await response.json()
+        expect(response.status).toBe(200)
+        expect(data.output.result).toBe('kept-result')
+        expect(data.output.stdout).toBe('kept-stdout')
+        expect(data.resources).toEqual([
+          expect.objectContaining({ type: 'file', path: 'files/report.txt' }),
+        ])
+        expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledWith(
+          expect.objectContaining({
+            workspaceId: 'workspace-1',
+            target: expect.objectContaining({ path: 'files/report.txt', mode: 'create' }),
+            buffer: bytes,
+          })
+        )
+        expect(mockUploadFile).not.toHaveBeenCalled()
+      }
+    )
+
+    it('still requires workflow context for ordinary Function file harvests', async () => {
+      envFlagsMock.isRemoteSandboxEnabled = true
+      mockExecuteInSandbox.mockResolvedValueOnce({
+        result: null,
+        stdout: '',
+        sandboxId: 'sbx',
+        collectedFiles: [
+          {
+            path: '/tmp/sim/outputs/a.txt',
+            relativePath: 'a.txt',
+            contentBase64: 'YQ==',
+            byteLength: 1,
+          },
+        ],
+      })
+      const response = await POST(
+        createMockRequest('POST', { code: 'x', language: 'python', workspaceId: 'workspace-1' })
+      )
+      expect(response.status).toBe(400)
+      expect((await response.json()).error).toContain('workflow, and execution context')
+      expect(mockWriteWorkspaceFileByPath).not.toHaveBeenCalled()
+    })
+
     it('gives overlapping calls in one persistent workbench distinct automatic export directories', async () => {
       envFlagsMock.isMothershipSandboxEnabled = true
       hybridAuthMockFns.mockCheckInternalAuth.mockResolvedValue({
@@ -2352,7 +2426,7 @@ describe('Function execution request', () => {
       const archiveBase64 =
         'UEsDBBQAAAAIAAAAIQAcWyFBIAAAAB8AAAAMAAAAcHJldmlldy5odG1ss8kwtHNLzcnJLy9WcM4vzUvOzFEIT03Nzqm00QdKAQBQSwECFAMUAAAACAAAACEAHFshQSAAAAAfAAAADAAAAAAAAAAAAAAAgAEAAAAAcHJldmlldy5odG1sUEsFBgAAAAABAAEAOgAAAEoAAAAAAA=='
       const source = readFileSync(
-        resolve(process.cwd(), 'lib/execution/remote-sandbox/fixtures/fellows-council-weekly.py'),
+        new URL('../execution/remote-sandbox/fixtures/fellows-council-weekly.py', import.meta.url),
         'utf8'
       )
       mockExecuteInSandbox.mockResolvedValueOnce({

apps/sim/lib/function-execution/execute-request.ts

@@ -1518,7 +1518,7 @@ async function maybeExportSandboxFileToWorkspace(args: {
   outputSandboxPath?: string
   overwriteFileId?: string
   outputMode?: 'create' | 'overwrite'
-  exportedFileContent?: string
+  exportedFileContent?: string | Buffer
   result: unknown
   stdout: string
   executionTime: number
@@ -1599,9 +1599,9 @@ async function maybeExportSandboxFileToWorkspace(args: {
       cost
     )
   }
-  const fileBuffer = isBinary
-    ? Buffer.from(exportedFileContent, 'base64')
-    : Buffer.from(exportedFileContent, 'utf-8')
+  const fileBuffer = Buffer.isBuffer(exportedFileContent)
+    ? exportedFileContent
+    : Buffer.from(exportedFileContent, isBinary ? 'base64' : 'utf-8')
   const secretProvenance = await getOutputFileSecretProvenance(fileBuffer, isBinary, routeContext, {
     userId: authUserId,
     workspaceId: resolvedWorkspaceId,
@@ -1691,8 +1691,8 @@ async function maybeExportSandboxFilesToWorkspace(args: {
   workflowId?: string
   workspaceId?: string
   outputFiles: OutputFileDeclaration[]
-  exportedFiles?: Record<string, string>
-  exportedFileContent?: string
+  exportedFiles?: Record<string, string | Buffer>
+  exportedFileContent?: string | Buffer
   result: unknown
   stdout: string
   executionTime: number
@@ -1780,7 +1780,9 @@ async function maybeExportSandboxFilesToWorkspace(args: {
         args.cost
       )
     }
-    const scanBuffer = isBinary ? Buffer.from(content, 'base64') : Buffer.from(content, 'utf-8')
+    const scanBuffer = Buffer.isBuffer(content)
+      ? content
+      : Buffer.from(content, isBinary ? 'base64' : 'utf-8')
     const secretProvenance = await getOutputFileSecretProvenance(
       scanBuffer,
       isBinary,
@@ -1846,9 +1848,9 @@ async function maybeExportSandboxFilesToWorkspace(args: {
   const writtenFiles = []
   try {
     for (const prepared of preparedFiles) {
-      const buffer = prepared.isBinary
-        ? Buffer.from(prepared.content, 'base64')
-        : Buffer.from(prepared.content, 'utf-8')
+      const buffer = Buffer.isBuffer(prepared.content)
+        ? prepared.content
+        : Buffer.from(prepared.content, prepared.isBinary ? 'base64' : 'utf-8')
       let previousSize: number | undefined
       let unchanged = false
       if (prepared.target.mode === 'overwrite') {
@@ -2007,20 +2009,40 @@ async function discardUploadedExecutionFiles(files: readonly UserFile[]): Promis
   }
 }
 
-async function collectExecutionOutputFiles(args: {
+async function collectSandboxOutputFiles(args: {
   routeContext: FunctionRouteExecutionContext
   authUserId: string
   workflowId?: string
   workspaceId?: string
   executionId?: string
   collectedFiles: SandboxCollectedFile[]
+  sandboxProfile?: 'mothership'
+  result: unknown
   stdout: string
   executionTime: number
   cost?: FunctionExecutionCost
 }): Promise<{ files: UserFile[] } | { response: NextResponse }> {
   const { routeContext, collectedFiles } = args
   if (collectedFiles.length === 0) return { files: [] }
 
+  /** Chat exports use workspace ownership; only workflow exports need an execution scope. */
+  if (args.sandboxProfile === 'mothership') {
+    const response = await maybeExportSandboxFilesToWorkspace({
+      ...args,
+      outputFiles: collectedFiles.map((file) => ({
+        path: `files/${collectedFileName(file.relativePath)}`,
+        sandboxPath: file.path,
+        mode: 'create',
+        mimeType: getMimeTypeFromExtension(getFileExtension(file.relativePath)),
+      })),
+      exportedFiles: Object.fromEntries(
+        collectedFiles.map((file) => [file.path, Buffer.from(file.contentBase64, 'base64')])
+      ),
+    })
+    if (!response) throw new Error('Collected sandbox files require an export response')
+    return { response }
+  }
+
   const resolvedWorkspaceId =
     args.workspaceId ||
     (args.workflowId ? (await getWorkflowById(args.workflowId))?.workspaceId : undefined)
@@ -2679,13 +2701,15 @@ export async function executeFunctionRequest(
         }
       }
 
-      const shellOutputFiles = await collectExecutionOutputFiles({
+      const shellOutputFiles = await collectSandboxOutputFiles({
         routeContext,
         authUserId: auth.attributedUserId,
         workflowId,
         workspaceId,
         executionId,
         collectedFiles: shellCollectedFiles ?? [],
+        sandboxProfile: auth.sandboxProfile,
+        result: shellResult,
         stdout: shellStdout,
         executionTime,
         cost: shellCost,
@@ -2843,13 +2867,15 @@ export async function executeFunctionRequest(
           }
         }
 
-        const jsOutputFiles = await collectExecutionOutputFiles({
+        const jsOutputFiles = await collectSandboxOutputFiles({
           routeContext,
           authUserId: auth.attributedUserId,
           workflowId,
           workspaceId,
           executionId,
           collectedFiles: jsCollectedFiles ?? [],
+          sandboxProfile: auth.sandboxProfile,
+          result: e2bResult,
           stdout,
           executionTime,
           cost: sandboxCost,
@@ -2969,13 +2995,15 @@ export async function executeFunctionRequest(
         }
       }
 
-      const pythonOutputFiles = await collectExecutionOutputFiles({
+      const pythonOutputFiles = await collectSandboxOutputFiles({
         routeContext,
         authUserId: auth.attributedUserId,
         workflowId,
         workspaceId,
         executionId,
         collectedFiles: pythonCollectedFiles ?? [],
+        sandboxProfile: auth.sandboxProfile,
+        result: e2bResult,
         stdout,
         executionTime,
         cost: sandboxCost,

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

@@ -104,15 +104,15 @@ describe.each([
   afterEach(clearHandlers)
 
   it.each([
-    { timeout: undefined, expected: 10_000 },
+    { timeout: undefined, expected: 60_000 },
     { timeout: 0.5, expected: 500 },
     { timeout: 7, expected: 7_000 },
     { timeout: '90', expected: 90_000 },
     { timeout: 300, expected: 300_000 },
     { timeout: 601, expected: 300_000 },
     { timeout: 45_000, expected: 300_000 },
-    { timeout: 0, expected: 10_000 },
-    { timeout: 'invalid', expected: 10_000 },
+    { timeout: 0, expected: 60_000 },
+    { timeout: 'invalid', expected: 60_000 },
   ])('maps $timeout seconds to $expected milliseconds once', async ({ timeout, expected }) => {
     const params = { code: 'return 1', language: 'javascript', timeout }
     const result = await executeTool(toolId, params, {

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

@@ -54,7 +54,8 @@ import { executeTool as executeAppTool } from '@/tools'
 const logger = createLogger('CopilotFunctionExecute')
 
 const MAX_MOUNTED_FILES = 500
-const DEFAULT_TIMEOUT_SECONDS = 10
+/** The call budget includes cold allocation and mounting, before the user's code starts. */
+const DEFAULT_TIMEOUT_SECONDS = 60
 const MAX_TIMEOUT_SECONDS = 300
 
 /** Keeps mount bytes, canonical classification and budget together until they enter the runtime. */