From 41836a4a7bfb0d7602839c081a4e3c4792e869df Mon Sep 17 00:00:00 2001
From: Jonathan Lipps <jlipps@gmail.com>
Date: Fri, 21 Jul 2023 12:28:31 -0700
Subject: [PATCH] fix(appium): ensure plugin commands reset newCommandTimeout

---
 packages/appium/lib/appium.js               | 32 +++++++++++++++++++--
 packages/appium/test/e2e/plugin.e2e.spec.js | 25 ++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/packages/appium/lib/appium.js b/packages/appium/lib/appium.js
index 740404059..360e919ae 100644
--- a/packages/appium/lib/appium.js
+++ b/packages/appium/lib/appium.js
@@ -231,8 +231,10 @@ class AppiumDriver extends DriverCore {
     jsonwpCaps = _.cloneDeep(jsonwpCaps);
     const jwpSettings = {...defaultSettings, ...pullSettings(jsonwpCaps)};
     w3cCapabilities = _.cloneDeep(w3cCapabilities);
-    if (!_.isPlainObject(w3cCapabilities)
-        || !(_.isArray(w3cCapabilities?.firstMatch) || _.isPlainObject(w3cCapabilities?.alwaysMatch))) {
+    if (
+      !_.isPlainObject(w3cCapabilities) ||
+      !(_.isArray(w3cCapabilities?.firstMatch) || _.isPlainObject(w3cCapabilities?.alwaysMatch))
+    ) {
       throw makeNonW3cCapsError();
     }
     // It is possible that the client only provides caps using JSONWP standard,
@@ -638,6 +640,16 @@ class AppiumDriver extends DriverCore {
     // get any plugins which are registered as handling this command
     const plugins = this.pluginsToHandleCmd(cmd, sessionId);
 
+    // if any plugins are going to handle this command, we can't guarantee that the default
+    // driver's executeCommand method will be called, which means we can't guarantee that the
+    // newCommandTimeout will be cleared. So we do it here as well.
+    if (plugins.length && dstSession) {
+      this.log.debug(
+        'Clearing new command timeout pre-emptively since plugin(s) will handle this command'
+      );
+      await dstSession.clearNewCommandTimeout();
+    }
+
     // now we define a 'cmdHandledBy' object which will keep track of which plugins have handled this
     // command. we care about this because (a) multiple plugins can handle the same command, and
     // (b) there's no guarantee that a plugin will actually call the next() method which runs the
@@ -703,6 +715,22 @@ class AppiumDriver extends DriverCore {
     // handling the command and which didn't
     this.logPluginHandlerReport(plugins, {cmd, cmdHandledBy});
 
+    // if we had plugins, and if they did not ultimately call the default handler, this means our
+    // new command timeout was not restarted by the default handler's executeCommand call, so
+    // restart it here using the same logic as in BaseDriver's executeCommand
+    if (
+      dstSession &&
+      !cmdHandledBy.default &&
+      dstSession.isCommandsQueueEnabled &&
+      cmd !== DELETE_SESSION_COMMAND
+    ) {
+      this.log.debug(
+        'Restarting new command timeout via umbrella driver since plugin did not ' +
+          'allow default handler to execute'
+      );
+      await dstSession.startNewCommandTimeout();
+    }
+
     // And finally, if the command was createSession, we want to migrate any plugins which were
     // previously sessionless to use the new sessionId, so that plugins can share state between
     // their createSession method and other instance methods
diff --git a/packages/appium/test/e2e/plugin.e2e.spec.js b/packages/appium/test/e2e/plugin.e2e.spec.js
index f4f33f97b..ff2bf9be8 100644
--- a/packages/appium/test/e2e/plugin.e2e.spec.js
+++ b/packages/appium/test/e2e/plugin.e2e.spec.js
@@ -275,6 +275,31 @@ describe('FakePlugin w/ FakeDriver via HTTP', function () {
         }
         shutdownErr.message.should.match(/either terminated or not started/);
       });
+
+      it('should allow plugin handled commands to reset newCommandTimeout', async function () {
+        /** @type {import('webdriverio').RemoteOptions} */
+        const newOpts = {...wdOpts};
+        newOpts.capabilities = {
+          ...(newOpts.capabilities ?? {}),
+          'appium:newCommandTimeout': 2,
+        };
+        const driver = await wdio(newOpts);
+        const {sessionId} = driver;
+        try {
+          const start = Date.now();
+          for (let i = 0; i < 5; i++) {
+            await B.delay(500);
+            await driver.getPageSource();
+          }
+          // prove that we went beyond the new command timeout as a result of sending commands
+          (Date.now() - start).should.be.above(2500);
+          await driver
+            .getPageSource()
+            .should.eventually.eql(`<Fake>${JSON.stringify([sessionId])}</Fake>`);
+        } finally {
+          await driver.deleteSession();
+        }
+      });
     });
   }
   describe('cli args handling for plugin args', function () {