Support running string of generator code in quick-js sandbox

Author: jahooma

backend/package.json

@@ -38,6 +38,8 @@
     "express": "4.19.2",
     "gpt-tokenizer": "2.8.1",
     "ignore": "5.3.2",
+    "@jitl/quickjs-wasmfile-release-sync": "0.31.0",
+    "quickjs-emscripten-core": "0.31.0",
     "lodash": "*",
     "openai": "^4.78.1",
     "pino": "9.4.0",

backend/src/__tests__/sandbox-generator.test.ts

@@ -0,0 +1,127 @@
+import { describe, test, expect, beforeEach, afterEach } from 'bun:test'
+import {
+  runProgrammaticStep,
+  clearAgentGeneratorCache,
+} from '../run-programmatic-step'
+import { AgentState } from '@codebuff/common/types/session-state'
+import { AgentTemplate } from '../templates/types'
+import { WebSocket } from 'ws'
+import { mockFileContext, MockWebSocket } from './test-utils'
+
+describe('QuickJS Sandbox Generator', () => {
+  let mockAgentState: AgentState
+  let mockParams: any
+  let mockTemplate: AgentTemplate
+
+  beforeEach(() => {
+    clearAgentGeneratorCache()
+
+    // Reuse common test data structure
+    mockAgentState = {
+      agentId: 'test-agent-123',
+      agentType: 'test-vm-agent',
+      messageHistory: [],
+      report: {},
+      agentContext: {},
+      subagents: [],
+      stepsRemaining: 10,
+    }
+
+    // Base template structure - will be customized per test
+    mockTemplate = {
+      id: 'test-vm-agent',
+      name: 'Test VM Agent',
+      purpose: 'Test VM isolation',
+      model: 'anthropic/claude-4-sonnet-20250522',
+      outputMode: 'report',
+      includeMessageHistory: false,
+      toolNames: ['update_report'],
+      spawnableAgents: [],
+      promptSchema: {},
+      systemPrompt: '',
+      userInputPrompt: '',
+      agentStepPrompt: '',
+      initialAssistantMessage: '',
+      initialAssistantPrefix: '',
+      stepAssistantMessage: '',
+      stepAssistantPrefix: '',
+      handleStep: '', // Will be set per test
+    }
+
+    // Common params structure
+    mockParams = {
+      template: mockTemplate,
+      prompt: 'Test prompt',
+      params: { testParam: 'value' },
+      userId: 'test-user',
+      userInputId: 'test-input',
+      clientSessionId: 'test-session',
+      fingerprintId: 'test-fingerprint',
+      onResponseChunk: () => {},
+      agentType: 'test-vm-agent',
+      fileContext: mockFileContext,
+      assistantMessage: undefined,
+      assistantPrefix: undefined,
+      ws: new MockWebSocket() as unknown as WebSocket,
+    }
+  })
+
+  afterEach(() => {
+    clearAgentGeneratorCache()
+  })
+
+  test('should execute string-based generator in QuickJS sandbox', async () => {
+    // Customize template for this test
+    mockTemplate.handleStep = `
+      function* ({ agentState, prompt, params }) {
+        yield {
+          toolName: 'update_report',
+          args: {
+            json_update: {
+              message: 'Hello from QuickJS sandbox!',
+              prompt: prompt,
+              agentId: agentState.agentId
+            }
+          }
+        }
+      }
+    `
+    mockParams.template = mockTemplate
+
+    const result = await runProgrammaticStep(mockAgentState, mockParams)
+
+    expect(result.agentState.report).toEqual({
+      message: 'Hello from QuickJS sandbox!',
+      prompt: 'Test prompt',
+      agentId: 'test-agent-123',
+    })
+    expect(result.endTurn).toBe(true)
+  })
+
+  test('should handle QuickJS sandbox errors gracefully', async () => {
+    // Customize for error test
+    mockTemplate.id = 'test-vm-agent-error'
+    mockTemplate.name = 'Test VM Agent Error'
+    mockTemplate.purpose = 'Test QuickJS error handling'
+    mockTemplate.toolNames = []
+    mockTemplate.handleStep = `
+      function* ({ agentState, prompt, params }) {
+        throw new Error('QuickJS error test')
+      }
+    `
+
+    mockAgentState.agentId = 'test-agent-error-123'
+    mockAgentState.agentType = 'test-vm-agent-error'
+
+    mockParams.template = mockTemplate
+    mockParams.agentType = 'test-vm-agent-error'
+    mockParams.params = {}
+
+    const result = await runProgrammaticStep(mockAgentState, mockParams)
+
+    expect(result.endTurn).toBe(true)
+    expect(result.agentState.report.error).toContain(
+      'Error executing programmatic agent'
+    )
+  })
+})

