chore(memory): expose a tool for getting the inital data. (#1909)

This refactors the code to extract the Id logic from the PageCollector
and provide it in the heapsnapshot.
We need to use an UID as we need the internal ClassKey to query the
heapsnapshot, but that is a strange string (usually looking like
`,ClassName`) which may get the LLM confused as we use comma separated
output.

Author: Lightning00Blade

README.md

@@ -497,11 +497,10 @@ If you run into any issues, checkout our [troubleshooting guide](./docs/troubles
 - **Emulation** (2 tools)
   - [`emulate`](docs/tool-reference.md#emulate)
   - [`resize_page`](docs/tool-reference.md#resize_page)
-- **Performance** (4 tools)
+- **Performance** (3 tools)
   - [`performance_analyze_insight`](docs/tool-reference.md#performance_analyze_insight)
   - [`performance_start_trace`](docs/tool-reference.md#performance_start_trace)
   - [`performance_stop_trace`](docs/tool-reference.md#performance_stop_trace)
-  - [`take_memory_snapshot`](docs/tool-reference.md#take_memory_snapshot)
 - **Network** (2 tools)
   - [`get_network_request`](docs/tool-reference.md#get_network_request)
   - [`list_network_requests`](docs/tool-reference.md#list_network_requests)
@@ -512,6 +511,8 @@ If you run into any issues, checkout our [troubleshooting guide](./docs/troubles
   - [`list_console_messages`](docs/tool-reference.md#list_console_messages)
   - [`take_screenshot`](docs/tool-reference.md#take_screenshot)
   - [`take_snapshot`](docs/tool-reference.md#take_snapshot)
+- **Memory** (1 tools)
+  - [`take_memory_snapshot`](docs/tool-reference.md#take_memory_snapshot)
 
 <!-- END AUTO GENERATED TOOLS -->
 

docs/tool-reference.md

@@ -22,11 +22,10 @@
 - **[Emulation](#emulation)** (2 tools)
   - [`emulate`](#emulate)
   - [`resize_page`](#resize_page)
-- **[Performance](#performance)** (4 tools)
+- **[Performance](#performance)** (3 tools)
   - [`performance_analyze_insight`](#performance_analyze_insight)
   - [`performance_start_trace`](#performance_start_trace)
   - [`performance_stop_trace`](#performance_stop_trace)
-  - [`take_memory_snapshot`](#take_memory_snapshot)
 - **[Network](#network)** (2 tools)
   - [`get_network_request`](#get_network_request)
   - [`list_network_requests`](#list_network_requests)
@@ -37,6 +36,8 @@
   - [`list_console_messages`](#list_console_messages)
   - [`take_screenshot`](#take_screenshot)
   - [`take_snapshot`](#take_snapshot)
+- **[Memory](#memory)** (1 tools)
+  - [`take_memory_snapshot`](#take_memory_snapshot)
 
 ## Input automation
 
@@ -276,16 +277,6 @@
 
 ---
 
-### `take_memory_snapshot`
-
-**Description:** Capture a heap snapshot of the currently selected page. Use to analyze the memory distribution of JavaScript objects and debug memory leaks.
-
-**Parameters:**
-
-- **filePath** (string) **(required)**: A path to a .heapsnapshot file to save the heapsnapshot to.
-
----
-
 ## Network
 
 ### `get_network_request`
@@ -398,3 +389,15 @@ in the DevTools Elements panel (if any).
 - **verbose** (boolean) _(optional)_: Whether to include all possible information available in the full a11y tree. Default is false.
 
 ---
+
+## Memory
+
+### `take_memory_snapshot`
+
+**Description:** Capture a heap snapshot of the currently selected page. Use to analyze the memory distribution of JavaScript objects and debug memory leaks.
+
+**Parameters:**
+
+- **filePath** (string) **(required)**: A path to a .heapsnapshot file to save the heapsnapshot to.
+
+---

src/HeapSnapshotManager.ts

@@ -8,13 +8,25 @@ import fsSync from 'node:fs';
 import path from 'node:path';
 
 import {DevTools} from './third_party/index.js';
+import {
+  createIdGenerator,
+  stableIdSymbol,
+  type WithSymbolId,
+} from './utils/id.js';
+
+export type AggregatedInfoWithUid =
+  WithSymbolId<DevTools.HeapSnapshotModel.HeapSnapshotModel.AggregatedInfo>;
 
 export class HeapSnapshotManager {
   #snapshots = new Map<
     string,
     {
       snapshot: DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotProxy;
       worker: DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotWorkerProxy;
+      // TODO: use a multimap
+      uidToClassKey: Map<number, string>;
+      classKeyToUid: Map<string, number>;
+      idGenerator: () => number;
     }
   >();
 
@@ -28,20 +40,35 @@ export class HeapSnapshotManager {
     }
 
     const {snapshot, worker} = await this.#loadSnapshot(absolutePath);
-    this.#snapshots.set(absolutePath, {snapshot, worker});
+    this.#snapshots.set(absolutePath, {
+      snapshot,
+      worker,
+      uidToClassKey: new Map<number, string>(),
+      classKeyToUid: new Map<string, number>(),
+      idGenerator: createIdGenerator(),
+    });
 
     return snapshot;
   }
 
   async getAggregates(
     filePath: string,
-  ): Promise<
-    Record<string, DevTools.HeapSnapshotModel.HeapSnapshotModel.AggregatedInfo>
-  > {
+  ): Promise<Record<string, AggregatedInfoWithUid>> {
     const snapshot = await this.getSnapshot(filePath);
     const filter =
       new DevTools.HeapSnapshotModel.HeapSnapshotModel.NodeFilter();
-    return await snapshot.aggregatesWithFilter(filter);
+    const aggregates: Record<string, AggregatedInfoWithUid> =
+      await snapshot.aggregatesWithFilter(filter);
+
+    for (const key of Object.keys(aggregates)) {
+      const uid = await this.getOrCreateUidForClassKey(filePath, key);
+      const aggregate = aggregates[key];
+      if (aggregate) {
+        aggregate[stableIdSymbol] = uid;
+      }
+    }
+
+    return aggregates;
   }
 
   async getStats(
@@ -58,6 +85,29 @@ export class HeapSnapshotManager {
     return snapshot.staticData;
   }
 
+  async getOrCreateUidForClassKey(
+    filePath: string,
+    classKey: string,
+  ): Promise<number> {
+    const cached = this.#getCachedSnapshot(filePath);
+    let uid = cached.classKeyToUid.get(classKey);
+    if (!uid) {
+      uid = cached.idGenerator();
+      cached.classKeyToUid.set(classKey, uid);
+      cached.uidToClassKey.set(uid, classKey);
+    }
+    return uid;
+  }
+
+  #getCachedSnapshot(filePath: string) {
+    const absolutePath = path.resolve(filePath);
+    const cached = this.#snapshots.get(absolutePath);
+    if (!cached) {
+      throw new Error(`Snapshot not loaded for ${filePath}`);
+    }
+    return cached;
+  }
+
   async #loadSnapshot(absolutePath: string): Promise<{
     snapshot: DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotProxy;
     worker: DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotWorkerProxy;

src/McpContext.ts

@@ -10,6 +10,7 @@ import path from 'node:path';
 import type {TargetUniverse} from './DevtoolsUtils.js';
 import {UniverseManager} from './DevtoolsUtils.js';
 import {HeapSnapshotManager} from './HeapSnapshotManager.js';
+import type {AggregatedInfoWithUid} from './HeapSnapshotManager.js';
 import {McpPage} from './McpPage.js';
 import {
   NetworkCollector,
@@ -906,9 +907,7 @@ export class McpContext implements Context {
 
   async getHeapSnapshotAggregates(
     filePath: string,
-  ): Promise<
-    Record<string, DevTools.HeapSnapshotModel.HeapSnapshotModel.AggregatedInfo>
-  > {
+  ): Promise<Record<string, AggregatedInfoWithUid>> {
     return await this.#heapSnapshotManager.getAggregates(filePath);
   }
 

src/PageCollector.ts

@@ -22,6 +22,11 @@ import {
   type Page,
   type PageEvents as PuppeteerPageEvents,
 } from './third_party/index.js';
+import {
+  createIdGenerator,
+  stableIdSymbol,
+  type WithSymbolId,
+} from './utils/id.js';
 
 export class UncaughtError {
   readonly details: Protocol.Runtime.ExceptionDetails;
@@ -42,21 +47,6 @@ export type ListenerMap<EventMap extends PageEvents = PageEvents> = {
   [K in keyof EventMap]?: (event: EventMap[K]) => void;
 };
 
-function createIdGenerator() {
-  let i = 1;
-  return () => {
-    if (i === Number.MAX_SAFE_INTEGER) {
-      i = 0;
-    }
-    return i++;
-  };
-}
-
-export const stableIdSymbol = Symbol('stableIdSymbol');
-type WithSymbolId<T> = T & {
-  [stableIdSymbol]?: number;
-};
-
 export class PageCollector<T> {
   #browser: Browser;
   #listenersInitializer: (

src/bin/cliDefinitions.ts

@@ -590,7 +590,7 @@ export const commands: Commands = {
   take_memory_snapshot: {
     description:
       'Capture a heap snapshot of the currently selected page. Use to analyze the memory distribution of JavaScript objects and debug memory leaks.',
-    category: 'Performance',
+    category: 'Memory',
     args: {
       filePath: {
         name: 'filePath',

src/formatters/HeapSnapshotFormatter.ts

@@ -4,41 +4,39 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import type {AggregatedInfoWithUid} from '../HeapSnapshotManager.js';
 import type {DevTools} from '../third_party/index.js';
+import {stableIdSymbol} from '../utils/id.js';
 
 export interface FormattedSnapshotEntry {
   className: string;
+  classUid?: number;
   count: number;
   selfSize: number;
   retainedSize: number;
 }
 
 export class HeapSnapshotFormatter {
-  #aggregates: Record<
-    string,
-    DevTools.HeapSnapshotModel.HeapSnapshotModel.AggregatedInfo
-  >;
+  #aggregates: Record<string, AggregatedInfoWithUid>;
 
-  constructor(
-    aggregates: Record<
-      string,
-      DevTools.HeapSnapshotModel.HeapSnapshotModel.AggregatedInfo
-    >,
-  ) {
+  constructor(aggregates: Record<string, AggregatedInfoWithUid>) {
     this.#aggregates = aggregates;
   }
 
-  #getSortedAggregates(): DevTools.HeapSnapshotModel.HeapSnapshotModel.AggregatedInfo[] {
+  #getSortedAggregates(): AggregatedInfoWithUid[] {
     return Object.values(this.#aggregates).sort((a, b) => b.self - a.self);
   }
 
   toString(): string {
     const sorted = this.#getSortedAggregates();
     const lines: string[] = [];
-    lines.push('className,count,selfSize,maxRetainedSize');
+    lines.push('uid,className,count,selfSize,maxRetainedSize');
 
     for (const info of sorted) {
-      lines.push(`"${info.name}",${info.count},${info.self},${info.maxRet}`);
+      const uid = info[stableIdSymbol] ?? '';
+      lines.push(
+        `${uid},"${info.name}",${info.count},${info.self},${info.maxRet}`,
+      );
     }
 
     return lines.join('\n');
@@ -47,6 +45,7 @@ export class HeapSnapshotFormatter {
   toJSON(): FormattedSnapshotEntry[] {
     const sorted = this.#getSortedAggregates();
     return sorted.map(info => ({
+      uid: info[stableIdSymbol],
       className: info.name,
       count: info.count,
       selfSize: info.self,

src/telemetry/tool_call_metrics.json

@@ -569,5 +569,22 @@
         "argType": "number"
       }
     ]
+  },
+  {
+    "name": "get_memory_snapshot_details",
+    "args": [
+      {
+        "name": "file_path_length",
+        "argType": "number"
+      },
+      {
+        "name": "page_idx",
+        "argType": "number"
+      },
+      {
+        "name": "page_size",
+        "argType": "number"
+      }
+    ]
   }
 ]

src/tools/ToolDefinition.ts

@@ -5,6 +5,7 @@
  */
 
 import type {ParsedArguments} from '../bin/chrome-devtools-mcp-cli-options.js';
+import type {AggregatedInfoWithUid} from '../HeapSnapshotManager.js';
 import type {McpPage} from '../McpPage.js';
 import {zod} from '../third_party/index.js';
 import type {
@@ -227,9 +228,7 @@ export type Context = Readonly<{
   ): string | undefined;
   getHeapSnapshotAggregates(
     filePath: string,
-  ): Promise<
-    Record<string, DevTools.HeapSnapshotModel.HeapSnapshotModel.AggregatedInfo>
-  >;
+  ): Promise<Record<string, AggregatedInfoWithUid>>;
   getHeapSnapshotStats(
     filePath: string,
   ): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.Statistics>;

src/tools/memory.ts

@@ -14,7 +14,7 @@ export const takeMemorySnapshot = definePageTool({
   name: 'take_memory_snapshot',
   description: `Capture a heap snapshot of the currently selected page. Use to analyze the memory distribution of JavaScript objects and debug memory leaks.`,
   annotations: {
-    category: ToolCategory.PERFORMANCE,
+    category: ToolCategory.MEMORY,
     readOnlyHint: false,
   },
   schema: {
@@ -38,7 +38,7 @@ export const takeMemorySnapshot = definePageTool({
 export const exploreMemorySnapshot = defineTool({
   name: 'load_memory_snapshot',
   description:
-    'Loads a memory heapsnapshot and returns snapshot summary stats.  ',
+    'Loads a memory heapsnapshot and returns snapshot summary stats.',
   annotations: {
     category: ToolCategory.MEMORY,
     readOnlyHint: true,
@@ -56,3 +56,35 @@ export const exploreMemorySnapshot = defineTool({
     response.setHeapSnapshotStats(stats, staticData);
   },
 });
+
+export const getMemorySnapshotDetails = defineTool({
+  name: 'get_memory_snapshot_details',
+  description:
+    'Loads a memory heapsnapshot and returns all available information including statistics, static data, and aggregated node information. Supports pagination for aggregates.',
+  annotations: {
+    category: ToolCategory.MEMORY,
+    readOnlyHint: true,
+    conditions: ['experimentalMemory'],
+  },
+  schema: {
+    filePath: zod.string().describe('A path to a .heapsnapshot file to read.'),
+    pageIdx: zod
+      .number()
+      .optional()
+      .describe('The page index for pagination of aggregates.'),
+    pageSize: zod
+      .number()
+      .optional()
+      .describe('The page size for pagination of aggregates.'),
+  },
+  handler: async (request, response, context) => {
+    const aggregates = await context.getHeapSnapshotAggregates(
+      request.params.filePath,
+    );
+
+    response.setHeapSnapshotAggregates(aggregates, {
+      pageIdx: request.params.pageIdx,
+      pageSize: request.params.pageSize,
+    });
+  },
+});