feat: Implement network log preservation features and related tools

Author: nicetrykiddo

src/McpContext.ts

@@ -19,7 +19,7 @@ import type {
   PredefinedNetworkConditions,
 } from 'puppeteer-core';
 
-import {NetworkCollector, PageCollector} from './PageCollector.js';
+import {NetworkCollector, PageCollector, PreservedNetworkRequest} from './PageCollector.js';
 import {listPages} from './tools/pages.js';
 import {takeSnapshot} from './tools/snapshot.js';
 import {CLOSE_PAGE_ERROR} from './tools/ToolDefinition.js';
@@ -134,6 +134,32 @@ export class McpContext implements Context {
     return this.#networkCollector.getData(page);
   }
 
+  getPreservedNetworkRequests(): PreservedNetworkRequest[] {
+    const page = this.getSelectedPage();
+    return this.#networkCollector.getPreservedData(page);
+  }
+
+  enableNetworkLogPreservation(options?: {
+    includeRequestBodies?: boolean;
+    includeResponseBodies?: boolean;
+    maxRequests?: number;
+  }): void {
+    this.#networkCollector.enablePreservation(options);
+  }
+
+  disableNetworkLogPreservation(): void {
+    this.#networkCollector.disablePreservation();
+  }
+
+  isNetworkLogPreservationEnabled(): boolean {
+    return this.#networkCollector.isPreservationEnabled();
+  }
+
+  clearPreservedNetworkLogs(): void {
+    const page = this.getSelectedPage();
+    this.#networkCollector.clearPreservedData(page);
+  }
+
   getConsoleData(): Array<ConsoleMessage | Error> {
     const page = this.getSelectedPage();
     return this.#consoleCollector.getData(page);

src/PageCollector.ts

@@ -4,16 +4,18 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type {Browser, HTTPRequest, Page} from 'puppeteer-core';
+import type {Browser, HTTPRequest, Page, HTTPResponse} from 'puppeteer-core';
+
+export interface PreservedNetworkRequest {
+  request: HTTPRequest;
+  timestamp: number;
+  requestBody?: string;
+  responseBody?: string;
+}
 
 export class PageCollector<T> {
   #browser: Browser;
   #initializer: (page: Page, collector: (item: T) => void) => void;
-  /**
-   * The Array in this map should only be set once
-   * As we use the reference to it.
-   * Use methods that manipulate the array in place.
-   */
   protected storage = new WeakMap<Page, T[]>();
 
   constructor(
@@ -24,6 +26,10 @@ export class PageCollector<T> {
     this.#initializer = initializer;
   }
 
+  protected getBrowser(): Browser {
+    return this.#browser;
+  }
+
   async init() {
     const pages = await this.#browser.pages();
     for (const page of pages) {
@@ -77,7 +83,86 @@ export class PageCollector<T> {
 }
 
 export class NetworkCollector extends PageCollector<HTTPRequest> {
+  #preservationEnabled = false;
+  #includeRequestBodies = true;
+  #includeResponseBodies = true;
+  #maxRequests?: number;
+  #preservedData = new WeakMap<Page, PreservedNetworkRequest[]>();
+
+  enablePreservation(options?: {
+    includeRequestBodies?: boolean;
+    includeResponseBodies?: boolean;
+    maxRequests?: number;
+  }): void {
+    this.#preservationEnabled = true;
+    this.#includeRequestBodies = options?.includeRequestBodies ?? true;
+    this.#includeResponseBodies = options?.includeResponseBodies ?? true;
+    this.#maxRequests = options?.maxRequests;
+  }
+
+  disablePreservation(): void {
+    this.#preservationEnabled = false;
+  }
+
+  isPreservationEnabled(): boolean {
+    return this.#preservationEnabled;
+  }
+
+  clearPreservedData(page: Page): void {
+    const preserved = this.#preservedData.get(page);
+    if (preserved) {
+      preserved.length = 0;
+    }
+  }
+
+  getPreservedData(page: Page): PreservedNetworkRequest[] {
+    return this.#preservedData.get(page) ?? [];
+  }
+
+  async #captureRequestData(request: HTTPRequest): Promise<PreservedNetworkRequest> {
+    const preserved: PreservedNetworkRequest = {
+      request,
+      timestamp: Date.now(),
+    };
+
+    if (this.#includeRequestBodies) {
+      try {
+        const postData = request.postData();
+        if (postData) {
+          preserved.requestBody = postData;
+        }
+      } catch (error) {
+      }
+    }
+
+    if (this.#includeResponseBodies) {
+      try {
+        const response = request.response();
+        if (response) {
+          const buffer = await response.buffer();
+          const contentType = response.headers()['content-type'] || '';
+          
+          if (contentType.includes('text/') || 
+              contentType.includes('application/json') ||
+              contentType.includes('application/xml') ||
+              contentType.includes('application/javascript')) {
+            preserved.responseBody = buffer.toString('utf-8');
+          } else {
+            preserved.responseBody = `[Binary data: ${contentType}, ${buffer.length} bytes]`;
+          }
+        }
+      } catch (error) {
+      }
+    }
+
+    return preserved;
+  }
+
   override cleanup(page: Page) {
+    if (this.#preservationEnabled) {
+      return;
+    }
+
     const requests = this.storage.get(page) ?? [];
     if (!requests) {
       return;
@@ -87,9 +172,41 @@ export class NetworkCollector extends PageCollector<HTTPRequest> {
         ? request.isNavigationRequest()
         : false;
     });
-    // Keep all requests since the last navigation request including that
-    // navigation request itself.
-    // Keep the reference
     requests.splice(0, Math.max(lastRequestIdx, 0));
   }
+
+  public override addPage(page: Page): void {
+    super.addPage(page);
+    
+    if (this.#preservationEnabled) {
+      if (!this.#preservedData.has(page)) {
+        this.#preservedData.set(page, []);
+      }
+
+      page.on('requestfinished', async (request: HTTPRequest) => {
+        const preserved = this.#preservedData.get(page);
+        if (!preserved) return;
+
+        const data = await this.#captureRequestData(request);
+        preserved.push(data);
+
+        if (this.#maxRequests && preserved.length > this.#maxRequests) {
+          preserved.shift();
+        }
+      });
+    }
+  }
+
+  override async init() {
+    await super.init();
+    
+    if (this.#preservationEnabled) {
+      const pages = await this.getBrowser().pages();
+      for (const page of pages) {
+        if (!this.#preservedData.has(page)) {
+          this.#preservedData.set(page, []);
+        }
+      }
+    }
+  }
 }

src/tools/ToolDefinition.ts

@@ -7,6 +7,7 @@
 import type {Dialog, ElementHandle, Page} from 'puppeteer-core';
 import z from 'zod';
 
+import type {PreservedNetworkRequest} from '../PageCollector.js';
 import type {TraceResult} from '../trace-processing/parse.js';
 
 import type {ToolCategories} from './categories.js';
@@ -79,6 +80,15 @@ export type Context = Readonly<{
     filename: string,
   ): Promise<{filename: string}>;
   waitForEventsAfterAction(action: () => Promise<unknown>): Promise<void>;
+  enableNetworkLogPreservation(options?: {
+    includeRequestBodies?: boolean;
+    includeResponseBodies?: boolean;
+    maxRequests?: number;
+  }): void;
+  disableNetworkLogPreservation(): void;
+  isNetworkLogPreservationEnabled(): boolean;
+  clearPreservedNetworkLogs(): void;
+  getPreservedNetworkRequests(): PreservedNetworkRequest[];
 }>;
 
 export function defineTool<Schema extends z.ZodRawShape>(

src/tools/network.ts

@@ -86,3 +86,101 @@ export const getNetworkRequest = defineTool({
     response.attachNetworkRequest(request.params.url);
   },
 });
+
+export const enableNetworkLogPreservation = defineTool({
+  name: 'enable_network_log_preservation',
+  description: `Enable network log preservation mode. When enabled, all network requests are preserved across page navigations and response bodies are automatically captured. By default, logs are cleaned on navigation for performance.`,
+  annotations: {
+    category: ToolCategories.NETWORK,
+    readOnlyHint: false,
+  },
+  schema: {
+    includeRequestBodies: z
+      .boolean()
+      .optional()
+      .default(true)
+      .describe('Whether to capture and cache request bodies. Default: true'),
+    includeResponseBodies: z
+      .boolean()
+      .optional()
+      .default(true)
+      .describe('Whether to capture and cache response bodies. Default: true'),
+    maxRequests: z
+      .number()
+      .int()
+      .positive()
+      .optional()
+      .describe(
+        'Maximum number of requests to preserve. Older requests are automatically removed when limit is reached. When omitted, no limit is applied.',
+      ),
+  },
+  handler: async (request, response, context) => {
+    context.enableNetworkLogPreservation({
+      includeRequestBodies: request.params.includeRequestBodies,
+      includeResponseBodies: request.params.includeResponseBodies,
+      maxRequests: request.params.maxRequests,
+    });
+    response.appendResponseLine(
+      'Network log preservation enabled. All network requests will be preserved across navigations.',
+    );
+    if (request.params.includeRequestBodies) {
+      response.appendResponseLine('Request bodies will be captured.');
+    }
+    if (request.params.includeResponseBodies) {
+      response.appendResponseLine('Response bodies will be captured.');
+    }
+    if (request.params.maxRequests) {
+      response.appendResponseLine(
+        `Maximum ${request.params.maxRequests} requests will be preserved.`,
+      );
+    }
+  },
+});
+
+export const disableNetworkLogPreservation = defineTool({
+  name: 'disable_network_log_preservation',
+  description: `Disable network log preservation mode and optionally clear existing preserved logs. After disabling, network logs will be cleaned on navigation (default behavior).`,
+  annotations: {
+    category: ToolCategories.NETWORK,
+    readOnlyHint: false,
+  },
+  schema: {
+    clearExisting: z
+      .boolean()
+      .optional()
+      .default(true)
+      .describe(
+        'Whether to clear existing preserved logs. Default: true',
+      ),
+  },
+  handler: async (request, response, context) => {
+    context.disableNetworkLogPreservation();
+    if (request.params.clearExisting) {
+      context.clearPreservedNetworkLogs();
+      response.appendResponseLine(
+        'Network log preservation disabled and existing logs cleared.',
+      );
+    } else {
+      response.appendResponseLine(
+        'Network log preservation disabled. Existing logs retained.',
+      );
+    }
+  },
+});
+
+export const clearPreservedNetworkLogs = defineTool({
+  name: 'clear_preserved_network_logs',
+  description: `Clear all preserved network logs for the currently selected page. This does not disable preservation mode.`,
+  annotations: {
+    category: ToolCategories.NETWORK,
+    readOnlyHint: false,
+  },
+  schema: {},
+  handler: async (_request, response, context) => {
+    const preservedCount = context.getPreservedNetworkRequests().length;
+    context.clearPreservedNetworkLogs();
+    response.appendResponseLine(
+      `Cleared ${preservedCount} preserved network request(s).`,
+    );
+  },
+});