rm legacy tool call and separate tools into individual files

Author: Kylejeong2

stagehand/src/context.ts

@@ -2,7 +2,6 @@ import { Stagehand } from "@browserbasehq/stagehand";
 import { Server } from "@modelcontextprotocol/sdk/server/index.js";
 import type { Config } from "../config.js";
 import { CallToolResult, TextContent, ImageContent } from "@modelcontextprotocol/sdk/types.js";
-import { handleToolCall } from "./tools/tools.js";
 import { listResources, readResource } from "./resources.js";
 import { 
   ensureLogDirectory, 
@@ -172,10 +171,8 @@ export class Context {
     try {
       log(`Executing tool: ${tool.schema.name} with args: ${JSON.stringify(args)}`, 'info');
 
-      // Check if this tool has a handle method (new session tools)
-      // Only use handle method for session create and close tools
-      if ("handle" in tool && typeof tool.handle === "function" && 
-          (tool.schema.name === "browserbase_session_create" || tool.schema.name === "browserbase_session_close")) {
+      // Check if this tool has a handle method (new tool system)
+      if ("handle" in tool && typeof tool.handle === "function") {
         const toolResult = await tool.handle(this as any, args);
 
         if (toolResult?.action) {
@@ -193,10 +190,8 @@ export class Context {
           };
         }
       } else {
-        const stagehand = await this.getStagehand();
-        const result = await handleToolCall(tool.schema.name, args, stagehand);
-        log(`Tool ${tool.schema.name} completed successfully`, 'info');
-        return result;
+        // Fallback for any legacy tools without handle method
+        throw new Error(`Tool ${tool.schema.name} does not have a handle method`);
       }
     } catch (error) {
       const errorMessage = error instanceof Error ? error.message : String(error);

stagehand/src/index.ts

@@ -8,7 +8,7 @@ import { CallToolRequestSchema, ListResourcesRequestSchema, ListToolsRequestSche
 import { z } from "zod";
 import { zodToJsonSchema } from "zod-to-json-schema";
 import { Context } from "./context.js";
-import { TOOLS } from "./tools/tools.js";
+import { TOOLS } from "./tools/index.js";
 
 export async function createServer(config: Config): Promise<Server> {
   // Create the server

stagehand/src/tools/act.ts

@@ -0,0 +1,73 @@
+import { z } from "zod";
+import type { Tool, ToolSchema, ToolResult } from "./tool.js";
+import type { Context } from "../context.js";
+import type { ToolActionResult } from "../context.js";
+
+const ActInputSchema = z.object({
+  action: z.string().describe(
+    "The action to perform. Should be as atomic and specific as possible, " +
+    "i.e. 'Click the sign in button' or 'Type 'hello' into the search input'. AVOID actions that are more than one " +
+    "step, i.e. 'Order me pizza' or 'Send an email to Paul asking him to call me'. The instruction should be just as specific as possible, " +
+    "and have a strong correlation to the text on the page. If unsure, use observe before using act."
+  ),
+  variables: z.object({}).optional().describe(
+    "Variables used in the action template. ONLY use variables if you're dealing " +
+    "with sensitive data or dynamic content. For example, if you're logging in to a website, " +
+    "you can use a variable for the password. When using variables, you MUST have the variable " +
+    "key in the action template. For example: {\"action\": \"Fill in the password\", \"variables\": {\"password\": \"123456\"}}"
+  ),
+});
+
+type ActInput = z.infer<typeof ActInputSchema>;
+
+const actSchema: ToolSchema<typeof ActInputSchema> = {
+  name: "stagehand_act",
+  description: 
+    "Performs an action on a web page element. Act actions should be as atomic and " +
+    "specific as possible, i.e. \"Click the sign in button\" or \"Type 'hello' into the search input\". " +
+    "AVOID actions that are more than one step, i.e. \"Order me pizza\" or \"Send an email to Paul " +
+    "asking him to call me\".",
+  inputSchema: ActInputSchema,
+};
+
+async function handleAct(
+  context: Context,
+  params: ActInput
+): Promise<ToolResult> {
+  const action = async (): Promise<ToolActionResult> => {
+    try {
+      const stagehand = await context.getStagehand();
+      await stagehand.page.act({
+        action: params.action,
+        variables: params.variables,
+      });
+      
+      return {
+        content: [
+          {
+            type: "text",
+            text: `Action performed: ${params.action}`,
+          },
+        ],
+      };
+    } catch (error) {
+      const errorMsg = error instanceof Error ? error.message : String(error);
+      throw new Error(`Failed to perform action: ${errorMsg}`);
+    }
+  };
+
+  return {
+    action,
+    captureSnapshot: false,
+    code: [],
+    waitForNetwork: false,
+  };
+}
+
+const actTool: Tool<typeof ActInputSchema> = {
+  capability: "core",
+  schema: actSchema,
+  handle: handleAct,
+};
+
+export default actTool; 

stagehand/src/tools/extract.ts

@@ -0,0 +1,80 @@
+import { z } from "zod";
+import type { Tool, ToolSchema, ToolResult } from "./tool.js";
+import type { Context } from "../context.js";
+import type { ToolActionResult } from "../context.js";
+
+const ExtractInputSchema = z.object({
+  random_string: z.string().optional().describe("Dummy parameter for no-parameter tools"),
+});
+
+type ExtractInput = z.infer<typeof ExtractInputSchema>;
+
+const extractSchema: ToolSchema<typeof ExtractInputSchema> = {
+  name: "stagehand_extract",
+  description: "Extracts all of the text from the current page.",
+  inputSchema: ExtractInputSchema,
+};
+
+async function handleExtract(
+  context: Context,
+  params: ExtractInput
+): Promise<ToolResult> {
+  const action = async (): Promise<ToolActionResult> => {
+    try {
+      const page = await context.getActivePage();
+      if (!page) {
+        throw new Error("No active page available");
+      }
+
+      const bodyText = await page.evaluate(() => document.body.innerText);
+      const content = bodyText
+        .split("\n")
+        .map((line) => line.trim())
+        .filter((line) => {
+          if (!line) return false;
+
+          if (
+            (line.includes("{") && line.includes("}")) ||
+            line.includes("@keyframes") || // Remove CSS animations
+            line.match(/^\.[a-zA-Z0-9_-]+\s*{/) || // Remove CSS lines starting with .className {
+            line.match(/^[a-zA-Z-]+:[a-zA-Z0-9%\s\(\)\.,-]+;$/) // Remove lines like "color: blue;" or "margin: 10px;"
+          ) {
+            return false;
+          }
+          return true;
+        })
+        .map((line) => {
+          return line.replace(/\\u([0-9a-fA-F]{4})/g, (_, hex) =>
+            String.fromCharCode(parseInt(hex, 16))
+          );
+        });
+
+      return {
+        content: [
+          {
+            type: "text",
+            text: `Extracted content:\n${content.join("\n")}`,
+          },
+        ],
+      };
+    } catch (error) {
+      const errorMsg = error instanceof Error ? error.message : String(error);
+      throw new Error(`Failed to extract content: ${errorMsg}`);
+    }
+  };
+
+  return {
+    action,
+    captureSnapshot: false,
+    code: [],
+    waitForNetwork: false,
+  };
+}
+
+const extractTool: Tool<typeof ExtractInputSchema> = {
+  capability: "core",
+  schema: extractSchema,
+  handle: handleExtract,
+};
+
+export default extractTool; 

stagehand/src/tools/index.ts

@@ -0,0 +1,36 @@
+// Import all individual tools
+import navigateTool from "./navigate.js";
+import actTool from "./act.js";
+import extractTool from "./extract.js";
+import observeTool from "./observe.js";
+import screenshotTool from "./screenshot.js";
+import sessionTools from "./session.js";
+
+// Export individual tools
+export { default as navigateTool } from "./navigate.js";
+export { default as actTool } from "./act.js";
+export { default as extractTool } from "./extract.js";
+export { default as observeTool } from "./observe.js";
+export { default as screenshotTool } from "./screenshot.js";
+export { default as sessionTools } from "./session.js";
+
+// Export all tools as array
+export const TOOLS = [
+  ...sessionTools,
+  navigateTool,
+  actTool,
+  extractTool,
+  observeTool,
+  screenshotTool,
+];
+
+// Export tools by category
+export const coreTools = [
+  navigateTool,
+  actTool,
+  extractTool,
+  observeTool,
+  screenshotTool,
+];
+
+export const sessionManagementTools = sessionTools; 

stagehand/src/tools/navigate.ts

@@ -0,0 +1,64 @@
+import { z } from "zod";
+import type { Tool, ToolSchema, ToolResult } from "./tool.js";
+import type { Context } from "../context.js";
+import type { ToolActionResult } from "../context.js";
+import { operationLogs } from "../logging.js";
+
+const NavigateInputSchema = z.object({
+  url: z.string().describe("The URL to navigate to"),
+});
+
+type NavigateInput = z.infer<typeof NavigateInputSchema>;
+
+const navigateSchema: ToolSchema<typeof NavigateInputSchema> = {
+  name: "stagehand_navigate",
+  description:
+    "Navigate to a URL in the browser. Only use this tool with URLs you're confident will work and stay up to date. Otheriwse use https://google.com as the starting point",
+  inputSchema: NavigateInputSchema,
+};
+
+async function handleNavigate(
+  context: Context,
+  params: NavigateInput
+): Promise<ToolResult> {
+  const action = async (): Promise<ToolActionResult> => {
+    try {
+      const page = await context.getActivePage();
+      if (!page) {
+        throw new Error("No active page available");
+      }
+      await page.goto(params.url, { waitUntil: "domcontentloaded" });
+      
+      return {
+        content: [
+          {
+            type: "text",
+            text: `Navigated to: ${params.url}`,
+          },
+          {
+            type: "text",
+            text: `View the live session here: https://browserbase.com/sessions/${context.currentSessionId}`,
+          },
+        ],
+      };
+    } catch (error) {
+      const errorMsg = error instanceof Error ? error.message : String(error);
+      throw new Error(`Failed to navigate: ${errorMsg}`);
+    }
+  };
+
+  return {
+    action,
+    captureSnapshot: false,
+    code: [],
+    waitForNetwork: false,
+  };
+}
+
+const navigateTool: Tool<typeof NavigateInputSchema> = {
+  capability: "core",
+  schema: navigateSchema,
+  handle: handleNavigate,
+};
+
+export default navigateTool; 

stagehand/src/tools/observe.ts

@@ -0,0 +1,61 @@
+import { z } from "zod";
+import type { Tool, ToolSchema, ToolResult } from "./tool.js";
+import type { Context } from "../context.js";
+import type { ToolActionResult } from "../context.js";
+
+const ObserveInputSchema = z.object({
+  instruction: z.string().describe(
+    "Instruction for observation (e.g., 'find the login button'). This instruction must be extremely specific."
+  ),
+});
+
+type ObserveInput = z.infer<typeof ObserveInputSchema>;
+
+const observeSchema: ToolSchema<typeof ObserveInputSchema> = {
+  name: "stagehand_observe",
+  description:
+    "Observes elements on the web page. Use this tool to observe elements that you can later use in an action. Use observe instead of extract when dealing with actionable (interactable) elements rather than text. More often than not, you'll want to use extract instead of observe when dealing with scraping or extracting structured text.",
+  inputSchema: ObserveInputSchema,
+};
+
+async function handleObserve(
+  context: Context,
+  params: ObserveInput
+): Promise<ToolResult> {
+  const action = async (): Promise<ToolActionResult> => {
+    try {
+      const stagehand = await context.getStagehand();
+      const observations = await stagehand.page.observe({
+        instruction: params.instruction,
+        returnAction: false,
+      });
+
+      return {
+        content: [
+          {
+            type: "text",
+            text: `Observations: ${JSON.stringify(observations)}`,
+          },
+        ],
+      };
+    } catch (error) {
+      const errorMsg = error instanceof Error ? error.message : String(error);
+      throw new Error(`Failed to observe: ${errorMsg}`);
+    }
+  };
+
+  return {
+    action,
+    captureSnapshot: false,
+    code: [],
+    waitForNetwork: false,
+  };
+}
+
+const observeTool: Tool<typeof ObserveInputSchema> = {
+  capability: "core",
+  schema: observeSchema,
+  handle: handleObserve,
+};
+
+export default observeTool;