backend/src/__tests__/test-utils.ts

@@ -13,6 +13,7 @@ export const mockFileContext: ProjectFileContext = {
   fileTree: [],
   fileTokenScores: {},
   knowledgeFiles: {},
+  userKnowledgeFiles: {},
   agentTemplates: {},
   gitChanges: {
     status: '',

backend/src/run-programmatic-step.ts

@@ -10,6 +10,10 @@ import { CodebuffToolCall } from './tools/constants'
 import { executeToolCall } from './tools/tool-executor'
 import { logger } from './util/logger'
 import { getRequestContext } from './websockets/request-context'
+import { SandboxManager } from './util/quickjs-sandbox'
+
+// Global sandbox manager for QuickJS contexts
+const sandboxManager = new SandboxManager()
 
 // Maintains generator state for all agents. Generator state can't be serialized, so we store it in memory.
 const agentIdToGenerator: Record<
@@ -22,6 +26,8 @@ export function clearAgentGeneratorCache() {
   for (const key in agentIdToGenerator) {
     delete agentIdToGenerator[key]
   }
+  // Clean up QuickJS sandboxes
+  sandboxManager.dispose()
 }
 
 // Function to handle programmatic agents
@@ -53,6 +59,9 @@ export async function runProgrammaticStep(
     fingerprintId,
     fileContext,
   } = params
+  if (!template.handleStep) {
+    throw new Error('No step handler found for agent template ' + template.id)
+  }
 
   logger.info(
     {
@@ -64,18 +73,34 @@ export async function runProgrammaticStep(
     'Running programmatic step'
   )
 
+  // Run with either a generator or a sandbox.
   let generator = agentIdToGenerator[agentState.agentId]
-  if (!generator) {
-    if (!template.handleStep) {
-      throw new Error('No step handler found for agent template ' + template.id)
+  let sandbox = sandboxManager.getSandbox(agentState.agentId)
+
+  // Check if we need to initialize a generator (either native or QuickJS-based)
+  if (!generator && !sandbox) {
+    if (typeof template.handleStep === 'string') {
+      // Initialize QuickJS sandbox for string-based generator
+      sandbox = await sandboxManager.getOrCreateSandbox(
+        agentState.agentId,
+        template.handleStep,
+        {
+          agentState,
+          prompt: params.prompt,
+          params: params.params,
+        }
+      )
+    } else {
+      // Initialize native generator
+      generator = template.handleStep({
+        agentState,
+        prompt: params.prompt,
+        params: params.params,
+      })
+      agentIdToGenerator[agentState.agentId] = generator
     }
-    generator = template.handleStep({
-      agentState,
-      prompt: params.prompt,
-      params: params.params,
-    })
-    agentIdToGenerator[agentState.agentId] = generator
   }
+
   if (generator === 'STEP_ALL') {
     return { agentState, endTurn: false }
   }
@@ -105,10 +130,16 @@ export async function runProgrammaticStep(
   try {
     // Execute tools synchronously as the generator yields them
     do {
-      let result = generator.next({
-        agentState: { ...state.agentState },
-        toolResult,
-      })
+      const result = sandbox
+        ? await sandbox.executeStep({
+            agentState: { ...state.agentState },
+            toolResult,
+          })
+        : generator!.next({
+            agentState: { ...state.agentState },
+            toolResult,
+          })
+
       if (result.done) {
         endTurn = true
         break
@@ -185,5 +216,10 @@ export async function runProgrammaticStep(
       agentState: state.agentState,
       endTurn: true,
     }
+  } finally {
+    // Clean up QuickJS sandbox if execution is complete
+    if (endTurn && sandbox) {
+      sandboxManager.removeSandbox(agentState.agentId)
+    }
   }
 }

backend/src/templates/types.ts

@@ -38,7 +38,7 @@ export type AgentTemplate<
   userInputPrompt: string
   agentStepPrompt: string
 
-  handleStep?: StepHandler<P, T>
+  handleStep?: StepHandler<P, T> | string // Function or string of the generator code for running in a sandbox
 }
 
 export type StepGenerator = Generator<