chore(memory): expose way to get the nodes (#1916)

The next step for the memory debugging is to provide the nodes that the
aggregates have, from which you can get the retainers of the object.

Author: Lightning00Blade

src/HeapSnapshotManager.ts

@@ -99,6 +99,23 @@ export class HeapSnapshotManager {
     return uid;
   }
 
+  async getNodesByUid(
+    filePath: string,
+    uid: number,
+  ): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.ItemsRange> {
+    const snapshot = await this.getSnapshot(filePath);
+    const filter =
+      new DevTools.HeapSnapshotModel.HeapSnapshotModel.NodeFilter();
+    const className = await this.resolveClassKeyFromUid(filePath, uid);
+    if (!className) {
+      throw new Error(`Class with UID ${uid} not found in heap snapshot`);
+    }
+    const provider = snapshot.createNodesProviderForClass(className, filter);
+
+    const range = await provider.serializeItemsRange(0, 1);
+    return await provider.serializeItemsRange(0, range.totalLength);
+  }
+
   #getCachedSnapshot(filePath: string) {
     const absolutePath = path.resolve(filePath);
     const cached = this.#snapshots.get(absolutePath);
@@ -108,6 +125,14 @@ export class HeapSnapshotManager {
     return cached;
   }
 
+  async resolveClassKeyFromUid(
+    filePath: string,
+    uid: number,
+  ): Promise<string | undefined> {
+    const cached = this.#getCachedSnapshot(filePath);
+    return cached.uidToClassKey.get(uid);
+  }
+
   async #loadSnapshot(absolutePath: string): Promise<{
     snapshot: DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotProxy;
     worker: DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotWorkerProxy;

src/McpContext.ts

@@ -765,4 +765,11 @@ export class McpContext implements Context {
   ): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.StaticData | null> {
     return await this.#heapSnapshotManager.getStaticData(filePath);
   }
+
+  async getHeapSnapshotNodesByUid(
+    filePath: string,
+    uid: number,
+  ): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.ItemsRange> {
+    return await this.#heapSnapshotManager.getNodesByUid(filePath, uid);
+  }
 }

src/McpResponse.ts

@@ -179,6 +179,7 @@ export class McpResponse implements Response {
     pagination?: PaginationOptions;
     stats?: DevTools.HeapSnapshotModel.HeapSnapshotModel.Statistics;
     staticData?: DevTools.HeapSnapshotModel.HeapSnapshotModel.StaticData | null;
+    nodes?: DevTools.HeapSnapshotModel.HeapSnapshotModel.ItemsRange;
   };
   #networkRequestsOptions?: {
     include: boolean;
@@ -404,6 +405,18 @@ export class McpResponse implements Response {
     };
   }
 
+  setHeapSnapshotNodes(
+    nodes: DevTools.HeapSnapshotModel.HeapSnapshotModel.ItemsRange,
+    options?: PaginationOptions,
+  ) {
+    this.#heapSnapshotOptions = {
+      ...this.#heapSnapshotOptions,
+      include: true,
+      nodes,
+      pagination: options,
+    };
+  }
+
   attachImage(value: ImageContentData): void {
     this.#images.push(value);
   }
@@ -701,6 +714,7 @@ export class McpResponse implements Response {
         staticData?: object;
       };
       heapSnapshotData?: object[];
+      heapSnapshotNodes?: readonly object[];
       extensionServiceWorkers?: object[];
       extensionPages?: object[];
     } = {};
