Fix stagehand.metrics

Author: miguelg719

.changeset/sharp-suns-fall.md

@@ -0,0 +1,5 @@
+---
+"@browserbasehq/stagehand": minor
+---
+
+Update stagehand.metrics to work on every environment

lib/api.ts

@@ -17,6 +17,7 @@ import {
   ExtractResult,
   ObserveOptions,
   ObserveResult,
+  StagehandMetrics,
 } from "../types/stagehand";
 import { AgentExecuteOptions, AgentResult } from ".";
 import {
@@ -165,6 +166,99 @@ export class StagehandAPI {
     return response;
   }
 
+  async getReplayMetrics(): Promise<StagehandMetrics> {
+    if (!this.sessionId) {
+      throw new Error("sessionId is required to fetch metrics.");
+    }
+
+    const response = await this.request(`/sessions/${this.sessionId}/replay`, {
+      method: "GET",
+    });
+
+    if (response.status !== 200) {
+      const errorText = await response.text();
+      this.logger({
+        category: "api",
+        message: `[HTTP ERROR] Failed to fetch metrics. Status ${response.status}: ${errorText}`,
+        level: 1,
+      });
+      throw new Error(
+        `Failed to fetch metrics with status ${response.status}: ${errorText}`,
+      );
+    }
+
+    const data = await response.json();
+
+    if (!data.success) {
+      throw new Error(
+        `Failed to fetch metrics: ${data.error || "Unknown error"}`,
+      );
+    }
+
+    // Parse the API data into StagehandMetrics format
+    const apiData = data.data || {};
+    const metrics: StagehandMetrics = {
+      actPromptTokens: 0,
+      actCompletionTokens: 0,
+      actInferenceTimeMs: 0,
+      extractPromptTokens: 0,
+      extractCompletionTokens: 0,
+      extractInferenceTimeMs: 0,
+      observePromptTokens: 0,
+      observeCompletionTokens: 0,
+      observeInferenceTimeMs: 0,
+      agentPromptTokens: 0,
+      agentCompletionTokens: 0,
+      agentInferenceTimeMs: 0,
+      totalPromptTokens: 0,
+      totalCompletionTokens: 0,
+      totalInferenceTimeMs: 0,
+    };
+
+    // Parse pages and their actions
+    const pages = apiData.pages || [];
+    for (const page of pages) {
+      const actions = page.actions || [];
+      for (const action of actions) {
+        // Get method name and token usage
+        const method = (action.method || "").toLowerCase();
+        const tokenUsage = action.tokenUsage || {};
+
+        if (tokenUsage) {
+          const inputTokens = tokenUsage.inputTokens || 0;
+          const outputTokens = tokenUsage.outputTokens || 0;
+          const timeMs = tokenUsage.timeMs || 0;
+
+          // Map method to metrics fields
+          if (method === "act") {
+            metrics.actPromptTokens += inputTokens;
+            metrics.actCompletionTokens += outputTokens;
+            metrics.actInferenceTimeMs += timeMs;
+          } else if (method === "extract") {
+            metrics.extractPromptTokens += inputTokens;
+            metrics.extractCompletionTokens += outputTokens;
+            metrics.extractInferenceTimeMs += timeMs;
+          } else if (method === "observe") {
+            metrics.observePromptTokens += inputTokens;
+            metrics.observeCompletionTokens += outputTokens;
+            metrics.observeInferenceTimeMs += timeMs;
+          } else if (method === "agent") {
+            metrics.agentPromptTokens += inputTokens;
+            metrics.agentCompletionTokens += outputTokens;
+            metrics.agentInferenceTimeMs += timeMs;
+          }
+
+          // Always update totals for any method with token usage
+          metrics.totalPromptTokens += inputTokens;
+          metrics.totalCompletionTokens += outputTokens;
+          metrics.totalInferenceTimeMs += timeMs;
+        }
+      }
+    }
+
+    return metrics;
+  }
+
   private async execute<T>({
     method,
     args,

lib/index.ts

@@ -458,8 +458,21 @@ export class Stagehand {
     totalInferenceTimeMs: 0,
   };
 
-  public get metrics(): StagehandMetrics {
-    return this.stagehandMetrics;
+  public get metrics(): Promise<StagehandMetrics> {
+    if (this.usingAPI && this.apiClient) {
+      // Fetch metrics from the API
+      return this.apiClient.getReplayMetrics().catch((error) => {
+        this.logger({
+          category: "metrics",
+          message: `Failed to fetch metrics from API: ${error}`,
+          level: 0,
+        });
+        // Fall back to local metrics on error
+        return this.stagehandMetrics;
+      });
+    }
+    // Return local metrics wrapped in a Promise for consistency
+    return Promise.resolve(this.stagehandMetrics);
   }
 
   public get isClosed(): boolean {