chore: structured performance responses (#809)

This PR moves formatting out of the tools and to the McpResponse.
Drive-by: clears up past trace results.

Author: OrKoN

src/McpContext.ts

@@ -667,6 +667,8 @@ export class McpContext implements Context {
   }
 
   storeTraceRecording(result: TraceResult): void {
+    // Clear the trace results because we only consume the latest trace currently.
+    this.#traceResults = [];
     this.#traceResults.push(result);
   }
 

src/McpResponse.ts

@@ -23,9 +23,17 @@ import type {
   Response,
   SnapshotParams,
 } from './tools/ToolDefinition.js';
+import type {InsightName, TraceResult} from './trace-processing/parse.js';
+import {getInsightOutput, getTraceSummary} from './trace-processing/parse.js';
 import {paginate} from './utils/pagination.js';
 import type {PaginationOptions} from './utils/types.js';
 
+interface TraceInsightData {
+  trace: TraceResult;
+  insightSetId: string;
+  insightName: InsightName;
+}
+
 export class McpResponse implements Response {
   #includePages = false;
   #snapshotParams?: SnapshotParams;
@@ -35,6 +43,8 @@ export class McpResponse implements Response {
     responseFilePath?: string;
   };
   #attachedConsoleMessageId?: number;
+  #attachedTraceSummary?: TraceResult;
+  #attachedTraceInsight?: TraceInsightData;
   #textResponseLines: string[] = [];
   #images: ImageContentData[] = [];
   #networkRequestsOptions?: {
@@ -137,10 +147,34 @@ export class McpResponse implements Response {
     this.#attachedConsoleMessageId = msgid;
   }
 
+  attachTraceSummary(result: TraceResult): void {
+    this.#attachedTraceSummary = result;
+  }
+
+  attachTraceInsight(
+    trace: TraceResult,
+    insightSetId: string,
+    insightName: InsightName,
+  ): void {
+    this.#attachedTraceInsight = {
+      trace,
+      insightSetId,
+      insightName,
+    };
+  }
+
   get includePages(): boolean {
     return this.#includePages;
   }
 
+  get attachedTraceSummary(): TraceResult | undefined {
+    return this.#attachedTraceSummary;
+  }
+
+  get attachedTracedInsight(): TraceInsightData | undefined {
+    return this.#attachedTraceInsight;
+  }
+
   get includeNetworkRequests(): boolean {
     return this.#networkRequestsOptions?.include ?? false;
   }
@@ -359,6 +393,8 @@ export class McpResponse implements Response {
       snapshot,
       detailedNetworkRequest,
       networkRequests,
+      traceInsight: this.#attachedTraceInsight,
+      traceSummary: this.#attachedTraceSummary,
     });
   }
 
