Shreypandya/print stagehand metrics

src/mcp/usage.ts

@@ -0,0 +1,115 @@
+import type { Stagehand } from "@browserbasehq/stagehand";
+
+export type StagehandUsageOperation = string;
+
+export type StagehandUsageKey = {
+  sessionId: string;
+  toolName: string;
+  operation: StagehandUsageOperation;
+};
+
+export type StagehandOperationStats = {
+  callCount: number;
+  toolCallCounts: Record<string, number>;
+};
+
+export type StagehandSessionUsage = {
+  operations: Record<StagehandUsageOperation, StagehandOperationStats>;
+};
+
+export type StagehandUsageSnapshot = {
+  global: Record<StagehandUsageOperation, StagehandOperationStats>;
+  perSession: Record<string, StagehandSessionUsage>;
+};
+
+const globalUsage: Record<string, StagehandOperationStats> = {};
+const perSessionUsage: Record<string, StagehandSessionUsage> = {};
+
+function getOrCreateOperationStats(
+  container: Record<StagehandUsageOperation, StagehandOperationStats>,
+  operation: StagehandUsageOperation,
+): StagehandOperationStats {
+  if (!container[operation]) {
+    container[operation] = {
+      callCount: 0,
+      toolCallCounts: {},
+    };
+  }
+  return container[operation];
+}
+
+async function logStagehandMetrics(
+  stagehand: Stagehand | undefined,
+  key: StagehandUsageKey,
+): Promise<void> {
+  if (!stagehand) return;
+
+
+  const rawMetrics: any = (stagehand as any).metrics;
+  const metrics =
+    rawMetrics && typeof rawMetrics.then === "function"
+      ? await rawMetrics
+      : rawMetrics;
+
+  if (!metrics) return;
+
+  // Keep this as a structured JSON line so it’s easy to grep/pipe elsewhere.
+
+  console.log(
+    JSON.stringify(
+      {
+        source: "stagehand-mcp",
+        event: "stagehand_metrics",
+        ...key,
+        metrics,
+      },
+      null,
+      2,
+    ),
+  );
+}
+
+export async function recordStagehandCall(
+  args: StagehandUsageKey & { stagehand?: Stagehand },
+): Promise<void> {
+  const { sessionId, toolName, operation, stagehand } = args;
+
+  // Update global aggregate
+  const globalStats = getOrCreateOperationStats(globalUsage, operation);
+  globalStats.callCount += 1;
+  globalStats.toolCallCounts[toolName] =
+    (globalStats.toolCallCounts[toolName] ?? 0) + 1;
+
+  // Update per-session usage
+  if (!perSessionUsage[sessionId]) {
+    perSessionUsage[sessionId] = { operations: {} };
+  }
+
+  const sessionStats = getOrCreateOperationStats(
+    perSessionUsage[sessionId].operations,
+    operation,
+  );
+  sessionStats.callCount += 1;
+  sessionStats.toolCallCounts[toolName] =
+    (sessionStats.toolCallCounts[toolName] ?? 0) + 1;
+
+  await logStagehandMetrics(stagehand, { sessionId, toolName, operation });
+}
+
+export function getUsageSnapshot(): StagehandUsageSnapshot {
+  return {
+    global: globalUsage,
+    perSession: perSessionUsage,
+  };
+}
+
+export function resetUsage(): void {
+  for (const key of Object.keys(globalUsage)) {
+
+    delete globalUsage[key];
+  }
+  for (const key of Object.keys(perSessionUsage)) {
+
+    delete perSessionUsage[key];
+  }
+}

src/tools/act.ts

