fix(tracing): stop recording websocket frames outside of chunks (#41398)

Author: dgozman

packages/playwright-core/src/server/har/harRecorder.ts

@@ -57,8 +57,8 @@ export class HarRecorder implements HarTracerDelegate {
       recordRequestOverrides: true,
       waitForContentOnStop: true,
       urlFilter: urlFilterRe ?? options.urlGlob,
-      omitWebSocketFrames: !!process.env.PLAYWRIGHT_HAR_NO_WEBSOCKET_FRAMES,
     });
+    this._tracer.setOmitWebSocketFrames(!!process.env.PLAYWRIGHT_HAR_NO_WEBSOCKET_FRAMES);
     this._tracer.start({ omitScripts: false });
   }
 

packages/playwright-core/src/server/har/harTracer.ts

@@ -62,7 +62,6 @@ type HarTracerOptions = {
   omitPages?: boolean;
   omitSizes?: boolean;
   omitScripts?: boolean;
-  omitWebSocketFrames?: boolean;
 };
 
 export class HarTracer {
@@ -77,6 +76,7 @@ export class HarTracer {
   private _pageEntrySymbol: symbol;
   private _baseURL: string | undefined;
   private _page: Page | null;
+  private _omitWebSocketFrames = true;
 
   constructor(context: BrowserContext | APIRequestContext, page: Page | null, delegate: HarTracerDelegate, options: HarTracerOptions) {
     this._context = context;
@@ -96,6 +96,10 @@ export class HarTracer {
     this._baseURL = context instanceof APIRequestContext ? context._defaultOptions().baseURL : context._options.baseURL;
   }
 
+  setOmitWebSocketFrames(omitWebSocketFrames: boolean) {
+    this._omitWebSocketFrames = omitWebSocketFrames;
+  }
+
   start(options: { omitScripts: boolean }) {
     if (this._started)
       return;
@@ -443,7 +447,7 @@ export class HarTracer {
 
     let sha1: string | undefined = undefined;
     const recordMessage = (type: 'send' | 'receive', opcode: number, data: string, wallTimeMs: number) => {
-      if (this._options.omitWebSocketFrames)
+      if (this._omitWebSocketFrames)
         return;
       const message = { type, time: this._options.omitTiming ? -1 : wallTimeMs, opcode, data };
       if (this._options.content === 'embed') {

packages/playwright-core/src/server/trace/recorder/tracing.ts

@@ -108,7 +108,6 @@ export class Tracing extends SdkObject implements InstrumentationListener, Snaps
       includeTraceInfo: true,
       recordRequestOverrides: false,
       waitForContentOnStop: false,
-      omitWebSocketFrames: !!process.env.PLAYWRIGHT_TRACING_NO_WEBSOCKET_FRAMES,
     });
     const testIdAttributeName = ('selectors' in context) ? context.selectors().testIdAttributeName() : undefined;
     this._contextCreatedEvent = {
@@ -223,6 +222,7 @@ export class Tracing extends SdkObject implements InstrumentationListener, Snaps
     );
     if (this._state.options.screenshots)
       this._startScreencast();
+    this._harTracer.setOmitWebSocketFrames(!!process.env.PLAYWRIGHT_TRACING_NO_WEBSOCKET_FRAMES);
     if (this._state.options.snapshots)
       await this._snapshotter?.start(progress);
     return { traceName: this._state.traceName };
@@ -394,7 +394,9 @@ export class Tracing extends SdkObject implements InstrumentationListener, Snaps
     eventsHelper.removeEventListeners(this._eventListeners);
     if (this._state.options.screenshots)
       this._stopScreencast();
-
+    // We don't need websocket frames outside of the recording window. This also
+    // stops updating websocket content blobs, which we want to stay unchanged for zipping.
+    this._harTracer.setOmitWebSocketFrames(true);
     if (this._state.options.snapshots)
       this._snapshotter?.stop();
 

tests/library/tracing.spec.ts

@@ -857,6 +857,61 @@ test('should not emit after w/o before', async ({ browserType, mode }, testInfo)
   expect(call2after).toBe(call2before);
 });
 
+test('should save trace while a WebSocket keeps streaming frames', {
+  annotation: { type: 'issue', description: 'https://github.com/microsoft/playwright/issues/41351' }
+}, async ({ context, page, server }, testInfo) => {
+  let streaming = true;
+  server.onceWebSocketConnection(ws => {
+    const timer = setInterval(() => {
+      if (streaming && ws.readyState === ws.OPEN)
+        ws.send('x'.repeat(16 * 1024), () => {});
+    }, 1);
+    const stop = () => clearInterval(timer);
+    ws.on('close', stop);
+    ws.on('error', stop);
+  });
+
+  await context.tracing.start({ snapshots: true });
+
+  await context.tracing.startChunk();
+  await page.goto(server.EMPTY_PAGE);
+  await page.evaluate(url => {
+    (window as any).ws = new WebSocket(url);
+    return new Promise<void>(resolve => (window as any).ws.addEventListener('open', () => resolve()));
+  }, `ws://${server.HOST}/ws`);
+  await page.waitForTimeout(100);
+  const tracePath1 = testInfo.outputPath('trace1.zip');
+  await context.tracing.stopChunk({ path: tracePath1 });
+
+  streaming = false;
+  await context.tracing.startChunk();
+  await page.waitForTimeout(100);
+  const tracePath2 = testInfo.outputPath('trace2.zip');
+  await context.tracing.stopChunk({ path: tracePath2 });
+
+  await page.evaluate(() => new Promise<void>(resolve => {
+    const ws = (window as any).ws as WebSocket;
+    if (ws.readyState === WebSocket.CLOSED) {
+      resolve();
+      return;
+    }
+    ws.addEventListener('close', () => resolve(), { once: true });
+    ws.close();
+  }));
+
+  const webSocketLines = await Promise.all([tracePath1, tracePath2].map(async path => {
+    const { resources } = await parseTraceRaw(path);
+    const websocketResource = Array.from(resources).find(([name, buffer]) => name.endsWith('.jsonl'))!;
+    const lines = websocketResource[1].toString().split('\n').filter(Boolean);
+    expect(lines.length).toBeGreaterThan(0);
+    for (const line of lines)
+      expect(() => JSON.parse(line)).not.toThrow();
+    return { path, lines };
+  }));
+  expect(webSocketLines[0].path).not.toEqual(webSocketLines[1].path);
+  expect(webSocketLines[1].lines).toEqual(webSocketLines[1].lines);
+});
+
 function expectRed(pixels: Buffer, offset: number) {
   const r = pixels.readUInt8(offset);
   const g = pixels.readUInt8(offset + 1);