@@ -371,6 +407,8 @@ export class McpResponse implements Response {
       snapshot: SnapshotFormatter | string | undefined;
       detailedNetworkRequest?: NetworkFormatter;
       networkRequests?: NetworkFormatter[];
+      traceSummary?: TraceResult;
+      traceInsight?: TraceInsightData;
     },
   ): {content: Array<TextContent | ImageContent>; structuredContent: object} {
     const response = [`# ${toolName} response`];
@@ -434,12 +472,42 @@ Call ${handleDialog.name} to handle it before continuing.`);
       networkRequests?: object[];
       consoleMessage?: object;
       consoleMessages?: object[];
+      traceSummary?: string;
+      traceInsights?: Array<{insightName: string; insightKey: string}>;
     } = {};
 
     if (this.#tabId) {
       structuredContent.tabId = this.#tabId;
     }
 
+    if (data.traceSummary) {
+      const summary = getTraceSummary(data.traceSummary);
+      response.push(summary);
+      structuredContent.traceSummary = summary;
+      structuredContent.traceInsights = [];
+      for (const insightSet of data.traceSummary.insights?.values() ?? []) {
+        for (const [insightName, model] of Object.entries(insightSet.model)) {
+          structuredContent.traceInsights.push({
+            insightName,
+            insightKey: model.insightKey,
+          });
+        }
+      }
+    }
+
+    if (data.traceInsight) {
+      const insightOutput = getInsightOutput(
+        data.traceInsight.trace,
+        data.traceInsight.insightSetId,
+        data.traceInsight.insightName,
+      );
+      if ('error' in insightOutput) {
+        response.push(insightOutput.error);
+      } else {
+        response.push(insightOutput.output);
+      }
+    }
+
     if (data.snapshot) {
       if (typeof data.snapshot === 'string') {
         response.push(`Saved snapshot to ${data.snapshot}.`);

src/tools/ToolDefinition.ts

@@ -12,7 +12,7 @@ import type {
   Page,
   Viewport,
 } from '../third_party/index.js';
-import type {TraceResult} from '../trace-processing/parse.js';
+import type {InsightName, TraceResult} from '../trace-processing/parse.js';
 import type {PaginationOptions} from '../utils/types.js';
 
 import type {ToolCategory} from './categories.js';
@@ -86,6 +86,12 @@ export interface Response {
   // Allows re-using DevTools data queried by some tools.
   attachDevToolsData(data: DevToolsData): void;
   setTabId(tabId: string): void;
+  attachTraceSummary(trace: TraceResult): void;
+  attachTraceInsight(
+    trace: TraceResult,
+    insightSetId: string,
+    insightName: InsightName,
+  ): void;
 }
 
 /**

src/tools/performance.ts

@@ -6,13 +6,10 @@
 
 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';
 import type {InsightName} from '../trace-processing/parse.js';
 import {
-  getInsightOutput,
-  getTraceSummary,
   parseRawTraceBuffer,
   traceResultIsSuccess,
 } from '../trace-processing/parse.js';
@@ -168,17 +165,11 @@ export const analyzeInsight = defineTool({
       return;
     }
 
-    const insightOutput = getInsightOutput(
+    response.attachTraceInsight(
       lastRecording,
       request.params.insightSetId,
       request.params.insightName as InsightName,
     );
-    if ('error' in insightOutput) {
-      response.appendResponseLine(insightOutput.error);
-      return;
-    }
-
-    response.appendResponseLine(insightOutput.output);
   },
 });
 
@@ -212,21 +203,12 @@ async function stopTracingAndAppendOutput(
     response.appendResponseLine('The performance trace has been stopped.');
     if (traceResultIsSuccess(result)) {
       context.storeTraceRecording(result);
-      const traceSummaryText = getTraceSummary(result);
-      response.appendResponseLine(traceSummaryText);
+      response.attachTraceSummary(result);
     } else {
-      response.appendResponseLine(
-        'There was an unexpected error parsing the trace:',
+      throw new Error(
+        `There was an unexpected error parsing the trace: ${result.error}`,
       );
-      response.appendResponseLine(result.error);
     }
-  } catch (e) {
-    const errorText = e instanceof Error ? e.message : JSON.stringify(e);
-    logger(`Error stopping performance trace: ${errorText}`);
-    response.appendResponseLine(
-      'An error occurred generating the response for this trace:',
-    );
-    response.appendResponseLine(errorText);
   } finally {
     context.setIsRunningPerformanceTrace(false);
   }

tests/McpContext.test.ts

@@ -41,13 +41,13 @@ describe('McpContext', () => {
     });
   });
 
-  it('can store and retrieve performance traces', async () => {
+  it('can store and retrieve the latest performance trace', async () => {
     await withMcpContext(async (_response, context) => {
       const fakeTrace1 = {} as unknown as TraceResult;
       const fakeTrace2 = {} as unknown as TraceResult;
       context.storeTraceRecording(fakeTrace1);
       context.storeTraceRecording(fakeTrace2);
-      assert.deepEqual(context.recordedTraces(), [fakeTrace1, fakeTrace2]);
+      assert.deepEqual(context.recordedTraces(), [fakeTrace2]);
     });
   });
 

tests/McpResponse.test.js.snapshot

@@ -10,6 +10,13 @@ import {tmpdir} from 'node:os';
 import {join} from 'node:path';
 import {describe, it} from 'node:test';
 
+import type {InsightName} from '../src/trace-processing/parse.js';
+import {
+  parseRawTraceBuffer,
+  traceResultIsSuccess,
+} from '../src/trace-processing/parse.js';
+
+import {loadTraceAsBuffer} from './trace-processing/fixtures/load.js';
 import {
   getImageContent,
   getMockAggregatedIssue,
@@ -529,4 +536,74 @@ describe('McpResponse network pagination', () => {
       assert.ok(text.includes('Showing 1-2 of 5 (Page 1 of 3).'));
     });
   });
+
+  describe('trace summaries', () => {
+    it('includes the trace summary text and structured data', async t => {
+      const rawData = loadTraceAsBuffer('web-dev-with-commit.json.gz');
+      const result = await parseRawTraceBuffer(rawData);
+      if (!traceResultIsSuccess(result)) {
+        throw new Error(result.error);
+      }
+
+      await withMcpContext(async (response, context) => {
+        response.attachTraceSummary(result);
+        const {content, structuredContent} = await response.handle(
+          'test',
+          context,
+        );
+
+        t.assert.snapshot?.(getTextContent(content[0]));
+        const typedStructuredContent = structuredContent as {
+          traceSummary?: string;
+          traceInsights?: unknown[];
+        };
+        t.assert.snapshot?.(
+          JSON.stringify(typedStructuredContent.traceSummary, null, 2),
+        );
+        t.assert.snapshot?.(
+          JSON.stringify(typedStructuredContent.traceInsights, null, 2),
+        );
+      });
+    });
+  });
+
+  describe('trace insights', () => {
+    it('includes the trace insight output', async t => {
+      const rawData = loadTraceAsBuffer('web-dev-with-commit.json.gz');
+      const result = await parseRawTraceBuffer(rawData);
+      if (!traceResultIsSuccess(result)) {
+        throw new Error(result.error);
+      }
+
+      await withMcpContext(async (response, context) => {
+        response.attachTraceInsight(
+          result,
+          'NAVIGATION_0',
+          'LCPBreakdown' as InsightName,
+        );
+        const {content} = await response.handle('test', context);
+
+        t.assert.snapshot?.(getTextContent(content[0]));
+      });
+    });
+
+    it('includes error if insight not found', async t => {
+      const rawData = loadTraceAsBuffer('web-dev-with-commit.json.gz');
+      const result = await parseRawTraceBuffer(rawData);
+      if (!traceResultIsSuccess(result)) {
+        throw new Error(result.error);
+      }
+
+      await withMcpContext(async (response, context) => {
+        response.attachTraceInsight(
+          result,
+          'BAD_ID',
+          'LCPBreakdown' as InsightName,
+        );
+        const {content} = await response.handle('test', context);
+
+        t.assert.snapshot?.(getTextContent(content[0]));
+      });
+    });
+  });
 });

tests/McpResponse.test.ts

@@ -194,7 +194,7 @@ describe('performance', () => {
       return result;
     }
 
-    it('returns the information on the insight', async t => {
+    it('returns the information on the insight', async () => {
       const trace = await parseTrace('web-dev-with-commit.json.gz');
       await withMcpContext(async (response, context) => {
         context.storeTraceRecording(trace);
@@ -211,31 +211,7 @@ describe('performance', () => {
           context,
         );
 
-        t.assert.snapshot?.(response.responseLines.join('\n'));
-      });
-    });
-
-    it('returns an error if the insight does not exist', async () => {
-      const trace = await parseTrace('web-dev-with-commit.json.gz');
-      await withMcpContext(async (response, context) => {
-        context.storeTraceRecording(trace);
-        context.setIsRunningPerformanceTrace(false);
-
-        await analyzeInsight.handler(
-          {
-            params: {
-              insightSetId: '8463DF94CD61B265B664E7F768183DE3',
-              insightName: 'MadeUpInsightName',
-            },
-          },
-          response,
-          context,
-        );
-        assert.ok(
-          response.responseLines
-            .join('\n')
-            .match(/No Performance Insights for the given insight set id/),
-        );
+        assert.ok(response.attachedTracedInsight);
       });
     });
 
@@ -295,28 +271,18 @@ describe('performance', () => {
       });
     });
 
-    it('returns an error message if parsing the trace buffer fails', async t => {
+    it('throws an error if parsing the trace buffer fails', async () => {
       await withMcpContext(async (response, context) => {
         context.setIsRunningPerformanceTrace(true);
         const selectedPage = context.getSelectedPage();
         sinon
           .stub(selectedPage.tracing, 'stop')
           .returns(Promise.resolve(undefined));
-        await stopTrace.handler({params: {}}, response, context);
-        t.assert.snapshot?.(response.responseLines.join('\n'));
-      });
-    });
 
-    it('returns the high level summary of the performance trace', async t => {
-      const rawData = loadTraceAsBuffer('web-dev-with-commit.json.gz');
-      await withMcpContext(async (response, context) => {
-        context.setIsRunningPerformanceTrace(true);
-        const selectedPage = context.getSelectedPage();
-        sinon.stub(selectedPage.tracing, 'stop').callsFake(async () => {
-          return rawData;
-        });
-        await stopTrace.handler({params: {}}, response, context);
-        t.assert.snapshot?.(response.responseLines.join('\n'));
+        await assert.rejects(
+          stopTrace.handler({params: {}}, response, context),
+          /There was an unexpected error parsing the trace/,
+        );
       });
     });