add hooks to agent

.changeset/cold-humans-work.md

@@ -0,0 +1,5 @@
+---
+"@browserbasehq/stagehand": minor
+---
+
+Added support for execute-level hooks on agents. This includes `onStep(action)`, `onSuccess(reuslt)`, and `onFailure(error)` inputs.

examples/cua-example.ts

@@ -27,7 +27,7 @@ async function main() {
 
     const agent = stagehand.agent({
       provider: "openai",
-      model: "computer-use-preview-2025-02-04",
+      model: "computer-use-preview",
       instructions: `You are a helpful assistant that can use a web browser.
       You are currently on the following page: ${page.url()}.
       Do not ask follow up questions, the user will trust your judgement.`,

lib/agent/AgentProvider.ts

@@ -1,12 +1,12 @@
-import { LogLine } from "@/types/log";
-import { AgentClient } from "./AgentClient";
 import { AgentType } from "@/types/agent";
-import { OpenAICUAClient } from "./OpenAICUAClient";
-import { AnthropicCUAClient } from "./AnthropicCUAClient";
+import { LogLine } from "@/types/log";
 import {
   UnsupportedModelError,
   UnsupportedModelProviderError,
 } from "@/types/stagehandErrors";
+import { AgentClient } from "./AgentClient";
+import { AnthropicCUAClient } from "./AnthropicCUAClient";
+import { OpenAICUAClient } from "./OpenAICUAClient";
 
 // Map model names to their provider types
 const modelToAgentProviderMap: Record<string, AgentType> = {
@@ -22,7 +22,6 @@ const modelToAgentProviderMap: Record<string, AgentType> = {
  */
 export class AgentProvider {
   private logger: (message: LogLine) => void;
-
   /**
    * Create a new agent provider
    */

lib/agent/AnthropicCUAClient.ts

@@ -1,18 +1,19 @@
-import Anthropic from "@anthropic-ai/sdk";
-import { LogLine } from "@/types/log";
 import {
   AgentAction,
+  AgentExecuteOptions,
+  AgentExecutionOptions,
   AgentResult,
   AgentType,
-  AgentExecutionOptions,
-  ToolUseItem,
-  AnthropicMessage,
   AnthropicContentBlock,
+  AnthropicMessage,
   AnthropicTextBlock,
   AnthropicToolResult,
+  ToolUseItem,
 } from "@/types/agent";
-import { AgentClient } from "./AgentClient";
+import { LogLine } from "@/types/log";
 import { AgentScreenshotProviderError } from "@/types/stagehandErrors";
+import Anthropic from "@anthropic-ai/sdk";
+import { AgentClient } from "./AgentClient";
 
 export type ResponseInputItem = AnthropicMessage | AnthropicToolResult;
 
@@ -116,7 +117,7 @@ export class AnthropicCUAClient extends AgentClient {
           level: 2,
         });
 
-        const result = await this.executeStep(inputItems, logger);
+        const result = await this.executeStep(inputItems, logger, options);
 
         // Add actions to the list
         if (result.actions.length > 0) {
@@ -153,12 +154,16 @@ export class AnthropicCUAClient extends AgentClient {
       });
 
       // Return the final result
-      return {
+      const result = {
         success: completed,
         actions,
         message: finalMessage,
         completed,
       };
+
+      options.onSuccess?.(result);
+
+      return result;
     } catch (error) {
       const errorMessage =
         error instanceof Error ? error.message : String(error);
@@ -168,6 +173,8 @@ export class AnthropicCUAClient extends AgentClient {
         level: 0,
       });
 
+      options.onFailure?.(error);
+
       return {
         success: false,
         actions,
@@ -180,6 +187,7 @@ export class AnthropicCUAClient extends AgentClient {
   async executeStep(
     inputItems: ResponseInputItem[],
     logger: (message: LogLine) => void,
+    options: AgentExecuteOptions,
   ): Promise<{
     actions: AgentAction[];
     message: string;
@@ -270,6 +278,8 @@ export class AnthropicCUAClient extends AgentClient {
               message: `Executing action: ${action.type}`,
               level: 1,
             });
+
+            await options.onStep?.(action);
             await this.actionHandler(action);
           } catch (error) {
             const errorMessage =

lib/agent/OpenAICUAClient.ts

@@ -1,17 +1,18 @@
-import OpenAI from "openai";
-import { LogLine } from "../../types/log";
 import {
   AgentAction,
+  AgentExecuteOptions,
+  AgentExecutionOptions,
   AgentResult,
   AgentType,
-  AgentExecutionOptions,
-  ResponseInputItem,
-  ResponseItem,
   ComputerCallItem,
   FunctionCallItem,
+  ResponseInputItem,
+  ResponseItem,
 } from "@/types/agent";
-import { AgentClient } from "./AgentClient";
 import { AgentScreenshotProviderError } from "@/types/stagehandErrors";
+import OpenAI from "openai";
+import { LogLine } from "../../types/log";
+import { AgentClient } from "./AgentClient";
 
 /**
  * Client for OpenAI's Computer Use Assistant API
@@ -111,6 +112,7 @@ export class OpenAICUAClient extends AgentClient {
           inputItems,
           previousResponseId,
           logger,
+          options,
         );
 
         // Add actions to the list
@@ -137,13 +139,17 @@ export class OpenAICUAClient extends AgentClient {
         currentStep++;
       }
 
-      // Return the final result
-      return {
+      const result = {
         success: completed,
         actions,
         message: finalMessage,
         completed,
       };
+
+      options.onSuccess?.(result);
+
+      // Return the final result
+      return result;
     } catch (error) {
       const errorMessage =
         error instanceof Error ? error.message : String(error);
@@ -153,6 +159,8 @@ export class OpenAICUAClient extends AgentClient {
         level: 0,
       });
 
+      options.onFailure?.(error as Error);
+
       return {
         success: false,
         actions,
@@ -170,6 +178,7 @@ export class OpenAICUAClient extends AgentClient {
     inputItems: ResponseInputItem[],
     previousResponseId: string | undefined,
     logger: (message: LogLine) => void,
+    options: AgentExecuteOptions,
   ): Promise<{
     actions: AgentAction[];
     message: string;
@@ -224,7 +233,7 @@ export class OpenAICUAClient extends AgentClient {
       }
 
       // Take actions and get results
-      const nextInputItems = await this.takeAction(output, logger);
+      const nextInputItems = await this.takeAction(output, logger, options);
 
       // Check if completed
       const completed =
@@ -334,11 +343,14 @@ export class OpenAICUAClient extends AgentClient {
   async takeAction(
     output: ResponseItem[],
     logger: (message: LogLine) => void,
+    options: AgentExecuteOptions,
   ): Promise<ResponseInputItem[]> {
     const nextInputItems: ResponseInputItem[] = [];
 
     // Add any computer calls to process
     for (const item of output) {
+      await options.onStep?.(item);
+
       if (item.type === "computer_call" && this.isComputerCallItem(item)) {
         // Execute the action
         try {

lib/handlers/agentHandler.ts

@@ -1,15 +1,15 @@
-import { StagehandPage } from "../StagehandPage";
-import { AgentProvider } from "../agent/AgentProvider";
-import { StagehandAgent } from "../agent/StagehandAgent";
-import { AgentClient } from "../agent/AgentClient";
-import { LogLine } from "../../types/log";
 import {
-  AgentExecuteOptions,
+  ActionExecutionResult,
   AgentAction,
-  AgentResult,
+  AgentExecuteOptions,
   AgentHandlerOptions,
-  ActionExecutionResult,
+  AgentResult,
 } from "@/types/agent";
+import { LogLine } from "../../types/log";
+import { StagehandPage } from "../StagehandPage";
+import { AgentClient } from "../agent/AgentClient";
+import { AgentProvider } from "../agent/AgentProvider";
+import { StagehandAgent } from "../agent/StagehandAgent";
 
 export class StagehandAgentHandler {
   private stagehandPage: StagehandPage;
@@ -27,7 +27,6 @@ export class StagehandAgentHandler {
     this.stagehandPage = stagehandPage;
     this.logger = logger;
     this.options = options;
-
     // Initialize the provider
     this.provider = new AgentProvider(logger);