diff --git a/src/mcp/usage.ts b/src/mcp/usage.ts new file mode 100644 index 0000000..9b2b9bf --- /dev/null +++ b/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; +}; + +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]; +} + +async function logStagehandMetrics( + stagehand: Stagehand | undefined, + key: StagehandUsageKey, +): Promise { + 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 { + 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]; + } +} diff --git a/src/tools/act.ts b/src/tools/act.ts index d395a66..2d22102 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,13 @@ async function handleAct( variables: params.variables, }); + await recordStagehandCall({ + sessionId: context.currentSessionId, + toolName: actSchema.name, + operation: "act", + stagehand, + }); + return { content: [ { diff --git a/src/tools/agent.ts b/src/tools/agent.ts index e333079..2a62313 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,13 @@ async function handleAgent( maxSteps: 20, }); + await recordStagehandCall({ + sessionId: context.currentSessionId, + toolName: agentSchema.name, + operation: "agent.execute", + stagehand, + }); + return { content: [ { diff --git a/src/tools/extract.ts b/src/tools/extract.ts index 7bfb56b..2c37942 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,13 @@ async function handleExtract( const extraction = await stagehand.extract(params.instruction); + await recordStagehandCall({ + sessionId: context.currentSessionId, + toolName: extractSchema.name, + operation: "extract", + stagehand, + }); + 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..11ece5f 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,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: [ { diff --git a/src/tools/observe.ts b/src/tools/observe.ts index b473300..9905b73 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,13 @@ async function handleObserve( const observations = await stagehand.observe(params.instruction); + await recordStagehandCall({ + sessionId: context.currentSessionId, + toolName: observeSchema.name, + operation: "observe", + stagehand, + }); + return { content: [ { diff --git a/src/tools/url.ts b/src/tools/url.ts index da7e124..f73a296 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,13 @@ async function handleGetUrl( const currentUrl = page.url(); + await recordStagehandCall({ + sessionId: context.currentSessionId, + toolName: getUrlSchema.name, + operation: "get_url", + stagehand, + }); + 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;