pass stagehand, instead of stagehandPage to agent

tkattkat · tkattkat · commit 3ea4178b605f · 2025-09-17T15:56:36.000-07:00
diff --git a/lib/agent/tools/act.ts b/lib/agent/tools/act.ts
@@ -1,12 +1,9 @@
 import { tool } from "ai";
 import { z } from "zod/v3";
-import { StagehandPage } from "../../StagehandPage";
+import { Stagehand } from "../../index";
 import { buildActObservePrompt } from "../../prompt";
 import { SupportedPlaywrightAction } from "@/types/act";
-export const createActTool = (
-  stagehandPage: StagehandPage,
-  executionModel?: string,
-) =>
+export const createActTool = (stagehand: Stagehand, executionModel?: string) =>
   tool({
     description: "Perform an action on the page (click, type)",
     parameters: z.object({
@@ -34,7 +31,7 @@ export const createActTool = (
               instruction: builtPrompt,
             };
 
-        const observeResults = await stagehandPage.page.observe(observeOptions);
+        const observeResults = await stagehand.page.observe(observeOptions);
 
         if (!observeResults || observeResults.length === 0) {
           return {
@@ -60,7 +57,7 @@ export const createActTool = (
               };
 
           const iframeObserveResults =
-            await stagehandPage.page.observe(iframeObserveOptions);
+            await stagehand.page.observe(iframeObserveOptions);
 
           if (!iframeObserveResults || iframeObserveResults.length === 0) {
             return {
@@ -71,7 +68,7 @@ export const createActTool = (
           }
 
           const iframeObserveResult = iframeObserveResults[0];
-          const fallback = await stagehandPage.page.act(iframeObserveResult);
+          const fallback = await stagehand.page.act(iframeObserveResult);
 
           return {
             success: fallback.success,
@@ -86,7 +83,7 @@ export const createActTool = (
           };
         }
 
-        const result = await stagehandPage.page.act(observeResult);
+        const result = await stagehand.page.act(observeResult);
         const playwrightArguments = {
           description: observeResult.description,
           method: observeResult.method,
diff --git a/lib/agent/tools/ariaTree.ts b/lib/agent/tools/ariaTree.ts
@@ -1,15 +1,15 @@
 import { tool } from "ai";
 import { z } from "zod/v3";
-import { StagehandPage } from "../../StagehandPage";
+import { Stagehand } from "../../index";
 
-export const createAriaTreeTool = (stagehandPage: StagehandPage) =>
+export const createAriaTreeTool = (stagehand: Stagehand) =>
   tool({
     description:
       "gets the accessibility (ARIA) tree from the current page. this is useful for understanding the page structure and accessibility features. it should provide full context of what is on the page",
     parameters: z.object({}),
     execute: async () => {
-      const { page_text } = await stagehandPage.page.extract();
-      const pageUrl = stagehandPage.page.url();
+      const { page_text } = await stagehand.page.extract();
+      const pageUrl = stagehand.page.url();
 
       let content = page_text;
       const MAX_CHARACTERS = 70000;
diff --git a/lib/agent/tools/extract.ts b/lib/agent/tools/extract.ts
@@ -1,6 +1,6 @@
 import { tool } from "ai";
 import { z } from "zod/v3";
-import { StagehandPage } from "../../StagehandPage";
+import { Stagehand } from "../../index";
 import { LogLine } from "@/types/log";
 
 /**
@@ -33,7 +33,7 @@ function evaluateZodSchema(
 }
 
 export const createExtractTool = (
-  stagehandPage: StagehandPage,
+  stagehand: Stagehand,
   executionModel?: string,
   logger?: (message: LogLine) => void,
 ) =>
@@ -80,7 +80,7 @@ export const createExtractTool = (
             : z.object({ result: zodSchema });
 
         // Extract with the schema - only pass modelName if executionModel is explicitly provided
-        const result = await stagehandPage.page.extract({
+        const result = await stagehand.page.extract({
           instruction,
           schema: schemaObject,
           ...(executionModel && { modelName: executionModel }),
diff --git a/lib/agent/tools/fillform.ts b/lib/agent/tools/fillform.ts
@@ -1,9 +1,9 @@
 import { tool } from "ai";
 import { z } from "zod/v3";
-import { StagehandPage } from "../../StagehandPage";
+import { Stagehand } from "../../index";
 
 export const createFillFormTool = (
-  stagehandPage: StagehandPage,
+  stagehand: Stagehand,
   executionModel?: string,
 ) =>
   tool({
@@ -54,15 +54,15 @@ export const createFillFormTool = (
         .join(", ")}`;
 
       const observeResults = executionModel
-        ? await stagehandPage.page.observe({
+        ? await stagehand.page.observe({
             instruction,
             modelName: executionModel,
           })
-        : await stagehandPage.page.observe(instruction);
+        : await stagehand.page.observe(instruction);
 
       const completedActions = [];
       for (const result of observeResults) {
-        const action = await stagehandPage.page.act(result);
+        const action = await stagehand.page.act(result);
         completedActions.push(action);
       }
 
diff --git a/lib/agent/tools/goto.ts b/lib/agent/tools/goto.ts
@@ -1,16 +1,16 @@
 import { tool } from "ai";
 import { z } from "zod/v3";
-import { StagehandPage } from "../../StagehandPage";
+import { Stagehand } from "../../index";
 
-export const createGotoTool = (stagehandPage: StagehandPage) =>
+export const createGotoTool = (stagehand: Stagehand) =>
   tool({
     description: "Navigate to a specific URL",
     parameters: z.object({
       url: z.string().describe("The URL to navigate to"),
     }),
     execute: async ({ url }) => {
       try {
-        await stagehandPage.page.goto(url, { waitUntil: "load" });
+        await stagehand.page.goto(url, { waitUntil: "load" });
         return { success: true, url };
       } catch (error) {
         return { success: false, error: error.message };
diff --git a/lib/agent/tools/index.ts b/lib/agent/tools/index.ts
@@ -7,7 +7,7 @@ import { createCloseTool } from "./close";
 import { createAriaTreeTool } from "./ariaTree";
 import { createFillFormTool } from "./fillform";
 import { createScrollTool } from "./scroll";
-import { StagehandPage } from "../../StagehandPage";
+import { Stagehand } from "../../index";
 import { LogLine } from "@/types/log";
 import { createExtractTool } from "./extract";
 
@@ -17,21 +17,21 @@ export interface AgentToolOptions {
 }
 
 export function createAgentTools(
-  stagehandPage: StagehandPage,
+  stagehand: Stagehand,
   options?: AgentToolOptions,
 ) {
   const executionModel = options?.executionModel;
 
   return {
-    act: createActTool(stagehandPage, executionModel),
-    ariaTree: createAriaTreeTool(stagehandPage),
+    act: createActTool(stagehand, executionModel),
+    ariaTree: createAriaTreeTool(stagehand),
     close: createCloseTool(),
-    extract: createExtractTool(stagehandPage, executionModel, options?.logger),
-    fillForm: createFillFormTool(stagehandPage, executionModel),
-    goto: createGotoTool(stagehandPage),
-    navback: createNavBackTool(stagehandPage),
-    screenshot: createScreenshotTool(stagehandPage),
-    scroll: createScrollTool(stagehandPage),
+    extract: createExtractTool(stagehand, executionModel, options?.logger),
+    fillForm: createFillFormTool(stagehand, executionModel),
+    goto: createGotoTool(stagehand),
+    navback: createNavBackTool(stagehand),
+    screenshot: createScreenshotTool(stagehand),
+    scroll: createScrollTool(stagehand),
     wait: createWaitTool(),
   };
 }
diff --git a/lib/agent/tools/navback.ts b/lib/agent/tools/navback.ts
@@ -1,15 +1,15 @@
 import { tool } from "ai";
 import { z } from "zod/v3";
-import { StagehandPage } from "../../StagehandPage";
+import { Stagehand } from "../../index";
 
-export const createNavBackTool = (stagehandPage: StagehandPage) =>
+export const createNavBackTool = (stagehand: Stagehand) =>
   tool({
     description: "Navigate back to the previous page",
     parameters: z.object({
       reasoning: z.string().describe("Why you're going back"),
     }),
     execute: async () => {
-      await stagehandPage.page.goBack();
+      await stagehand.page.goBack();
       return { success: true };
     },
   });
diff --git a/lib/agent/tools/screenshot.ts b/lib/agent/tools/screenshot.ts
@@ -1,18 +1,18 @@
 import { tool } from "ai";
 import { z } from "zod/v3";
-import { StagehandPage } from "../../StagehandPage";
+import { Stagehand } from "../../index";
 
-export const createScreenshotTool = (stagehandPage: StagehandPage) =>
+export const createScreenshotTool = (stagehand: Stagehand) =>
   tool({
     description:
       "Takes a screenshot of the current page. Use this tool to learn where you are on the page, or to get context of elements on the page",
     parameters: z.object({}),
     execute: async () => {
-      const screenshotBuffer = await stagehandPage.page.screenshot({
+      const screenshotBuffer = await stagehand.page.screenshot({
         fullPage: false,
         type: "jpeg",
       });
-      const pageUrl = stagehandPage.page.url();
+      const pageUrl = stagehand.page.url();
 
       return {
         base64: screenshotBuffer.toString("base64"),
diff --git a/lib/agent/tools/scroll.ts b/lib/agent/tools/scroll.ts
@@ -1,16 +1,16 @@
 import { tool } from "ai";
 import { z } from "zod/v3";
-import { StagehandPage } from "../../StagehandPage";
+import { Stagehand } from "../../index";
 
-export const createScrollTool = (stagehandPage: StagehandPage) =>
+export const createScrollTool = (stagehand: Stagehand) =>
   tool({
     description: "Scroll the page",
     parameters: z.object({
       pixels: z.number().describe("Number of pixels to scroll up or down"),
       direction: z.enum(["up", "down"]).describe("Direction to scroll"),
     }),
     execute: async ({ pixels, direction }) => {
-      await stagehandPage.page.mouse.wheel(
+      await stagehand.page.mouse.wheel(
         0,
         direction === "up" ? -pixels : pixels,
       );
diff --git a/lib/handlers/stagehandAgentHandler.ts b/lib/handlers/stagehandAgentHandler.ts
@@ -5,31 +5,31 @@ import {
   ActToolResult,
 } from "@/types/agent";
 import { LogLine } from "@/types/log";
-import { StagehandPage } from "../StagehandPage";
 import { LLMClient } from "../llm/LLMClient";
 import { CoreMessage, wrapLanguageModel } from "ai";
 import { LanguageModel } from "ai";
 import { processMessages } from "../agent/utils/messageProcessing";
 import { createAgentTools } from "../agent/tools";
 import { ToolSet } from "ai";
+import { Stagehand } from "../index";
 
 export class StagehandAgentHandler {
-  private stagehandPage: StagehandPage;
+  private stagehand: Stagehand;
   private logger: (message: LogLine) => void;
   private llmClient: LLMClient;
   private executionModel?: string;
   private systemInstructions?: string;
   private mcpTools?: ToolSet;
 
   constructor(
-    stagehandPage: StagehandPage,
+    stagehand: Stagehand,
     logger: (message: LogLine) => void,
     llmClient: LLMClient,
     executionModel?: string,
     systemInstructions?: string,
     mcpTools?: ToolSet,
   ) {
-    this.stagehandPage = stagehandPage;
+    this.stagehand = stagehand;
     this.logger = logger;
     this.llmClient = llmClient;
     this.executionModel = executionModel;
@@ -239,7 +239,7 @@ For each action, provide clear reasoning about why you're taking that step.`;
   }
 
   private createTools() {
-    return createAgentTools(this.stagehandPage, {
+    return createAgentTools(this.stagehand, {
       executionModel: this.executionModel,
       logger: this.logger,
     });
diff --git a/lib/index.ts b/lib/index.ts
@@ -943,7 +943,7 @@ export class Stagehand {
             ? await resolveTools(options?.integrations, options?.tools)
             : (options?.tools ?? {});
           return new StagehandAgentHandler(
-            this.stagehandPage,
+            this,
             this.logger,
             this.llmClient,
             executionModel,
