From c0017593f30c7a3bbd928db5124d5fe578814092 Mon Sep 17 00:00:00 2001 From: Shrey Pandya Date: Fri, 14 Nov 2025 09:06:08 -0800 Subject: [PATCH] add usage tool --- README.md | 52 ++++++++++++++++++++ src/mcp/usage.ts | 80 +++++++++++++++++++++++++++++++ src/tools/act.ts | 7 +++ src/tools/agent.ts | 7 +++ src/tools/extract.ts | 7 +++ src/tools/index.ts | 3 ++ src/tools/navigate.ts | 7 +++ src/tools/observe.ts | 7 +++ src/tools/screenshot.ts | 7 +++ src/tools/url.ts | 7 +++ src/tools/usage.ts | 102 ++++++++++++++++++++++++++++++++++++++++ 11 files changed, 286 insertions(+) create mode 100644 src/mcp/usage.ts create mode 100644 src/tools/usage.ts diff --git a/README.md b/README.md index 640def0..5f4764b 100644 --- a/README.md +++ b/README.md @@ -420,6 +420,58 @@ The server provides access to screenshot resources: 1. **Screenshots** (`screenshot://`) - PNG images of captured screenshots +### Stagehand usage metrics (experimental) + +This server exposes an internal **Stagehand usage** tool that you can call via MCP to understand how often Stagehand-backed operations run inside the Browserbase MCP process. + +- **Tool name**: `browserbase_usage_stats` +- **Capability**: `core` +- **Input schema**: + - `sessionId?: string` – optional internal MCP session ID to filter per-session stats + - `scope?: "global" | "perSession" | "all"` – which part of the snapshot to return (defaults to `"all"`) + - `reset?: boolean` – when `true`, clears the in-memory counters after returning the snapshot + +Example `call_tool` request (conceptual): + +```json +{ + "name": "browserbase_usage_stats", + "arguments": { + "scope": "all", + "reset": false + } +} +``` + +Example response `content` (truncated): + +```json +{ + "global": { + "agent.execute": { + "callCount": 3, + "toolCallCounts": { + "browserbase_stagehand_agent": 3 + } + } + }, + "perSession": { + "browserbase_session_default_...": { + "operations": { + "agent.execute": { + "callCount": 2, + "toolCallCounts": { + "browserbase_stagehand_agent": 2 + } + } + } + } + } +} +``` + +You can have a Claude Agent SDK-based agent invoke this tool at the end of a run to fetch Stagehand usage for that run, then optionally call it again with `"reset": true` to clear counters for the next run. + ## Key Features - **AI-Powered Automation**: Natural language commands for web interactions diff --git a/src/mcp/usage.ts b/src/mcp/usage.ts new file mode 100644 index 0000000..7ce9393 --- /dev/null +++ b/src/mcp/usage.ts @@ -0,0 +1,80 @@ +export type StagehandUsageOperation = string; + +export type StagehandUsageKey = { + sessionId: string; + toolName: string; + operation: StagehandUsageOperation; +}; + +export type StagehandOperationStats = { + callCount: number; + toolCallCounts: Record; +}; + +export type StagehandSessionUsage = { + operations: Record; +}; + +export type StagehandUsageSnapshot = { + global: Record; + perSession: Record; +}; + +const globalUsage: Record = {}; +const perSessionUsage: Record = {}; + +function getOrCreateOperationStats( + container: Record, + operation: StagehandUsageOperation, +): StagehandOperationStats { + if (!container[operation]) { + container[operation] = { + callCount: 0, + toolCallCounts: {}, + }; + } + return container[operation]; +} + +export function recordStagehandCall({ + sessionId, + toolName, + operation, +}: StagehandUsageKey): void { + // 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; +} + +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]; + } +} diff --git a/src/tools/act.ts b/src/tools/act.ts index d395a66..481be2c 100644 --- a/src/tools/act.ts +++ b/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,12 @@ async function handleAct( variables: params.variables, }); + recordStagehandCall({ + sessionId: context.currentSessionId, + toolName: actSchema.name, + operation: "act", + }); + return { content: [ { diff --git a/src/tools/agent.ts b/src/tools/agent.ts index e333079..c7622b5 100644 --- a/src/tools/agent.ts +++ b/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,12 @@ async function handleAgent( maxSteps: 20, }); + recordStagehandCall({ + sessionId: context.currentSessionId, + toolName: agentSchema.name, + operation: "agent.execute", + }); + return { content: [ { diff --git a/src/tools/extract.ts b/src/tools/extract.ts index 7bfb56b..dc0490d 100644 --- a/src/tools/extract.ts +++ b/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,12 @@ async function handleExtract( const extraction = await stagehand.extract(params.instruction); + recordStagehandCall({ + sessionId: context.currentSessionId, + toolName: extractSchema.name, + operation: "extract", + }); + return { content: [ { diff --git a/src/tools/index.ts b/src/tools/index.ts index f0a19da..5ed63f6 100644 --- a/src/tools/index.ts +++ b/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; diff --git a/src/tools/navigate.ts b/src/tools/navigate.ts index c6309a0..586baa7 100644 --- a/src/tools/navigate.ts +++ b/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,12 @@ async function handleNavigate( throw new Error("No Browserbase session ID available"); } + recordStagehandCall({ + sessionId: context.currentSessionId, + toolName: navigateSchema.name, + operation: "navigate.goto", + }); + return { content: [ { diff --git a/src/tools/observe.ts b/src/tools/observe.ts index b473300..c8ddc07 100644 --- a/src/tools/observe.ts +++ b/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,12 @@ async function handleObserve( const observations = await stagehand.observe(params.instruction); + recordStagehandCall({ + sessionId: context.currentSessionId, + toolName: observeSchema.name, + operation: "observe", + }); + return { content: [ { diff --git a/src/tools/screenshot.ts b/src/tools/screenshot.ts index 5a96bf9..62df31a 100644 --- a/src/tools/screenshot.ts +++ b/src/tools/screenshot.ts @@ -3,6 +3,7 @@ import type { Tool, ToolSchema, ToolResult } from "./tool.js"; import type { Context } from "../context.js"; import type { ToolActionResult } from "../types/types.js"; import { registerScreenshot } from "../mcp/resources.js"; +import { recordStagehandCall } from "../mcp/usage.js"; /** * Screenshot @@ -62,6 +63,12 @@ async function handleScreenshot( }); } + recordStagehandCall({ + sessionId: context.currentSessionId, + toolName: screenshotSchema.name, + operation: "screenshot", + }); + return { content: [ { diff --git a/src/tools/url.ts b/src/tools/url.ts index da7e124..dd1da8a 100644 --- a/src/tools/url.ts +++ b/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,12 @@ async function handleGetUrl( const currentUrl = page.url(); + recordStagehandCall({ + sessionId: context.currentSessionId, + toolName: getUrlSchema.name, + operation: "get_url", + }); + return { content: [ { diff --git a/src/tools/usage.ts b/src/tools/usage.ts new file mode 100644 index 0000000..0f6e37b --- /dev/null +++ b/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; + +const usageSchema: ToolSchema = { + 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 { + const action = async (): Promise => { + 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 = { + capability: "core", + schema: usageSchema, + handle: handleUsage, +}; + +export default usageTool;