chore(memory): expose way to get the nodes

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

@@ -922,4 +922,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

@@ -178,6 +178,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;
@@ -403,6 +404,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);
   }
@@ -704,6 +717,7 @@ export class McpResponse implements Response {
         staticData?: object;
       };
       heapSnapshotData?: object[];
+      heapSnapshotNodes?: object[];
       extensionServiceWorkers?: object[];
       extensionPages?: object[];
     } = {};
@@ -932,6 +946,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 = nodes.items;
+      }
     }
 
     if (data.detailedNetworkRequest) {

src/formatters/HeapSnapshotFormatter.ts

@@ -16,13 +16,46 @@ export interface FormattedSnapshotEntry {
   retainedSize: number;
 }
 
+export interface FormattedNodeEntry {
+  id: number;
+  name: string;
+  type: string;
+  distance: number;
+  selfSize: number;
+  retainedSize: number;
+}
+
 export class HeapSnapshotFormatter {
   #aggregates: Record<string, AggregatedInfoWithUid>;
 
   constructor(aggregates: Record<string, AggregatedInfoWithUid>) {
     this.#aggregates = aggregates;
   }
 
+  static formatNodes(items: readonly unknown[]): string {
+    const lines: string[] = [];
+    lines.push('id,name,type,distance,selfSize,retainedSize');
+
+    for (const item of items) {
+      if (typeof item === 'object' && item !== null) {
+        if (
+          'id' in item &&
+          'name' in item &&
+          'type' in item &&
+          'distance' in item &&
+          'selfSize' in item &&
+          'retainedSize' in 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/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,
@@ -235,6 +239,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,56 @@ 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'},
+        );
+      });
+    });
+  });
 });