fix(tracing): include source for tracing.group(location) (#39307)

Author: Skn0tt

packages/playwright-core/src/client/tracing.ts

@@ -22,6 +22,7 @@ import type * as channels from '@protocol/channels';
 
 export class Tracing extends ChannelOwner<channels.TracingChannel> implements api.Tracing {
   private _includeSources = false;
+  private _additionalSources = new Set<string>();
   private _isLive = false;
   _tracesDir: string | undefined;
   private _stacksId: string | undefined;
@@ -58,6 +59,8 @@ export class Tracing extends ChannelOwner<channels.TracingChannel> implements ap
   }
 
   async group(name: string, options: { location?: { file: string, line?: number, column?: number } } = {}) {
+    if (options.location)
+      this._additionalSources.add(options.location.file);
     await this._channel.tracingGroup({ name, location: options.location });
   }
 
@@ -90,6 +93,9 @@ export class Tracing extends ChannelOwner<channels.TracingChannel> implements ap
   private async _doStopChunk(filePath: string | undefined) {
     this._resetStackCounter();
 
+    const additionalSources = [...this._additionalSources];
+    this._additionalSources.clear();
+
     if (!filePath) {
       // Not interested in artifacts.
       await this._channel.tracingStopChunk({ mode: 'discard' });
@@ -106,7 +112,7 @@ export class Tracing extends ChannelOwner<channels.TracingChannel> implements ap
 
     if (isLocal) {
       const result = await this._channel.tracingStopChunk({ mode: 'entries' });
-      await localUtils.zip({ zipFile: filePath, entries: result.entries!, mode: 'write', stacksId: this._stacksId, includeSources: this._includeSources });
+      await localUtils.zip({ zipFile: filePath, entries: result.entries!, mode: 'write', stacksId: this._stacksId, includeSources: this._includeSources, additionalSources });
       return;
     }
 
@@ -124,7 +130,7 @@ export class Tracing extends ChannelOwner<channels.TracingChannel> implements ap
     await artifact.saveAs(filePath);
     await artifact.delete();
 
-    await localUtils.zip({ zipFile: filePath, entries: [], mode: 'append', stacksId: this._stacksId, includeSources: this._includeSources });
+    await localUtils.zip({ zipFile: filePath, entries: [], mode: 'append', stacksId: this._stacksId, includeSources: this._includeSources, additionalSources });
   }
 
   _resetStackCounter() {

packages/playwright-core/src/protocol/validator.ts

@@ -278,6 +278,7 @@ scheme.LocalUtilsZipParams = tObject({
   stacksId: tOptional(tString),
   mode: tEnum(['write', 'append']),
   includeSources: tBoolean,
+  additionalSources: tOptional(tArray(tString)),
 });
 scheme.LocalUtilsZipResult = tOptional(tObject({}));
 scheme.LocalUtilsHarOpenParams = tObject({

packages/playwright-core/src/server/localUtils.ts

@@ -67,15 +67,15 @@ export async function zip(progress: Progress, stackSessions: Map<string, StackSe
 
   // Collect sources from stacks.
   if (params.includeSources) {
-    const sourceFiles = new Set<string>();
+    const sourceFiles = new Set<string>(params.additionalSources);
     for (const { stack } of stackSession?.callStacks || []) {
       if (!stack)
         continue;
       for (const { file } of stack)
         sourceFiles.add(file);
     }
     for (const sourceFile of sourceFiles)
-      addFile(sourceFile, 'resources/src@' + await calculateSha1(sourceFile) + '.txt');
+      addFile(sourceFile, 'resources/src@' + calculateSha1(sourceFile) + '.txt');
   }
 
   if (params.mode === 'write') {

packages/protocol/src/channels.d.ts

@@ -477,9 +477,11 @@ export type LocalUtilsZipParams = {
   stacksId?: string,
   mode: 'write' | 'append',
   includeSources: boolean,
+  additionalSources?: string[],
 };
 export type LocalUtilsZipOptions = {
   stacksId?: string,
+  additionalSources?: string[],
 };
 export type LocalUtilsZipResult = void;
 export type LocalUtilsHarOpenParams = {

packages/protocol/src/protocol.yml

@@ -652,6 +652,9 @@ LocalUtils:
             - write
             - append
         includeSources: boolean
+        additionalSources:
+          type: array?
+          items: string
 
     harOpen:
       internal: true

tests/library/trace-viewer.spec.ts

@@ -113,6 +113,10 @@ test('should open trace viewer on specific host', async ({ showTraceViewer }, te
 
 test('should show tracing.group in the action list with location', async ({ runAndTrace, page, context }) => {
   test.info().annotations.push({ type: 'issue', description: 'https://github.com/microsoft/playwright/issues/36483' });
+  test.info().annotations.push({ type: 'issue', description: 'https://github.com/microsoft/playwright/issues/39302' });
+
+  const sourceFile = test.info().outputPath('source.js');
+  await fs.promises.writeFile(sourceFile, 'buddy beaver');
 
   const traceViewer = await test.step('create trace with groups', async () => {
     await page.context().tracing.group('ignored group');
@@ -125,15 +129,21 @@ test('should show tracing.group in the action list with location', async ({ runA
       await context.tracing.group('inner group 2');
       await expect(page.getByText('Hello')).toBeVisible();
       await context.tracing.groupEnd();
+      await context.tracing.group('inner group 3', { location: { file: sourceFile, line: 1, column: 1 } });
+      await expect(page.getByText('Hello')).toBeVisible();
+      await context.tracing.groupEnd();
       await context.tracing.groupEnd();
     });
   });
 
+  await fs.promises.rm(sourceFile);
+
   await expect(traceViewer.actionTitles).toHaveText([
     /outer group/,
     /Navigate/,
     /inner group 1 {{ eager_beaver }}/,
     /inner group 2/,
+    /inner group 3/,
     /toBeVisible/,
   ]);
 
@@ -145,12 +155,16 @@ test('should show tracing.group in the action list with location', async ({ runA
     /inner group 1 {{ eager_beaver }}/,
     /Click.*locator/,
     /inner group 2/,
+    /inner group 3/,
   ]);
   await traceViewer.showSourceTab();
   await expect(traceViewer.sourceCodeTab.locator('.source-line-running')).toHaveText(/DO NOT TOUCH THIS LINE/);
 
   await traceViewer.selectAction('inner group 2');
   await expect(traceViewer.sourceCodeTab.locator('.source-line-running')).toContainText("await context.tracing.group('inner group 2');");
+
+  await traceViewer.selectAction('inner group 3');
+  await expect(traceViewer.sourceCodeTab.locator('.source-line-running')).toContainText('buddy beaver');
 });
 
 test('should open simple trace viewer', async ({ showTraceViewer }) => {