fix: parent step for API calls inside waitForEvent callback (#33409)

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

@@ -22,14 +22,14 @@ import { Worker } from './worker';
 import type { Headers, RemoteAddr, SecurityDetails, WaitForEventOptions } from './types';
 import fs from 'fs';
 import { mime } from '../utilsBundle';
-import { assert, isString, headersObjectToArray, isRegExp, rewriteErrorMessage } from '../utils';
+import { assert, isString, headersObjectToArray, isRegExp, rewriteErrorMessage, MultiMap, urlMatches, zones } from '../utils';
+import type { URLMatch, Zone } from '../utils';
 import { ManualPromise, LongStandingScope } from '../utils/manualPromise';
 import { Events } from './events';
 import type { Page } from './page';
 import { Waiter } from './waiter';
 import type * as api from '../../types/types';
 import type { HeadersArray } from '../common/types';
-import { MultiMap, urlMatches, type URLMatch } from '../utils';
 import { APIResponse } from './fetch';
 import type { Serializable } from '../../types/structs';
 import type { BrowserContext } from './browserContext';
@@ -811,12 +811,14 @@ export class RouteHandler {
   readonly handler: RouteHandlerCallback;
   private _ignoreException: boolean = false;
   private _activeInvocations: Set<{ complete: Promise<void>, route: Route }> = new Set();
+  private _svedZone: Zone;
 
   constructor(baseURL: string | undefined, url: URLMatch, handler: RouteHandlerCallback, times: number = Number.MAX_SAFE_INTEGER) {
     this._baseURL = baseURL;
     this._times = times;
     this.url = url;
     this.handler = handler;
+    this._svedZone = zones.currentZone();
   }
 
   static prepareInterceptionPatterns(handlers: RouteHandler[]) {
@@ -840,6 +842,10 @@ export class RouteHandler {
   }
 
   public async handle(route: Route): Promise<boolean> {
+    return await this._svedZone.run(async () => this._handleImpl(route));
+  }
+
+  private async _handleImpl(route: Route): Promise<boolean> {
     const handlerInvocation = { complete: new ManualPromise(), route } ;
     this._activeInvocations.add(handlerInvocation);
     try {

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

@@ -17,7 +17,8 @@
 import type { EventEmitter } from 'events';
 import { rewriteErrorMessage } from '../utils/stackTrace';
 import { TimeoutError } from './errors';
-import { createGuid } from '../utils';
+import { createGuid, zones } from '../utils';
+import type { Zone } from '../utils';
 import type * as channels from '@protocol/channels';
 import type { ChannelOwner } from './channelOwner';
 
@@ -29,10 +30,13 @@ export class Waiter {
   private _channelOwner: ChannelOwner<channels.EventTargetChannel>;
   private _waitId: string;
   private _error: string | undefined;
+  private _savedZone: Zone;
 
   constructor(channelOwner: ChannelOwner<channels.EventTargetChannel>, event: string) {
     this._waitId = createGuid();
     this._channelOwner = channelOwner;
+    this._savedZone = zones.currentZone();
+
     this._channelOwner._channel.waitForEventInfo({ info: { waitId: this._waitId, phase: 'before', event } }).catch(() => {});
     this._dispose = [
       () => this._channelOwner._wrapApiCall(async () => {
@@ -46,12 +50,12 @@ export class Waiter {
   }
 
   async waitForEvent<T = void>(emitter: EventEmitter, event: string, predicate?: (arg: T) => boolean | Promise<boolean>): Promise<T> {
-    const { promise, dispose } = waitForEvent(emitter, event, predicate);
+    const { promise, dispose } = waitForEvent(emitter, event, this._savedZone, predicate);
     return await this.waitForPromise(promise, dispose);
   }
 
   rejectOnEvent<T = void>(emitter: EventEmitter, event: string, error: Error | (() => Error), predicate?: (arg: T) => boolean | Promise<boolean>) {
-    const { promise, dispose } = waitForEvent(emitter, event, predicate);
+    const { promise, dispose } = waitForEvent(emitter, event, this._savedZone, predicate);
     this._rejectOn(promise.then(() => { throw (typeof error === 'function' ? error() : error); }), dispose);
   }
 
@@ -103,19 +107,21 @@ export class Waiter {
   }
 }
 
-function waitForEvent<T = void>(emitter: EventEmitter, event: string, predicate?: (arg: T) => boolean | Promise<boolean>): { promise: Promise<T>, dispose: () => void } {
+function waitForEvent<T = void>(emitter: EventEmitter, event: string, savedZone: Zone, predicate?: (arg: T) => boolean | Promise<boolean>): { promise: Promise<T>, dispose: () => void } {
   let listener: (eventArg: any) => void;
   const promise = new Promise<T>((resolve, reject) => {
     listener = async (eventArg: any) => {
-      try {
+      await savedZone.run(async () => {
-        if (predicate && !(await predicate(eventArg)))
+        try {
-          return;
+          if (predicate && !(await predicate(eventArg)))
-        emitter.removeListener(event, listener);
+            return;
-        resolve(eventArg);
+          emitter.removeListener(event, listener);
-      } catch (e) {
+          resolve(eventArg);
-        emitter.removeListener(event, listener);
+        } catch (e) {
-        reject(e);
+          emitter.removeListener(event, listener);
-      }
+          reject(e);
+        }
+      });
     };
     emitter.addListener(event, listener);
   });

packages/playwright-core/src/utils/zones.ts

@@ -19,49 +19,55 @@ import { AsyncLocalStorage } from 'async_hooks';
 export type ZoneType = 'apiZone' | 'expectZone' | 'stepZone';
 
 class ZoneManager {
-  private readonly _asyncLocalStorage = new AsyncLocalStorage<Zone<unknown>|undefined>();
+  private readonly _asyncLocalStorage = new AsyncLocalStorage<Zone|undefined>();
 
   run<T, R>(type: ZoneType, data: T, func: () => R): R {
-    const previous = this._asyncLocalStorage.getStore();
+    const zone = Zone._createWithData(this._asyncLocalStorage, type, data);
-    const zone = new Zone(previous, type, data);
     return this._asyncLocalStorage.run(zone, func);
   }
 
   zoneData<T>(type: ZoneType): T | undefined {
-    for (let zone = this._asyncLocalStorage.getStore(); zone; zone = zone.previous) {
+    const zone = this._asyncLocalStorage.getStore();
-      if (zone.type === type)
+    return zone?.get(type);
-        return zone.data as T;
+  }
-    }
+
-    return undefined;
+  currentZone(): Zone {
+    return this._asyncLocalStorage.getStore() ?? Zone._createEmpty(this._asyncLocalStorage);
   }
 
   exitZones<R>(func: () => R): R {
     return this._asyncLocalStorage.run(undefined, func);
   }
-
-  printZones() {
-    const zones = [];
-    for (let zone = this._asyncLocalStorage.getStore(); zone; zone = zone.previous) {
-      let str = zone.type;
-      if (zone.type === 'apiZone')
-        str += `(${(zone.data as any).apiName})`;
-      zones.push(str);
-
-    }
-    // eslint-disable-next-line no-console
-    console.log('zones: ', zones.join(' -> '));
-  }
 }
 
-class Zone<T> {
+export class Zone {
-  readonly type: ZoneType;
+  private readonly _asyncLocalStorage: AsyncLocalStorage<Zone | undefined>;
-  readonly data: T;
+  private readonly _data: Map<ZoneType, unknown>;
-  readonly previous: Zone<unknown> | undefined;
 
-  constructor(previous: Zone<unknown> | undefined, type: ZoneType, data: T) {
+  static _createWithData(asyncLocalStorage: AsyncLocalStorage<Zone|undefined>, type: ZoneType, data: unknown) {
-    this.type = type;
+    const store = new Map(asyncLocalStorage.getStore()?._data);
-    this.data = data;
+    store.set(type, data);
-    this.previous = previous;
+    return new Zone(asyncLocalStorage, store);
+  }
+
+  static _createEmpty(asyncLocalStorage: AsyncLocalStorage<Zone|undefined>) {
+    return new Zone(asyncLocalStorage, new Map());
+  }
+
+  private constructor(asyncLocalStorage: AsyncLocalStorage<Zone|undefined>, store: Map<ZoneType, unknown>) {
+    this._asyncLocalStorage = asyncLocalStorage;
+    this._data = store;
+  }
+
+  run<R>(func: () => R): R {
+    // Reset apiZone and expectZone, but restore stepZone.
+    const entries = [...this._data.entries()].filter(([type]) => (type !== 'apiZone' && type !== 'expectZone'));
+    const resetZone = new Zone(this._asyncLocalStorage, new Map(entries));
+    return this._asyncLocalStorage.run(resetZone, func);
+  }
+
+  get<T>(type: ZoneType): T | undefined {
+    return this._data.get(type) as T | undefined;
   }
 }
 

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

@@ -1310,3 +1310,101 @@ fixture   |  fixture: page
 fixture   |  fixture: context
 `);
 });
+
+test('calls from waitForEvent callback should be under its parent step', {
+  annotation: { type: 'issue', description: 'https://github.com/microsoft/playwright/issues/33186' }
+}, async ({ runInlineTest, server }) => {
+  const result = await runInlineTest({
+    'reporter.ts': stepIndentReporter,
+    'playwright.config.ts': `module.exports = { reporter: './reporter' };`,
+    'a.test.ts': `
+      import { test, expect } from '@playwright/test';
+      test('waitForResponse step nesting', async ({ page }) => {
+        await page.goto('${server.EMPTY_PAGE}');
+        await page.setContent('<div onclick="fetch(\\'/simple.json\\').then(r => r.text());">Go!</div>');
+        const responseJson = await test.step('custom step', async () => {
+          const responsePromise = page.waitForResponse(async response => {
+            const text = await response.text();
+            expect(text).toBeTruthy();
+            return true;
+          });
+
+          await page.click('div');
+          const response = await responsePromise;
+          return await response.text();
+        });
+        expect(responseJson).toBe('{"foo": "bar"}\\n');
+      });
+      `
+  }, { reporter: '', workers: 1, timeout: 3000 });
+
+  expect(result.exitCode).toBe(0);
+  expect(result.passed).toBe(0);
+  expect(result.output).not.toContain('Internal error');
+  expect(stripAnsi(result.output)).toBe(`
+hook      |Before Hooks
+fixture   |  fixture: browser
+pw:api    |    browserType.launch
+fixture   |  fixture: context
+pw:api    |    browser.newContext
+fixture   |  fixture: page
+pw:api    |    browserContext.newPage
+pw:api    |page.goto(${server.EMPTY_PAGE}) @ a.test.ts:4
+pw:api    |page.setContent @ a.test.ts:5
+test.step |custom step @ a.test.ts:6
+pw:api    |  page.waitForResponse @ a.test.ts:7
+pw:api    |  page.click(div) @ a.test.ts:13
+pw:api    |  response.text @ a.test.ts:8
+expect    |  expect.toBeTruthy @ a.test.ts:9
+pw:api    |  response.text @ a.test.ts:15
+expect    |expect.toBe @ a.test.ts:17
+hook      |After Hooks
+fixture   |  fixture: page
+fixture   |  fixture: context
+`);
+});
+
+test('calls from page.route callback should be under its parent step', {
+  annotation: { type: 'issue', description: 'https://github.com/microsoft/playwright/issues/33186' }
+}, async ({ runInlineTest, server }) => {
+  const result = await runInlineTest({
+    'reporter.ts': stepIndentReporter,
+    'playwright.config.ts': `module.exports = { reporter: './reporter' };`,
+    'a.test.ts': `
+      import { test, expect } from '@playwright/test';
+      test('waitForResponse step nesting', async ({ page }) => {
+        await test.step('custom step', async () => {
+          await page.route('**/empty.html', async route => {
+            const response = await route.fetch();
+            const text = await response.text();
+            expect(text).toBe('');
+            await route.fulfill({ response })
+          });
+          await page.goto('${server.EMPTY_PAGE}');
+        });
+      });
+      `
+  }, { reporter: '', workers: 1, timeout: 3000 });
+
+  expect(result.exitCode).toBe(0);
+  expect(result.passed).toBe(0);
+  expect(result.output).not.toContain('Internal error');
+  expect(stripAnsi(result.output)).toBe(`
+hook      |Before Hooks
+fixture   |  fixture: browser
+pw:api    |    browserType.launch
+fixture   |  fixture: context
+pw:api    |    browser.newContext
+fixture   |  fixture: page
+pw:api    |    browserContext.newPage
+test.step |custom step @ a.test.ts:4
+pw:api    |  page.route @ a.test.ts:5
+pw:api    |  page.goto(${server.EMPTY_PAGE}) @ a.test.ts:11
+pw:api    |  apiResponse.text @ a.test.ts:7
+expect    |  expect.toBe @ a.test.ts:8
+hook      |After Hooks
+fixture   |  fixture: page
+fixture   |  fixture: context
+`);
+});
+