chore: move artifacts recording to TestLifecycleInstrumentation (#30935)

The spirit of this change is reverting #23153. Since that time, we have
moved tracing and `artifactsDir` lifetime into the test runner, so the
reason for revert is mitigated.

Fixes #30287, fixes #30718, fixes #30959.

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

@@ -34,8 +34,8 @@ export class Tracing extends ChannelOwner<channels.TracingChannel> implements ap
   }
 
   async start(options: { name?: string, title?: string, snapshots?: boolean, screenshots?: boolean, sources?: boolean, _live?: boolean } = {}) {
-    this._includeSources = !!options.sources;
+    await this._wrapApiCall(async () => {
-    const traceName = await this._wrapApiCall(async () => {
+      this._includeSources = !!options.sources;
       await this._channel.tracingStart({
         name: options.name,
         snapshots: options.snapshots,
@@ -43,14 +43,15 @@ export class Tracing extends ChannelOwner<channels.TracingChannel> implements ap
         live: options._live,
       });
       const response = await this._channel.tracingStartChunk({ name: options.name, title: options.title });
-      return response.traceName;
+      await this._startCollectingStacks(response.traceName);
     }, true);
-    await this._startCollectingStacks(traceName);
   }
 
   async startChunk(options: { name?: string, title?: string } = {}) {
-    const { traceName } = await this._channel.tracingStartChunk(options);
+    await this._wrapApiCall(async () => {
-    await this._startCollectingStacks(traceName);
+      const { traceName } = await this._channel.tracingStartChunk(options);
+      await this._startCollectingStacks(traceName);
+    }, true);
   }
 
   private async _startCollectingStacks(traceName: string) {

packages/playwright/src/common/globals.ts

@@ -42,3 +42,19 @@ export function setIsWorkerProcess() {
 export function isWorkerProcess() {
   return _isWorkerProcess;
 }
+
+export interface TestLifecycleInstrumentation {
+  onTestBegin?(): Promise<void>;
+  onTestFunctionEnd?(): Promise<void>;
+  onTestEnd?(): Promise<void>;
+}
+
+let _testLifecycleInstrumentation: TestLifecycleInstrumentation | undefined;
+
+export function setTestLifecycleInstrumentation(instrumentation: TestLifecycleInstrumentation | undefined) {
+  _testLifecycleInstrumentation = instrumentation;
+}
+
+export function testLifecycleInstrumentation() {
+  return _testLifecycleInstrumentation;
+}

packages/playwright/src/index.ts

@@ -16,15 +16,14 @@
 
 import * as fs from 'fs';
 import * as path from 'path';
-import type { APIRequestContext, BrowserContext, Browser, BrowserContextOptions, LaunchOptions, Page, Tracing, Video } from 'playwright-core';
+import type { APIRequestContext, BrowserContext, Browser, BrowserContextOptions, LaunchOptions, Page, Tracing, Video, PageScreenshotOptions } from 'playwright-core';
-import * as playwrightLibrary from 'playwright-core';
 import { createGuid, debugMode, addInternalStackPrefix, isString, asLocator, jsonStringifyForceASCII } from 'playwright-core/lib/utils';
 import type { Fixtures, PlaywrightTestArgs, PlaywrightTestOptions, PlaywrightWorkerArgs, PlaywrightWorkerOptions, ScreenshotMode, TestInfo, TestType, VideoMode } from '../types/test';
 import type { TestInfoImpl } from './worker/testInfo';
 import { rootTestType } from './common/testType';
 import type { ContextReuseMode } from './common/config';
 import type { ClientInstrumentation, ClientInstrumentationListener } from '../../playwright-core/src/client/clientInstrumentation';
-import { currentTestInfo } from './common/globals';
+import { currentTestInfo, setTestLifecycleInstrumentation, type TestLifecycleInstrumentation } from './common/globals';
 export { expect } from './matchers/expect';
 export const _baseTest: TestType<{}, {}> = rootTestType.test;
 
@@ -45,11 +44,12 @@ if ((process as any)['__pw_initiator__']) {
 type TestFixtures = PlaywrightTestArgs & PlaywrightTestOptions & {
   _combinedContextOptions: BrowserContextOptions,
   _setupContextOptions: void;
-  _setupArtifacts: void;
   _contextFactory: (options?: BrowserContextOptions) => Promise<BrowserContext>;
 };
 
 type WorkerFixtures = PlaywrightWorkerArgs & PlaywrightWorkerOptions & {
+  // Same as "playwright", but exposed so that our internal tests can override it.
+  _playwrightImpl: PlaywrightWorkerArgs['playwright'];
   _browserOptions: LaunchOptions;
   _optionContextReuseMode: ContextReuseMode,
   _optionConnectOptions: PlaywrightWorkerOptions['connectOptions'],
@@ -59,9 +59,14 @@ type WorkerFixtures = PlaywrightWorkerArgs & PlaywrightWorkerOptions & {
 const playwrightFixtures: Fixtures<TestFixtures, WorkerFixtures> = ({
   defaultBrowserType: ['chromium', { scope: 'worker', option: true }],
   browserName: [({ defaultBrowserType }, use) => use(defaultBrowserType), { scope: 'worker', option: true }],
-  playwright: [async ({}, use) => {
+  _playwrightImpl: [({}, use) => use(require('playwright-core')), { scope: 'worker' }],
-    await use(require('playwright-core'));
+
+  playwright: [async ({ _playwrightImpl, screenshot }, use) => {
+    await connector.setPlaywright(_playwrightImpl, screenshot);
+    await use(_playwrightImpl);
+    await connector.setPlaywright(undefined, screenshot);
   }, { scope: 'worker', _hideStep: true } as any],
+
   headless: [({ launchOptions }, use) => use(launchOptions.headless ?? true), { scope: 'worker', option: true }],
   channel: [({ launchOptions }, use) => use(launchOptions.channel), { scope: 'worker', option: true }],
   launchOptions: [{}, { scope: 'worker', option: true }],
@@ -222,7 +227,7 @@ const playwrightFixtures: Fixtures<TestFixtures, WorkerFixtures> = ({
 
   _setupContextOptions: [async ({ playwright, _combinedContextOptions, actionTimeout, navigationTimeout, testIdAttribute }, use, testInfo) => {
     if (testIdAttribute)
-      playwrightLibrary.selectors.setTestIdAttribute(testIdAttribute);
+      playwright.selectors.setTestIdAttribute(testIdAttribute);
     testInfo.snapshotSuffix = process.platform;
     if (debugMode())
       testInfo.setTimeout(0);
@@ -243,58 +248,6 @@ const playwrightFixtures: Fixtures<TestFixtures, WorkerFixtures> = ({
     }
   }, { auto: 'all-hooks-included',  _title: 'context configuration' } as any],
 
-  _setupArtifacts: [async ({ playwright, screenshot }, use, testInfo) => {
-    const artifactsRecorder = new ArtifactsRecorder(playwright, tracing().artifactsDir(), screenshot);
-    await artifactsRecorder.willStartTest(testInfo as TestInfoImpl);
-    const csiListener: ClientInstrumentationListener = {
-      onApiCallBegin: (apiName: string, params: Record<string, any>, frames: StackFrame[], userData: any, out: { stepId?: string }) => {
-        const testInfo = currentTestInfo();
-        if (!testInfo || apiName.includes('setTestIdAttribute'))
-          return { userObject: null };
-        const step = testInfo._addStep({
-          location: frames[0] as any,
-          category: 'pw:api',
-          title: renderApiCall(apiName, params),
-          apiName,
-          params,
-        });
-        userData.userObject = step;
-        out.stepId = step.stepId;
-      },
-      onApiCallEnd: (userData: any, error?: Error) => {
-        const step = userData.userObject;
-        step?.complete({ error });
-      },
-      onWillPause: () => {
-        currentTestInfo()?.setTimeout(0);
-      },
-      runAfterCreateBrowserContext: async (context: BrowserContext) => {
-        await artifactsRecorder?.didCreateBrowserContext(context);
-        const testInfo = currentTestInfo();
-        if (testInfo)
-          attachConnectedHeaderIfNeeded(testInfo, context.browser());
-      },
-      runAfterCreateRequestContext: async (context: APIRequestContext) => {
-        await artifactsRecorder?.didCreateRequestContext(context);
-      },
-      runBeforeCloseBrowserContext: async (context: BrowserContext) => {
-        await artifactsRecorder?.willCloseBrowserContext(context);
-      },
-      runBeforeCloseRequestContext: async (context: APIRequestContext) => {
-        await artifactsRecorder?.willCloseRequestContext(context);
-      },
-    };
-
-    const clientInstrumentation = (playwright as any)._instrumentation as ClientInstrumentation;
-    clientInstrumentation.addListener(csiListener);
-
-    await use();
-
-    clientInstrumentation.removeListener(csiListener);
-    await artifactsRecorder.didFinishTest();
-
-  }, { auto: 'all-hooks-included',  _title: 'trace recording' } as any],
-
   _contextFactory: [async ({ browser, video, _reuseContext }, use, testInfo) => {
     const testInfoImpl = testInfo as TestInfoImpl;
     const videoMode = normalizeVideoMode(video);
@@ -471,7 +424,7 @@ class ArtifactsRecorder {
   private _playwright: Playwright;
   private _artifactsDir: string;
   private _screenshotMode: ScreenshotMode;
-  private _screenshotOptions: { mode: ScreenshotMode } & Pick<playwrightLibrary.PageScreenshotOptions, 'fullPage' | 'omitBackground'> | undefined;
+  private _screenshotOptions: { mode: ScreenshotMode } & Pick<PageScreenshotOptions, 'fullPage' | 'omitBackground'> | undefined;
   private _temporaryScreenshots: string[] = [];
   private _temporaryArtifacts: string[] = [];
   private _reusedContexts = new Set<BrowserContext>();
@@ -496,7 +449,6 @@ class ArtifactsRecorder {
 
   async willStartTest(testInfo: TestInfoImpl) {
     this._testInfo = testInfo;
-    testInfo._onDidFinishTestFunction = () => this.didFinishTestFunction();
 
     // Since beforeAll(s), test and afterAll(s) reuse the same TestInfo, make sure we do not
     // overwrite previous screenshots.
@@ -678,6 +630,101 @@ function tracing() {
   return (test.info() as TestInfoImpl)._tracing;
 }
 
+class InstrumentationConnector implements TestLifecycleInstrumentation, ClientInstrumentationListener {
+  private _playwright: PlaywrightWorkerArgs['playwright'] | undefined;
+  private _screenshot: ScreenshotOption = 'off';
+  private _artifactsRecorder: ArtifactsRecorder | undefined;
+  private _testIsRunning = false;
+
+  constructor() {
+    setTestLifecycleInstrumentation(this);
+  }
+
+  async setPlaywright(playwright: PlaywrightWorkerArgs['playwright'] | undefined, screenshot: ScreenshotOption) {
+    if (this._playwright) {
+      if (this._testIsRunning) {
+        // When "playwright" is destroyed during a test, collect artifacts immediately.
+        await this.onTestEnd();
+      }
+      const clientInstrumentation = (this._playwright as any)._instrumentation as ClientInstrumentation;
+      clientInstrumentation.removeListener(this);
+    }
+    this._playwright = playwright;
+    this._screenshot = screenshot;
+    if (this._playwright) {
+      const clientInstrumentation = (this._playwright as any)._instrumentation as ClientInstrumentation;
+      clientInstrumentation.addListener(this);
+      if (this._testIsRunning) {
+        // When "playwright" is created during a test, wire it up immediately.
+        await this.onTestBegin();
+      }
+    }
+  }
+
+  async onTestBegin() {
+    this._testIsRunning = true;
+    if (this._playwright) {
+      this._artifactsRecorder = new ArtifactsRecorder(this._playwright, tracing().artifactsDir(), this._screenshot);
+      await this._artifactsRecorder.willStartTest(currentTestInfo() as TestInfoImpl);
+    }
+  }
+
+  async onTestFunctionEnd() {
+    await this._artifactsRecorder?.didFinishTestFunction();
+  }
+
+  async onTestEnd() {
+    await this._artifactsRecorder?.didFinishTest();
+    this._artifactsRecorder = undefined;
+    this._testIsRunning = false;
+  }
+
+  onApiCallBegin(apiName: string, params: Record<string, any>, frames: StackFrame[], userData: any, out: { stepId?: string }) {
+    const testInfo = currentTestInfo();
+    if (!testInfo || apiName.includes('setTestIdAttribute'))
+      return { userObject: null };
+    const step = testInfo._addStep({
+      location: frames[0] as any,
+      category: 'pw:api',
+      title: renderApiCall(apiName, params),
+      apiName,
+      params,
+    });
+    userData.userObject = step;
+    out.stepId = step.stepId;
+  }
+
+  onApiCallEnd(userData: any, error?: Error) {
+    const step = userData.userObject;
+    step?.complete({ error });
+  }
+
+  onWillPause() {
+    currentTestInfo()?.setTimeout(0);
+  }
+
+  async runAfterCreateBrowserContext(context: BrowserContext) {
+    await this._artifactsRecorder?.didCreateBrowserContext(context);
+    const testInfo = currentTestInfo();
+    if (testInfo)
+      attachConnectedHeaderIfNeeded(testInfo, context.browser());
+  }
+
+  async runAfterCreateRequestContext(context: APIRequestContext) {
+    await this._artifactsRecorder?.didCreateRequestContext(context);
+  }
+
+  async runBeforeCloseBrowserContext(context: BrowserContext) {
+    await this._artifactsRecorder?.willCloseBrowserContext(context);
+  }
+
+  async runBeforeCloseRequestContext(context: APIRequestContext) {
+    await this._artifactsRecorder?.willCloseRequestContext(context);
+  }
+}
+
+const connector = new InstrumentationConnector();
+
 export const test = _baseTest.extend<TestFixtures, WorkerFixtures>(playwrightFixtures);
 
 export { defineConfig } from './common/configLoader';

packages/playwright/src/worker/testInfo.ts

@@ -68,7 +68,6 @@ export class TestInfoImpl implements TestInfo {
   readonly _projectInternal: FullProjectInternal;
   readonly _configInternal: FullConfigInternal;
   private readonly _steps: TestStepInternal[] = [];
-  _onDidFinishTestFunction: (() => Promise<void>) | undefined;
   private readonly _stages: TestStage[] = [];
   _hasNonRetriableError = false;
   _hasUnhandledError = false;

packages/playwright/src/worker/workerMain.ts

@@ -17,7 +17,7 @@
 import { colors } from 'playwright-core/lib/utilsBundle';
 import { debugTest, relativeFilePath, serializeError } from '../util';
 import { type TestBeginPayload, type TestEndPayload, type RunPayload, type DonePayload, type WorkerInitParams, type TeardownErrorsPayload, stdioChunkToParams } from '../common/ipc';
-import { setCurrentTestInfo, setIsWorkerProcess } from '../common/globals';
+import { setCurrentTestInfo, setIsWorkerProcess, testLifecycleInstrumentation } from '../common/globals';
 import { deserializeConfig } from '../common/configLoader';
 import type { Suite, TestCase } from '../common/test';
 import type { Annotation, FullConfigInternal, FullProjectInternal } from '../common/config';
@@ -304,10 +304,11 @@ export class WorkerMain extends ProcessRunner {
     if (this._lastRunningTests.length > 10)
       this._lastRunningTests.shift();
     let shouldRunAfterEachHooks = false;
+    const tracingSlot = { timeout: this._project.project.timeout, elapsed: 0 };
 
     testInfo._allowSkips = true;
     await testInfo._runAsStage({ title: 'setup and test' }, async () => {
-      await testInfo._runAsStage({ title: 'start tracing', runnable: { type: 'test' } }, async () => {
+      await testInfo._runAsStage({ title: 'start tracing', runnable: { type: 'test', slot: tracingSlot } }, async () => {
         // Ideally, "trace" would be an config-level option belonging to the
         // test runner instead of a fixture belonging to Playwright.
         // However, for backwards compatibility, we have to read it from a fixture today.
@@ -318,6 +319,7 @@ export class WorkerMain extends ProcessRunner {
         if (typeof traceFixtureRegistration.fn === 'function')
           throw new Error(`"trace" option cannot be a function`);
         await testInfo._tracing.startIfNeeded(traceFixtureRegistration.fn);
+        await testLifecycleInstrumentation()?.onTestBegin?.();
       });
 
       if (this._isStopped || isSkipped) {
@@ -372,10 +374,10 @@ export class WorkerMain extends ProcessRunner {
 
       try {
         // Run "immediately upon test function finish" callback.
-        await testInfo._runAsStage({ title: 'on-test-function-finish', runnable: { type: 'test', slot: afterHooksSlot } }, async () => testInfo._onDidFinishTestFunction?.());
+        await testInfo._runAsStage({ title: 'on-test-function-finish', runnable: { type: 'test', slot: tracingSlot } }, async () => {
+          await testLifecycleInstrumentation()?.onTestFunctionEnd?.();
+        });
       } catch (error) {
-        if (error instanceof TimeoutManagerError)
-          didTimeoutInAfterHooks = true;
         firstAfterHooksError = firstAfterHooksError ?? error;
       }
 
@@ -458,8 +460,8 @@ export class WorkerMain extends ProcessRunner {
       }).catch(() => {});  // Ignore the top-level error, it is already inside TestInfo.errors.
     }
 
-    const tracingSlot = { timeout: this._project.project.timeout, elapsed: 0 };
     await testInfo._runAsStage({ title: 'stop tracing', runnable: { type: 'test', slot: tracingSlot } }, async () => {
+      await testLifecycleInstrumentation()?.onTestEnd?.();
       await testInfo._tracing.stopIfNeeded();
     }).catch(() => {});  // Ignore the top-level error, it is already inside TestInfo.errors.
 

tests/config/testModeFixtures.ts

@@ -30,11 +30,12 @@ export type TestModeTestFixtures = {
 export type TestModeWorkerFixtures = {
   toImplInWorkerScope: (rpcObject?: any) => any;
   playwright: typeof import('@playwright/test');
+  _playwrightImpl: typeof import('@playwright/test');
 };
 
 export const testModeTest = test.extend<TestModeTestFixtures, TestModeWorkerOptions & TestModeWorkerFixtures>({
   mode: ['default', { scope: 'worker', option: true }],
-  playwright: [async ({ mode }, run) => {
+  _playwrightImpl: [async ({ mode }, run) => {
     const testMode = {
       'default': new DefaultTestMode(),
       'service': new DefaultTestMode(),

tests/playwright-test/playwright.artifacts.spec.ts

@@ -151,10 +151,8 @@ test('should work with screenshot: on', async ({ runInlineTest }, testInfo) => {
     '  test-finished-1.png',
     'artifacts-shared-shared-failing',
     '  test-failed-1.png',
-    '  test-failed-2.png',
     'artifacts-shared-shared-passing',
     '  test-finished-1.png',
-    '  test-finished-2.png',
     'artifacts-two-contexts',
     '  test-finished-1.png',
     '  test-finished-2.png',
@@ -185,7 +183,6 @@ test('should work with screenshot: only-on-failure', async ({ runInlineTest }, t
     '  test-failed-1.png',
     'artifacts-shared-shared-failing',
     '  test-failed-1.png',
-    '  test-failed-2.png',
     'artifacts-two-contexts-failing',
     '  test-failed-1.png',
     '  test-failed-2.png',

tests/playwright-test/playwright.trace.spec.ts

@@ -569,14 +569,19 @@ test('should opt out of attachments', async ({ runInlineTest, server }, testInfo
   expect([...trace.resources.keys()].filter(f => f.startsWith('resources/'))).toHaveLength(0);
 });
 
-test('should record with custom page fixture', async ({ runInlineTest }, testInfo) => {
+test('should record with custom page fixture that closes the context', async ({ runInlineTest }, testInfo) => {
+  // Note that original issue did not close the context, but we do not support such usecase.
+  test.info().annotations.push({ type: 'issue', description: 'https://github.com/microsoft/playwright/issues/23220' });
+
   const result = await runInlineTest({
     'a.spec.ts': `
       import { test as base, expect } from '@playwright/test';
 
       const test = base.extend({
         myPage: async ({ browser }, use) => {
-          await use(await browser.newPage());
+          const page = await browser.newPage();
+          await use(page);
+          await page.close();
         },
       });
 
@@ -1112,3 +1117,121 @@ test('trace:retain-on-first-failure should create trace if request context is di
   expect(trace.apiNames).toContain('apiRequestContext.get');
   expect(result.failed).toBe(1);
 });
+
+test('should record trace in workerStorageState', async ({ runInlineTest }) => {
+  test.info().annotations.push({ type: 'issue', description: 'https://github.com/microsoft/playwright/issues/30287' });
+
+  const result = await runInlineTest({
+    'a.spec.ts': `
+      import { test as base, expect } from '@playwright/test';
+      const test = base.extend({
+        storageState: ({ workerStorageState }, use) => use(workerStorageState),
+        workerStorageState: [async ({ browser }, use) => {
+          const page = await browser.newPage({ storageState: undefined });
+          await page.setContent('<div>hello</div>');
+          await page.close();
+          await use(undefined);
+        }, { scope: 'worker' }],
+      })
+      test('pass', async ({ page }) => {
+        await page.goto('data:text/html,<div>hi</div>');
+      });
+    `,
+  }, { trace: 'on' });
+  expect(result.exitCode).toBe(0);
+  expect(result.passed).toBe(1);
+
+  const tracePath = test.info().outputPath('test-results', 'a-pass', 'trace.zip');
+  const trace = await parseTrace(tracePath);
+  expect(trace.actionTree).toEqual([
+    'Before Hooks',
+    '  fixture: browser',
+    '    browserType.launch',
+    '  fixture: workerStorageState',
+    '    browser.newPage',
+    '    page.setContent',
+    '    page.close',
+    '  fixture: context',
+    '    browser.newContext',
+    '  fixture: page',
+    '    browserContext.newPage',
+    'page.goto',
+    'After Hooks',
+    '  fixture: page',
+    '  fixture: context',
+  ]);
+});
+
+test('should record trace after fixture teardown timeout', async ({ runInlineTest }) => {
+  test.info().annotations.push({ type: 'issue', description: 'https://github.com/microsoft/playwright/issues/30718' });
+
+  const result = await runInlineTest({
+    'a.spec.ts': `
+      import { test as base, expect } from '@playwright/test';
+      const test = base.extend({
+        fixture: async ({}, use) => {
+          await use('foo');
+          await new Promise(() => {});
+        },
+      })
+      test('fails', async ({ fixture, page }) => {
+        await page.evaluate(() => console.log('from the page'));
+      });
+    `,
+  }, { trace: 'on', timeout: '4000' });
+  expect(result.exitCode).toBe(1);
+  expect(result.failed).toBe(1);
+
+  const tracePath = test.info().outputPath('test-results', 'a-fails', 'trace.zip');
+  const trace = await parseTrace(tracePath);
+  expect(trace.actionTree).toEqual([
+    'Before Hooks',
+    '  fixture: fixture',
+    '  fixture: browser',
+    '    browserType.launch',
+    '  fixture: context',
+    '    browser.newContext',
+    '  fixture: page',
+    '    browserContext.newPage',
+    'page.evaluate',
+    'After Hooks',
+    '  fixture: page',
+    '  fixture: context',
+    '  fixture: fixture',
+    'Worker Cleanup',
+    '  fixture: browser',
+  ]);
+  // Check console events to make sure that library trace is recorded.
+  expect(trace.events).toContainEqual(expect.objectContaining({ type: 'console', text: 'from the page' }));
+});
+
+test('should take a screenshot-on-failure in workerStorageState', async ({ runInlineTest }) => {
+  test.info().annotations.push({ type: 'issue', description: 'https://github.com/microsoft/playwright/issues/30959' });
+
+  const result = await runInlineTest({
+    'playwright.config.ts': `
+      export default {
+        use: {
+          screenshot: 'only-on-failure',
+        },
+      };
+    `,
+    'a.spec.ts': `
+      import { test as base, expect } from '@playwright/test';
+      const test = base.extend({
+        storageState: ({ workerStorageState }, use) => use(workerStorageState),
+        workerStorageState: [async ({ browser }, use) => {
+          const page = await browser.newPage({ storageState: undefined });
+          await page.setContent('hello world!');
+          throw new Error('Failed!');
+          await use(undefined);
+        }, { scope: 'worker' }],
+      })
+      test('fail', async ({ page }) => {
+      });
+    `,
+  });
+  expect(result.exitCode).toBe(1);
+  expect(result.failed).toBe(1);
+  expect(fs.existsSync(test.info().outputPath('test-results', 'a-fail', 'test-failed-1.png'))).toBeTruthy();
+});