@@ -929,6 +943,20 @@ Call ${handleDialog.name} to handle it before continuing.`);
         response.push(formatter.toString());
         structuredContent.heapSnapshotData = formatter.toJSON();
       }
+      const nodes = this.#heapSnapshotOptions.nodes;
+      if (nodes) {
+        const paginationData = this.#dataWithPagination(
+          nodes.items,
+          this.#heapSnapshotOptions.pagination,
+        );
+
+        response.push(HeapSnapshotFormatter.formatNodes(paginationData.items));
+
+        structuredContent.pagination = paginationData.pagination;
+        response.push(...paginationData.info);
+
+        structuredContent.heapSnapshotNodes = paginationData.items;
+      }
     }
 
     if (data.detailedNetworkRequest) {

src/formatters/HeapSnapshotFormatter.ts

@@ -16,13 +16,44 @@ export interface FormattedSnapshotEntry {
   retainedSize: number;
 }
 
+function isNodeLike(
+  item: unknown,
+): item is DevTools.HeapSnapshotModel.HeapSnapshotModel.Node {
+  return (
+    typeof item === 'object' && item !== null && 'id' in item && 'name' in item
+  );
+}
+
 export class HeapSnapshotFormatter {
   #aggregates: Record<string, AggregatedInfoWithUid>;
 
   constructor(aggregates: Record<string, AggregatedInfoWithUid>) {
     this.#aggregates = aggregates;
   }
 
+  static formatNodes(
+    items: ReadonlyArray<
+      | DevTools.HeapSnapshotModel.HeapSnapshotModel.Node
+      | DevTools.HeapSnapshotModel.HeapSnapshotModel.Edge
+    >,
+  ): string {
+    const lines: string[] = [];
+
+    if (items.length > 0 && isNodeLike(items[0])) {
+      lines.push('id,name,type,distance,selfSize,retainedSize');
+    }
+
+    for (const item of items) {
+      if (isNodeLike(item)) {
+        lines.push(
+          `${item.id},"${item.name}",${item.type},${item.distance},${item.selfSize},${item.retainedSize}`,
+        );
+      }
+    }
+
+    return lines.join('\n');
+  }
+
   #getSortedAggregates(): AggregatedInfoWithUid[] {
     return Object.values(this.#aggregates).sort((a, b) => b.self - a.self);
   }

src/telemetry/tool_call_metrics.json

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

src/tools/ToolDefinition.ts

@@ -112,6 +112,10 @@ export interface Response {
     stats: DevTools.HeapSnapshotModel.HeapSnapshotModel.Statistics,
     staticData: DevTools.HeapSnapshotModel.HeapSnapshotModel.StaticData | null,
   ): void;
+  setHeapSnapshotNodes(
+    nodes: DevTools.HeapSnapshotModel.HeapSnapshotModel.ItemsRange,
+    options?: PaginationOptions,
+  ): void;
   setIncludePages(value: boolean): void;
   setIncludeNetworkRequests(
     value: boolean,
@@ -234,6 +238,10 @@ export type Context = Readonly<{
   getHeapSnapshotStaticData(
     filePath: string,
   ): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.StaticData | null>;
+  getHeapSnapshotNodesByUid(
+    filePath: string,
+    uid: number,
+  ): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.ItemsRange>;
 }>;
 
 export type ContextPage = Readonly<{

src/tools/memory.ts

@@ -88,3 +88,35 @@ export const getMemorySnapshotDetails = defineTool({
     });
   },
 });
+
+export const getNodesByClass = defineTool({
+  name: 'get_nodes_by_class',
+  description:
+    'Loads a memory heapsnapshot and returns instances of a specific class with their stable IDs.',
+  annotations: {
+    category: ToolCategory.MEMORY,
+    readOnlyHint: true,
+    conditions: ['experimentalMemory'],
+  },
+  schema: {
+    filePath: zod.string().describe('A path to a .heapsnapshot file to read.'),
+    uid: zod
+      .number()
+      .describe(
+        'The unique UID for the class, obtained from aggregates listing.',
+      ),
+    pageIdx: zod.number().optional().describe('The page index for pagination.'),
+    pageSize: zod.number().optional().describe('The page size for pagination.'),
+  },
+  handler: async (request, response, context) => {
+    const nodes = await context.getHeapSnapshotNodesByUid(
+      request.params.filePath,
+      request.params.uid,
+    );
+
+    response.setHeapSnapshotNodes(nodes, {
+      pageIdx: request.params.pageIdx,
+      pageSize: request.params.pageSize,
+    });
+  },
+});

tests/tools/memory.test.js.snapshot

@@ -161,6 +161,20 @@ uid,className,count,selfSize,maxRetainedSize
 108,"HTMLBodyElement (internal cache) / https://example.com",1,16,16
 `;
 
+exports[`memory > get_nodes_by_class > with default options 1`] = `
+## Heap Snapshot Data
+id,name,type,distance,selfSize,retainedSize
+25307,"Array",object,2,192,2056
+33187,"Array",object,2,192,1664
+36255,"Array",object,2,192,1664
+45899,"Array",object,5,56,56
+45901,"Array",object,5,88,88
+46149,"Array",object,5,56,56
+46151,"Array",object,5,88,88
+46355,"Array",object,2,192,2056
+Showing 1-8 of 8 (Page 1 of 1).
+`;
+
 exports[`memory > load_memory_snapshot > with default options 1`] = `
 ## Heap Snapshot Data
 Statistics: {

tests/tools/memory.test.ts

@@ -15,6 +15,7 @@ import {
   takeMemorySnapshot,
   exploreMemorySnapshot,
   getMemorySnapshotDetails,
+  getNodesByClass,
 } from '../../src/tools/memory.js';
 import {withMcpContext} from '../utils.js';
 
@@ -97,4 +98,54 @@ describe('memory', () => {
       });
     });
   });
+
+  describe('get_nodes_by_class', () => {
+    it('with default options', async t => {
+      await withMcpContext(async (response, context) => {
+        const filePath = join(
+          process.cwd(),
+          'tests/fixtures/example.heapsnapshot',
+        );
+
+        await context.getHeapSnapshotAggregates(filePath);
+
+        await getNodesByClass.handler(
+          {params: {filePath, uid: 19}},
+          response,
+          context,
+        );
+
+        const responseData = await response.handle(
+          getNodesByClass.name,
+          context,
+        );
+
+        const output = responseData.content
+          .map(c => (c.type === 'text' ? c.text : ''))
+          .join('\n');
+
+        t.assert.snapshot?.(output);
+      });
+    });
+
+    it('with non-existent class name', async () => {
+      await withMcpContext(async (response, context) => {
+        const filePath = join(
+          process.cwd(),
+          'tests/fixtures/example.heapsnapshot',
+        );
+
+        await context.getHeapSnapshotAggregates(filePath);
+
+        await assert.rejects(
+          getNodesByClass.handler(
+            {params: {filePath, uid: 999999}},
+            response,
+            context,
+          ),
+          {message: 'Class with UID 999999 not found in heap snapshot'},
+        );
+      });
+    });
+  });
 });