feat: added logger in handleSteps

brandonkachen · brandonkachen · commit 2407e2b21507 · 2025-09-24T22:28:57.000-07:00
🤖 Generated with Codebuff Co-Authored-By: Codebuff
&lt;noreply@codebuff.com&gt;
diff --git a/.agents/types/agent-definition.ts b/.agents/types/agent-definition.ts
@@ -179,7 +179,7 @@ export interface AgentDefinition {
    * }
    * }
    */
-  handleSteps?: (context: AgentStepContext) => Generator<
+  handleSteps?: (context: AgentStepContext, logger?: Logger) => Generator<
     ToolCall | 'STEP' | 'STEP_ALL',
     void,
     {
@@ -194,6 +194,13 @@ export interface AgentDefinition {
 // Supporting Types
 // ============================================================================
 
+export interface Logger {
+  debug: (data: any, msg?: string) => void
+  info: (data: any, msg?: string) => void
+  warn: (data: any, msg?: string) => void
+  error: (data: any, msg?: string) => void
+}
+
 export interface AgentState {
   agentId: string
   runId: string
diff --git a/backend/src/run-programmatic-step.ts b/backend/src/run-programmatic-step.ts
@@ -74,6 +74,18 @@ export async function runProgrammaticStep(
     stepNumber: number
   },
 ): Promise<{ agentState: AgentState; endTurn: boolean; stepNumber: number }> {
+  logger.info(
+    {
+      agentType: template.id,
+      runId: agentState.runId,
+      hasHandleSteps: !!template.handleSteps,
+      handleStepsType: typeof template.handleSteps,
+      stepNumber,
+      stepsComplete,
+    },
+    'runProgrammaticStep: Starting programmatic step execution',
+  )
+
   if (!template.handleSteps) {
     throw new Error('No step handler found for agent template ' + template.id)
   }
@@ -86,9 +98,37 @@ export async function runProgrammaticStep(
   let generator = runIdToGenerator[agentState.runId]
   let sandbox = sandboxManager.getSandbox(agentState.runId)
 
-  // Check if we need to initialize a generator (either native or QuickJS-based)
+  // Check if we need to initialize a generator
   if (!generator && !sandbox) {
+    const createLogMethod =
+      (level: 'debug' | 'info' | 'warn' | 'error') =>
+      (data: any, msg?: string) => {
+        logger[level](data, msg) // Log to backend
+        sendAction(ws, {
+          type: 'handlesteps-log-chunk',
+          userInputId,
+          agentId: agentState.agentId,
+          level,
+          data,
+          message: msg,
+        })
+      }
+
+    const streamingLogger = {
+      debug: createLogMethod('debug'),
+      info: createLogMethod('info'),
+      warn: createLogMethod('warn'),
+      error: createLogMethod('error'),
+    }
+
     if (typeof template.handleSteps === 'string') {
+      logger.info(
+        {
+          agentType: template.id,
+          runId: agentState.runId,
+        },
+        'runProgrammaticStep: Initializing QuickJS sandbox for string-based generator',
+      )
       // Initialize QuickJS sandbox for string-based generator
       sandbox = await sandboxManager.getOrCreateSandbox(
         agentState.runId,
@@ -98,15 +138,35 @@ export async function runProgrammaticStep(
           prompt,
           params,
         },
+        undefined, // config
+        streamingLogger, // pass the streaming logger instance
       )
     } else {
+      logger.info(
+        {
+          agentType: template.id,
+          runId: agentState.runId,
+        },
+        'runProgrammaticStep: Initializing native JavaScript generator',
+      )
       // Initialize native generator
-      generator = template.handleSteps({
-        agentState,
-        prompt,
-        params,
-      })
+      generator = (template.handleSteps as any)(
+        {
+          agentState,
+          prompt,
+          params,
+        },
+        streamingLogger,
+      )
       runIdToGenerator[agentState.runId] = generator
+      logger.info(
+        {
+          agentType: template.id,
+          runId: agentState.runId,
+          generatorInitialized: !!generator,
+        },
+        'runProgrammaticStep: Native generator initialized successfully',
+      )
     }
   }
 
@@ -168,6 +228,17 @@ export async function runProgrammaticStep(
       creditsBefore = state.agentState.directCreditsUsed
       childrenBefore = state.agentState.childRunIds.length
 
+      logger.info(
+        {
+          agentType: template.id,
+          runId: agentState.runId,
+          usingSandbox: !!sandbox,
+          usingGenerator: !!generator,
+          stepsComplete,
+        },
+        'runProgrammaticStep: About to execute generator step',
+      )
+
       const result = sandbox
         ? await sandbox.executeStep({
             agentState: getPublicAgentState(state.agentState),
@@ -180,6 +251,16 @@ export async function runProgrammaticStep(
             stepsComplete,
           })
 
+      logger.info(
+        {
+          agentType: template.id,
+          runId: agentState.runId,
+          resultDone: result.done,
+          resultValue: result.value,
+        },
+        'runProgrammaticStep: Generator step executed, got result',
+      )
+
       if (result.done) {
         endTurn = true
         break
diff --git a/backend/src/util/quickjs-sandbox.ts b/backend/src/util/quickjs-sandbox.ts
@@ -52,6 +52,12 @@ export class QuickJSSandbox {
     generatorCode: string,
     initialInput: any,
     config: SandboxConfig = {},
+    logger?: {
+      debug: (data: any, msg?: string) => void
+      info: (data: any, msg?: string) => void
+      warn: (data: any, msg?: string) => void
+      error: (data: any, msg?: string) => void
+    },
   ): Promise<QuickJSSandbox> {
     const {
       memoryLimit = 1024 * 1024 * 20, // 20MB
@@ -80,6 +86,45 @@ export class QuickJSSandbox {
     const context = runtime.newContext()
 
     try {
+      // Set up logger handler
+      const loggerHandler = context.newFunction(
+        '_loggerHandler',
+        (level, data, msg) => {
+          try {
+            const levelStr = context.getString(level)
+            let dataObj: any
+            let msgStr: string | undefined
+
+            try {
+              dataObj = data ? JSON.parse(context.getString(data)) : undefined
+            } catch {
+              dataObj = context.getString(data)
+            }
+
+            msgStr = msg ? context.getString(msg) : undefined
+
+            if (logger) {
+              if (levelStr === 'debug' && logger.debug) {
+                logger.debug(dataObj, msgStr)
+              } else if (levelStr === 'info' && logger.info) {
+                logger.info(dataObj, msgStr)
+              } else if (levelStr === 'warn' && logger.warn) {
+                logger.warn(dataObj, msgStr)
+              } else if (levelStr === 'error' && logger.error) {
+                logger.error(dataObj, msgStr)
+              }
+            }
+          } catch (err) {
+            // Fallback for logging errors
+            if (logger?.error) {
+              logger.error({ error: err }, 'Logger handler error')
+            }
+          }
+        },
+      )
+
+      context.setProp(context.global, '_loggerHandler', loggerHandler)
+      loggerHandler.dispose()
       // Inject safe globals and the generator function
       const setupCode = `
         // Safe console implementation
@@ -89,11 +134,24 @@ export class QuickJSSandbox {
           warn: (...args) => undefined
         };
         
+        // Logger implementation
+        const createLogMethod = (level) => (data, msg) => 
+          globalThis._loggerHandler(level, 
+            typeof data === 'object' ? JSON.stringify(data) : String(data), 
+            msg ? String(msg) : undefined);
+        
+        const logger = {
+          debug: createLogMethod('debug'),
+          info: createLogMethod('info'),
+          warn: createLogMethod('warn'),
+          error: createLogMethod('error')
+        };
+        
         // Agent function
         const handleSteps = ${generatorCode};
         
         // Create generator instance
-        let generator = handleSteps(${JSON.stringify(initialInput)});
+        let generator = handleSteps(${JSON.stringify(initialInput)}, logger);
         
         // Generator management
         globalThis._generator = generator;
@@ -213,6 +271,12 @@ export class SandboxManager {
     generatorCode: string,
     initialInput: any,
     config?: SandboxConfig,
+    logger?: {
+      debug: (data: any, msg?: string) => void
+      info: (data: any, msg?: string) => void
+      warn: (data: any, msg?: string) => void
+      error: (data: any, msg?: string) => void
+    },
   ): Promise<QuickJSSandbox> {
     const existing = this.sandboxes.get(runId)
     if (existing && existing.isInitialized()) {
@@ -229,6 +293,7 @@ export class SandboxManager {
       generatorCode,
       initialInput,
       config,
+      logger,
     )
     this.sandboxes.set(runId, sandbox)
     return sandbox
diff --git a/common/src/actions.ts b/common/src/actions.ts
@@ -130,6 +130,14 @@ export const SERVER_ACTION_SCHEMA = z.discriminatedUnion('type', [
     chunk: z.string(),
     prompt: z.string().optional(),
   }),
+  z.object({
+    type: z.literal('handlesteps-log-chunk'),
+    userInputId: z.string(),
+    agentId: z.string(),
+    level: z.enum(['debug', 'info', 'warn', 'error']),
+    data: z.any(),
+    message: z.string().optional(),
+  }),
   PromptResponseSchema,
   z.object({
     type: z.literal('read-files'),
diff --git a/common/src/types/dynamic-agent-template.ts b/common/src/types/dynamic-agent-template.ts
@@ -68,6 +68,26 @@ export type PromptField = z.infer<typeof PromptFieldSchema>
 
 const functionSchema = <T extends z.core.$ZodFunction>(schema: T) =>
   z.custom<Parameters<T['implement']>[0]>((fn: any) => schema.implement(fn))
+// Schema for the Logger interface
+const LoggerSchema = z.object({
+  debug: z.function({
+    input: [z.any(), z.string().optional()],
+    output: z.void(),
+  }),
+  info: z.function({
+    input: [z.any(), z.string().optional()],
+    output: z.void(),
+  }),
+  warn: z.function({
+    input: [z.any(), z.string().optional()],
+    output: z.void(),
+  }),
+  error: z.function({
+    input: [z.any(), z.string().optional()],
+    output: z.void(),
+  }),
+})
+
 // Schema for validating handleSteps function signature
 const HandleStepsSchema = functionSchema(
   z.function({
@@ -81,6 +101,7 @@ const HandleStepsSchema = functionSchema(
         prompt: z.string().optional(),
         params: z.any().optional(),
       }),
+      LoggerSchema.optional(),
     ],
     output: z.any(),
   }),
diff --git a/npm-app/src/client.ts b/npm-app/src/client.ts
@@ -951,6 +951,46 @@ export class Client {
       // Refresh display if we're currently viewing this agent
       refreshSubagentDisplay(agentId)
     })
+
+    // Handle handleSteps log streaming
+    this.webSocket.subscribe('handlesteps-log-chunk', (action) => {
+      const { agentId, level, data, message } = action
+
+      // Format the log message for display
+      const formattedMessage = this.formatLogMessage(level, data, message)
+
+      // Display the log message immediately
+      if (formattedMessage) {
+        process.stdout.write(formattedMessage + '\n')
+      }
+    })
+  }
+
+  private formatLogMessage(level: string, data: any, message?: string): string {
+    const timestamp = new Date().toISOString().substring(11, 23) // HH:MM:SS.mmm
+    const levelColors = { debug: blue, info: green, warn: yellow, error: red }
+    const levelColor =
+      levelColors[level as keyof typeof levelColors] || ((s: string) => s)
+
+    const timeTag = `[${timestamp}]`
+    const levelTag = levelColor(`[${level.toUpperCase()}]`)
+    const dataStr = this.serializeLogData(data)
+
+    return [timeTag, levelTag, message, dataStr].filter(Boolean).join(' ')
+  }
+
+  private serializeLogData(data: any): string {
+    if (data === undefined || data === null) return ''
+
+    if (typeof data === 'object') {
+      try {
+        return JSON.stringify(data, null, 2)
+      } catch {
+        return String(data)
+      }
+    }
+
+    return String(data)
   }
 
   private showUsageWarning() {
