fix: parent step for API calls inside waitForEvent callback

Fixes https://github.com/microsoft/playwright/issues/33186

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

@@ -22,7 +22,7 @@ import type * as api from '../../types/types';
 import { serializeError, isTargetClosedError, TargetClosedError } from './errors';
 import { TimeoutSettings } from '../common/timeoutSettings';
 import type * as channels from '@protocol/channels';
-import { assert, headersObjectToArray, isObject, isRegExp, isString, LongStandingScope, urlMatches, urlMatchesEqual, mkdirIfNeeded, trimStringWithEllipsis, type URLMatch } from '../utils';
+import { assert, headersObjectToArray, isObject, isRegExp, isString, LongStandingScope, urlMatches, urlMatchesEqual, mkdirIfNeeded, trimStringWithEllipsis, type URLMatch, zones } from '../utils';
 import { Accessibility } from './accessibility';
 import { Artifact } from './artifact';
 import type { BrowserContext } from './browserContext';

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

@@ -17,7 +17,7 @@
 import type { EventEmitter } from 'events';
 import { rewriteErrorMessage } from '../utils/stackTrace';
 import { TimeoutError } from './errors';
-import { createGuid } from '../utils';
+import { createGuid, ZoneReference, zones } from '../utils';
 import type * as channels from '@protocol/channels';
 import type { ChannelOwner } from './channelOwner';
 
@@ -29,10 +29,14 @@ export class Waiter {
   private _channelOwner: ChannelOwner<channels.EventTargetChannel>;
   private _waitId: string;
   private _error: string | undefined;
+  private _savedZone: ZoneReference;
 
   constructor(channelOwner: ChannelOwner<channels.EventTargetChannel>, event: string) {
     this._waitId = createGuid();
     this._channelOwner = channelOwner;
+    // Save current chain before the wait for event API call (we don't nest API calls)
+    // and restore it later to find proper parent for the event listener.
+    this._savedZone = zones.currentZoneBefore('apiZone');
     this._channelOwner._channel.waitForEventInfo({ info: { waitId: this._waitId, phase: 'before', event } }).catch(() => {});
     this._dispose = [
       () => this._channelOwner._wrapApiCall(async () => {
@@ -41,17 +45,17 @@ export class Waiter {
     ];
   }
 
-  static createForEvent(channelOwner: ChannelOwner<channels.EventTargetChannel>, event: string) {
+  static createForEvent(channelOwner: ChannelOwner<channels.EventTargetChannel>, event: string): Waiter {
     return new Waiter(channelOwner, event);
   }
 
   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: ZoneReference, 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 {
+      savedZone.runInZone(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

@@ -39,6 +39,17 @@ class ZoneManager {
     return this._asyncLocalStorage.run(undefined, func);
   }
 
+  currentZoneBefore(type: ZoneType): ZoneReference {
+    let zone = this._asyncLocalStorage.getStore();
+    while (zone && zone.type === type)
+      zone = zone.previous;
+    return new ZoneReference(zone, this._asyncLocalStorage);
+  }
+
+  currentZone(): ZoneReference {
+    return new ZoneReference(this._asyncLocalStorage.getStore(), this._asyncLocalStorage);
+  }
+
   printZones() {
     const zones = [];
     for (let zone = this._asyncLocalStorage.getStore(); zone; zone = zone.previous) {
@@ -53,6 +64,20 @@ class ZoneManager {
   }
 }
 
+export class ZoneReference {
+  private _zone: Zone<unknown> | undefined;
+  private _asyncLocalStorage: AsyncLocalStorage<Zone<unknown>|undefined>;
+
+  constructor(zone: Zone<unknown> | undefined,  asyncLocalStorage: AsyncLocalStorage<Zone<unknown>|undefined>) {
+    this._zone = zone;
+    this._asyncLocalStorage = asyncLocalStorage;
+  }
+
+  runInZone<R>(func: () => R): R {
+    return this._asyncLocalStorage.run(this._zone, func);
+  }
+}
+
 class Zone<T> {
   readonly type: ZoneType;
   readonly data: T;

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

@@ -1310,3 +1310,56 @@ 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
+`);
+});