fix(tracing): copy WebSocket message files before zipping

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

@@ -75,6 +75,7 @@ type RecordingState = {
   chunkOrdinal: number,
   networkSha1s: Set<string>,
   traceSha1s: Set<string>,
+  appendableSha1s: Set<string>,
   recording: boolean;
   callIds: Set<string>;
   groupStack: string[];
@@ -172,6 +173,7 @@ export class Tracing extends SdkObject implements InstrumentationListener, Snaps
       chunkOrdinal: 0,
       traceSha1s: new Set(),
       networkSha1s: new Set(),
+      appendableSha1s: new Set(),
       recording: false,
       callIds: new Set(),
       groupStack: [],
@@ -411,8 +413,15 @@ export class Tracing extends SdkObject implements InstrumentationListener, Snaps
     const entries: NameValue[] = [];
     entries.push({ name: 'trace.trace', value: this._state.traceFile });
     entries.push({ name: 'trace.network', value: newNetworkFile });
-    for (const sha1 of new Set([...this._state.traceSha1s, ...this._state.networkSha1s]))
-      entries.push({ name: path.join('resources', sha1), value: path.join(this._state.resourcesDir, sha1) });
+    for (const sha1 of new Set([...this._state.traceSha1s, ...this._state.networkSha1s])) {
+      let value = path.join(this._state.resourcesDir, sha1);
+      if (params.mode === 'entries' && this._state.appendableSha1s.has(sha1)) {
+        const copy = path.join(this._state.tracesDir, `${this._state.traceName}-pwnetcopy-${this._state.chunkOrdinal}-${sha1}`);
+        this._fs.copyFile(value, copy);
+        value = copy;
+      }
+      entries.push({ name: path.join('resources', sha1), value });
+    }
 
     // Only reset trace sha1s, network resources are preserved between chunks.
     this._state.traceSha1s = new Set();
@@ -543,6 +552,10 @@ export class Tracing extends SdkObject implements InstrumentationListener, Snaps
     const event: trace.ResourceSnapshotTraceEvent = { type: 'resource-snapshot', snapshot: entry };
     const visited = visitTraceEvent(event, this._state!.networkSha1s);
     this._fs.appendFile(this._state!.networkFile, JSON.stringify(visited) + '\n', true /* flush */);
+
+    const sha1 = entry.response.content._sha1;
+    if (sha1)
+      this._state!.appendableSha1s.delete(sha1);
   }
 
   flushHarEntries() {
@@ -562,8 +575,8 @@ export class Tracing extends SdkObject implements InstrumentationListener, Snaps
   }
 
   onContentBlobAppend(sha1: string, text: string) {
-    if (!this._allResources.has(sha1))
-      this._allResources.add(sha1);
+    this._allResources.add(sha1);
+    this._state!.appendableSha1s.add(sha1);
     this._fs.appendFile(path.join(this._state!.resourcesDir, sha1), text, this._state!.options.live /* flush */);
   }
 

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);