Skip to content

Commit 583888f

Browse files
Lightning00BladeLightning00Blade
committed
add caching
1 parent edb6a05 commit 583888f

5 files changed

Lines changed: 184 additions & 77 deletions

File tree

‎src/HeapSnapshotManager.ts‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
/**
2+
* @license
3+
* Copyright 2025 Google LLC
4+
* SPDX-License-Identifier: Apache-2.0
5+
*/
6+
7+
import fsSync from 'node:fs';
8+
import path from 'node:path';
9+
10+
import {DevTools} from './third_party/index.js';
11+
12+
export class HeapSnapshotManager {
13+
#snapshots = new Map<
14+
string,
15+
{
16+
snapshot: DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotProxy;
17+
worker: DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotWorkerProxy;
18+
}
19+
>();
20+
21+
async getSnapshot(
22+
filePath: string,
23+
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotProxy> {
24+
const absolutePath = path.resolve(filePath);
25+
const cached = this.#snapshots.get(absolutePath);
26+
if (cached) {
27+
return cached.snapshot;
28+
}
29+
30+
const {snapshot, worker} = await this.#loadSnapshot(absolutePath);
31+
this.#snapshots.set(absolutePath, {snapshot, worker});
32+
33+
return snapshot;
34+
}
35+
36+
async getAggregates(
37+
filePath: string,
38+
): Promise<
39+
Record<string, DevTools.HeapSnapshotModel.HeapSnapshotModel.AggregatedInfo>
40+
> {
41+
const snapshot = await this.getSnapshot(filePath);
42+
const filter =
43+
new DevTools.HeapSnapshotModel.HeapSnapshotModel.NodeFilter();
44+
return await snapshot.aggregatesWithFilter(filter);
45+
}
46+
47+
async getStats(
48+
filePath: string,
49+
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.Statistics> {
50+
const snapshot = await this.getSnapshot(filePath);
51+
return await snapshot.getStatistics();
52+
}
53+
54+
async getStaticData(
55+
filePath: string,
56+
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.StaticData | null> {
57+
const snapshot = await this.getSnapshot(filePath);
58+
return snapshot.staticData;
59+
}
60+
61+
async #loadSnapshot(absolutePath: string): Promise<{
62+
snapshot: DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotProxy;
63+
worker: DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotWorkerProxy;
64+
}> {
65+
const workerProxy =
66+
new DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotWorkerProxy(
67+
() => {
68+
/* noop */
69+
},
70+
import.meta.resolve('./third_party/devtools-heap-snapshot-worker.js'),
71+
);
72+
73+
const {promise: snapshotPromise, resolve: resolveSnapshot} =
74+
Promise.withResolvers<DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotProxy>();
75+
76+
const loaderProxy = workerProxy.createLoader(1, snapshotProxy => {
77+
resolveSnapshot(snapshotProxy);
78+
});
79+
80+
const fileStream = fsSync.createReadStream(absolutePath, {
81+
encoding: 'utf-8',
82+
highWaterMark: 1024 * 1024,
83+
});
84+
85+
for await (const chunk of fileStream) {
86+
await loaderProxy.write(chunk);
87+
}
88+
89+
await loaderProxy.close();
90+
91+
const snapshot = await snapshotPromise;
92+
return {snapshot, worker: workerProxy};
93+
}
94+
95+
dispose(filePath: string): void {
96+
const absolutePath = path.resolve(filePath);
97+
const cached = this.#snapshots.get(absolutePath);
98+
if (cached) {
99+
cached.worker.dispose();
100+
this.#snapshots.delete(absolutePath);
101+
}
102+
}
103+
}

‎src/McpContext.ts‎

Lines changed: 19 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,12 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66

7-
import fsSync from 'node:fs';
87
import fs from 'node:fs/promises';
98
import path from 'node:path';
109

1110
import type {TargetUniverse} from './DevtoolsUtils.js';
1211
import {UniverseManager} from './DevtoolsUtils.js';
12+
import {HeapSnapshotManager} from './HeapSnapshotManager.js';
1313
import {McpPage} from './McpPage.js';
1414
import {
1515
NetworkCollector,
@@ -29,7 +29,7 @@ import type {
2929
Viewport,
3030
Target,
3131
} from './third_party/index.js';
32-
import {DevTools} from './third_party/index.js';
32+
import type {DevTools} from './third_party/index.js';
3333
import {Locator} from './third_party/index.js';
3434
import {PredefinedNetworkConditions} from './third_party/index.js';
3535
import {listPages} from './tools/pages.js';
@@ -100,6 +100,7 @@ export class McpContext implements Context {
100100

101101
#locatorClass: typeof Locator;
102102
#options: McpContextOptions;
103+
#heapSnapshotManager = new HeapSnapshotManager();
103104

104105
private constructor(
105106
browser: Browser,
@@ -915,44 +916,23 @@ export class McpContext implements Context {
915916
return this.#extensionRegistry.getById(id);
916917
}
917918

918-
async getHeapSnapshotProxy(
919-
heapsnapshotPath: string,
920-
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotProxy> {
921-
const workerProxy =
922-
new DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotWorkerProxy(
923-
() => {
924-
/* noop */
925-
},
926-
import.meta.resolve('./third_party/devtools-heap-snapshot-worker.js'),
927-
);
928-
929-
const absolutePath = path.resolve(heapsnapshotPath);
930-
931-
const {promise: snapshotPromise, resolve: resolveSnapshot} =
932-
Promise.withResolvers<DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotProxy>();
933-
934-
const loaderProxy = workerProxy.createLoader(1, snapshotProxy => {
935-
resolveSnapshot(snapshotProxy);
936-
});
937-
938-
const fileStream = fsSync.createReadStream(absolutePath, {
939-
encoding: 'utf-8',
940-
highWaterMark: 1024 * 1024,
941-
});
942-
943-
for await (const chunk of fileStream) {
944-
await loaderProxy.write(chunk);
945-
}
946-
947-
await loaderProxy.close();
948-
949-
const snapshot = await snapshotPromise;
919+
async getHeapSnapshotAggregates(
920+
filePath: string,
921+
): Promise<
922+
Record<string, DevTools.HeapSnapshotModel.HeapSnapshotModel.AggregatedInfo>
923+
> {
924+
return await this.#heapSnapshotManager.getAggregates(filePath);
925+
}
950926

951-
//TODO Figure ot how to dispose the workeProxy.
952-
// Also the snapshot methods hangs if the workerProxy is
953-
// duposed instead of throwing
927+
async getHeapSnapshotStats(
928+
filePath: string,
929+
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.Statistics> {
930+
return await this.#heapSnapshotManager.getStats(filePath);
931+
}
954932

955-
// workerProxy.dispose();
956-
return snapshot;
933+
async getHeapSnapshotStaticData(
934+
filePath: string,
935+
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.StaticData | null> {
936+
return await this.#heapSnapshotManager.getStaticData(filePath);
957937
}
958938
}

‎src/McpResponse.ts‎

Lines changed: 43 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -171,11 +171,13 @@ export class McpResponse implements Response {
171171
#images: ImageContentData[] = [];
172172
#heapSnapshotOptions?: {
173173
include: boolean;
174-
aggregates: Record<
174+
aggregates?: Record<
175175
string,
176176
DevTools.HeapSnapshotModel.HeapSnapshotModel.AggregatedInfo
177177
>;
178178
pagination?: PaginationOptions;
179+
stats?: DevTools.HeapSnapshotModel.HeapSnapshotModel.Statistics;
180+
staticData?: DevTools.HeapSnapshotModel.HeapSnapshotModel.StaticData | null;
179181
};
180182
#networkRequestsOptions?: {
181183
include: boolean;
@@ -374,20 +376,33 @@ export class McpResponse implements Response {
374376
this.#textResponseLines.push(value);
375377
}
376378

377-
setHeapSnapshot(
379+
setHeapSnapshotAggregates(
378380
aggregates: Record<
379381
string,
380382
DevTools.HeapSnapshotModel.HeapSnapshotModel.AggregatedInfo
381383
>,
382384
options?: PaginationOptions,
383385
) {
384386
this.#heapSnapshotOptions = {
387+
...this.#heapSnapshotOptions,
385388
include: true,
386389
aggregates,
387390
pagination: options,
388391
};
389392
}
390393

394+
setHeapSnapshotStats(
395+
stats: DevTools.HeapSnapshotModel.HeapSnapshotModel.Statistics,
396+
staticData: DevTools.HeapSnapshotModel.HeapSnapshotModel.StaticData | null,
397+
) {
398+
this.#heapSnapshotOptions = {
399+
...this.#heapSnapshotOptions,
400+
include: true,
401+
stats,
402+
staticData,
403+
};
404+
}
405+
391406
attachImage(value: ImageContentData): void {
392407
this.#images.push(value);
393408
}
@@ -685,6 +700,8 @@ export class McpResponse implements Response {
685700
pages?: object[];
686701
pagination?: object;
687702
heapSnapshot?: object[];
703+
heapSnapshotStats?: object;
704+
heapSnapshotStaticData?: object;
688705
extensionServiceWorkers?: object[];
689706
extensionPages?: object[];
690707
} = {};
@@ -882,23 +899,35 @@ Call ${handleDialog.name} to handle it before continuing.`);
882899
}
883900

884901
if (this.#heapSnapshotOptions?.include) {
902+
const stats = this.#heapSnapshotOptions.stats;
903+
const staticData = this.#heapSnapshotOptions.staticData;
904+
if (stats) {
905+
response.push(`Statistics: ${JSON.stringify(stats, null, 2)}`);
906+
structuredContent.heapSnapshotStats = stats;
907+
}
908+
if (staticData) {
909+
response.push(`Static Data: ${JSON.stringify(staticData, null, 2)}`);
910+
structuredContent.heapSnapshotStaticData = staticData;
911+
}
885912
const aggregates = this.#heapSnapshotOptions.aggregates;
886-
const entries = Object.entries(aggregates);
887-
const sortedEntries = entries.sort((a, b) => b[1].self - a[1].self);
913+
if (aggregates) {
914+
const entries = Object.entries(aggregates);
915+
const sortedEntries = entries.sort((a, b) => b[1].self - a[1].self);
888916

889-
const paginationData = this.#dataWithPagination(
890-
sortedEntries,
891-
this.#heapSnapshotOptions.pagination,
892-
);
917+
const paginationData = this.#dataWithPagination(
918+
sortedEntries,
919+
this.#heapSnapshotOptions.pagination,
920+
);
893921

894-
structuredContent.pagination = paginationData.pagination;
895-
response.push(...paginationData.info);
922+
structuredContent.pagination = paginationData.pagination;
923+
response.push(...paginationData.info);
896924

897-
const paginatedRecord = Object.fromEntries(paginationData.items);
898-
const formatter = new HeapSnapshotFormatter(paginatedRecord);
925+
const paginatedRecord = Object.fromEntries(paginationData.items);
926+
const formatter = new HeapSnapshotFormatter(paginatedRecord);
899927

900-
response.push(formatter.toString());
901-
structuredContent.heapSnapshot = formatter.toJSON();
928+
response.push(formatter.toString());
929+
structuredContent.heapSnapshot = formatter.toJSON();
930+
}
902931
}
903932

904933
if (data.detailedNetworkRequest) {

‎src/tools/ToolDefinition.ts‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -100,13 +100,17 @@ export interface DevToolsData {
100100

101101
export interface Response {
102102
appendResponseLine(value: string): void;
103-
setHeapSnapshot(
103+
setHeapSnapshotAggregates(
104104
aggregates: Record<
105105
string,
106106
DevTools.HeapSnapshotModel.HeapSnapshotModel.AggregatedInfo
107107
>,
108108
options?: PaginationOptions,
109109
): void;
110+
setHeapSnapshotStats(
111+
stats: DevTools.HeapSnapshotModel.HeapSnapshotModel.Statistics,
112+
staticData: DevTools.HeapSnapshotModel.HeapSnapshotModel.StaticData | null,
113+
): void;
110114
setIncludePages(value: boolean): void;
111115
setIncludeNetworkRequests(
112116
value: boolean,
@@ -221,9 +225,13 @@ export type Context = Readonly<{
221225
getExtensionServiceWorkerId(
222226
extensionServiceWorker: ExtensionServiceWorker,
223227
): string | undefined;
224-
getHeapSnapshotProxy(
225-
heapsnapshotPath: string,
226-
): Promise<DevTools.HeapSnapshotModel.HeapSnapshotProxy.HeapSnapshotProxy>;
228+
getHeapSnapshotAggregates(
229+
filePath: string,
230+
): Promise<
231+
Record<string, DevTools.HeapSnapshotModel.HeapSnapshotModel.AggregatedInfo>
232+
>;
233+
getHeapSnapshotStats(filePath: string): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.Statistics>;
234+
getHeapSnapshotStaticData(filePath: string): Promise<DevTools.HeapSnapshotModel.HeapSnapshotModel.StaticData | null>;
227235
}>;
228236

229237
export type ContextPage = Readonly<{

‎src/tools/memory.ts‎

Lines changed: 7 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66

7-
import {zod, DevTools} from '../third_party/index.js';
7+
import {zod} from '../third_party/index.js';
88
import {ensureExtension} from '../utils/files.js';
99

1010
import {ToolCategory} from './categories.js';
@@ -36,35 +36,22 @@ export const takeMemorySnapshot = definePageTool({
3636
});
3737

3838
export const exploreMemorySnapshot = defineTool({
39-
name: 'explored_memory_snapshot',
40-
description: 'Explose ',
39+
name: 'load_memory_snapshot',
40+
description:
41+
'Loads a memory heapsnapshot and returns snapshot summary stats. ',
4142
annotations: {
4243
category: ToolCategory.PERFORMANCE,
4344
readOnlyHint: true,
4445
},
4546
schema: {
4647
filePath: zod.string().describe('A path to a .heapsnapshot file to read.'),
47-
pageSize: zod.number().optional().describe('Page size for pagination.'),
48-
pageIdx: zod.number().optional().describe('Page index for pagination.'),
4948
},
5049
handler: async (request, response, context) => {
51-
const snapshot = await context.getHeapSnapshotProxy(
50+
const stats = await context.getHeapSnapshotStats(request.params.filePath);
51+
const staticData = await context.getHeapSnapshotStaticData(
5252
request.params.filePath,
5353
);
54-
const stats = await snapshot.getStatistics();
5554

56-
response.appendResponseLine(
57-
`Statistics: ${JSON.stringify(stats, null, 2)}`,
58-
);
59-
response.appendResponseLine(
60-
`Static Data: ${JSON.stringify(snapshot.staticData, null, 2)}`,
61-
);
62-
63-
const filter =
64-
new DevTools.HeapSnapshotModel.HeapSnapshotModel.NodeFilter();
65-
const aggregates = await snapshot.aggregatesWithFilter(filter);
66-
67-
const {pageSize, pageIdx} = request.params;
68-
response.setHeapSnapshot(aggregates, {pageSize, pageIdx});
55+
response.setHeapSnapshotStats(stats, staticData);
6956
},
7057
});

0 commit comments

Comments
 (0)