diff --git a/src/HeapSnapshotManager.ts b/src/HeapSnapshotManager.ts index ea9c022a5..81d875e3e 100644 --- a/src/HeapSnapshotManager.ts +++ b/src/HeapSnapshotManager.ts @@ -99,6 +99,23 @@ export class HeapSnapshotManager { return uid; } + async getNodesByUid( + filePath: string, + uid: number, + ): Promise { + 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 { + 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; diff --git a/src/McpContext.ts b/src/McpContext.ts index 1aacb391c..9d6bdadd9 100644 --- a/src/McpContext.ts +++ b/src/McpContext.ts @@ -922,4 +922,11 @@ export class McpContext implements Context { ): Promise { return await this.#heapSnapshotManager.getStaticData(filePath); } + + async getHeapSnapshotNodesByUid( + filePath: string, + uid: number, + ): Promise { + return await this.#heapSnapshotManager.getNodesByUid(filePath, uid); + } } diff --git a/src/McpResponse.ts b/src/McpResponse.ts index 0d5f926ec..ea8098934 100644 --- a/src/McpResponse.ts +++ b/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?: readonly 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 = paginationData.items; + } } if (data.detailedNetworkRequest) { diff --git a/src/formatters/HeapSnapshotFormatter.ts b/src/formatters/HeapSnapshotFormatter.ts index f5bcda0c9..9fa6afe0b 100644 --- a/src/formatters/HeapSnapshotFormatter.ts +++ b/src/formatters/HeapSnapshotFormatter.ts @@ -16,6 +16,14 @@ 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; @@ -23,6 +31,29 @@ export class HeapSnapshotFormatter { 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); } diff --git a/src/telemetry/tool_call_metrics.json b/src/telemetry/tool_call_metrics.json index 14d8c13bf..1b35a86c8 100644 --- a/src/telemetry/tool_call_metrics.json +++ b/src/telemetry/tool_call_metrics.json @@ -586,5 +586,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" + } + ] } ] diff --git a/src/tools/ToolDefinition.ts b/src/tools/ToolDefinition.ts index 1c42915d5..e9d8a8a77 100644 --- a/src/tools/ToolDefinition.ts +++ b/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; + getHeapSnapshotNodesByUid( + filePath: string, + uid: number, + ): Promise; }>; export type ContextPage = Readonly<{ diff --git a/src/tools/memory.ts b/src/tools/memory.ts index 67281b029..db61dfb7b 100644 --- a/src/tools/memory.ts +++ b/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, + }); + }, +}); diff --git a/tests/tools/memory.test.js.snapshot b/tests/tools/memory.test.js.snapshot index 4d091bff4..d2fe2c9cb 100644 --- a/tests/tools/memory.test.js.snapshot +++ b/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: { diff --git a/tests/tools/memory.test.ts b/tests/tools/memory.test.ts index f7421ddb0..d6c1931f5 100644 --- a/tests/tools/memory.test.ts +++ b/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'}, + ); + }); + }); + }); });