feat: adding stagehand agent tool

Author: Kylejeong2

src/tools/agent.ts

@@ -0,0 +1,91 @@
+import { z } from "zod";
+import type { Tool, ToolSchema, ToolResult } from "./tool.js";
+import type { Context } from "../context.js";
+import type { ToolActionResult } from "../types/types.js";
+
+/**
+ * Stagehand Agent
+ * Docs: https://docs.stagehand.dev/basics/agent
+ *
+ * This tool uses Gemini Computer Use to autonomously complete web-based tasks.
+ * The agent will navigate, interact, and complete the task described in the prompt.
+ */
+
+const AgentInputSchema = z.object({
+  prompt: z.string().describe(
+    `The task prompt describing what you want the agent to accomplish.
+    Be clear and specific about the goal. For example:
+    'Go to Hacker News and find the most controversial post from today, then summarize the top 3 comments'.
+    The agent will autonomously navigate and interact with web pages to complete this task.`,
+  ),
+});
+
+type AgentInput = z.infer<typeof AgentInputSchema>;
+
+const agentSchema: ToolSchema<typeof AgentInputSchema> = {
+  name: "browserbase_stagehand_agent",
+  description: `Execute a task autonomously using Gemini Computer Use agent. The agent will navigate and interact with web pages to complete the given task.`,
+  inputSchema: AgentInputSchema,
+};
+
+async function handleAgent(
+  context: Context,
+  params: AgentInput,
+): Promise<ToolResult> {
+  const action = async (): Promise<ToolActionResult> => {
+    try {
+      const stagehand = await context.getStagehand();
+
+      // You need to provide GOOGLE_GENERATIVE_AI_API_KEY
+      const agent = stagehand.agent({
+        cua: true,
+        model: {
+          modelName: "google/gemini-2.5-computer-use-preview-10-2025",
+          apiKey:
+            process.env.GOOGLE_GENERATIVE_AI_API_KEY ||
+            process.env.GOOGLE_API_KEY ||
+            process.env.GEMINI_API_KEY,
+        },
+      });
+
+      // Execute the task
+      const result = await agent.execute({
+        instruction: params.prompt,
+        maxSteps: 20,
+      });
+
+      // Format response with both steps and result
+      // The result structure may vary, so we handle it flexibly
+      const resultData = result as unknown as Record<string, unknown>;
+      const response = {
+        result: resultData.result || result,
+        steps: resultData.steps || resultData.trace || [],
+      };
+
+      return {
+        content: [
+          {
+            type: "text",
+            text: `Agent execution completed:\n${JSON.stringify(response, null, 2)}`,
+          },
+        ],
+      };
+    } catch (error) {
+      const errorMsg = error instanceof Error ? error.message : String(error);
+      throw new Error(`Failed to execute agent task: ${errorMsg}`);
+    }
+  };
+
+  return {
+    action,
+    waitForNetwork: false,
+  };
+}
+
+const agentTool: Tool<typeof AgentInputSchema> = {
+  capability: "core",
+  schema: agentSchema,
+  handle: handleAgent,
+};
+
+export default agentTool;

src/tools/index.ts

@@ -5,6 +5,7 @@ import observeTool from "./observe.js";
 import screenshotTool from "./screenshot.js";
 import sessionTools from "./session.js";
 import getUrlTool from "./url.js";
+import agentTool from "./agent.js";
 
 // Export individual tools
 export { default as navigateTool } from "./navigate.js";
@@ -14,6 +15,7 @@ export { default as observeTool } from "./observe.js";
 export { default as screenshotTool } from "./screenshot.js";
 export { default as sessionTools } from "./session.js";
 export { default as getUrlTool } from "./url.js";
+export { default as agentTool } from "./agent.js";
 
 // Export all tools as array
 export const TOOLS = [
@@ -24,6 +26,7 @@ export const TOOLS = [
   observeTool,
   screenshotTool,
   getUrlTool,
+  agentTool,
 ];
 
 export const sessionManagementTools = sessionTools;