Agent callbacks

.changeset/shiny-wings-reply.md

@@ -0,0 +1,5 @@
+---
+"@browserbasehq/stagehand": patch
+---
+
+Add support for callbacks in stagehand agent

packages/core/lib/v3/cache/AgentCache.ts

@@ -20,7 +20,7 @@ import type {
   AgentResult,
   AgentStreamResult,
   AgentConfig,
-  AgentExecuteOptions,
+  AgentExecuteOptionsBase,
   Logger,
 } from "../types/public";
 import type { Page } from "../understudy/page";
@@ -74,7 +74,7 @@ export class AgentCache {
   }
 
   sanitizeExecuteOptions(
-    options?: AgentExecuteOptions,
+    options?: AgentExecuteOptionsBase,
   ): SanitizedAgentExecuteOptions {
     if (!options) return {};
     const sanitized: SanitizedAgentExecuteOptions = {};

packages/core/lib/v3/handlers/v3AgentHandler.ts

@@ -13,10 +13,14 @@
 import { LLMClient } from "../llm/LLMClient";
 import {
   AgentExecuteOptions,
+  AgentStreamExecuteOptions,
+  AgentExecuteOptionsBase,
   AgentResult,
   AgentContext,
   AgentState,
   AgentStreamResult,
+  AgentExecuteCallbacks,
+  AgentStreamCallbacks,
 } from "../types/public/agent";
 import { V3FunctionName } from "../types/public/methods";
 import { mapToolResultToActions } from "../agent/utils/actionMapping";
@@ -47,7 +51,7 @@
   }
 
   private async prepareAgent(
-    instructionOrOptions: string | AgentExecuteOptions,
+    instructionOrOptions: string | AgentExecuteOptionsBase,
   ): Promise<AgentContext> {
     try {
       const options =
@@ -102,7 +106,10 @@
     }
   }
 
-  private createStepHandler(state: AgentState) {
+  private createStepHandler(
+    state: AgentState,
+    userCallback?: (event: StepResult<ToolSet>) => PromiseLike<void> | void,
+  ) {
     return async (event: StepResult<ToolSet>) => {
       this.logger({
         category: "agent",
@@ -150,6 +157,10 @@
         }
         state.currentPageUrl = (await this.v3.context.awaitActivePage()).url();
       }
+
+      if (userCallback) {
+        await userCallback(event);
+      }
     };
   }
 
@@ -158,6 +169,7 @@
   ): Promise<AgentResult> {
     const startTime = Date.now();
     const {
+      options,
       maxSteps,
       systemPrompt,
       allTools,
@@ -166,6 +178,7 @@
       initialPageUrl,
     } = await this.prepareAgent(instructionOrOptions);
 
+    const callbacks = (instructionOrOptions as AgentExecuteOptions).callbacks;
     const state: AgentState = {
       collectedReasoning: [],
       actions: [],
@@ -183,7 +196,8 @@
         stopWhen: (result) => this.handleStop(result, maxSteps),
         temperature: 1,
         toolChoice: "auto",
-        onStepFinish: this.createStepHandler(state),
+        prepareStep: callbacks?.prepareStep,
+        onStepFinish: this.createStepHandler(state, callbacks?.onStepFinish),
       });
 
       return this.consolidateMetricsAndResult(startTime, state, result);
@@ -204,7 +218,7 @@
   }
 
   public async stream(
-    instructionOrOptions: string | AgentExecuteOptions,
+    instructionOrOptions: string | AgentStreamExecuteOptions,
   ): Promise<AgentStreamResult> {
     const {
       maxSteps,
@@ -215,6 +229,9 @@
       initialPageUrl,
     } = await this.prepareAgent(instructionOrOptions);
 
+    const callbacks = (instructionOrOptions as AgentStreamExecuteOptions)
+      .callbacks as AgentStreamCallbacks | undefined;
+
     const state: AgentState = {
       collectedReasoning: [],
       actions: [],
@@ -250,11 +267,19 @@
       stopWhen: (result) => this.handleStop(result, maxSteps),
       temperature: 1,
       toolChoice: "auto",
-      onStepFinish: this.createStepHandler(state),
-      onError: ({ error }) => {
-        handleError(error);
+      prepareStep: callbacks?.prepareStep,
+      onStepFinish: this.createStepHandler(state, callbacks?.onStepFinish),
+      onError: (event) => {
+        if (callbacks?.onError) {
+          callbacks.onError(event);
+        }
+        handleError(event.error);
       },
+      onChunk: callbacks?.onChunk,
       onFinish: (event) => {
+        if (callbacks?.onFinish) {
+          callbacks.onFinish(event);
+        }
         try {
           const result = this.consolidateMetricsAndResult(
             startTime,