From 4674f21a708c8221b9cdb87d80c964eb451c6fb1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mateusz=20Burzy=C5=84ski?= <mateuszburzynski@gmail.com>
Date: Tue, 14 May 2024 15:13:43 +0200
Subject: [PATCH] Await `onApiCallEnd`/`onApiCallBegin` hooks

---
 packages/playwright-core/src/client/channelOwner.ts       | 7 ++++---
 .../playwright-core/src/client/clientInstrumentation.ts   | 8 ++++----
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/packages/playwright-core/src/client/channelOwner.ts b/packages/playwright-core/src/client/channelOwner.ts
index 8670ebbea2..04af8c5aea 100644
--- a/packages/playwright-core/src/client/channelOwner.ts
+++ b/packages/playwright-core/src/client/channelOwner.ts
@@ -146,7 +146,8 @@ export abstract class ChannelOwner<T extends channels.Channel = channels.Channel
                 const { apiName, frames, csi, callCookie, wallTime } = apiZone.reported ? { apiName: undefined, csi: undefined, callCookie: undefined, frames: [], wallTime: undefined } : apiZone;
                 apiZone.reported = true;
                 if (csi && apiName)
-                  csi.onApiCallBegin(apiName, params, frames, wallTime, callCookie);
+                  // this is meant to be awaited so all async hooks can be completed before the actual command gets dispatched
+                  await csi.onApiCallBegin(apiName, params, frames, wallTime, callCookie);
                 return await this._connection.sendMessageToServer(this, prop, validator(params, '', { tChannelImpl: tChannelImplToWire, binary: this._connection.rawBuffers() ? 'buffer' : 'toBase64' }), apiName, frames, wallTime);
               });
             };
@@ -188,7 +189,7 @@ export abstract class ChannelOwner<T extends channels.Channel = channels.Channel
       logApiCall(logger, `=> ${apiName} started`, isInternal);
       const apiZone: ApiZone = { apiName, frames, isInternal, reported: false, csi, callCookie, wallTime };
       const result = await zones.run('apiZone', apiZone, async () => await func(apiZone));
-      csi?.onApiCallEnd(callCookie);
+      await csi?.onApiCallEnd(callCookie);
       logApiCall(logger, `<= ${apiName} succeeded`, isInternal);
       return result;
     } catch (e) {
@@ -200,7 +201,7 @@ export abstract class ChannelOwner<T extends channels.Channel = channels.Channel
         e.stack = e.message + stackFrames;
       else
         e.stack = '';
-      csi?.onApiCallEnd(callCookie, e);
+      await csi?.onApiCallEnd(callCookie, e);
       logApiCall(logger, `<= ${apiName} failed`, isInternal);
       throw e;
     }
diff --git a/packages/playwright-core/src/client/clientInstrumentation.ts b/packages/playwright-core/src/client/clientInstrumentation.ts
index 2b5dd00970..c34715d670 100644
--- a/packages/playwright-core/src/client/clientInstrumentation.ts
+++ b/packages/playwright-core/src/client/clientInstrumentation.ts
@@ -22,8 +22,8 @@ export interface ClientInstrumentation {
   addListener(listener: ClientInstrumentationListener): void;
   removeListener(listener: ClientInstrumentationListener): void;
   removeAllListeners(): void;
-  onApiCallBegin(apiCall: string, params: Record<string, any>, frames: StackFrame[], wallTime: number, userData: any): void;
-  onApiCallEnd(userData: any, error?: Error): void;
+  onApiCallBegin(apiCall: string, params: Record<string, any>, frames: StackFrame[], wallTime: number, userData: any): Promise<void> | void;
+  onApiCallEnd(userData: any, error?: Error): Promise<void> | void;
   onDidCreateBrowserContext(context: BrowserContext): Promise<void>;
   onDidCreateRequestContext(context: APIRequestContext): Promise<void>;
   onWillPause(): void;
@@ -32,8 +32,8 @@ export interface ClientInstrumentation {
 }
 
 export interface ClientInstrumentationListener {
-  onApiCallBegin?(apiName: string, params: Record<string, any>, frames: StackFrame[], wallTime: number, userData: any): void;
-  onApiCallEnd?(userData: any, error?: Error): void;
+  onApiCallBegin?(apiName: string, params: Record<string, any>, frames: StackFrame[], wallTime: number, userData: any): Promise<void> | void;
+  onApiCallEnd?(userData: any, error?: Error): Promise<void> | void;
   onDidCreateBrowserContext?(context: BrowserContext): Promise<void>;
   onDidCreateRequestContext?(context: APIRequestContext): Promise<void>;
   onWillPause?(): void;