chore: explicitly reset apiZone instead of everything (#34265)

Author: dgozman

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

@@ -138,7 +138,7 @@ export class Connection extends EventEmitter {
       this._localUtils?._channel.addStackToTracingNoReply({ callData: { stack: frames, id } }).catch(() => {});
     // We need to exit zones before calling into the server, otherwise
     // when we receive events from the server, we would be in an API zone.
-    zones.exitZones(() => this.onmessage({ ...message, metadata }));
+    zones.empty().run(() => this.onmessage({ ...message, metadata }));
     return await new Promise((resolve, reject) => this._callbacks.set(id, { resolve, reject, apiName, type, method }));
   }
 

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

@@ -820,7 +820,7 @@ export class RouteHandler {
     this._times = times;
     this.url = url;
     this.handler = handler;
-    this._svedZone = zones.currentZone();
+    this._svedZone = zones.current().without('apiZone');
   }
 
   static prepareInterceptionPatterns(handlers: RouteHandler[]) {

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

@@ -35,7 +35,7 @@ export class Waiter {
   constructor(channelOwner: ChannelOwner<channels.EventTargetChannel>, event: string) {
     this._waitId = createGuid();
     this._channelOwner = channelOwner;
-    this._savedZone = zones.currentZone();
+    this._savedZone = zones.current().without('apiZone');
 
     this._channelOwner._channel.waitForEventInfo({ info: { waitId: this._waitId, phase: 'before', event } }).catch(() => {});
     this._dispose = [

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

@@ -19,54 +19,54 @@ import { AsyncLocalStorage } from 'async_hooks';
 export type ZoneType = 'apiZone' | 'expectZone' | 'stepZone';
 
 class ZoneManager {
-  private readonly _asyncLocalStorage = new AsyncLocalStorage<Zone|undefined>();
+  private readonly _asyncLocalStorage = new AsyncLocalStorage<Zone | undefined>();
+  private readonly _emptyZone = Zone.createEmpty(this._asyncLocalStorage);
 
   run<T, R>(type: ZoneType, data: T, func: () => R): R {
-    const zone = Zone._createWithData(this._asyncLocalStorage, type, data);
-    return this._asyncLocalStorage.run(zone, func);
+    return this.current().with(type, data).run(func);
   }
 
   zoneData<T>(type: ZoneType): T | undefined {
-    const zone = this._asyncLocalStorage.getStore();
-    return zone?.get(type);
+    return this.current().data(type);
   }
 
-  currentZone(): Zone {
-    return this._asyncLocalStorage.getStore() ?? Zone._createEmpty(this._asyncLocalStorage);
+  current(): Zone {
+    return this._asyncLocalStorage.getStore() ?? this._emptyZone;
   }
 
-  exitZones<R>(func: () => R): R {
-    return this._asyncLocalStorage.run(undefined, func);
+  empty(): Zone {
+    return this._emptyZone;
   }
 }
 
 export class Zone {
   private readonly _asyncLocalStorage: AsyncLocalStorage<Zone | undefined>;
-  private readonly _data: Map<ZoneType, unknown>;
+  private readonly _data: ReadonlyMap<ZoneType, unknown>;
 
-  static _createWithData(asyncLocalStorage: AsyncLocalStorage<Zone|undefined>, type: ZoneType, data: unknown) {
-    const store = new Map(asyncLocalStorage.getStore()?._data);
-    store.set(type, data);
-    return new Zone(asyncLocalStorage, store);
-  }
-
-  static _createEmpty(asyncLocalStorage: AsyncLocalStorage<Zone|undefined>) {
+  static createEmpty(asyncLocalStorage: AsyncLocalStorage<Zone | undefined>) {
     return new Zone(asyncLocalStorage, new Map());
   }
 
-  private constructor(asyncLocalStorage: AsyncLocalStorage<Zone|undefined>, store: Map<ZoneType, unknown>) {
+  private constructor(asyncLocalStorage: AsyncLocalStorage<Zone | undefined>, store: Map<ZoneType, unknown>) {
     this._asyncLocalStorage = asyncLocalStorage;
     this._data = store;
   }
 
+  with(type: ZoneType, data: unknown): Zone {
+    return new Zone(this._asyncLocalStorage, new Map(this._data).set(type, data));
+  }
+
+  without(type?: ZoneType): Zone {
+    const data = type ? new Map(this._data) : new Map();
+    data.delete(type);
+    return new Zone(this._asyncLocalStorage, data);
+  }
+
   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);
+    return this._asyncLocalStorage.run(this, func);
   }
 
-  get<T>(type: ZoneType): T | undefined {
+  data<T>(type: ZoneType): T | undefined {
     return this._data.get(type) as T | undefined;
   }
 }

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

@@ -1376,8 +1376,9 @@ test('calls from waitForEvent callback should be under its parent step', {
         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();
+            await page.content();
+            await page.content();  // second time a charm!
+            await expect(page.locator('div')).toContainText('Go');
             return true;
           });
 
@@ -1405,9 +1406,11 @@ 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
-expect    |  expect.toBeTruthy @ a.test.ts:9
-expect    |expect.toBe @ a.test.ts:17
+pw:api    |  page.click(div) @ a.test.ts:14
+pw:api    |  page.content @ a.test.ts:8
+pw:api    |  page.content @ a.test.ts:9
+expect    |  expect.toContainText @ a.test.ts:10
+expect    |expect.toBe @ a.test.ts:18
 hook      |After Hooks
 fixture   |  fixture: page
 fixture   |  fixture: context
@@ -1464,7 +1467,8 @@ test('calls from page.route callback should be under its parent step', {
             const response = await route.fetch();
             const text = await response.text();
             expect(text).toBe('');
-            await route.fulfill({ response })
+            await response.text();  // second time a charm!
+            await route.fulfill({ response });
           });
           await page.goto('${server.EMPTY_PAGE}');
         });
@@ -1485,9 +1489,10 @@ 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    |  page.goto(${server.EMPTY_PAGE}) @ a.test.ts:12
 pw:api    |  apiResponse.text @ a.test.ts:7
 expect    |  expect.toBe @ a.test.ts:8
+pw:api    |  apiResponse.text @ a.test.ts:9
 hook      |After Hooks
 fixture   |  fixture: page
 fixture   |  fixture: context