feat(test): refactor the memory usage test to use metrics from CLI process instead of test runner (#25708)

2026-04-23 03:24:42 -07:00 · 2026-04-21 10:06:22 -07:00
parent 2c14954010
commit aee2cde1a3
5 changed files with 284 additions and 223 deletions
@@ -1475,7 +1475,7 @@ export class TestRig {
  readMetric(metricName: string): TelemetryMetric | null {
    const logs = this._readAndParseTelemetryLog();
    for (const logData of logs) {
-      if (logData.scopeMetrics) {
+      if (logData && logData.scopeMetrics) {
        for (const scopeMetric of logData.scopeMetrics) {
          for (const metric of scopeMetric.metrics) {
            if (metric.descriptor.name === `gemini_cli.${metricName}`) {
@@ -1488,6 +1488,133 @@ export class TestRig {
    return null;
  }

+  readMemoryMetrics(strategy: 'peak' | 'last' = 'peak'): {
+    timestamp: number;
+    heapUsed: number;
+    heapTotal: number;
+    rss: number;
+    external: number;
+  } {
+    const snapshots = this._getMemorySnapshots();
+    if (snapshots.length === 0) {
+      return {
+        timestamp: Date.now(),
+        heapUsed: 0,
+        heapTotal: 0,
+        rss: 0,
+        external: 0,
+      };
+    }
+
+    if (strategy === 'last') {
+      const last = snapshots[snapshots.length - 1];
+      return {
+        timestamp: last.timestamp,
+        heapUsed: last.heapUsed,
+        heapTotal: last.heapTotal,
+        rss: last.rss,
+        external: last.external,
+      };
+    }
+
+    // Find the snapshot with the highest RSS
+    let peak = snapshots[0];
+    for (const snapshot of snapshots) {
+      if (snapshot.rss > peak.rss) {
+        peak = snapshot;
+      }
+    }
+
+    // Fallback: if we didn't find any RSS but found heap, use the max heap
+    if (peak.rss === 0) {
+      for (const snapshot of snapshots) {
+        if (snapshot.heapUsed > peak.heapUsed) {
+          peak = snapshot;
+        }
+      }
+    }
+
+    return {
+      timestamp: peak.timestamp,
+      heapUsed: peak.heapUsed,
+      heapTotal: peak.heapTotal,
+      rss: peak.rss,
+      external: peak.external,
+    };
+  }
+
+  readAllMemorySnapshots(): {
+    timestamp: number;
+    heapUsed: number;
+    heapTotal: number;
+    rss: number;
+    external: number;
+  }[] {
+    return this._getMemorySnapshots();
+  }
+
+  private _getMemorySnapshots(): {
+    timestamp: number;
+    heapUsed: number;
+    heapTotal: number;
+    rss: number;
+    external: number;
+  }[] {
+    const snapshots: Record<
+      string,
+      {
+        timestamp: number;
+        heapUsed: number;
+        heapTotal: number;
+        rss: number;
+        external: number;
+      }
+    > = {};
+
+    const logs = this._readAndParseTelemetryLog();
+    for (const logData of logs) {
+      if (logData && logData.scopeMetrics) {
+        for (const scopeMetric of logData.scopeMetrics) {
+          for (const metric of scopeMetric.metrics) {
+            if (metric.descriptor.name === 'gemini_cli.memory.usage') {
+              for (const dp of metric.dataPoints) {
+                const sessionId =
+                  (dp.attributes?.['session.id'] as string) || 'unknown';
+                const component =
+                  (dp.attributes?.['component'] as string) || 'unknown';
+                const seconds = dp.startTime?.[0] || 0;
+                const nanos = dp.startTime?.[1] || 0;
+                const timeKey = `${sessionId}-${component}-${seconds}-${nanos}`;
+
+                if (!snapshots[timeKey]) {
+                  snapshots[timeKey] = {
+                    timestamp: seconds * 1000 + Math.floor(nanos / 1000000),
+                    rss: 0,
+                    heapUsed: 0,
+                    heapTotal: 0,
+                    external: 0,
+                  };
+                }
+
+                const type = dp.attributes?.['memory_type'];
+                const value = dp.value?.max ?? dp.value?.sum ?? 0;
+
+                if (type === 'heap_used') snapshots[timeKey].heapUsed = value;
+                else if (type === 'heap_total')
+                  snapshots[timeKey].heapTotal = value;
+                else if (type === 'rss') snapshots[timeKey].rss = value;
+                else if (type === 'external')
+                  snapshots[timeKey].external = value;
+              }
+            }
+          }
+        }
+      }
+    }
+
+    return Object.values(snapshots).sort((a, b) => a.timestamp - b.timestamp);
+  }
+
  async runInteractive(options?: {
    args?: string | string[];
    approvalMode?: 'default' | 'auto_edit' | 'yolo' | 'plan';