feat: support filePath in performance tools (#686)

Closes #680

Author: OrKoN

docs/tool-reference.md

@@ -234,14 +234,17 @@
 
 - **autoStop** (boolean) **(required)**: Determines if the trace recording should be automatically stopped.
 - **reload** (boolean) **(required)**: Determines if, once tracing has started, the page should be automatically reloaded.
+- **filePath** (string) _(optional)_: The absolute file path, or a file path relative to the current working directory, to save the raw trace data. For example, trace.json.gz (compressed) or trace.json (uncompressed).
 
 ---
 
 ### `performance_stop_trace`
 
 **Description:** Stops the active performance trace recording on the selected page.
 
-**Parameters:** None
+**Parameters:**
+
+- **filePath** (string) _(optional)_: The absolute file path, or a file path relative to the current working directory, to save the raw trace data. For example, trace.json.gz (compressed) or trace.json (uncompressed).
 
 ---
 

src/tools/performance.ts

@@ -4,6 +4,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import zlib from 'node:zlib';
+
 import {logger} from '../logger.js';
 import {zod} from '../third_party/index.js';
 import type {Page} from '../third_party/index.js';
@@ -19,13 +21,20 @@ import {ToolCategory} from './categories.js';
 import type {Context, Response} from './ToolDefinition.js';
 import {defineTool} from './ToolDefinition.js';
 
+const filePathSchema = zod
+  .string()
+  .optional()
+  .describe(
+    'The absolute file path, or a file path relative to the current working directory, to save the raw trace data. For example, trace.json.gz (compressed) or trace.json (uncompressed).',
+  );
+
 export const startTrace = defineTool({
   name: 'performance_start_trace',
   description:
     '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.',
   annotations: {
     category: ToolCategory.PERFORMANCE,
-    readOnlyHint: true,
+    readOnlyHint: false,
   },
   schema: {
     reload: zod
@@ -38,6 +47,7 @@ export const startTrace = defineTool({
       .describe(
         'Determines if the trace recording should be automatically stopped.',
       ),
+    filePath: filePathSchema,
   },
   handler: async (request, response, context) => {
     if (context.isRunningPerformanceTrace()) {
@@ -91,7 +101,12 @@ export const startTrace = defineTool({
 
     if (request.params.autoStop) {
       await new Promise(resolve => setTimeout(resolve, 5_000));
-      await stopTracingAndAppendOutput(page, response, context);
+      await stopTracingAndAppendOutput(
+        page,
+        response,
+        context,
+        request.params.filePath,
+      );
     } else {
       response.appendResponseLine(
         `The performance trace is being recorded. Use performance_stop_trace to stop it.`,
@@ -106,15 +121,22 @@ export const stopTrace = defineTool({
     'Stops the active performance trace recording on the selected page.',
   annotations: {
     category: ToolCategory.PERFORMANCE,
-    readOnlyHint: true,
+    readOnlyHint: false,
+  },
+  schema: {
+    filePath: filePathSchema,
   },
-  schema: {},
-  handler: async (_request, response, context) => {
+  handler: async (request, response, context) => {
     if (!context.isRunningPerformanceTrace()) {
       return;
     }
     const page = context.getSelectedPage();
-    await stopTracingAndAppendOutput(page, response, context);
+    await stopTracingAndAppendOutput(
+      page,
+      response,
+      context,
+      request.params.filePath,
+    );
   },
 });
 
@@ -165,9 +187,28 @@ async function stopTracingAndAppendOutput(
   page: Page,
   response: Response,
   context: Context,
+  filePath?: string,
 ): Promise<void> {
   try {
     const traceEventsBuffer = await page.tracing.stop();
+    if (filePath && traceEventsBuffer) {
+      let dataToWrite: Uint8Array = traceEventsBuffer;
+      if (filePath.endsWith('.gz')) {
+        dataToWrite = await new Promise((resolve, reject) => {
+          zlib.gzip(traceEventsBuffer, (error, result) => {
+            if (error) {
+              reject(error);
+            } else {
+              resolve(result);
+            }
+          });
+        });
+      }
+      const file = await context.saveFile(dataToWrite, filePath);
+      response.appendResponseLine(
+        `The raw trace data was saved to ${file.filename}.`,
+      );
+    }
     const result = await parseRawTraceBuffer(traceEventsBuffer);
     response.appendResponseLine('The performance trace has been stopped.');
     if (traceResultIsSuccess(result)) {

tests/tools/performance.test.ts

@@ -6,6 +6,7 @@
 
 import assert from 'node:assert';
 import {describe, it, afterEach} from 'node:test';
+import zlib from 'node:zlib';
 
 import sinon from 'sinon';
 
@@ -138,6 +139,49 @@ describe('performance', () => {
         );
       });
     });
+
+    it.only('supports filePath', async () => {
+      const rawData = loadTraceAsBuffer('basic-trace.json.gz');
+      // rawData is the decompressed buffer (based on loadTraceAsBuffer implementation).
+      // We want to simulate saving it as a .gz file, so the tool should compress it.
+      const expectedCompressedData = zlib.gzipSync(rawData);
+
+      await withMcpContext(async (response, context) => {
+        const filePath = 'test-trace.json.gz';
+        const selectedPage = context.getSelectedPage();
+        sinon.stub(selectedPage, 'url').callsFake(() => 'https://www.test.com');
+        sinon.stub(selectedPage, 'goto').callsFake(() => Promise.resolve(null));
+        sinon.stub(selectedPage.tracing, 'start');
+        sinon.stub(selectedPage.tracing, 'stop').resolves(rawData);
+        const saveFileStub = sinon
+          .stub(context, 'saveFile')
+          .resolves({filename: filePath});
+
+        const handlerPromise = startTrace.handler(
+          {params: {reload: true, autoStop: true, filePath}},
+          response,
+          context,
+        );
+        // In the handler we wait 5 seconds after the page load event (which is
+        // what DevTools does), hence we now fake-progress time to allow
+        // the handler to complete. We allow extra time because the Trace
+        // Engine also uses some timers to yield updates and we need those to
+        // execute.
+        await handlerPromise;
+
+        assert.ok(
+          response.responseLines.includes(
+            `The raw trace data was saved to ${filePath}.`,
+          ),
+        );
+        sinon.assert.calledOnce(saveFileStub);
+        const [savedData, savedPath] = saveFileStub.firstCall.args;
+        assert.strictEqual(savedPath, filePath);
+        // Compare the saved data with expected compressed data
+        // We can't compare buffers directly with strictEqual easily if they are different instances, but deepStrictEqual works for Buffers.
+        assert.deepStrictEqual(savedData, expectedCompressedData);
+      });
+    });
   });
 
   describe('performance_analyze_insight', () => {
@@ -275,5 +319,31 @@ describe('performance', () => {
         t.assert.snapshot?.(response.responseLines.join('\n'));
       });
     });
+
+    it('supports filePath', async () => {
+      const rawData = loadTraceAsBuffer('basic-trace.json.gz');
+      await withMcpContext(async (response, context) => {
+        const filePath = 'test-trace.json';
+        context.setIsRunningPerformanceTrace(true);
+        const selectedPage = context.getSelectedPage();
+        const stopTracingStub = sinon
+          .stub(selectedPage.tracing, 'stop')
+          .resolves(rawData);
+        const saveFileStub = sinon
+          .stub(context, 'saveFile')
+          .resolves({filename: filePath});
+
+        await stopTrace.handler({params: {filePath}}, response, context);
+
+        sinon.assert.calledOnce(stopTracingStub);
+        sinon.assert.calledOnce(saveFileStub);
+        sinon.assert.calledWith(saveFileStub, rawData, filePath);
+        assert.ok(
+          response.responseLines.includes(
+            `The raw trace data was saved to ${filePath}.`,
+          ),
+        );
+      });
+    });
   });
 });