Skip to content

Commit 2c9ea1c

Browse files
committed
feat: support filePath in performance tools
1 parent 14ff400 commit 2c9ea1c

3 files changed

Lines changed: 119 additions & 7 deletions

File tree

docs/tool-reference.md

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -234,14 +234,17 @@
234234

235235
- **autoStop** (boolean) **(required)**: Determines if the trace recording should be automatically stopped.
236236
- **reload** (boolean) **(required)**: Determines if, once tracing has started, the page should be automatically reloaded.
237+
- **filePath** (string) _(optional)_: The absolute path, or a path relative to the current working directory, to save the raw trace data.
237238

238239
---
239240

240241
### `performance_stop_trace`
241242

242243
**Description:** Stops the active performance trace recording on the selected page.
243244

244-
**Parameters:** None
245+
**Parameters:**
246+
247+
- **filePath** (string) _(optional)_: The absolute path, or a path relative to the current working directory, to save the raw trace data.
245248

246249
---
247250

src/tools/performance.ts

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

7+
import zlib from 'node:zlib';
8+
import util from 'node:util';
9+
710
import {logger} from '../logger.js';
811
import {zod} from '../third_party/index.js';
912
import type {Page} from '../third_party/index.js';
@@ -25,7 +28,7 @@ export const startTrace = defineTool({
2528
'Starts a performance trace recording on the selected page. This can be used to look for performance problems and insights to improve the performance of the page. It will also report Core Web Vital (CWV) scores for the page.',
2629
annotations: {
2730
category: ToolCategory.PERFORMANCE,
28-
readOnlyHint: true,
31+
readOnlyHint: false,
2932
},
3033
schema: {
3134
reload: zod
@@ -38,6 +41,12 @@ export const startTrace = defineTool({
3841
.describe(
3942
'Determines if the trace recording should be automatically stopped.',
4043
),
44+
filePath: zod
45+
.string()
46+
.optional()
47+
.describe(
48+
'The absolute path, or a path relative to the current working directory, to save the raw trace data.',
49+
),
4150
},
4251
handler: async (request, response, context) => {
4352
if (context.isRunningPerformanceTrace()) {
@@ -91,7 +100,12 @@ export const startTrace = defineTool({
91100

92101
if (request.params.autoStop) {
93102
await new Promise(resolve => setTimeout(resolve, 5_000));
94-
await stopTracingAndAppendOutput(page, response, context);
103+
await stopTracingAndAppendOutput(
104+
page,
105+
response,
106+
context,
107+
request.params.filePath,
108+
);
95109
} else {
96110
response.appendResponseLine(
97111
`The performance trace is being recorded. Use performance_stop_trace to stop it.`,
@@ -106,15 +120,27 @@ export const stopTrace = defineTool({
106120
'Stops the active performance trace recording on the selected page.',
107121
annotations: {
108122
category: ToolCategory.PERFORMANCE,
109-
readOnlyHint: true,
123+
readOnlyHint: false,
110124
},
111-
schema: {},
112-
handler: async (_request, response, context) => {
125+
schema: {
126+
filePath: zod
127+
.string()
128+
.optional()
129+
.describe(
130+
'The absolute path, or a path relative to the current working directory, to save the raw trace data.',
131+
),
132+
},
133+
handler: async (request, response, context) => {
113134
if (!context.isRunningPerformanceTrace()) {
114135
return;
115136
}
116137
const page = context.getSelectedPage();
117-
await stopTracingAndAppendOutput(page, response, context);
138+
await stopTracingAndAppendOutput(
139+
page,
140+
response,
141+
context,
142+
request.params.filePath,
143+
);
118144
},
119145
});
120146

@@ -165,9 +191,21 @@ async function stopTracingAndAppendOutput(
165191
page: Page,
166192
response: Response,
167193
context: Context,
194+
filePath?: string,
168195
): Promise<void> {
169196
try {
170197
const traceEventsBuffer = await page.tracing.stop();
198+
if (filePath && traceEventsBuffer) {
199+
let dataToWrite: Uint8Array = traceEventsBuffer;
200+
if (filePath.endsWith('.gz')) {
201+
const gzip = util.promisify(zlib.gzip);
202+
dataToWrite = await gzip(traceEventsBuffer);
203+
}
204+
const file = await context.saveFile(dataToWrite, filePath);
205+
response.appendResponseLine(
206+
`The raw trace data was saved to ${file.filename}.`,
207+
);
208+
}
171209
const result = await parseRawTraceBuffer(traceEventsBuffer);
172210
response.appendResponseLine('The performance trace has been stopped.');
173211
if (traceResultIsSuccess(result)) {

tests/tools/performance.test.ts

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
*/
66

77
import assert from 'node:assert';
8+
import zlib from 'node:zlib';
89
import {describe, it, afterEach} from 'node:test';
910

1011
import sinon from 'sinon';
@@ -138,6 +139,50 @@ describe('performance', () => {
138139
);
139140
});
140141
});
142+
143+
it('supports filePath', async () => {
144+
const rawData = loadTraceAsBuffer('basic-trace.json.gz');
145+
// rawData is the decompressed buffer (based on loadTraceAsBuffer implementation).
146+
// We want to simulate saving it as a .gz file, so the tool should compress it.
147+
const expectedCompressedData = zlib.gzipSync(rawData);
148+
149+
await withMcpContext(async (response, context) => {
150+
const filePath = 'test-trace.json.gz';
151+
const selectedPage = context.getSelectedPage();
152+
sinon.stub(selectedPage, 'url').callsFake(() => 'https://www.test.com');
153+
sinon.stub(selectedPage, 'goto').callsFake(() => Promise.resolve(null));
154+
sinon.stub(selectedPage.tracing, 'start');
155+
sinon.stub(selectedPage.tracing, 'stop').resolves(rawData);
156+
const saveFileStub = sinon
157+
.stub(context, 'saveFile')
158+
.resolves({filename: filePath});
159+
160+
const clock = sinon.useFakeTimers({shouldClearNativeTimers: true});
161+
try {
162+
const handlerPromise = startTrace.handler(
163+
{params: {reload: true, autoStop: true, filePath}},
164+
response,
165+
context,
166+
);
167+
await clock.tickAsync(6_000);
168+
await handlerPromise;
169+
170+
assert.ok(
171+
response.responseLines.includes(
172+
`The raw trace data was saved to ${filePath}.`,
173+
),
174+
);
175+
sinon.assert.calledOnce(saveFileStub);
176+
const [savedData, savedPath] = saveFileStub.firstCall.args;
177+
assert.strictEqual(savedPath, filePath);
178+
// Compare the saved data with expected compressed data
179+
// We can't compare buffers directly with strictEqual easily if they are different instances, but deepStrictEqual works for Buffers.
180+
assert.deepStrictEqual(savedData, expectedCompressedData);
181+
} finally {
182+
clock.restore();
183+
}
184+
});
185+
});
141186
});
142187

143188
describe('performance_analyze_insight', () => {
@@ -275,5 +320,31 @@ describe('performance', () => {
275320
t.assert.snapshot?.(response.responseLines.join('\n'));
276321
});
277322
});
323+
324+
it('supports filePath', async () => {
325+
const rawData = loadTraceAsBuffer('basic-trace.json.gz');
326+
await withMcpContext(async (response, context) => {
327+
const filePath = 'test-trace.json';
328+
context.setIsRunningPerformanceTrace(true);
329+
const selectedPage = context.getSelectedPage();
330+
const stopTracingStub = sinon
331+
.stub(selectedPage.tracing, 'stop')
332+
.resolves(rawData);
333+
const saveFileStub = sinon
334+
.stub(context, 'saveFile')
335+
.resolves({filename: filePath});
336+
337+
await stopTrace.handler({params: {filePath}}, response, context);
338+
339+
sinon.assert.calledOnce(stopTracingStub);
340+
sinon.assert.calledOnce(saveFileStub);
341+
sinon.assert.calledWith(saveFileStub, rawData, filePath);
342+
assert.ok(
343+
response.responseLines.includes(
344+
`The raw trace data was saved to ${filePath}.`,
345+
),
346+
);
347+
});
348+
});
278349
});
279350
});

0 commit comments

Comments
 (0)