@@ -2,6 +2,7 @@ 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";
+import { recordStagehandCall } from "../mcp/usage.js";
 
 /**
  * Stagehand Act
@@ -45,6 +46,13 @@ async function handleAct(
         variables: params.variables,
       });
 
+      await recordStagehandCall({
+        sessionId: context.currentSessionId,
+        toolName: actSchema.name,
+        operation: "act",
+        stagehand,
+      });
+
       return {
         content: [
           {

src/tools/agent.ts

@@ -2,6 +2,7 @@ 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";
+import { recordStagehandCall } from "../mcp/usage.js";
 
 /**
  * Stagehand Agent
@@ -54,6 +55,13 @@ async function handleAgent(
         maxSteps: 20,
       });
 
+      await recordStagehandCall({
+        sessionId: context.currentSessionId,
+        toolName: agentSchema.name,
+        operation: "agent.execute",
+        stagehand,
+      });
+
       return {
         content: [
           {

src/tools/extract.ts

@@ -2,6 +2,7 @@ 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";
+import { recordStagehandCall } from "../mcp/usage.js";
 
 /**
  * Stagehand Extract
@@ -39,6 +40,13 @@ async function handleExtract(
 
       const extraction = await stagehand.extract(params.instruction);
 
+      await recordStagehandCall({
+        sessionId: context.currentSessionId,
+        toolName: extractSchema.name,
+        operation: "extract",
+        stagehand,
+      });
+
       return {
         content: [
           {

src/tools/index.ts

@@ -6,6 +6,7 @@ import screenshotTool from "./screenshot.js";
 import sessionTools from "./session.js";
 import getUrlTool from "./url.js";
 import agentTool from "./agent.js";
+import usageTool from "./usage.js";
 
 // Export individual tools
 export { default as navigateTool } from "./navigate.js";
@@ -16,6 +17,7 @@ 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 { default as usageTool } from "./usage.js";
 
 // Export all tools as array
 export const TOOLS = [
@@ -27,6 +29,7 @@ export const TOOLS = [
   screenshotTool,
   getUrlTool,
   agentTool,
+  usageTool,
 ];
 
 export const sessionManagementTools = sessionTools;

src/tools/navigate.ts

@@ -2,6 +2,7 @@ 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";
+import { recordStagehandCall } from "../mcp/usage.js";
 
 const NavigateInputSchema = z.object({
   url: z.string().describe("The URL to navigate to"),
@@ -37,6 +38,13 @@ async function handleNavigate(
         throw new Error("No Browserbase session ID available");
       }
 
+      await recordStagehandCall({
+        sessionId: context.currentSessionId,
+        toolName: navigateSchema.name,
+        operation: "navigate.goto",
+        stagehand,
+      });
+
       return {
         content: [
           {

src/tools/observe.ts

@@ -2,6 +2,7 @@ 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";
+import { recordStagehandCall } from "../mcp/usage.js";
 
 /**
  * Stagehand Observe
@@ -42,6 +43,13 @@ async function handleObserve(
 
       const observations = await stagehand.observe(params.instruction);
 
+      await recordStagehandCall({
+        sessionId: context.currentSessionId,
+        toolName: observeSchema.name,
+        operation: "observe",
+        stagehand,
+      });
+
       return {
         content: [
           {

src/tools/url.ts

@@ -2,6 +2,7 @@ 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";
+import { recordStagehandCall } from "../mcp/usage.js";
 
 /**
  * Stagehand Get URL
@@ -37,6 +38,13 @@ async function handleGetUrl(
 
       const currentUrl = page.url();
 
+      await recordStagehandCall({
+        sessionId: context.currentSessionId,
+        toolName: getUrlSchema.name,
+        operation: "get_url",
+        stagehand,
+      });
+
       return {
         content: [
           {

src/tools/usage.ts

@@ -0,0 +1,102 @@
+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";
+import { getUsageSnapshot, resetUsage } from "../mcp/usage.js";
+
+const UsageInputSchema = z
+  .object({
+    sessionId: z
+      .string()
+      .optional()
+      .describe(
+        "Optional: filter per-session stats to a specific internal MCP session ID.",
+      ),
+    scope: z
+      .enum(["global", "perSession", "all"])
+      .optional()
+      .describe(
+        'Optional: which portion of the snapshot to return: "global", "perSession", or "all" (default).',
+      ),
+    reset: z
+      .boolean()
+      .optional()
+      .describe(
+        "Optional: when true, reset accumulated usage counters after returning the snapshot.",
+      ),
+  })
+  .optional()
+  .default({});
+
+type UsageInput = z.infer<typeof UsageInputSchema>;
+
+const usageSchema: ToolSchema<typeof UsageInputSchema> = {
+  name: "browserbase_usage_stats",
+  description:
+    "Return a snapshot of Stagehand usage metrics (call counts) for this MCP process, optionally filtered by session.",
+  inputSchema: UsageInputSchema,
+};
+
+async function handleUsage(
+
+  context: Context,
+  params: UsageInput,
+): Promise<ToolResult> {
+  const action = async (): Promise<ToolActionResult> => {
+    const snapshot = getUsageSnapshot();
+
+    const scope = params.scope ?? "all";
+    let result: unknown = snapshot;
+
+    if (scope === "global") {
+      result = { global: snapshot.global };
+    } else if (scope === "perSession") {
+      if (params.sessionId) {
+        result = {
+          perSession: {
+            [params.sessionId]: snapshot.perSession[params.sessionId] ?? {
+              operations: {},
+            },
+          },
+        };
+      } else {
+        result = { perSession: snapshot.perSession };
+      }
+    } else if (scope === "all" && params.sessionId) {
+      result = {
+        global: snapshot.global,
+        perSession: {
+          [params.sessionId]: snapshot.perSession[params.sessionId] ?? {
+            operations: {},
+          },
+        },
+      };
+    }
+
+    if (params.reset) {
+      resetUsage();
+    }
+
+    return {
+      content: [
+        {
+          type: "text",
+          text: JSON.stringify(result, null, 2),
+        },
+      ],
+    };
+  };
+
+  return {
+    action,
+    waitForNetwork: false,
+  };
+}
+
+const usageTool: Tool<typeof UsageInputSchema> = {
+  capability: "core",
+  schema: usageSchema,
+  handle: handleUsage,
+};
+
+export default usageTool;