Get handleSteps fully working from template files!

Author: jahooma

.agents/templates/example-handlesteps-agent.ts

@@ -0,0 +1,74 @@
+import { DynamicAgentConfig } from '@codebuff/common/types/dynamic-agent-template'
+
+export default {
+  id: 'example-handlesteps-agent',
+  version: '1.0.0',
+  name: 'Example HandleSteps Agent',
+  purpose:
+    'Demonstrates how to use handleSteps generator functions for programmatic agent control',
+  model: 'claude-3-5-sonnet-20241022',
+  outputMode: 'json',
+  toolNames: ['spawn_agents', 'set_output', 'end_turn'],
+  spawnableAgents: ['file_picker'],
+
+  systemPrompt:
+    'You are an example agent that demonstrates handleSteps functionality.',
+  userInputPrompt: 'User request: {prompt}',
+  agentStepPrompt: 'Continue processing the request.',
+
+  // Generator function that defines the agent's execution flow
+  handleSteps: function* ({ agentState, prompt, params }) {
+    // Step 1: Spawn a file picker to find relevant files
+    const { toolResult: filePickerResult } = yield {
+      toolName: 'spawn_agents',
+      args: {
+        agents: [
+          {
+            agent_type: 'file_picker',
+            prompt: prompt || 'Find relevant files for the user request',
+          },
+        ],
+      },
+    }
+
+    // Step 2: Process the results and extract file paths using regex
+    let message = 'File picker completed'
+    let files: string[] = []
+
+    if (filePickerResult) {
+      // Extract the actual response from the agent_report wrapper
+      const resultText = filePickerResult.result || ''
+
+      // Extract file paths from backticks in the text
+      const filePathRegex = /`([^`]+\.[a-zA-Z0-9]+)`/g
+      const matches: string[] = []
+      let match
+
+      while ((match = filePathRegex.exec(resultText)) !== null) {
+        const filePath: string | undefined = match[1]
+        if (filePath && !matches.includes(filePath)) {
+          matches.push(filePath)
+        }
+      }
+
+      if (matches.length > 0) {
+        files = matches
+        message = `Found ${files.length} file paths mentioned in response`
+      } else {
+        message = `File picker completed but no file paths found`
+      }
+    }
+
+    // Step 3: Set the final output
+    yield {
+      toolName: 'set_output',
+      args: {
+        message,
+        files,
+        prompt: prompt || 'No prompt provided',
+        params: params || {},
+        agentId: agentState.agentId,
+      },
+    }
+  },
+} satisfies DynamicAgentConfig

backend/src/__tests__/handlesteps-parsing.test.ts

@@ -0,0 +1,237 @@
+import { describe, test, expect, beforeEach, afterEach } from 'bun:test'
+import {
+  DynamicAgentConfigSchema,
+  DynamicAgentTemplate,
+} from '@codebuff/common/types/dynamic-agent-template'
+import { dynamicAgentService } from '../templates/dynamic-agent-service'
+import { AgentState } from '@codebuff/common/types/session-state'
+import { ProjectFileContext } from '@codebuff/common/util/file'
+
+describe('handleSteps Parsing Tests', () => {
+  let mockFileContext: ProjectFileContext
+  let mockAgentTemplate: DynamicAgentTemplate
+
+  beforeEach(() => {
+    dynamicAgentService.reset()
+
+    // Setup common mock data
+    mockFileContext = {
+      projectRoot: '/test',
+      cwd: '/test',
+      fileTree: [],
+      fileTokenScores: {},
+      knowledgeFiles: {},
+      agentTemplates: {},
+      gitChanges: {
+        status: '',
+        diff: '',
+        diffCached: '',
+        lastCommitMessages: '',
+      },
+      changesSinceLastChat: {},
+      shellConfigFiles: {},
+      systemInfo: {
+        platform: 'test',
+        shell: 'test',
+        nodeVersion: 'test',
+        arch: 'test',
+        homedir: '/test',
+        cpus: 1,
+      },
+      tokenCallers: {},
+    }
+
+    mockAgentTemplate = {
+      id: 'test-agent',
+      version: '1.0.0',
+      name: 'Test Agent',
+      purpose: 'Testing',
+      model: 'claude-3-5-sonnet-20241022',
+      outputMode: 'json' as const,
+      toolNames: ['set_output'],
+      spawnableAgents: [],
+      override: false as const,
+      includeMessageHistory: true,
+      systemPrompt: 'Test system prompt',
+      userInputPrompt: 'Test user prompt',
+      agentStepPrompt: 'Test agent step prompt',
+      initialAssistantMessage: '',
+      initialAssistantPrefix: '',
+      stepAssistantMessage: '',
+      stepAssistantPrefix: '',
+    }
+  })
+
+  afterEach(() => {
+    dynamicAgentService.reset()
+  })
+
+  test('should validate agent config with handleSteps function', () => {
+    const agentConfig = {
+      id: 'test-agent',
+      version: '1.0.0',
+      name: 'Test Agent',
+      purpose: 'Testing handleSteps',
+      model: 'claude-3-5-sonnet-20241022',
+      outputMode: 'json' as const,
+      toolNames: ['set_output'],
+      systemPrompt: 'You are a test agent',
+      userInputPrompt: 'Process: {prompt}',
+      agentStepPrompt: 'Continue processing',
+      handleSteps: function* ({
+        agentState,
+        prompt,
+        params,
+      }: {
+        agentState: AgentState
+        prompt?: string
+        params?: any
+      }) {
+        yield {
+          toolName: 'set_output',
+          args: { message: 'Test completed' },
+        }
+      },
+    }
+
+    const result = DynamicAgentConfigSchema.safeParse(agentConfig)
+    expect(result.success).toBe(true)
+
+    if (result.success) {
+      expect(typeof result.data.handleSteps).toBe('function')
+    }
+  })
+
+  test('should convert handleSteps function to string', async () => {
+    const handleStepsFunction = function* ({
+      agentState,
+      prompt,
+      params,
+    }: {
+      agentState: AgentState
+      prompt?: string
+      params?: any
+    }) {
+      yield {
+        toolName: 'set_output',
+        args: { message: 'Hello from generator' },
+      }
+    }
+
+    const agentTemplates = {
+      'test-agent': {
+        ...mockAgentTemplate,
+        handleSteps: handleStepsFunction.toString(),
+      },
+    }
+
+    const fileContext: ProjectFileContext = {
+      ...mockFileContext,
+      agentTemplates,
+    }
+
+    const result = await dynamicAgentService.loadAgents(fileContext)
+
+    expect(result.validationErrors).toHaveLength(0)
+    expect(result.templates['test-agent']).toBeDefined()
+    expect(typeof result.templates['test-agent'].handleSteps).toBe('string')
+  })
+
+  test('should require set_output tool for handleSteps with json output mode', () => {
+    const {
+      DynamicAgentTemplateSchema,
+    } = require('@codebuff/common/types/dynamic-agent-template')
+
+    const agentConfig = {
+      id: 'test-agent',
+      version: '1.0.0',
+      name: 'Test Agent',
+      purpose: 'Testing',
+      model: 'claude-3-5-sonnet-20241022',
+      outputMode: 'json' as const,
+      toolNames: ['end_turn'], // Missing set_output
+      spawnableAgents: [],
+      systemPrompt: 'Test',
+      userInputPrompt: 'Test',
+      agentStepPrompt: 'Test',
+      initialAssistantMessage: '',
+      initialAssistantPrefix: '',
+      stepAssistantMessage: '',
+      stepAssistantPrefix: '',
+      handleSteps:
+        'function* () { yield { toolName: "set_output", args: {} } }',
+    }
+
+    const result = DynamicAgentTemplateSchema.safeParse(agentConfig)
+    expect(result.success).toBe(false)
+    if (!result.success) {
+      const errorMessage = result.error.issues[0]?.message || ''
+      expect(errorMessage).toContain('set_output')
+    }
+  })
+
+  test('should validate that handleSteps is a generator function', async () => {
+    const agentTemplates = {
+      'test-agent': {
+        ...mockAgentTemplate,
+        handleSteps: 'function () { return "not a generator" }', // Missing *
+      },
+    }
+
+    const fileContext: ProjectFileContext = {
+      ...mockFileContext,
+      agentTemplates,
+    }
+
+    const result = await dynamicAgentService.loadAgents(fileContext)
+
+    expect(result.validationErrors.length).toBeGreaterThan(0)
+    expect(result.validationErrors[0].message).toContain('generator function')
+    expect(result.validationErrors[0].message).toContain('function*')
+  })
+
+  test('should verify loaded template handleSteps matches original function toString', async () => {
+    // Create a generator function
+    const originalFunction = function* ({
+      agentState,
+      prompt,
+      params,
+    }: {
+      agentState: AgentState
+      prompt?: string
+      params?: any
+    }) {
+      yield {
+        toolName: 'set_output',
+        args: { message: 'Test output', data: params },
+      }
+    }
+
+    // Get the string representation
+    const expectedStringified = originalFunction.toString()
+
+    // Create agent templates with the function
+    const agentTemplates = {
+      'test-agent': {
+        ...mockAgentTemplate,
+        handleSteps: expectedStringified,
+      },
+    }
+
+    const fileContext: ProjectFileContext = {
+      ...mockFileContext,
+      agentTemplates,
+    }
+
+    // Load agents through the service
+    const result = await dynamicAgentService.loadAgents(fileContext)
+
+    // Verify no validation errors
+    expect(result.validationErrors).toHaveLength(0)
+    expect(result.templates['test-agent']).toBeDefined()
+    
+    // Verify the loaded template's handleSteps field matches the original toString
+    expect(result.templates['test-agent'].handleSteps).toBe(expectedStringified)
+    expect(typeof result.templates['test-agent'].handleSteps).toBe('string')
+  })
+})

backend/src/templates/dynamic-agent-service.ts

@@ -230,6 +230,18 @@ export class DynamicAgentService {
         }
       }
 
+      // Validate handleSteps if present
+      if (content.handleSteps) {
+        if (!content.handleSteps.includes('function*')) {
+          this.validationErrors.push({
+            filePath,
+            message: `handleSteps must be a generator function: "function* (params) { ... }". Found: ${content.handleSteps.substring(0, 50)}...`,
+            details: 'handleSteps should start with "function*" to be a valid generator function',
+          })
+          return
+        }
+      }
+
       // Determine outputMode: default to 'json' if outputSchema is present, otherwise 'last_message'
       const outputMode =
         content.outputMode ?? (content.outputSchema ? 'json' : 'last_message')