From 44f6489a523f6dc917ab10b031f5ea355059e8ca Mon Sep 17 00:00:00 2001
From: Oj <oj@oojmed.com>
Date: Thu, 5 Jan 2023 18:41:45 +0000
Subject: [PATCH] inject: add window.loaded promise

---
 src/launcher/inject.js | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/launcher/inject.js b/src/launcher/inject.js
index bb2ade3..03bf4bd 100644
--- a/src/launcher/inject.js
+++ b/src/launcher/inject.js
@@ -1,10 +1,12 @@
 import IPCApi from '../lib/ipc.js';
 
 export default async (CDP, proc, injectionType = 'browser', { browserName } = { browserName: 'unknown' }) => {
-  let pageLoadCallback = () => {}, onWindowMessage = () => {};
+  let pageLoadCallback, pageLoadPromise = new Promise(res => pageLoadCallback = res);
+  let frameLoadCallback = () => {}, onWindowMessage = () => {};
   CDP.onMessage(msg => {
     if (msg.method === 'Runtime.bindingCalled' && msg.params.name === '_gluonSend') onWindowMessage(JSON.parse(msg.params.payload));
-    if (msg.method === 'Page.frameStoppedLoading') pageLoadCallback(msg.params);
+    if (msg.method === 'Page.frameStoppedLoading') frameLoadCallback(msg.params);
+    if (msg.method === 'Page.loadEventFired') pageLoadCallback();
     if (msg.method === 'Runtime.executionContextCreated') injectIPC(); // ensure IPC injection again
   });
 
@@ -42,13 +44,17 @@ export default async (CDP, proc, injectionType = 'browser', { browserName } = {
   }, {
     evalInWindow,
     evalOnNewDocument: source => CDP.sendMessage('Page.addScriptToEvaluateOnNewDocument', { source }, sessionId),
-    pageLoadPromise: new Promise(res => pageLoadCallback = res)
+    pageLoadPromise: new Promise(res => frameLoadCallback = res)
   });
 
   onWindowMessage = ipcMessageCallback;
 
   log('finished setup');
 
+  evalInWindow('document.readyState').then(readyState => {
+    if (readyState === 'complete' || readyState === 'ready') pageLoadCallback();
+  });
+
   const generateVersionInfo = (name, version) => ({
     name,
     version,
@@ -58,6 +64,7 @@ export default async (CDP, proc, injectionType = 'browser', { browserName } = {
   return {
     window: {
       eval: evalInWindow,
+      loaded: pageLoadPromise
     },
 
     ipc: IPC,