chore: use id generation for memory aggragates

Author: Lightning00Blade

src/HeapSnapshotManager.ts

@@ -8,13 +8,24 @@ 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;
+      uidToClassKey: Map<number, string>;
+      classKeyToUid: Map<string, number>;
+      idGenerator: () => number;
     }
   >();
 
@@ -28,20 +39,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(

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,
@@ -918,9 +919,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/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/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,

src/utils/id.ts

@@ -0,0 +1,21 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export function createIdGenerator() {
+  let i = 1;
+  return () => {
+    if (i === Number.MAX_SAFE_INTEGER) {
+      i = 0;
+    }
+    return i++;
+  };
+}
+
+export const stableIdSymbol = Symbol('stableIdSymbol');
+
+export type WithSymbolId<T> = T & {
+  [stableIdSymbol]?: number;
+};

tests/utils.ts

@@ -24,8 +24,8 @@ import sinon from 'sinon';
 import type {ParsedArguments} from '../src/bin/chrome-devtools-mcp-cli-options.js';
 import {McpContext} from '../src/McpContext.js';
 import {McpResponse} from '../src/McpResponse.js';
-import {stableIdSymbol} from '../src/PageCollector.js';
 import {DevTools} from '../src/third_party/index.js';
+import {stableIdSymbol} from '../src/utils/id.js';
 
 export function getTextContent(
   content: CallToolResult['content'][number],